[SPARK-58761][SQL][TEST] Add dedicated e2e test coverage for GROUP BY over nanosecond-precision timestamp keys

sql/core/src/test/resources/sql-tests/analyzer-results/timestamp-ltz-nanos.sql.out

@@ -740,6 +740,21 @@ Sort [c#x ASC NULLS FIRST], true
       +- LocalRelation [c#x]
 
 
+-- !query
+SELECT k, count(*), sum(v) FROM VALUES
+  (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 1),
+  (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 2),
+  (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC', 3),
+  (CAST(NULL AS timestamp_ltz(9)), 4),
+  (CAST(NULL AS timestamp_ltz(9)), 5) AS t(k, v)
+  GROUP BY k ORDER BY k
+-- !query analysis
+Sort [k#x ASC NULLS FIRST], true
++- Aggregate [k#x], [k#x, count(1) AS count(1)#xL, sum(v#x) AS sum(v)#xL]
+   +- SubqueryAlias t
+      +- LocalRelation [k#x, v#x]
+
+
 -- !query
 SELECT mode(c) FROM VALUES
   (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'),

sql/core/src/test/resources/sql-tests/analyzer-results/timestamp-ntz-nanos.sql.out

@@ -668,6 +668,21 @@ Sort [c#x ASC NULLS FIRST], true
       +- LocalRelation [c#x]
 
 
+-- !query
+SELECT k, count(*), sum(v) FROM VALUES
+  (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 1),
+  (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 2),
+  (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', 3),
+  (CAST(NULL AS timestamp_ntz(9)), 4),
+  (CAST(NULL AS timestamp_ntz(9)), 5) AS t(k, v)
+  GROUP BY k ORDER BY k
+-- !query analysis
+Sort [k#x ASC NULLS FIRST], true
++- Aggregate [k#x], [k#x, count(1) AS count(1)#xL, sum(v#x) AS sum(v)#xL]
+   +- SubqueryAlias t
+      +- LocalRelation [k#x, v#x]
+
+
 -- !query
 SELECT mode(c) FROM VALUES
   (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'),

sql/core/src/test/resources/sql-tests/inputs/timestamp-ltz-nanos.sql

@@ -219,6 +219,17 @@ SELECT c, count(*) FROM VALUES
   (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC'),
   (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC') AS t(c)
   GROUP BY c ORDER BY c;
+-- GROUP BY a nanosecond key with aggregates and a NULL group: exact-duplicate keys collapse, two
+-- keys sharing epochMicros but differing within the microsecond stay in separate groups, and all
+-- NULL keys group together (unlike an equi-join). Three groups: .000000001 (count 2, sum 3),
+-- .000000999 (count 1, sum 3), NULL (count 2, sum 9). Values render in the session time zone.
+SELECT k, count(*), sum(v) FROM VALUES
+  (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 1),
+  (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 2),
+  (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC', 3),
+  (CAST(NULL AS timestamp_ltz(9)), 4),
+  (CAST(NULL AS timestamp_ltz(9)), 5) AS t(k, v)
+  GROUP BY k ORDER BY k;
 
 -- SPARK-56822: mode over nanosecond-precision TIMESTAMP_LTZ. Frequencies are counted on the full
 -- nanos value, so the most-frequent value is selected down to the sub-microsecond and the result

sql/core/src/test/resources/sql-tests/inputs/timestamp-ntz-nanos.sql

@@ -192,6 +192,17 @@ SELECT c, count(*) FROM VALUES
   (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999'),
   (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001') AS t(c)
   GROUP BY c ORDER BY c;
+-- GROUP BY a nanosecond key with aggregates and a NULL group: exact-duplicate keys collapse, two
+-- keys sharing epochMicros but differing within the microsecond stay in separate groups, and all
+-- NULL keys group together (unlike an equi-join). Three groups: .000000001 (count 2, sum 3),
+-- .000000999 (count 1, sum 3), NULL (count 2, sum 9).
+SELECT k, count(*), sum(v) FROM VALUES
+  (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 1),
+  (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 2),
+  (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', 3),
+  (CAST(NULL AS timestamp_ntz(9)), 4),
+  (CAST(NULL AS timestamp_ntz(9)), 5) AS t(k, v)
+  GROUP BY k ORDER BY k;
 
 -- SPARK-56822: mode over nanosecond-precision TIMESTAMP_NTZ. Frequencies are counted on the full
 -- nanos value, so the most-frequent value is selected down to the sub-microsecond and the result

sql/core/src/test/resources/sql-tests/results/timestamp-ltz-nanos.sql.out

@@ -832,6 +832,22 @@ struct<c:timestamp_ltz(9),count(1):bigint>
 2019-12-31 16:00:00.000000999	1
 
 
+-- !query
+SELECT k, count(*), sum(v) FROM VALUES
+  (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 1),
+  (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC', 2),
+  (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000999 UTC', 3),
+  (CAST(NULL AS timestamp_ltz(9)), 4),
+  (CAST(NULL AS timestamp_ltz(9)), 5) AS t(k, v)
+  GROUP BY k ORDER BY k
+-- !query schema
+struct<k:timestamp_ltz(9),count(1):bigint,sum(v):bigint>
+-- !query output
+NULL	2	9
+2019-12-31 16:00:00.000000001	2	3
+2019-12-31 16:00:00.000000999	1	3
+
+
 -- !query
 SELECT mode(c) FROM VALUES
   (TIMESTAMP_LTZ '2020-01-01 00:00:00.000000001 UTC'),

sql/core/src/test/resources/sql-tests/results/timestamp-ntz-nanos.sql.out

@@ -750,6 +750,22 @@ struct<c:timestamp_ntz(9),count(1):bigint>
 2020-01-01 00:00:00.000000999	1
 
 
+-- !query
+SELECT k, count(*), sum(v) FROM VALUES
+  (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 1),
+  (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001', 2),
+  (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000999', 3),
+  (CAST(NULL AS timestamp_ntz(9)), 4),
+  (CAST(NULL AS timestamp_ntz(9)), 5) AS t(k, v)
+  GROUP BY k ORDER BY k
+-- !query schema
+struct<k:timestamp_ntz(9),count(1):bigint,sum(v):bigint>
+-- !query output
+NULL	2	9
+2020-01-01 00:00:00.000000001	2	3
+2020-01-01 00:00:00.000000999	1	3
+
+
 -- !query
 SELECT mode(c) FROM VALUES
   (TIMESTAMP_NTZ '2020-01-01 00:00:00.000000001'),

sql/core/src/test/scala/org/apache/spark/sql/TimestampNanosGroupBySuiteBase.scala

@@ -0,0 +1,190 @@
+/*
+ * 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.spark.sql
+
+import java.time.{Instant, LocalDateTime}
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types._
+
+/**
+ * End-to-end `GROUP BY` correctness tests over the nanosecond-precision timestamp types
+ * `TIMESTAMP_NTZ(p)` / `TIMESTAMP_LTZ(p)` (`p` in `[7, 9]`), part of the nanosecond timestamp
+ * preview (SPARK-56822). Grouping over these types executes today with no production change -- it
+ * rides on the nanos hashing and equality the types already implement (SPARK-57103 extended
+ * `Murmur3Hash` / `XxHash64` / `HiveHash` to hash the carrier's two fields, `epochMicros: Long`
+ * and `nanosWithinMicro: Short in [0, 999]`). There is no coverage organized around grouping by
+ * column type, so this per-type suite adds it and locks the regression.
+ *
+ * Headline risk: SUB-MICROSECOND key correctness, mirroring the join suite. The carrier is
+ * `TimestampNanosVal = (epochMicros: Long, nanosWithinMicro: Short in [0, 999])`. Every key in
+ * these tables shares the SAME epochMicros (1577836800000000, = 2020-01-01T00:00:00Z) and differs
+ * ONLY in nanosWithinMicro. So the micro-level path alone cannot tell the keys apart -- correct
+ * grouping MUST be driven by the full nanos value:
+ *   - two rows equal in epochMicros but DIFFERENT in nanosWithinMicro must NOT collapse into one
+ *     group,
+ *   - exact duplicates (equal incl. the sub-microsecond remainder) MUST collapse,
+ *   - NULL keys, unlike an equi-join, DO group together -- all NULL rows land in a single group.
+ * If the sub-microsecond remainder were ignored, the two distinct sub-microsecond keys here would
+ * fold into one, dropping the group count from 3 to 2 and failing these tests loudly.
+ *
+ * Precision-safety: all sub-microsecond remainders are multiples of 100ns (100 / 900), which are
+ * exact at every p in [7, 9] (`createDataFrame` floors nanosWithinMicro to (n/100)*100 at p=7 and
+ * (n/10)*10 at p=8). So the SAME inputs and the SAME expected results are valid verbatim at all
+ * three precisions, and the two distinct remainders never collide even at the coarsest p=7.
+ *
+ * Each test runs under whole-stage codegen on and off, so the same sub-microsecond grouping is
+ * proven on the nanos hash path in both codegen modes, for NTZ and LTZ. The mixed-precision test
+ * additionally pins that a `UNION` of two different nanos precisions widens the key to the higher
+ * precision (`findWiderDateTimeType`) and preserves the distinction.
+ *
+ * The nanosecond timestamp types are gated behind a preview flag enabled by default under tests
+ * (`Utils.isTesting`), so it is not set here. The session time zone is fixed so the
+ * `TIMESTAMP_LTZ` (`Instant`) values are deterministic. The two subclasses run every test with
+ * ANSI mode on and off.
+ */
+abstract class TimestampNanosGroupBySuiteBase extends QueryTest with SharedSparkSession {
+
+  override def sparkConf: SparkConf = super.sparkConf
+    .set(SQLConf.SESSION_LOCAL_TIMEZONE.key, "America/Los_Angeles")
+
+  // Whole-stage codegen on (CODEGEN_ONLY) vs off (NO_CODEGEN). The hash-aggregate exec is the same
+  // in both modes; only the WholeStageCodegenExec wrapper differs. Mirrors the join and functions
+  // suites.
+  protected val codegenModes: Seq[Seq[(String, String)]] = Seq(
+    Seq(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+      SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY"),
+    Seq(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
+      SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN"))
+
+  private def cgLabel(cgConf: Seq[(String, String)]): String =
+    if (cgConf.contains(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true")) "codegen on"
+    else "codegen off"
+
+  // ---- relation builders: key column "k" of the given nanos type + an Int value column "v" ----
+  //
+  // Five rows, all sharing epochMicros (2020-01-01T00:00:00Z) and differing only within the
+  // microsecond:
+  //   - two rows with the 100ns key (an exact duplicate): must collapse into one group,
+  //   - one row with the 900ns key: same microsecond, distinct sub-microsecond remainder,
+  //   - two NULL-key rows: group together into a single NULL group.
+  // The value column distinguishes the group signatures: sum(v) is 1+2=3 for the 100ns group,
+  // 3 for the 900ns group, and 4+5=9 for the NULL group, so the three groups are all identifiable
+  // by their (count, sum) pair even without inspecting the key.
+
+  private def ntzDF(p: Int): DataFrame =
+    spark.createDataFrame(
+      spark.sparkContext.parallelize(Seq(
+        Row(LocalDateTime.parse("2020-01-01T00:00:00.000000100"), 1),
+        Row(LocalDateTime.parse("2020-01-01T00:00:00.000000100"), 2),
+        Row(LocalDateTime.parse("2020-01-01T00:00:00.000000900"), 3),
+        Row(null, 4),
+        Row(null, 5))),
+      new StructType().add("k", TimestampNTZNanosType(p)).add("v", IntegerType))
+
+  private def ltzDF(p: Int): DataFrame =
+    spark.createDataFrame(
+      spark.sparkContext.parallelize(Seq(
+        Row(Instant.parse("2020-01-01T00:00:00.000000100Z"), 1),
+        Row(Instant.parse("2020-01-01T00:00:00.000000100Z"), 2),
+        Row(Instant.parse("2020-01-01T00:00:00.000000900Z"), 3),
+        Row(null, 4),
+        Row(null, 5))),
+      new StructType().add("k", TimestampLTZNanosType(p)).add("v", IntegerType))
+
+  // Each family as (label, builder(p), widerType(p)). GROUP BY is type-agnostic, so the family only
+  // picks the data builder and (for the mixed-precision widening assertion) the expected widened
+  // key type.
+  protected val families: Seq[(String, Int => DataFrame, Int => DataType)] = Seq(
+    ("NTZ", ntzDF, (p: Int) => TimestampNTZNanosType(p)),
+    ("LTZ", ltzDF, (p: Int) => TimestampLTZNanosType(p)))
+
+  // Mixed-precision UNION pairs; the key widens to max(pl, pr).
+  private val mixedPrecisionPairs: Seq[(Int, Int)] = Seq((7, 9), (7, 8), (8, 9))
+
+  for {
+    (family, builder, widerType) <- families
+    cgConf <- codegenModes
+  } {
+    // ========================================================================================
+    // GROUP BY a nanos key: exact duplicates collapse, but two keys sharing epochMicros and
+    // differing only within the microsecond stay in separate groups. Aggregates (count, sum) are
+    // computed per group. Expected three groups: 100ns (count 2, sum 3), 900ns (count 1, sum 3),
+    // NULL (count 2, sum 9). If the sub-microsecond remainder were ignored, 100ns and 900ns would
+    // merge into a single (count 3, sum 6) group and this checkAnswer would fail.
+    // ========================================================================================
+    test(s"$family nanos GROUP BY distinguishes the sub-microsecond remainder - " +
+      s"${cgLabel(cgConf)}") {
+      withSQLConf(cgConf: _*) {
+        Seq(7, 8, 9).foreach { p =>
+          val grouped = builder(p).groupBy("k")
+            .agg(count(lit(1)).as("c"), sum(col("v")).as("s"))
+          checkAnswer(grouped.select("c", "s"),
+            Seq(Row(2L, 3L), Row(1L, 3L), Row(2L, 9L)))
+        }
+      }
+    }
+
+    // ========================================================================================
+    // GROUP BY over a mixed-precision key: a UNION of the same data at two different nanos
+    // precisions widens the key to max(pl, pr) (findWiderDateTimeType). Grouping the widened column
+    // still separates the two sub-microsecond keys (all remainders are multiples of 100ns, exact at
+    // every p in [7, 9]). The union doubles every row, so counts/sums double: 100ns (4, 6),
+    // 900ns (2, 6), NULL (4, 18).
+    // ========================================================================================
+    test(s"$family nanos GROUP BY across mixed precisions widens key - ${cgLabel(cgConf)}") {
+      withSQLConf(cgConf: _*) {
+        mixedPrecisionPairs.foreach { case (pl, pr) =>
+          val union = builder(pl).union(builder(pr))
+          assert(union.schema("k").dataType == widerType(math.max(pl, pr)),
+            s"expected key widened to ${widerType(math.max(pl, pr))} for ($pl, $pr)")
+          val grouped = union.groupBy("k")
+            .agg(count(lit(1)).as("c"), sum(col("v")).as("s"))
+          checkAnswer(grouped.select("c", "s"),
+            Seq(Row(4L, 6L), Row(2L, 6L), Row(4L, 18L)))
+        }
+      }
+    }
+
+    // ========================================================================================
+    // NULL grouping: unlike an equi-join (where NULL never matches NULL), GROUP BY treats NULL as a
+    // single group. Both NULL-key rows land in one group (count 2).
+    // ========================================================================================
+    test(s"$family nanos GROUP BY collapses NULL keys into a single group - ${cgLabel(cgConf)}") {
+      withSQLConf(cgConf: _*) {
+        Seq(7, 8, 9).foreach { p =>
+          val nullGroup = builder(p).groupBy("k").count().filter(col("k").isNull)
+          checkAnswer(nullGroup.select("count"), Seq(Row(2L)))
+        }
+      }
+    }
+  }
+}
+
+// Runs the nanosecond timestamp GROUP BY tests with ANSI mode enabled explicitly.
+class TimestampNanosGroupByAnsiOnSuite extends TimestampNanosGroupBySuiteBase {
+  override def sparkConf: SparkConf = super.sparkConf.set(SQLConf.ANSI_ENABLED.key, "true")
+}
+
+// Runs the nanosecond timestamp GROUP BY tests with ANSI mode disabled explicitly.
+class TimestampNanosGroupByAnsiOffSuite extends TimestampNanosGroupBySuiteBase {
+  override def sparkConf: SparkConf = super.sparkConf.set(SQLConf.ANSI_ENABLED.key, "false")
+}