perf: default multi COUNT(DISTINCT) logical optimizer rewrite

Add MultiDistinctCountRewrite in datafusion-optimizer and register it in Optimizer::new() after SingleDistinctToGroupBy. Rewrites 2+ simple COUNT(DISTINCT) on different args into a join of two-phase aggregates; filter distinct_arg IS NOT NULL on each branch for correct NULL semantics.

✅ Unit tests in datafusion-optimizer; ✅ SQL integration test (NULLs) in core_integration.

Author: yash

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

@@ -1024,3 +1024,4 @@ pub fn split_fuzz_timestamp_data_into_batches(
 
 pub mod basic;
 pub mod dict_nulls;
+pub mod multi_distinct_count_rewrite;

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

@@ -0,0 +1,58 @@
+// 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.
+
+//! End-to-end SQL tests for the multi-`COUNT(DISTINCT)` logical optimizer rewrite.
+
+use super::*;
+use arrow::array::{Int32Array, StringArray};
+use datafusion::common::test_util::batches_to_sort_string;
+use datafusion_catalog::MemTable;
+
+#[tokio::test]
+async fn multi_count_distinct_matches_expected_with_nulls() -> Result<()> {
+    let ctx = SessionContext::new();
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("g", DataType::Int32, false),
+        Field::new("b", DataType::Utf8, true),
+        Field::new("c", DataType::Utf8, true),
+    ]));
+    let batch = RecordBatch::try_new(
+        schema.clone(),
+        vec![
+            Arc::new(Int32Array::from(vec![1, 1, 1])),
+            Arc::new(StringArray::from(vec![Some("x"), None, Some("x")])),
+            Arc::new(StringArray::from(vec![None, Some("y"), Some("y")])),
+        ],
+    )?;
+    let provider = MemTable::try_new(schema, vec![vec![batch]])?;
+    ctx.register_table("t", Arc::new(provider))?;
+
+    let sql =
+        "SELECT g, COUNT(DISTINCT b) AS cb, COUNT(DISTINCT c) AS cc 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 | cb | cc |\n\
+         +---+----+----+\n\
+         | 1 | 1  | 1  |\n\
+         +---+----+----+"
+    );
+    Ok(())
+}

datafusion/optimizer/src/lib.rs

@@ -59,6 +59,7 @@ pub mod eliminate_outer_join;
 pub mod extract_equijoin_predicate;
 pub mod extract_leaf_expressions;
 pub mod filter_null_join_keys;
+pub mod multi_distinct_count_rewrite;
 pub mod optimize_projections;
 pub mod optimize_unions;
 pub mod optimizer;