Sort Aggregate

ballista/rust/core/src/serde/physical_plan/from_proto.rs

@@ -36,9 +36,10 @@ use datafusion::logical_plan::{
     window_frames::WindowFrame, DFSchema, Expr, JoinConstraint, JoinType,
 };
 use datafusion::physical_plan::aggregates::{create_aggregate_expr, AggregateFunction};
+use datafusion::physical_plan::aggregations::hash_aggregate::HashAggregateExec;
+use datafusion::physical_plan::aggregations::AggregateMode;
 use datafusion::physical_plan::avro::{AvroExec, AvroReadOptions};
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
-use datafusion::physical_plan::hash_aggregate::{AggregateMode, HashAggregateExec};
 use datafusion::physical_plan::joins::cross_join::CrossJoinExec;
 use datafusion::physical_plan::joins::hash_join::{HashJoinExec, PartitionMode};
 use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet;

ballista/rust/core/src/serde/physical_plan/mod.rs

@@ -30,11 +30,12 @@ mod roundtrip_tests {
         },
         logical_plan::{JoinType, Operator},
         physical_plan::{
+            aggregations::hash_aggregate::HashAggregateExec,
+            aggregations::AggregateMode,
             empty::EmptyExec,
             expressions::{binary, col, lit, InListExpr, NotExpr},
             expressions::{Avg, Column, PhysicalSortExpr},
             filter::FilterExec,
-            hash_aggregate::{AggregateMode, HashAggregateExec},
             limit::{GlobalLimitExec, LocalLimitExec},
             sorts::sort::SortExec,
             AggregateExpr, ColumnarValue, Distribution, ExecutionPlan, Partitioning,

ballista/rust/core/src/serde/physical_plan/to_proto.rs

@@ -27,14 +27,14 @@ use std::{
 };
 
 use datafusion::logical_plan::JoinType;
+use datafusion::physical_plan::aggregations::AggregateMode;
 use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec;
 use datafusion::physical_plan::csv::CsvExec;
 use datafusion::physical_plan::expressions::{
     CaseExpr, InListExpr, IsNotNullExpr, IsNullExpr, NegativeExpr, NotExpr,
 };
 use datafusion::physical_plan::expressions::{CastExpr, TryCastExpr};
 use datafusion::physical_plan::filter::FilterExec;
-use datafusion::physical_plan::hash_aggregate::AggregateMode;
 use datafusion::physical_plan::joins::cross_join::CrossJoinExec;
 use datafusion::physical_plan::joins::hash_join::{HashJoinExec, PartitionMode};
 use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
@@ -53,7 +53,7 @@ use datafusion::physical_plan::{
 };
 use datafusion::physical_plan::{AggregateExpr, ExecutionPlan, PhysicalExpr};
 
-use datafusion::physical_plan::hash_aggregate::HashAggregateExec;
+use datafusion::physical_plan::aggregations::hash_aggregate::HashAggregateExec;
 use protobuf::physical_plan_node::PhysicalPlanType;
 
 use crate::execution_plans::{

ballista/rust/core/src/utils.rs

@@ -49,13 +49,13 @@ use datafusion::logical_plan::{LogicalPlan, Operator};
 use datafusion::physical_optimizer::coalesce_batches::CoalesceBatches;
 use datafusion::physical_optimizer::merge_exec::AddCoalescePartitionsExec;
 use datafusion::physical_optimizer::optimizer::PhysicalOptimizerRule;
+use datafusion::physical_plan::aggregations::hash_aggregate::HashAggregateExec;
 use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec;
 use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion::physical_plan::csv::CsvExec;
 use datafusion::physical_plan::empty::EmptyExec;
 use datafusion::physical_plan::expressions::{BinaryExpr, Column, Literal};
 use datafusion::physical_plan::filter::FilterExec;
-use datafusion::physical_plan::hash_aggregate::HashAggregateExec;
 use datafusion::physical_plan::joins::hash_join::HashJoinExec;
 use datafusion::physical_plan::parquet::ParquetExec;
 use datafusion::physical_plan::projection::ProjectionExec;

ballista/rust/scheduler/src/planner.rs

@@ -249,8 +249,9 @@ mod test {
     use ballista_core::error::BallistaError;
     use ballista_core::execution_plans::UnresolvedShuffleExec;
     use ballista_core::serde::protobuf;
+    use datafusion::physical_plan::aggregations::hash_aggregate::HashAggregateExec;
+    use datafusion::physical_plan::aggregations::AggregateMode;
     use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec;
-    use datafusion::physical_plan::hash_aggregate::{AggregateMode, HashAggregateExec};
     use datafusion::physical_plan::joins::hash_join::HashJoinExec;
     use datafusion::physical_plan::sorts::sort::SortExec;
     use datafusion::physical_plan::{

datafusion/src/lib.rs

@@ -165,7 +165,7 @@
 //!
 //! * Projection: [`ProjectionExec`](physical_plan::projection::ProjectionExec)
 //! * Filter: [`FilterExec`](physical_plan::filter::FilterExec)
-//! * Hash and Grouped aggregations: [`HashAggregateExec`](physical_plan::hash_aggregate::HashAggregateExec)
+//! * Hash and Grouped aggregations: [`HashAggregateExec`](physical_plan::aggregations::hash_aggregate::HashAggregateExec)
 //! * Sort: [`SortExec`](physical_plan::sort::SortExec)
 //! * Coalesce partitions: [`CoalescePartitionsExec`](physical_plan::coalesce_partitions::CoalescePartitionsExec)
 //! * Limit: [`LocalLimitExec`](physical_plan::limit::LocalLimitExec) and [`GlobalLimitExec`](physical_plan::limit::GlobalLimitExec)

datafusion/src/physical_optimizer/aggregate_statistics.rs

@@ -21,8 +21,9 @@ use std::sync::Arc;
 use arrow::datatypes::Schema;
 
 use crate::execution::context::ExecutionConfig;
+use crate::physical_plan::aggregations::hash_aggregate::HashAggregateExec;
+use crate::physical_plan::aggregations::AggregateMode;
 use crate::physical_plan::empty::EmptyExec;
-use crate::physical_plan::hash_aggregate::{AggregateMode, HashAggregateExec};
 use crate::physical_plan::projection::ProjectionExec;
 use crate::physical_plan::{
     expressions, AggregateExpr, ColumnStatistics, ExecutionPlan, Statistics,
@@ -220,11 +221,11 @@ mod tests {
 
     use crate::error::Result;
     use crate::logical_plan::Operator;
+    use crate::physical_plan::aggregations::hash_aggregate::HashAggregateExec;
     use crate::physical_plan::coalesce_partitions::CoalescePartitionsExec;
     use crate::physical_plan::common;
     use crate::physical_plan::expressions::Count;
     use crate::physical_plan::filter::FilterExec;
-    use crate::physical_plan::hash_aggregate::HashAggregateExec;
     use crate::physical_plan::memory::MemoryExec;
 
     /// Mock data using a MemoryExec which has an exact count statistic

datafusion/src/physical_plan/hash_aggregate.rs → datafusion/src/physical_plan/aggregations/hash_aggregate.rs

@@ -34,15 +34,14 @@ use crate::physical_plan::{
 };
 use crate::{
     error::{DataFusionError, Result},
-    execution::memory_management::MemoryConsumerId,
     scalar::ScalarValue,
 };
 
 use arrow::{
     array::*,
     buffer::MutableBuffer,
     compute,
-    datatypes::{DataType, Field, Schema, SchemaRef},
+    datatypes::{DataType, Schema, SchemaRef},
     error::{ArrowError, Result as ArrowResult},
     record_batch::RecordBatch,
 };
@@ -51,27 +50,15 @@ use pin_project_lite::pin_project;
 
 use async_trait::async_trait;
 
-use super::metrics::{
+use crate::physical_plan::aggregations::{
+    aggregate_expressions, create_accumulators, evaluate, evaluate_many, AggregateMode,
+};
+use crate::physical_plan::metrics::{
     self, BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput,
 };
-use super::Statistics;
-use super::{expressions::Column, RecordBatchStream, SendableRecordBatchStream};
-
-/// Hash aggregate modes
-#[derive(Debug, Copy, Clone, PartialEq, Eq)]
-pub enum AggregateMode {
-    /// Partial aggregate that can be applied in parallel across input partitions
-    Partial,
-    /// Final aggregate that produces a single partition of output
-    Final,
-    /// Final aggregate that works on pre-partitioned data.
-    ///
-    /// This requires the invariant that all rows with a particular
-    /// grouping key are in the same partitions, such as is the case
-    /// with Hash repartitioning on the group keys. If a group key is
-    /// duplicated, duplicate groups would be produced
-    FinalPartitioned,
-}
+use crate::physical_plan::Statistics;
+use crate::physical_plan::{RecordBatchStream, SendableRecordBatchStream};
+use arrow::datatypes::Field;
 
 /// Hash aggregate execution plan
 #[derive(Debug)]
@@ -213,11 +200,8 @@ impl ExecutionPlan for HashAggregateExec {
 
         let baseline_metrics = BaselineMetrics::new(&self.metrics, partition);
 
-        let streamer_id = MemoryConsumerId::new(partition);
-
         if self.group_expr.is_empty() {
             Ok(Box::pin(HashAggregateStream::new(
-                streamer_id,
                 self.mode,
                 self.schema.clone(),
                 self.aggr_expr.clone(),
@@ -666,75 +650,8 @@ impl RecordBatchStream for GroupedHashAggregateStream {
     }
 }
 
-/// Evaluates expressions against a record batch.
-fn evaluate(
-    expr: &[Arc<dyn PhysicalExpr>],
-    batch: &RecordBatch,
-) -> Result<Vec<ArrayRef>> {
-    expr.iter()
-        .map(|expr| expr.evaluate(batch))
-        .map(|r| r.map(|v| v.into_array(batch.num_rows())))
-        .collect::<Result<Vec<_>>>()
-}
-
-/// Evaluates expressions against a record batch.
-fn evaluate_many(
-    expr: &[Vec<Arc<dyn PhysicalExpr>>],
-    batch: &RecordBatch,
-) -> Result<Vec<Vec<ArrayRef>>> {
-    expr.iter()
-        .map(|expr| evaluate(expr, batch))
-        .collect::<Result<Vec<_>>>()
-}
-
-/// uses `state_fields` to build a vec of physical column expressions required to merge the
-/// AggregateExpr' accumulator's state.
-///
-/// `index_base` is the starting physical column index for the next expanded state field.
-fn merge_expressions(
-    index_base: usize,
-    expr: &Arc<dyn AggregateExpr>,
-) -> Result<Vec<Arc<dyn PhysicalExpr>>> {
-    Ok(expr
-        .state_fields()?
-        .iter()
-        .enumerate()
-        .map(|(idx, f)| {
-            Arc::new(Column::new(f.name(), index_base + idx)) as Arc<dyn PhysicalExpr>
-        })
-        .collect::<Vec<_>>())
-}
-
-/// returns physical expressions to evaluate against a batch
-/// The expressions are different depending on `mode`:
-/// * Partial: AggregateExpr::expressions
-/// * Final: columns of `AggregateExpr::state_fields()`
-fn aggregate_expressions(
-    aggr_expr: &[Arc<dyn AggregateExpr>],
-    mode: &AggregateMode,
-    col_idx_base: usize,
-) -> Result<Vec<Vec<Arc<dyn PhysicalExpr>>>> {
-    match mode {
-        AggregateMode::Partial => {
-            Ok(aggr_expr.iter().map(|agg| agg.expressions()).collect())
-        }
-        // in this mode, we build the merge expressions of the aggregation
-        AggregateMode::Final | AggregateMode::FinalPartitioned => {
-            let mut col_idx_base = col_idx_base;
-            Ok(aggr_expr
-                .iter()
-                .map(|agg| {
-                    let exprs = merge_expressions(col_idx_base, agg)?;
-                    col_idx_base += exprs.len();
-                    Ok(exprs)
-                })
-                .collect::<Result<Vec<_>>>()?)
-        }
-    }
-}
-
 pin_project! {
-    struct HashAggregateStream {
+    pub(crate) struct HashAggregateStream {
         schema: SchemaRef,
         #[pin]
         output: futures::channel::oneshot::Receiver<ArrowResult<RecordBatch>>,
@@ -744,7 +661,6 @@ pin_project! {
 
 /// Special case aggregate with no groups
 async fn compute_hash_aggregate(
-    _id: MemoryConsumerId,
     mode: AggregateMode,
     schema: SchemaRef,
     aggr_expr: Vec<Arc<dyn AggregateExpr>>,
@@ -781,7 +697,6 @@ async fn compute_hash_aggregate(
 impl HashAggregateStream {
     /// Create a new HashAggregateStream
     pub fn new(
-        id: MemoryConsumerId,
         mode: AggregateMode,
         schema: SchemaRef,
         aggr_expr: Vec<Arc<dyn AggregateExpr>>,
@@ -794,7 +709,6 @@ impl HashAggregateStream {
         let elapsed_compute = baseline_metrics.elapsed_compute().clone();
         tokio::spawn(async move {
             let result = compute_hash_aggregate(
-                id,
                 mode,
                 schema_clone,
                 aggr_expr,
@@ -983,15 +897,6 @@ fn create_batch_from_map(
     RecordBatch::try_new(Arc::new(output_schema.to_owned()), columns)
 }
 
-fn create_accumulators(
-    aggr_expr: &[Arc<dyn AggregateExpr>],
-) -> Result<Vec<AccumulatorItem>> {
-    aggr_expr
-        .iter()
-        .map(|expr| expr.create_accumulator())
-        .collect::<Result<Vec<_>>>()
-}
-
 /// returns a vector of ArrayRefs, where each entry corresponds to either the
 /// final value (mode = Final) or states (mode = Partial)
 fn finalize_aggregation(
@@ -1026,7 +931,7 @@ fn finalize_aggregation(
 mod tests {
 
     use arrow::array::{Float64Array, UInt32Array};
-    use arrow::datatypes::DataType;
+    use arrow::datatypes::{DataType, Field};
 
     use super::*;
     use crate::physical_plan::expressions::{col, Avg};