feat: pull input batches in parallel

viirya · Sep 16, 2024 · 5efd919 · 5efd919
1 parent 2108180
commit 5efd919
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 1 deletion.
diff --git a/native/Cargo.lock b/native/Cargo.lock
diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml
@@ -77,6 +77,7 @@ crc32fast = "1.3.2"
 simd-adler32 = "0.3.7"
 datafusion-comet-spark-expr = { workspace = true }
 datafusion-comet-proto = { workspace = true }
+rayon = "1.10.0"
 
 [dev-dependencies]
 pprof = { version = "0.13.0", features = ["flamegraph"] }

diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs
@@ -39,6 +39,8 @@ use jni::{
 };
 use std::{collections::HashMap, sync::Arc, task::Poll};
 
+use rayon::prelude::*;
+
 use super::{serde, utils::SparkArrowConvert, CometMemoryPool};
 
 use crate::{
@@ -322,7 +324,7 @@ fn prepare_output(
 /// operators before polling the stream,
 #[inline]
 fn pull_input_batches(exec_context: &mut ExecutionContext) -> Result<(), CometError> {
-    exec_context.scans.iter_mut().try_for_each(|scan| {
+    exec_context.scans.par_iter_mut().try_for_each(|scan| {
         scan.get_next_batch()?;
         Ok::<(), CometError>(())
     })