Do not drive all streams to error

Author: alamb

datafusion/core/src/physical_plan/stream.rs

@@ -117,16 +117,26 @@ impl RecordBatchReceiverStreamBuilder {
                 Ok(stream) => stream,
             };
 
+            // Transfer batches from inner stream to the output tx
+            // immediately.
             while let Some(item) = stream.next().await {
-                // If send fails, plan being torn down,
-                // there is no place to send the error.
+                let is_err = item.is_err();
+
+                // If send fails, plan being torn down, there is no
+                // place to send the error and no reason to continue.
                 if output.send(item).await.is_err() {
                     debug!(
                         "Stopping execution: output is gone, plan cancelling: {}",
                         displayable(input.as_ref()).one_line()
                     );
                     return;
                 }
+
+                // stop after the first error is encontered (don't
+                // drive all streams to completion)
+                if is_err {
+                    return;
+                }
             }
         });
     }
@@ -332,7 +342,7 @@ mod test {
 
     #[tokio::test]
     #[should_panic(expected = "PanickingStream did panic: 1")]
-    async fn record_batch_receiver_stream_propagates_panics_one() {
+    async fn record_batch_receiver_stream_propagates_panics_early_shutdown() {
         let schema = schema();
 
         // make 2 partitions, second panics before the first
@@ -341,7 +351,12 @@ mod test {
             .with_partition_panic(0, 10)
             .with_partition_panic(1, 3); // partition 1 should panic first (after 3 )
 
-        let max_batches = 5; // expect to read every other batch: (0,1,0,1,0,panic)
+        // ensure that the panic results in an early shutdown (that
+        // everything stops after the first panic).
+
+        // Since the stream reads every other batch: (0,1,0,1,0,panic)
+        // so should not exceed 5 batches prior to the panic
+        let max_batches = 5;
         consume(input, max_batches).await
     }
 
@@ -378,10 +393,6 @@ mod test {
         let task_ctx = session_ctx.task_ctx();
         let schema = schema();
 
-        // Make an input that will not proceed
-        let blocking_input = BlockingExec::new(schema.clone(), 1);
-        let refs = blocking_input.refs();
-
         // make an input that will error
         let error_stream = MockExec::new(
             vec![
@@ -392,28 +403,17 @@ mod test {
         )
         .with_use_task(false);
 
-        // Configure a RecordBatchReceiverStream to consume the
-        // blocking input (which will never advance) and the stream
-        // that will error.
-
         let mut builder = RecordBatchReceiverStream::builder(schema, 2);
-        builder.run_input(Arc::new(blocking_input), 0, task_ctx.clone());
         builder.run_input(Arc::new(error_stream), 0, task_ctx.clone());
         let mut stream = builder.build();
 
-        // first input input should be present
-        assert!(std::sync::Weak::strong_count(&refs) > 0);
-
         // get the first result, which should be an error
         let first_batch = stream.next().await.unwrap();
         let first_err = first_batch.unwrap_err();
         assert_eq!(first_err.to_string(), "Execution error: Test1");
 
         // There should be no more batches produced (should not get the second error)
         assert!(stream.next().await.is_none());
-
-        // And the other inputs should be cleaned up (even before stream is dropped)
-        assert_strong_count_converges_to_zero(refs).await;
     }
 
     /// Consumes all the input's partitions into a

datafusion/core/src/test/exec.rs

@@ -116,8 +116,8 @@ impl RecordBatchStream for TestStream {
     }
 }
 
-/// A Mock ExecutionPlan that can be used for writing tests of other ExecutionPlans
-///
+/// A Mock ExecutionPlan that can be used for writing tests of other
+/// ExecutionPlans
 #[derive(Debug)]
 pub struct MockExec {
     /// the results to send back
@@ -129,15 +129,18 @@ pub struct MockExec {
 }
 
 impl MockExec {
-    /// Create a new exec with a single partition that returns the
-    /// record batches in this Exec. Note the batches are not produced
-    /// immediately (the caller has to actually yield and another task
-    /// must run) to ensure any poll loops are correct.
+    /// Create a new `MockExec` with a single partition that returns
+    /// the specified `Results`s.
+    ///
+    /// By default, the batches are not produced immediately (the
+    /// caller has to actually yield and another task must run) to
+    /// ensure any poll loops are correct. This behavior can be
+    /// changed with `with_use_task`
     pub fn new(data: Vec<Result<RecordBatch>>, schema: SchemaRef) -> Self {
         Self {
             data,
             schema,
-            use_task: false,
+            use_task: true,
         }
     }
 
@@ -198,9 +201,9 @@ impl ExecutionPlan for MockExec {
 
         if self.use_task {
             let mut builder = RecordBatchReceiverStream::builder(self.schema(), 2);
-            // send data in order but in a separate
-            // thread (to ensure the batches are not available without the
-            // DelayedStream yielding).
+            // send data in order but in a separate task (to ensure
+            // the batches are not available without the stream
+            // yielding).
             let tx = builder.tx();
             builder.spawn(async move {
                 for batch in data {