fix: park the native scan loop instead of busy-polling while waiting on native I/O

docs/source/contributor-guide/development.md

@@ -43,12 +43,15 @@ onto a tokio worker thread and batches are delivered to the executor thread via
 The executor thread parks in `blocking_recv()` until the next batch is ready. This avoids
 busy-polling on I/O-bound workloads.
 
-**JVM data source path (ScanExec present):** The executor thread calls `block_on()` and polls the
-DataFusion stream directly, interleaving `pull_input_batches()` calls on `Poll::Pending` to feed
-data from the JVM into ScanExec operators.
-
-In both cases, DataFusion operators execute on **tokio worker threads**, not on the Spark executor
-task thread. All Spark tasks on an executor share one tokio runtime.
+**JVM data source path (ScanExec or ShuffleScanExec present):** The executor thread calls
+`block_on()` and polls the DataFusion stream directly. On `Poll::Pending` it calls
+`pull_input_batches()` to feed data from the JVM into the ScanExec and ShuffleScanExec operators,
+whose streams register the poll's waker and are woken by the refill, then parks until a waker
+fires, so a stream that is waiting on native I/O sleeps instead of busy-polling.
+
+On the async I/O path, DataFusion operators execute on **tokio worker threads**. On the JVM data
+source path, `block_on()` polls them on the Spark executor task thread, and any tasks they spawn
+run on the shared runtime. All Spark tasks on an executor share one tokio runtime.
 
 ### Rules for native code
 

native/core/src/execution/jni_api.rs

@@ -78,7 +78,7 @@ use datafusion_spark::function::url::try_url_decode::TryUrlDecode as SparkTryUrl
 use datafusion_spark::function::url::url_decode::UrlDecode as SparkUrlDecode;
 use datafusion_spark::function::url::url_encode::UrlEncode as SparkUrlEncode;
 use futures::poll;
-use futures::stream::StreamExt;
+use futures::stream::{Stream, StreamExt};
 use futures::FutureExt;
 use jni::objects::JByteBuffer;
 use jni::sys::{jlongArray, JNI_FALSE};
@@ -96,7 +96,7 @@ use prost::Message;
 use std::collections::HashMap;
 use std::path::PathBuf;
 use std::time::{Duration, Instant};
-use std::{sync::Arc, task::Poll};
+use std::{future::poll_fn, sync::Arc, task::Poll};
 use tokio::runtime::{Handle, Runtime};
 use tokio::sync::mpsc;
 
@@ -463,8 +463,6 @@ struct ExecutionContext {
     pub metrics_update_interval: Option<Duration>,
     // The last update time of metrics
     pub metrics_last_update_time: Instant,
-    /// Counter to avoid checking time on every poll iteration (reduces syscalls)
-    pub poll_count_since_metrics_check: u32,
     /// The time it took to create the native plan and configure the context
     pub plan_creation_time: Duration,
     /// DataFusion SessionContext
@@ -680,7 +678,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan(
                 metrics,
                 metrics_update_interval,
                 metrics_last_update_time: Instant::now(),
-                poll_count_since_metrics_check: 0,
                 plan_creation_time,
                 session_ctx: session,
                 debug_native,
@@ -966,14 +963,50 @@ 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| {
+    for scan in exec_context.scans.iter_mut() {
         scan.get_next_batch()?;
-        Ok::<(), CometError>(())
-    })?;
-    exec_context.shuffle_scans.iter_mut().try_for_each(|scan| {
+    }
+    for scan in exec_context.shuffle_scans.iter_mut() {
         scan.get_next_batch()?;
-        Ok::<(), CometError>(())
+    }
+    Ok(())
+}
+
+/// Yields once, so the `block_on` thread sleeps until a waker registered by an earlier poll
+/// fires: a JVM-fed scan refilled by `pull_input_batches`, or native I/O that completed.
+async fn park_until_woken() {
+    let mut polled = false;
+    poll_fn(|_| {
+        if std::mem::replace(&mut polled, true) {
+            Poll::Ready(())
+        } else {
+            Poll::Pending
+        }
     })
+    .await
+}
+
+/// Drives `stream` to its next item. JVM-fed scans return `Pending` until `on_pending` refills
+/// them, so every pending poll runs it and then parks instead of polling again at once.
+async fn next_batch<S>(
+    stream: &mut S,
+    mut on_pending: impl FnMut() -> Result<(), CometError>,
+) -> Result<Option<RecordBatch>, CometError>
+where
+    S: Stream<Item = DataFusionResult<RecordBatch>> + Unpin,
+{
+    loop {
+        match poll!(stream.next()) {
+            Poll::Ready(item) => return Ok(item.transpose()?),
+            Poll::Pending => {
+                // JNI call to pull batches from JVM into ScanExec operators.
+                // block_in_place lets tokio move other tasks off this worker
+                // while we wait for JVM data.
+                tokio::task::block_in_place(&mut on_pending)?;
+                park_until_woken().await;
+            }
+        }
+    }
 }
 
 /// Accept serialized query plan and the addresses of Arrow Arrays from Spark,
@@ -1113,54 +1146,30 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan(
                 }
             }
 
-            // ScanExec path: busy-poll to interleave JVM batch pulls with stream polling
-            get_runtime().block_on(async {
-                loop {
-                    let next_item = exec_context.stream.as_mut().unwrap().next();
-                    let poll_output = poll!(next_item);
-
-                    // Only check time/tracing every 100 polls to reduce overhead
-                    exec_context.poll_count_since_metrics_check += 1;
-                    if exec_context.poll_count_since_metrics_check >= 100 {
-                        exec_context.poll_count_since_metrics_check = 0;
-                        if let Some(interval) = exec_context.metrics_update_interval {
-                            let now = Instant::now();
-                            if now - exec_context.metrics_last_update_time >= interval {
-                                update_metrics(env, exec_context)?;
-                                exec_context.metrics_last_update_time = now;
-                            }
-                        }
-                        if exec_context.tracing_enabled {
-                            log_memory_usage(
-                                &exec_context.tracing_memory_metric_name,
-                                total_reserved_for_thread(exec_context.rust_thread_id) as u64,
-                            );
-                        }
-                    }
-
-                    match poll_output {
-                        Poll::Ready(Some(output)) => {
-                            return prepare_output(
-                                env,
-                                array_addrs,
-                                schema_addrs,
-                                output?,
-                                exec_context.debug_native,
-                            );
-                        }
-                        Poll::Ready(None) => {
-                            log_plan_metrics(exec_context, stage_id, partition);
-                            return Ok(-1);
-                        }
-                        Poll::Pending => {
-                            // JNI call to pull batches from JVM into ScanExec operators.
-                            // block_in_place lets tokio move other tasks off this worker
-                            // while we wait for JVM data.
-                            tokio::task::block_in_place(|| pull_input_batches(exec_context))?;
-                        }
-                    }
+            // ScanExec path: JVM-fed scans return `Pending` until `pull_input_batches` refills
+            // them and wakes the stream, so a poll that is still pending after the pull waits on
+            // native I/O and the loop parks for it.
+            let mut stream = exec_context.stream.take().unwrap();
+            let next = get_runtime().block_on(next_batch(&mut stream, || {
+                pull_input_batches(exec_context)?;
+                update_metrics_on_interval(env, exec_context)
+            }));
+            exec_context.stream = Some(stream);
+            let next = next?;
+            update_metrics_on_interval(env, exec_context)?;
+            match next {
+                Some(batch) => prepare_output(
+                    env,
+                    array_addrs,
+                    schema_addrs,
+                    batch,
+                    exec_context.debug_native,
+                ),
+                None => {
+                    log_plan_metrics(exec_context, stage_id, partition);
+                    Ok(-1)
                 }
-            })
+            }
         });
 
         if exec_context.tracing_enabled {
@@ -1224,6 +1233,30 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan(
     })
 }
 
+/// Runs `update_metrics` once the configured interval has passed and, with tracing on, samples
+/// this thread's pool reservation at the same cadence.
+fn update_metrics_on_interval(
+    env: &mut Env,
+    exec_context: &mut ExecutionContext,
+) -> CometResult<()> {
+    let Some(interval) = exec_context.metrics_update_interval else {
+        return Ok(());
+    };
+    let now = Instant::now();
+    if now - exec_context.metrics_last_update_time < interval {
+        return Ok(());
+    }
+    update_metrics(env, exec_context)?;
+    exec_context.metrics_last_update_time = now;
+    if exec_context.tracing_enabled {
+        log_memory_usage(
+            &exec_context.tracing_memory_metric_name,
+            total_reserved_for_thread(exec_context.rust_thread_id) as u64,
+        );
+    }
+    Ok(())
+}
+
 fn update_metrics(env: &mut Env, exec_context: &mut ExecutionContext) -> CometResult<()> {
     if let Some(native_query) = &exec_context.root_op {
         let metrics = exec_context.metrics.as_obj();
@@ -1764,15 +1797,22 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowClose(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::execution::operators::InputBatch;
+    use crate::execution::planner::TEST_EXEC_CONTEXT_ID;
+    use arrow::array::{ArrayRef, Int32Array};
     use arrow::datatypes::{DataType, Field, Schema};
     use datafusion::execution::memory_pool::{
         MemoryConsumer, MemoryReservation, UnboundedMemoryPool,
     };
     use datafusion::execution::FunctionRegistry;
+    use datafusion::execution::TaskContext;
     use datafusion::logical_expr::ReturnFieldArgs;
+    use datafusion::physical_plan::ExecutionPlan;
     use datafusion_comet_proto::spark_expression;
     use datafusion_comet_proto::spark_expression::{AggExpr, Count, Expr, Sum};
     use datafusion_comet_proto::spark_operator::{HashAggregate, ShuffleWriter};
+    use std::cell::Cell;
+    use std::future::Future;
 
     #[test]
     fn skip_partial_eligibility_is_fail_closed() {
@@ -2143,4 +2183,67 @@ mod tests {
             assert_eq!(ret.data_type(), &DataType::Int32, "length({input})");
         }
     }
+    fn single_worker_runtime() -> Runtime {
+        tokio::runtime::Builder::new_multi_thread()
+            .worker_threads(1)
+            .enable_all()
+            .build()
+            .unwrap()
+    }
+
+    /// Fails instead of hanging the suite when a wake is lost.
+    async fn within_ten_seconds<F: Future>(future: F) -> F::Output {
+        tokio::time::timeout(Duration::from_secs(10), future)
+            .await
+            .expect("timed out: a wake was lost")
+    }
+
+    #[test]
+    fn next_batch_parks_while_the_stream_waits_on_native_io() {
+        let batch = RecordBatch::new_empty(Arc::new(Schema::empty()));
+        let mut stream = futures::stream::once(async move {
+            tokio::time::sleep(Duration::from_millis(50)).await;
+            Ok::<_, DataFusionError>(batch)
+        })
+        .boxed();
+        let mut pulls = 0;
+        let next = single_worker_runtime()
+            .block_on(within_ten_seconds(next_batch(&mut stream, || {
+                pulls += 1;
+                Ok(())
+            })))
+            .unwrap();
+        assert!(next.is_some());
+        assert!(
+            pulls < 5,
+            "the loop pulled {pulls} times during one 50 ms wait"
+        );
+    }
+
+    #[test]
+    fn next_batch_resumes_on_a_refill_and_stops_pulling_after_eof() {
+        let mut scan =
+            ScanExec::new(TEST_EXEC_CONTEXT_ID, None, "", vec![DataType::Int32]).unwrap();
+        let mut stream = scan.execute(0, Arc::new(TaskContext::default())).unwrap();
+        let column: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3]));
+        let mut inputs = vec![InputBatch::new(vec![column], Some(3)), InputBatch::EOF].into_iter();
+        let pulls = Cell::new(0);
+        let mut pull = || {
+            pulls.set(pulls.get() + 1);
+            if let Some(input) = inputs.next() {
+                scan.set_input_batch(input);
+            }
+            Ok::<(), CometError>(())
+        };
+        // Only the refill's wake ends each park.
+        single_worker_runtime().block_on(within_ten_seconds(async {
+            let first = next_batch(&mut stream, &mut pull).await.unwrap();
+            assert_eq!(first.unwrap().num_rows(), 3);
+            assert_eq!(pulls.get(), 1);
+            assert!(next_batch(&mut stream, &mut pull).await.unwrap().is_none());
+            assert_eq!(pulls.get(), 2);
+            assert!(next_batch(&mut stream, &mut pull).await.unwrap().is_none());
+            assert_eq!(pulls.get(), 2);
+        }));
+    }
 }