(datafusion) 01/01: fix: surface BufferExec input panics instead of silently truncating output (#23243)

[gi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23243-01bf68cdc46ffbf4c226e760b0ba8a63fdb89a6b
in repository https://gitbox.apache.org/repos/asf/datafusion.git

commit 6ef3a3bdccdb3a6d0a75079d0be7d1d5deeec1ba
Author: Wenqi Mou <we...@gmail.com>
AuthorDate: Tue Jun 30 03:38:30 2026 -0400

    fix: surface BufferExec input panics instead of silently truncating output (#23243)
    
    ## Which issue does this close?
    
    Closes #23242.
    
    ## Rationale for this change
    
    When the input to a `BufferExec` panics, the panic unwinds the
    background producer task (`MemoryBufferedStream`), tokio catches it at
    the task boundary, and the dropped sender looks like a clean
    end-of-stream to the consumer. A panic isn't a `Result::Err`, so it also
    skips the `batch_tx.send(Err(..))` path — the partition gets silently
    truncated and the query returns partial results instead of failing. More
    detail in #23242.
    
    ## What changes are included in this PR?
    
    Catch the panic at the input poll inside `MemoryBufferedStream` and
    forward it as a `DataFusionError` over the channel the consumer already
    drains, so it propagates instead of being swallowed.
    
    I went with surfacing it as an error rather than re-raising via
    `join_unwind` (the way `RecordBatchReceiverStream` does), since the
    consumer already handles `Some(Err(..))` and this fails only the query.
    Happy to switch to a re-raise if that's preferred for consistency.
    
    ## Are these changes tested?
    
    Yes — added `panic_in_input_is_propagated`, which feeds a stream that
    panics partway through and asserts the buffered stream yields an error
    instead of finishing cleanly.
    
    ## Are there any user-facing changes?
    
    A query whose input panics under a `BufferExec` now fails with an error
    instead of silently returning truncated results. No API changes.
---
 datafusion/physical-plan/src/buffer.rs | 49 +++++++++++++++++++++++++++++-----
 1 file changed, 43 insertions(+), 6 deletions(-)

diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs
index 6d1fb69635..5e220b7e48 100644
--- a/datafusion/physical-plan/src/buffer.rs
+++ b/datafusion/physical-plan/src/buffer.rs
@@ -41,9 +41,10 @@ use datafusion_physical_expr_common::metrics::{
 };
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
-use futures::{Stream, StreamExt, TryStreamExt};
+use futures::{FutureExt, Stream, StreamExt, TryStreamExt};
 use pin_project_lite::pin_project;
 use std::fmt;
+use std::panic::AssertUnwindSafe;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicUsize, Ordering};
@@ -337,11 +338,24 @@ impl<T: Send + SizedMessage + 'static> MemoryBufferedStream<T> {
                 let item_or_err = tokio::select! {
                     biased;
                     _ = batch_tx.closed() => break,
-                    item_or_err = input.next() => {
-                        let Some(item_or_err) = item_or_err else {
-                            break; // stream finished
-                        };
-                        item_or_err
+                    // Catch a panic in the input poll so it surfaces as a stream error
+                    // instead of dropping `batch_tx` and looking like a clean EOF.
+                    polled = AssertUnwindSafe(input.next()).catch_unwind() => {
+                        match polled {
+                            Ok(Some(item_or_err)) => item_or_err,
+                            Ok(None) => break, // stream finished
+                            Err(panic) => {
+                                let msg = panic
+                                    .downcast_ref::<&str>()
+                                    .map(|s| s.to_string())
+                                    .or_else(|| panic.downcast_ref::<String>().cloned())
+                                    .unwrap_or_else(|| "unknown panic".to_string());
+                                let _ = batch_tx.send(internal_err!(
+                                    "BufferExec input stream panicked: {msg}"
+                                ));
+                                break;
+                            }
+                        }
                     }
                 };
 
@@ -554,6 +568,29 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn panic_in_input_is_propagated() -> Result<(), Box<dyn Error>> {
+        // A panic while polling the input must surface as a stream error, not a
+        // silent end-of-stream that drops the rest of the partition's output.
+        let input = futures::stream::iter([1, 2, 3, 4]).map(|v| {
+            if v == 3 {
+                panic!("boom on 3");
+            }
+            Ok(v)
+        });
+        let (_, res) = memory_pool_and_reservation();
+
+        let mut buffered = MemoryBufferedStream::new(input, 10, res);
+        wait_for_buffering().await;
+
+        pull_ok_msg(&mut buffered).await?;
+        pull_ok_msg(&mut buffered).await?;
+        let err = pull_err_msg(&mut buffered).await?;
+        assert_contains!(err.to_string(), "panicked");
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn memory_gets_released_if_stream_drops() -> Result<(), Box<dyn Error>> {
         let input = futures::stream::iter([1, 2, 3, 4]).map(Ok);


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org