add test case for spill compression

Author: ding-young

Cargo.toml

@@ -97,7 +97,7 @@ arrow-flight = { version = "55.1.0", features = [
     "flight-sql-experimental",
 ] }
 arrow-ipc = { version = "55.0.0", default-features = false, features = [
-    "lz4",
+    "lz4", "zstd",
 ] }
 arrow-ord = { version = "55.0.0", default-features = false }
 arrow-schema = { version = "55.0.0", default-features = false }

datafusion/common/src/config.rs

@@ -280,7 +280,7 @@ config_namespace! {
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum SpillCompression {
     Zstd,
-    Lz4_frame,
+    Lz4Frame,
     Uncompressed,
 }
 
@@ -290,7 +290,7 @@ impl FromStr for SpillCompression {
     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),
+            "lz4_frame" => Ok(Self::Lz4Frame),
             "uncompressed" | "" => Ok(Self::Uncompressed),
             other => Err(DataFusionError::Execution(format!(
                 "Invalid Spill file compression type: {other}. Expected one of: zstd, lz4, uncompressed"
@@ -313,8 +313,8 @@ impl ConfigField for SpillCompression {
 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::Zstd => "zstd",
+            Self::Lz4Frame => "lz4_frame",
             Self::Uncompressed => "",
         };
         write!(f, "{str}")
@@ -325,7 +325,7 @@ 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::Lz4Frame => Some(CompressionType::LZ4_FRAME),
             SpillCompression::Uncompressed => None,
         }
     }

datafusion/physical-plan/benches/spill_io.rs

@@ -83,7 +83,8 @@ 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, SpillCompression::Uncompressed);
+    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/spill/mod.rs

@@ -31,8 +31,8 @@ use std::task::{Context, Poll};
 use arrow::array::ArrayData;
 use arrow::datatypes::{Schema, SchemaRef};
 use arrow::ipc::writer::IpcWriteOptions;
+use arrow::ipc::MetadataVersion;
 use arrow::ipc::{reader::StreamReader, writer::StreamWriter};
-use arrow::ipc::{CompressionType, MetadataVersion};
 use arrow::record_batch::RecordBatch;
 
 use datafusion_common::config::SpillCompression;
@@ -349,6 +349,7 @@ mod tests {
     use crate::metrics::SpillMetrics;
     use crate::spill::spill_manager::SpillManager;
     use crate::test::build_table_i32;
+    use arrow::array::ArrayRef;
     use arrow::array::{Float64Array, Int32Array, ListArray, StringArray};
     use arrow::compute::cast;
     use arrow::datatypes::{DataType, Field, Int32Type, Schema};
@@ -502,6 +503,121 @@ mod tests {
         Ok(())
     }
 
+    fn build_compressible_batch() -> RecordBatch {
+        let schema = Arc::new(Schema::new(vec![
+            Field::new("a", DataType::Utf8, false),
+            Field::new("b", DataType::Int32, false),
+            Field::new("c", DataType::Int32, true),
+        ]));
+
+        let a: ArrayRef = Arc::new(StringArray::from_iter_values(
+            std::iter::repeat("repeated").take(100),
+        ));
+        let b: ArrayRef = Arc::new(Int32Array::from(vec![1; 100]));
+        let c: ArrayRef = Arc::new(Int32Array::from(vec![2; 100]));
+
+        RecordBatch::try_new(schema, vec![a, b, c]).unwrap()
+    }
+
+    async fn validate(
+        spill_manager: &SpillManager,
+        spill_file: RefCountedTempFile,
+        num_rows: usize,
+        schema: SchemaRef,
+        batch_count: usize,
+    ) -> Result<()> {
+        let spilled_rows = spill_manager.metrics.spilled_rows.value();
+        assert_eq!(spilled_rows, num_rows);
+
+        let stream = spill_manager.read_spill_as_stream(spill_file)?;
+        assert_eq!(stream.schema(), schema);
+
+        let batches = collect(stream).await?;
+        assert_eq!(batches.len(), batch_count);
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_spill_compression() -> Result<()> {
+        let batch = build_compressible_batch();
+        let num_rows = batch.num_rows();
+        let schema = batch.schema();
+        let batch_count = 1;
+        let batches = [batch];
+
+        // Construct SpillManager
+        let env = Arc::new(RuntimeEnv::default());
+        let uncompressed_metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
+        let lz4_metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
+        let zstd_metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0);
+        let uncompressed_spill_manager = SpillManager::new(
+            env.clone(),
+            uncompressed_metrics,
+            Arc::clone(&schema),
+            SpillCompression::Uncompressed,
+        );
+        let lz4_spill_manager = SpillManager::new(
+            env.clone(),
+            lz4_metrics,
+            Arc::clone(&schema),
+            SpillCompression::Lz4Frame,
+        );
+        let zstd_spill_manager = SpillManager::new(
+            env,
+            zstd_metrics,
+            Arc::clone(&schema),
+            SpillCompression::Zstd,
+        );
+        let uncompressed_spill_file = uncompressed_spill_manager
+            .spill_record_batch_and_finish(&batches, "Test")?
+            .unwrap();
+        let lz4_spill_file = lz4_spill_manager
+            .spill_record_batch_and_finish(&batches, "Lz4_Test")?
+            .unwrap();
+        let zstd_spill_file = zstd_spill_manager
+            .spill_record_batch_and_finish(&batches, "ZSTD_Test")?
+            .unwrap();
+        assert!(uncompressed_spill_file.path().exists());
+        assert!(lz4_spill_file.path().exists());
+        assert!(zstd_spill_file.path().exists());
+
+        let lz4_spill_size = std::fs::metadata(lz4_spill_file.path())?.len();
+        let zstd_spill_size = std::fs::metadata(zstd_spill_file.path())?.len();
+        let uncompressed_spill_size =
+            std::fs::metadata(uncompressed_spill_file.path())?.len();
+        
+        assert!(uncompressed_spill_size > lz4_spill_size);
+        assert!(uncompressed_spill_size > zstd_spill_size);
+
+        // TODO validate with function
+        validate(
+            &lz4_spill_manager,
+            lz4_spill_file,
+            num_rows,
+            schema.clone(),
+            batch_count,
+        )
+        .await?;
+        validate(
+            &zstd_spill_manager,
+            zstd_spill_file,
+            num_rows,
+            schema.clone(),
+            batch_count,
+        )
+        .await?;
+        validate(
+            &uncompressed_spill_manager,
+            uncompressed_spill_file,
+            num_rows,
+            schema,
+            batch_count,
+        )
+        .await?;
+        Ok(())
+    }
+
     #[test]
     fn test_get_record_batch_memory_size() {
         // Create a simple record batch with two columns