Jimexist commented on a change in pull request #375:
URL: https://github.com/apache/arrow-datafusion/pull/375#discussion_r639280266
##########
File path: datafusion/src/physical_plan/windows.rs
##########
@@ -186,10 +271,259 @@ impl ExecutionPlan for WindowAggExec {
));
}
- // let input = self.input.execute(0).await?;
+ let input = self.input.execute(partition).await?;
+
+ let stream = Box::pin(WindowAggStream::new(
+ self.schema.clone(),
+ self.window_expr.clone(),
+ input,
+ ));
+ Ok(stream)
+ }
+}
+
+pin_project! {
+ /// stream for window aggregation plan
+ pub struct WindowAggStream {
+ schema: SchemaRef,
+ #[pin]
+ output: futures::channel::oneshot::Receiver<ArrowResult<RecordBatch>>,
+ finished: bool,
+ }
+}
+
+type WindowAccumulatorItem = Box<dyn WindowAccumulator>;
+
+fn window_expressions(
+ window_expr: &[Arc<dyn WindowExpr>],
+) -> Result<Vec<Vec<Arc<dyn PhysicalExpr>>>> {
+ Ok(window_expr
+ .iter()
+ .map(|expr| expr.expressions())
+ .collect::<Vec<_>>())
+}
+
+fn window_aggregate_batch(
+ batch: &RecordBatch,
+ window_accumulators: &mut [WindowAccumulatorItem],
+ expressions: &[Vec<Arc<dyn PhysicalExpr>>],
+) -> Result<Vec<Option<ArrayRef>>> {
+ // 1.1 iterate accumulators and respective expressions together
+ // 1.2 evaluate expressions
+ // 1.3 update / merge window accumulators with the expressions' values
+
+ // 1.1
+ window_accumulators
+ .iter_mut()
+ .zip(expressions)
+ .map(|(window_acc, expr)| {
+ // 1.2
+ let values = &expr
+ .iter()
+ .map(|e| e.evaluate(batch))
+ .map(|r| r.map(|v| v.into_array(batch.num_rows())))
+ .collect::<Result<Vec<_>>>()?;
+
+ window_acc.scan_batch(batch.num_rows(), values)
+ })
+ .into_iter()
+ .collect::<Result<Vec<_>>>()
+}
+
+/// returns a vector of ArrayRefs, where each entry corresponds to either the
+/// final value (mode = Final) or states (mode = Partial)
Review comment:
updated
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
