mixermt commented on code in PR #6092:
URL: https://github.com/apache/datafusion-comet/pull/6092#discussion_r4070568034
##########
native/core/src/execution/jni_api.rs:
##########
@@ -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<bool,
CometError> {
+ 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;
Review Comment:
Done. Both streams register `cx.waker()` on an empty buffer and
`get_next_batch` wakes it after the refill, so every `Pending` carries a waker.
The timeout and the `bool` return are gone, so there is no firing left to log
or count.
##########
native/core/src/execution/jni_api.rs:
##########
@@ -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<bool,
CometError> {
+ 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<()> {
Review Comment:
Done. The struct is replaced by `poll_fn` with a local `bool`, imported next
to `task::Poll`.
##########
native/core/src/execution/jni_api.rs:
##########
@@ -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() {
Review Comment:
Done. The loop is now `next_batch(stream, on_pending)`, with
`update_metrics` and `prepare_output` left in `executePlan`.
`next_batch_parks_while_the_stream_waits_on_native_io` drives it with a stream
pending on `tokio::time::sleep`; with the park deleted the closure ran 135,311
times in one 50 ms wait.
`next_batch_resumes_on_a_refill_and_stops_pulling_after_eof` covers the refill
wake and the re-poll after EOF.
##########
native/core/src/execution/jni_api.rs:
##########
@@ -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 {
Review Comment:
Done. The sample sits behind the same interval check as `update_metrics`, in
`update_metrics_on_interval`, which runs on every pending poll and once per
returned batch.
##########
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<bool, CometError> {
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)
Review Comment:
Done. EOF stays buffered in both streams, so a re-poll returns `Ready(None)`
again and `get_next_batch` is a no-op. The refill test fails if the EOF is
cleared again.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
