[SPARK-40812][CONNECT] Add Deduplicate to Connect proto and DSL

connector/connect/src/main/protobuf/spark/connect/relations.proto

@@ -43,6 +43,7 @@ message Relation {
     LocalRelation local_relation = 11;
     Sample sample = 12;
     Offset offset = 13;
+    Deduplicate deduplicate = 14;
 
     Unknown unknown = 999;
   }
@@ -181,6 +182,14 @@ message Sort {
   }
 }
 
+// Relation of type [[Deduplicate]] which have duplicate rows removed, could consider either only
+// the subset of columns or all the columns.
+message Deduplicate {
+  Relation input = 1;
+  repeated string column_names = 2;
+  bool all_columns_as_keys = 3;
+}
+
 message LocalRelation {
   repeated Expression.QualifiedAttribute attributes = 1;
   // TODO: support local data.

connector/connect/src/main/scala/org/apache/spark/sql/connect/dsl/package.scala

@@ -215,6 +215,26 @@ package object dsl {
           .build()
       }
 
+      def deduplicate(colNames: Seq[String]): proto.Relation =
+        proto.Relation
+          .newBuilder()
+          .setDeduplicate(
+            proto.Deduplicate
+              .newBuilder()
+              .setInput(logicalPlan)
+              .addAllColumnNames(colNames.asJava))
+          .build()
+
+      def distinct(): proto.Relation =
+        proto.Relation
+          .newBuilder()
+          .setDeduplicate(
+            proto.Deduplicate
+              .newBuilder()
+              .setInput(logicalPlan)
+              .setAllColumnsAsKeys(true))
+          .build()
+
       def join(
           otherPlan: proto.Relation,
           joinType: JoinType = JoinType.JOIN_TYPE_INNER,

connector/connect/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala

@@ -27,8 +27,9 @@ import org.apache.spark.sql.catalyst.expressions
 import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, Expression}
 import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
 import org.apache.spark.sql.catalyst.plans.{logical, FullOuter, Inner, JoinType, LeftAnti, LeftOuter, LeftSemi, RightOuter}
-import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Sample, SubqueryAlias}
+import org.apache.spark.sql.catalyst.plans.logical.{Deduplicate, LogicalPlan, Sample, SubqueryAlias}
 import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap
+import org.apache.spark.sql.execution.QueryExecution
 import org.apache.spark.sql.types._
 
 final case class InvalidPlanInput(
@@ -60,6 +61,7 @@ class SparkConnectPlanner(plan: proto.Relation, session: SparkSession) {
       case proto.Relation.RelTypeCase.OFFSET => transformOffset(rel.getOffset)
       case proto.Relation.RelTypeCase.JOIN => transformJoin(rel.getJoin)
       case proto.Relation.RelTypeCase.UNION => transformUnion(rel.getUnion)
+      case proto.Relation.RelTypeCase.DEDUPLICATE => transformDeduplicate(rel.getDeduplicate)
       case proto.Relation.RelTypeCase.SORT => transformSort(rel.getSort)
       case proto.Relation.RelTypeCase.AGGREGATE => transformAggregate(rel.getAggregate)
       case proto.Relation.RelTypeCase.SQL => transformSql(rel.getSql)
@@ -91,6 +93,37 @@ class SparkConnectPlanner(plan: proto.Relation, session: SparkSession) {
       transformRelation(rel.getInput))
   }
 
+  private def transformDeduplicate(rel: proto.Deduplicate): LogicalPlan = {
+    if (!rel.hasInput) {
+      throw InvalidPlanInput("Deduplicate needs a plan input")
+    }
+    if (rel.getAllColumnsAsKeys && rel.getColumnNamesCount > 0) {
+      throw InvalidPlanInput("Cannot deduplicate on both all columns and a subset of columns")
+    }
+    if (!rel.getAllColumnsAsKeys && rel.getColumnNamesCount == 0) {
+      throw InvalidPlanInput(
+        "Deduplicate requires to either deduplicate on all columns or a subset of columns")
+    }
+    val queryExecution = new QueryExecution(session, transformRelation(rel.getInput))
+    val resolver = session.sessionState.analyzer.resolver
+    val allColumns = queryExecution.analyzed.output
+    if (rel.getAllColumnsAsKeys) {
+      Deduplicate(allColumns, queryExecution.analyzed)
+    } else {
+      val toGroupColumnNames = rel.getColumnNamesList.asScala.toSeq
+      val groupCols = toGroupColumnNames.flatMap { (colName: String) =>
+        // It is possibly there are more than one columns with the same name,
+        // so we call filter instead of find.
+        val cols = allColumns.filter(col => resolver(col.name, colName))
+        if (cols.isEmpty) {
+          throw InvalidPlanInput(s"Invalid deduplicate column ${colName}")
+        }
+        cols
+      }
+      Deduplicate(groupCols, queryExecution.analyzed)
+    }
+  }
+
   private def transformLocalRelation(rel: proto.LocalRelation): LogicalPlan = {
     val attributes = rel.getAttributesList.asScala.map(transformAttribute(_)).toSeq
     new org.apache.spark.sql.catalyst.plans.logical.LocalRelation(attributes)

connector/connect/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectDeduplicateSuite.scala

@@ -0,0 +1,68 @@
+/*
+ * 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.connect.planner
+
+import org.apache.spark.sql.{Dataset, Row, SparkSession}
+import org.apache.spark.sql.catalyst.expressions.AttributeReference
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+
+/**
+ * [[SparkConnectPlanTestWithSparkSession]] contains a SparkSession for the connect planner.
+ *
+ * It is not recommended to use Catalyst DSL along with this trait because `SharedSparkSession`
+ * has also defined implicits over Catalyst LogicalPlan which will cause ambiguity with the
+ * implicits defined in Catalyst DSL.
+ */
+trait SparkConnectPlanTestWithSparkSession extends SharedSparkSession with SparkConnectPlanTest {
+  override def getSession(): SparkSession = spark
+}
+
+class SparkConnectDeduplicateSuite extends SparkConnectPlanTestWithSparkSession {
+  lazy val connectTestRelation = createLocalRelationProto(
+    Seq(
+      AttributeReference("id", IntegerType)(),
+      AttributeReference("key", StringType)(),
+      AttributeReference("value", StringType)()))
+
+  lazy val sparkTestRelation = {
+    spark.createDataFrame(
+      new java.util.ArrayList[Row](),
+      StructType(
+        Seq(
+          StructField("id", IntegerType),
+          StructField("key", StringType),
+          StructField("value", StringType))))
+  }
+
+  test("Test basic deduplicate") {
+    val connectPlan = {
+      import org.apache.spark.sql.connect.dsl.plans._
+      Dataset.ofRows(spark, transform(connectTestRelation.distinct()))
+    }
+
+    val sparkPlan = sparkTestRelation.distinct()
+    comparePlans(connectPlan.queryExecution.analyzed, sparkPlan.queryExecution.analyzed, false)
+
+    val connectPlan2 = {
+      import org.apache.spark.sql.connect.dsl.plans._
+      Dataset.ofRows(spark, transform(connectTestRelation.deduplicate(Seq("key", "value"))))
+    }
+    val sparkPlan2 = sparkTestRelation.dropDuplicates(Seq("key", "value"))
+    comparePlans(connectPlan2.queryExecution.analyzed, sparkPlan2.queryExecution.analyzed, false)
+  }
+}

connector/connect/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectPlannerSuite.scala

@@ -31,8 +31,11 @@ import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
  * test cases.
  */
 trait SparkConnectPlanTest {
+
+  def getSession(): SparkSession = None.orNull
+
   def transform(rel: proto.Relation): LogicalPlan = {
-    new SparkConnectPlanner(rel, None.orNull).transform()
+    new SparkConnectPlanner(rel, getSession()).transform()
   }
 
   def readRel: proto.Relation =
@@ -72,8 +75,6 @@ trait SparkConnectPlanTest {
  */
 class SparkConnectPlannerSuite extends SparkFunSuite with SparkConnectPlanTest {
 
-  protected var spark: SparkSession = null
-
   test("Simple Limit") {
     assertThrows[IndexOutOfBoundsException] {
       new SparkConnectPlanner(
@@ -266,4 +267,26 @@ class SparkConnectPlannerSuite extends SparkFunSuite with SparkConnectPlanTest {
           .build()))
     assert(e.getMessage.contains("DataSource requires a format"))
   }
+
+  test("Test invalid deduplicate") {
+    val deduplicate = proto.Deduplicate
+      .newBuilder()
+      .setInput(readRel)
+      .setAllColumnsAsKeys(true)
+      .addColumnNames("test")
+
+    val e = intercept[InvalidPlanInput] {
+      transform(proto.Relation.newBuilder.setDeduplicate(deduplicate).build())
+    }
+    assert(
+      e.getMessage.contains("Cannot deduplicate on both all columns and a subset of columns"))
+
+    val deduplicate2 = proto.Deduplicate
+      .newBuilder()
+      .setInput(readRel)
+    val e2 = intercept[InvalidPlanInput] {
+      transform(proto.Relation.newBuilder.setDeduplicate(deduplicate2).build())
+    }
+    assert(e2.getMessage.contains("either deduplicate on all columns or a subset of columns"))
+  }
 }