Add UPDATE support for DataSource V2

Author: xy_xin

docs/sql-keywords.md

@@ -280,6 +280,7 @@ Below is a list of all the keywords in Spark SQL.
   <tr><td>UNKNOWN</td><td>reserved</td><td>non-reserved</td><td>reserved</td></tr>
   <tr><td>UNLOCK</td><td>non-reserved</td><td>non-reserved</td><td>non-reserved</td></tr>
   <tr><td>UNSET</td><td>non-reserved</td><td>non-reserved</td><td>non-reserved</td></tr>
+  <tr><td>UPDATE</td><td>non-reserved</td><td>non-reserved</td><td>reserved</td></tr>
   <tr><td>USE</td><td>non-reserved</td><td>non-reserved</td><td>non-reserved</td></tr>
   <tr><td>USER</td><td>reserved</td><td>non-reserved</td><td>reserved</td></tr>
   <tr><td>USING</td><td>reserved</td><td>strict-non-reserved</td><td>reserved</td></tr>

sql/catalyst/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBase.g4

@@ -217,6 +217,7 @@ statement
     | SET .*?                                                          #setConfiguration
     | RESET                                                            #resetConfiguration
     | DELETE FROM multipartIdentifier tableAlias whereClause?          #deleteFromTable
+    | UPDATE multipartIdentifier tableAlias setClause whereClause?     #updateTable
     | unsupportedHiveNativeCommands .*?                                #failNativeCommand
     ;
 
@@ -476,6 +477,14 @@ selectClause
     : SELECT (hints+=hint)* setQuantifier? namedExpressionSeq
     ;
 
+setClause
+    : SET assign (',' assign)*
+    ;
+
+assign
+    : key=multipartIdentifier EQ value=expression
+    ;
+
 whereClause
     : WHERE booleanExpression
     ;
@@ -1085,6 +1094,7 @@ ansiNonReserved
     | UNCACHE
     | UNLOCK
     | UNSET
+    | UPDATE
     | USE
     | VALUES
     | VIEW
@@ -1355,6 +1365,7 @@ nonReserved
     | UNKNOWN
     | UNLOCK
     | UNSET
+    | UPDATE
     | USE
     | USER
     | VALUES
@@ -1622,6 +1633,7 @@ UNIQUE: 'UNIQUE';
 UNKNOWN: 'UNKNOWN';
 UNLOCK: 'UNLOCK';
 UNSET: 'UNSET';
+UPDATE: 'UPDATE';
 USE: 'USE';
 USER: 'USER';
 USING: 'USING';

sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/SupportsUpdate.java

@@ -0,0 +1,56 @@
+/*
+ * 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.connector.catalog;
+
+import java.util.Map;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.Expression;
+import org.apache.spark.sql.sources.Filter;
+
+/**
+ * A mix-in interface for {@link Table} update support. Data sources can implement this
+ * interface to provide the ability to update data that matches filter expressions
+ * with the given sets.
+ */
+@Experimental
+public interface SupportsUpdate {
+
+  /**
+   * Update data that matches filter expressions with the given sets for a data source table.
+   * <p>
+   * Rows will be updated with the given values iff all of the filter expressions match.
+   * That is, the expressions must be interpreted as a set of filters that are ANDed together.
+   * <p>
+   * Implementations may reject a update operation if the update isn't possible without significant
+   * effort or it cannot deal with the sets expression. For example, partitioned data sources may
+   * reject updates that do not filter by partition columns because the filter may require
+   * rewriting files. The update may also be rejected if the update requires a complex computation
+   * that the data source does not support.
+   * To reject an update implementations should throw {@link IllegalArgumentException} with a clear
+   * error message that identifies which expression was rejected.
+   *
+   * @param sets the fields to be updated and the corresponding updated values in form of
+   *             key-value pairs in a map. The value can be a literal, or a simple expression
+   *             like {{{ originalValue + 1 }}}.
+   * @param filters filter expressions, used to select rows to update when all expressions match
+   * @throws IllegalArgumentException If the update is rejected due to required effort
+   *                                  or unsupported update expression
+   */
+  void updateWhere(Map<String, Expression> sets, Filter[] filters);
+}

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/Analyzer.scala

@@ -1781,6 +1781,8 @@ class Analyzer(
         resolveSubQueries(q, q.children)
       case d: DeleteFromTable if d.childrenResolved =>
         resolveSubQueries(d, d.children)
+      case u: UpdateTable if u.childrenResolved =>
+        resolveSubQueries(u, u.children)
     }
   }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/CheckAnalysis.scala

@@ -587,19 +587,20 @@ trait CheckAnalysis extends PredicateHelper {
           // Only certain operators are allowed to host subquery expression containing
           // outer references.
           plan match {
-            case _: Filter | _: Aggregate | _: Project | _: DeleteFromTable => // Ok
+            case _: Filter | _: Aggregate | _: Project | _: DeleteFromTable | _: UpdateTable =>
+              // Ok
             case other => failAnalysis(
               "Correlated scalar sub-queries can only be used in a " +
-                s"Filter/Aggregate/Project: $plan")
+                s"Filter/Aggregate/Project/DeleteFromTable/UpdateTable: $plan")
           }
         }
 
       case inSubqueryOrExistsSubquery =>
         plan match {
-          case _: Filter | _: DeleteFromTable => // Ok
+          case _: Filter | _: DeleteFromTable | _: UpdateTable => // Ok
           case _ =>
             failAnalysis(s"IN/EXISTS predicate sub-queries can only be used in" +
-                s" Filter/DeleteFromTable: $plan")
+                s" Filter/DeleteFromTable/UpdateTable: $plan")
         }
     }
 

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala

@@ -36,7 +36,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.{First, Last}
 import org.apache.spark.sql.catalyst.parser.SqlBaseParser._
 import org.apache.spark.sql.catalyst.plans._
 import org.apache.spark.sql.catalyst.plans.logical._
-import org.apache.spark.sql.catalyst.plans.logical.sql.{AlterTableAddColumnsStatement, AlterTableAlterColumnStatement, AlterTableDropColumnsStatement, AlterTableRenameColumnStatement, AlterTableSetLocationStatement, AlterTableSetPropertiesStatement, AlterTableUnsetPropertiesStatement, AlterViewSetPropertiesStatement, AlterViewUnsetPropertiesStatement, CreateTableAsSelectStatement, CreateTableStatement, DeleteFromStatement, DescribeColumnStatement, DescribeTableStatement, DropTableStatement, DropViewStatement, InsertIntoStatement, QualifiedColType, ReplaceTableAsSelectStatement, ReplaceTableStatement, ShowNamespacesStatement, ShowTablesStatement}
+import org.apache.spark.sql.catalyst.plans.logical.sql.{AlterTableAddColumnsStatement, AlterTableAlterColumnStatement, AlterTableDropColumnsStatement, AlterTableRenameColumnStatement, AlterTableSetLocationStatement, AlterTableSetPropertiesStatement, AlterTableUnsetPropertiesStatement, AlterViewSetPropertiesStatement, AlterViewUnsetPropertiesStatement, CreateTableAsSelectStatement, CreateTableStatement, DeleteFromStatement, DescribeColumnStatement, DescribeTableStatement, DropTableStatement, DropViewStatement, InsertIntoStatement, QualifiedColType, ReplaceTableAsSelectStatement, ReplaceTableStatement, ShowNamespacesStatement, ShowTablesStatement, UpdateTableStatement}
 import org.apache.spark.sql.catalyst.util.DateTimeUtils.{getZoneId, stringToDate, stringToTimestamp}
 import org.apache.spark.sql.connector.expressions.{ApplyTransform, BucketTransform, DaysTransform, Expression => V2Expression, FieldReference, HoursTransform, IdentityTransform, LiteralValue, MonthsTransform, Transform, YearsTransform}
 import org.apache.spark.sql.internal.SQLConf
@@ -361,6 +361,36 @@ class AstBuilder(conf: SQLConf) extends SqlBaseBaseVisitor[AnyRef] with Logging
     DeleteFromStatement(tableId, tableAlias, predicate)
   }
 
+  override def visitUpdateTable(ctx: UpdateTableContext): LogicalPlan = withOrigin(ctx) {
+    val tableId = visitMultipartIdentifier(ctx.multipartIdentifier)
+    val tableAlias = if (ctx.tableAlias() != null) {
+      val ident = ctx.tableAlias().strictIdentifier()
+      // We do not allow columns aliases after table alias.
+      if (ctx.tableAlias().identifierList() != null) {
+        throw new ParseException("Columns aliases is not allowed in UPDATE.",
+          ctx.tableAlias().identifierList())
+      }
+      if (ident != null) Some(ident.getText) else None
+    } else {
+      None
+    }
+    val (attrs, values) = ctx.setClause().assign().asScala.map {
+      kv => visitMultipartIdentifier(kv.key) -> expression(kv.value)
+    }.unzip
+    val predicate = if (ctx.whereClause() != null) {
+      Some(expression(ctx.whereClause().booleanExpression()))
+    } else {
+      None
+    }
+
+    UpdateTableStatement(
+      tableId,
+      tableAlias,
+      attrs,
+      values,
+      predicate)
+  }
+
   /**
    * Create a partition specification map.
    */

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala

@@ -584,6 +584,21 @@ case class DeleteFromTable(
   override def children: Seq[LogicalPlan] = child :: Nil
 }
 
+/**
+ * Update the data of table that specified with the condition with the given updated values.
+ * NOTE: Considering nested fields, we let the type of `attrs` to be `Seq[Expression]` becuase
+ * nested fields may be resolved to objects of type like `GetStructField`. However, currently
+ * we only support top-level fields, in which case `attrs` is actually of type `Seq[Attribute]`.
+ */
+case class UpdateTable(
+    child: LogicalPlan,
+    attrs: Seq[Expression],
+    values: Seq[Expression],
+    condition: Option[Expression]) extends Command {
+
+  override def children: Seq[LogicalPlan] = child :: Nil
+}
+
 /**
  * Drop a table.
  */

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/sql/UpdateTableStatement.scala

@@ -0,0 +1,27 @@
+/*
+ * 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.catalyst.plans.logical.sql
+
+import org.apache.spark.sql.catalyst.expressions.Expression
+
+case class UpdateTableStatement(
+    tableName: Seq[String],
+    tableAlias: Option[String],
+    attrs: Seq[Seq[String]],
+    values: Seq[Expression],
+    condition: Option[Expression]) extends ParsedStatement

sql/catalyst/src/main/scala/org/apache/spark/sql/connector/expressions/expressions.scala

@@ -18,6 +18,7 @@
 package org.apache.spark.sql.connector.expressions
 
 import org.apache.spark.sql.catalyst
+import org.apache.spark.sql.catalyst.expressions.{Expression => CatalystExpression, Literal => CatalystLiteral}
 import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{DataType, IntegerType, StringType}
@@ -57,6 +58,22 @@ private[sql] object LogicalExpressions {
   def days(column: String): DaysTransform = DaysTransform(reference(column))
 
   def hours(column: String): HoursTransform = HoursTransform(reference(column))
+
+  /**
+   * Tries to translate a Catalyst [[CatalystExpression]] into data source [[Expression]].
+   *
+   * @return a `Some(catalyst.expression.Expression)` if the input [[CatalystExpression]]
+   *         is convertible, otherwise a `None`.
+   */
+  private[sql] def translateExpression(value: CatalystExpression): Option[Expression] = {
+    // TODO: currently we only support literal. Will add more public expression in future.
+    value match {
+      case l: CatalystLiteral =>
+        Some(LiteralValue(l.value, l.dataType))
+      case _ =>
+        None
+    }
+  }
 }
 
 /**

sql/catalyst/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2Implicits.scala

@@ -18,7 +18,7 @@
 package org.apache.spark.sql.execution.datasources.v2
 
 import org.apache.spark.sql.AnalysisException
-import org.apache.spark.sql.connector.catalog.{SupportsDelete, SupportsRead, SupportsWrite, Table, TableCapability}
+import org.apache.spark.sql.connector.catalog.{SupportsDelete, SupportsRead, SupportsUpdate, SupportsWrite, Table, TableCapability}
 
 object DataSourceV2Implicits {
   implicit class TableHelper(table: Table) {
@@ -49,6 +49,15 @@ object DataSourceV2Implicits {
       }
     }
 
+    def asUpdatable: SupportsUpdate = {
+      table match {
+        case support: SupportsUpdate =>
+          support
+        case _ =>
+          throw new AnalysisException(s"Table does not support updates: ${table.name}")
+      }
+    }
+
     def supports(capability: TableCapability): Boolean = table.capabilities.contains(capability)
 
     def supportsAny(capabilities: TableCapability*): Boolean = capabilities.exists(supports)