[DF] Add support for filtered groupby aggregations

dask_planner/Cargo.toml

@@ -12,11 +12,11 @@ rust-version = "1.62"
 tokio = { version = "1.0", features = ["macros", "rt", "rt-multi-thread", "sync", "fs", "parking_lot"] }
 rand = "0.7"
 pyo3 = { version = "0.17.1", features = ["extension-module", "abi3", "abi3-py38"] }
-arrow = { version = "22.0.0", features = ["prettyprint"] }
-datafusion-sql = "12.0.0"
-datafusion-expr = "12.0.0"
-datafusion-common = "12.0.0"
-datafusion-optimizer = "12.0.0"
+arrow = { version = "23.0.0", features = ["prettyprint"] }
+datafusion-sql = { git="https://github.com/apache/arrow-datafusion/", rev = "1261741af2a5e142fa0c7916e759859cc18ea59a" }
+datafusion-expr = { git="https://github.com/apache/arrow-datafusion/", rev = "1261741af2a5e142fa0c7916e759859cc18ea59a" }
+datafusion-common = { git="https://github.com/apache/arrow-datafusion/", rev = "1261741af2a5e142fa0c7916e759859cc18ea59a" }
+datafusion-optimizer = { git="https://github.com/apache/arrow-datafusion/", rev = "1261741af2a5e142fa0c7916e759859cc18ea59a" }
 uuid = { version = "0.8", features = ["v4"] }
 mimalloc = { version = "*", default-features = false }
 parking_lot = "0.12"

dask_planner/src/dialect.rs

@@ -33,4 +33,8 @@ impl Dialect for DaskDialect {
     fn is_proper_identifier_inside_quotes(&self, mut _chars: Peekable<Chars<'_>>) -> bool {
         true
     }
+    /// Determine if FILTER (WHERE ...) filters are allowed during aggregations
+    fn supports_filter_during_aggregation(&self) -> bool {
+        true
+    }
 }

dask_planner/src/expression.rs

@@ -526,6 +526,23 @@ impl PyExpr {
         }))
     }
 
+    #[pyo3(name = "getFilterExpr")]
+    pub fn get_filter_expr(&self) -> PyResult<Option<PyExpr>> {
+        match &self.expr {
+            Expr::AggregateFunction { filter, .. } | Expr::AggregateUDF { filter, .. } => {
+                match filter {
+                    Some(filter) => {
+                        Ok(Some(PyExpr::from(*filter.clone(), self.input_plan.clone())))
+                    }
+                    None => Ok(None),
+                }
+            }
+            _ => Err(py_type_err(
+                "getFilterExpr() - Non-aggregate expression encountered",
+            )),
+        }
+    }
+
     /// TODO: I can't express how much I dislike explicity listing all of these methods out
     /// but PyO3 makes it necessary since its annotations cannot be used in trait impl blocks
     #[pyo3(name = "getFloat32Value")]

dask_sql/physical/rel/logical/aggregate.py

@@ -336,24 +336,6 @@ def _collect_aggregations(
         new_columns = {}
         new_mappings = {}
 
-        # convert and assign any input columns that don't currently exist
-        for expr in agg.getNamedAggCalls():
-            for expr in agg.getArgs(expr):
-                key = expr.column_name(input_rel)
-                if key in cc._frontend_backend_mapping:
-                    continue
-                random_name = new_temporary_column(df)
-                new_columns[random_name] = RexConverter.convert(
-                    input_rel, expr, dc, context=context
-                )
-                new_mappings[key] = random_name
-
-        if new_columns:
-            df = df.assign(**new_columns)
-
-        for key, backend_column_name in new_mappings.items():
-            cc = cc.add(key, backend_column_name)
-
         for expr in agg.getNamedAggCalls():
             # Determine the aggregation function to use
             assert expr.getExprType() in {
@@ -364,9 +346,20 @@ def _collect_aggregations(
             schema_name = context.schema_name
             aggregation_name = agg.getAggregationFuncName(expr).lower()
 
-            # Gather information about the input column
+            # Gather information about input columns
             inputs = agg.getArgs(expr)
 
+            # Compute any input columns that don't yet exist
+            for input_expr in inputs:
+                input_col = input_expr.column_name(input_rel)
+                if input_col in cc._frontend_backend_mapping:
+                    continue
+                random_name = new_temporary_column(df)
+                new_columns[random_name] = RexConverter.convert(
+                    input_rel, expr, dc, context=context
+                )
+                new_mappings[input_col] = random_name
+
             # TODO: This if statement is likely no longer needed but left here for the time being just in case
             if aggregation_name == "regr_count":
                 is_null = IsNullOperation()
@@ -396,10 +389,21 @@ def _collect_aggregations(
             else:
                 raise NotImplementedError("Can not cope with more than one input")
 
-            # TODO: DataFusion does not yet have the concept of "filters" in aggregations
-            filter_column = None
-            # if expr.hasFilter():
-            #     filter_column = cc.get_backend_by_frontend_index(expr.filterArg)
+            # Compute filter column if it doesn't yet exist
+            filter_expr = expr.getFilterExpr()
+            if filter_expr is not None:
+                filter_col = filter_expr.column_name(input_rel)
+                if filter_col not in cc._frontend_backend_mapping:
+                    random_name = new_temporary_column(df)
+                    new_columns[random_name] = RexConverter.convert(
+                        input_rel, filter_expr, dc, context=context
+                    )
+                    new_mappings[filter_col] = random_name
+                    filter_col = random_name
+                else:
+                    filter_col = cc.get_backend_by_frontend_name(filter_col)
+            else:
+                filter_col = None
 
             try:
                 aggregation_function = self.AGGREGATION_MAPPING[aggregation_name]
@@ -422,11 +426,17 @@ def _collect_aggregations(
             output_col = expr.toString()
 
             # Store the aggregation
-            key = filter_column
-            value = (input_col, output_col, aggregation_function)
-            collected_aggregations[key].append(value)
+            collected_aggregations[filter_col].append(
+                (input_col, output_col, aggregation_function)
+            )
             output_column_order.append(output_col)
 
+        # Add any required columns to table and column container
+        if new_columns:
+            df = df.assign(**new_columns)
+        for key, backend_column_name in new_mappings.items():
+            cc = cc.add(key, backend_column_name)
+
         return collected_aggregations, output_column_order, df, cc
 
     def _perform_aggregation(

tests/integration/test_groupby.py

@@ -110,9 +110,6 @@ def test_group_by_all(c, df):
     assert_eq(result_df, expected_df)
 
 
-@pytest.mark.skip(
-    reason="WIP DataFusion - https://github.com/dask-contrib/dask-sql/issues/463"
-)
 def test_group_by_filtered(c):
     return_df = c.sql(
         """