Add Arrow IPC format options

Cargo.toml

@@ -371,7 +371,7 @@ unused_qualifications = "deny"
 # `CARGO_PROFILE_DEV_DEBUG=2 cargo build` when you need to step through code in a
 # debugger.
 [profile.dev]
-debug = "line-tables-only"
+#debug = "line-tables-only"
 
 [profile.release]
 codegen-units = 1

datafusion/common/src/config.rs

@@ -2947,6 +2947,7 @@ pub enum ConfigFileType {
     #[cfg(feature = "parquet")]
     PARQUET,
     JSON,
+    ARROW,
 }
 
 /// Represents the configuration options available for handling different table formats within a data processing application.
@@ -2965,6 +2966,10 @@ pub struct TableOptions {
     /// Configuration options for JSON file handling.
     pub json: JsonOptions,
 
+    /// Configuration options for Arrow IPC file handling. This includes settings for the
+    /// compression codec used for record batches written to disk.
+    pub arrow: ArrowOptions,
+
     /// The current file format that the table operations should assume. This option allows
     /// for dynamic switching between the supported file types (e.g., CSV, Parquet, JSON).
     pub current_format: Option<ConfigFileType>,
@@ -2988,11 +2993,13 @@ impl ConfigField for TableOptions {
                 ConfigFileType::PARQUET => self.parquet.visit(v, "format", ""),
                 ConfigFileType::CSV => self.csv.visit(v, "format", ""),
                 ConfigFileType::JSON => self.json.visit(v, "format", ""),
+                ConfigFileType::ARROW => self.arrow.visit(v, "format", ""),
             }
         } else {
             self.csv.visit(v, "csv", "");
             self.parquet.visit(v, "parquet", "");
             self.json.visit(v, "json", "");
+            self.arrow.visit(v, "arrow", "");
         }
     }
 
@@ -3024,6 +3031,7 @@ impl ConfigField for TableOptions {
                     ConfigFileType::PARQUET => self.parquet.set(rem, value),
                     ConfigFileType::CSV => self.csv.set(rem, value),
                     ConfigFileType::JSON => self.json.set(rem, value),
+                    ConfigFileType::ARROW => self.arrow.set(rem, value),
                 }
             }
             _ => _config_err!("Config value \"{key}\" not found on TableOptions"),
@@ -3885,6 +3893,80 @@ impl EncryptionFactoryOptions {
     }
 }
 
+/// Compression codec used for the record batch bodies of Arrow IPC files.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ArrowIpcCompression {
+    /// No compression
+    Uncompressed,
+    /// LZ4 (frame format)
+    Lz4Frame,
+    /// Zstd
+    Zstd,
+}
+
+impl FromStr for ArrowIpcCompression {
+    type Err = DataFusionError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.replace('\'', "").to_ascii_lowercase().as_str() {
+            "uncompressed" | "" => Ok(Self::Uncompressed),
+            "lz4_frame" => Ok(Self::Lz4Frame),
+            "zstd" => Ok(Self::Zstd),
+            other => Err(DataFusionError::Configuration(format!(
+                "Invalid Arrow IPC compression type: {other}. \
+                Expected one of: uncompressed, lz4_frame, zstd"
+            ))),
+        }
+    }
+}
+
+impl Display for ArrowIpcCompression {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let str = match self {
+            Self::Uncompressed => "uncompressed",
+            Self::Lz4Frame => "lz4_frame",
+            Self::Zstd => "zstd",
+        };
+        write!(f, "{str}")
+    }
+}
+
+impl ConfigField for ArrowIpcCompression {
+    fn visit<V: Visit>(&self, v: &mut V, key: &str, description: &'static str) {
+        v.some(key, self, description)
+    }
+
+    fn set(&mut self, _: &str, value: &str) -> Result<()> {
+        *self = ArrowIpcCompression::from_str(value)?;
+        Ok(())
+    }
+}
+
+impl From<ArrowIpcCompression> for Option<CompressionType> {
+    fn from(c: ArrowIpcCompression) -> Self {
+        match c {
+            ArrowIpcCompression::Uncompressed => None,
+            ArrowIpcCompression::Lz4Frame => Some(CompressionType::LZ4_FRAME),
+            ArrowIpcCompression::Zstd => Some(CompressionType::ZSTD),
+        }
+    }
+}
+
+config_namespace! {
+    /// Options controlling Arrow IPC format
+    pub struct ArrowOptions {
+        /// Compression codec used for the record batch bodies of Arrow IPC files.
+        /// One of: "uncompressed", "lz4_frame", "zstd"
+        pub compression: ArrowIpcCompression, default = ArrowIpcCompression::Lz4Frame
+        /// Compression level to use with the "zstd" codec. Not supported by "lz4_frame"
+        /// or "uncompressed". If not specified, the default zstd compression level is used.
+        pub compression_level: Option<i32>, default = None
+        /// Byte alignment to pad buffers to when writing Arrow IPC files.
+        /// Must be one of 8, 16, 32, or 64.
+        pub alignment: usize, default = 64
+    }
+}
+
 config_namespace! {
     /// Options controlling CSV format
     pub struct CsvOptions {
@@ -4140,7 +4222,7 @@ pub enum OutputFormat {
     #[cfg(feature = "parquet")]
     PARQUET(TableParquetOptions),
     AVRO,
-    ARROW,
+    ARROW(ArrowOptions),
 }
 
 impl Display for OutputFormat {
@@ -4151,7 +4233,7 @@ impl Display for OutputFormat {
             #[cfg(feature = "parquet")]
             OutputFormat::PARQUET(_) => "parquet",
             OutputFormat::AVRO => "avro",
-            OutputFormat::ARROW => "arrow",
+            OutputFormat::ARROW(_) => "arrow",
         };
         write!(f, "{out}")
     }

datafusion/core/src/datasource/file_format/options.rs

@@ -821,7 +821,8 @@ impl ReadOptions<'_> for ArrowReadOptions<'_> {
         _config: &SessionConfig,
         _table_options: TableOptions,
     ) -> ListingOptions {
-        let file_format = ArrowFormat;
+        // There are no read options for Arrow, so we can ignore the `TableOptions`
+        let file_format = ArrowFormat::default();
 
         ListingOptions::new(Arc::new(file_format))
             .with_file_extension(self.file_extension)

datafusion/datasource-arrow/src/file_format.rs

@@ -28,8 +28,9 @@ use arrow::datatypes::{Schema, SchemaRef};
 use arrow::error::ArrowError;
 use arrow::ipc::convert::try_fb_to_schema;
 use arrow::ipc::reader::{FileReader, StreamReader};
+use arrow::ipc::root_as_message;
 use arrow::ipc::writer::IpcWriteOptions;
-use arrow::ipc::{CompressionType, root_as_message};
+use datafusion_common::config::{ArrowOptions, ConfigField, ConfigFileType};
 use datafusion_common::error::Result;
 use datafusion_common::parsers::CompressionTypeVariant;
 use datafusion_common::{
@@ -76,26 +77,52 @@ const BUFFER_FLUSH_BYTES: usize = 1024000;
 
 /// Factory struct used to create [`ArrowFormat`]
 #[derive(Default, Debug)]
-pub struct ArrowFormatFactory;
+pub struct ArrowFormatFactory {
+    /// the write options for arrow
+    pub options: Option<ArrowOptions>,
+}
 
 impl ArrowFormatFactory {
     /// Creates an instance of [ArrowFormatFactory]
     pub fn new() -> Self {
-        Self {}
+        Self { options: None }
+    }
+
+    /// Creates an instance of [`ArrowFormatFactory`] with customized default options
+    pub fn new_with_options(options: ArrowOptions) -> Self {
+        Self {
+            options: Some(options),
+        }
     }
 }
 
 impl FileFormatFactory for ArrowFormatFactory {
     fn create(
         &self,
-        _state: &dyn Session,
-        _format_options: &HashMap<String, String>,
+        state: &dyn Session,
+        format_options: &HashMap<String, String>,
     ) -> Result<Arc<dyn FileFormat>> {
-        Ok(Arc::new(ArrowFormat))
+        let arrow_options = match &self.options {
+            None => {
+                let mut table_options = state.default_table_options();
+                table_options.set_config_format(ConfigFileType::ARROW);
+                table_options.alter_with_string_hash_map(format_options)?;
+                table_options.arrow
+            }
+            Some(arrow_options) => {
+                let mut arrow_options = arrow_options.clone();
+                for (k, v) in format_options {
+                    arrow_options.set(k, v)?;
+                }
+                arrow_options
+            }
+        };
+
+        Ok(Arc::new(ArrowFormat::default().with_options(arrow_options)))
     }
 
     fn default(&self) -> Arc<dyn FileFormat> {
-        Arc::new(ArrowFormat)
+        Arc::new(ArrowFormat::default())
     }
 }
 
@@ -108,7 +135,22 @@ impl GetExt for ArrowFormatFactory {
 
 /// Arrow [`FileFormat`] implementation.
 #[derive(Default, Debug)]
-pub struct ArrowFormat;
+pub struct ArrowFormat {
+    options: ArrowOptions,
+}
+
+impl ArrowFormat {
+    /// Set the arrow options
+    pub fn with_options(mut self, options: ArrowOptions) -> Self {
+        self.options = options;
+        self
+    }
+
+    /// Retrieve the arrow options
+    pub fn options(&self) -> &ArrowOptions {
+        &self.options
+    }
+}
 
 #[async_trait]
 impl FileFormat for ArrowFormat {
@@ -234,7 +276,7 @@ impl FileFormat for ArrowFormat {
             return not_impl_err!("Overwrites are not implemented yet for Arrow format");
         }
 
-        let sink = Arc::new(ArrowFileSink::new(conf));
+        let sink = Arc::new(ArrowFileSink::new(conf, self.options.clone()));
 
         Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _)
     }
@@ -247,11 +289,12 @@ impl FileFormat for ArrowFormat {
 /// Implements [`FileSink`] for Arrow IPC files
 struct ArrowFileSink {
     config: FileSinkConfig,
+    options: ArrowOptions,
 }
 
 impl ArrowFileSink {
-    fn new(config: FileSinkConfig) -> Self {
-        Self { config }
+    fn new(config: FileSinkConfig, options: ArrowOptions) -> Self {
+        Self { config, options }
     }
 }
 
@@ -271,9 +314,13 @@ impl FileSink for ArrowFileSink {
         let mut file_write_tasks: JoinSet<std::result::Result<usize, DataFusionError>> =
             JoinSet::new();
 
-        let ipc_options =
-            IpcWriteOptions::try_new(64, false, arrow_ipc::MetadataVersion::V5)?
-                .try_with_compression(Some(CompressionType::LZ4_FRAME))?;
+        let ipc_options = IpcWriteOptions::try_new(
+            self.options.alignment,
+            false,
+            arrow_ipc::MetadataVersion::V5,
+        )?
+        .try_with_compression(self.options.compression.into())?
+        .try_with_compression_level(self.options.compression_level)?;
         while let Some((file_metadata, mut rx)) = file_stream_rx.recv().await {
             let shared_buffer = SharedBuffer::new(INITIAL_BUFFER_BYTES);
             let mut arrow_writer = arrow_ipc::writer::FileWriter::try_new_with_options(
@@ -660,7 +707,7 @@ mod tests {
                 version: None,
             };
 
-            let arrow_format = ArrowFormat {};
+            let arrow_format = ArrowFormat::default();
             let expected = vec!["f0: Int64", "f1: Utf8", "f2: Boolean"];
 
             // Test chunk sizes where too small so we keep having to read more bytes
@@ -704,7 +751,7 @@ mod tests {
                 version: None,
             };
 
-            let arrow_format = ArrowFormat {};
+            let arrow_format = ArrowFormat::default();
 
             let store = Arc::new(ChunkedStore::new(in_memory_store.clone(), 7));
             let err = arrow_format

datafusion/ffi/src/session/mod.rs

@@ -329,6 +329,7 @@ fn table_options_to_rhash(mut options: TableOptions) -> SVec<(SString, SString)>
                 #[cfg(feature = "parquet")]
                 ConfigFileType::PARQUET => "parquet",
                 ConfigFileType::CSV => "csv",
+                ConfigFileType::ARROW => "arrow",
             }
             .into(),
         );
@@ -633,6 +634,7 @@ fn table_options_from_rhashmap(options: SVec<(SString, SString)>) -> TableOption
         ConfigFileType::JSON,
         #[cfg(feature = "parquet")]
         ConfigFileType::PARQUET,
+        ConfigFileType::ARROW,
     ];
     for format in formats {
         // It is imperative that if new enum variants are added below that they be
@@ -642,6 +644,7 @@ fn table_options_from_rhashmap(options: SVec<(SString, SString)>) -> TableOption
             #[cfg(feature = "parquet")]
             ConfigFileType::PARQUET => "parquet",
             ConfigFileType::JSON => "json",
+            ConfigFileType::ARROW => "arrow",
         };
         let format_options: HashMap<String, String> = options
             .iter()
@@ -665,7 +668,7 @@ fn table_options_from_rhashmap(options: SVec<(SString, SString)>) -> TableOption
         .iter()
         .filter_map(|(k, v)| {
             let (prefix, _) = k.split_once('.')?;
-            if !["json", "parquet", "csv"].contains(&prefix) {
+            if !["json", "parquet", "csv", "arrow"].contains(&prefix) {
                 Some((k.to_owned(), v.to_owned()))
             } else {
                 None
@@ -684,6 +687,7 @@ fn table_options_from_rhashmap(options: SVec<(SString, SString)>) -> TableOption
             #[cfg(feature = "parquet")]
             "parquet" => Some(ConfigFileType::PARQUET),
             "json" => Some(ConfigFileType::JSON),
+            "arrow" => Some(ConfigFileType::ARROW),
             _ => None,
         });
     table_options

datafusion/proto/src/logical_plan/mod.rs

@@ -649,7 +649,7 @@ impl AsLogicalPlan for LogicalPlanNode {
                             }
                         }
                         FileFormatType::Arrow(..) => {
-                            Arc::new(ArrowFormat)
+                            Arc::new(ArrowFormat::default())
                         }
                     };