Main upstream dd 04 02 2026

Cargo.toml

@@ -51,7 +51,7 @@ arrow = { version = "57.1.0", optional = true }
 hyper-util = { version = "0.1.16", optional = true }
 pretty_assertions = { version = "1.4", optional = true }
 reqwest = { version = "0.12", optional = true }
-zip = { version = "6.0", optional = true }
+zip = { version = "4.0", optional = true }
 tonic-prost = "0.14.2"
 
 [features]
@@ -83,5 +83,5 @@ tokio-stream = "0.1.17"
 hyper-util = "0.1.16"
 pretty_assertions = "1.4"
 reqwest = "0.12"
-zip = "6.0"
+zip = "4.0"
 test-case = "3.3.1"

src/distributed_ext.rs

@@ -469,6 +469,24 @@ pub trait DistributedExt: Sized {
         &mut self,
         compression: Option<CompressionType>,
     ) -> Result<(), DataFusionError>;
+
+    /// How many rows to collect in each record batch before sending it over the wire in a
+    /// shuffle operation. This value defaults to the same as `datafusion.execution.batch_size`.
+    ///
+    /// Setting it to something smaller than `datafusion.execution.batch_size` has no effect.
+    ///
+    /// It's preferable to set `datafusion.execution.batch_size` directly instead of this
+    /// parameter if the specific use case allows it.
+    fn with_distributed_shuffle_batch_size(
+        self,
+        batch_size: usize,
+    ) -> Result<Self, DataFusionError>;
+
+    /// Same as [DistributedExt::with_distributed_shuffle_batch_size] but with an in-place mutation.
+    fn set_distributed_shuffle_batch_size(
+        &mut self,
+        batch_size: usize,
+    ) -> Result<(), DataFusionError>;
 }
 
 impl DistributedExt for SessionConfig {
@@ -565,6 +583,15 @@ impl DistributedExt for SessionConfig {
         Ok(())
     }
 
+    fn set_distributed_shuffle_batch_size(
+        &mut self,
+        batch_size: usize,
+    ) -> Result<(), DataFusionError> {
+        let d_cfg = DistributedConfig::from_config_options_mut(self.options_mut())?;
+        d_cfg.shuffle_batch_size = batch_size;
+        Ok(())
+    }
+
     delegate! {
         to self {
             #[call(set_distributed_option_extension)]
@@ -618,6 +645,10 @@ impl DistributedExt for SessionConfig {
             #[call(set_distributed_compression)]
             #[expr($?;Ok(self))]
             fn with_distributed_compression(mut self, compression: Option<CompressionType>) -> Result<Self, DataFusionError>;
+
+            #[call(set_distributed_shuffle_batch_size)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_shuffle_batch_size(mut self, batch_size: usize) -> Result<Self, DataFusionError>;
         }
     }
 }
@@ -689,6 +720,11 @@ impl DistributedExt for SessionStateBuilder {
             #[call(set_distributed_compression)]
             #[expr($?;Ok(self))]
             fn with_distributed_compression(mut self, compression: Option<CompressionType>) -> Result<Self, DataFusionError>;
+
+            fn set_distributed_shuffle_batch_size(&mut self, batch_size: usize) -> Result<(), DataFusionError>;
+            #[call(set_distributed_shuffle_batch_size)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_shuffle_batch_size(mut self, batch_size: usize) -> Result<Self, DataFusionError>;
         }
     }
 }
@@ -760,6 +796,11 @@ impl DistributedExt for SessionState {
             #[call(set_distributed_compression)]
             #[expr($?;Ok(self))]
             fn with_distributed_compression(mut self, compression: Option<CompressionType>) -> Result<Self, DataFusionError>;
+
+            fn set_distributed_shuffle_batch_size(&mut self, batch_size: usize) -> Result<(), DataFusionError>;
+            #[call(set_distributed_shuffle_batch_size)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_shuffle_batch_size(mut self, batch_size: usize) -> Result<Self, DataFusionError>;
         }
     }
 }
@@ -831,6 +872,11 @@ impl DistributedExt for SessionContext {
             #[call(set_distributed_compression)]
             #[expr($?;Ok(self))]
             fn with_distributed_compression(self, compression: Option<CompressionType>) -> Result<Self, DataFusionError>;
+
+            fn set_distributed_shuffle_batch_size(&mut self, batch_size: usize) -> Result<(), DataFusionError>;
+            #[call(set_distributed_shuffle_batch_size)]
+            #[expr($?;Ok(self))]
+            fn with_distributed_shuffle_batch_size(self, batch_size: usize) -> Result<Self, DataFusionError>;
         }
     }
 }