From a39c36fcdaaa465b62e55012854db6c34b05aaf3 Mon Sep 17 00:00:00 2001
From: Rui Wang <rui.wang@databricks.com>
Date: Thu, 2 Mar 2023 14:13:24 -0800
Subject: [PATCH 1/2] [SPARK-42609][CONNECT] Add tests for grouping() and
 grouping_id() functions.

---
 .../spark/sql/PlanGenerationTestSuite.scala   |   6 ++
 .../grouping_and_grouping_id.explain          |   4 ++
 .../queries/grouping_and_grouping_id.json     |  57 ++++++++++++++++++
 .../grouping_and_grouping_id.proto.bin        | Bin 0 -> 138 bytes
 4 files changed, 67 insertions(+)
 create mode 100644 connector/connect/common/src/test/resources/query-tests/explain-results/grouping_and_grouping_id.explain
 create mode 100644 connector/connect/common/src/test/resources/query-tests/queries/grouping_and_grouping_id.json
 create mode 100644 connector/connect/common/src/test/resources/query-tests/queries/grouping_and_grouping_id.proto.bin

diff --git a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
index 80ca5b4362251..67dc92a747233 100644
--- a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
+++ b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
@@ -1966,6 +1966,12 @@ class PlanGenerationTestSuite
     simple.cube("a", "b").count()
   }
 
+  test("grouping and grouping_id") {
+    simple
+      .cube("a", "b")
+      .agg(fn.grouping("a"), fn.grouping("b"), fn.grouping_id("a", "b"))
+  }
+
   test("pivot") {
     simple.groupBy(Column("id")).pivot("a", Seq(1, 2, 3)).agg(functions.count(Column("b")))
   }
diff --git a/connector/connect/common/src/test/resources/query-tests/explain-results/grouping_and_grouping_id.explain b/connector/connect/common/src/test/resources/query-tests/explain-results/grouping_and_grouping_id.explain
new file mode 100644
index 0000000000000..3b7d6fb2b7072
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/explain-results/grouping_and_grouping_id.explain
@@ -0,0 +1,4 @@
+Aggregate [a#0, b#0, spark_grouping_id#0L], [a#0, b#0, cast((shiftright(spark_grouping_id#0L, 1) & 1) as tinyint) AS grouping(a)#0, cast((shiftright(spark_grouping_id#0L, 0) & 1) as tinyint) AS grouping(b)#0, spark_grouping_id#0L AS grouping_id(a, b)#0L]
++- Expand [[id#0L, a#0, b#0, a#0, b#0, 0], [id#0L, a#0, b#0, a#0, null, 1], [id#0L, a#0, b#0, null, b#0, 2], [id#0L, a#0, b#0, null, null, 3]], [id#0L, a#0, b#0, a#0, b#0, spark_grouping_id#0L]
+   +- Project [id#0L, a#0, b#0, a#0 AS a#0, b#0 AS b#0]
+      +- LocalRelation <empty>, [id#0L, a#0, b#0]
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/grouping_and_grouping_id.json b/connector/connect/common/src/test/resources/query-tests/queries/grouping_and_grouping_id.json
new file mode 100644
index 0000000000000..07bbd315a5fe9
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/queries/grouping_and_grouping_id.json
@@ -0,0 +1,57 @@
+{
+  "common": {
+    "planId": "1"
+  },
+  "aggregate": {
+    "input": {
+      "common": {
+        "planId": "0"
+      },
+      "localRelation": {
+        "schema": "struct\u003cid:bigint,a:int,b:double\u003e"
+      }
+    },
+    "groupType": "GROUP_TYPE_CUBE",
+    "groupingExpressions": [{
+      "unresolvedAttribute": {
+        "unparsedIdentifier": "a"
+      }
+    }, {
+      "unresolvedAttribute": {
+        "unparsedIdentifier": "b"
+      }
+    }],
+    "aggregateExpressions": [{
+      "unresolvedFunction": {
+        "functionName": "grouping",
+        "arguments": [{
+          "unresolvedAttribute": {
+            "unparsedIdentifier": "a"
+          }
+        }]
+      }
+    }, {
+      "unresolvedFunction": {
+        "functionName": "grouping",
+        "arguments": [{
+          "unresolvedAttribute": {
+            "unparsedIdentifier": "b"
+          }
+        }]
+      }
+    }, {
+      "unresolvedFunction": {
+        "functionName": "grouping_id",
+        "arguments": [{
+          "unresolvedAttribute": {
+            "unparsedIdentifier": "a"
+          }
+        }, {
+          "unresolvedAttribute": {
+            "unparsedIdentifier": "b"
+          }
+        }]
+      }
+    }]
+  }
+}
\ No newline at end of file
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/grouping_and_grouping_id.proto.bin b/connector/connect/common/src/test/resources/query-tests/queries/grouping_and_grouping_id.proto.bin
new file mode 100644
index 0000000000000000000000000000000000000000..88b3f0593132809e30279854b1a3c60cf07738ad
GIT binary patch
literal 138
zcmd;L5@7UdX5`Y~ViI78QW8=qE-5NaF0si>u}aEJ&&(^)NwflyNmeQOrAaxdb^^>&
xtU}COjEN8`Nl92rkc%U|D8IBIGcR2TB8N>dNl8{pnu{BzGCnf}su)5g0RWwz9&i8v

literal 0
HcmV?d00001


From eafcb8fbafe8d486e8fd54ae97094c15671efdef Mon Sep 17 00:00:00 2001
From: Rui Wang <rui.wang@databricks.com>
Date: Thu, 2 Mar 2023 18:13:14 -0800
Subject: [PATCH 2/2] Empty-Commit