fix: preserve aggregate order in multi-distinct COUNT rewrite and add tests

✅ Fix projection after join so output columns match the original aggregate list when COUNT(DISTINCT …) and non-distinct aggs are interleaved (schema-compatible with mixed BI-style queries).

✅ Add internal_err guard for inconsistent aggregate index mapping.

✅ Optimizer tests: three grouped COUNT(DISTINCT), non-distinct between distincts, CAST(distinct) args, no rewrite for GROUPING SETS.

✅ SQL integration: COUNT(*) + two COUNT(DISTINCT); two GROUP BY keys with expected results.

❌ Grouping-set / filtered-distinct cases remain explicitly out of scope for this rule (covered by unchanged-plan tests where applicable).

Made-with: Cursor

Author: yash

datafusion/core/tests/sql/aggregates/multi_distinct_count_rewrite.rs

@@ -56,3 +56,78 @@ async fn multi_count_distinct_matches_expected_with_nulls() -> Result<()> {
     );
     Ok(())
 }
+
+/// `COUNT(*)` + two `COUNT(DISTINCT …)` per group (BI-style); must match non-rewritten semantics.
+#[tokio::test]
+async fn multi_count_distinct_with_count_star_matches_expected() -> Result<()> {
+    let ctx = SessionContext::new();
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("g", DataType::Int32, false),
+        Field::new("b", DataType::Int32, false),
+        Field::new("c", DataType::Int32, false),
+    ]));
+    let batch = RecordBatch::try_new(
+        schema.clone(),
+        vec![
+            Arc::new(Int32Array::from(vec![1, 1, 1])),
+            Arc::new(Int32Array::from(vec![1, 2, 1])),
+            Arc::new(Int32Array::from(vec![10, 20, 30])),
+        ],
+    )?;
+    let provider = MemTable::try_new(schema, vec![vec![batch]])?;
+    ctx.register_table("t", Arc::new(provider))?;
+
+    let sql = "SELECT g, COUNT(*) AS n, COUNT(DISTINCT b) AS db, COUNT(DISTINCT c) AS dc \
+               FROM t GROUP BY g";
+    let batches = ctx.sql(sql).await?.collect().await?;
+    let out = batches_to_sort_string(&batches);
+
+    assert_eq!(
+        out,
+        "+---+---+----+----+\n\
+         | g | n | db | dc |\n\
+         +---+---+----+----+\n\
+         | 1 | 3 | 2  | 3  |\n\
+         +---+---+----+----+"
+    );
+    Ok(())
+}
+
+/// Multiple `GROUP BY` keys: join must align on all keys.
+#[tokio::test]
+async fn multi_count_distinct_two_group_keys_matches_expected() -> Result<()> {
+    let ctx = SessionContext::new();
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("g1", DataType::Int32, false),
+        Field::new("g2", DataType::Int32, false),
+        Field::new("b", DataType::Int32, false),
+        Field::new("c", DataType::Int32, false),
+    ]));
+    let batch = RecordBatch::try_new(
+        schema.clone(),
+        vec![
+            Arc::new(Int32Array::from(vec![1, 1, 1])),
+            Arc::new(Int32Array::from(vec![1, 1, 2])),
+            Arc::new(Int32Array::from(vec![1, 1, 3])),
+            Arc::new(Int32Array::from(vec![1, 2, 3])),
+        ],
+    )?;
+    let provider = MemTable::try_new(schema, vec![vec![batch]])?;
+    ctx.register_table("t", Arc::new(provider))?;
+
+    let sql = "SELECT g1, g2, COUNT(DISTINCT b) AS db, COUNT(DISTINCT c) AS dc \
+               FROM t GROUP BY g1, g2";
+    let batches = ctx.sql(sql).await?.collect().await?;
+    let out = batches_to_sort_string(&batches);
+
+    assert_eq!(
+        out,
+        "+----+----+----+----+\n\
+         | g1 | g2 | db | dc |\n\
+         +----+----+----+----+\n\
+         | 1  | 1  | 1  | 2  |\n\
+         | 1  | 2  | 1  | 1  |\n\
+         +----+----+----+----+"
+    );
+    Ok(())
+}

datafusion/optimizer/src/multi_distinct_count_rewrite.rs

@@ -25,7 +25,7 @@ use crate::optimizer::ApplyOrder;
 use crate::{OptimizerConfig, OptimizerRule};
 
 use datafusion_common::{
-    Column, JoinConstraint, NullEquality, Result, tree_node::Transformed,
+    Column, JoinConstraint, NullEquality, Result, internal_err, tree_node::Transformed,
 };
 use datafusion_expr::builder::project;
 use datafusion_expr::expr::{AggregateFunction, AggregateFunctionParams, ScalarFunction};
@@ -205,10 +205,9 @@ impl OptimizerRule for MultiDistinctCountRewrite {
                     base_aggr_exprs,
                 )?);
                 let base_alias = config.alias_generator().next("mdc_base");
-                Some(Arc::new(LogicalPlan::SubqueryAlias(SubqueryAlias::try_new(
-                    Arc::new(base_plan),
-                    &base_alias,
-                )?)))
+                Some(Arc::new(LogicalPlan::SubqueryAlias(
+                    SubqueryAlias::try_new(Arc::new(base_plan), &base_alias)?,
+                )))
             };
 
         let mut current = base_plan_opt;
@@ -285,9 +284,12 @@ impl OptimizerRule for MultiDistinctCountRewrite {
             };
         }
 
-        let current = current.expect("distinct_list non-empty implies at least one branch");
+        let current =
+            current.expect("distinct_list non-empty implies at least one branch");
         let join_schema = current.schema();
 
+        let base_field_count = group_size + other_list.len();
+
         let mut proj_exprs: Vec<Expr> = vec![];
         for i in 0..group_size {
             let (q, f) = schema.qualified_field(i);
@@ -296,23 +298,36 @@ impl OptimizerRule for MultiDistinctCountRewrite {
             let c = Expr::Column(Column::new(join_q.cloned(), join_f.name()));
             proj_exprs.push(c.alias_qualified(q.cloned(), orig_name));
         }
-        for (field_idx, (_, schema_aggr_idx)) in other_list.iter().enumerate() {
-            let (q, f) = schema.qualified_field(*schema_aggr_idx);
+        // Preserve original aggregate column order (distinct and non-distinct may be interleaved).
+        for aggr_i in 0..aggr_expr.len() {
+            let schema_idx = group_size + aggr_i;
+            let (q, f) = schema.qualified_field(schema_idx);
             let orig_name = f.name();
-            let join_idx = group_size + field_idx;
-            let (join_q, join_f) = join_schema.qualified_field(join_idx);
-            let c = Expr::Column(Column::new(join_q.cloned(), join_f.name()));
-            proj_exprs.push(c.alias_qualified(q.cloned(), orig_name));
-        }
-        let base_field_count = group_size + other_list.len();
-        for (idx, (_, schema_aggr_idx, _)) in distinct_list.iter().enumerate() {
-            let (q, f) = schema.qualified_field(*schema_aggr_idx);
-            let orig_name = f.name();
-            let branch_start_idx = base_field_count + idx * (group_size + 1);
-            let branch_aggr_idx = branch_start_idx + group_size;
-            let (join_q, join_f) = join_schema.qualified_field(branch_aggr_idx);
-            let c = Expr::Column(Column::new(join_q.cloned(), join_f.name()));
-            proj_exprs.push(c.alias_qualified(q.cloned(), orig_name));
+
+            if let Some((dist_idx, (_, _, _))) = distinct_list
+                .iter()
+                .enumerate()
+                .find(|(_, (_, idx, _))| *idx == schema_idx)
+            {
+                let branch_start_idx = base_field_count + dist_idx * (group_size + 1);
+                let branch_aggr_idx = branch_start_idx + group_size;
+                let (join_q, join_f) = join_schema.qualified_field(branch_aggr_idx);
+                let c = Expr::Column(Column::new(join_q.cloned(), join_f.name()));
+                proj_exprs.push(c.alias_qualified(q.cloned(), orig_name));
+            } else if let Some((other_idx, _)) = other_list
+                .iter()
+                .enumerate()
+                .find(|(_, (_, idx))| *idx == schema_idx)
+            {
+                let join_idx = group_size + other_idx;
+                let (join_q, join_f) = join_schema.qualified_field(join_idx);
+                let c = Expr::Column(Column::new(join_q.cloned(), join_f.name()));
+                proj_exprs.push(c.alias_qualified(q.cloned(), orig_name));
+            } else {
+                return internal_err!(
+                    "aggregate index {aggr_i} (schema index {schema_idx}) is neither distinct nor other"
+                );
+            }
         }
 
         let out = project((*current).clone(), proj_exprs)?;
@@ -327,8 +342,13 @@ mod tests {
     use crate::OptimizerContext;
     use crate::OptimizerRule;
     use crate::test::*;
+    use arrow::datatypes::DataType;
+    use datafusion_expr::GroupingSet;
     use datafusion_expr::LogicalPlan;
+    use datafusion_expr::expr_fn::cast;
+    use datafusion_expr::logical_plan::Aggregate;
     use datafusion_expr::logical_plan::builder::LogicalPlanBuilder;
+    use datafusion_expr::{Expr, col};
     use datafusion_functions_aggregate::expr_fn::{count, count_distinct};
 
     fn optimize_with_rule(
@@ -403,6 +423,53 @@ mod tests {
         Ok(())
     }
 
+    /// Grouped query with multiple `COUNT(DISTINCT …)` **and** non-distinct aggregates (typical BI).
+    /// Non-distinct aggs live in `mdc_base`; each distinct column gets a branch + join on keys.
+    #[test]
+    fn rewrites_two_count_distinct_with_non_distinct_count() -> Result<()> {
+        let table_scan = test_table_scan()?;
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .aggregate(
+                vec![col("a")],
+                vec![
+                    count_distinct(col("b")),
+                    count_distinct(col("c")),
+                    count(col("a")),
+                ],
+            )?
+            .build()?;
+
+        let optimized =
+            optimize_with_rule(plan, Arc::new(MultiDistinctCountRewrite::new()))?;
+        let s = optimized.display_indent_schema().to_string();
+        assert!(s.contains("Inner Join"), "expected join rewrite, got:\n{s}");
+        assert!(
+            s.contains("SubqueryAlias: mdc_base"),
+            "expected base aggregate for non-distinct aggs, got:\n{s}"
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn does_not_rewrite_two_count_distinct_same_column() -> Result<()> {
+        let table_scan = test_table_scan()?;
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .aggregate(
+                vec![col("a")],
+                vec![
+                    count_distinct(col("b")).alias("cd1"),
+                    count_distinct(col("b")).alias("cd2"),
+                ],
+            )?
+            .build()?;
+        let before = plan.display_indent_schema().to_string();
+        let optimized =
+            optimize_with_rule(plan, Arc::new(MultiDistinctCountRewrite::new()))?;
+        let after = optimized.display_indent_schema().to_string();
+        assert_eq!(before, after);
+        Ok(())
+    }
+
     #[test]
     fn does_not_rewrite_single_count_distinct() -> Result<()> {
         let table_scan = test_table_scan()?;
@@ -417,6 +484,107 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn rewrites_three_count_distinct_grouped() -> Result<()> {
+        let table_scan = test_table_scan()?;
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .aggregate(
+                vec![col("a")],
+                vec![
+                    count_distinct(col("b")),
+                    count_distinct(col("c")),
+                    count_distinct(col("a")),
+                ],
+            )?
+            .build()?;
+
+        let optimized =
+            optimize_with_rule(plan, Arc::new(MultiDistinctCountRewrite::new()))?;
+        let s = optimized.display_indent_schema().to_string();
+        assert!(
+            s.matches("Inner Join").count() >= 2,
+            "expected two joins for three branches, got:\n{s}"
+        );
+        assert!(
+            s.contains("SubqueryAlias: mdc_base"),
+            "expected base aggregate, got:\n{s}"
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn rewrites_interleaved_non_distinct_between_distincts() -> Result<()> {
+        let table_scan = test_table_scan()?;
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .aggregate(
+                vec![col("a")],
+                vec![
+                    count_distinct(col("b")),
+                    count(col("a")),
+                    count_distinct(col("c")),
+                ],
+            )?
+            .build()?;
+
+        let optimized =
+            optimize_with_rule(plan, Arc::new(MultiDistinctCountRewrite::new()))?;
+        let s = optimized.display_indent_schema().to_string();
+        assert!(s.contains("Inner Join"), "expected join rewrite, got:\n{s}");
+        assert!(
+            s.contains("SubqueryAlias: mdc_base"),
+            "expected base for middle count(a), got:\n{s}"
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn rewrites_count_distinct_on_cast_exprs() -> Result<()> {
+        let table_scan = test_table_scan()?;
+        let plan = LogicalPlanBuilder::from(table_scan)
+            .aggregate(
+                vec![col("a")],
+                vec![
+                    count_distinct(cast(col("b"), DataType::Int64)),
+                    count_distinct(cast(col("c"), DataType::Int64)),
+                ],
+            )?
+            .build()?;
+
+        let optimized =
+            optimize_with_rule(plan, Arc::new(MultiDistinctCountRewrite::new()))?;
+        let s = optimized.display_indent_schema().to_string();
+        assert!(s.contains("Inner Join"), "expected join rewrite, got:\n{s}");
+        assert!(
+            s.contains("Filter: CAST(test.b AS Int64) IS NOT NULL"),
+            "expected null filter on cast(b), got:\n{s}"
+        );
+        assert!(
+            s.contains("Filter: CAST(test.c AS Int64) IS NOT NULL"),
+            "expected null filter on cast(c), got:\n{s}"
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn does_not_rewrite_grouping_sets_multi_distinct() -> Result<()> {
+        let table_scan = test_table_scan()?;
+        let group_expr = vec![Expr::GroupingSet(GroupingSet::GroupingSets(vec![vec![
+            col("a"),
+        ]]))];
+        let aggr_expr = vec![count_distinct(col("b")), count_distinct(col("c"))];
+        let plan = LogicalPlan::Aggregate(Aggregate::try_new(
+            Arc::new(table_scan),
+            group_expr,
+            aggr_expr,
+        )?);
+        let before = plan.display_indent_schema().to_string();
+        let optimized =
+            optimize_with_rule(plan, Arc::new(MultiDistinctCountRewrite::new()))?;
+        let after = optimized.display_indent_schema().to_string();
+        assert_eq!(before, after);
+        Ok(())
+    }
+
     #[test]
     fn does_not_rewrite_mixed_agg() -> Result<()> {
         let table_scan = test_table_scan()?;