mbutrovich commented on code in PR #6092:
URL: https://github.com/apache/datafusion-comet/pull/6092#discussion_r4067016907
##########
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:
The timeout guards against a `Pending` with no waker. In `native/`, only
`ScanStream`
([scan.rs](https://github.com/apache/datafusion-comet/blob/2b25d4a03992caf45d5f9d39d8b0d13f4cd29d28/native/core/src/execution/operators/scan.rs#L329-L339))
and `ShuffleScanStream`
([shuffle_scan.rs](https://github.com/apache/datafusion-comet/blob/2b25d4a03992caf45d5f9d39d8b0d13f4cd29d28/native/core/src/execution/operators/shuffle_scan.rs#L367-L377))
do that. Could they store `cx.waker()` when the buffer is empty, and could
`get_next_batch` wake it after a refill? Then every `Pending` carries a waker,
the park needs no timeout, and `get_next_batch` doesn't need to return whether
it pulled. Both streams would also follow the `Stream` contract for any task
that polls them.
As written, a lost wake-up fails nothing. It adds 100 ms per wait, which
tests won't catch and production will report as a slow scan. If you keep the
timeout, please log or count each time it fires.
##########
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:
Could `Future`, `Pin`, and `Context` be imported next to `task::Poll` at
[line
99](https://github.com/apache/datafusion-comet/blob/2b25d4a03992caf45d5f9d39d8b0d13f4cd29d28/native/core/src/execution/jni_api.rs#L99),
as the rest of the file does? `std::future::poll_fn` with a local `bool` would
also replace the struct.
##########
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:
This test covers `park_until_woken` alone. It still passes with the `if
!pulled` branch at line 1177 deleted, which reintroduces #6091. Could the loop
get a test that fails on `main`? One way is to move the poll, pull, and park
steps into a function that takes the stream and a pull closure returning
`Result<bool, CometError>`, leaving `update_metrics` and `prepare_output` in
`executePlan`. A test can then drive it with a stream pending on
`tokio::time::sleep` and a closure that returns `false`, and assert the closure
runs a few times, not thousands.
--
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]
