fix

Author: viirya

native/core/src/execution/datafusion/shuffle_writer.rs

@@ -56,6 +56,7 @@ use datafusion::{
 };
 use datafusion_physical_expr::EquivalenceProperties;
 use futures::{lock::Mutex, Stream, StreamExt, TryFutureExt, TryStreamExt};
+use futures::executor::block_on;
 use itertools::Itertools;
 use simd_adler32::Adler32;
 use tokio::task;
@@ -1111,7 +1112,11 @@ async fn external_shuffle(
     );
 
     while let Some(batch) = input.next().await {
-        repartitioner.insert_batch(batch?).await?;
+        // Block on the repartitioner to insert the batch and shuffle the rows
+        // into the corresponding partition buffer.
+        // Otherwise, pull the next batch from the input stream might overwrite the
+        // current batch in the repartitioner.
+        block_on(repartitioner.insert_batch(batch?))?;
     }
     repartitioner.shuffle_write().await
 }