From 2b25d4a03992caf45d5f9d39d8b0d13f4cd29d28 Mon Sep 17 00:00:00 2001 From: Michael Taranov Date: Tue, 22 Sep 2026 00:01:28 +0300 Subject: [PATCH 1/3] fix: park the native scan loop instead of busy-polling while waiting on native I/O executePlan's ScanExec path re-polled the plan's stream in a tight loop whenever it returned Pending, relying on pull_input_batches to block on the JVM iterators in between. Once every JVM-fed scan holds a batch or has reached EOF that pull is a no-op, so a stream pending on native I/O (a Parquet or Iceberg scan reading from S3 or HDFS) spun the executor thread at 100% CPU for the whole read. A broadcast hash join over a native scan hits this on every probe-side read. pull_input_batches and the two scan operators now report whether a buffer was refilled. When the stream is Pending and nothing was pulled, the loop parks the block_on task until a waker registered by that poll fires, bounded by a short safety timeout, then re-enters the loop so JVM-fed scans still get refilled. Awaiting the stream directly is not safe: ScanExec returns Pending without a waker when an operator drains and re-polls it within one poll. The metrics interval is checked every iteration now that iterations are no longer spins. Closes #6091 Co-Authored-By: Claude Fable 5.1 --- native/core/Cargo.toml | 2 +- native/core/src/execution/jni_api.rs | 102 ++++++++++++------ native/core/src/execution/operators/scan.rs | 23 ++-- .../src/execution/operators/shuffle_scan.rs | 32 +++--- 4 files changed, 103 insertions(+), 56 deletions(-) diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index 501592945d2..d0d9d5c97d3 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -43,7 +43,7 @@ futures = { workspace = true } mimalloc = { version = "*", default-features = false, optional = true } tikv-jemallocator = { version = "0.6.1", optional = true, features = ["disable_initial_exec_tls"] } tikv-jemalloc-ctl = { version = "0.6.1", optional = true, features = ["disable_initial_exec_tls", "stats"] } -tokio = { version = "1", features = ["rt-multi-thread"] } +tokio = { version = "1", features = ["rt-multi-thread", "time"] } async-trait = { workspace = true } log = "0.4" log4rs = "1.4.0" diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 7652428410e..7c2753948e7 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -463,8 +463,6 @@ struct ExecutionContext { pub metrics_update_interval: Option, // 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, @@ -965,15 +962,38 @@ fn prepare_output( /// Java exception. So we pull input batches here and insert them into scan /// 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| { - scan.get_next_batch()?; - Ok::<(), CometError>(()) - })?; - exec_context.shuffle_scans.iter_mut().try_for_each(|scan| { - scan.get_next_batch()?; - Ok::<(), CometError>(()) - }) +fn pull_input_batches(exec_context: &mut ExecutionContext) -> Result { + let mut pulled = false; + for scan in exec_context.scans.iter_mut() { + pulled |= scan.get_next_batch()?; + } + for scan in exec_context.shuffle_scans.iter_mut() { + pulled |= scan.get_next_batch()?; + } + Ok(pulled) +} + +/// Safety net in case a stream ever returns Pending without registering a waker. +const PARK_TIMEOUT: Duration = Duration::from_millis(100); + +/// Sleeps until a waker registered by an earlier poll fires. +async fn park_until_woken() { + struct Park(bool); + + impl std::future::Future for Park { + type Output = (); + + fn poll(mut self: std::pin::Pin<&mut Self>, _: &mut std::task::Context<'_>) -> Poll<()> { + if self.0 { + Poll::Ready(()) + } else { + self.0 = true; + Poll::Pending + } + } + } + + let _ = tokio::time::timeout(PARK_TIMEOUT, Park(false)).await; } /// Accept serialized query plan and the addresses of Arrow Arrays from Spark, @@ -1113,30 +1133,26 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( } } - // ScanExec path: busy-poll to interleave JVM batch pulls with stream polling + // ScanExec path: JVM-fed scans return Pending without a waker and are refilled here. + // Nothing pulled means the stream waits on native I/O, so park instead of spinning. 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, - ); + 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)) => { @@ -1156,7 +1172,11 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( // 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))?; + let pulled = + tokio::task::block_in_place(|| pull_input_batches(exec_context))?; + if !pulled { + park_until_woken().await; + } } } } @@ -2143,4 +2163,26 @@ mod tests { assert_eq!(ret.data_type(), &DataType::Int32, "length({input})"); } } + #[test] + fn park_until_woken_ends_on_a_registered_waker_or_the_timeout() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .unwrap(); + runtime.block_on(async { + let (tx, mut rx) = tokio::sync::oneshot::channel::<()>(); + assert!(poll!(&mut rx).is_pending()); + tokio::spawn(async move { + let _ = tx.send(()); + }); + let start = Instant::now(); + park_until_woken().await; + assert!(start.elapsed() < PARK_TIMEOUT); + + let start = Instant::now(); + park_until_woken().await; + assert!(start.elapsed() >= PARK_TIMEOUT); + }); + } } diff --git a/native/core/src/execution/operators/scan.rs b/native/core/src/execution/operators/scan.rs index b1bf1990681..c6ec2d020b0 100644 --- a/native/core/src/execution/operators/scan.rs +++ b/native/core/src/execution/operators/scan.rs @@ -112,23 +112,26 @@ impl ScanExec { *self.batch.try_lock().unwrap() = Some(input); } - /// Pull next input batch from the upstream `ArrowArrayStreamReader`. - pub fn get_next_batch(&mut self) -> Result<(), CometError> { + /// Pulls the next input batch from the upstream `ArrowArrayStreamReader` unless one is + /// already buffered; returns whether it did. + pub fn get_next_batch(&mut self) -> Result { if self.input_source.is_none() { // This is a unit test. Input batches are seeded via `set_input_batch`. - return Ok(()); + return Ok(false); } let mut current_batch = self.batch.try_lock().unwrap(); - if current_batch.is_none() { - let mut timer = self.baseline_metrics.elapsed_compute().timer(); - let next_batch = - ScanExec::pull_next(self.exec_context_id, self.input_source.as_ref().unwrap())?; - *current_batch = Some(next_batch); - timer.stop(); + if current_batch.is_some() { + return Ok(false); } - Ok(()) + let mut timer = self.baseline_metrics.elapsed_compute().timer(); + let next_batch = + ScanExec::pull_next(self.exec_context_id, self.input_source.as_ref().unwrap())?; + *current_batch = Some(next_batch); + timer.stop(); + + Ok(true) } /// Pull the next `RecordBatch` from the stream and convert it to an `InputBatch`. Dictionary diff --git a/native/core/src/execution/operators/shuffle_scan.rs b/native/core/src/execution/operators/shuffle_scan.rs index b5a282f5a39..406521d9c5e 100644 --- a/native/core/src/execution/operators/shuffle_scan.rs +++ b/native/core/src/execution/operators/shuffle_scan.rs @@ -123,30 +123,32 @@ impl ShuffleScanExec { *self.batch.try_lock().unwrap() = Some(input); } - /// Pull next input batch from JVM. Called externally before poll_next() - /// because JNI calls cannot happen from within poll_next on tokio threads. - pub fn get_next_batch(&mut self) -> Result<(), CometError> { + /// Pulls the next input batch from the JVM unless one is already buffered; returns whether it + /// did. Called externally before poll_next() because JNI calls cannot happen from within + /// poll_next on tokio threads. + pub fn get_next_batch(&mut self) -> Result { if self.input_source.is_none() { // Unit test mode - no JNI calls needed. - return Ok(()); + return Ok(false); } - let mut timer = self.baseline_metrics.elapsed_compute().timer(); let mut current_batch = self.batch.try_lock().unwrap(); - if current_batch.is_none() { - let next_batch = Self::get_next( - self.exec_context_id, - self.input_source.as_ref().unwrap().as_obj(), - &self.data_types, - &self.decode_time, - self.requires_validation, - )?; - *current_batch = Some(next_batch); + if current_batch.is_some() { + return Ok(false); } + let mut timer = self.baseline_metrics.elapsed_compute().timer(); + let next_batch = Self::get_next( + self.exec_context_id, + self.input_source.as_ref().unwrap().as_obj(), + &self.data_types, + &self.decode_time, + self.requires_validation, + )?; + *current_batch = Some(next_batch); timer.stop(); - Ok(()) + Ok(true) } /// Invokes JNI calls to get the next compressed shuffle block and decode it. From 772d1b72515bc3cea922dc5ecf8742e5259b89ca Mon Sep 17 00:00:00 2001 From: Michael Taranov Date: Tue, 22 Sep 2026 12:42:27 +0300 Subject: [PATCH 2/3] fix: wake the scan streams on refill and park the loop without a timeout ScanStream and ShuffleScanStream now register the poll's waker when their buffer is empty and get_next_batch wakes it after refilling, so every Pending from the plan carries a waker. The loop in executePlan moves into next_batch(stream, on_pending): poll, refill and check metrics on Pending, then park until a waker fires. The 100 ms timeout, the bool from get_next_batch and the tokio time feature are gone. EOF stays buffered so a re-poll of an exhausted scan returns Ready(None) again instead of another JNI round trip. The tracing memory sample sits behind the metrics interval, and development.md describes the loop. Tests drive next_batch with a stream pending on a sleep (with the park removed it pulls 135,311 times in 50 ms) and with a ScanExec refilled by the pull closure under a timeout, plus a ShuffleScanStream waker test. Co-Authored-By: Claude Fable 5.1 --- docs/source/contributor-guide/development.md | 5 +- native/core/Cargo.toml | 2 +- native/core/src/execution/jni_api.rs | 227 +++++++++++------- native/core/src/execution/operators/scan.rs | 45 ++-- .../src/execution/operators/shuffle_scan.rs | 85 +++++-- 5 files changed, 234 insertions(+), 130 deletions(-) diff --git a/docs/source/contributor-guide/development.md b/docs/source/contributor-guide/development.md index b5752e08630..c5fa0113122 100644 --- a/docs/source/contributor-guide/development.md +++ b/docs/source/contributor-guide/development.md @@ -44,8 +44,9 @@ The executor thread parks in `blocking_recv()` until the next batch is ready. Th 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. +DataFusion stream directly. On `Poll::Pending` it calls `pull_input_batches()` to feed data from +the JVM into ScanExec operators, which wakes the stream, then parks until a waker fires, so a +stream that is waiting on native I/O sleeps instead of busy-polling. 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. diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index d0d9d5c97d3..501592945d2 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -43,7 +43,7 @@ futures = { workspace = true } mimalloc = { version = "*", default-features = false, optional = true } tikv-jemallocator = { version = "0.6.1", optional = true, features = ["disable_initial_exec_tls"] } tikv-jemalloc-ctl = { version = "0.6.1", optional = true, features = ["disable_initial_exec_tls", "stats"] } -tokio = { version = "1", features = ["rt-multi-thread", "time"] } +tokio = { version = "1", features = ["rt-multi-thread"] } async-trait = { workspace = true } log = "0.4" log4rs = "1.4.0" diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 7c2753948e7..5d89992add4 100644 --- a/native/core/src/execution/jni_api.rs +++ b/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; @@ -962,38 +962,51 @@ fn prepare_output( /// Java exception. So we pull input batches here and insert them into scan /// operators before polling the stream, #[inline] -fn pull_input_batches(exec_context: &mut ExecutionContext) -> Result { - let mut pulled = false; +fn pull_input_batches(exec_context: &mut ExecutionContext) -> Result<(), CometError> { for scan in exec_context.scans.iter_mut() { - pulled |= scan.get_next_batch()?; + scan.get_next_batch()?; } for scan in exec_context.shuffle_scans.iter_mut() { - pulled |= scan.get_next_batch()?; + scan.get_next_batch()?; } - Ok(pulled) + Ok(()) } -/// Safety net in case a stream ever returns Pending without registering a waker. -const PARK_TIMEOUT: Duration = Duration::from_millis(100); - -/// Sleeps until a waker registered by an earlier poll fires. +/// 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() { - struct Park(bool); - - impl std::future::Future for Park { - type Output = (); + let mut polled = false; + poll_fn(|_| { + if std::mem::replace(&mut polled, true) { + Poll::Ready(()) + } else { + Poll::Pending + } + }) + .await +} - fn poll(mut self: std::pin::Pin<&mut Self>, _: &mut std::task::Context<'_>) -> Poll<()> { - if self.0 { - Poll::Ready(()) - } else { - self.0 = true; - Poll::Pending +/// 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( + stream: &mut S, + mut on_pending: impl FnMut() -> Result<(), CometError>, +) -> Result, CometError> +where + S: Stream> + 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; } } } - - let _ = tokio::time::timeout(PARK_TIMEOUT, Park(false)).await; } /// Accept serialized query plan and the addresses of Arrow Arrays from Spark, @@ -1133,54 +1146,30 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( } } - // ScanExec path: JVM-fed scans return Pending without a waker and are refilled here. - // Nothing pulled means the stream waits on native I/O, so park instead of spinning. - get_runtime().block_on(async { - loop { - let next_item = exec_context.stream.as_mut().unwrap().next(); - let poll_output = poll!(next_item); - - 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. - let pulled = - tokio::task::block_in_place(|| pull_input_batches(exec_context))?; - if !pulled { - park_until_woken().await; - } - } - } + // 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 { @@ -1244,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(); @@ -1784,15 +1797,21 @@ 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; #[test] fn skip_partial_eligibility_is_fail_closed() { @@ -2163,26 +2182,64 @@ mod tests { assert_eq!(ret.data_type(), &DataType::Int32, "length({input})"); } } - #[test] - fn park_until_woken_ends_on_a_registered_waker_or_the_timeout() { - let runtime = tokio::runtime::Builder::new_multi_thread() + fn single_worker_runtime() -> Runtime { + tokio::runtime::Builder::new_multi_thread() .worker_threads(1) .enable_all() .build() + .unwrap() + } + + #[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(next_batch(&mut stream, || { + pulls += 1; + Ok(()) + })) .unwrap(); - runtime.block_on(async { - let (tx, mut rx) = tokio::sync::oneshot::channel::<()>(); - assert!(poll!(&mut rx).is_pending()); - tokio::spawn(async move { - let _ = tx.send(()); - }); - let start = Instant::now(); - park_until_woken().await; - assert!(start.elapsed() < PARK_TIMEOUT); + assert!(next.is_some()); + assert!( + pulls < 5, + "the loop pulled {pulls} times during one 50 ms wait" + ); + } - let start = Instant::now(); - park_until_woken().await; - assert!(start.elapsed() >= PARK_TIMEOUT); + #[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, so a lost wake would hang here. + single_worker_runtime().block_on(async { + tokio::time::timeout(Duration::from_secs(10), 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); + }) + .await + .unwrap(); }); } } diff --git a/native/core/src/execution/operators/scan.rs b/native/core/src/execution/operators/scan.rs index c6ec2d020b0..665f6d3e346 100644 --- a/native/core/src/execution/operators/scan.rs +++ b/native/core/src/execution/operators/scan.rs @@ -32,7 +32,7 @@ use datafusion::{ physical_plan::{ExecutionPlan, *}, }; use datafusion_comet_common::decode_string_arrays; -use futures::Stream; +use futures::{task::AtomicWaker, Stream}; use itertools::Itertools; use std::{ pin::Pin, @@ -57,6 +57,8 @@ pub struct ScanExec { /// Used in unit tests to mock the input batch; otherwise written by `pull_next` on each /// poll. pub batch: Arc>>, + /// Woken when `batch` is refilled, so a poll that found it empty is repeated. + waker: Arc, cache: Arc, metrics: ExecutionPlanMetricsSet, baseline_metrics: BaselineMetrics, @@ -90,6 +92,7 @@ impl ScanExec { input_source_description: input_source_description.to_string(), data_types, batch: Arc::new(Mutex::new(None)), + waker: Arc::new(AtomicWaker::new()), cache, metrics: metrics_set, baseline_metrics, @@ -110,19 +113,20 @@ impl ScanExec { /// Feeds input batch into this `Scan`. Only used in unit test. pub fn set_input_batch(&mut self, input: InputBatch) { *self.batch.try_lock().unwrap() = Some(input); + self.waker.wake(); } /// Pulls the next input batch from the upstream `ArrowArrayStreamReader` unless one is - /// already buffered; returns whether it did. - pub fn get_next_batch(&mut self) -> Result { + /// already buffered, then wakes the stream waiting for it. + pub fn get_next_batch(&mut self) -> Result<(), CometError> { if self.input_source.is_none() { // This is a unit test. Input batches are seeded via `set_input_batch`. - return Ok(false); + return Ok(()); } let mut current_batch = self.batch.try_lock().unwrap(); if current_batch.is_some() { - return Ok(false); + return Ok(()); } let mut timer = self.baseline_metrics.elapsed_compute().timer(); @@ -130,8 +134,10 @@ impl ScanExec { ScanExec::pull_next(self.exec_context_id, self.input_source.as_ref().unwrap())?; *current_batch = Some(next_batch); timer.stop(); + drop(current_batch); + self.waker.wake(); - Ok(true) + Ok(()) } /// Pull the next `RecordBatch` from the stream and convert it to an `InputBatch`. Dictionary @@ -326,28 +332,27 @@ impl ScanStream<'_> { impl Stream for ScanStream<'_> { type Item = DataFusionResult; - fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut timer = self.baseline_metrics.elapsed_compute().timer(); let mut scan_batch = self.scan.batch.try_lock().unwrap(); - let input_batch = &*scan_batch; - let input_batch = if let Some(batch) = input_batch { - batch - } else { - timer.stop(); - return Poll::Pending; - }; - - let result = match input_batch { - InputBatch::EOF => Poll::Ready(None), - InputBatch::Batch(columns, num_rows) => { + let result = match &*scan_batch { + None => { + self.scan.waker.register(cx.waker()); + Poll::Pending + } + // EOF stays buffered: a re-poll ends the stream again and `get_next_batch` has + // nothing to pull. + Some(InputBatch::EOF) => Poll::Ready(None), + Some(InputBatch::Batch(columns, num_rows)) => { self.baseline_metrics.record_output(*num_rows); let maybe_batch = self.build_record_batch(columns, *num_rows); Poll::Ready(Some(maybe_batch)) } }; - - *scan_batch = None; + if matches!(result, Poll::Ready(Some(_))) { + *scan_batch = None; + } timer.stop(); diff --git a/native/core/src/execution/operators/shuffle_scan.rs b/native/core/src/execution/operators/shuffle_scan.rs index 406521d9c5e..05f26583e54 100644 --- a/native/core/src/execution/operators/shuffle_scan.rs +++ b/native/core/src/execution/operators/shuffle_scan.rs @@ -38,7 +38,7 @@ use datafusion::{ physical_plan::{ExecutionPlan, *}, }; use datafusion_comet_common::cast_and_stamp_schema; -use futures::Stream; +use futures::{task::AtomicWaker, Stream}; use jni::objects::{Global, JByteBuffer, JObject}; use std::{ pin::Pin, @@ -63,6 +63,8 @@ pub struct ShuffleScanExec { pub schema: SchemaRef, /// The current input batch, populated by get_next_batch() before poll_next(). pub batch: Arc>>, + /// Woken when `batch` is refilled, so a poll that found it empty is repeated. + waker: Arc, /// Cache of plan properties. cache: Arc, /// Metrics collector. @@ -109,6 +111,7 @@ impl ShuffleScanExec { input_source, data_types, batch: Arc::new(Mutex::new(None)), + waker: Arc::new(AtomicWaker::new()), cache, metrics: metrics_set, baseline_metrics, @@ -121,20 +124,21 @@ impl ShuffleScanExec { /// Feeds input batch into this scan. Only used in unit tests. pub fn set_input_batch(&mut self, input: InputBatch) { *self.batch.try_lock().unwrap() = Some(input); + self.waker.wake(); } - /// Pulls the next input batch from the JVM unless one is already buffered; returns whether it - /// did. Called externally before poll_next() because JNI calls cannot happen from within - /// poll_next on tokio threads. - pub fn get_next_batch(&mut self) -> Result { + /// Pulls the next input batch from the JVM unless one is already buffered, then wakes the + /// stream waiting for it. Called externally before poll_next() because JNI calls cannot + /// happen from within poll_next on tokio threads. + pub fn get_next_batch(&mut self) -> Result<(), CometError> { if self.input_source.is_none() { // Unit test mode - no JNI calls needed. - return Ok(false); + return Ok(()); } let mut current_batch = self.batch.try_lock().unwrap(); if current_batch.is_some() { - return Ok(false); + return Ok(()); } let mut timer = self.baseline_metrics.elapsed_compute().timer(); @@ -147,8 +151,10 @@ impl ShuffleScanExec { )?; *current_batch = Some(next_batch); timer.stop(); + drop(current_batch); + self.waker.wake(); - Ok(true) + Ok(()) } /// Invokes JNI calls to get the next compressed shuffle block and decode it. @@ -364,21 +370,19 @@ impl ShuffleScanStream { impl Stream for ShuffleScanStream { type Item = DataFusionResult; - fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut timer = self.baseline_metrics.elapsed_compute().timer(); let mut scan_batch = self.shuffle_scan.batch.try_lock().unwrap(); - let input_batch = &*scan_batch; - let input_batch = if let Some(batch) = input_batch { - batch - } else { - timer.stop(); - return Poll::Pending; - }; - - let result = match input_batch { - InputBatch::EOF => Poll::Ready(None), - InputBatch::Batch(columns, num_rows) => { + let result = match &*scan_batch { + None => { + self.shuffle_scan.waker.register(cx.waker()); + Poll::Pending + } + // EOF stays buffered: a re-poll ends the stream again and `get_next_batch` has + // nothing to pull. + Some(InputBatch::EOF) => Poll::Ready(None), + Some(InputBatch::Batch(columns, num_rows)) => { self.baseline_metrics.record_output(*num_rows); // Reconcile the decoded block with the catalyst-declared schema rather than // stamping it on, so that nested field nullability drift is absorbed here the way @@ -393,8 +397,9 @@ impl Stream for ShuffleScanStream { Poll::Ready(Some(maybe_batch)) } }; - - *scan_batch = None; + if matches!(result, Poll::Ready(Some(_))) { + *scan_batch = None; + } timer.stop(); @@ -759,4 +764,40 @@ mod tests { assert!(err.contains("col_0: expected Struct"), "{err}"); }); } + + #[test] + fn refill_wakes_the_pending_poll_and_eof_stays_buffered() { + use super::*; + use crate::execution::planner::TEST_EXEC_CONTEXT_ID; + use datafusion::physical_plan::ExecutionPlan; + use futures::task::{waker, ArcWake}; + use std::sync::atomic::{AtomicBool, Ordering}; + + struct Woken(AtomicBool); + impl ArcWake for Woken { + fn wake_by_ref(arc_self: &Arc) { + arc_self.0.store(true, Ordering::SeqCst); + } + } + let woken = Arc::new(Woken(AtomicBool::new(false))); + let waker = waker(Arc::clone(&woken)); + let mut cx = Context::from_waker(&waker); + + let mut scan = + ShuffleScanExec::new(TEST_EXEC_CONTEXT_ID, None, vec![DataType::Int32]).unwrap(); + let mut stream = scan.execute(0, Arc::new(TaskContext::default())).unwrap(); + + assert!(stream.as_mut().poll_next(&mut cx).is_pending()); + assert!(!woken.0.load(Ordering::SeqCst)); + scan.set_input_batch(InputBatch::EOF); + assert!(woken.0.load(Ordering::SeqCst)); + assert!(matches!( + stream.as_mut().poll_next(&mut cx), + Poll::Ready(None) + )); + assert!(matches!( + stream.as_mut().poll_next(&mut cx), + Poll::Ready(None) + )); + } } From 013d566fa7b8cb45ea36681e9fbb800fe0d8c9e6 Mon Sep 17 00:00:00 2001 From: Michael Taranov Date: Tue, 22 Sep 2026 13:47:18 +0300 Subject: [PATCH 3/3] fix: bound the sleep-based loop test and correct the threading notes The JVM data source path polls operators on the Spark executor thread inside block_on; only tasks they spawn run on tokio workers. The heading now names ShuffleScanExec as well, since pull_input_batches feeds both streams and both register a waker. The native-wait test gets the same ten second bound as the refill test, so a lost wake fails instead of hanging the suite. Co-Authored-By: Claude Fable 5.1 --- docs/source/contributor-guide/development.md | 16 +++++---- native/core/src/execution/jni_api.rs | 36 +++++++++++--------- 2 files changed, 29 insertions(+), 23 deletions(-) diff --git a/docs/source/contributor-guide/development.md b/docs/source/contributor-guide/development.md index c5fa0113122..26295f93bca 100644 --- a/docs/source/contributor-guide/development.md +++ b/docs/source/contributor-guide/development.md @@ -43,13 +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. On `Poll::Pending` it calls `pull_input_batches()` to feed data from -the JVM into ScanExec operators, which wakes the stream, then parks until a waker fires, so a -stream that is waiting on native I/O sleeps instead of busy-polling. - -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 diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 5d89992add4..0f2bbec16b8 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -1812,6 +1812,7 @@ mod tests { 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() { @@ -2190,6 +2191,13 @@ mod tests { .unwrap() } + /// Fails instead of hanging the suite when a wake is lost. + async fn within_ten_seconds(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())); @@ -2200,10 +2208,10 @@ mod tests { .boxed(); let mut pulls = 0; let next = single_worker_runtime() - .block_on(next_batch(&mut stream, || { + .block_on(within_ten_seconds(next_batch(&mut stream, || { pulls += 1; Ok(()) - })) + }))) .unwrap(); assert!(next.is_some()); assert!( @@ -2227,19 +2235,15 @@ mod tests { } Ok::<(), CometError>(()) }; - // Only the refill's wake ends each park, so a lost wake would hang here. - single_worker_runtime().block_on(async { - tokio::time::timeout(Duration::from_secs(10), 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); - }) - .await - .unwrap(); - }); + // 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); + })); } }