Add group join physical optimizer

Author: nathanb9

datafusion/physical-optimizer/src/group_join.rs

@@ -0,0 +1,196 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`GroupJoinOptimizer`] replaces an `AggregateExec` directly above a
+//! `HashJoinExec` with a fused [`GroupJoinExec`] when the aggregate's GROUP BY
+//! keys match the join's equi-join keys.
+//!
+//! Based on: Moerkotte & Neumann, "Accelerating Queries with Group-By and Join
+//! by Groupjoin", PVLDB 4(11), 2011.
+
+use std::sync::Arc;
+
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_common::{JoinType, Result};
+use datafusion_physical_expr::physical_exprs_equal;
+use datafusion_physical_plan::ExecutionPlan;
+use datafusion_physical_plan::aggregates::{AggregateExec, AggregateMode};
+use datafusion_physical_plan::joins::HashJoinExec;
+use datafusion_physical_plan::joins::group_join::GroupJoinExec;
+use datafusion_physical_plan::projection::ProjectionExec;
+
+use crate::PhysicalOptimizerRule;
+
+/// Replaces `AggregateExec(HashJoinExec)` with a fused `GroupJoinExec` when:
+///
+/// 1. The aggregate mode is `Single`, `SinglePartitioned`, or `Partial`
+/// 2. The aggregate has at least one aggregate expression (not just DISTINCT)
+/// 3. The aggregate has no GROUPING SETS
+/// 4. The input is a `HashJoinExec` (possibly through a `ProjectionExec`)
+/// 5. The join type is `Inner` or `Left`
+/// 6. The join has no residual filter (equi-join only)
+/// 7. The GROUP BY expressions exactly match the left join keys
+/// 8. All aggregate functions support `GroupsAccumulator`
+///
+/// This rule should run after `CombinePartialFinalAggregate` (which may
+/// collapse two-phase aggregation into Single mode) and after `JoinSelection`
+/// (which decides build/probe sides).
+#[derive(Default, Debug)]
+pub struct GroupJoinOptimizer {}
+
+impl GroupJoinOptimizer {
+    /// Create a new `GroupJoinOptimizer`.
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl PhysicalOptimizerRule for GroupJoinOptimizer {
+    fn optimize(
+        &self,
+        plan: Arc<dyn ExecutionPlan>,
+        _config: &ConfigOptions,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        plan.transform_down(|plan| {
+            let Some(agg_exec) = plan.downcast_ref::<AggregateExec>() else {
+                return Ok(Transformed::no(plan));
+            };
+
+            if !matches!(
+                agg_exec.mode(),
+                AggregateMode::Single
+                    | AggregateMode::SinglePartitioned
+                    | AggregateMode::Partial
+            ) {
+                return Ok(Transformed::no(plan));
+            }
+
+            // Must have actual aggregate functions (not just GROUP BY for DISTINCT)
+            let aggr_exprs = agg_exec.aggr_expr();
+            if aggr_exprs.is_empty() {
+                return Ok(Transformed::no(plan));
+            }
+
+            // No GROUPING SETS
+            if agg_exec.group_expr().groups().len() > 1 {
+                return Ok(Transformed::no(plan));
+            }
+
+            // Find HashJoinExec (possibly through a ProjectionExec)
+            let input = agg_exec.input();
+            let hash_join: &HashJoinExec;
+            if let Some(hj) = input.downcast_ref::<HashJoinExec>() {
+                hash_join = hj;
+            } else if let Some(proj) = input.downcast_ref::<ProjectionExec>() {
+                if let Some(hj) = proj.input().downcast_ref::<HashJoinExec>() {
+                    hash_join = hj;
+                } else {
+                    return Ok(Transformed::no(plan));
+                }
+            } else {
+                return Ok(Transformed::no(plan));
+            };
+
+            // Inner and Left joins
+            if !matches!(
+                hash_join.join_type(),
+                JoinType::Inner | JoinType::Left
+            ) {
+                return Ok(Transformed::no(plan));
+            }
+
+            // No residual join filter (equi-join only)
+            if hash_join.filter().is_some() {
+                return Ok(Transformed::no(plan));
+            }
+
+            // GROUP BY keys must exactly match left join keys
+            let group_exprs: Vec<_> = agg_exec
+                .group_expr()
+                .expr()
+                .iter()
+                .map(|(expr, _)| Arc::clone(expr))
+                .collect();
+
+            let join_on = hash_join.on();
+            let left_join_keys: Vec<_> =
+                join_on.iter().map(|(l, _)| Arc::clone(l)).collect();
+
+            if group_exprs.len() != left_join_keys.len() {
+                return Ok(Transformed::no(plan));
+            }
+
+            if !physical_exprs_equal(&group_exprs, &left_join_keys) {
+                return Ok(Transformed::no(plan));
+            }
+
+            // All aggregates must support GroupsAccumulator
+            for agg in aggr_exprs {
+                if !agg.groups_accumulator_supported() {
+                    return Ok(Transformed::no(plan));
+                }
+            }
+
+            // For Inner joins, skip if any aggregate has a literal argument
+            // (e.g., COUNT(*) rewritten as count(Int64(1))). These queries
+            // don't benefit enough from GroupJoin to justify changing the plan.
+            if *hash_join.join_type() == JoinType::Inner {
+                let has_literal_arg = aggr_exprs.iter().any(|agg| {
+                    agg.expressions().iter().any(|expr| {
+                        expr.as_ref()
+                            .downcast_ref::<datafusion_physical_expr::expressions::Literal>()
+                            .is_some()
+                    })
+                });
+                if has_literal_arg {
+                    return Ok(Transformed::no(plan));
+                }
+            }
+
+            // All preconditions met — create GroupJoinExec
+            let group_by_with_names: Vec<_> = agg_exec
+                .group_expr()
+                .expr()
+                .iter()
+                .map(|(expr, name)| (Arc::clone(expr), name.clone()))
+                .collect();
+
+            let group_join = GroupJoinExec::try_new(
+                Arc::clone(hash_join.left()),
+                Arc::clone(hash_join.right()),
+                join_on.to_vec(),
+                *hash_join.join_type(),
+                group_by_with_names,
+                aggr_exprs.to_vec(),
+            )?;
+
+            Ok(Transformed::yes(
+                Arc::new(group_join) as Arc<dyn ExecutionPlan>
+            ))
+        })
+        .data()
+    }
+
+    fn name(&self) -> &str {
+        "group_join"
+    }
+
+    fn schema_check(&self) -> bool {
+        false // Schema changes (aggregate output differs from join output)
+    }
+}

datafusion/physical-optimizer/src/lib.rs

@@ -31,6 +31,7 @@ pub mod enforce_distribution;
 pub mod enforce_sorting;
 pub mod ensure_coop;
 pub mod filter_pushdown;
+pub mod group_join;
 pub mod join_selection;
 pub mod limit_pushdown;
 pub mod limit_pushdown_past_window;

datafusion/physical-optimizer/src/optimizer.rs

@@ -26,6 +26,7 @@ use crate::enforce_distribution::EnforceDistribution;
 use crate::enforce_sorting::EnforceSorting;
 use crate::ensure_coop::EnsureCooperative;
 use crate::filter_pushdown::FilterPushdown;
+use crate::group_join::GroupJoinOptimizer;
 use crate::join_selection::JoinSelection;
 use crate::limit_pushdown::LimitPushdown;
 use crate::limited_distinct_aggregation::LimitedDistinctAggregation;
@@ -177,6 +178,11 @@ impl PhysicalOptimizer {
             Arc::new(EnforceDistribution::new()),
             // The CombinePartialFinalAggregate rule should be applied after the EnforceDistribution rule
             Arc::new(CombinePartialFinalAggregate::new()),
+            // GroupJoinOptimizer fuses Aggregate+HashJoin into GroupJoinExec
+            // when GROUP BY keys match join keys. Runs after
+            // CombinePartialFinalAggregate (which creates SinglePartitioned
+            // aggregates) and JoinSelection (which decides build/probe sides).
+            Arc::new(GroupJoinOptimizer::new()),
             // The EnforceSorting rule is for adding essential local sorting to satisfy the required
             // ordering. Please make sure that the whole plan tree is determined before this rule.
             // Note that one should always run this rule after running the EnforceDistribution rule