Co-authored-by: Sarvesh Tandon <[email protected]>

eggstrain/src/execution/operators/aggregate/datafusion_aggregate.rs

@@ -0,0 +1,92 @@
+use arrow::array::{ArrayRef, RecordBatch};
+use datafusion::common::Result;
+use datafusion::physical_plan::{aggregates::PhysicalGroupBy, PhysicalExpr};
+use std::sync::Arc;
+
+/// Evaluates expressions against a record batch.
+pub(crate) fn evaluate(
+    expr: &[Arc<dyn PhysicalExpr>],
+    batch: &RecordBatch,
+) -> Result<Vec<ArrayRef>> {
+    expr.iter()
+        .map(|expr| {
+            expr.evaluate(batch)
+                .and_then(|v| v.into_array(batch.num_rows()))
+        })
+        .collect()
+}
+
+/// Evaluates expressions against a record batch.
+pub(crate) fn evaluate_many(
+    expr: &[Vec<Arc<dyn PhysicalExpr>>],
+    batch: &RecordBatch,
+) -> Result<Vec<Vec<ArrayRef>>> {
+    expr.iter().map(|expr| evaluate(expr, batch)).collect()
+}
+
+pub(crate) fn evaluate_optional(
+    expr: &[Option<Arc<dyn PhysicalExpr>>],
+    batch: &RecordBatch,
+) -> Result<Vec<Option<ArrayRef>>> {
+    expr.iter()
+        .map(|expr| {
+            expr.as_ref()
+                .map(|expr| {
+                    expr.evaluate(batch)
+                        .and_then(|v| v.into_array(batch.num_rows()))
+                })
+                .transpose()
+        })
+        .collect()
+}
+
+/// Evaluate a group by expression against a `RecordBatch`
+///
+/// Arguments:
+/// `group_by`: the expression to evaluate
+/// `batch`: the `RecordBatch` to evaluate against
+///
+/// Returns: A Vec of Vecs of Array of results
+/// The outer Vect appears to be for grouping sets
+/// The inner Vect contains the results per expression
+/// The inner-inner Array contains the results per row
+pub(crate) fn evaluate_group_by(
+    group_by: &PhysicalGroupBy,
+    batch: &RecordBatch,
+) -> Result<Vec<Vec<ArrayRef>>> {
+    let exprs: Vec<ArrayRef> = group_by
+        .expr()
+        .iter()
+        .map(|(expr, _)| {
+            let value = expr.evaluate(batch)?;
+            value.into_array(batch.num_rows())
+        })
+        .collect::<Result<Vec<_>>>()?;
+
+    let null_exprs: Vec<ArrayRef> = group_by
+        .null_expr()
+        .iter()
+        .map(|(expr, _)| {
+            let value = expr.evaluate(batch)?;
+            value.into_array(batch.num_rows())
+        })
+        .collect::<Result<Vec<_>>>()?;
+
+    Ok(group_by
+        .groups()
+        .iter()
+        .map(|group| {
+            group
+                .iter()
+                .enumerate()
+                .map(|(idx, is_null)| {
+                    if *is_null {
+                        null_exprs[idx].clone()
+                    } else {
+                        exprs[idx].clone()
+                    }
+                })
+                .collect()
+        })
+        .collect())
+}

eggstrain/src/execution/operators/aggregate/mod.rs

@@ -0,0 +1,10 @@
+mod datafusion_aggregate;
+use datafusion_aggregate::*;
+
+mod operator;
+use operator::*;
+
+
+
+
+