[SPARK-41169][CONNECT][PYTHON] Implement DataFrame.drop

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

@@ -50,6 +50,7 @@ message Relation {
     RenameColumnsBySameLengthNames rename_columns_by_same_length_names = 18;
     RenameColumnsByNameToNameMap rename_columns_by_name_to_name_map = 19;
     ShowString show_string = 20;
+    Drop drop = 21;
 
     // NA functions
     NAFill fill_na = 90;
@@ -252,6 +253,19 @@ message Sort {
   }
 }
 
+
+// Drop specified columns.
+message Drop {
+  // (Required) The input relation.
+  Relation input = 1;
+
+  // (Required) columns to drop.
+  //
+  // Should contain at least 1 item.
+  repeated Expression cols = 2;
+}
+
+
 // Relation of type [[Deduplicate]] which have duplicate rows removed, could consider either only
 // the subset of columns or all the columns.
 message Deduplicate {

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

@@ -510,6 +510,28 @@ package object dsl {
           .build()
       }
 
+      def drop(columns: String*): Relation = {
+        assert(columns.nonEmpty)
+
+        val cols = columns.map(col =>
+          Expression.newBuilder
+            .setUnresolvedAttribute(
+              Expression.UnresolvedAttribute.newBuilder
+                .setUnparsedIdentifier(col)
+                .build())
+            .build())
+
+        Relation
+          .newBuilder()
+          .setDrop(
+            Drop
+              .newBuilder()
+              .setInput(logicalPlan)
+              .addAllCols(cols.asJava)
+              .build())
+          .build()
+      }
+
       def groupBy(groupingExprs: Expression*)(aggregateExprs: Expression*): Relation = {
         val agg = Aggregate.newBuilder()
         agg.setInput(logicalPlan)

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

@@ -25,7 +25,7 @@ import com.google.common.collect.{Lists, Maps}
 import org.apache.spark.api.python.{PythonEvalType, SimplePythonFunction}
 import org.apache.spark.connect.proto
 import org.apache.spark.connect.proto.WriteOperation
-import org.apache.spark.sql.{Dataset, SparkSession}
+import org.apache.spark.sql.{Column, Dataset, SparkSession}
 import org.apache.spark.sql.catalyst.AliasIdentifier
 import org.apache.spark.sql.catalyst.analysis.{GlobalTempView, LocalTempView, MultiAlias, UnresolvedAlias, UnresolvedAttribute, UnresolvedFunction, UnresolvedRelation, UnresolvedStar}
 import org.apache.spark.sql.catalyst.expressions
@@ -69,6 +69,7 @@ class SparkConnectPlanner(session: SparkSession) {
       case proto.Relation.RelTypeCase.DEDUPLICATE => transformDeduplicate(rel.getDeduplicate)
       case proto.Relation.RelTypeCase.SET_OP => transformSetOperation(rel.getSetOp)
       case proto.Relation.RelTypeCase.SORT => transformSort(rel.getSort)
+      case proto.Relation.RelTypeCase.DROP => transformDrop(rel.getDrop)
       case proto.Relation.RelTypeCase.AGGREGATE => transformAggregate(rel.getAggregate)
       case proto.Relation.RelTypeCase.SQL => transformSql(rel.getSql)
       case proto.Relation.RelTypeCase.LOCAL_RELATION =>
@@ -523,6 +524,19 @@ class SparkConnectPlanner(session: SparkSession) {
       sameOrderExpressions = Seq.empty)
   }
 
+  private def transformDrop(rel: proto.Drop): LogicalPlan = {
+    assert(rel.getColsCount > 0, s"cols must contains at least 1 item!")
+
+    val cols = rel.getColsList.asScala.toArray.map { expr =>
+      Column(transformExpression(expr))
+    }
+
+    Dataset
+      .ofRows(session, transformRelation(rel.getInput))
+      .drop(cols.head, cols.tail: _*)
+      .logicalPlan
+  }
+
   private def transformAggregate(rel: proto.Aggregate): LogicalPlan = {
     assert(rel.hasInput)
 

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

@@ -148,6 +148,23 @@ class SparkConnectProtoSuite extends PlanTest with SparkConnectPlanTest {
     comparePlans(connectPlan2, sparkPlan2)
   }
 
+  test("SPARK-41169: Test drop") {
+    // single column
+    val connectPlan = connectTestRelation.drop("id")
+    val sparkPlan = sparkTestRelation.drop("id")
+    comparePlans(connectPlan, sparkPlan)
+
+    // all columns
+    val connectPlan2 = connectTestRelation.drop("id", "name")
+    val sparkPlan2 = sparkTestRelation.drop("id", "name")
+    comparePlans(connectPlan2, sparkPlan2)
+
+    // non-existing column
+    val connectPlan3 = connectTestRelation.drop("id2", "name")
+    val sparkPlan3 = sparkTestRelation.drop("id2", "name")
+    comparePlans(connectPlan3, sparkPlan3)
+  }
+
   test("SPARK-40809: column alias") {
     // Simple Test.
     val connectPlan = connectTestRelation.select("id".protoAttr.as("id2"))

python/pyspark/sql/connect/dataframe.py

@@ -257,10 +257,21 @@ def distinct(self) -> "DataFrame":
         )
 
     def drop(self, *cols: "ColumnOrString") -> "DataFrame":
-        all_cols = self.columns
-        dropped = set([c.name() if isinstance(c, Column) else self[c].name() for c in cols])
-        dropped_cols = filter(lambda x: x in dropped, all_cols)
-        return DataFrame.withPlan(plan.Project(self._plan, *dropped_cols), session=self._session)
+        _cols = list(cols)
+        if any(not isinstance(c, (str, Column)) for c in _cols):
+            raise TypeError(
+                f"'cols' must contains strings or Columns, but got {type(cols).__name__}"
+            )
+        if len(_cols) == 0:
+            raise ValueError("'cols' must be non-empty")
+
+        return DataFrame.withPlan(
+            plan.Drop(
+                child=self._plan,
+                columns=_cols,
+            ),
+            session=self._session,
+        )
 
     def filter(self, condition: Expression) -> "DataFrame":
         return DataFrame.withPlan(

python/pyspark/sql/connect/plan.py

@@ -464,6 +464,49 @@ def _repr_html_(self) -> str:
         """
 
 
+class Drop(LogicalPlan):
+    def __init__(
+        self,
+        child: Optional["LogicalPlan"],
+        columns: List[Union[Column, str]],
+    ) -> None:
+        super().__init__(child)
+        assert len(columns) > 0 and all(isinstance(c, (Column, str)) for c in columns)
+        self.columns = columns
+
+    def _convert_to_expr(
+        self, col: Union[Column, str], session: "RemoteSparkSession"
+    ) -> proto.Expression:
+        expr = proto.Expression()
+        if isinstance(col, Column):
+            expr.CopyFrom(col.to_plan(session))
+        else:
+            expr.CopyFrom(self.unresolved_attr(col))
+        return expr
+
+    def plan(self, session: "RemoteSparkSession") -> proto.Relation:
+        assert self._child is not None
+        plan = proto.Relation()
+        plan.drop.input.CopyFrom(self._child.plan(session))
+        plan.drop.cols.extend([self._convert_to_expr(c, session) for c in self.columns])
+        return plan
+
+    def print(self, indent: int = 0) -> str:
+        c_buf = self._child.print(indent + LogicalPlan.INDENT) if self._child else ""
+        return f"{' ' * indent}<Drop columns={self.columns}>\n{c_buf}"
+
+    def _repr_html_(self) -> str:
+        return f"""
+        <ul>
+            <li>
+                <b>Drop</b><br />
+                columns: {self.columns} <br />
+                {self._child_repr_()}
+            </li>
+        </uL>
+        """
+
+
 class Sample(LogicalPlan):
     def __init__(
         self,

python/pyspark/sql/connect/proto/relations_pb2.pyi

@@ -79,6 +79,7 @@ class Relation(google.protobuf.message.Message):
     RENAME_COLUMNS_BY_SAME_LENGTH_NAMES_FIELD_NUMBER: builtins.int
     RENAME_COLUMNS_BY_NAME_TO_NAME_MAP_FIELD_NUMBER: builtins.int
     SHOW_STRING_FIELD_NUMBER: builtins.int
+    DROP_FIELD_NUMBER: builtins.int
     FILL_NA_FIELD_NUMBER: builtins.int
     SUMMARY_FIELD_NUMBER: builtins.int
     CROSSTAB_FIELD_NUMBER: builtins.int
@@ -124,6 +125,8 @@ class Relation(google.protobuf.message.Message):
     @property
     def show_string(self) -> global___ShowString: ...
     @property
+    def drop(self) -> global___Drop: ...
+    @property
     def fill_na(self) -> global___NAFill:
         """NA functions"""
     @property
@@ -156,6 +159,7 @@ class Relation(google.protobuf.message.Message):
         rename_columns_by_same_length_names: global___RenameColumnsBySameLengthNames | None = ...,
         rename_columns_by_name_to_name_map: global___RenameColumnsByNameToNameMap | None = ...,
         show_string: global___ShowString | None = ...,
+        drop: global___Drop | None = ...,
         fill_na: global___NAFill | None = ...,
         summary: global___StatSummary | None = ...,
         crosstab: global___StatCrosstab | None = ...,
@@ -172,6 +176,8 @@ class Relation(google.protobuf.message.Message):
             b"crosstab",
             "deduplicate",
             b"deduplicate",
+            "drop",
+            b"drop",
             "fill_na",
             b"fill_na",
             "filter",
@@ -227,6 +233,8 @@ class Relation(google.protobuf.message.Message):
             b"crosstab",
             "deduplicate",
             b"deduplicate",
+            "drop",
+            b"drop",
             "fill_na",
             b"fill_na",
             "filter",
@@ -293,6 +301,7 @@ class Relation(google.protobuf.message.Message):
         "rename_columns_by_same_length_names",
         "rename_columns_by_name_to_name_map",
         "show_string",
+        "drop",
         "fill_na",
         "summary",
         "crosstab",
@@ -961,6 +970,42 @@ class Sort(google.protobuf.message.Message):
 
 global___Sort = Sort
 
+class Drop(google.protobuf.message.Message):
+    """Drop specified columns."""
+
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+    INPUT_FIELD_NUMBER: builtins.int
+    COLS_FIELD_NUMBER: builtins.int
+    @property
+    def input(self) -> global___Relation:
+        """(Required) The input relation."""
+    @property
+    def cols(
+        self,
+    ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[
+        pyspark.sql.connect.proto.expressions_pb2.Expression
+    ]:
+        """(Required) columns to drop.
+
+        Should contain at least 1 item.
+        """
+    def __init__(
+        self,
+        *,
+        input: global___Relation | None = ...,
+        cols: collections.abc.Iterable[pyspark.sql.connect.proto.expressions_pb2.Expression]
+        | None = ...,
+    ) -> None: ...
+    def HasField(
+        self, field_name: typing_extensions.Literal["input", b"input"]
+    ) -> builtins.bool: ...
+    def ClearField(
+        self, field_name: typing_extensions.Literal["cols", b"cols", "input", b"input"]
+    ) -> None: ...
+
+global___Drop = Drop
+
 class Deduplicate(google.protobuf.message.Message):
     """Relation of type [[Deduplicate]] which have duplicate rows removed, could consider either only
     the subset of columns or all the columns.

python/pyspark/sql/tests/connect/test_connect_basic.py

@@ -188,6 +188,33 @@ def test_take(self) -> None:
         df2 = self.connect.read.table(self.tbl_name_empty)
         self.assertEqual(0, len(df2.take(5)))
 
+    def test_drop(self):
+        # SPARK-41169: test drop
+        query = """
+            SELECT * FROM VALUES
+            (false, 1, NULL), (false, NULL, 2), (NULL, 3, 3)
+            AS tab(a, b, c)
+            """
+
+        cdf = self.connect.sql(query)
+        sdf = self.spark.sql(query)
+        self.assert_eq(
+            cdf.drop("a").toPandas(),
+            sdf.drop("a").toPandas(),
+        )
+        self.assert_eq(
+            cdf.drop("a", "b").toPandas(),
+            sdf.drop("a", "b").toPandas(),
+        )
+        self.assert_eq(
+            cdf.drop("a", "x").toPandas(),
+            sdf.drop("a", "x").toPandas(),
+        )
+        self.assert_eq(
+            cdf.drop(cdf.a, cdf.x).toPandas(),
+            sdf.drop("a", "x").toPandas(),
+        )
+
     def test_subquery_alias(self) -> None:
         # SPARK-40938: test subquery alias.
         plan_text = (

python/pyspark/sql/tests/connect/test_connect_plan_only.py

@@ -181,6 +181,22 @@ def test_sort(self):
         )
         self.assertEqual(plan.root.sort.is_global, False)
 
+    def test_drop(self):
+        # SPARK-41169: test drop
+        df = self.connect.readTable(table_name=self.tbl_name)
+
+        plan = df.filter(df.col_name > 3).drop("col_a", "col_b")._plan.to_proto(self.connect)
+        self.assertEqual(
+            [f.unresolved_attribute.unparsed_identifier for f in plan.root.drop.cols],
+            ["col_a", "col_b"],
+        )
+
+        plan = df.filter(df.col_name > 3).drop(df.col_x, "col_b")._plan.to_proto(self.connect)
+        self.assertEqual(
+            [f.unresolved_attribute.unparsed_identifier for f in plan.root.drop.cols],
+            ["col_x", "col_b"],
+        )
+
     def test_deduplicate(self):
         df = self.connect.readTable(table_name=self.tbl_name)