Substrait: handle identical grouping expressions (#16189)

* Substrait: handle identical grouping expressions

* fix

Author: cht42

datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs

@@ -15,8 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::logical_plan::consumer::SubstraitConsumer;
 use crate::logical_plan::consumer::{from_substrait_agg_func, from_substrait_sorts};
+use crate::logical_plan::consumer::{NameTracker, SubstraitConsumer};
 use datafusion::common::{not_impl_err, DFSchemaRef};
 use datafusion::logical_expr::{Expr, GroupingSet, LogicalPlan, LogicalPlanBuilder};
 use substrait::proto::aggregate_function::AggregationInvocation;
@@ -113,6 +113,14 @@ pub async fn from_aggregate_rel(
             };
             aggr_exprs.push(agg_func?.as_ref().clone());
         }
+
+        // Ensure that all expressions have a unique name
+        let mut name_tracker = NameTracker::new();
+        let group_exprs = group_exprs
+            .iter()
+            .map(|e| name_tracker.get_uniquely_named_expr(e.clone()))
+            .collect::<Result<Vec<Expr>, _>>()?;
+
         input.aggregate(group_exprs, aggr_exprs)?.build()
     } else {
         not_impl_err!("Aggregate without an input is not valid")

datafusion/substrait/tests/cases/roundtrip_logical_plan.rs

@@ -854,6 +854,22 @@ async fn aggregate_wo_projection_sorted_consume() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn aggregate_identical_grouping_expressions() -> Result<()> {
+    let proto_plan =
+        read_json("tests/testdata/test_plans/aggregate_identical_grouping_expressions.substrait.json");
+
+    let plan = generate_plan_from_substrait(proto_plan).await?;
+    assert_snapshot!(
+    plan,
+    @r#"
+    Aggregate: groupBy=[[Int32(1) AS grouping_col_1, Int32(1) AS grouping_col_2]], aggr=[[]]
+      TableScan: data projection=[]
+    "#
+            );
+    Ok(())
+}
+
 #[tokio::test]
 async fn simple_intersect_consume() -> Result<()> {
     let proto_plan = read_json("tests/testdata/test_plans/intersect.substrait.json");

datafusion/substrait/tests/testdata/test_plans/aggregate_identical_grouping_expressions.substrait.json

@@ -0,0 +1,53 @@
+{
+  "extensionUris": [],
+  "extensions": [],
+  "relations": [
+    {
+      "root": {
+        "input": {
+          "aggregate": {
+            "input": {
+              "read": {
+                "common": {
+                  "direct": {}
+                },
+                "baseSchema": {
+                  "names": [],
+                  "struct": {
+                    "types": [],
+                    "nullability": "NULLABILITY_NULLABLE"
+                  }
+                },
+                "namedTable": {
+                  "names": ["data"]
+                }
+              }
+            },
+            "groupings": [
+              {
+                "groupingExpressions": [
+                  {
+                    "literal": {
+                      "i32": 1
+                    }
+                  },
+                  {
+                    "literal": {
+                      "i32": 1
+                    }
+                  }
+                ]
+              }
+            ],
+            "measures": []
+          }
+        },
+        "names": ["grouping_col_1", "grouping_col_2"]
+      }
+    }
+  ],
+  "version": {
+    "minorNumber": 54,
+    "producer": "manual"
+  }
+}