datafusion_optimizer/
eliminate_group_by_constant.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! [`EliminateGroupByConstant`] removes constant and functionally redundant
19//! expressions from `GROUP BY` clause
20use crate::optimizer::ApplyOrder;
21use crate::{OptimizerConfig, OptimizerRule};
22
23use std::collections::HashSet;
24
25use datafusion_common::Result;
26use datafusion_common::tree_node::Transformed;
27use datafusion_expr::{Aggregate, Expr, LogicalPlan, LogicalPlanBuilder, Volatility};
28
29/// Optimizer rule that removes constant expressions from `GROUP BY` clause
30/// and places additional projection on top of aggregation, to preserve
31/// original schema
32#[derive(Default, Debug)]
33pub struct EliminateGroupByConstant {}
34
35impl EliminateGroupByConstant {
36    pub fn new() -> Self {
37        Self {}
38    }
39}
40
41impl OptimizerRule for EliminateGroupByConstant {
42    fn supports_rewrite(&self) -> bool {
43        true
44    }
45
46    fn rewrite(
47        &self,
48        plan: LogicalPlan,
49        _config: &dyn OptimizerConfig,
50    ) -> Result<Transformed<LogicalPlan>> {
51        match plan {
52            LogicalPlan::Aggregate(aggregate) => {
53                // Collect bare column references in GROUP BY
54                let group_by_columns: HashSet<&datafusion_common::Column> = aggregate
55                    .group_expr
56                    .iter()
57                    .filter_map(|expr| match expr {
58                        Expr::Column(c) => Some(c),
59                        _ => None,
60                    })
61                    .collect();
62
63                let (redundant, required): (Vec<_>, Vec<_>) = aggregate
64                    .group_expr
65                    .iter()
66                    .partition(|expr| is_redundant_group_expr(expr, &group_by_columns));
67                // Return now if no simplification can be done. We also bail out
68                // if applying the optimization would eliminate all of the
69                // grouping expressions (e.g., GROUP BY on only constant
70                // expressions): this would turn a grouped aggregate into an
71                // ungrouped aggregate, which changes query semantics (grouped
72                // aggregates produce an empty result set on an empty input,
73                // whereas ungrouped aggregates return a single row).
74                if redundant.is_empty() || required.is_empty() {
75                    return Ok(Transformed::no(LogicalPlan::Aggregate(aggregate)));
76                }
77
78                let simplified_aggregate = LogicalPlan::Aggregate(Aggregate::try_new(
79                    aggregate.input,
80                    required.into_iter().cloned().collect(),
81                    aggregate.aggr_expr.clone(),
82                )?);
83
84                let projection_expr =
85                    aggregate.group_expr.into_iter().chain(aggregate.aggr_expr);
86
87                let projection = LogicalPlanBuilder::from(simplified_aggregate)
88                    .project(projection_expr)?
89                    .build()?;
90
91                Ok(Transformed::yes(projection))
92            }
93            _ => Ok(Transformed::no(plan)),
94        }
95    }
96
97    fn name(&self) -> &str {
98        "eliminate_group_by_constant"
99    }
100
101    fn apply_order(&self) -> Option<ApplyOrder> {
102        Some(ApplyOrder::BottomUp)
103    }
104}
105
106/// Checks if a GROUP BY expression is redundant (can be removed without
107/// changing grouping semantics). An expression is redundant if it is a
108/// deterministic function of constants and columns already present as bare
109/// column references in the GROUP BY.
110fn is_redundant_group_expr(
111    expr: &Expr,
112    group_by_columns: &HashSet<&datafusion_common::Column>,
113) -> bool {
114    // Bare column references are never redundant - they define the grouping
115    if matches!(expr, Expr::Column(_)) {
116        return false;
117    }
118    is_deterministic_of(expr, group_by_columns)
119}
120
121/// Returns true if `expr` is a deterministic expression whose only column
122/// references are contained in `known_columns`.
123fn is_deterministic_of(
124    expr: &Expr,
125    known_columns: &HashSet<&datafusion_common::Column>,
126) -> bool {
127    match expr {
128        Expr::Alias(e) => is_deterministic_of(&e.expr, known_columns),
129        Expr::Column(c) => known_columns.contains(c),
130        Expr::Literal(_, _) => true,
131        Expr::BinaryExpr(e) => {
132            is_deterministic_of(&e.left, known_columns)
133                && is_deterministic_of(&e.right, known_columns)
134        }
135        Expr::ScalarFunction(e) => {
136            matches!(
137                e.func.signature().volatility,
138                Volatility::Immutable | Volatility::Stable
139            ) && e
140                .args
141                .iter()
142                .all(|arg| is_deterministic_of(arg, known_columns))
143        }
144        Expr::Cast(e) => is_deterministic_of(&e.expr, known_columns),
145        Expr::TryCast(e) => is_deterministic_of(&e.expr, known_columns),
146        Expr::Negative(e) => is_deterministic_of(e, known_columns),
147        _ => false,
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    use crate::OptimizerContext;
155    use crate::assert_optimized_plan_eq_snapshot;
156    use crate::test::*;
157
158    use arrow::datatypes::DataType;
159    use datafusion_expr::expr::ScalarFunction;
160    use datafusion_expr::{
161        ColumnarValue, LogicalPlanBuilder, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl,
162        Signature, TypeSignature, col, lit,
163    };
164
165    use datafusion_functions_aggregate::expr_fn::count;
166
167    use std::sync::Arc;
168
169    macro_rules! assert_optimized_plan_equal {
170        (
171            $plan:expr,
172            @ $expected:literal $(,)?
173        ) => {{
174            let optimizer_ctx = OptimizerContext::new().with_max_passes(1);
175            let rules: Vec<Arc<dyn crate::OptimizerRule + Send + Sync>> = vec![Arc::new(EliminateGroupByConstant::new())];
176            assert_optimized_plan_eq_snapshot!(
177                optimizer_ctx,
178                rules,
179                $plan,
180                @ $expected,
181            )
182        }};
183    }
184
185    #[derive(Debug, PartialEq, Eq, Hash)]
186    struct ScalarUDFMock {
187        signature: Signature,
188    }
189
190    impl ScalarUDFMock {
191        fn new_with_volatility(volatility: Volatility) -> Self {
192            Self {
193                signature: Signature::new(TypeSignature::Any(1), volatility),
194            }
195        }
196    }
197
198    impl ScalarUDFImpl for ScalarUDFMock {
199        fn name(&self) -> &str {
200            "scalar_fn_mock"
201        }
202        fn signature(&self) -> &Signature {
203            &self.signature
204        }
205        fn return_type(&self, _args: &[DataType]) -> Result<DataType> {
206            Ok(DataType::Int32)
207        }
208        fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
209            unimplemented!()
210        }
211    }
212
213    #[test]
214    fn test_eliminate_gby_literal() -> Result<()> {
215        let scan = test_table_scan()?;
216        let plan = LogicalPlanBuilder::from(scan)
217            .aggregate(vec![col("a"), lit(1u32)], vec![count(col("c"))])?
218            .build()?;
219
220        assert_optimized_plan_equal!(plan, @r"
221        Projection: test.a, UInt32(1), count(test.c)
222          Aggregate: groupBy=[[test.a]], aggr=[[count(test.c)]]
223            TableScan: test
224        ")
225    }
226
227    #[test]
228    fn test_no_op_only_constant_with_aggregate() -> Result<()> {
229        let scan = test_table_scan()?;
230        let plan = LogicalPlanBuilder::from(scan)
231            .aggregate(vec![lit("test"), lit(123u32)], vec![count(col("c"))])?
232            .build()?;
233
234        assert_optimized_plan_equal!(plan, @r#"
235        Aggregate: groupBy=[[Utf8("test"), UInt32(123)]], aggr=[[count(test.c)]]
236          TableScan: test
237        "#)
238    }
239
240    #[test]
241    fn test_no_op_no_constants() -> Result<()> {
242        let scan = test_table_scan()?;
243        let plan = LogicalPlanBuilder::from(scan)
244            .aggregate(vec![col("a"), col("b")], vec![count(col("c"))])?
245            .build()?;
246
247        assert_optimized_plan_equal!(plan, @r"
248        Aggregate: groupBy=[[test.a, test.b]], aggr=[[count(test.c)]]
249          TableScan: test
250        ")
251    }
252
253    #[test]
254    fn test_no_op_only_constant() -> Result<()> {
255        let scan = test_table_scan()?;
256        let plan = LogicalPlanBuilder::from(scan)
257            .aggregate(vec![lit(123u32)], Vec::<Expr>::new())?
258            .build()?;
259
260        assert_optimized_plan_equal!(plan, @r"
261        Aggregate: groupBy=[[UInt32(123)]], aggr=[[]]
262          TableScan: test
263        ")
264    }
265
266    #[test]
267    fn test_eliminate_constant_with_alias() -> Result<()> {
268        let scan = test_table_scan()?;
269        let plan = LogicalPlanBuilder::from(scan)
270            .aggregate(
271                vec![lit(123u32).alias("const"), col("a")],
272                vec![count(col("c"))],
273            )?
274            .build()?;
275
276        assert_optimized_plan_equal!(plan, @r"
277        Projection: UInt32(123) AS const, test.a, count(test.c)
278          Aggregate: groupBy=[[test.a]], aggr=[[count(test.c)]]
279            TableScan: test
280        ")
281    }
282
283    #[test]
284    fn test_eliminate_scalar_fn_with_constant_arg() -> Result<()> {
285        let udf = ScalarUDF::new_from_impl(ScalarUDFMock::new_with_volatility(
286            Volatility::Immutable,
287        ));
288        let udf_expr =
289            Expr::ScalarFunction(ScalarFunction::new_udf(udf.into(), vec![lit(123u32)]));
290        let scan = test_table_scan()?;
291        let plan = LogicalPlanBuilder::from(scan)
292            .aggregate(vec![udf_expr, col("a")], vec![count(col("c"))])?
293            .build()?;
294
295        assert_optimized_plan_equal!(plan, @r"
296        Projection: scalar_fn_mock(UInt32(123)), test.a, count(test.c)
297          Aggregate: groupBy=[[test.a]], aggr=[[count(test.c)]]
298            TableScan: test
299        ")
300    }
301
302    #[test]
303    fn test_eliminate_deterministic_expr_of_group_by_column() -> Result<()> {
304        let scan = test_table_scan()?;
305        // GROUP BY a, a - 1, a - 2, a - 3  ->  GROUP BY a
306        let plan = LogicalPlanBuilder::from(scan)
307            .aggregate(
308                vec![
309                    col("a"),
310                    col("a") - lit(1u32),
311                    col("a") - lit(2u32),
312                    col("a") - lit(3u32),
313                ],
314                vec![count(col("c"))],
315            )?
316            .build()?;
317
318        assert_optimized_plan_equal!(plan, @r"
319        Projection: test.a, test.a - UInt32(1), test.a - UInt32(2), test.a - UInt32(3), count(test.c)
320          Aggregate: groupBy=[[test.a]], aggr=[[count(test.c)]]
321            TableScan: test
322        ")
323    }
324
325    #[test]
326    fn test_no_eliminate_independent_columns() -> Result<()> {
327        // GROUP BY a, b - 1 should NOT eliminate b - 1 (b is not a group by column)
328        let scan = test_table_scan()?;
329        let plan = LogicalPlanBuilder::from(scan)
330            .aggregate(vec![col("a"), col("b") - lit(1u32)], vec![count(col("c"))])?
331            .build()?;
332
333        assert_optimized_plan_equal!(plan, @r"
334        Aggregate: groupBy=[[test.a, test.b - UInt32(1)]], aggr=[[count(test.c)]]
335          TableScan: test
336        ")
337    }
338
339    #[test]
340    fn test_no_op_volatile_scalar_fn_with_constant_arg() -> Result<()> {
341        let udf = ScalarUDF::new_from_impl(ScalarUDFMock::new_with_volatility(
342            Volatility::Volatile,
343        ));
344        let udf_expr =
345            Expr::ScalarFunction(ScalarFunction::new_udf(udf.into(), vec![lit(123u32)]));
346        let scan = test_table_scan()?;
347        let plan = LogicalPlanBuilder::from(scan)
348            .aggregate(vec![udf_expr, col("a")], vec![count(col("c"))])?
349            .build()?;
350
351        assert_optimized_plan_equal!(plan, @r"
352        Aggregate: groupBy=[[scalar_fn_mock(UInt32(123)), test.a]], aggr=[[count(test.c)]]
353          TableScan: test
354        ")
355    }
356}
datafusion_optimizer/eliminate_group_by_constant.rs

datafusion_optimizer/
eliminate_group_by_constant.rs
