/*
 * 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.
 */

package org.apache.comet.exec

import scala.util.Random

import org.apache.hadoop.fs.Path
import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
import org.apache.spark.sql.catalyst.expressions.Cast
import org.apache.spark.sql.catalyst.optimizer.EliminateSorts
import org.apache.spark.sql.comet.CometHashAggregateExec
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.functions.{avg, col, count_distinct, sum}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{DataTypes, StructField, StructType}

import org.apache.comet.CometConf
import org.apache.comet.CometConf.COMET_EXEC_STRICT_FLOATING_POINT
import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, ParquetGenerator, SchemaGenOptions}

/**
 * Test suite dedicated to Comet native aggregate operator
 */
class CometAggregateSuite extends CometTestBase with AdaptiveSparkPlanHelper {
  import testImplicits._

  test("min/max floating point with negative zero") {
    val r = new Random(42)
    val schema = StructType(
      Seq(
        StructField("float_col", DataTypes.FloatType, nullable = true),
        StructField("double_col", DataTypes.DoubleType, nullable = true)))
    val df = FuzzDataGenerator.generateDataFrame(
      r,
      spark,
      schema,
      1000,
      DataGenOptions(generateNegativeZero = true))
    df.createOrReplaceTempView("tbl")

    for (col <- Seq("float_col", "double_col")) {
      // assert that data contains positive and negative zero
      assert(spark.sql(s"select * from tbl where cast($col as string) = '0.0'").count() > 0)
      assert(spark.sql(s"select * from tbl where cast($col as string) = '-0.0'").count() > 0)
      for (agg <- Seq("min", "max")) {
        withSQLConf(COMET_EXEC_STRICT_FLOATING_POINT.key -> "true") {
          checkSparkAnswerAndFallbackReasons(
            s"select $agg($col) from tbl where cast($col as string) in ('0.0', '-0.0')",
            Set(
              "Unsupported aggregate expression(s)",
              s"floating-point not supported when ${COMET_EXEC_STRICT_FLOATING_POINT.key}=true"))
        }
        checkSparkAnswer(
          s"select $col, count(*) from tbl " +
            s"where cast($col as string) in ('0.0', '-0.0') group by $col")
      }
    }
  }

  test("avg decimal") {
    withTempDir { dir =>
      val path = new Path(dir.toURI.toString, "test.parquet")
      val filename = path.toString
      val random = new Random(42)
      withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
        ParquetGenerator.makeParquetFile(
          random,
          spark,
          filename,
          10000,
          SchemaGenOptions(),
          DataGenOptions())
      }
      val tableName = "avg_decimal"
      withTable(tableName) {
        val table = spark.read.parquet(filename).coalesce(1)
        table.createOrReplaceTempView(tableName)
        // we fall back to Spark for avg on decimal due to the following issue
        // https://github.com/apache/datafusion-comet/issues/1371
        // once this is fixed, we should change this test to
        // checkSparkAnswerAndNumOfAggregates
        checkSparkAnswer(s"SELECT c1, avg(c7) FROM $tableName GROUP BY c1 ORDER BY c1")
      }
    }
  }

  test("stddev_pop should return NaN for some cases") {
    withSQLConf(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true") {
      Seq(true, false).foreach { nullOnDivideByZero =>
        withSQLConf("spark.sql.legacy.statisticalAggregate" -> nullOnDivideByZero.toString) {

          val data: Seq[(Float, Int)] = Seq((Float.PositiveInfinity, 1))
          withParquetTable(data, "tbl", false) {
            val df = sql("SELECT stddev_pop(_1), stddev_pop(_2) FROM tbl")
            checkSparkAnswerAndOperator(df)
          }
        }
      }
    }
  }

  test("count with aggregation filter") {
    withSQLConf(
      CometConf.COMET_ENABLED.key -> "true",
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
      val df1 = sql("SELECT count(DISTINCT 2), count(DISTINCT 2,3)")
      checkSparkAnswer(df1)

      val df2 = sql("SELECT count(DISTINCT 2), count(DISTINCT 3,2)")
      checkSparkAnswer(df2)
    }
  }

  test("multiple column distinct count") {
    withSQLConf(
      CometConf.COMET_ENABLED.key -> "true",
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
      val df1 = Seq(
        ("a", "b", "c"),
        ("a", "b", "c"),
        ("a", "b", "d"),
        ("x", "y", "z"),
        ("x", "q", null.asInstanceOf[String]))
        .toDF("key1", "key2", "key3")

      checkSparkAnswer(df1.agg(count_distinct($"key1", $"key2")))
      checkSparkAnswer(df1.agg(count_distinct($"key1", $"key2", $"key3")))
      checkSparkAnswer(df1.groupBy($"key1").agg(count_distinct($"key2", $"key3")))
    }
  }

  test("Only trigger Comet Final aggregation on Comet partial aggregation") {
    withTempView("lowerCaseData") {
      lowerCaseData.createOrReplaceTempView("lowerCaseData")
      withSQLConf(
        CometConf.COMET_ENABLED.key -> "true",
        CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
        CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
        val df = sql("SELECT LAST(n) FROM lowerCaseData")
        checkSparkAnswer(df)
      }
    }
  }

  test(
    "Average expression in Comet Final should handle " +
      "all null inputs from partial Spark aggregation") {
    withTempView("allNulls") {
      allNulls.createOrReplaceTempView("allNulls")
      withSQLConf(
        CometConf.COMET_ENABLED.key -> "true",
        CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
        CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
        val df = sql("select sum(a), avg(a) from allNulls")
        checkSparkAnswer(df)
      }
    }
  }

  test("Aggregation without aggregate expressions should use correct result expressions") {
    withSQLConf(
      CometConf.COMET_ENABLED.key -> "true",
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "test")
        makeParquetFile(path, 10000, 10, false)
        withParquetTable(path.toUri.toString, "tbl") {
          val df = sql("SELECT _g5 FROM tbl GROUP BY _g1, _g2, _g3, _g4, _g5")
          checkSparkAnswerAndOperator(df)
        }
      }
    }
  }

  test("Final aggregation should not bind to the input of partial aggregation") {
    withSQLConf(
      CometConf.COMET_ENABLED.key -> "true",
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
      Seq(true, false).foreach { dictionaryEnabled =>
        withTempDir { dir =>
          val path = new Path(dir.toURI.toString, "test")
          makeParquetFile(path, 10000, 10, dictionaryEnabled)
          withParquetTable(path.toUri.toString, "tbl") {
            val df = sql("SELECT * FROM tbl").groupBy("_g1").agg(sum($"_3" + $"_g3"))
            checkSparkAnswerAndOperator(df)
          }
        }
      }
    }
  }

  test("Ensure traversed operators during finding first partial aggregation are all native") {
    withTable("lineitem", "part") {
      withSQLConf(
        CometConf.COMET_ENABLED.key -> "true",
        CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
        CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {

        sql(
          "CREATE TABLE lineitem(l_extendedprice DOUBLE, l_quantity DOUBLE, l_partkey STRING) USING PARQUET")
        sql("INSERT INTO TABLE lineitem VALUES (1.0, 1.0, '1')")

        sql(
          "CREATE TABLE part(p_partkey STRING, p_brand STRING, p_container STRING) USING PARQUET")
        sql("INSERT INTO TABLE part VALUES ('1', 'Brand#23', 'MED BOX')")

        val df = sql("""select
            sum(l_extendedprice) / 7.0 as avg_yearly
            from
            lineitem,
            part
              where
              p_partkey = l_partkey
              and p_brand = 'Brand#23'
          and p_container = 'MED BOX'
          and l_quantity < (
            select
          0.2 * avg(l_quantity)
          from
          lineitem
          where
          l_partkey = p_partkey
          )""")
        checkAnswer(df, Row(null))
      }
    }
  }

  test("SUM decimal supports emit.first") {
    withSQLConf(
      SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> EliminateSorts.ruleName,
      CometConf.COMET_ENABLED.key -> "true",
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
      Seq(true, false).foreach { dictionaryEnabled =>
        withTempDir { dir =>
          val path = new Path(dir.toURI.toString, "test")
          makeParquetFile(path, 10000, 10, dictionaryEnabled)
          withParquetTable(path.toUri.toString, "tbl") {
            checkSparkAnswerAndOperator(
              sql("SELECT * FROM tbl").sort("_g1").groupBy("_g1").agg(sum("_8")))
          }
        }
      }
    }
  }

  test("AVG decimal supports emit.first") {
    withSQLConf(
      SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> EliminateSorts.ruleName,
      CometConf.COMET_ENABLED.key -> "true",
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
      Seq(true, false).foreach { dictionaryEnabled =>
        withTempDir { dir =>
          val path = new Path(dir.toURI.toString, "test")
          makeParquetFile(path, 10000, 10, dictionaryEnabled)
          withParquetTable(path.toUri.toString, "tbl") {
            checkSparkAnswerAndOperator(
              sql("SELECT * FROM tbl").sort("_g1").groupBy("_g1").agg(avg("_8")))
          }
        }
      }
    }
  }

  test("Fix NPE in partial decimal sum") {
    val table = "tbl"
    withTable(table) {
      withSQLConf(
        CometConf.COMET_ENABLED.key -> "true",
        CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "false",
        CometConf.COMET_SHUFFLE_MODE.key -> "native") {
        withTable(table) {
          sql(s"CREATE TABLE $table(col DECIMAL(5, 2)) USING PARQUET")
          sql(s"INSERT INTO TABLE $table VALUES (CAST(12345.01 AS DECIMAL(5, 2)))")
          val df = sql(s"SELECT SUM(col + 100000.01) FROM $table")
          checkAnswer(df, Row(null))
        }
      }
    }
  }

  test("fix: Decimal Average should not enable native final aggregation") {
    withSQLConf(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true") {
      Seq(true, false).foreach { dictionaryEnabled =>
        withTempDir { dir =>
          val path = new Path(dir.toURI.toString, "test")
          makeParquetFile(path, 1000, 10, dictionaryEnabled)
          withParquetTable(path.toUri.toString, "tbl") {
            checkSparkAnswer("SELECT _g1, AVG(_7) FROM tbl GROUP BY _g1")
            checkSparkAnswer("SELECT _g1, AVG(_8) FROM tbl GROUP BY _g1")
            checkSparkAnswer("SELECT _g1, AVG(_9) FROM tbl GROUP BY _g1")
          }
        }
      }
    }
  }

  test("trivial case") {
    Seq(true, false).foreach { dictionaryEnabled =>
      withParquetTable((0 until 5).map(i => (i, i)), "tbl", dictionaryEnabled) {
        val df1 = sql("SELECT _2, SUM(_1) FROM tbl GROUP BY _2")
        checkAnswer(df1, Row(0, 0) :: Row(1, 1) :: Row(2, 2) :: Row(3, 3) :: Row(4, 4) :: Nil)

        val df2 = sql("SELECT _2, COUNT(_1) FROM tbl GROUP BY _2")
        checkAnswer(df2, Row(0, 1) :: Row(1, 1) :: Row(2, 1) :: Row(3, 1) :: Row(4, 1) :: Nil)

        val df3 = sql("SELECT COUNT(_1), COUNT(_2) FROM tbl")
        checkAnswer(df3, Row(5, 5) :: Nil)

        checkSparkAnswerAndOperator("SELECT _2, MIN(_1), MAX(_1) FROM tbl GROUP BY _2")
      }
    }
  }

  test("avg") {
    Seq(true, false).foreach { dictionaryEnabled =>
      withParquetTable(
        (0 until 10).map(i => ((i + 1) * (i + 1), (i + 1) / 2)),
        "tbl",
        dictionaryEnabled) {

        checkSparkAnswerAndOperator("SELECT _2, AVG(_1) FROM tbl GROUP BY _2")
        checkSparkAnswerAndOperator("SELECT AVG(_2) FROM tbl")
      }
    }
  }

  test("count, avg with null") {
    Seq(false, true).foreach { dictionary =>
      withSQLConf("parquet.enable.dictionary" -> dictionary.toString) {
        val table = "test"
        withTable(table) {
          sql(s"create table $table(col1 int, col2 int) using parquet")
          sql(s"insert into $table values(1, 1), (2, 1), (3, 2), (null, 2), (null, 1)")
          checkSparkAnswerAndOperator(s"SELECT COUNT(col1) FROM $table")
          checkSparkAnswerAndOperator(s"SELECT col2, COUNT(col1) FROM $table GROUP BY col2")
          checkSparkAnswerAndOperator(s"SELECT avg(col1) FROM $table")
          checkSparkAnswerAndOperator(s"SELECT col2, avg(col1) FROM $table GROUP BY col2")
        }
      }
    }
  }

  test("SUM/AVG non-decimal overflow") {
    Seq(true, false).foreach { dictionaryEnabled =>
      withParquetTable(Seq((0, 100.toLong), (0, Long.MaxValue)), "tbl", dictionaryEnabled) {
        checkSparkAnswerAndOperator("SELECT SUM(_2) FROM tbl GROUP BY _1")
        checkSparkAnswerAndOperator("SELECT AVG(_2) FROM tbl GROUP BY _1")
      }
    }
  }

  test("simple SUM, COUNT, MIN, MAX, AVG with non-distinct group keys") {
    Seq(true, false).foreach { dictionaryEnabled =>
      withParquetTable((0 until 5).map(i => (i, i % 2)), "tbl", dictionaryEnabled) {
        val df1 = sql("SELECT _2, SUM(_1) FROM tbl GROUP BY _2")
        checkAnswer(df1, Row(0, 6) :: Row(1, 4) :: Nil)
        val df2 = sql("SELECT _2, COUNT(_1) FROM tbl GROUP BY _2")
        checkAnswer(df2, Row(0, 3) :: Row(1, 2) :: Nil)
        checkSparkAnswerAndOperator("SELECT _2, MIN(_1), MAX(_1) FROM tbl GROUP BY _2")
        checkSparkAnswerAndOperator("SELECT _2, AVG(_1) FROM tbl GROUP BY _2")
      }
    }
  }

  test("group-by on variable length types") {
    Seq(true, false).foreach { nativeShuffleEnabled =>
      Seq(true, false).foreach { dictionaryEnabled =>
        withSQLConf(
          CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> nativeShuffleEnabled.toString,
          CometConf.COMET_SHUFFLE_MODE.key -> "native") {
          withParquetTable(
            (0 until 100).map(i => (i, (i % 10).toString)),
            "tbl",
            dictionaryEnabled) {
            val n = if (nativeShuffleEnabled) 2 else 1
            checkSparkAnswerAndNumOfAggregates("SELECT _2, SUM(_1) FROM tbl GROUP BY _2", n)
            checkSparkAnswerAndNumOfAggregates("SELECT _2, COUNT(_1) FROM tbl GROUP BY _2", n)
            checkSparkAnswerAndNumOfAggregates("SELECT _2, MIN(_1) FROM tbl GROUP BY _2", n)
            checkSparkAnswerAndNumOfAggregates("SELECT _2, MAX(_1) FROM tbl GROUP BY _2", n)
            checkSparkAnswerAndNumOfAggregates("SELECT _2, AVG(_1) FROM tbl GROUP BY _2", n)
          }
        }
      }
    }
  }

  test("simple SUM, COUNT, MIN, MAX, AVG with non-distinct + null group keys") {
    Seq(true, false).foreach { dictionaryEnabled =>
      withParquetTable(
        (0 until 10).map { i =>
          (i, if (i % 3 == 0) null.asInstanceOf[Int] else i % 3)
        },
        "tbl",
        dictionaryEnabled) {
        val df1 = sql("SELECT _2, SUM(_1) FROM tbl GROUP BY _2")
        checkAnswer(df1, Row(null.asInstanceOf[Int], 18) :: Row(1, 12) :: Row(2, 15) :: Nil)

        val df2 = sql("SELECT _2, COUNT(_1) FROM tbl GROUP BY _2")
        checkAnswer(df2, Row(null.asInstanceOf[Int], 4) :: Row(1, 3) :: Row(2, 3) :: Nil)

        val df3 = sql("SELECT _2, MIN(_1), MAX(_1) FROM tbl GROUP BY _2")
        checkAnswer(df3, Row(null.asInstanceOf[Int], 0, 9) :: Row(1, 1, 7) :: Row(2, 2, 8) :: Nil)
        checkSparkAnswerAndOperator(sql("SELECT _2, AVG(_1) FROM tbl GROUP BY _2"))
      }
    }
  }

  test("simple SUM, COUNT, MIN, MAX, AVG with null aggregates") {
    Seq(true, false).foreach { dictionaryEnabled =>
      withParquetTable(
        (0 until 10).map { i =>
          (
            if (i % 2 == 0) null.asInstanceOf[Int] else i,
            if (i % 3 == 0) null.asInstanceOf[Int] else i % 3)
        },
        "tbl",
        dictionaryEnabled) {
        val df1 = sql("SELECT _2, SUM(_1) FROM tbl GROUP BY _2")
        checkAnswer(df1, Row(null.asInstanceOf[Int], 12) :: Row(1, 8) :: Row(2, 5) :: Nil)

        val df2 = sql("SELECT _2, COUNT(_1) FROM tbl GROUP BY _2")
        checkAnswer(df2, Row(null.asInstanceOf[Int], 4) :: Row(1, 3) :: Row(2, 3) :: Nil)

        val df3 = sql("SELECT _2, MIN(_1), MAX(_1) FROM tbl GROUP BY _2")
        checkAnswer(df3, Row(null.asInstanceOf[Int], 0, 9) :: Row(1, 0, 7) :: Row(2, 0, 5) :: Nil)

        checkSparkAnswerAndOperator(sql("SELECT _2, AVG(_1) FROM tbl GROUP BY _2"))
      }
    }
  }

  test("simple SUM, MIN, MAX, AVG with all nulls") {
    Seq(true, false).foreach { dictionaryEnabled =>
      withParquetTable(
        (0 until 10).map { i =>
          (null.asInstanceOf[Int], if (i % 3 == 0) null.asInstanceOf[Int] else i % 3)
        },
        "tbl",
        dictionaryEnabled) {
        val df = sql("SELECT _2, SUM(_1) FROM tbl GROUP BY _2")
        checkAnswer(
          df,
          Seq(
            Row(null.asInstanceOf[Int], null.asInstanceOf[Int]),
            Row(1, null.asInstanceOf[Int]),
            Row(2, null.asInstanceOf[Int])))

        val df2 = sql("SELECT _2, MIN(_1), MAX(_1) FROM tbl GROUP BY _2")
        checkAnswer(
          df2,
          Seq(
            Row(null.asInstanceOf[Int], null.asInstanceOf[Int], null.asInstanceOf[Int]),
            Row(1, null.asInstanceOf[Int], null.asInstanceOf[Int]),
            Row(2, null.asInstanceOf[Int], null.asInstanceOf[Int])))
        checkSparkAnswerAndOperator(sql("SELECT _2, SUM(_1) FROM tbl GROUP BY _2"))
      }
    }
  }

  test("SUM, COUNT, MIN, MAX, AVG on float & double") {
    Seq(true, false).foreach { dictionaryEnabled =>
      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "test")
        makeParquetFile(path, 1000, 10, dictionaryEnabled)
        withParquetTable(path.toUri.toString, "tbl") {
          checkSparkAnswerAndOperator(
            "SELECT _g5, SUM(_5), COUNT(_5), MIN(_5), MAX(_5), AVG(_5) FROM tbl GROUP BY _g5")
          checkSparkAnswerAndOperator(
            "SELECT _g6, SUM(_6), COUNT(_6), MIN(_6), MAX(_6), AVG(_6) FROM tbl GROUP BY _g6")
        }
      }
    }
  }

  test("SUM, MIN, MAX, AVG for NaN, -0.0 and 0.0") {
    // NaN should be grouped together, and -0.0 and 0.0 should also be grouped together
    Seq(true, false).foreach { dictionaryEnabled =>
      val data: Seq[(Float, Int)] = Seq(
        (Float.NaN, 1),
        (-0.0.asInstanceOf[Float], 2),
        (0.0.asInstanceOf[Float], 3),
        (Float.NaN, 4))
      withSQLConf(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "false") {
        withParquetTable(data, "tbl", dictionaryEnabled) {
          checkSparkAnswer("SELECT SUM(_2), MIN(_2), MAX(_2), _1 FROM tbl GROUP BY _1")
          checkSparkAnswer("SELECT MIN(_1), MAX(_1), MIN(_2), MAX(_2) FROM tbl")
          checkSparkAnswer("SELECT AVG(_2), _1 FROM tbl GROUP BY _1")
          checkSparkAnswer("SELECT AVG(_1), AVG(_2) FROM tbl")
        }
      }
    }
  }

  test("SUM/MIN/MAX/AVG on decimal") {
    Seq(true, false).foreach { dictionaryEnabled =>
      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "test")
        makeParquetFile(path, 1000, 10, dictionaryEnabled)
        withParquetTable(path.toUri.toString, "tbl") {
          checkSparkAnswer("SELECT _g1, SUM(_7), MIN(_7), MAX(_7), AVG(_7) FROM tbl GROUP BY _g1")
          checkSparkAnswer("SELECT _g1, SUM(_8), MIN(_8), MAX(_8), AVG(_8) FROM tbl GROUP BY _g1")
          checkSparkAnswer("SELECT _g1, SUM(_9), MIN(_9), MAX(_9), AVG(_9) FROM tbl GROUP BY _g1")
        }
      }
    }
  }

  test("multiple SUM/MIN/MAX/AVG on decimal and non-decimal") {
    Seq(true, false).foreach { dictionaryEnabled =>
      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "test")
        makeParquetFile(path, 1000, 10, dictionaryEnabled)
        withParquetTable(path.toUri.toString, "tbl") {
          checkSparkAnswer(
            "SELECT _g1, COUNT(_6), COUNT(_7), SUM(_6), SUM(_7), MIN(_6), MIN(_7), MAX(_6), MAX(_7), AVG(_6), AVG(_7) FROM tbl GROUP BY _g1")
          checkSparkAnswer(
            "SELECT _g1, COUNT(_7), COUNT(_8), SUM(_7), SUM(_8), MIN(_7), MIN(_8), MAX(_7), MAX(_8), AVG(_7), AVG(_8) FROM tbl GROUP BY _g1")
          checkSparkAnswer(
            "SELECT _g1, COUNT(_8), COUNT(_9), SUM(_8), SUM(_9), MIN(_8), MIN(_9), MAX(_8), MAX(_9), AVG(_8), AVG(_9) FROM tbl GROUP BY _g1")
          checkSparkAnswer(
            "SELECT _g1, COUNT(_9), COUNT(_1), SUM(_9), SUM(_1), MIN(_9), MIN(_1), MAX(_9), MAX(_1), AVG(_9), AVG(_1) FROM tbl GROUP BY _g1")
        }
      }
    }
  }

  test("SUM/AVG on decimal with different precisions") {
    Seq(true, false).foreach { dictionaryEnabled =>
      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "test")
        makeParquetFile(path, 1000, 10, dictionaryEnabled)
        withParquetTable(path.toUri.toString, "tbl") {
          Seq("SUM", "AVG").foreach { FN =>
            checkSparkAnswerAndOperator(
              s"SELECT _g1, $FN(_8 + CAST(1 AS DECIMAL(20, 10))) FROM tbl GROUP BY _g1")
            checkSparkAnswerAndOperator(
              s"SELECT _g1, $FN(_8 - CAST(-1 AS DECIMAL(10, 3))) FROM tbl GROUP BY _g1")
            checkSparkAnswerAndOperator(
              s"SELECT _g1, $FN(_9 * CAST(3.14 AS DECIMAL(4, 3))) FROM tbl GROUP BY _g1")
            checkSparkAnswerAndOperator(
              s"SELECT _g1, $FN(_9 / CAST(1.2345 AS DECIMAL(35, 10))) FROM tbl GROUP BY _g1")
          }
        }
      }
    }
  }

  test("SUM decimal with DF") {
    Seq(true, false).foreach { dictionaryEnabled =>
      Seq(true, false).foreach { nativeShuffleEnabled =>
        withSQLConf(
          CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> nativeShuffleEnabled.toString,
          CometConf.COMET_SHUFFLE_MODE.key -> "native") {
          withTempDir { dir =>
            val path = new Path(dir.toURI.toString, "test")
            makeParquetFile(path, 1000, 20, dictionaryEnabled)
            withParquetTable(path.toUri.toString, "tbl") {
              val expectedNumOfCometAggregates = if (nativeShuffleEnabled) 2 else 1

              checkSparkAnswerAndNumOfAggregates(
                "SELECT _g2, SUM(_7) FROM tbl GROUP BY _g2",
                expectedNumOfCometAggregates)
              checkSparkAnswerAndNumOfAggregates(
                "SELECT _g3, SUM(_8) FROM tbl GROUP BY _g3",
                expectedNumOfCometAggregates)
              checkSparkAnswerAndNumOfAggregates(
                "SELECT _g4, SUM(_9) FROM tbl GROUP BY _g4",
                expectedNumOfCometAggregates)
              checkSparkAnswerAndNumOfAggregates(
                "SELECT SUM(_7) FROM tbl",
                expectedNumOfCometAggregates)
              checkSparkAnswerAndNumOfAggregates(
                "SELECT SUM(_8) FROM tbl",
                expectedNumOfCometAggregates)
              checkSparkAnswerAndNumOfAggregates(
                "SELECT SUM(_9) FROM tbl",
                expectedNumOfCometAggregates)
            }
          }
        }
      }
    }
  }

  test("COUNT/MIN/MAX on date, timestamp") {
    Seq(true, false).foreach { dictionaryEnabled =>
      withTempDir { dir =>
        val path = new Path(dir.toURI.toString, "test")
        makeParquetFile(path, 1000, 10, dictionaryEnabled)
        withParquetTable(path.toUri.toString, "tbl") {
          checkSparkAnswerAndOperator(
            "SELECT _g1, COUNT(_10), MIN(_10), MAX(_10) FROM tbl GROUP BY _g1")
          checkSparkAnswerAndOperator(
            "SELECT _g1, COUNT(_11), MIN(_11), MAX(_11) FROM tbl GROUP BY _g1")
          checkSparkAnswerAndOperator(
            "SELECT _g1, COUNT(_12), MIN(_12), MAX(_12) FROM tbl GROUP BY _g1")
        }
      }
    }
  }

  test("single group-by column + aggregate column, multiple batches, no null") {
    val numValues = 10000

    Seq(1, 100, 10000).foreach { numGroups =>
      Seq(128, 1024, numValues + 1).foreach { batchSize =>
        Seq(true, false).foreach { dictionaryEnabled =>
          withSQLConf(
            SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
            CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
            withParquetTable(
              (0 until numValues).map(i => (i, Random.nextInt() % numGroups)),
              "tbl",
              dictionaryEnabled) {
              withView("v") {
                sql("CREATE TEMP VIEW v AS SELECT _1, _2 FROM tbl ORDER BY _1")
                checkSparkAnswerAndOperator(
                  "SELECT _2, SUM(_1), SUM(DISTINCT _1), MIN(_1), MAX(_1), COUNT(_1)," +
                    " COUNT(DISTINCT _1), AVG(_1)" +
                    " FROM v GROUP BY _2 ORDER BY _2")
              }
            }
          }
        }
      }
    }
  }

  // FIRST/LAST are order-dependent aggregates whose merge result depends on hash table
  // processing order. In PartialMerge mode, DataFusion's hash table may process rows
  // in a different order than Spark's, so we fall back to Spark for correctness.
  // https://github.com/apache/datafusion-comet/issues/4131
  test("partialMerge - FIRST/LAST with distinct aggregates falls back") {
    val numValues = 10000
    Seq(100).foreach { numGroups =>
      Seq(128).foreach { batchSize =>
        withSQLConf(
          SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
          CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
          withParquetTable(
            (0 until numValues).map(i => (i, Random.nextInt() % numGroups)),
            "tbl",
            false) {
            withView("v") {
              sql("CREATE TEMP VIEW v AS SELECT _1, _2 FROM tbl ORDER BY _1")
              checkSparkAnswerAndFallbackReason(
                "SELECT _2, FIRST(_1), LAST(_1), COUNT(DISTINCT _1)" +
                  " FROM v GROUP BY _2 ORDER BY _2",
                "PartialMerge not supported for aggregates: first, last")
            }
          }
        }
      }
    }
  }

  test("partialMerge - cnt distinct + sum") {
    withTempDir(dir => {
      withSQLConf("spark.comet.enabled" -> "false") {
        withView("t") {
          sql("""
              CREATE OR REPLACE TEMP VIEW t (v, v1, i) AS
              VALUES
                ('c',  'a',  1),
                ('c1', 'a1', 1),
                ('c2', 'a2', 2),
                ('c3', 'a3', 2),
                ('c4', 'a4', 2),
                ('c',  'a',  1),
                ('c1', 'a1', 1),
                ('c2', 'a2', 2),
                ('c3', 'a3', 2),
                ('c4', 'a4', 2),
                ('c',  'a',  1),
                ('c1', 'a1', 1),
                ('c2', 'a2', 2),
                ('c3', 'a3', 2),
                ('c4', 'a4', 2)
              """)
          sql("select * from t")
            .repartition(3)
            .write
            .mode("overwrite")
            .parquet(dir.getAbsolutePath)
        }
      }

      withSQLConf(
        CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
        "spark.comet.exec.shuffle.fallbackToColumnar" -> "false",
        "spark.comet.enabled" -> "true") {
        spark.read.parquet(dir.getAbsolutePath).createOrReplaceTempView("t2")
        checkSparkAnswerAndOperator("SELECT i, sum(v1), count(distinct v) FROM t2 group by i")
      }
    })
  }

  test("partialMerge - distinct + non-distinct aggregates (Expand pattern)") {
    withParquetTable((1 to 100).map(i => (i, i.toString)), "tbl", false) {
      checkSparkAnswerAndOperator("SELECT avg(_1), sum(_1), count(distinct _1) FROM tbl")

      checkSparkAnswerAndOperator(
        "SELECT max(_1), count(distinct _1), sum(distinct _1), sum(1) FROM tbl")

      // Exercises the Expand plan for a non-COUNT distinct (AVG(DISTINCT)) mixed with a
      // non-distinct aggregate, so the PartialMerge stage has to survive for a float-producing
      // distinct aggregate.
      checkSparkAnswerAndOperator("SELECT avg(distinct _1), sum(_1) FROM tbl")
    }
  }

  test("multiple group-by columns + single aggregate column (first/last), with nulls") {
    val numValues = 10000

    Seq(1, 100, numValues).foreach { numGroups =>
      Seq(128, numValues + 100).foreach { batchSize =>
        Seq(true, false).foreach { dictionaryEnabled =>
          withSQLConf(
            SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
            CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
            withTempPath { dir =>
              val path = new Path(dir.toURI.toString, "test.parquet")
              makeParquetFile(path, numValues, numGroups, dictionaryEnabled)
              withParquetTable(path.toUri.toString, "tbl") {
                withView("v") {
                  sql("CREATE TEMP VIEW v AS SELECT _g1, _g2, _3 FROM tbl ORDER BY _3")
                  checkSparkAnswerAndOperator(
                    "SELECT _g1, _g2, FIRST(_3) FROM v GROUP BY _g1, _g2 ORDER BY _g1, _g2")
                  checkSparkAnswerAndOperator(
                    "SELECT _g1, _g2, LAST(_3) FROM v GROUP BY _g1, _g2 ORDER BY _g1, _g2")
                  checkSparkAnswerAndOperator(
                    "SELECT _g1, _g2, FIRST(_3) IGNORE NULLS FROM v GROUP BY _g1, _g2 ORDER BY _g1, _g2")
                  checkSparkAnswerAndOperator(
                    "SELECT _g1, _g2, LAST(_3) IGNORE NULLS FROM v GROUP BY _g1, _g2 ORDER BY _g1, _g2")
                }
              }
            }
          }
        }
      }
    }
  }

  test("multiple group-by columns + single aggregate column, with nulls") {
    val numValues = 10000

    Seq(1, 100, numValues).foreach { numGroups =>
      Seq(128, numValues + 100).foreach { batchSize =>
        Seq(true, false).foreach { dictionaryEnabled =>
          withSQLConf(
            SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
            CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
            withTempPath { dir =>
              val path = new Path(dir.toURI.toString, "test.parquet")
              makeParquetFile(path, numValues, numGroups, dictionaryEnabled)
              withParquetTable(path.toUri.toString, "tbl") {
                checkSparkAnswer(
                  "SELECT _g1, _g2, SUM(_3) FROM tbl GROUP BY _g1, _g2 ORDER BY _g1, _g2")
                checkSparkAnswer(
                  "SELECT _g1, _g2, COUNT(_3) FROM tbl GROUP BY _g1, _g2 ORDER BY _g1, _g2")
                checkSparkAnswer(
                  "SELECT _g1, _g2, SUM(DISTINCT _3) FROM tbl GROUP BY _g1, _g2 ORDER BY _g1, _g2")
                checkSparkAnswer(
                  "SELECT _g1, _g2, COUNT(DISTINCT _3) FROM tbl GROUP BY _g1, _g2 ORDER BY _g1, _g2")
                checkSparkAnswer(
                  "SELECT _g1, _g2, MIN(_3), MAX(_3) FROM tbl GROUP BY _g1, _g2 ORDER BY _g1, _g2")
                checkSparkAnswer(
                  "SELECT _g1, _g2, AVG(_3) FROM tbl GROUP BY _g1, _g2 ORDER BY _g1, _g2")
              }
            }
          }
        }
      }
    }
  }

  test("string should be supported") {
    withTable("t") {
      sql("CREATE TABLE t(v VARCHAR(3), i INT) USING PARQUET")
      sql("INSERT INTO t VALUES ('c', 1)")
      withSQLConf(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "false") {
        checkSparkAnswerAndNumOfAggregates("SELECT v, sum(i) FROM t GROUP BY v ORDER BY v", 1)
      }
    }
  }

  test("multiple group-by columns + multiple aggregate column (first/last), with nulls") {
    val numValues = 10000

    Seq(1, 100, numValues).foreach { numGroups =>
      Seq(128, numValues + 100).foreach { batchSize =>
        Seq(true, false).foreach { dictionaryEnabled =>
          withSQLConf(
            SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
            CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
            withTempPath { dir =>
              val path = new Path(dir.toURI.toString, "test.parquet")
              makeParquetFile(path, numValues, numGroups, dictionaryEnabled)
              withParquetTable(path.toUri.toString, "tbl") {
                withView("v") {
                  sql("CREATE TEMP VIEW v AS SELECT _g3, _g4, _3, _4 FROM tbl ORDER BY _3, _4")
                  checkSparkAnswerAndOperator(
                    "SELECT _g3, _g4, FIRST(_3), FIRST(_4) FROM v GROUP BY _g3, _g4 ORDER BY _g3, _g4")
                  checkSparkAnswerAndOperator(
                    "SELECT _g3, _g4, LAST(_3), LAST(_4) FROM v GROUP BY _g3, _g4 ORDER BY _g3, _g4")
                }
              }
            }
          }
        }
      }
    }

  }

  test("multiple group-by columns + multiple aggregate column, with nulls") {
    val numValues = 10000

    Seq(1, 100, numValues).foreach { numGroups =>
      Seq(128, numValues + 100).foreach { batchSize =>
        Seq(true, false).foreach { dictionaryEnabled =>
          withSQLConf(
            SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
            CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) {
            withTempPath { dir =>
              val path = new Path(dir.toURI.toString, "test.parquet")
              makeParquetFile(path, numValues, numGroups, dictionaryEnabled)
              withParquetTable(path.toUri.toString, "tbl") {
                checkSparkAnswer(
                  "SELECT _g3, _g4, SUM(_3), SUM(_4) FROM tbl GROUP BY _g3, _g4 ORDER BY _g3, _g4")
                checkSparkAnswer(
                  "SELECT _g3, _g4, SUM(DISTINCT _3), SUM(DISTINCT _4) FROM tbl GROUP BY _g3, _g4 ORDER BY _g3, _g4")
                checkSparkAnswer(
                  "SELECT _g3, _g4, COUNT(_3), COUNT(_4) FROM tbl GROUP BY _g3, _g4 ORDER BY _g3, _g4")
                checkSparkAnswer(
                  "SELECT _g3, _g4, COUNT(DISTINCT _3), COUNT(DISTINCT _4) FROM tbl GROUP BY _g3, _g4 ORDER BY _g3, _g4")
                checkSparkAnswer(
                  "SELECT _g3, _g4, MIN(_3), MAX(_3), MIN(_4), MAX(_4) FROM tbl GROUP BY _g3, _g4 ORDER BY _g3, _g4")
                checkSparkAnswer(
                  "SELECT _g3, _g4, AVG(_3), AVG(_4) FROM tbl GROUP BY _g3, _g4 ORDER BY _g3, _g4")
              }
            }
          }
        }
      }
    }
  }

  test("all types first/last, with nulls") {
    val numValues = 2048

    Seq(1, 100, numValues).foreach { numGroups =>
      Seq(true, false).foreach { dictionaryEnabled =>
        withTempPath { dir =>
          val path = new Path(dir.toURI.toString, "test.parquet")
          makeParquetFile(path, numValues, numGroups, dictionaryEnabled)
          withParquetTable(path.toUri.toString, "tbl") {
            Seq(128, numValues + 100).foreach { batchSize =>
              withSQLConf(
                CometConf.COMET_BATCH_SIZE.key -> batchSize.toString,
                CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true") {

                // Test all combinations of different aggregation & group-by types
                (1 to 14).foreach { gCol =>
                  withView("v") {
                    sql(s"CREATE TEMP VIEW v AS SELECT _g$gCol, _1, _2, _3, _4 " +
                      "FROM tbl ORDER BY _1, _2, _3, _4")
                    checkSparkAnswerAndOperator(
                      s"SELECT _g$gCol, FIRST(_1), FIRST(_2), FIRST(_3), " +
                        s"FIRST(_4), LAST(_1), LAST(_2), LAST(_3), LAST(_4) FROM v GROUP BY _g$gCol ORDER BY _g$gCol")
                  }
                }
              }
            }
          }
        }
      }
    }
  }

  test("first/last with ignore null") {
    val data = Range(0, 8192).flatMap(n => Seq((n, 1), (n, 2))).toDF("a", "b")
    withTempDir { dir =>
      val filename = s"${dir.getAbsolutePath}/first_last_ignore_null.parquet"
      data.write.parquet(filename)
      withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "100") {
        spark.read.parquet(filename).createOrReplaceTempView("t1")
        for (expr <- Seq("first", "last")) {
          // deterministic query that should return one non-null value per group
          val df = spark.sql(
            s"SELECT a, $expr(IF(b==1,null,b)) IGNORE NULLS FROM t1 GROUP BY a ORDER BY a")
          checkSparkAnswerAndOperator(df)
        }
      }
    }
  }

  test("all types, with nulls") {
    val numValues = 2048

    Seq(1, 100, numValues).foreach { numGroups =>
      Seq(true, false).foreach { dictionaryEnabled =>
        withTempPath { dir =>
          val path = new Path(dir.toURI.toString, "test.parquet")
          makeParquetFile(path, numValues, numGroups, dictionaryEnabled)
          withParquetTable(path.toUri.toString, "tbl") {
            Seq(128, numValues + 100).foreach { batchSize =>
              withSQLConf(
                CometConf.COMET_BATCH_SIZE.key -> batchSize.toString,
                CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "false") {

                // Test all combinations of different aggregation & group-by types
                (1 to 14).foreach { gCol =>
                  checkSparkAnswer(s"SELECT _g$gCol, SUM(_1), SUM(_2), COUNT(_3), COUNT(_4), " +
                    s"MIN(_1), MAX(_4), AVG(_2), AVG(_4) FROM tbl GROUP BY _g$gCol ORDER BY _g$gCol")
                  checkSparkAnswer(
                    s"SELECT _g$gCol, SUM(DISTINCT _3) FROM tbl GROUP BY _g$gCol ORDER BY _g$gCol")
                  checkSparkAnswer(
                    s"SELECT _g$gCol, COUNT(DISTINCT _1) FROM tbl GROUP BY _g$gCol ORDER BY _g$gCol")
                }
              }
            }
          }
        }
      }
    }
  }

  test("test final count") {
    withSQLConf(
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_SHUFFLE_MODE.key -> "native") {
      Seq(false, true).foreach { dictionaryEnabled =>
        withParquetTable((0 until 5).map(i => (i, i % 2)), "tbl", dictionaryEnabled) {
          checkSparkAnswerAndNumOfAggregates("SELECT _2, COUNT(_1) FROM tbl GROUP BY _2", 2)
          checkSparkAnswerAndNumOfAggregates("select count(_1) from tbl", 2)
          checkSparkAnswerAndNumOfAggregates(
            "SELECT _2, COUNT(_1), SUM(_1) FROM tbl GROUP BY _2",
            2)
          checkSparkAnswerAndNumOfAggregates("SELECT COUNT(_1), COUNT(_2) FROM tbl", 2)
        }
      }
    }
  }

  test("test final min/max") {
    withSQLConf(
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_SHUFFLE_MODE.key -> "native") {
      Seq(true, false).foreach { dictionaryEnabled =>
        withParquetTable((0 until 5).map(i => (i, i % 2)), "tbl", dictionaryEnabled) {
          checkSparkAnswerAndNumOfAggregates(
            "SELECT _2, MIN(_1), MAX(_1), COUNT(_1) FROM tbl GROUP BY _2",
            2)
          checkSparkAnswerAndNumOfAggregates("SELECT MIN(_1), MAX(_1), COUNT(_1) FROM tbl", 2)
          checkSparkAnswerAndNumOfAggregates(
            "SELECT _2, MIN(_1), MAX(_1), COUNT(_1), SUM(_1) FROM tbl GROUP BY _2",
            2)
          checkSparkAnswerAndNumOfAggregates(
            "SELECT MIN(_1), MIN(_2), MAX(_1), MAX(_2), COUNT(_1), COUNT(_2) FROM tbl",
            2)
        }
      }
    }
  }

  test("test final min/max/count with result expressions") {
    withSQLConf(
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_SHUFFLE_MODE.key -> "native") {
      Seq(true, false).foreach { dictionaryEnabled =>
        withParquetTable((0 until 5).map(i => (i, i % 2)), "tbl", dictionaryEnabled) {
          checkSparkAnswerAndNumOfAggregates(
            "SELECT _2, MIN(_1) + 2, COUNT(_1) FROM tbl GROUP BY _2",
            2)
          checkSparkAnswerAndNumOfAggregates("SELECT _2, COUNT(_1) + 2 FROM tbl GROUP BY _2", 2)
          checkSparkAnswerAndNumOfAggregates("SELECT _2 + 2, COUNT(_1) FROM tbl GROUP BY _2", 2)
          checkSparkAnswerAndNumOfAggregates(
            "SELECT _2, MIN(_1) + MAX(_1) FROM tbl GROUP BY _2",
            2)
          checkSparkAnswerAndNumOfAggregates("SELECT _2, MIN(_1) + _2 FROM tbl GROUP BY _2", 2)
          checkSparkAnswerAndNumOfAggregates(
            "SELECT _2 + 2, MIN(_1), MAX(_1), COUNT(_1) FROM tbl GROUP BY _2",
            2)
          checkSparkAnswerAndNumOfAggregates(
            "SELECT _2, MIN(_1), MAX(_1) + 2, COUNT(_1) FROM tbl GROUP BY _2",
            2)
          checkSparkAnswerAndNumOfAggregates("SELECT _2, SUM(_1) + 2 FROM tbl GROUP BY _2", 2)
          checkSparkAnswerAndNumOfAggregates("SELECT _2 + 2, SUM(_1) FROM tbl GROUP BY _2", 2)
          checkSparkAnswerAndNumOfAggregates("SELECT _2, SUM(_1 + 1) FROM tbl GROUP BY _2", 2)

          // result expression is unsupported by Comet, so only partial aggregation should be used
          val df = sql(
            "SELECT _2, MIN(_1) + java_method('java.lang.Math', 'random') " +
              "FROM tbl GROUP BY _2")
          assert(getNumCometHashAggregate(df) == 1)
        }
      }
    }
  }

  test("test final sum") {
    withSQLConf(
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_SHUFFLE_MODE.key -> "native") {
      Seq(false, true).foreach { dictionaryEnabled =>
        withParquetTable((0L until 5L).map(i => (i, i % 2)), "tbl", dictionaryEnabled) {
          checkSparkAnswerAndNumOfAggregates(
            "SELECT _2, SUM(_1), MIN(_1) FROM tbl GROUP BY _2",
            2)
          checkSparkAnswerAndNumOfAggregates("SELECT SUM(_1) FROM tbl", 2)
          checkSparkAnswerAndNumOfAggregates(
            "SELECT _2, MIN(_1), MAX(_1), COUNT(_1), SUM(_1), AVG(_1) FROM tbl GROUP BY _2",
            2)
          checkSparkAnswerAndNumOfAggregates(
            "SELECT MIN(_1), MIN(_2), MAX(_1), MAX(_2), COUNT(_1), COUNT(_2), SUM(_1), SUM(_2) FROM tbl",
            2)
        }
      }
    }
  }

  test("avg/sum overflow on decimal(38, _)") {
    val table = "overflow_decimal_38"
    withTable(table) {
      sql(s"create table $table(a decimal(38, 2), b INT) using parquet")
      sql(s"insert into $table values(42.00, 1), (999999999999999999999999999999999999.99, 1)")
      checkSparkAnswerAndNumOfAggregates(s"select sum(a) from $table", 2)
      sql(s"insert into $table values(42.00, 2), (99999999999999999999999999999999.99, 2)")
      sql(s"insert into $table values(999999999999999999999999999999999999.99, 3)")
      sql(s"insert into $table values(99999999999999999999999999999999.99, 4)")
      checkSparkAnswerAndNumOfAggregates(
        s"select avg(a), sum(a) from $table group by b order by b",
        2)
    }
  }

  test("test final avg") {
    withSQLConf(
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_SHUFFLE_MODE.key -> "native") {
      Seq(true, false).foreach { dictionaryEnabled =>
        withParquetTable(
          (0 until 5).map(i => (i.toDouble, i.toDouble % 2)),
          "tbl",
          dictionaryEnabled) {
          checkSparkAnswerAndNumOfAggregates("SELECT _2 , AVG(_1) FROM tbl GROUP BY _2", 2)
          checkSparkAnswerAndNumOfAggregates("SELECT AVG(_1) FROM tbl", 2)
          checkSparkAnswerAndNumOfAggregates(
            "SELECT _2, MIN(_1), MAX(_1), COUNT(_1), SUM(_1), AVG(_1) FROM tbl GROUP BY _2",
            2)
          checkSparkAnswerAndNumOfAggregates(
            "SELECT MIN(_1), MIN(_2), MAX(_1), MAX(_2), COUNT(_1), COUNT(_2), SUM(_1), SUM(_2), AVG(_1), AVG(_2) FROM tbl",
            2)
        }
      }
    }
  }

  test("final decimal avg") {
    withSQLConf(
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_SHUFFLE_MODE.key -> "native") {
      Seq(true, false).foreach { dictionaryEnabled =>
        withSQLConf("parquet.enable.dictionary" -> dictionaryEnabled.toString) {
          val table = s"final_decimal_avg_$dictionaryEnabled"
          withTable(table) {
            sql(s"create table $table(a decimal(38, 37), b INT) using parquet")
            sql(s"insert into $table values(-0.0000000000000000000000000000000000002, 1)")
            sql(s"insert into $table values(-0.0000000000000000000000000000000000002, 1)")
            sql(s"insert into $table values(-0.0000000000000000000000000000000000004, 2)")
            sql(s"insert into $table values(-0.0000000000000000000000000000000000004, 2)")
            sql(s"insert into $table values(-0.00000000000000000000000000000000000002, 3)")
            sql(s"insert into $table values(-0.00000000000000000000000000000000000002, 3)")
            sql(s"insert into $table values(-0.00000000000000000000000000000000000004, 4)")
            sql(s"insert into $table values(-0.00000000000000000000000000000000000004, 4)")
            sql(s"insert into $table values(0.13344406545919155429936259114971302408, 5)")
            sql(s"insert into $table values(0.13344406545919155429936259114971302408, 5)")

            checkSparkAnswerAndNumOfAggregates(s"SELECT b , AVG(a) FROM $table GROUP BY b", 2)
            checkSparkAnswerAndNumOfAggregates(s"SELECT AVG(a) FROM $table", 2)
            checkSparkAnswerAndNumOfAggregates(
              s"SELECT b, MIN(a), MAX(a), COUNT(a), SUM(a), AVG(a) FROM $table GROUP BY b",
              2)
            checkSparkAnswerAndNumOfAggregates(
              s"SELECT MIN(a), MAX(a), COUNT(a), SUM(a), AVG(a) FROM $table",
              2)
          }
        }
      }
    }
  }

  test("test partial avg") {
    Seq(true, false).foreach { dictionaryEnabled =>
      withParquetTable(
        (0 until 5).map(i => (i.toDouble, i.toDouble % 2)),
        "tbl",
        dictionaryEnabled) {
        withSQLConf(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "false") {
          checkSparkAnswerAndNumOfAggregates("SELECT _2 , AVG(_1) FROM tbl GROUP BY _2", 1)
        }
      }
    }
  }

  test("avg null handling") {
    withSQLConf(
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_SHUFFLE_MODE.key -> "native") {
      val table = "avg_null_handling"
      withTable(table) {
        sql(s"create table $table(a double, b double) using parquet")
        sql(s"insert into $table values(1, 1.0)")
        sql(s"insert into $table values(null, null)")
        sql(s"insert into $table values(1, 2.0)")
        sql(s"insert into $table values(null, null)")
        sql(s"insert into $table values(2, null)")
        sql(s"insert into $table values(2, null)")

        val query = sql(s"select a, AVG(b) from $table GROUP BY a")
        checkSparkAnswerAndOperator(query)
      }
    }
  }

  test("Decimal Avg with DF") {
    Seq(true, false).foreach { dictionaryEnabled =>
      Seq(true, false).foreach { nativeShuffleEnabled =>
        withSQLConf(
          CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> nativeShuffleEnabled.toString,
          CometConf.COMET_SHUFFLE_MODE.key -> "native",
          CometConf.getExprAllowIncompatConfigKey(classOf[Cast]) -> "true") {
          withTempDir { dir =>
            val path = new Path(dir.toURI.toString, "test")
            makeParquetFile(path, 1000, 20, dictionaryEnabled)
            withParquetTable(path.toUri.toString, "tbl") {
              val expectedNumOfCometAggregates = if (nativeShuffleEnabled) 2 else 1

              checkSparkAnswerAndNumOfAggregates(
                "SELECT _g2, AVG(_7) FROM tbl GROUP BY _g2",
                expectedNumOfCometAggregates)

              checkSparkAnswerWithTolerance("SELECT _g3, AVG(_8) FROM tbl GROUP BY _g3")
              assert(getNumCometHashAggregate(
                sql("SELECT _g3, AVG(_8) FROM tbl GROUP BY _g3")) == expectedNumOfCometAggregates)

              checkSparkAnswerAndNumOfAggregates(
                "SELECT _g4, AVG(_9) FROM tbl GROUP BY _g4",
                expectedNumOfCometAggregates)

              checkSparkAnswerAndNumOfAggregates(
                "SELECT AVG(_7) FROM tbl",
                expectedNumOfCometAggregates)

              checkSparkAnswerWithTolerance("SELECT AVG(_8) FROM tbl")
              assert(getNumCometHashAggregate(
                sql("SELECT AVG(_8) FROM tbl")) == expectedNumOfCometAggregates)

              checkSparkAnswerAndNumOfAggregates(
                "SELECT AVG(_9) FROM tbl",
                expectedNumOfCometAggregates)
            }
          }
        }
      }
    }
  }

  test("distinct") {
    withSQLConf(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true") {
      Seq("native", "jvm").foreach { cometShuffleMode =>
        withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> cometShuffleMode) {
          Seq(true, false).foreach { dictionary =>
            withSQLConf("parquet.enable.dictionary" -> dictionary.toString) {
              val table = "test"
              withTable(table) {
                sql(s"create table $table(col1 int, col2 int, col3 int) using parquet")
                sql(
                  s"insert into $table values(1, 1, 1), (1, 1, 1), (1, 3, 1), (1, 4, 2), (5, 3, 2)")

                checkSparkAnswerAndOperator(s"SELECT DISTINCT(col2) FROM $table")

                // Keep one explicit aggregate-count assertion so a silent regression in the
                // middle PartialMerge stage (which used to fall back before this PR) is caught:
                // all four Spark aggregate stages for COUNT(DISTINCT ...) must remain native.
                checkSparkAnswerAndNumOfAggregates(s"SELECT COUNT(distinct col2) FROM $table", 4)

                checkSparkAnswerAndOperator(
                  s"SELECT COUNT(distinct col2), col1 FROM $table group by col1")

                checkSparkAnswerAndOperator(s"SELECT SUM(distinct col2) FROM $table")

                checkSparkAnswerAndOperator(
                  s"SELECT SUM(distinct col2), col1 FROM $table group by col1")

                checkSparkAnswerAndOperator(
                  "SELECT COUNT(distinct col2), SUM(distinct col2), col1, COUNT(distinct col2)," +
                    s" SUM(distinct col2) FROM $table group by col1")

                checkSparkAnswerAndOperator(
                  "SELECT COUNT(col2), MIN(col2), COUNT(DISTINCT col2), SUM(col2)," +
                    s" SUM(DISTINCT col2), COUNT(DISTINCT col2), col1 FROM $table group by col1")
              }
            }
          }
        }
      }
    }
  }

  test("first/last") {
    withSQLConf(
      SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
      Seq(true, false).foreach { dictionary =>
        withSQLConf("parquet.enable.dictionary" -> dictionary.toString) {
          val table = "test"
          withTable(table) {
            sql(s"create table $table(col1 int, col2 int, col3 int) using parquet")
            sql(
              s"insert into $table values(4, 1, 1), (4, 1, 1), (3, 3, 1)," +
                " (2, 4, 2), (1, 3, 2), (null, 1, 1)")
            withView("t") {
              sql("CREATE VIEW t AS SELECT col1, col3 FROM test ORDER BY col1")

              checkSparkAnswerAndOperator("SELECT FIRST(col1), LAST(col1) FROM t")

              checkSparkAnswerAndOperator(
                "SELECT FIRST(col1), LAST(col1), MIN(col1), COUNT(col1) FROM t")

              checkSparkAnswerAndOperator(
                "SELECT FIRST(col1), LAST(col1), col3 FROM t GROUP BY col3")

              checkSparkAnswerAndOperator(
                "SELECT FIRST(col1), LAST(col1), MIN(col1), COUNT(col1), col3 FROM t GROUP BY col3")

              checkSparkAnswerAndOperator("SELECT FIRST(col1, true), LAST(col1) FROM t")

              checkSparkAnswerAndOperator(
                "SELECT FIRST(col1), LAST(col1, true), col3 FROM t GROUP BY col3")
            }
          }
        }
      }
    }
  }

  test("test bool_and/bool_or") {
    withSQLConf(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true") {
      Seq("native", "jvm").foreach { cometShuffleMode =>
        withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> cometShuffleMode) {
          Seq(true, false).foreach { dictionary =>
            withSQLConf("parquet.enable.dictionary" -> dictionary.toString) {
              val table = "test"
              withTable(table) {
                sql(s"create table $table(a boolean, b int) using parquet")
                sql(s"insert into $table values(true, 1)")
                sql(s"insert into $table values(false, 2)")
                sql(s"insert into $table values(true, 3)")
                sql(s"insert into $table values(true, 3)")
                // Spark maps BOOL_AND to MIN and BOOL_OR to MAX
                checkSparkAnswerAndNumOfAggregates(
                  s"SELECT MIN(a), MAX(a), BOOL_AND(a), BOOL_OR(a) FROM $table",
                  2)
              }
            }
          }
        }
      }
    }
  }

  test("bitwise aggregate") {
    withSQLConf(
      CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true",
      CometConf.COMET_SHUFFLE_MODE.key -> "jvm") {
      Seq(true, false).foreach { dictionary =>
        withSQLConf("parquet.enable.dictionary" -> dictionary.toString) {
          val table = "test"
          withTable(table) {
            sql(s"create table $table(col1 long, col2 int, col3 short, col4 byte) using parquet")
            sql(
              s"insert into $table values(4, 1, 1, 3), (4, 1, 1, 3), (3, 3, 1, 4)," +
                " (2, 4, 2, 5), (1, 3, 2, 6), (null, 1, 1, 7)")
            val expectedNumOfCometAggregates = 2
            checkSparkAnswerAndNumOfAggregates(
              "SELECT BIT_AND(col1), BIT_OR(col1), BIT_XOR(col1)," +
                " BIT_AND(col2), BIT_OR(col2), BIT_XOR(col2)," +
                " BIT_AND(col3), BIT_OR(col3), BIT_XOR(col3)," +
                " BIT_AND(col4), BIT_OR(col4), BIT_XOR(col4) FROM test",
              expectedNumOfCometAggregates)

            // Make sure the combination of BITWISE aggregates and other aggregates work OK
            checkSparkAnswerAndNumOfAggregates(
              "SELECT BIT_AND(col1), BIT_OR(col1), BIT_XOR(col1)," +
                " BIT_AND(col2), BIT_OR(col2), BIT_XOR(col2)," +
                " BIT_AND(col3), BIT_OR(col3), BIT_XOR(col3)," +
                " BIT_AND(col4), BIT_OR(col4), BIT_XOR(col4), MIN(col1), COUNT(col1) FROM test",
              expectedNumOfCometAggregates)

            checkSparkAnswerAndNumOfAggregates(
              "SELECT BIT_AND(col1), BIT_OR(col1), BIT_XOR(col1)," +
                " BIT_AND(col2), BIT_OR(col2), BIT_XOR(col2)," +
                " BIT_AND(col3), BIT_OR(col3), BIT_XOR(col3)," +
                " BIT_AND(col4), BIT_OR(col4), BIT_XOR(col4), col3 FROM test GROUP BY col3",
              expectedNumOfCometAggregates)

            // Make sure the combination of BITWISE aggregates and other aggregates work OK
            // with group by
            checkSparkAnswerAndNumOfAggregates(
              "SELECT BIT_AND(col1), BIT_OR(col1), BIT_XOR(col1)," +
                " BIT_AND(col2), BIT_OR(col2), BIT_XOR(col2)," +
                " BIT_AND(col3), BIT_OR(col3), BIT_XOR(col3)," +
                " BIT_AND(col4), BIT_OR(col4), BIT_XOR(col4)," +
                " MIN(col1), COUNT(col1), col3 FROM test GROUP BY col3",
              expectedNumOfCometAggregates)
          }
        }
      }
    }
  }

  def setupAndTestAggregates(
      table: String,
      data: Seq[(Any, Any, Any)],
      dataTypes: (String, String, String),
      aggregates: String): Unit = {
    val (type1, type2, type3) = dataTypes
    withTable(table) {
      sql(s"create table $table(col1 $type1, col2 $type2, col3 $type3) using parquet")
      val values = data
        .map { case (c1, c2, c3) =>
          s"($c1, $c2, $c3)"
        }
        .mkString(", ")
      sql(s"insert into $table values $values")

      val expectedNumOfCometAggregates = 2

      checkSparkAnswerWithTolAndNumOfAggregates(
        s"SELECT $aggregates FROM $table",
        expectedNumOfCometAggregates)

      checkSparkAnswerWithTolAndNumOfAggregates(
        s"SELECT $aggregates FROM $table GROUP BY col3",
        expectedNumOfCometAggregates)
    }
  }

  test("covariance & correlation") {
    withSQLConf(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true") {
      Seq("jvm", "native").foreach { cometShuffleMode =>
        withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> cometShuffleMode) {
          Seq(true, false).foreach { dictionary =>
            withSQLConf("parquet.enable.dictionary" -> dictionary.toString) {
              Seq(true, false).foreach { nullOnDivideByZero =>
                withSQLConf(
                  "spark.sql.legacy.statisticalAggregate" -> nullOnDivideByZero.toString) {
                  val table = "test"
                  val aggregates =
                    "covar_samp(col1, col2), covar_pop(col1, col2), corr(col1, col2)"
                  setupAndTestAggregates(
                    table,
                    Seq((1, 4, 1), (2, 5, 1), (3, 6, 2)),
                    ("double", "double", "double"),
                    aggregates)
                  setupAndTestAggregates(
                    table,
                    Seq((1, 4, 3), (2, -5, 3), (3, 6, 1)),
                    ("double", "double", "double"),
                    aggregates)
                  setupAndTestAggregates(
                    table,
                    Seq((1.1, 4.1, 2.3), (2, 5, 1.5), (3, 6, 2.3)),
                    ("double", "double", "double"),
                    aggregates)
                  setupAndTestAggregates(
                    table,
                    Seq(
                      (1, 4, 1),
                      (2, 5, 2),
                      (3, 6, 3),
                      (1.1, 4.4, 1),
                      (2.2, 5.5, 2),
                      (3.3, 6.6, 3)),
                    ("double", "double", "double"),
                    aggregates)
                  setupAndTestAggregates(
                    table,
                    Seq((1, 4, 1), (2, 5, 2), (3, 6, 3)),
                    ("int", "int", "int"),
                    aggregates)
                  setupAndTestAggregates(
                    table,
                    Seq((1, 4, 2), (null, null, 2), (3, 6, 1), (3, 3, 1)),
                    ("int", "int", "int"),
                    aggregates)
                  setupAndTestAggregates(
                    table,
                    Seq((1, 4, 1), (null, 5, 1), (2, 5, 2), (9, null, 2), (3, 6, 2)),
                    ("int", "int", "int"),
                    aggregates)
                  setupAndTestAggregates(
                    table,
                    Seq((null, null, 1), (1, 2, 1), (null, null, 2)),
                    ("int", "int", "int"),
                    aggregates)
                  setupAndTestAggregates(
                    table,
                    Seq((null, null, 1), (null, null, 1), (null, null, 2)),
                    ("int", "int", "int"),
                    aggregates)
                }
              }
            }
          }
        }
      }
    }
  }

  test("corr - nan/null") {
    Seq(true, false).foreach { nullOnDivideByZero =>
      withSQLConf("spark.sql.legacy.statisticalAggregate" -> nullOnDivideByZero.toString) {
        withTable("t") {
          sql("""create table t using parquet as
              select cast(null as float) f1, CAST('NaN' AS float) f2, cast(null as double) d1, CAST('NaN' AS double) d2
              from range(1)
            """)

          checkSparkAnswerAndOperator("""
              |select
              | corr(f1, f2) c1,
              | corr(f1, f1) c2,
              | corr(f2, f1) c3,
              | corr(f2, f2) c4,
              | corr(d1, d2) c5,
              | corr(d1, d1) c6,
              | corr(d2, d1) c7,
              | corr(d2, d2) c8
              | FROM t""".stripMargin)
        }
      }
    }
  }

  test("var_pop and var_samp") {
    withSQLConf(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true") {
      Seq("native", "jvm").foreach { cometShuffleMode =>
        withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> cometShuffleMode) {
          Seq(true, false).foreach { dictionary =>
            withSQLConf("parquet.enable.dictionary" -> dictionary.toString) {
              Seq(true, false).foreach { nullOnDivideByZero =>
                withSQLConf(
                  "spark.sql.legacy.statisticalAggregate" -> nullOnDivideByZero.toString) {
                  val table = "test"
                  withTable(table) {
                    sql(s"create table $table(col1 int, col2 int, col3 int, col4 float, col5 double, col6 int) using parquet")
                    sql(s"insert into $table values(1, null, null, 1.1, 2.2, 1)," +
                      " (2, null, null, 3.4, 5.6, 1), (3, null, 4, 7.9, 2.4, 2)")
                    val expectedNumOfCometAggregates = 2
                    checkSparkAnswerWithTolAndNumOfAggregates(
                      "SELECT var_samp(col1), var_samp(col2), var_samp(col3), var_samp(col4), var_samp(col5) FROM test",
                      expectedNumOfCometAggregates)
                    checkSparkAnswerWithTolAndNumOfAggregates(
                      "SELECT var_pop(col1), var_pop(col2), var_pop(col3), var_pop(col4), var_samp(col5) FROM test",
                      expectedNumOfCometAggregates)
                    checkSparkAnswerAndNumOfAggregates(
                      "SELECT var_samp(col1), var_samp(col2), var_samp(col3), var_samp(col4), var_samp(col5)" +
                        " FROM test GROUP BY col6",
                      expectedNumOfCometAggregates)
                    checkSparkAnswerAndNumOfAggregates(
                      "SELECT var_pop(col1), var_pop(col2), var_pop(col3), var_pop(col4), var_samp(col5)" +
                        " FROM test GROUP BY col6",
                      expectedNumOfCometAggregates)
                  }
                }
              }
            }
          }
        }
      }
    }
  }

  test("stddev_pop and stddev_samp") {
    withSQLConf(CometConf.COMET_EXEC_SHUFFLE_ENABLED.key -> "true") {
      Seq("native", "jvm").foreach { cometShuffleMode =>
        withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> cometShuffleMode) {
          Seq(true, false).foreach { dictionary =>
            withSQLConf("parquet.enable.dictionary" -> dictionary.toString) {
              Seq(true, false).foreach { nullOnDivideByZero =>
                withSQLConf(
                  "spark.sql.legacy.statisticalAggregate" -> nullOnDivideByZero.toString) {
                  val table = "test"
                  withTable(table) {
                    sql(s"create table $table(col1 int, col2 int, col3 int, col4 float, " +
                      "col5 double, col6 int) using parquet")
                    sql(s"insert into $table values(1, null, null, 1.1, 2.2, 1), " +
                      "(2, null, null, 3.4, 5.6, 1), (3, null, 4, 7.9, 2.4, 2)")
                    val expectedNumOfCometAggregates = 2
                    checkSparkAnswerWithTolAndNumOfAggregates(
                      "SELECT stddev_samp(col1), stddev_samp(col2), stddev_samp(col3), " +
                        "stddev_samp(col4), stddev_samp(col5) FROM test",
                      expectedNumOfCometAggregates)
                    checkSparkAnswerWithTolAndNumOfAggregates(
                      "SELECT stddev_pop(col1), stddev_pop(col2), stddev_pop(col3), " +
                        "stddev_pop(col4), stddev_pop(col5) FROM test",
                      expectedNumOfCometAggregates)
                    checkSparkAnswerAndNumOfAggregates(
                      "SELECT stddev_samp(col1), stddev_samp(col2), stddev_samp(col3), " +
                        "stddev_samp(col4), stddev_samp(col5) FROM test GROUP BY col6",
                      expectedNumOfCometAggregates)
                    checkSparkAnswerWithTolAndNumOfAggregates(
                      "SELECT stddev_pop(col1), stddev_pop(col2), stddev_pop(col3), " +
                        "stddev_pop(col4), stddev_pop(col5) FROM test GROUP BY col6",
                      expectedNumOfCometAggregates)
                  }
                }
              }
            }
          }
        }
      }
    }
  }

  test("AVG and try_avg - basic functionality") {
    withParquetTable(
      Seq(
        (10L, 1),
        (20L, 1),
        (null.asInstanceOf[Long], 1),
        (100L, 2),
        (200L, 2),
        (null.asInstanceOf[Long], 3)),
      "tbl") {

      Seq(true, false).foreach({ ansiMode =>
        // without GROUP BY
        withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiMode.toString) {
          val res = sql("SELECT avg(_1) FROM tbl")
          checkSparkAnswerAndOperator(res)
        }

        // with GROUP BY
        withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiMode.toString) {
          val res = sql("SELECT _2, avg(_1) FROM tbl GROUP BY _2")
          checkSparkAnswerAndOperator(res)
        }
      })

      // try_avg without GROUP BY
      val resTry = sql("SELECT try_avg(_1) FROM tbl")
      checkSparkAnswerAndOperator(resTry)

      // try_avg with GROUP BY
      val resTryGroup = sql("SELECT _2, try_avg(_1) FROM tbl GROUP BY _2")
      checkSparkAnswerAndOperator(resTryGroup)

    }
  }

  test("AVG and try_avg - special numbers") {

    val negativeNumbers: Seq[(Long, Int)] = Seq(
      (-1L, 1),
      (-123L, 1),
      (-456L, 1),
      (Long.MinValue, 1),
      (Long.MinValue, 1),
      (Long.MinValue, 2),
      (Long.MinValue, 2),
      (null.asInstanceOf[Long], 3))

    val zeroSeq: Seq[(Long, Int)] =
      Seq((0L, 1), (-0L, 1), (+0L, 2), (+0L, 2), (null.asInstanceOf[Long], 3))

    val highValNumbers: Seq[(Long, Int)] = Seq(
      (Long.MaxValue, 1),
      (Long.MaxValue, 1),
      (Long.MaxValue, 2),
      (Long.MaxValue, 2),
      (null.asInstanceOf[Long], 3))

    val inputs = Seq(negativeNumbers, highValNumbers, zeroSeq)
    inputs.foreach(inputSeq => {
      withParquetTable(inputSeq, "tbl") {
        Seq(true, false).foreach({ ansiMode =>
          // without GROUP BY
          withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiMode.toString) {
            checkSparkAnswerAndOperator("SELECT avg(_1) FROM tbl")
          }

          // with GROUP BY
          withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiMode.toString) {
            checkSparkAnswerAndOperator("SELECT _2, avg(_1) FROM tbl GROUP BY _2")
          }
        })

        // try_avg without GROUP BY
        checkSparkAnswerAndOperator("SELECT try_avg(_1) FROM tbl")

        // try_avg with GROUP BY
        checkSparkAnswerAndOperator("SELECT _2, try_avg(_1) FROM tbl GROUP BY _2")

      }
    })
  }

  test("ANSI support for sum - null test") {
    Seq(true, false).foreach { ansiEnabled =>
      withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
        withParquetTable(
          Seq((null.asInstanceOf[java.lang.Long], "a"), (null.asInstanceOf[java.lang.Long], "b")),
          "null_tbl") {
          val res = sql("SELECT sum(_1) FROM null_tbl")
          checkSparkAnswerAndOperator(res)
        }
      }
    }
  }

  test("ANSI support for decimal sum - null test") {
    Seq(true, false).foreach { ansiEnabled =>
      withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
        withParquetTable(
          Seq(
            (null.asInstanceOf[java.math.BigDecimal], "a"),
            (null.asInstanceOf[java.math.BigDecimal], "b")),
          "null_tbl") {
          val res = sql("SELECT sum(_1) FROM null_tbl")
          checkSparkAnswerAndOperator(res)
          assert(res.collect() === Array(Row(null)))
        }
      }
    }
  }

  test("ANSI support for try_sum - null test") {
    Seq(true, false).foreach { ansiEnabled =>
      withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
        withParquetTable(
          Seq((null.asInstanceOf[java.lang.Long], "a"), (null.asInstanceOf[java.lang.Long], "b")),
          "null_tbl") {
          val res = sql("SELECT try_sum(_1) FROM null_tbl")
          checkSparkAnswerAndOperator(res)
        }
      }
    }
  }

  test("ANSI support for try_sum decimal - null test") {
    Seq(true, false).foreach { ansiEnabled =>
      withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
        withParquetTable(
          Seq(
            (null.asInstanceOf[java.math.BigDecimal], "a"),
            (null.asInstanceOf[java.math.BigDecimal], "b")),
          "null_tbl") {
          val res = sql("SELECT try_sum(_1) FROM null_tbl")
          checkSparkAnswerAndOperator(res)
          assert(res.collect() === Array(Row(null)))
        }
      }
    }
  }

  test("ANSI support for sum - null test (group by)") {
    Seq(true, false).foreach { ansiEnabled =>
      withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
        withParquetTable(
          Seq(
            (null.asInstanceOf[java.lang.Long], "a"),
            (null.asInstanceOf[java.lang.Long], "a"),
            (null.asInstanceOf[java.lang.Long], "b"),
            (null.asInstanceOf[java.lang.Long], "b"),
            (null.asInstanceOf[java.lang.Long], "b")),
          "tbl") {
          val res = sql("SELECT _2, sum(_1) FROM tbl group by 1")
          checkSparkAnswerAndOperator(res)
          assert(res.orderBy(col("_2")).collect() === Array(Row("a", null), Row("b", null)))
        }
      }
    }
  }

  test("ANSI support for decimal sum - null test (group by)") {
    Seq(true, false).foreach { ansiEnabled =>
      withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
        withParquetTable(
          Seq(
            (null.asInstanceOf[java.math.BigDecimal], "a"),
            (null.asInstanceOf[java.math.BigDecimal], "a"),
            (null.asInstanceOf[java.math.BigDecimal], "b"),
            (null.asInstanceOf[java.math.BigDecimal], "b"),
            (null.asInstanceOf[java.math.BigDecimal], "b")),
          "tbl") {
          val res = sql("SELECT _2, sum(_1) FROM tbl group by 1")
          checkSparkAnswerAndOperator(res)
          assert(res.orderBy(col("_2")).collect() === Array(Row("a", null), Row("b", null)))
        }
      }
    }
  }

  test("ANSI support for try_sum - null test (group by)") {
    Seq(true, false).foreach { ansiEnabled =>
      withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
        withParquetTable(
          Seq(
            (null.asInstanceOf[java.lang.Long], "a"),
            (null.asInstanceOf[java.lang.Long], "a"),
            (null.asInstanceOf[java.lang.Long], "b"),
            (null.asInstanceOf[java.lang.Long], "b"),
            (null.asInstanceOf[java.lang.Long], "b")),
          "tbl") {
          val res = sql("SELECT _2, try_sum(_1) FROM tbl group by 1")
          checkSparkAnswerAndOperator(res)
        }
      }
    }
  }

  test("ANSI support for try_sum decimal - null test (group by)") {
    Seq(true, false).foreach { ansiEnabled =>
      withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
        withParquetTable(
          Seq(
            (null.asInstanceOf[java.math.BigDecimal], "a"),
            (null.asInstanceOf[java.math.BigDecimal], "a"),
            (null.asInstanceOf[java.math.BigDecimal], "b"),
            (null.asInstanceOf[java.math.BigDecimal], "b"),
            (null.asInstanceOf[java.math.BigDecimal], "b")),
          "tbl") {
          val res = sql("SELECT _2, try_sum(_1) FROM tbl group by 1")
          checkSparkAnswerAndOperator(res)
        }
      }
    }
  }

  protected def generateOverflowDecimalInputs: Seq[(java.math.BigDecimal, Int)] = {
    val maxDec38_0 = new java.math.BigDecimal("99999999999999999999")
    (1 to 50).flatMap(_ => Seq((maxDec38_0, 1)))
  }

  test("ANSI support - SUM function") {
    Seq(true, false).foreach { ansiEnabled =>
      withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
        // Test long overflow
        withParquetTable(Seq((Long.MaxValue, 1L), (100L, 1L)), "tbl") {
          val res = sql("SELECT SUM(_1) FROM tbl")
          if (ansiEnabled) {
            checkSparkAnswerMaybeThrows(res) match {
              case (Some(sparkExc), Some(cometExc)) =>
                // make sure that the error message throws overflow exception only
                assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
                assert(cometExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
              case _ => fail("Exception should be thrown for Long overflow in ANSI mode")
            }
          } else {
            checkSparkAnswerAndOperator(res)
          }
        }
        // Test long underflow
        withParquetTable(Seq((Long.MinValue, 1L), (-100L, 1L)), "tbl") {
          val res = sql("SELECT SUM(_1) FROM tbl")
          if (ansiEnabled) {
            checkSparkAnswerMaybeThrows(res) match {
              case (Some(sparkExc), Some(cometExc)) =>
                assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
                assert(cometExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
              case _ => fail("Exception should be thrown for Long underflow in ANSI mode")
            }
          } else {
            checkSparkAnswerAndOperator(res)
          }
        }
        // Test Int SUM (should not overflow)
        withParquetTable(Seq((Int.MaxValue, 1), (Int.MaxValue, 1), (100, 1)), "tbl") {
          val res = sql("SELECT SUM(_1) FROM tbl")
          checkSparkAnswerAndOperator(res)
        }
        // Test Short SUM (should not overflow)
        withParquetTable(
          Seq((Short.MaxValue, 1.toShort), (Short.MaxValue, 1.toShort), (100.toShort, 1.toShort)),
          "tbl") {
          val res = sql("SELECT SUM(_1) FROM tbl")
          checkSparkAnswerAndOperator(res)
        }
        // Test Byte SUM (should not overflow)
        withParquetTable(
          Seq((Byte.MaxValue, 1.toByte), (Byte.MaxValue, 1.toByte), (10.toByte, 1.toByte)),
          "tbl") {
          val res = sql("SELECT SUM(_1) FROM tbl")
          checkSparkAnswerAndOperator(res)
        }
      }
    }
  }

  test("ANSI support for decimal SUM function") {
    Seq(true, false).foreach { ansiEnabled =>
      withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
        withParquetTable(generateOverflowDecimalInputs, "tbl") {
          val res = sql("SELECT SUM(_1) FROM tbl")
          if (ansiEnabled) {
            checkSparkAnswerMaybeThrows(res) match {
              case (Some(sparkExc), Some(cometExc)) =>
                assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
                assert(cometExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
              case _ =>
                fail("Exception should be thrown for decimal overflow in ANSI mode")
            }
          } else {
            checkSparkAnswerAndOperator(res)
          }
        }
      }
    }
  }

  test("ANSI support for SUM - GROUP BY") {
    Seq(true, false).foreach { ansiEnabled =>
      withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
        withParquetTable(
          Seq((Long.MaxValue, 1), (100L, 1), (Long.MaxValue, 2), (200L, 2)),
          "tbl") {
          val res = sql("SELECT _2, SUM(_1) FROM tbl GROUP BY _2").repartition(2)
          if (ansiEnabled) {
            checkSparkAnswerMaybeThrows(res) match {
              case (Some(sparkExc), Some(cometExc)) =>
                assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
                assert(cometExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
              case _ =>
                fail("Exception should be thrown for Long overflow with GROUP BY in ANSI mode")
            }
          } else {
            checkSparkAnswerAndOperator(res)
          }
        }

        withParquetTable(
          Seq((Long.MinValue, 1), (-100L, 1), (Long.MinValue, 2), (-200L, 2)),
          "tbl") {
          val res = sql("SELECT _2, SUM(_1) FROM tbl GROUP BY _2")
          if (ansiEnabled) {
            checkSparkAnswerMaybeThrows(res) match {
              case (Some(sparkExc), Some(cometExc)) =>
                assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
                assert(cometExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
              case _ =>
                fail("Exception should be thrown for Long underflow with GROUP BY in ANSI mode")
            }
          } else {
            checkSparkAnswerAndOperator(res)
          }
        }
        // Test Int with GROUP BY
        withParquetTable(Seq((Int.MaxValue, 1), (Int.MaxValue, 1), (100, 2), (200, 2)), "tbl") {
          val res = sql("SELECT _2, SUM(_1) FROM tbl GROUP BY _2")
          checkSparkAnswerAndOperator(res)
        }
        // Test Short with GROUP BY
        withParquetTable(
          Seq((Short.MaxValue, 1), (Short.MaxValue, 1), (100.toShort, 2), (200.toShort, 2)),
          "tbl") {
          val res = sql("SELECT _2, SUM(_1) FROM tbl GROUP BY _2")
          checkSparkAnswerAndOperator(res)
        }
        // Test Byte with GROUP BY
        withParquetTable(
          Seq((Byte.MaxValue, 1), (Byte.MaxValue, 1), (10.toByte, 2), (20.toByte, 2)),
          "tbl") {
          val res = sql("SELECT _2, SUM(_1) FROM tbl GROUP BY _2")
          checkSparkAnswerAndOperator(res)
        }
      }
    }
  }

  test("ANSI support for decimal SUM - GROUP BY") {
    Seq(true, false).foreach { ansiEnabled =>
      withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
        withParquetTable(generateOverflowDecimalInputs, "tbl") {
          val res =
            sql("SELECT _2, SUM(_1) FROM tbl GROUP BY _2").repartition(2)
          if (ansiEnabled) {
            checkSparkAnswerMaybeThrows(res) match {
              case (Some(sparkExc), Some(cometExc)) =>
                assert(sparkExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
                assert(cometExc.getMessage.contains("ARITHMETIC_OVERFLOW"))
              case _ =>
                fail("Exception should be thrown for decimal overflow with GROUP BY in ANSI mode")
            }
          } else {
            checkSparkAnswerAndOperator(res)
          }
        }
      }
    }
  }

  test("try_sum overflow - with GROUP BY") {
    // Test Long overflow with GROUP BY - some groups overflow while some don't
    withParquetTable(Seq((Long.MaxValue, 1), (100L, 1), (200L, 2), (300L, 2)), "tbl") {
      // repartition to trigger merge batch and state checks
      val res = sql("SELECT _2, try_sum(_1) FROM tbl GROUP BY _2").repartition(2, col("_2"))
      // first group should return NULL (overflow) and group 2 should return 500
      checkSparkAnswerAndOperator(res)
    }

    // Test Long underflow with GROUP BY
    withParquetTable(Seq((Long.MinValue, 1), (-100L, 1), (-200L, 2), (-300L, 2)), "tbl") {
      val res = sql("SELECT _2, try_sum(_1) FROM tbl GROUP BY _2").repartition(2, col("_2"))
      // first group should return NULL (underflow), second group should return neg 500
      checkSparkAnswerAndOperator(res)
    }

    // Test all groups overflow
    withParquetTable(Seq((Long.MaxValue, 1), (100L, 1), (Long.MaxValue, 2), (100L, 2)), "tbl") {
      val res = sql("SELECT _2, try_sum(_1) FROM tbl GROUP BY _2").repartition(2, col("_2"))
      // Both groups should return NULL
      checkSparkAnswerAndOperator(res)
    }

    // Test Short with GROUP BY (should NOT overflow)
    withParquetTable(
      Seq((Short.MaxValue, 1), (Short.MaxValue, 1), (100.toShort, 2), (200.toShort, 2)),
      "tbl") {
      val res = sql("SELECT _2, try_sum(_1) FROM tbl GROUP BY _2").repartition(2, col("_2"))
      checkSparkAnswerAndOperator(res)
    }

    // Test Byte with GROUP BY (no overflow)
    withParquetTable(
      Seq((Byte.MaxValue, 1), (Byte.MaxValue, 1), (10.toByte, 2), (20.toByte, 2)),
      "tbl") {
      val res = sql("SELECT _2, try_sum(_1) FROM tbl GROUP BY _2").repartition(2, col("_2"))
      checkSparkAnswerAndOperator(res)
    }
  }

  test("try_sum decimal overflow") {
    withParquetTable(generateOverflowDecimalInputs, "tbl") {
      val res = sql("SELECT try_sum(_1) FROM tbl")
      checkSparkAnswerAndOperator(res)
    }
  }

  test("try_sum decimal overflow - with GROUP BY") {
    withParquetTable(generateOverflowDecimalInputs, "tbl") {
      val res = sql("SELECT _2, try_sum(_1) FROM tbl GROUP BY _2").repartition(2, col("_2"))
      checkSparkAnswerAndOperator(res)
    }
  }

  test("try_sum decimal partial overflow - with GROUP BY") {
    // Group 1 overflows, Group 2 succeeds
    val data: Seq[(java.math.BigDecimal, Int)] = generateOverflowDecimalInputs ++ Seq(
      (new java.math.BigDecimal(300), 2),
      (new java.math.BigDecimal(200), 2))
    withParquetTable(data, "tbl") {
      val res = sql("SELECT _2, try_sum(_1) FROM tbl GROUP BY _2")
      // Group 1 should be NULL, Group 2 should be 500
      checkSparkAnswerAndOperator(res)
    }
  }

  test("SumDecimal and AvgDecimal nullable should always be true") {
    // SumDecimal and AvgDecimal currently hardcode nullable=true.
    // This matches Spark's Sum.nullable and Average.nullable which always return true,
    // regardless of ANSI mode or input nullability.
    val nonNullableData: Seq[(java.math.BigDecimal, Int)] = Seq(
      (new java.math.BigDecimal("10.00"), 1),
      (new java.math.BigDecimal("20.00"), 1),
      (new java.math.BigDecimal("30.00"), 2))

    val nullableData: Seq[(java.math.BigDecimal, Int)] = Seq(
      (new java.math.BigDecimal("10.00"), 1),
      (null.asInstanceOf[java.math.BigDecimal], 1),
      (new java.math.BigDecimal("30.00"), 2))

    Seq(true, false).foreach { ansiEnabled =>
      withSQLConf(SQLConf.ANSI_ENABLED.key -> ansiEnabled.toString) {
        withParquetTable(nonNullableData, "tbl") {
          val sumRes = sql("SELECT _2, sum(_1) FROM tbl GROUP BY _2")
          checkSparkAnswerAndOperator(sumRes)
          assert(sumRes.schema.fields(1).nullable == true)

          val avgRes = sql("SELECT _2, avg(_1) FROM tbl GROUP BY _2")
          checkSparkAnswerAndOperator(avgRes)
          assert(avgRes.schema.fields(1).nullable == true)
        }

        withParquetTable(nullableData, "tbl") {
          val sumRes = sql("SELECT _2, sum(_1) FROM tbl GROUP BY _2")
          checkSparkAnswerAndOperator(sumRes)
          assert(sumRes.schema.fields(1).nullable == true)

          val avgRes = sql("SELECT _2, avg(_1) FROM tbl GROUP BY _2")
          checkSparkAnswerAndOperator(avgRes)
          assert(avgRes.schema.fields(1).nullable == true)
        }
      }
    }
  }

  protected def checkSparkAnswerAndNumOfAggregates(query: String, numAggregates: Int): Unit = {
    val df = sql(query)
    checkSparkAnswer(df)
    val actualNumAggregates = getNumCometHashAggregate(df)
    assert(
      actualNumAggregates == numAggregates,
      s"Expected $numAggregates Comet aggregate operators, but found $actualNumAggregates")
  }

  protected def checkSparkAnswerWithTolAndNumOfAggregates(
      query: String,
      numAggregates: Int,
      absTol: Double = 1e-6): Unit = {
    val df = sql(query)
    checkSparkAnswerWithTolerance(df, absTol)
    val actualNumAggregates = getNumCometHashAggregate(df)
    assert(
      actualNumAggregates == numAggregates,
      s"Expected $numAggregates Comet aggregate operators, but found $actualNumAggregates")
  }

  def getNumCometHashAggregate(df: DataFrame): Int = {
    val sparkPlan = stripAQEPlan(df.queryExecution.executedPlan)
    sparkPlan.collect { case s: CometHashAggregateExec => s }.size
  }

}
