Add groupByRelationKey.

sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala

@@ -1725,6 +1725,30 @@ class Dataset[T] private[sql](
       withGroupingKey.newColumns)
   }
 
+  /**
+   * Returns a [[KeyValueGroupedDataset]] where the data is grouped by a row of given
+   * columns.
+   *
+   * @group typedrel
+   * @since 3.0.0
+   */
+  @scala.annotation.varargs
+  def groupByRelationKey(
+      col1: String,
+      cols: String*): KeyValueGroupedDataset[Row, Row] = {
+    val colNames: Seq[String] = col1 +: cols
+    val keyAttrs = colNames.map(colName => resolve(colName).toAttribute)
+    val keySchema = StructType.fromAttributes(keyAttrs)
+    val keyEncoder = RowEncoder(keySchema)
+    val valEncoder = RowEncoder(schema)
+    new KeyValueGroupedDataset(
+      keyEncoder,
+      valEncoder,
+      queryExecution,
+      logicalPlan.output,
+      keyAttrs)
+  }
+
   /**
    * (Java-specific)
    * Returns a [[KeyValueGroupedDataset]] where the data is grouped by the given key `func`.

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

@@ -2203,4 +2203,20 @@ class DataFrameSuite extends QueryTest with SharedSparkSession {
           |*(1) Range (0, 10, step=1, splits=2)""".stripMargin))
     }
   }
+
+  test("groupByRelationKey") {
+    val df1 = Seq((1, 2, 3), (2, 3, 4)).toDF("a", "b", "c")
+      .repartition($"a", $"b").sortWithinPartitions("a", "b")
+    val df2 = Seq((1, 2, 4), (2, 3, 5)).toDF("a", "b", "c")
+      .repartition($"a", $"b").sortWithinPartitions("a", "b")
+
+    val df3 = df1.groupByRelationKey("a", "b")
+      .cogroup(df2.groupByRelationKey("a", "b")) { case (key, data1, data2) =>
+         data1.zip(data2).map { p =>
+           p._1.getInt(2) + p._2.getInt(2)
+         }
+      }.toDF
+
+    checkAnswer(df3, Row(7) :: Row(9) :: Nil)
+  }
 }