Add compression option to SpillManager

Author: ding-young

datafusion/common/src/config.rs

@@ -17,6 +17,8 @@
 
 //! Runtime configuration, via [`ConfigOptions`]
 
+use arrow_ipc::CompressionType;
+
 use crate::error::_config_err;
 use crate::parsers::CompressionTypeVariant;
 use crate::utils::get_available_parallelism;
@@ -274,6 +276,61 @@ config_namespace! {
     }
 }
 
+// TODO where should we parse?
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum SpillCompression {
+    Zstd,
+    Lz4_frame,
+    Uncompressed,
+}
+
+impl FromStr for SpillCompression {
+    type Err = DataFusionError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.to_ascii_lowercase().as_str() {
+            "zstd" => Ok(Self::Zstd),
+            "lz4_frame" => Ok(Self::Lz4_frame),
+            "uncompressed" | "" => Ok(Self::Uncompressed),
+            other => Err(DataFusionError::Execution(format!(
+                "Invalid Spill file compression type: {other}. Expected one of: zstd, lz4, uncompressed"
+            ))),
+        }
+    }
+}
+
+impl ConfigField for SpillCompression {
+    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 = SpillCompression::from_str(value)?;
+        Ok(())
+    }
+}
+
+impl Display for SpillCompression {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let str = match self {
+            Self::Zstd => "Zstd",
+            Self::Lz4_frame => "Lz4",
+            Self::Uncompressed => "",
+        };
+        write!(f, "{str}")
+    }
+}
+
+impl From<SpillCompression> for Option<CompressionType> {
+    fn from(c: SpillCompression) -> Self {
+        match c {
+            SpillCompression::Zstd => Some(CompressionType::ZSTD),
+            SpillCompression::Lz4_frame => Some(CompressionType::LZ4_FRAME),
+            SpillCompression::Uncompressed => None,
+        }
+    }
+}
+
 config_namespace! {
     /// Options related to query execution
     ///
@@ -330,6 +387,11 @@ config_namespace! {
         /// the new schema verification step.
         pub skip_physical_aggregate_schema_check: bool, default = false
 
+        /// TODO do we may need to support compression level setting Then we may change where parsing config happens
+        /// Sets default compression codec for spill file.
+        /// Valid values are: uncompressed, zstd, lz4_frame
+        pub spill_compression: SpillCompression, default = SpillCompression::Uncompressed
+
         /// Specifies the reserved memory for each spillable sort operation to
         /// facilitate an in-memory merge.
         ///

datafusion/execution/src/config.rs

@@ -23,7 +23,7 @@ use std::{
 };
 
 use datafusion_common::{
-    config::{ConfigExtension, ConfigOptions},
+    config::{ConfigExtension, ConfigOptions, SpillCompression},
     Result, ScalarValue,
 };
 
@@ -258,6 +258,11 @@ impl SessionConfig {
         self.options.execution.collect_statistics
     }
 
+    /// Compression codec for spill file
+    pub fn spill_compression(&self) -> SpillCompression {
+        self.options.execution.spill_compression
+    }
+
     /// Selects a name for the default catalog and schema
     pub fn with_default_catalog_and_schema(
         mut self,

datafusion/physical-plan/benches/spill_io.rs

@@ -20,6 +20,7 @@ use arrow::array::{
 };
 use arrow::datatypes::{DataType, Field, Schema};
 use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
+use datafusion_common::config::SpillCompression;
 use datafusion_execution::runtime_env::RuntimeEnv;
 use datafusion_physical_plan::common::collect;
 use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, SpillMetrics};
@@ -82,7 +83,7 @@ fn bench_spill_io(c: &mut Criterion) {
         Field::new("c2", DataType::Date32, true),
         Field::new("c3", DataType::Decimal128(11, 2), true),
     ]));
-    let spill_manager = SpillManager::new(env, metrics, schema);
+    let spill_manager = SpillManager::new(env, metrics, schema, SpillCompression::Uncompressed);
 
     let mut group = c.benchmark_group("spill_io");
     let rt = Runtime::new().unwrap();

datafusion/physical-plan/src/aggregates/row_hash.rs

@@ -552,6 +552,7 @@ impl GroupedHashAggregateStream {
             context.runtime_env(),
             metrics::SpillMetrics::new(&agg.metrics, partition),
             Arc::clone(&partial_agg_schema),
+            context.session_config().spill_compression(),
         );
 
         let spill_state = SpillState {

datafusion/physical-plan/src/joins/sort_merge_join.rs

@@ -61,6 +61,7 @@ use arrow::compute::{
 use arrow::datatypes::{DataType, SchemaRef, TimeUnit};
 use arrow::error::ArrowError;
 use arrow::ipc::reader::StreamReader;
+use datafusion_common::config::SpillCompression;
 use datafusion_common::{
     exec_err, internal_err, not_impl_err, plan_err, DataFusionError, HashSet, JoinSide,
     JoinType, Result,
@@ -495,6 +496,7 @@ impl ExecutionPlan for SortMergeJoinExec {
 
         // create join stream
         Ok(Box::pin(SortMergeJoinStream::try_new(
+            context.session_config().spill_compression(),
             Arc::clone(&self.schema),
             self.sort_options.clone(),
             self.null_equals_null,
@@ -1319,6 +1321,7 @@ impl Stream for SortMergeJoinStream {
 impl SortMergeJoinStream {
     #[allow(clippy::too_many_arguments)]
     pub fn try_new(
+        spill_compression: SpillCompression,
         schema: SchemaRef,
         sort_options: Vec<SortOptions>,
         null_equals_null: bool,
@@ -1339,6 +1342,7 @@ impl SortMergeJoinStream {
             Arc::clone(&runtime_env),
             join_metrics.spill_metrics.clone(),
             Arc::clone(&buffered_schema),
+            spill_compression,
         );
         Ok(Self {
             state: SortMergeJoinState::Init,

datafusion/physical-plan/src/sorts/sort.rs

@@ -47,6 +47,7 @@ use crate::{
 use arrow::array::{Array, RecordBatch, RecordBatchOptions, StringViewArray};
 use arrow::compute::{concat_batches, lexsort_to_indices, take_arrays};
 use arrow::datatypes::SchemaRef;
+use datafusion_common::config::SpillCompression;
 use datafusion_common::{internal_datafusion_err, internal_err, DataFusionError, Result};
 use datafusion_execution::disk_manager::RefCountedTempFile;
 use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
@@ -255,6 +256,7 @@ impl ExternalSorter {
         batch_size: usize,
         sort_spill_reservation_bytes: usize,
         sort_in_place_threshold_bytes: usize,
+        spill_compression: SpillCompression,
         metrics: &ExecutionPlanMetricsSet,
         runtime: Arc<RuntimeEnv>,
     ) -> Result<Self> {
@@ -271,6 +273,7 @@ impl ExternalSorter {
             Arc::clone(&runtime),
             metrics.spill_metrics.clone(),
             Arc::clone(&schema),
+            spill_compression,
         );
 
         Ok(Self {
@@ -1146,6 +1149,7 @@ impl ExecutionPlan for SortExec {
                     context.session_config().batch_size(),
                     execution_options.sort_spill_reservation_bytes,
                     execution_options.sort_in_place_threshold_bytes,
+                    context.session_config().spill_compression(),
                     &self.metrics_set,
                     context.runtime_env(),
                 )?;

datafusion/physical-plan/src/spill/in_progress_spill_file.rs

@@ -64,9 +64,11 @@ impl InProgressSpillFile {
         if self.writer.is_none() {
             let schema = batch.schema();
             if let Some(ref in_progress_file) = self.in_progress_file {
+                // TODO hmm passing compresion option again?
                 self.writer = Some(IPCStreamWriter::new(
                     in_progress_file.path(),
                     schema.as_ref(),
+                    self.spill_writer.compression,
                 )?);
 
                 // Update metrics

datafusion/physical-plan/src/spill/mod.rs

@@ -30,9 +30,12 @@ use std::task::{Context, Poll};
 
 use arrow::array::ArrayData;
 use arrow::datatypes::{Schema, SchemaRef};
+use arrow::ipc::writer::IpcWriteOptions;
 use arrow::ipc::{reader::StreamReader, writer::StreamWriter};
+use arrow::ipc::{CompressionType, MetadataVersion};
 use arrow::record_batch::RecordBatch;
 
+use datafusion_common::config::SpillCompression;
 use datafusion_common::{exec_datafusion_err, DataFusionError, HashSet, Result};
 use datafusion_common_runtime::SpawnedTask;
 use datafusion_execution::disk_manager::RefCountedTempFile;
@@ -194,7 +197,8 @@ pub fn spill_record_batch_by_size(
 ) -> Result<()> {
     let mut offset = 0;
     let total_rows = batch.num_rows();
-    let mut writer = IPCStreamWriter::new(&path, schema.as_ref())?;
+    let mut writer =
+        IPCStreamWriter::new(&path, schema.as_ref(), SpillCompression::Uncompressed)?;
 
     while offset < total_rows {
         let length = std::cmp::min(total_rows - offset, batch_size_rows);
@@ -292,15 +296,28 @@ struct IPCStreamWriter {
 
 impl IPCStreamWriter {
     /// Create new writer
-    pub fn new(path: &Path, schema: &Schema) -> Result<Self> {
+    pub fn new(
+        path: &Path,
+        schema: &Schema,
+        compression_type: SpillCompression,
+    ) -> Result<Self> {
         let file = File::create(path).map_err(|e| {
             exec_datafusion_err!("Failed to create partition file at {path:?}: {e:?}")
         })?;
+
+        // TODO what should be default metadata version & alignment?
+        let metadata_version = MetadataVersion::V5;
+        let alignment = 8;
+        let mut write_options =
+            IpcWriteOptions::try_new(alignment, false, metadata_version)?;
+        write_options = write_options.try_with_compression(compression_type.into())?;
+
+        let writer = StreamWriter::try_new_with_options(file, schema, write_options)?;
         Ok(Self {
             num_batches: 0,
             num_rows: 0,
             num_bytes: 0,
-            writer: StreamWriter::try_new(file, schema)?,
+            writer,
         })
     }
 
@@ -362,7 +379,12 @@ mod tests {
         // Construct SpillManager
         let env = Arc::new(RuntimeEnv::default());
         let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
-        let spill_manager = SpillManager::new(env, metrics, Arc::clone(&schema));
+        let spill_manager = SpillManager::new(
+            env,
+            metrics,
+            Arc::clone(&schema),
+            SpillCompression::Uncompressed,
+        );
 
         let spill_file = spill_manager
             .spill_record_batch_and_finish(&[batch1, batch2], "Test")?
@@ -426,7 +448,12 @@ mod tests {
         // Construct SpillManager
         let env = Arc::new(RuntimeEnv::default());
         let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
-        let spill_manager = SpillManager::new(env, metrics, Arc::clone(&dict_schema));
+        let spill_manager = SpillManager::new(
+            env,
+            metrics,
+            Arc::clone(&dict_schema),
+            SpillCompression::Uncompressed,
+        );
 
         let num_rows = batch1.num_rows() + batch2.num_rows();
         let spill_file = spill_manager
@@ -454,7 +481,12 @@ mod tests {
         let schema = batch1.schema();
         let env = Arc::new(RuntimeEnv::default());
         let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
-        let spill_manager = SpillManager::new(env, metrics, Arc::clone(&schema));
+        let spill_manager = SpillManager::new(
+            env,
+            metrics,
+            Arc::clone(&schema),
+            SpillCompression::Uncompressed,
+        );
 
         let spill_file = spill_manager
             .spill_record_batch_by_size(&batch1, "Test Spill", 1)?
@@ -608,7 +640,12 @@ mod tests {
             Field::new("b", DataType::Utf8, false),
         ]));
 
-        let spill_manager = SpillManager::new(env, metrics, Arc::clone(&schema));
+        let spill_manager = SpillManager::new(
+            env,
+            metrics,
+            Arc::clone(&schema),
+            SpillCompression::Uncompressed,
+        );
 
         let batch = RecordBatch::try_new(
             schema,
@@ -665,8 +702,12 @@ mod tests {
             Field::new("b", DataType::Utf8, false),
         ]));
 
-        let spill_manager =
-            Arc::new(SpillManager::new(env, metrics, Arc::clone(&schema)));
+        let spill_manager = Arc::new(SpillManager::new(
+            env,
+            metrics,
+            Arc::clone(&schema),
+            SpillCompression::Uncompressed,
+        ));
         let mut in_progress_file = spill_manager.create_in_progress_file("Test")?;
 
         let batch1 = RecordBatch::try_new(
@@ -712,8 +753,12 @@ mod tests {
             Field::new("b", DataType::Utf8, false),
         ]));
 
-        let spill_manager =
-            Arc::new(SpillManager::new(env, metrics, Arc::clone(&schema)));
+        let spill_manager = Arc::new(SpillManager::new(
+            env,
+            metrics,
+            Arc::clone(&schema),
+            SpillCompression::Uncompressed,
+        ));
 
         // Test write empty batch with interface `InProgressSpillFile` and `append_batch()`
         let mut in_progress_file = spill_manager.create_in_progress_file("Test")?;
@@ -758,7 +803,12 @@ mod tests {
                 // Construct SpillManager
                 let env = Arc::new(RuntimeEnv::default());
                 let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
-                let spill_manager = SpillManager::new(env, metrics, Arc::clone(&schema));
+                let spill_manager = SpillManager::new(
+                    env,
+                    metrics,
+                    Arc::clone(&schema),
+                    SpillCompression::Uncompressed,
+                );
                 let batches: [_; 10] = std::array::from_fn(|_| batch.clone());
 
                 let spill_file_1 = spill_manager

datafusion/physical-plan/src/spill/spill_manager.rs

@@ -23,7 +23,7 @@ use arrow::datatypes::SchemaRef;
 use arrow::record_batch::RecordBatch;
 use datafusion_execution::runtime_env::RuntimeEnv;
 
-use datafusion_common::Result;
+use datafusion_common::{config::SpillCompression, Result};
 use datafusion_execution::disk_manager::RefCountedTempFile;
 use datafusion_execution::SendableRecordBatchStream;
 
@@ -44,16 +44,23 @@ pub struct SpillManager {
     schema: SchemaRef,
     /// Number of batches to buffer in memory during disk reads
     batch_read_buffer_capacity: usize,
-    // TODO: Add general-purpose compression options
+    /// general-purpose compression options
+    pub(crate) compression: SpillCompression,
 }
 
 impl SpillManager {
-    pub fn new(env: Arc<RuntimeEnv>, metrics: SpillMetrics, schema: SchemaRef) -> Self {
+    pub fn new(
+        env: Arc<RuntimeEnv>,
+        metrics: SpillMetrics,
+        schema: SchemaRef,
+        compression: SpillCompression,
+    ) -> Self {
         Self {
             env,
             metrics,
             schema,
             batch_read_buffer_capacity: 2,
+            compression,
         }
     }