diff --git a/docs/source/contributor-guide/development.md b/docs/source/contributor-guide/development.md index 2bf3bd6ec8d..e817baddaa0 100644 --- a/docs/source/contributor-guide/development.md +++ b/docs/source/contributor-guide/development.md @@ -43,12 +43,19 @@ 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. The stream is polled with a +waker that also sets a flag. If the flag is still clear when the pull returns, the stream is +waiting on native I/O, and the thread parks until a waker fires instead of busy-polling. If it is +set, the loop polls again. It checks the flag rather than trusting the thread's parker because the +pull can run another Comet plan on the same thread, and that plan's `block_on()` shares the parker +and can consume the wake-up meant for the outer loop. + +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 58011c4a652..8c9e1da9984 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -77,8 +77,7 @@ use datafusion_spark::function::string::substring::SparkSubstring; use datafusion_spark::function::url::try_url_decode::TryUrlDecode as SparkTryUrlDecode; 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}; @@ -95,8 +94,13 @@ use parking_lot::Mutex; use prost::Message; use std::collections::HashMap; use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::{Duration, Instant}; -use std::{sync::Arc, task::Poll}; +use std::{ + future::poll_fn, + sync::Arc, + task::{Context, Poll, Wake, Waker}, +}; use tokio::runtime::{Handle, Runtime}; use tokio::sync::mpsc; @@ -508,8 +512,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 @@ -725,7 +727,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, @@ -1021,6 +1022,67 @@ fn pull_input_batches(exec_context: &mut ExecutionContext) -> Result<(), CometEr }) } +/// Forwards a wake-up to the `block_on` task and records that it happened, so `next_batch` can +/// tell whether the stream was woken even if something else took the wake-up from the thread's +/// parker. +struct WakeFlag { + woken: AtomicBool, + parent: Waker, +} + +impl Wake for WakeFlag { + fn wake(self: Arc) { + self.wake_by_ref() + } + + fn wake_by_ref(self: &Arc) { + self.woken.store(true, Ordering::Release); + self.parent.wake_by_ref(); + } +} + +/// Drives `stream` to its next item. JVM-fed scans return `Pending` until `on_pending` refills +/// them, so every pending poll runs it, and the refill wakes the stream. +/// +/// Each poll gives the stream a `WakeFlag` waker. If the flag is still clear when `on_pending` +/// returns, the stream is waiting on native I/O and `block_on` parks until it completes. +/// Otherwise `next_batch` wakes `block_on` itself so that it polls again at once. It can't rely on +/// the wake-up that set the flag, because `on_pending` can run another Comet plan on this thread, +/// as when a native writer's input is itself native. That plan's `block_on` shares this thread's +/// parker, which holds a single wake-up, and its park can take this one. +/// +/// It polls again by yielding to `block_on` rather than looping here, so that each poll starts +/// with a fresh coop budget. A stream that has spent its budget wakes itself and returns +/// `Pending`, and a loop here would only get past that because `block_in_place` happens to leave +/// this thread's budget unconstrained. +async fn next_batch( + stream: &mut S, + mut on_pending: impl FnMut() -> Result<(), CometError>, +) -> Result, CometError> +where + S: Stream> + Unpin, +{ + poll_fn(|cx| { + let flag = Arc::new(WakeFlag { + woken: AtomicBool::new(false), + parent: cx.waker().clone(), + }); + let waker = Waker::from(Arc::clone(&flag)); + if let Poll::Ready(item) = stream.poll_next_unpin(&mut Context::from_waker(&waker)) { + return Poll::Ready(Ok(item.transpose()?)); + } + // `on_pending` calls into the JVM, which can run another Comet plan on this thread. + // `block_in_place` exits the runtime context so that plan's `block_on` doesn't panic. + tokio::task::block_in_place(&mut on_pending)?; + if flag.woken.load(Ordering::Acquire) { + // Poll again at once: a nested `block_on` may have taken the wake-up. + cx.waker().wake_by_ref(); + } + Poll::Pending + }) + .await +} + /// Accept serialized query plan and the addresses of Arrow Arrays from Spark, /// then execute the query. Return addresses of arrow vector. /// # Safety @@ -1158,54 +1220,31 @@ 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. A poll that is still pending, with nothing having woken + // the stream by the end of the pull, waits on native I/O, and `next_batch` parks + // until it completes. + 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 { @@ -1267,6 +1306,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(); @@ -1829,15 +1892,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() { @@ -2314,4 +2384,104 @@ 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 when a wake is lost. The timeout keeps a lost wake from hanging the suite, but when + /// its timer fires it polls `future` again, which can finish it, so this also fails when + /// `future` took more than five seconds. + async fn without_a_lost_wake(future: F) -> F::Output { + let start = Instant::now(); + let output = tokio::time::timeout(Duration::from_secs(10), future) + .await + .expect("timed out: a wake was lost"); + let elapsed = start.elapsed(); + assert!( + elapsed < Duration::from_secs(5), + "took {elapsed:?}: a wake was lost, and only the timeout's timer woke the task" + ); + output + } + + #[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(without_a_lost_wake(next_batch(&mut stream, || { + // Every JVM-fed scan already holds a batch, so the pull wakes nothing. + 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 gets the stream polled again. + single_worker_runtime().block_on(without_a_lost_wake(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); + })); + } + + /// A pull that runs another Comet plan on this thread, whose `block_on` parks until after the + /// stream's native I/O has completed. The nested park takes the I/O's wake-up from the + /// thread's parker, so `next_batch` has to have seen the wake some other way, or it parks + /// until `without_a_lost_wake`'s timer wakes it. + #[test] + fn next_batch_polls_again_when_a_nested_block_on_took_the_wake_up() { + let runtime = single_worker_runtime(); + let handle = runtime.handle().clone(); + let batch = RecordBatch::new_empty(Arc::new(Schema::empty())); + let mut stream = futures::stream::once(async move { + tokio::time::sleep(Duration::from_millis(20)).await; + Ok::<_, DataFusionError>(batch) + }) + .boxed(); + let mut pulls = 0; + let next = runtime + .block_on(without_a_lost_wake(next_batch(&mut stream, || { + pulls += 1; + handle.block_on(tokio::time::sleep(Duration::from_millis(100))); + Ok(()) + }))) + .unwrap(); + assert!(next.is_some()); + assert_eq!(pulls, 1); + } } diff --git a/native/core/src/execution/operators/scan.rs b/native/core/src/execution/operators/scan.rs index b1bf1990681..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,9 +113,11 @@ 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(); } - /// Pull next input batch from the upstream `ArrowArrayStreamReader`. + /// Pulls the next input batch from the upstream `ArrowArrayStreamReader` unless one is + /// 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`. @@ -120,14 +125,18 @@ impl ScanExec { } 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(()); } + 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(); + drop(current_batch); + self.waker.wake(); + Ok(()) } @@ -323,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 b5a282f5a39..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,30 +124,35 @@ 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(); } - /// Pull next input batch from JVM. Called externally before poll_next() - /// because JNI calls cannot happen from within poll_next on tokio threads. + /// 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(()); } - 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(()); } + 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(); + drop(current_batch); + self.waker.wake(); Ok(()) } @@ -362,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 @@ -391,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(); @@ -757,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) + )); + } }