add streams

Author: Jiayu Liu

datafusion/src/physical_plan/aggregates.rs

@@ -105,7 +105,7 @@ pub fn return_type(fun: &AggregateFunction, arg_types: &[DataType]) -> Result<Da
 
 /// Create a physical (function) expression.
 /// This function errors when `args`' can't be coerced to a valid argument type of the function.
-pub fn create_aggregate_expr(
+pub(super) fn create_aggregate_expr(
     fun: &AggregateFunction,
     distinct: bool,
     args: &[Arc<dyn PhysicalExpr>],

datafusion/src/physical_plan/hash_aggregate.rs

@@ -712,7 +712,7 @@ impl GroupedHashAggregateStream {
             tx.send(result)
         });
 
-        GroupedHashAggregateStream {
+        Self {
             schema,
             output: rx,
             finished: false,
@@ -825,6 +825,7 @@ fn aggregate_expressions(
 }
 
 pin_project! {
+    /// stream struct for hash aggregation
     pub struct HashAggregateStream {
         schema: SchemaRef,
         #[pin]

datafusion/src/physical_plan/sort.rs

@@ -227,6 +227,7 @@ fn sort_batches(
 }
 
 pin_project! {
+    /// stream for sort plan
     struct SortStream {
         #[pin]
         output: futures::channel::oneshot::Receiver<ArrowResult<Option<RecordBatch>>>,

datafusion/src/physical_plan/windows.rs

@@ -19,13 +19,24 @@
 
 use crate::error::{DataFusionError, Result};
 use crate::physical_plan::{
-    aggregates, window_functions::WindowFunction, AggregateExpr, Distribution,
-    ExecutionPlan, Partitioning, PhysicalExpr, SendableRecordBatchStream, WindowExpr,
+    aggregates, expressions::Column, window_functions::WindowFunction, AggregateExpr,
+    Distribution, ExecutionPlan, Partitioning, PhysicalExpr, RecordBatchStream,
+    SendableRecordBatchStream, WindowExpr,
+};
+use arrow::{
+    array::{Array, UInt32Builder},
+    datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit},
+    error::{ArrowError, Result as ArrowResult},
+    record_batch::RecordBatch,
 };
-use arrow::datatypes::{Field, Schema, SchemaRef};
 use async_trait::async_trait;
+use futures::stream::Stream;
+use futures::stream::StreamExt;
+use pin_project_lite::pin_project;
 use std::any::Any;
+use std::pin::Pin;
 use std::sync::Arc;
+use std::task::{Context, Poll};
 
 /// Window execution plan
 #[derive(Debug)]
@@ -195,43 +206,89 @@ impl ExecutionPlan for WindowAggExec {
 
         let input = self.input.execute(partition).await?;
 
-
-        
-
         Err(DataFusionError::NotImplemented(
             "WindowAggExec::execute".to_owned(),
         ))
     }
 }
 
-// struct WindowAggStream {
-//     scheme: SchemaRef,
-//     window_expr: Vec<Arc<dyn WindowExpr>>,
-//     input: SendableRecordBatchStream,
-// }
-
-// impl Stream for WindowAggStream {
-//     type Item = ArrowResult<RecordBatch>;
-
-//     fn poll_next(
-//         mut self: Pin<&mut Self>,
-//         cx: &mut Context<'_>,
-//     ) -> Poll<Option<Self::Item>> {
-//         self.input.poll_next_unpin(cx).map(|x| match x {
-//             Some(Ok(batch)) => Some(batch_project(&batch, &self.expr, &self.schema)),
-//             other => other,
-//         })
-//     }
-
-//     fn size_hint(&self) -> (usize, Option<usize>) {
-//         // same number of record batches
-//         self.input.size_hint()
-//     }
-// }
-
-// impl RecordBatchStream for WindowAggStream {
-//     /// Get the schema
-//     fn schema(&self) -> SchemaRef {
-//         self.schema.clone()
-//     }
-// }
+pin_project! {
+    /// stream for window aggregation plan
+    pub struct WindowAggStream {
+        #[pin]
+        output: futures::channel::oneshot::Receiver<ArrowResult<Option<RecordBatch>>>,
+        finished: bool,
+        schema: SchemaRef,
+    }
+}
+
+async fn compute_window_aggregate(
+    schema: SchemaRef,
+    window_expr: Vec<Arc<dyn WindowExpr>>,
+    mut input: SendableRecordBatchStream,
+) -> ArrowResult<Option<RecordBatch>> {
+    unimplemented!("not implemented")
+}
+
+impl WindowAggStream {
+    fn new(
+        schema: SchemaRef,
+        input: SendableRecordBatchStream,
+        window_expr: Vec<Arc<dyn WindowExpr>>,
+    ) -> Self {
+        let (tx, rx) = futures::channel::oneshot::channel();
+        let schema_clone = schema.clone();
+        tokio::spawn(async move {
+            let result = compute_window_aggregate(schema_clone, window_expr, input).await;
+            tx.send(result)
+        });
+
+        Self {
+            output: rx,
+            finished: false,
+            schema,
+        }
+    }
+}
+
+impl Stream for WindowAggStream {
+    type Item = ArrowResult<RecordBatch>;
+
+    fn poll_next(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Self::Item>> {
+        if self.finished {
+            return Poll::Ready(None);
+        }
+
+        // is the output ready?
+        let this = self.project();
+        let output_poll = this.output.poll(cx);
+
+        match output_poll {
+            Poll::Ready(result) => {
+                *this.finished = true;
+                // check for error in receiving channel and unwrap actual result
+                let result = match result {
+                    Err(e) => Err(ArrowError::ExternalError(Box::new(e))), // error receiving
+                    Ok(result) => result,
+                };
+                Poll::Ready(Some(result))
+            }
+            Poll::Pending => Poll::Pending,
+        }
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        // same number of record batches
+        self.input.size_hint()
+    }
+}
+
+impl RecordBatchStream for WindowAggStream {
+    /// Get the schema
+    fn schema(&self) -> SchemaRef {
+        self.schema.clone()
+    }
+}