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]