[SPARK-40737][CONNECT] Add basic support for DataFrameWriter

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

@@ -17,6 +17,8 @@
 
 syntax = 'proto3';
 
+import "spark/connect/expressions.proto";
+import "spark/connect/relations.proto";
 import "spark/connect/types.proto";
 
 package spark.connect;
@@ -29,6 +31,7 @@ option java_package = "org.apache.spark.connect.proto";
 message Command {
   oneof command_type {
     CreateScalarFunction create_function = 1;
+    WriteOperation write_operation = 2;
   }
 }
 
@@ -62,3 +65,29 @@ message CreateScalarFunction {
     FUNCTION_LANGUAGE_SCALA = 3;
   }
 }
+
+// As writes are not directly handled during analysis and planning, they are modeled as commands.
+message WriteOperation {
+  Relation input = 1;
+  string format = 2;
+
+  oneof save_type {
+    string path = 3;
+    string table_name = 4;
+  }
+  string mode = 5;
+  repeated string sortColumnNames = 6;
+  repeated string partitionByColumns = 7;
+  BucketBy bucketBy = 8;
+  repeated WriteOptions options = 9;
+
+  message BucketBy {
+    repeated string columns = 1;
+    int32 bucketCount = 2;
+  }
+
+  message WriteOptions {
+    string key = 1;
+    string value = 2;
+  }
+}

connector/connect/src/main/scala/org/apache/spark/sql/connect/command/SparkConnectCommandPlanner.scala

@@ -24,10 +24,16 @@ import com.google.common.collect.{Lists, Maps}
 import org.apache.spark.annotation.{Since, Unstable}
 import org.apache.spark.api.python.{PythonEvalType, SimplePythonFunction}
 import org.apache.spark.connect.proto
-import org.apache.spark.sql.SparkSession
+import org.apache.spark.connect.proto.WriteOperation
+import org.apache.spark.sql.{Dataset, SparkSession}
+import org.apache.spark.sql.connect.planner.SparkConnectPlanner
 import org.apache.spark.sql.execution.python.UserDefinedPythonFunction
 import org.apache.spark.sql.types.StringType
 
+final case class InvalidCommandInput(
+    private val message: String = "",
+    private val cause: Throwable = None.orNull)
+    extends Exception(message, cause)
 
 @Unstable
 @Since("3.4.0")
@@ -40,17 +46,19 @@ class SparkConnectCommandPlanner(session: SparkSession, command: proto.Command)
     command.getCommandTypeCase match {
       case proto.Command.CommandTypeCase.CREATE_FUNCTION =>
         handleCreateScalarFunction(command.getCreateFunction)
+      case proto.Command.CommandTypeCase.WRITE_OPERATION =>
+        handleWriteOperation(command.getWriteOperation)
       case _ => throw new UnsupportedOperationException(s"$command not supported.")
     }
   }
 
   /**
    * This is a helper function that registers a new Python function in the SparkSession.
    *
-   * Right now this function is very rudimentary and bare-bones just to showcase how it
-   * is possible to remotely serialize a Python function and execute it on the Spark cluster.
-   * If the Python version on the client and server diverge, the execution of the function that
-   * is serialized will most likely fail.
+   * Right now this function is very rudimentary and bare-bones just to showcase how it is
+   * possible to remotely serialize a Python function and execute it on the Spark cluster. If the
+   * Python version on the client and server diverge, the execution of the function that is
+   * serialized will most likely fail.
    *
    * @param cf
    */
@@ -74,4 +82,60 @@ class SparkConnectCommandPlanner(session: SparkSession, command: proto.Command)
     session.udf.registerPython(cf.getPartsList.asScala.head, udf)
   }
 
+  /**
+   * Transforms the write operation and executes it.
+   *
+   * The input write operation contains a reference to the input plan and transforms it to the
+   * corresponding logical plan. Afterwards, creates the DataFrameWriter and translates the
+   * parameters of the WriteOperation into the corresponding methods calls.
+   *
+   * @param writeOperation
+   */
+  def handleWriteOperation(writeOperation: WriteOperation): Unit = {
+    // Transform the input plan into the logical plan.
+    val planner = new SparkConnectPlanner(writeOperation.getInput, session)
+    val plan = planner.transform()
+    // And create a Dataset from the plan.
+    val dataset = Dataset.ofRows(session, logicalPlan = plan)
+
+    val w = dataset.write
+    if (writeOperation.getOptionsCount > 0) {
+      writeOperation.getOptionsList.asScala.foreach(x => w.option(x.getKey, x.getValue))
+    }
+
+    if (writeOperation.getSortColumnNamesCount > 0) {
+      val names = writeOperation.getSortColumnNamesList.asScala
+      w.sortBy(names.head, names.tail.toSeq: _*)
+    }
+
+    if (writeOperation.hasBucketBy) {
+      val op = writeOperation.getBucketBy
+      val cols = op.getColumnsList.asScala
+      if (op.getBucketCount <= 0) {
+        throw InvalidCommandInput(
+          s"BucketBy must specify a bucket count > 0, received ${op.getBucketCount} instead.")
+      }
+      w.bucketBy(op.getBucketCount, cols.head, cols.tail.toSeq: _*)
+    }
+
+    if (writeOperation.getPartitionByColumnsCount > 0) {
+      val names = writeOperation.getPartitionByColumnsList.asScala
+      w.partitionBy(names.toSeq: _*)
+    }
+
+    if (writeOperation.getFormat != null) {
+      w.format(writeOperation.getFormat)
+    }
+
+    writeOperation.getSaveTypeCase match {
+      case proto.WriteOperation.SaveTypeCase.PATH => w.save(writeOperation.getPath)
+      case proto.WriteOperation.SaveTypeCase.TABLE_NAME =>
+        w.saveAsTable(writeOperation.getTableName)
+      case _ =>
+        throw new UnsupportedOperationException(
+          s"WriteOperation:SaveTypeCase not supported "
+            + "${writeOperation.getSaveTypeCase.getNumber}")
+    }
+  }
+
 }

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

@@ -33,9 +33,11 @@ package object dsl {
       val identifier = CatalystSqlParser.parseMultipartIdentifier(s)
 
       def protoAttr: proto.Expression =
-        proto.Expression.newBuilder()
+        proto.Expression
+          .newBuilder()
           .setUnresolvedAttribute(
-            proto.Expression.UnresolvedAttribute.newBuilder()
+            proto.Expression.UnresolvedAttribute
+              .newBuilder()
               .addAllParts(identifier.asJava)
               .build())
           .build()
@@ -47,15 +49,53 @@ package object dsl {
     }
   }
 
+  object commands { // scalastyle:ignore
+    implicit class DslCommands(val logicalPlan: proto.Relation) {
+      def write(
+          format: Option[String] = None,
+          path: Option[String] = None,
+          tableName: Option[String] = None,
+          mode: Option[String] = None,
+          sortByColumns: Seq[String] = Seq.empty,
+          partitionByCols: Seq[String] = Seq.empty,
+          bucketByCols: Seq[String] = Seq.empty,
+          numBuckets: Option[Int] = None): proto.Command = {
+        val writeOp = proto.WriteOperation.newBuilder()
+        format.foreach(writeOp.setFormat(_))
+        mode.foreach(writeOp.setMode(_))
+
+        if (tableName.nonEmpty) {
+          tableName.foreach(writeOp.setTableName(_))
+        } else {
+          path.foreach(writeOp.setPath(_))
+        }
+        sortByColumns.foreach(writeOp.addSortColumnNames(_))
+        partitionByCols.foreach(writeOp.addPartitionByColumns(_))
+
+        if (numBuckets.nonEmpty && bucketByCols.nonEmpty) {
+          val op = proto.WriteOperation.BucketBy.newBuilder()
+          numBuckets.foreach(op.setBucketCount(_))
+          bucketByCols.foreach(op.addColumns(_))
+          writeOp.setBucketBy(op.build())
+        }
+        writeOp.setInput(logicalPlan)
+        proto.Command.newBuilder().setWriteOperation(writeOp.build()).build()
+      }
+    }
+  }
+
   object plans { // scalastyle:ignore
     implicit class DslLogicalPlan(val logicalPlan: proto.Relation) {
       def select(exprs: proto.Expression*): proto.Relation = {
-        proto.Relation.newBuilder().setProject(
-          proto.Project.newBuilder()
-            .setInput(logicalPlan)
-            .addAllExpressions(exprs.toIterable.asJava)
-            .build()
-        ).build()
+        proto.Relation
+          .newBuilder()
+          .setProject(
+            proto.Project
+              .newBuilder()
+              .setInput(logicalPlan)
+              .addAllExpressions(exprs.toIterable.asJava)
+              .build())
+          .build()
       }
 
       def join(

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

@@ -26,6 +26,7 @@ import org.apache.spark.sql.catalyst.analysis.{UnresolvedAlias, UnresolvedAttrib
 import org.apache.spark.sql.catalyst.expressions
 import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, Expression}
 import org.apache.spark.sql.catalyst.plans.{logical, FullOuter, Inner, JoinType, LeftAnti, LeftOuter, LeftSemi, RightOuter}
+import org.apache.spark.sql.catalyst.plans.logical
 import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, SubqueryAlias}
 import org.apache.spark.sql.types._
 
@@ -60,7 +61,8 @@ class SparkConnectPlanner(plan: proto.Relation, session: SparkSession) {
       case proto.Relation.RelTypeCase.SORT => transformSort(rel.getSort)
       case proto.Relation.RelTypeCase.AGGREGATE => transformAggregate(rel.getAggregate)
       case proto.Relation.RelTypeCase.SQL => transformSql(rel.getSql)
-      case proto.Relation.RelTypeCase.LOCAL_RELATION => transformLocalRelation(rel.getLocalRelation)
+      case proto.Relation.RelTypeCase.LOCAL_RELATION =>
+        transformLocalRelation(rel.getLocalRelation)
       case proto.Relation.RelTypeCase.RELTYPE_NOT_SET =>
         throw new IndexOutOfBoundsException("Expected Relation to be set, but is empty.")
       case _ => throw InvalidPlanInput(s"${rel.getUnknown} not supported.")
@@ -109,10 +111,10 @@ class SparkConnectPlanner(plan: proto.Relation, session: SparkSession) {
     // TODO: support the target field for *.
     val projection =
       if (rel.getExpressionsCount == 1 && rel.getExpressions(0).hasUnresolvedStar) {
-      Seq(UnresolvedStar(Option.empty))
-    } else {
-      rel.getExpressionsList.asScala.map(transformExpression).map(UnresolvedAlias(_))
-    }
+        Seq(UnresolvedStar(Option.empty))
+      } else {
+        rel.getExpressionsList.asScala.map(transformExpression).map(UnresolvedAlias(_))
+      }
     val project = logical.Project(projectList = projection.toSeq, child = baseRel)
     if (common.nonEmpty && common.get.getAlias.nonEmpty) {
       logical.SubqueryAlias(identifier = common.get.getAlias, child = project)
@@ -141,7 +143,7 @@ class SparkConnectPlanner(plan: proto.Relation, session: SparkSession) {
    * Transforms the protocol buffers literals into the appropriate Catalyst literal expression.
    *
    * TODO(SPARK-40533): Missing support for Instant, BigDecimal, LocalDate, LocalTimestamp,
-   *   Duration, Period.
+   * Duration, Period.
    * @param lit
    * @return
    *   Expression
@@ -167,9 +169,10 @@ class SparkConnectPlanner(plan: proto.Relation, session: SparkSession) {
       // Days since UNIX epoch.
       case proto.Expression.Literal.LiteralTypeCase.DATE =>
         expressions.Literal(lit.getDate, DateType)
-      case _ => throw InvalidPlanInput(
-        s"Unsupported Literal Type: ${lit.getLiteralTypeCase.getNumber}" +
-          s"(${lit.getLiteralTypeCase.name})")
+      case _ =>
+        throw InvalidPlanInput(
+          s"Unsupported Literal Type: ${lit.getLiteralTypeCase.getNumber}" +
+            s"(${lit.getLiteralTypeCase.name})")
     }
   }
 
@@ -188,7 +191,8 @@ class SparkConnectPlanner(plan: proto.Relation, session: SparkSession) {
    *
    * TODO(SPARK-40546) We need to homogenize the function names for binary operators.
    *
-   * @param fun Proto representation of the function call.
+   * @param fun
+   *   Proto representation of the function call.
    * @return
    */
   private def transformScalarFunction(fun: proto.Expression.UnresolvedFunction): Expression = {