Skip to content

Commit

Permalink
Do not drive all streams to error
Browse files Browse the repository at this point in the history
  • Loading branch information
alamb committed Jun 6, 2023
1 parent a531afe commit 56a26eb
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 29 deletions.
38 changes: 19 additions & 19 deletions datafusion/core/src/physical_plan/stream.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
});
}
Expand Down Expand Up @@ -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
Expand All @@ -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
}

Expand Down Expand Up @@ -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![
Expand All @@ -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
Expand Down
23 changes: 13 additions & 10 deletions datafusion/core/src/test/exec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
}
}

Expand Down Expand Up @@ -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 {
Expand Down

0 comments on commit 56a26eb

Please sign in to comment.