feat: add compression level configuration for JSON/CSV writers

Adds `compression_level` option to `JsonOptions` and `CsvOptions` allowing
users to specify compression level for ZSTD, GZIP, BZIP2, and XZ algorithms.

- Add compression_level field to JsonOptions and CsvOptions in config.rs
- Add convert_async_writer_with_level method (non-breaking, extends API)
- Keep original convert_async_writer for backward compatibility
- Update JsonWriterOptions and CsvWriterOptions with compression_level
- Update ObjectWriterBuilder to support compression level
- Update JSON and CSV sinks to pass compression level through
- Update proto definitions and conversions for serialization

Closes #18947

Author: Smotrov

datafusion/common/src/config.rs

@@ -1661,6 +1661,7 @@ config_field!(bool, value => default_config_transform(value.to_lowercase().as_st
 config_field!(usize);
 config_field!(f64);
 config_field!(u64);
+config_field!(u32);
 
 impl ConfigField for u8 {
     fn visit<V: Visit>(&self, v: &mut V, key: &str, description: &'static str) {
@@ -2786,6 +2787,14 @@ config_namespace! {
         /// The default behaviour depends on the `datafusion.catalog.newlines_in_values` setting.
         pub newlines_in_values: Option<bool>, default = None
         pub compression: CompressionTypeVariant, default = CompressionTypeVariant::UNCOMPRESSED
+        /// Compression level for the output file. The valid range depends on the
+        /// compression algorithm:
+        /// - ZSTD: 1 to 22 (default: 3)
+        /// - GZIP: 0 to 10 (default: varies by implementation)
+        /// - BZIP2: 0 to 9 (default: 6)
+        /// - XZ: 0 to 9 (default: 6)
+        /// If not specified, the default level for the compression algorithm is used.
+        pub compression_level: Option<u32>, default = None
         pub schema_infer_max_rec: Option<usize>, default = None
         pub date_format: Option<String>, default = None
         pub datetime_format: Option<String>, default = None
@@ -2908,6 +2917,14 @@ impl CsvOptions {
         self
     }
 
+    /// Set the compression level for the output file.
+    /// The valid range depends on the compression algorithm.
+    /// If not specified, the default level for the algorithm is used.
+    pub fn with_compression_level(mut self, level: u32) -> Self {
+        self.compression_level = Some(level);
+        self
+    }
+
     /// The delimiter character.
     pub fn delimiter(&self) -> u8 {
         self.delimiter
@@ -2933,6 +2950,14 @@ config_namespace! {
     /// Options controlling JSON format
     pub struct JsonOptions {
         pub compression: CompressionTypeVariant, default = CompressionTypeVariant::UNCOMPRESSED
+        /// Compression level for the output file. The valid range depends on the
+        /// compression algorithm:
+        /// - ZSTD: 1 to 22 (default: 3)
+        /// - GZIP: 0 to 10 (default: varies by implementation)
+        /// - BZIP2: 0 to 9 (default: 6)
+        /// - XZ: 0 to 9 (default: 6)
+        /// If not specified, the default level for the compression algorithm is used.
+        pub compression_level: Option<u32>, default = None
         pub schema_infer_max_rec: Option<usize>, default = None
     }
 }

datafusion/common/src/file_options/csv_writer.rs

@@ -31,6 +31,8 @@ pub struct CsvWriterOptions {
     /// Compression to apply after ArrowWriter serializes RecordBatches.
     /// This compression is applied by DataFusion not the ArrowWriter itself.
     pub compression: CompressionTypeVariant,
+    /// Compression level for the output file.
+    pub compression_level: Option<u32>,
 }
 
 impl CsvWriterOptions {
@@ -41,6 +43,20 @@ impl CsvWriterOptions {
         Self {
             writer_options,
             compression,
+            compression_level: None,
+        }
+    }
+
+    /// Create a new `CsvWriterOptions` with the specified compression level.
+    pub fn new_with_level(
+        writer_options: WriterBuilder,
+        compression: CompressionTypeVariant,
+        compression_level: Option<u32>,
+    ) -> Self {
+        Self {
+            writer_options,
+            compression,
+            compression_level,
         }
     }
 }
@@ -81,6 +97,7 @@ impl TryFrom<&CsvOptions> for CsvWriterOptions {
         Ok(CsvWriterOptions {
             writer_options: builder,
             compression: value.compression,
+            compression_level: value.compression_level,
         })
     }
 }

datafusion/common/src/file_options/json_writer.rs

@@ -27,11 +27,26 @@ use crate::{
 #[derive(Clone, Debug)]
 pub struct JsonWriterOptions {
     pub compression: CompressionTypeVariant,
+    pub compression_level: Option<u32>,
 }
 
 impl JsonWriterOptions {
     pub fn new(compression: CompressionTypeVariant) -> Self {
-        Self { compression }
+        Self {
+            compression,
+            compression_level: None,
+        }
+    }
+
+    /// Create a new `JsonWriterOptions` with the specified compression and level.
+    pub fn new_with_level(
+        compression: CompressionTypeVariant,
+        compression_level: Option<u32>,
+    ) -> Self {
+        Self {
+            compression,
+            compression_level,
+        }
     }
 }
 
@@ -41,6 +56,7 @@ impl TryFrom<&JsonOptions> for JsonWriterOptions {
     fn try_from(value: &JsonOptions) -> Result<Self> {
         Ok(JsonWriterOptions {
             compression: value.compression,
+            compression_level: value.compression_level,
         })
     }
 }

datafusion/datasource-csv/src/file_format.rs

@@ -780,6 +780,7 @@ impl FileSink for CsvSink {
             context,
             serializer,
             self.writer_options.compression.into(),
+            self.writer_options.compression_level,
             object_store,
             demux_task,
             file_stream_rx,

datafusion/datasource-json/src/file_format.rs

@@ -379,6 +379,7 @@ impl FileSink for JsonSink {
             context,
             serializer,
             self.writer_options.compression.into(),
+            self.writer_options.compression_level,
             object_store,
             demux_task,
             file_stream_rx,

datafusion/datasource/src/file_compression_type.rs

@@ -155,25 +155,63 @@ impl FileCompressionType {
     }
 
     /// Wrap the given `BufWriter` so that it performs compressed writes
-    /// according to this `FileCompressionType`.
+    /// according to this `FileCompressionType` using the default compression level.
     pub fn convert_async_writer(
         &self,
         w: BufWriter,
     ) -> Result<Box<dyn AsyncWrite + Send + Unpin>> {
+        self.convert_async_writer_with_level(w, None)
+    }
+
+    /// Wrap the given `BufWriter` so that it performs compressed writes
+    /// according to this `FileCompressionType`.
+    ///
+    /// If `compression_level` is `Some`, the encoder will use the specified
+    /// compression level. If `None`, the default level for each algorithm is used.
+    pub fn convert_async_writer_with_level(
+        &self,
+        w: BufWriter,
+        compression_level: Option<u32>,
+    ) -> Result<Box<dyn AsyncWrite + Send + Unpin>> {
+        #[cfg(feature = "compression")]
+        use async_compression::Level;
+
         Ok(match self.variant {
             #[cfg(feature = "compression")]
-            GZIP => Box::new(GzipEncoder::new(w)),
+            GZIP => match compression_level {
+                Some(level) => {
+                    Box::new(GzipEncoder::with_quality(w, Level::Precise(level as i32)))
+                }
+                None => Box::new(GzipEncoder::new(w)),
+            },
             #[cfg(feature = "compression")]
-            BZIP2 => Box::new(BzEncoder::new(w)),
+            BZIP2 => match compression_level {
+                Some(level) => {
+                    Box::new(BzEncoder::with_quality(w, Level::Precise(level as i32)))
+                }
+                None => Box::new(BzEncoder::new(w)),
+            },
             #[cfg(feature = "compression")]
-            XZ => Box::new(XzEncoder::new(w)),
+            XZ => match compression_level {
+                Some(level) => {
+                    Box::new(XzEncoder::with_quality(w, Level::Precise(level as i32)))
+                }
+                None => Box::new(XzEncoder::new(w)),
+            },
             #[cfg(feature = "compression")]
-            ZSTD => Box::new(ZstdEncoder::new(w)),
+            ZSTD => match compression_level {
+                Some(level) => {
+                    Box::new(ZstdEncoder::with_quality(w, Level::Precise(level as i32)))
+                }
+                None => Box::new(ZstdEncoder::new(w)),
+            },
             #[cfg(not(feature = "compression"))]
             GZIP | BZIP2 | XZ | ZSTD => {
+                // compression_level is not used when compression feature is disabled
+                let _ = compression_level;
                 return Err(DataFusionError::NotImplemented(
                     "Compression feature is not enabled".to_owned(),
-                ))
+                ));
             }
             UNCOMPRESSED => Box::new(w),
         })

datafusion/datasource/src/write/mod.rs

@@ -131,6 +131,8 @@ pub struct ObjectWriterBuilder {
     object_store: Arc<dyn ObjectStore>,
     /// The size of the buffer for the object writer.
     buffer_size: Option<usize>,
+    /// The compression level for the object writer.
+    compression_level: Option<u32>,
 }
 
 impl ObjectWriterBuilder {
@@ -145,6 +147,7 @@ impl ObjectWriterBuilder {
             location: location.clone(),
             object_store,
             buffer_size: None,
+            compression_level: None,
         }
     }
 
@@ -202,6 +205,22 @@ impl ObjectWriterBuilder {
         self.buffer_size
     }
 
+    /// Set compression level for object writer.
+    pub fn set_compression_level(&mut self, compression_level: Option<u32>) {
+        self.compression_level = compression_level;
+    }
+
+    /// Set compression level for object writer, returning the builder.
+    pub fn with_compression_level(mut self, compression_level: Option<u32>) -> Self {
+        self.compression_level = compression_level;
+        self
+    }
+
+    /// Currently specified compression level.
+    pub fn get_compression_level(&self) -> Option<u32> {
+        self.compression_level
+    }
+
     /// Return a writer object that writes to the object store location.
     ///
     /// If a buffer size has not been set, the default buffer buffer size will
@@ -215,13 +234,15 @@ impl ObjectWriterBuilder {
             location,
             object_store,
             buffer_size,
+            compression_level,
         } = self;
 
         let buf_writer = match buffer_size {
             Some(size) => BufWriter::with_capacity(object_store, location, size),
             None => BufWriter::new(object_store, location),
         };
 
-        file_compression_type.convert_async_writer(buf_writer)
+        file_compression_type
+            .convert_async_writer_with_level(buf_writer, compression_level)
     }
 }

datafusion/datasource/src/write/orchestration.rs

@@ -240,6 +240,7 @@ pub async fn spawn_writer_tasks_and_join(
     context: &Arc<TaskContext>,
     serializer: Arc<dyn BatchSerializer>,
     compression: FileCompressionType,
+    compression_level: Option<u32>,
     object_store: Arc<dyn ObjectStore>,
     demux_task: SpawnedTask<Result<()>>,
     mut file_stream_rx: DemuxedStreamReceiver,
@@ -265,6 +266,7 @@ pub async fn spawn_writer_tasks_and_join(
                         .execution
                         .objectstore_writer_buffer_size,
                 ))
+                .with_compression_level(compression_level)
                 .build()?;
 
         if tx_file_bundle

datafusion/expr/src/udf.rs

@@ -25,10 +25,9 @@ use crate::udf_eq::UdfEq;
 use crate::{ColumnarValue, Documentation, Expr, Signature};
 use arrow::datatypes::{DataType, Field, FieldRef};
 use datafusion_common::config::ConfigOptions;
-use datafusion_common::{
-    assert_or_internal_err, not_impl_err, DataFusionError, ExprSchema, Result,
-    ScalarValue,
-};
+#[cfg(debug_assertions)]
+use datafusion_common::{assert_or_internal_err, DataFusionError};
+use datafusion_common::{not_impl_err, ExprSchema, Result, ScalarValue};
 use datafusion_expr_common::dyn_eq::{DynEq, DynHash};
 use datafusion_expr_common::interval_arithmetic::Interval;
 use std::any::Any;

datafusion/proto-common/proto/datafusion_common.proto

@@ -461,12 +461,14 @@ message CsvOptions {
   bytes newlines_in_values = 16; // Indicates if newlines are supported in values
   bytes terminator = 17; // Optional terminator character as a byte
   bytes truncated_rows = 18; // Indicates if truncated rows are allowed
+  optional uint32 compression_level = 19; // Optional compression level
 }
 
 // Options controlling CSV format
 message JsonOptions {
   CompressionTypeVariant compression = 1; // Compression type
   optional uint64 schema_infer_max_rec = 2; // Optional max records for schema inference
+  optional uint32 compression_level = 3; // Optional compression level
 }
 
 message TableParquetOptions {