Rule ReplaceDistinctWithAggregate

datafusion/core/src/physical_plan/planner.rs

@@ -27,8 +27,8 @@ use crate::datasource::source_as_provider;
 use crate::execution::context::{ExecutionProps, SessionState};
 use crate::logical_expr::utils::generate_sort_key;
 use crate::logical_expr::{
-    Aggregate, Distinct, EmptyRelation, Join, Projection, Sort, SubqueryAlias, TableScan,
-    Unnest, Window,
+    Aggregate, EmptyRelation, Join, Projection, Sort, SubqueryAlias, TableScan, Unnest,
+    Window,
 };
 use crate::logical_expr::{
     CrossJoin, Expr, LogicalPlan, Partitioning as LogicalPartitioning, PlanType,
@@ -65,7 +65,6 @@ use datafusion_expr::expr::{
 };
 use datafusion_expr::expr_rewriter::unnormalize_cols;
 use datafusion_expr::logical_plan::builder::wrap_projection_for_join_if_necessary;
-use datafusion_expr::utils::expand_wildcard;
 use datafusion_expr::{logical_plan, StringifiedPlan};
 use datafusion_expr::{WindowFrame, WindowFrameBound};
 use datafusion_optimizer::utils::unalias;
@@ -710,17 +709,6 @@ impl DefaultPhysicalPlanner {
                         physical_input_schema.clone(),
                     )?))
                 }
-                LogicalPlan::Distinct(Distinct { input }) => {
-                    // Convert distinct to groupby with no aggregations
-                    let group_expr = expand_wildcard(input.schema(), input)?;
-                    let aggregate = LogicalPlan::Aggregate(Aggregate::try_new_with_schema(
-                        input.clone(),
-                        group_expr,
-                        vec![],
-                        input.schema().clone(), // input schema and aggregate schema are the same in this case
-                    )?);
-                    Ok(self.create_initial_plan(&aggregate, session_state).await?)
-                }
                 LogicalPlan::Projection(Projection { input, expr, .. }) => {
                     let input_exec = self.create_initial_plan(input, session_state).await?;
                     let input_schema = input.as_ref().schema();
@@ -1207,6 +1195,11 @@ impl DefaultPhysicalPlanner {
                 LogicalPlan::Explain(_) => Err(DataFusionError::Internal(
                     "Unsupported logical plan: Explain must be root of the plan".to_string(),
                 )),
+                LogicalPlan::Distinct(_) => {
+                    Err(DataFusionError::Internal(
+                        "Unsupported logical plan: Distinct should be replaced to Aggregate".to_string(),
+                    ))
+                }
                 LogicalPlan::Analyze(a) => {
                     let input = self.create_initial_plan(&a.input, session_state).await?;
                     let schema = SchemaRef::new((*a.schema).clone().into());

datafusion/core/tests/sql/union.rs

@@ -38,13 +38,15 @@ async fn union_with_except_input() -> Result<()> {
         "Explain [plan_type:Utf8, plan:Utf8]",
         "  Union [name:UInt8;N]",
         "    LeftAnti Join: t1.name = t2.name [name:UInt8;N]",
-        "      Distinct: [name:UInt8;N]",
-        "        TableScan: t1 projection=[name] [name:UInt8;N]",
+        "      Aggregate: groupBy=[[t1.name]], aggr=[[]] [name:UInt8;N]",
+        "        Projection: t1.name [name:UInt8;N]",
+        "          TableScan: t1 projection=[name] [name:UInt8;N]",
         "      Projection: t2.name [name:UInt8;N]",
         "        TableScan: t2 projection=[name] [name:UInt8;N]",
         "    LeftAnti Join: t2.name = t1.name [name:UInt8;N]",
-        "      Distinct: [name:UInt8;N]",
-        "        TableScan: t2 projection=[name] [name:UInt8;N]",
+        "      Aggregate: groupBy=[[t2.name]], aggr=[[]] [name:UInt8;N]",
+        "        Projection: t2.name [name:UInt8;N]",
+        "          TableScan: t2 projection=[name] [name:UInt8;N]",
         "      Projection: t1.name [name:UInt8;N]",
         "        TableScan: t1 projection=[name] [name:UInt8;N]",
     ];
@@ -79,14 +81,16 @@ async fn union_with_type_coercion() -> Result<()> {
         "Explain [plan_type:Utf8, plan:Utf8]",
         "  Union [id:Int32;N, name:UInt8;N]",
         "    LeftAnti Join: t1.id = CAST(t2.id AS Int32), t1.name = t2.name [id:Int32;N, name:UInt8;N]",
-        "      Distinct: [id:Int32;N, name:UInt8;N]",
-        "        TableScan: t1 projection=[id, name] [id:Int32;N, name:UInt8;N]",
+        "      Aggregate: groupBy=[[t1.id, t1.name]], aggr=[[]] [id:Int32;N, name:UInt8;N]",
+        "        Projection: t1.id, t1.name [id:Int32;N, name:UInt8;N]",
+        "          TableScan: t1 projection=[id, name] [id:Int32;N, name:UInt8;N]",
         "      Projection: t2.id, t2.name [id:UInt8;N, name:UInt8;N]",
         "        TableScan: t2 projection=[id, name] [id:UInt8;N, name:UInt8;N]",
         "    Projection: CAST(t2.id AS Int32) AS id, t2.name [id:Int32;N, name:UInt8;N]",
         "      LeftAnti Join: CAST(t2.id AS Int32) = t1.id, t2.name = t1.name [id:UInt8;N, name:UInt8;N]",
-        "        Distinct: [id:UInt8;N, name:UInt8;N]",
-        "          TableScan: t2 projection=[id, name] [id:UInt8;N, name:UInt8;N]",
+        "        Aggregate: groupBy=[[t2.id, t2.name]], aggr=[[]] [id:UInt8;N, name:UInt8;N]",
+        "          Projection: t2.id, t2.name [id:UInt8;N, name:UInt8;N]",
+        "            TableScan: t2 projection=[id, name] [id:UInt8;N, name:UInt8;N]",
         "        Projection: t1.id, t1.name [id:Int32;N, name:UInt8;N]",
         "          TableScan: t1 projection=[id, name] [id:Int32;N, name:UInt8;N]",
     ];

datafusion/optimizer/src/lib.rs

@@ -37,6 +37,7 @@ pub mod single_distinct_to_groupby;
 pub mod type_coercion;
 pub mod utils;
 
+pub mod replace_distinct_aggregate;
 pub mod rewrite_disjunctive_predicate;
 #[cfg(test)]
 pub mod test;

datafusion/optimizer/src/optimizer.rs

@@ -31,6 +31,7 @@ use crate::propagate_empty_relation::PropagateEmptyRelation;
 use crate::push_down_filter::PushDownFilter;
 use crate::push_down_limit::PushDownLimit;
 use crate::push_down_projection::PushDownProjection;
+use crate::replace_distinct_aggregate::ReplaceDistinctWithAggregate;
 use crate::rewrite_disjunctive_predicate::RewriteDisjunctivePredicate;
 use crate::scalar_subquery_to_join::ScalarSubqueryToJoin;
 use crate::simplify_expressions::SimplifyExpressions;
@@ -207,6 +208,7 @@ impl Optimizer {
             Arc::new(TypeCoercion::new()),
             Arc::new(SimplifyExpressions::new()),
             Arc::new(UnwrapCastInComparison::new()),
+            Arc::new(ReplaceDistinctWithAggregate::new()),
             Arc::new(DecorrelateWhereExists::new()),
             Arc::new(DecorrelateWhereIn::new()),
             Arc::new(ScalarSubqueryToJoin::new()),

datafusion/optimizer/src/replace_distinct_aggregate.rs

@@ -0,0 +1,103 @@
+// 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.
+
+use crate::optimizer::ApplyOrder;
+use crate::{OptimizerConfig, OptimizerRule};
+use datafusion_common::Result;
+use datafusion_expr::utils::expand_wildcard;
+use datafusion_expr::Distinct;
+use datafusion_expr::{Aggregate, LogicalPlan};
+use ApplyOrder::BottomUp;
+
+/// Optimizer that replaces logical [[Distinct]] with a logical [[Aggregate]]
+///
+/// ```text
+/// SELECT DISTINCT a, b FROM tab
+/// ```
+///
+/// Into
+/// ```text
+/// SELECT a, b FROM tab GROUP BY a, b
+/// ```
+
+/// Optimizer that replaces logical [[Distinct]] with a logical [[Aggregate]]
+#[derive(Default)]
+pub struct ReplaceDistinctWithAggregate {}
+
+impl ReplaceDistinctWithAggregate {
+    #[allow(missing_docs)]
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl OptimizerRule for ReplaceDistinctWithAggregate {
+    fn try_optimize(
+        &self,
+        plan: &LogicalPlan,
+        _config: &dyn OptimizerConfig,
+    ) -> Result<Option<LogicalPlan>> {
+        match plan {
+            LogicalPlan::Distinct(Distinct { input }) => {
+                let group_expr = expand_wildcard(input.schema(), input)?;
+                let aggregate = LogicalPlan::Aggregate(Aggregate::try_new_with_schema(
+                    input.clone(),
+                    group_expr,
+                    vec![],
+                    input.schema().clone(), // input schema and aggregate schema are the same in this case
+                )?);
+                Ok(Some(aggregate))
+            }
+            _ => Ok(None),
+        }
+    }
+
+    fn name(&self) -> &str {
+        "replace_distinct_aggregate"
+    }
+
+    fn apply_order(&self) -> Option<ApplyOrder> {
+        Some(BottomUp)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::replace_distinct_aggregate::ReplaceDistinctWithAggregate;
+    use crate::test::{assert_optimized_plan_eq, test_table_scan};
+    use datafusion_expr::{col, LogicalPlanBuilder};
+    use std::sync::Arc;
+
+    #[test]
+    fn replace_distinct() -> datafusion_common::Result<()> {
+        let table_scan = test_table_scan().unwrap();
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .project(vec![col("a"), col("b")])?
+            .distinct()?
+            .build()?;
+
+        let expected = "Aggregate: groupBy=[[test.a, test.b]], aggr=[[]]\
+                            \n  Projection: test.a, test.b\
+                            \n    TableScan: test";
+
+        assert_optimized_plan_eq(
+            Arc::new(ReplaceDistinctWithAggregate::new()),
+            &plan,
+            expected,
+        )
+    }
+}

datafusion/optimizer/tests/integration-test.rs

@@ -169,9 +169,9 @@ fn intersect() -> Result<()> {
     let plan = test_sql(sql)?;
     let expected =
         "LeftSemi Join: test.col_int32 = test.col_int32, test.col_utf8 = test.col_utf8\
-    \n  Distinct:\
+    \n  Aggregate: groupBy=[[test.col_int32, test.col_utf8]], aggr=[[]]\
     \n    LeftSemi Join: test.col_int32 = test.col_int32, test.col_utf8 = test.col_utf8\
-    \n      Distinct:\
+    \n      Aggregate: groupBy=[[test.col_int32, test.col_utf8]], aggr=[[]]\
     \n        TableScan: test projection=[col_int32, col_utf8]\
     \n      TableScan: test projection=[col_int32, col_utf8]\
     \n  TableScan: test projection=[col_int32, col_utf8]";

datafusion/substrait/tests/roundtrip_logical_plan.rs

@@ -231,7 +231,8 @@ mod tests {
             \n  Aggregate: groupBy=[[]], aggr=[[COUNT(Int16(1))]]\
             \n    LeftSemi Join: data.a = data2.a\
             \n      Aggregate: groupBy=[[data.a]], aggr=[[]]\
-            \n        TableScan: data projection=[a]\
+            \n        Projection: data.a\
+            \n          TableScan: data projection=[a]\
             \n      Projection: data2.a\
             \n        TableScan: data2 projection=[a]",
         )