kosiew commented on code in PR #25004:
URL: https://github.com/apache/datafusion/pull/25004#discussion_r3956124531
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -2011,6 +2213,41 @@ impl Stream for NestedLoopJoinStream {
cx: &mut std::task::Context<'_>,
) -> Poll<Option<Self::Item>> {
loop {
+ // A peer partition may have been dropped unfinished at any point,
+ // including while this stream was working on the final chunk. The
+ // coordinated execution cannot produce a complete result after
that,
+ // so fail rather than emit output from partial input.
+ if !matches!(self.state, NLJState::Done) {
+ // Poll a registered watcher rather than only reading the flag:
+ // this stream may be about to park on its own input, and the
+ // poll leaves a waker with the coordinator so a peer's
+ // cancellation actually reaches it.
+ if self.cancel_watch.is_none() {
+ self.cancel_watch = match &self.spill_state {
+ SpillState::Active(active) => {
+ Some(active.coordinator.cancellation_watcher())
+ }
+ SpillState::Pending {
Review Comment:
I think we need to be a bit more careful about cancellation while the stream
is `SpillState::Pending`.
`Pending` only means spilling is possible. Every eligible execution enters
this state before the shared left load determines whether the input actually
needs to spill. If `collect_left_input` resolves to `LeftLoad::InMemory`, there
is no coordinated fallback or shared chunk counter, so the right partitions are
still independent.
With the current behavior, dropping an unfinished partition while it is
still `Pending` can cancel the coordinator and cause a surviving partition to
fail even though the left side ultimately fits in memory.
Could we defer or separately record the pending drop, and only apply the
cancellation if the shared load resolves to `Spilled`?
I don't think limiting cancellation to `Active` streams is sufficient
either. One partition may already have transitioned to `Active` while another
is still `Pending`. If that pending partition is dropped, the active peer is
already relying on coordinated fallback and needs to be cancelled.
It would also be good to add a regression test with a sufficiently large
memory pool where one pending partition is dropped and another partition can
still be collected successfully.
##########
datafusion/physical-plan/src/joins/nested_loop_join.rs:
##########
@@ -5683,6 +5897,691 @@ pub(crate) mod tests {
}))
}
+ #[tokio::test]
+ async fn nlj_cancel_wakes_stream_parked_on_build_input() -> Result<()> {
+ let runtime = RuntimeEnvBuilder::new().build_arc()?;
+ let ctx = Arc::new(TaskContext::default().with_runtime(runtime));
+ let coordinator = Arc::new(FallbackCoordinator::new(2, true));
+ let spill = spill_left_for_test(build_left_table(),
Arc::clone(&ctx)).await?;
+ let _chunk = Arc::clone(&coordinator)
+ .next_chunk(0, Arc::clone(&spill), Arc::clone(&ctx), Time::new())
+ .await?
+ .expect("chunk");
+ let make_stream = || {
+ let right_schema = build_right_table().schema();
+ let (schema, columns) =
+ build_join_schema(&spill.schema, &right_schema,
&JoinType::Left);
+ NestedLoopJoinStream::new(
+ Arc::new(schema),
+ None,
+ JoinType::Left,
+ Box::pin(crate::stream::RecordBatchStreamAdapter::new(
+ right_schema,
+ futures::stream::pending::<Result<RecordBatch>>(),
+ )),
+ OnceFut::new(futures::future::pending::<Result<LeftLoad>>()),
+ columns,
+ NestedLoopJoinMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
+ 1,
+ SpillState::Pending {
+ task_context: Arc::clone(&ctx),
+ fallback_coordinator: Arc::clone(&coordinator),
+ },
+ )
+ };
+ let peer = make_stream();
+ let mut survivor = make_stream();
+ let wakes = WakeCount::new();
+ let waker = futures::task::waker(Arc::clone(&wakes));
+ let mut cx = std::task::Context::from_waker(&waker);
+ assert!(survivor.poll_next_unpin(&mut cx).is_pending());
+ assert!(matches!(survivor.state, NLJState::BufferingLeft));
+ wakes.reset();
+ drop(peer);
+ assert!(coordinator.is_cancelled());
+ assert!(
+ wakes.count() > 0,
+ "cancel must wake the survivor waiting on build input"
+ );
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn nlj_cancel_observed_when_watcher_polled_after_cancellation() {
+ let coordinator = Arc::new(FallbackCoordinator::new(2, true));
+ let before = coordinator.cancellation_watcher();
+ coordinator.cancel();
+ coordinator.cancel();
+ let after = coordinator.cancellation_watcher();
+ assert!(before.now_or_never().is_some());
+ assert!(after.now_or_never().is_some());
+ }
+
+ #[tokio::test]
+ async fn nlj_cancel_reaches_every_watcher_and_survives_waker_replacement()
{
+ let coordinator = Arc::new(FallbackCoordinator::new(2, true));
+ let mut first = coordinator.cancellation_watcher();
+ let mut second = coordinator.cancellation_watcher();
+ let old_count = WakeCount::new();
+ let new_count = WakeCount::new();
+ let other_count = WakeCount::new();
+ let old_waker = futures::task::waker(Arc::clone(&old_count));
+ let new_waker = futures::task::waker(Arc::clone(&new_count));
+ let other_waker = futures::task::waker(Arc::clone(&other_count));
+ assert!(
+ first
+ .poll_unpin(&mut std::task::Context::from_waker(&old_waker))
+ .is_pending()
+ );
+ assert!(
+ first
+ .poll_unpin(&mut std::task::Context::from_waker(&new_waker))
+ .is_pending()
+ );
+ assert!(
+ second
+ .poll_unpin(&mut std::task::Context::from_waker(&other_waker))
+ .is_pending()
+ );
+ coordinator.notify.notify_waiters();
+ assert_eq!(new_count.count(), 0);
+ assert_eq!(other_count.count(), 0);
+ coordinator.cancel();
+ coordinator.cancel();
+ assert!(new_count.count() > 0);
+ assert!(other_count.count() > 0);
+ assert!(first.now_or_never().is_some());
+ assert!(second.now_or_never().is_some());
+ }
+
+ #[tokio::test]
+ async fn nlj_normal_completion_does_not_cancel_peers() -> Result<()> {
+ tokio::time::timeout(Duration::from_secs(5), async {
+ let (plan, ctx) = cancellation_test_plan()?;
+ let mut rows = 0;
+ for partition in 0..2 {
+ let batches =
+ common::collect(plan.execute(partition,
Arc::clone(&ctx))?).await?;
+ rows +=
batches.iter().map(RecordBatch::num_rows).sum::<usize>();
+ assert!(!plan.fallback_coordinator.is_cancelled());
+ }
+ assert_eq!(rows, 9);
+ assert_eq!(ctx.memory_pool().reserved(), 0);
+ Ok(())
+ })
+ .await
+ .expect("normal execution hung")
+ }
+
+ /// A stream that ends in an error is unfinished, so dropping it cancels
the
+ /// peers.
+ ///
+ /// The error path reaches `Drop` without passing through `Done`. The build
+ /// side has to resolve for the poll to get as far as the failing right
input,
+ /// so this uses a `LeftLoad::Spilled` future rather than a pending one,
and
+ /// asserts the injected error actually surfaced before the drop --
otherwise
+ /// the test would silently degrade into "an unstarted stream cancels",
which
+ /// other tests already cover.
+ #[tokio::test]
+ async fn nlj_errored_stream_drop_cancels_peers() -> Result<()> {
+ let runtime = RuntimeEnvBuilder::new().build_arc()?;
+ let ctx = Arc::new(TaskContext::default().with_runtime(runtime));
+ let coordinator = Arc::new(FallbackCoordinator::new(2, true));
+ let spill = spill_left_for_test(build_left_table(),
Arc::clone(&ctx)).await?;
+
+ let right_schema = build_right_table().schema();
+ let (schema, columns) =
+ build_join_schema(&spill.schema, &right_schema, &JoinType::Left);
+ let left_spill = Arc::clone(&spill);
+ let mut failing = NestedLoopJoinStream::new(
+ Arc::new(schema),
+ None,
+ JoinType::Left,
+ Box::pin(crate::stream::RecordBatchStreamAdapter::new(
+ right_schema,
+ futures::stream::once(async {
+ exec_err!("injected right-input failure")
+ }),
+ )),
+ OnceFut::new(async move { Ok(LeftLoad::Spilled(left_spill)) }),
+ columns,
+ NestedLoopJoinMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
+ 1,
+ SpillState::Pending {
+ task_context: Arc::clone(&ctx),
+ fallback_coordinator: Arc::clone(&coordinator),
+ },
+ );
+
+ // Drive it until the injected error comes out. Bounded so a fixture
that
+ // stops reaching the right input fails instead of spinning.
+ let err = tokio::time::timeout(Duration::from_secs(5), async {
+ loop {
+ match failing.next().await {
+ Some(Err(e)) => return e,
+ Some(Ok(_)) => {}
+ None => panic!("stream finished without surfacing the
error"),
+ }
+ }
+ })
+ .await
+ .expect("fixture never reached its failing right input");
+ assert_contains!(err.to_string(), "injected right-input failure");
+ assert!(
+ !matches!(failing.state, NLJState::Done),
+ "an errored stream must not look finished"
+ );
+ assert!(!coordinator.is_cancelled());
+
+ drop(failing);
+ assert!(
+ coordinator.is_cancelled(),
+ "an unfinished stream must cancel its peers when dropped,
including \
+ one that ended in an error"
+ );
+ Ok(())
+ }
+
+ /// A survivor parked on the global-right replay must be woken by a peer's
+ /// cancellation.
+ ///
+ /// `EmitGlobalRightUnmatched` is only reachable with `Active` spill
state, so
+ /// the fixture installs one rather than assigning the state on top of
+ /// `Pending` -- otherwise the handler reaches its pending read through a
path
+ /// a real execution never takes, and the test would be checking a
`Pending`
+ /// watcher while claiming to check the replay configuration. The replay
+ /// reader is injected directly, which skips the reopen step; that is the
+ /// shortcut here, and reopening itself is covered elsewhere.
+ #[tokio::test]
+ async fn nlj_cancel_wakes_stream_parked_on_global_right_replay() ->
Result<()> {
+ let runtime = RuntimeEnvBuilder::new().build_arc()?;
+ let ctx = Arc::new(TaskContext::default().with_runtime(runtime));
+ let coordinator = Arc::new(FallbackCoordinator::new(2, true));
+ let spill = spill_left_for_test(build_left_table(),
Arc::clone(&ctx)).await?;
+ let _chunk = Arc::clone(&coordinator)
+ .next_chunk(0, Arc::clone(&spill), Arc::clone(&ctx), Time::new())
+ .await?
+ .expect("chunk");
+
+ let right_schema = build_right_table().schema();
+ let make_stream = || {
+ let (schema, columns) =
+ build_join_schema(&spill.schema, &right_schema,
&JoinType::Full);
+ let left_spill = Arc::clone(&spill);
+ NestedLoopJoinStream::new(
+ Arc::new(schema),
+ None,
+ JoinType::Full,
+ Box::pin(crate::stream::RecordBatchStreamAdapter::new(
+ Arc::clone(&right_schema),
+ futures::stream::pending::<Result<RecordBatch>>(),
+ )),
+ OnceFut::new(async move { Ok(LeftLoad::Spilled(left_spill)) }),
+ columns,
+ NestedLoopJoinMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
+ 1,
+ SpillState::Pending {
+ task_context: Arc::clone(&ctx),
+ fallback_coordinator: Arc::clone(&coordinator),
+ },
+ )
+ };
+ let peer = make_stream();
+ let mut survivor = make_stream();
+
+ // Reach `Active` the way execution does, by letting the spilled build
+ // side resolve, then park in the replay stage.
+ let wakes = WakeCount::new();
+ let waker = futures::task::waker(Arc::clone(&wakes));
+ let mut cx = std::task::Context::from_waker(&waker);
+ assert!(survivor.poll_next_unpin(&mut cx).is_pending());
+ assert!(
+ matches!(survivor.spill_state, SpillState::Active(_)),
+ "the global-right replay only exists with Active spill state"
+ );
+
+ survivor.state = NLJState::EmitGlobalRightUnmatched;
+ survivor.left_exhausted = true;
+ // Inject the replay reader so the handler parks on it rather than
+ // reopening the spill file, which other tests cover.
+ survivor.right_data =
+ Some(Box::pin(crate::stream::RecordBatchStreamAdapter::new(
+ Arc::clone(&right_schema),
+ futures::stream::pending::<Result<RecordBatch>>(),
+ )));
+ assert!(survivor.poll_next_unpin(&mut cx).is_pending());
+ assert!(matches!(survivor.state, NLJState::EmitGlobalRightUnmatched));
+ wakes.reset();
+
+ drop(peer);
+ assert!(coordinator.is_cancelled());
+ assert!(
+ wakes.count() > 0,
+ "cancel must wake a survivor parked on the global-right replay"
+ );
+
+ // And the wake must actually surface the cancellation, not just tick.
+ match survivor.poll_next_unpin(&mut cx) {
+ Poll::Ready(Some(Err(e))) => {
+ assert_contains!(e.to_string(), "cancelled");
+ }
+ _ => panic!("the woken survivor must report the cancellation"),
+ }
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn nlj_cancel_wakes_stream_parked_on_right_input() -> Result<()> {
+ let runtime = RuntimeEnvBuilder::new().build_arc()?;
+ let ctx = Arc::new(TaskContext::default().with_runtime(runtime));
+ let coordinator = Arc::new(FallbackCoordinator::new(2, true));
+ let spill = spill_left_for_test(build_left_table(),
Arc::clone(&ctx)).await?;
+ let _chunk = Arc::clone(&coordinator)
+ .next_chunk(0, Arc::clone(&spill), Arc::clone(&ctx), Time::new())
+ .await?
+ .expect("chunk");
+ let make_stream = || {
+ let right_schema = build_right_table().schema();
+ let (schema, columns) =
+ build_join_schema(&spill.schema, &right_schema,
&JoinType::Left);
+ let left_spill = Arc::clone(&spill);
+ NestedLoopJoinStream::new(
+ Arc::new(schema),
+ None,
+ JoinType::Left,
+ Box::pin(crate::stream::RecordBatchStreamAdapter::new(
+ right_schema,
+ futures::stream::pending::<Result<RecordBatch>>(),
+ )),
+ OnceFut::new(async move { Ok(LeftLoad::Spilled(left_spill)) }),
+ columns,
+ NestedLoopJoinMetrics::new(&ExecutionPlanMetricsSet::new(), 0),
+ 1,
+ SpillState::Pending {
+ task_context: Arc::clone(&ctx),
+ fallback_coordinator: Arc::clone(&coordinator),
+ },
+ )
+ };
+ let peer = make_stream();
+ let mut survivor = make_stream();
+ let wakes = WakeCount::new();
+ let waker = futures::task::waker(Arc::clone(&wakes));
+ let mut cx = std::task::Context::from_waker(&waker);
+ assert!(survivor.poll_next_unpin(&mut cx).is_pending());
+ assert!(matches!(survivor.state, NLJState::FetchingRight));
+ wakes.reset();
+ drop(peer);
+ assert!(coordinator.is_cancelled());
+ assert!(
+ wakes.count() > 0,
+ "cancel must wake the survivor waiting on right input"
+ );
+ Ok(())
+ }
+
+ #[tokio::test]
+ async fn nlj_stream_stops_producing_after_cancellation_error() ->
Result<()> {
Review Comment:
Small non-blocking test suggestion: this uses the shared fixture with
`batch_size = 1`, so it can't exercise cancellation while the coalescer has a
partial batch buffered.
It might be worth adding a focused test with a larger batch size, cancelling
while the survivor has buffered output, and checking that the terminal behavior
is still an error followed by `None`. This would give us a little more
confidence that buffered output can't escape after cancellation.
--
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]
