Add support for `COUNT(DISTINCT expr, expr1, ...)`

[andygrove]

What is the problem the feature request solves?

The expression COUNT(DISTINCT expr) is relatively common and it is used in TPC-H, so it would be good to be able to accelerate this in Comet.

Spark supports multiple expressions e.g. COUNT(DISTINCT a, b, c), but DataFusion does not, so we should only attempt to accelerate this if there is a single input expression.

Implementing this feature is not trivial because there are some design issues with how we currently support partial aggregates. Specifically, we do not report the correct output schema from the partial aggregate. For the aggregate expressions that we currently support it doesn't matter because the output of the partial and final aggregates is the same. For example SUM(int_column) will have the output type int for both partial and final. For COUNT(DISTINCT int_column) the output of the partial will be a list of int and the output of the final will be a long.

Describe the potential solution

No response

Additional context

No response

[comphead]

Test to reproduce

test("count distinct") {
    sql("drop table if exists t")
    sql("CREATE TABLE t(v VARCHAR(10), v1 VARCHAR(10), i INT) USING PARQUET")
    sql("INSERT INTO t VALUES ('c', 'a', 1)")
    sql("INSERT INTO t VALUES ('c1', 'a1', 1)")
    sql("INSERT INTO t VALUES ('c2', 'a2', 2)")
    sql("INSERT INTO t VALUES ('c3', 'a3', 2)")
    sql("INSERT INTO t VALUES ('c4', 'a4', 2)")
    sql("INSERT INTO t VALUES ('c', 'a', 1)")
    sql("INSERT INTO t VALUES ('c1', 'a1', 1)")
    sql("INSERT INTO t VALUES ('c2', 'a2', 2)")
    sql("INSERT INTO t VALUES ('c3', 'a3', 2)")
    sql("INSERT INTO t VALUES ('c4', 'a4', 2)")
    sql("INSERT INTO t VALUES ('c', 'a', 1)")
    sql("INSERT INTO t VALUES ('c1', 'a1', 1)")
    sql("INSERT INTO t VALUES ('c2', 'a2', 2)")
    sql("INSERT INTO t VALUES ('c3', 'a3', 2)")
    sql("INSERT INTO t VALUES ('c4', 'a4', 2)")
    sql("select * from t").repartition(3).write.mode("overwrite").parquet("/tmp/test_11")
    // sql("SELECT i, count(distinct v), count(distinct v1) FROM t GROUP BY i").show(false)
    withSQLConf(
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_BATCH_SIZE.key -> "1",
      "spark.comet.cast.allowIncompatible" -> "true",
      "spark.comet.enabled" -> "true",
      "spark.comet.exec.replaceSortMergeJoin" -> "true",
      "spark.comet.exec.shuffle.enableFastEncoding" -> "true",
      "spark.comet.exec.shuffle.enabled" -> "true",
      "spark.comet.exec.shuffle.fallbackToColumnar" -> "true",
      "spark.comet.explain.native.enabled" -> "true",
      "spark.comet.explainFallback.enabled" -> "true",
      "spark.shuffle.manager" -> "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager",
      "spark.comet.logFallbackReasons.enabled" -> "true") {
      spark.read.parquet("/tmp/test_11").createOrReplaceTempView("t2")
      checkSparkAnswerAndOperator("SELECT i, v1, count(distinct v) FROM t2 GROUP BY i, v1")
    }
  }