kosiew commented on code in PR #26365:
URL: https://github.com/apache/datafusion/pull/26365#discussion_r4535575784


##########
datafusion/common/src/config.rs:
##########
@@ -5459,6 +5479,51 @@ mod tests {
         assert_eq!(MaxRowGroupBytes::from_str("128").unwrap().get(), 128);
     }
 
+    #[test]
+    fn spill_compression_zstd_carries_its_level() {
+        use crate::config::SpillCompression;
+
+        for (input, expected, displayed) in [
+            (
+                "zstd",
+                SpillCompression::Zstd(SpillCompression::DEFAULT_ZSTD_LEVEL),
+                "zstd(3)",
+            ),
+            ("ZSTD(1)", SpillCompression::Zstd(1), "zstd(1)"),
+            ("zstd(22)", SpillCompression::Zstd(22), "zstd(22)"),

Review Comment:
   Consider adding `("zstd(0)", SpillCompression::Zstd(0), "zstd(0)")` to the 
round-trip cases. Zero is the documented lower boundary and default-level 
alias, so this would verify that parsing and display preserve the explicit zero 
setting.



##########
datafusion/physical-plan/src/spill/mod.rs:
##########
@@ -1991,6 +2002,47 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_spill_compression_zstd_level() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new("a", 
DataType::Int32, false)]));
+        let values =
+            (0..64 * 1024u32).map(|i| (i.wrapping_mul(2654435761) >> 7) as i32 
% 1024);
+        let batch = RecordBatch::try_new(
+            Arc::clone(&schema),
+            vec![Arc::new(Int32Array::from_iter_values(values))],
+        )?;
+        let num_rows = batch.num_rows();
+        let batches = [batch];
+        let env = Arc::new(RuntimeEnv::default());
+
+        let spill = |zstd_level: u32| -> Result<(SpillManager, Arc<dyn 
SpillFile>)> {
+            let spill_manager = SpillManager::new(
+                Arc::clone(&env),
+                SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
+                Arc::clone(&schema),
+            )
+            .with_compression_type(SpillCompression::Zstd(zstd_level));
+            let file = spill_manager
+                .spill_record_batch_and_finish(&batches, "ZSTD_Level_Test")?
+                .unwrap();
+            Ok((spill_manager, file))
+        };
+
+        let (_, default_file) = spill(SpillCompression::DEFAULT_ZSTD_LEVEL)?;
+        let default_bytes = std::fs::read(default_file.path().unwrap())?;
+
+        for level in [1, 19] {
+            let (spill_manager, file) = spill(level)?;
+            assert_ne!(
+                std::fs::read(file.path().unwrap())?,
+                default_bytes,
+                "level {level}"
+            );
+            validate(&spill_manager, file, num_rows, Arc::clone(&schema), 
1).await?;
+        }

Review Comment:
   Consider collecting the decoded stream and comparing its batches with 
`batches`, while retaining the existing metrics and schema checks. The current 
`validate` helper checks row metrics, schema, and batch count but not decoded 
values, so this would strengthen round-trip coverage for levels 1 and 19.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to