Struct datafusion_physical_plan::stream::RecordBatchReceiverStreamBuilder

pub struct RecordBatchReceiverStreamBuilder { /* private fields */ }

Builder for RecordBatchReceiverStream that propagates errors and panic’s correctly.

RecordBatchReceiverStreamBuilder is used to spawn one or more tasks that produce RecordBatches and send them to a single Receiver which can improve parallelism.

This also handles propagating panic`s and canceling the tasks.

Example

The following example spawns 2 tasks that will write RecordBatches to the tx end of the builder, after building the stream, we can receive those batches with calling .next()

let schema = Arc::new(Schema::new(vec![Field::new("foo", DataType::Int8, false)]));
let mut builder = RecordBatchReceiverStreamBuilder::new(Arc::clone(&schema), 10);

// task 1
let tx_1 = builder.tx();
let schema_1 = Arc::clone(&schema);
builder.spawn(async move {
    // Your task needs to send batches to the tx
    tx_1.send(Ok(RecordBatch::new_empty(schema_1))).await.unwrap();

    Ok(())
});

// task 2
let tx_2 = builder.tx();
let schema_2 = Arc::clone(&schema);
builder.spawn(async move {
    // Your task needs to send batches to the tx
    tx_2.send(Ok(RecordBatch::new_empty(schema_2))).await.unwrap();

    Ok(())
});

let mut stream = builder.build();
while let Some(res_batch) = stream.next().await {
    // `res_batch` can either from task 1 or 2

    // do something with `res_batch`
}

Implementations

impl RecordBatchReceiverStreamBuilder

pub fn new(schema: SchemaRef, capacity: usize) -> Self

create new channels with the specified buffer size

pub fn tx(&self) -> Sender<Result<RecordBatch>>

Get a handle for sending RecordBatch to the output

pub fn spawn<F>(&mut self, task: F)

where F: Future<Output = Result<()>> + Send + 'static,

Spawn task that will be aborted if this builder (or the stream built from it) are dropped

This is often used to spawn tasks that write to the sender retrieved from Self::tx, for examples, see the document of this type.

pub fn spawn_blocking<F>(&mut self, f: F)

where F: FnOnce() -> Result<()> + Send + 'static,

Spawn a blocking task that will be aborted if this builder (or the stream built from it) are dropped

This is often used to spawn tasks that write to the sender retrieved from Self::tx, for examples, see the document of this type.

pub fn build(self) -> SendableRecordBatchStream

Create a stream of all RecordBatch written to tx