Support Configuring Arrow RecordBatch Writers via SQL Statement Options (#7390)

* squash, merge main

* cargo fmt

* try to fix pyarrow compile failure

* Apply suggestions from code review

Co-authored-by: Andrew Lamb <[email protected]>

* review comments + add test cases

* update Display for FileTypeWriterOptions

* add unit tests

---------

Co-authored-by: Andrew Lamb <[email protected]>

Author: devinjdangelo

datafusion/common/Cargo.toml

@@ -35,7 +35,7 @@ path = "src/lib.rs"
 [features]
 avro = ["apache-avro"]
 compression = ["xz2", "bzip2", "flate2", "zstd", "async-compression"]
-default = ["compression"]
+default = ["compression", "parquet"]
 pyarrow = ["pyo3", "arrow/pyarrow"]
 
 [dependencies]

datafusion/common/src/error.rs

@@ -73,6 +73,9 @@ pub enum DataFusionError {
     /// This error happens whenever a plan is not valid. Examples include
     /// impossible casts.
     Plan(String),
+    /// This error happens when an invalid or unsupported option is passed
+    /// in a SQL statement
+    Configuration(String),
     /// This error happens with schema-related errors, such as schema inference not possible
     /// and non-unique column names.
     SchemaError(SchemaError),
@@ -288,6 +291,9 @@ impl Display for DataFusionError {
             DataFusionError::SQL(ref desc) => {
                 write!(f, "SQL error: {desc:?}")
             }
+            DataFusionError::Configuration(ref desc) => {
+                write!(f, "Invalid or Unsupported Configuration: {desc}")
+            }
             DataFusionError::NotImplemented(ref desc) => {
                 write!(f, "This feature is not implemented: {desc}")
             }
@@ -338,6 +344,7 @@ impl Error for DataFusionError {
             DataFusionError::SQL(e) => Some(e),
             DataFusionError::NotImplemented(_) => None,
             DataFusionError::Internal(_) => None,
+            DataFusionError::Configuration(_) => None,
             DataFusionError::Plan(_) => None,
             DataFusionError::SchemaError(e) => Some(e),
             DataFusionError::Execution(_) => None,

datafusion/common/src/file_options/arrow_writer.rs

@@ -0,0 +1,36 @@
+// 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.
+
+//! Options related to how Arrow files should be written
+
+use crate::{
+    config::ConfigOptions,
+    error::{DataFusionError, Result},
+};
+
+use super::StatementOptions;
+
+#[derive(Clone, Debug)]
+pub struct ArrowWriterOptions {}
+
+impl TryFrom<(&ConfigOptions, &StatementOptions)> for ArrowWriterOptions {
+    type Error = DataFusionError;
+
+    fn try_from(_value: (&ConfigOptions, &StatementOptions)) -> Result<Self> {
+        Ok(ArrowWriterOptions {})
+    }
+}

datafusion/common/src/file_options/avro_writer.rs

@@ -0,0 +1,36 @@
+// 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.
+
+//! Options related to how avro files should be written
+
+use crate::{
+    config::ConfigOptions,
+    error::{DataFusionError, Result},
+};
+
+use super::StatementOptions;
+
+#[derive(Clone, Debug)]
+pub struct AvroWriterOptions {}
+
+impl TryFrom<(&ConfigOptions, &StatementOptions)> for AvroWriterOptions {
+    type Error = DataFusionError;
+
+    fn try_from(_value: (&ConfigOptions, &StatementOptions)) -> Result<Self> {
+        Ok(AvroWriterOptions {})
+    }
+}

datafusion/common/src/file_options/csv_writer.rs

@@ -0,0 +1,107 @@
+// 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.
+
+//! Options related to how csv files should be written
+
+use std::str::FromStr;
+
+use arrow::csv::WriterBuilder;
+
+use crate::{
+    config::ConfigOptions,
+    error::{DataFusionError, Result},
+    parsers::CompressionTypeVariant,
+};
+
+use super::StatementOptions;
+
+/// Options for writing CSV files
+#[derive(Clone, Debug)]
+pub struct CsvWriterOptions {
+    /// Struct from the arrow crate which contains all csv writing related settings
+    pub writer_options: WriterBuilder,
+    /// Compression to apply after ArrowWriter serializes RecordBatches.
+    /// This compression is applied by DataFusion not the ArrowWriter itself.
+    pub compression: CompressionTypeVariant,
+    /// Indicates whether WriterBuilder.has_header() is set to true.
+    /// This is duplicative as WriterBuilder also stores this information.
+    /// However, WriterBuilder does not allow public read access to the
+    /// has_header parameter.
+    pub has_header: bool,
+    // TODO: expose a way to read has_header in arrow create
+    // https://github.com/apache/arrow-rs/issues/4735
+}
+
+impl TryFrom<(&ConfigOptions, &StatementOptions)> for CsvWriterOptions {
+    type Error = DataFusionError;
+
+    fn try_from(value: (&ConfigOptions, &StatementOptions)) -> Result<Self> {
+        let _configs = value.0;
+        let statement_options = value.1;
+        let mut has_header = true;
+        let mut builder = WriterBuilder::default();
+        let mut compression = CompressionTypeVariant::UNCOMPRESSED;
+        for (option, value) in &statement_options.options {
+            builder = match option.to_lowercase().as_str(){
+                "header" => {
+                    has_header = value.parse()
+                        .map_err(|_| DataFusionError::Configuration(format!("Unable to parse {value} as bool as required for {option}!")))?;
+                    builder.has_headers(has_header)
+                },
+                "date_format" => builder.with_date_format(value.to_owned()),
+                "datetime_format" => builder.with_datetime_format(value.to_owned()),
+                "timestamp_format" => builder.with_timestamp_format(value.to_owned()),
+                "time_format" => builder.with_time_format(value.to_owned()),
+                "rfc3339" => {
+                    let value_bool = value.parse()
+                        .map_err(|_| DataFusionError::Configuration(format!("Unable to parse {value} as bool as required for {option}!")))?;
+                    if value_bool{
+                        builder.with_rfc3339()
+                    } else{
+                        builder
+                    }
+                },
+                "null_value" => builder.with_null(value.to_owned()),
+                "compression" => {
+                    compression = CompressionTypeVariant::from_str(value.replace('\'', "").as_str())?;
+                    builder
+                },
+                "delimeter" => {
+                    // Ignore string literal single quotes passed from sql parsing
+                    let value = value.replace('\'', "");
+                    let chars: Vec<char> = value.chars().collect();
+                    if chars.len()>1{
+                        return Err(DataFusionError::Configuration(format!(
+                            "CSV Delimeter Option must be a single char, got: {}", value
+                        )))
+                    }
+                    builder.with_delimiter(chars[0].try_into().map_err(|_| {
+                        DataFusionError::Internal(
+                            "Unable to convert CSV delimiter into u8".into(),
+                        )
+                    })?)
+            },
+                _ => return Err(DataFusionError::Configuration(format!("Found unsupported option {option} with value {value} for CSV format!")))
+            }
+        }
+        Ok(CsvWriterOptions {
+            has_header,
+            writer_options: builder,
+            compression,
+        })
+    }
+}