datafusion_datasource_csv/

source.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Execution plan for reading CSV files

use datafusion_datasource::boundary_stream::AlignedBoundaryStream;
use datafusion_datasource::projection::{ProjectionOpener, SplitProjection};
use datafusion_physical_plan::projection::ProjectionExprs;
use std::fmt;
use std::io::Read;
use std::sync::Arc;

use datafusion_datasource::decoder::{DecoderDeserializer, deserialize_stream};
use datafusion_datasource::file_compression_type::FileCompressionType;
use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};
use datafusion_datasource::{
    FileRange, ListingTableUrl, PartitionedFile, TableSchema, as_file_source,
};

use arrow::csv;
use datafusion_common::config::CsvOptions;
use datafusion_common::tree_node::TreeNodeRecursion;
use datafusion_common::{DataFusionError, Result, exec_datafusion_err};
use datafusion_common_runtime::JoinSet;
use datafusion_datasource::file::FileSource;
use datafusion_datasource::file_scan_config::FileScanConfig;
use datafusion_execution::TaskContext;
use datafusion_physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet};
use datafusion_physical_plan::{
    DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
};

use crate::file_format::CsvDecoder;
use futures::{StreamExt, TryStreamExt};
use object_store::buffered::BufWriter;
use object_store::{GetOptions, GetResultPayload, ObjectStore};
use tokio::io::AsyncWriteExt;

/// A Config for [`CsvOpener`]
///
/// # Example: create a `DataSourceExec` for CSV
/// ```
/// # use std::sync::Arc;
/// # use arrow::datatypes::Schema;
/// # use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
/// # use datafusion_datasource::PartitionedFile;
/// # use datafusion_datasource_csv::source::CsvSource;
/// # use datafusion_execution::object_store::ObjectStoreUrl;
/// # use datafusion_datasource::source::DataSourceExec;
/// # use datafusion_common::config::CsvOptions;
///
/// # let object_store_url = ObjectStoreUrl::local_filesystem();
/// # let file_schema = Arc::new(Schema::empty());
///
/// let options = CsvOptions {
///     has_header: Some(true),
///     delimiter: b',',
///     quote: b'"',
///     newlines_in_values: Some(true), // The file contains newlines in values
///     ..Default::default()
/// };
/// let source = Arc::new(CsvSource::new(file_schema.clone())
///     .with_csv_options(options)
///     .with_terminator(Some(b'#'))
/// );
/// // Create a DataSourceExec for reading the first 100MB of `file1.csv`
/// let config = FileScanConfigBuilder::new(object_store_url, source)
///     .with_file(PartitionedFile::new("file1.csv", 100*1024*1024))
///     .build();
/// let exec = (DataSourceExec::from_data_source(config));
/// ```
#[derive(Debug, Clone)]
pub struct CsvSource {
    options: CsvOptions,
    batch_size: Option<usize>,
    table_schema: TableSchema,
    projection: SplitProjection,
    metrics: ExecutionPlanMetricsSet,
}

impl CsvSource {
    /// Returns a [`CsvSource`]
    pub fn new(table_schema: impl Into<TableSchema>) -> Self {
        let table_schema = table_schema.into();
        Self {
            options: CsvOptions::default(),
            projection: SplitProjection::unprojected(&table_schema),
            table_schema,
            batch_size: None,
            metrics: ExecutionPlanMetricsSet::new(),
        }
    }

    /// Sets the CSV options
    pub fn with_csv_options(mut self, options: CsvOptions) -> Self {
        self.options = options;
        self
    }

    /// true if the first line of each file is a header
    pub fn has_header(&self) -> bool {
        self.options.has_header.unwrap_or(true)
    }

    // true if rows length support truncate
    pub fn truncate_rows(&self) -> bool {
        self.options.truncated_rows.unwrap_or(false)
    }
    /// A column delimiter
    pub fn delimiter(&self) -> u8 {
        self.options.delimiter
    }

    /// The quote character
    pub fn quote(&self) -> u8 {
        self.options.quote
    }

    /// The line terminator
    pub fn terminator(&self) -> Option<u8> {
        self.options.terminator
    }

    /// Lines beginning with this byte are ignored.
    pub fn comment(&self) -> Option<u8> {
        self.options.comment
    }

    /// The escape character
    pub fn escape(&self) -> Option<u8> {
        self.options.escape
    }

    /// Initialize a CsvSource with escape
    pub fn with_escape(&self, escape: Option<u8>) -> Self {
        let mut conf = self.clone();
        conf.options.escape = escape;
        conf
    }

    /// Initialize a CsvSource with terminator
    pub fn with_terminator(&self, terminator: Option<u8>) -> Self {
        let mut conf = self.clone();
        conf.options.terminator = terminator;
        conf
    }

    /// Initialize a CsvSource with comment
    pub fn with_comment(&self, comment: Option<u8>) -> Self {
        let mut conf = self.clone();
        conf.options.comment = comment;
        conf
    }

    /// Whether to support truncate rows when read csv file
    pub fn with_truncate_rows(&self, truncate_rows: bool) -> Self {
        let mut conf = self.clone();
        conf.options.truncated_rows = Some(truncate_rows);
        conf
    }

    /// Whether values may contain newline characters
    pub fn newlines_in_values(&self) -> bool {
        self.options.newlines_in_values.unwrap_or(false)
    }
}

impl CsvSource {
    fn open<R: Read>(&self, reader: R) -> Result<csv::Reader<R>> {
        Ok(self.builder().build(reader)?)
    }

    fn builder(&self) -> csv::ReaderBuilder {
        let mut builder =
            csv::ReaderBuilder::new(Arc::clone(self.table_schema.file_schema()))
                .with_delimiter(self.delimiter())
                .with_batch_size(
                    self.batch_size
                        .expect("Batch size must be set before initializing builder"),
                )
                .with_header(self.has_header())
                .with_quote(self.quote())
                .with_truncated_rows(self.truncate_rows());
        if let Some(terminator) = self.terminator() {
            builder = builder.with_terminator(terminator);
        }
        builder = builder.with_projection(self.projection.file_indices.clone());
        if let Some(escape) = self.escape() {
            builder = builder.with_escape(escape)
        }
        if let Some(comment) = self.comment() {
            builder = builder.with_comment(comment);
        }

        builder
    }
}

/// A [`FileOpener`] that opens a CSV file and yields a [`FileOpenFuture`]
pub struct CsvOpener {
    config: Arc<CsvSource>,
    file_compression_type: FileCompressionType,
    object_store: Arc<dyn ObjectStore>,
    partition_index: usize,
}

impl CsvOpener {
    /// Returns a [`CsvOpener`]
    pub fn new(
        config: Arc<CsvSource>,
        file_compression_type: FileCompressionType,
        object_store: Arc<dyn ObjectStore>,
    ) -> Self {
        Self {
            config,
            file_compression_type,
            object_store,
            partition_index: 0,
        }
    }
}

impl From<CsvSource> for Arc<dyn FileSource> {
    fn from(source: CsvSource) -> Self {
        as_file_source(source)
    }
}

impl FileSource for CsvSource {
    fn create_file_opener(
        &self,
        object_store: Arc<dyn ObjectStore>,
        base_config: &FileScanConfig,
        partition_index: usize,
    ) -> Result<Arc<dyn FileOpener>> {
        let mut opener = Arc::new(CsvOpener {
            config: Arc::new(self.clone()),
            file_compression_type: base_config.file_compression_type,
            object_store,
            partition_index,
        }) as Arc<dyn FileOpener>;
        opener = ProjectionOpener::try_new(
            self.projection.clone(),
            Arc::clone(&opener),
            self.table_schema.file_schema(),
        )?;
        Ok(opener)
    }

    fn table_schema(&self) -> &TableSchema {
        &self.table_schema
    }

    fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource> {
        let mut conf = self.clone();
        conf.batch_size = Some(batch_size);
        Arc::new(conf)
    }

    fn try_pushdown_projection(
        &self,
        projection: &ProjectionExprs,
    ) -> Result<Option<Arc<dyn FileSource>>> {
        let mut source = self.clone();
        let new_projection = self.projection.source.try_merge(projection)?;
        let split_projection =
            SplitProjection::new(self.table_schema.file_schema(), &new_projection);
        source.projection = split_projection;
        Ok(Some(Arc::new(source)))
    }

    fn projection(&self) -> Option<&ProjectionExprs> {
        Some(&self.projection.source)
    }

    fn metrics(&self) -> &ExecutionPlanMetricsSet {
        &self.metrics
    }

    fn file_type(&self) -> &str {
        "csv"
    }

    fn supports_repartitioning(&self) -> bool {
        // Cannot repartition if values may contain newlines, as record
        // boundaries cannot be determined by byte offset alone
        !self.options.newlines_in_values.unwrap_or(false)
    }

    fn fmt_extra(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
        match t {
            DisplayFormatType::Default | DisplayFormatType::Verbose => {
                write!(f, ", has_header={}", self.has_header())
            }
            DisplayFormatType::TreeRender => Ok(()),
        }
    }

    fn apply_expressions(
        &self,
        f: &mut dyn FnMut(
            &Arc<dyn datafusion_physical_plan::PhysicalExpr>,
        ) -> Result<TreeNodeRecursion>,
    ) -> Result<TreeNodeRecursion> {
        datafusion_physical_plan::apply_expression_roots(self.projection.source.iter(), f)
    }

    /// Emit a `CsvScan` node wrapping the shared base config and CSV options.
    #[cfg(feature = "proto")]
    fn try_to_proto(
        &self,
        base: &FileScanConfig,
        ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
        use datafusion_proto_models::protobuf;
        use protobuf::physical_plan_node::PhysicalPlanType;

        let node = protobuf::CsvScanExecNode {
            base_conf: Some(base.try_to_proto(ctx)?),
            has_header: self.has_header(),
            delimiter: proto_byte_to_string(self.delimiter(), "delimiter")?,
            quote: proto_byte_to_string(self.quote(), "quote")?,
            optional_escape: self
                .escape()
                .map(|escape| {
                    Ok::<_, DataFusionError>(
                        protobuf::csv_scan_exec_node::OptionalEscape::Escape(
                            proto_byte_to_string(escape, "escape")?,
                        ),
                    )
                })
                .transpose()?,
            optional_comment: self
                .comment()
                .map(|comment| {
                    Ok::<_, DataFusionError>(
                        protobuf::csv_scan_exec_node::OptionalComment::Comment(
                            proto_byte_to_string(comment, "comment")?,
                        ),
                    )
                })
                .transpose()?,
            newlines_in_values: self.newlines_in_values(),
            truncate_rows: self.truncate_rows(),
        };
        Ok(Some(protobuf::PhysicalPlanNode {
            physical_plan_type: Some(PhysicalPlanType::CsvScan(node)),
        }))
    }
}

impl FileOpener for CsvOpener {
    /// Open a partitioned CSV file.
    ///
    /// If `file_meta.range` is `None`, the entire file is opened.
    /// If `file_meta.range` is `Some(FileRange {start, end})`, this signifies that the partition
    /// corresponds to the byte range [start, end) within the file.
    ///
    /// Note: `start` or `end` might be in the middle of some lines. In such cases, the following rules
    /// are applied to determine which lines to read:
    /// 1. The first line of the partition is the line in which the index of the first character >= `start`.
    /// 2. The last line of the partition is the line in which the byte at position `end - 1` resides.
    ///
    /// Examples:
    /// Consider the following partitions enclosed by braces `{}`:
    ///
    /// {A,1,2,3,4,5,6,7,8,9\n
    ///  A,1,2,3,4,5,6,7,8,9\n}
    ///  A,1,2,3,4,5,6,7,8,9\n
    ///  The lines read would be: [0, 1]
    ///
    ///  A,{1,2,3,4,5,6,7,8,9\n
    ///  A,1,2,3,4,5,6,7,8,9\n
    ///  A},1,2,3,4,5,6,7,8,9\n
    ///  The lines read would be: [1, 2]
    fn open(&self, partitioned_file: PartitionedFile) -> Result<FileOpenFuture> {
        // `self.config.has_header` controls whether to skip reading the 1st line header
        // If the .csv file is read in parallel and this `CsvOpener` is only reading some middle
        // partition, then don't skip first line
        let mut csv_has_header = self.config.has_header();
        if let Some(FileRange { start, .. }) = partitioned_file.range
            && start != 0
        {
            csv_has_header = false;
        }

        let mut config = (*self.config).clone();
        config.options.has_header = Some(csv_has_header);
        config.options.truncated_rows = Some(config.truncate_rows());

        let file_compression_type = self.file_compression_type.to_owned();

        if partitioned_file.range.is_some() {
            assert!(
                !file_compression_type.is_compressed(),
                "Reading compressed .csv in parallel is not supported"
            );
        }

        let store = Arc::clone(&self.object_store);
        let terminator = self.config.terminator();

        let baseline_metrics =
            BaselineMetrics::new(&self.config.metrics, self.partition_index);

        Ok(Box::pin(async move {
            // Current partition contains bytes [start_byte, end_byte) (might contain incomplete lines at boundaries)
            let file_size = partitioned_file.object_meta.size;
            let location = partitioned_file.object_meta.location;

            if let Some(file_range) = partitioned_file.range.as_ref() {
                let raw_start: u64 = file_range.start.try_into().map_err(|_| {
                    exec_datafusion_err!(
                        "Expected start range to fit in u64, got {}",
                        file_range.start
                    )
                })?;
                let raw_end: u64 = file_range.end.try_into().map_err(|_| {
                    exec_datafusion_err!(
                        "Expected end range to fit in u64, got {}",
                        file_range.end
                    )
                })?;

                let aligned_stream = AlignedBoundaryStream::new(
                    Arc::clone(&store),
                    location.clone(),
                    raw_start,
                    raw_end,
                    file_size,
                    terminator.unwrap_or(b'\n'),
                )
                .await?
                .map_err(DataFusionError::from);

                let decoder = config.builder().build_decoder();
                let input = file_compression_type
                    .convert_stream(aligned_stream.boxed())?
                    .fuse();
                let stream = deserialize_stream(
                    input,
                    DecoderDeserializer::new(CsvDecoder::new(decoder)),
                );
                return Ok(stream.map_err(Into::into).boxed());
            }

            // No range specified — read the entire file
            let options = GetOptions::default();
            let result = store.get_opts(&location, options).await?;

            match result.payload {
                #[cfg(not(target_arch = "wasm32"))]
                GetResultPayload::File(file, _) => {
                    let decoder = file_compression_type.convert_read(file)?;
                    let mut reader = config.open(decoder)?;

                    // Use std::iter::from_fn to wrap execution of iterator's next() method.
                    let iterator = std::iter::from_fn(move || {
                        let mut timer = baseline_metrics.elapsed_compute().timer();
                        let result = reader.next();
                        timer.stop();
                        result
                    });

                    Ok(futures::stream::iter(iterator)
                        .map(|r| r.map_err(Into::into))
                        .boxed())
                }
                GetResultPayload::Stream(s) => {
                    let decoder = config.builder().build_decoder();
                    let s = s.map_err(DataFusionError::from);
                    let input = file_compression_type.convert_stream(s.boxed())?.fuse();

                    let stream = deserialize_stream(
                        input,
                        DecoderDeserializer::new(CsvDecoder::new(decoder)),
                    );
                    Ok(stream.map_err(Into::into).boxed())
                }
            }
        }))
    }
}

pub async fn plan_to_csv(
    task_ctx: Arc<TaskContext>,
    plan: Arc<dyn ExecutionPlan>,
    path: impl AsRef<str>,
) -> Result<()> {
    let path = path.as_ref();
    let parsed = ListingTableUrl::parse(path)?;
    let object_store_url = parsed.object_store();
    let store = task_ctx.runtime_env().object_store(&object_store_url)?;
    let writer_buffer_size = task_ctx
        .session_config()
        .options()
        .execution
        .objectstore_writer_buffer_size;
    let mut join_set = JoinSet::new();
    for i in 0..plan.output_partitioning().partition_count() {
        let storeref = Arc::clone(&store);
        let plan: Arc<dyn ExecutionPlan> = Arc::clone(&plan);
        let filename = format!("{}/part-{i}.csv", parsed.prefix());
        let file = object_store::path::Path::parse(filename)?;

        let mut stream = plan.execute(i, Arc::clone(&task_ctx))?;
        join_set.spawn(async move {
            let mut buf_writer =
                BufWriter::with_capacity(storeref, file.clone(), writer_buffer_size);
            let mut buffer = Vec::with_capacity(1024);
            //only write headers on first iteration
            let mut write_headers = true;
            while let Some(batch) = stream.next().await.transpose()? {
                let mut writer = csv::WriterBuilder::new()
                    .with_header(write_headers)
                    .build(buffer);
                writer.write(&batch)?;
                buffer = writer.into_inner();
                buf_writer.write_all(&buffer).await?;
                buffer.clear();
                //prevent writing headers more than once
                write_headers = false;
            }
            buf_writer.shutdown().await.map_err(DataFusionError::from)
        });
    }

    while let Some(result) = join_set.join_next().await {
        match result {
            Ok(res) => res?, // propagate DataFusion error
            Err(e) => {
                if e.is_panic() {
                    std::panic::resume_unwind(e.into_panic());
                } else {
                    unreachable!();
                }
            }
        }
    }

    Ok(())
}

#[cfg(feature = "proto")]
fn proto_byte_to_string(b: u8, description: &str) -> Result<String> {
    let bytes = &[b];
    let s = std::str::from_utf8(bytes).map_err(|_| {
        datafusion_common::internal_datafusion_err!(
            "Invalid CSV {description}: can not represent {bytes:0x?} as utf8"
        )
    })?;
    Ok(s.to_owned())
}

#[cfg(feature = "proto")]
fn proto_str_to_byte(s: &str, description: &str) -> Result<u8> {
    datafusion_common::assert_eq_or_internal_err!(
        s.len(),
        1,
        "Invalid CSV {description}: expected single character, got {s}"
    );
    Ok(s.as_bytes()[0])
}

#[cfg(feature = "proto")]
impl CsvSource {
    /// Reconstructs a `DataSourceExec` from a protobuf `CsvScan`.
    ///
    /// Custom line terminators are not represented in the wire format.
    pub fn try_from_proto(
        node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
        ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        use datafusion_common::config::CsvOptions;
        use datafusion_datasource::file_compression_type::FileCompressionType;
        use datafusion_datasource::file_scan_config::{
            FileScanConfig, FileScanConfigBuilder,
        };
        use datafusion_datasource::source::DataSourceExec;
        use datafusion_proto_models::protobuf;

        let scan = match &node.physical_plan_type {
            Some(protobuf::physical_plan_node::PhysicalPlanType::CsvScan(scan)) => scan,
            _ => {
                return datafusion_common::internal_err!(
                    "PhysicalPlanNode is not a CsvScan"
                );
            }
        };

        let base_conf = scan.base_conf.as_ref().ok_or_else(|| {
            datafusion_common::internal_datafusion_err!(
                "CsvScanExecNode is missing required field 'base_conf'"
            )
        })?;

        let escape = match &scan.optional_escape {
            Some(protobuf::csv_scan_exec_node::OptionalEscape::Escape(escape)) => {
                Some(proto_str_to_byte(escape, "escape")?)
            }
            None => None,
        };
        let comment = match &scan.optional_comment {
            Some(protobuf::csv_scan_exec_node::OptionalComment::Comment(comment)) => {
                Some(proto_str_to_byte(comment, "comment")?)
            }
            None => None,
        };

        let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?;

        let csv_options = CsvOptions {
            has_header: Some(scan.has_header),
            delimiter: proto_str_to_byte(&scan.delimiter, "delimiter")?,
            quote: proto_str_to_byte(&scan.quote, "quote")?,
            newlines_in_values: Some(scan.newlines_in_values),
            truncated_rows: Some(scan.truncate_rows),
            ..Default::default()
        };
        let source = Arc::new(
            CsvSource::new(table_schema)
                .with_csv_options(csv_options)
                .with_escape(escape)
                .with_comment(comment),
        );

        // The compression type is not on the wire; CSV scans always
        // deserialize as uncompressed.
        let conf = FileScanConfigBuilder::from(FileScanConfig::try_from_proto(
            base_conf, ctx, source,
        )?)
        .with_file_compression_type(FileCompressionType::UNCOMPRESSED)
        .build();
        Ok(DataSourceExec::from_data_source(conf))
    }
}