diff --git a/docs/source/contributor-guide/development.md b/docs/source/contributor-guide/development.md index b5752e0863..26295f93bc 100644 --- a/docs/source/contributor-guide/development.md +++ b/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 diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 7652428410..0f2bbec16b 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; @@ -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, @@ -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( + 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; + } + } + } } /// 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(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); + })); + } } diff --git a/native/core/src/execution/operators/scan.rs b/native/core/src/execution/operators/scan.rs index b1bf199068..665f6d3e34 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 b5a282f5a3..05f26583e5 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) + )); + } }