Skip to content

Commit 7ff3cbe

Browse files
authored
perf: Use await instead of block_on in native shuffle writer (#2937)
1 parent 53e4092 commit 7ff3cbe

File tree

1 file changed

+2
-3
lines changed

1 file changed

+2
-3
lines changed

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

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,6 @@ use datafusion::{
4444
},
4545
};
4646
use datafusion_comet_spark_expr::hash_funcs::murmur3::create_murmur3_hashes;
47-
use futures::executor::block_on;
4847
use futures::{StreamExt, TryFutureExt, TryStreamExt};
4948
use itertools::Itertools;
5049
use std::borrow::Borrow;
@@ -254,11 +253,11 @@ async fn external_shuffle(
254253
};
255254

256255
while let Some(batch) = input.next().await {
257-
// Block on the repartitioner to insert the batch and shuffle the rows
256+
// Await the repartitioner to insert the batch and shuffle the rows
258257
// into the corresponding partition buffer.
259258
// Otherwise, pull the next batch from the input stream might overwrite the
260259
// current batch in the repartitioner.
261-
block_on(repartitioner.insert_batch(batch?))?;
260+
repartitioner.insert_batch(batch?).await?;
262261
}
263262

264263
repartitioner.shuffle_write()?;

0 commit comments

Comments
 (0)