feat: unbounded channel support for AdaptorStream

Allows the AdaptorStream to work with either a bounded, or unbounded
channel internally.

datafusion_util/src/lib.rs

@@ -16,8 +16,8 @@ use datafusion::{
     scalar::ScalarValue,
 };
 use futures::{Stream, StreamExt};
-use tokio::sync::mpsc::Receiver;
+use tokio::sync::mpsc::{Receiver, UnboundedReceiver};
-use tokio_stream::wrappers::ReceiverStream;
+use tokio_stream::wrappers::{ReceiverStream, UnboundedReceiverStream};
 
 /// Traits to help creating DataFusion [`Expr`]s
 pub trait AsExpr {
@@ -162,22 +162,24 @@ impl futures::Stream for MemoryStream {
 }
 
 #[derive(Debug)]
-/// Implements a [`SendableRecordBatchStream`] to help create DataFusion outputs from tokio channels.
+/// Implements a [`SendableRecordBatchStream`] to help create DataFusion outputs
+/// from tokio channels.
 ///
-/// It sends streams of RecordBatches from a tokio channel *and*
+/// It sends streams of RecordBatches from a tokio channel *and* crucially knows
-/// crucially knows up front the schema each batch will have be used.
+/// up front the schema each batch will have be used.
-pub struct AdapterStream {
+pub struct AdapterStream<T> {
     /// Schema
     schema: SchemaRef,
     /// channel for getting deduplicated batches
-    inner: ReceiverStream<ArrowResult<RecordBatch>>,
+    inner: T,
 }
 
-impl AdapterStream {
+impl AdapterStream<ReceiverStream<ArrowResult<RecordBatch>>> {
-    /// Create a new stream which wraps the `inner` channel which
+    /// Create a new stream which wraps the `inner` channel which produces
-    /// produces [`RecordBatch`]es that each have the specified schema
+    /// [`RecordBatch`]es that each have the specified schema
     ///
-    /// Not called `new` because it returns a pinned reference rather than the object itself
+    /// Not called `new` because it returns a pinned reference rather than the
+    /// object itself.
     pub fn adapt(
         schema: SchemaRef,
         rx: Receiver<ArrowResult<RecordBatch>>,
@@ -187,7 +189,25 @@ impl AdapterStream {
     }
 }
 
-impl Stream for AdapterStream {
+impl AdapterStream<UnboundedReceiverStream<ArrowResult<RecordBatch>>> {
+    /// Create a new stream which wraps the `inner` unbounded channel which
+    /// produces [`RecordBatch`]es that each have the specified schema
+    ///
+    /// Not called `new` because it returns a pinned reference rather than the
+    /// object itself.
+    pub fn adapt_unbounded(
+        schema: SchemaRef,
+        rx: UnboundedReceiver<ArrowResult<RecordBatch>>,
+    ) -> SendableRecordBatchStream {
+        let inner = UnboundedReceiverStream::new(rx);
+        Box::pin(Self { schema, inner })
+    }
+}
+
+impl<T> Stream for AdapterStream<T>
+where
+    T: Stream<Item = ArrowResult<RecordBatch>> + Unpin,
+{
     type Item = ArrowResult<RecordBatch>;
 
     fn poll_next(
@@ -198,7 +218,10 @@ impl Stream for AdapterStream {
     }
 }
 
-impl RecordBatchStream for AdapterStream {
+impl<T> RecordBatchStream for AdapterStream<T>
+where
+    T: Stream<Item = ArrowResult<RecordBatch>> + Unpin,
+{
     fn schema(&self) -> SchemaRef {
         Arc::clone(&self.schema)
     }