anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4476682894
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -237,4 +253,101 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Returns the table attributes that these assignments set to a new value,
skipping identity
+ * assignments.
+ */
+ private def collectUpdatedAttrs(assignments: Seq[Assignment]):
Seq[AttributeReference] = {
+ assignments.collect {
+ case Assignment(key: AttributeReference, value) if
!isIdentityAssignment(key, value) => key
Review Comment:
Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/LogicalWriteInfoImpl.scala:
##########
@@ -29,7 +29,8 @@ private[sql] case class LogicalWriteInfoImpl(
schema: StructType,
options: CaseInsensitiveStringMap,
override val rowIdSchema: Optional[StructType] =
Optional.empty[StructType],
- override val metadataSchema: Optional[StructType] =
Optional.empty[StructType])
+ override val metadataSchema: Optional[StructType] =
Optional.empty[StructType],
+ override val columnUpdateSchema: Optional[StructType] =
Optional.empty[StructType])
Review Comment:
Done in PR description.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java:
##########
@@ -37,4 +38,22 @@ public interface RowLevelOperationInfo {
* Returns the row-level SQL command (e.g. DELETE, UPDATE, MERGE).
*/
Command command();
+
+ /**
+ * Returns the columns being updated by this operation. Currently only
UPDATE populates it;
+ * other commands report an empty array.
+ * <p>
+ * A column is reported unless its assigned value is the column itself, so
identity assignments
+ * such as {@code SET a = a} are excluded. With the default CHAR and VARCHAR
handling, a
+ * self-assignment on a CHAR or VARCHAR column is still reported, because
Spark pads or
+ * length-checks the assigned value. Nested struct field updates are
reported at root-column
+ * granularity (e.g. {@code SET s.c1 = -1} returns {@code s}), even when the
field is assigned to
Review Comment:
Done.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,607 @@
+/*
+ * 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
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.{SparkRuntimeException,
SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.analysis.RewriteRowLevelCommand
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+import
org.apache.spark.sql.catalyst.util.RowDeltaUtils.ORIGINAL_ROW_ID_VALUE_PREFIX
+import org.apache.spark.sql.connector.catalog.{Delete, InMemoryBaseTable,
InMemoryRowLevelOperationTable, Reinsert}
+import org.apache.spark.sql.connector.expressions.{FieldReference,
NamedReference}
+import org.apache.spark.sql.connector.write.LogicalWriteInfo
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.errors.QueryCompilationErrors
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, updated rows contain only the
declared columns
+ * (LogicalWriteInfo.columnUpdateSchema()) rather than the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase
+ with ColumnUpdateTableSharedTests {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] =
{
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ override protected def reqAttrsTableProps(reqAttrs: String):
java.util.Map[String, String] = {
+ java.util.Map.of("column-update-req-attrs", reqAttrs)
+ }
+
+ override protected def fullWidthTableProps: java.util.Map[String, String] = {
+ java.util.Map.of("supports-deltas", "true")
+ }
+
+ override protected def scanMetadataColumns: String = ", _partition STRING"
+
+ private val splitOperationClass =
+
classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperation].getName
+
+ private val splitWithRowIdOperationClass =
+
classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithRowId].getName
+
+ private def splitReqAttrsTableProps(reqAttrs: String): java.util.Map[String,
String] = {
+ tablePropsWith(reqAttrsTableProps(reqAttrs), "split-updates" -> "true")
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData, reqAttrsTableProps(reqAttrs))
+ }
+
+ private val splitTableProps = tablePropsWith(extraTableProps,
"split-updates" -> "true")
+
+ private val splitPartitionRowIdTableProps =
+ java.util.Map.of("column-update-split-row-id", "_partition")
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String):
Unit = {
+ createAndInitTable(schemaString, jsonData, splitTableProps)
+ }
+
+ test("column-update: column update schema contains a single assigned
column") {
+ Seq(extraTableProps, splitTableProps).foreach { props =>
+ withTable(tableNameAsString) {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin,
+ props)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+ }
+ }
+
+ test("column-update: update handles metadata columns correctly") {
+ Seq(
+ "id IN (1, 100)",
+ "id IN (1, 20) AND dep IN (SELECT * FROM updated_dep)"
+ ).foreach { condition =>
+ withTable(tableNameAsString) {
+ withTempView("updated_dep") {
+ createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING",
+ """{ "pk": 1, "id": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "id": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+ import testImplicits._
+ Seq(Some("hr"),
Some("it")).toDF().createOrReplaceTempView("updated_dep")
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE $condition")
+
+ checkLastColumnUpdateWriteLog(
+ updateWriteLogEntry(id = 1, metadata = Row("hr", null), data =
Row(1, "hr", -1)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, -1, 100, "hr") :: Row(2, 2, 200, "software") :: Row(3, 3,
300, "hr") :: Nil)
+ }
+ }
+ }
+ }
+
+ test("column-update: CHAR and VARCHAR self-assignments are reported as
updated") {
+ createTable("pk INT NOT NULL, id INT, c CHAR(3), v VARCHAR(10), dep
STRING")
+ sql(s"INSERT INTO $tableNameAsString VALUES (1, 1, 'a', 'x', 'hr')")
+
+ sql(s"UPDATE $tableNameAsString SET id = id, c = c, v = v WHERE pk = 1")
+
+ checkLastUpdatedColumns("c", "v")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, 1, "a ", "x", "hr") :: Nil)
+ }
+
+ test("column-update: a LogicalWriteInfo that does not implement
columnUpdateSchema throws") {
+ val info = new LogicalWriteInfo {
+ override def options(): CaseInsensitiveStringMap =
CaseInsensitiveStringMap.empty()
+ override def queryId(): String = "query"
+ override def schema(): StructType = new StructType()
+ }
+
+ checkError(
+ exception =
intercept[SparkUnsupportedOperationException](info.columnUpdateSchema()),
+ condition = "DATA_SOURCE_COLUMN_UPDATE_SCHEMA_NOT_IMPLEMENTED",
+ parameters = Map("class" -> info.getClass.getName))
+ }
+
+ test("column-update: updatedColumns and the column update schema follow the
non-identity " +
+ "assignments") {
+ Seq(
+ ("id = -1", Seq("id"), Seq("pk", "dep", "id"), Row(1, -1, "hr")),
+ ("ID = -1", Seq("id"), Seq("pk", "dep", "id"), Row(1, -1, "hr")),
+ ("id = -1, dep = 'eng'", Seq("id", "dep"), Seq("pk", "dep", "id"),
Row(1, -1, "eng")),
+ ("id = -1, dep = dep", Seq("id"), Seq("pk", "dep", "id"), Row(1, -1,
"hr")),
+ ("id = id, dep = 'eng'", Seq("dep"), Seq("pk", "dep"), Row(1, 10,
"eng")),
+ ("id = id, dep = dep", Seq.empty, Seq("pk", "dep"), Row(1, 10, "hr")),
+ ("id = pk", Seq("id"), Seq("pk", "dep", "id"), Row(1, 1, "hr"))
+ ).foreach { case (assignments, updatedColumns, columnUpdateNames,
updatedRow) =>
+ withTable(tableNameAsString) {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 10, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET $assignments WHERE pk = 1")
+
+ checkLastUpdatedColumns(updatedColumns: _*)
+ checkLastColumnUpdateSchema(columnUpdateNames: _*)
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), updatedRow ::
Nil)
+ }
+ }
+ }
+
+ test("column-update split: row-level DELETE falls back to full-width rows") {
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin)
+ checkDeleteFallsBackToFullWidthRows(Row(2, 200, "hr"))
+ }
+
+ test("column-update split: row-level DELETE by a metadata row ID falls back
to full-width " +
+ "rows") {
+ // Each partition holds one row, so `_partition` identifies it.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ splitPartitionRowIdTableProps)
+ checkDeleteFallsBackToFullWidthRows(Row(2, 200, "software"))
+ }
+
+ test("column-update split: MERGE falls back to full-width rows") {
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+ checkMergeFallsBackToFullWidthRows()
+ }
+
+ test("column-update: undeclared write requirement columns are rejected") {
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key ->
ColumnPruning.ruleName
+ val splitReqAttrsOperationClass =
+
classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithReqAttrs]
+ .getName
+ Seq(
+ (reqAttrsTableProps("pk,salary"), "dep", reqAttrsOperationClass,
"[`dep`]"),
+ (splitReqAttrsTableProps("pk,salary"), "dep",
splitReqAttrsOperationClass, "[`dep`]"),
+ (tablePropsWith(reqAttrsTableProps("pk,salary"),
+ "column-update-cluster-by-nested-transform" -> "true"), "dep",
reqAttrsOperationClass,
+ "[`dep`]"),
+ (reqAttrsTableProps("pk,salary"), "id,dep", reqAttrsOperationClass,
"[`dep`, `id`]"),
+ (reqAttrsTableProps("pk,salary,s"), "s.missing", reqAttrsOperationClass,
+ "[`s`.`missing` (FIELD_NOT_FOUND)]")
+ ).foreach { case (props, clusterBy, operationClass, columns) =>
+ withTable(tableNameAsString) {
+ val tableProps = tablePropsWith(props, "column-update-cluster-by" ->
clusterBy)
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, s STRUCT<c1:
INT>, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 1, "s": { "c1": 1 }, "dep": "hr" }
+ |""".stripMargin,
+ tableProps)
+
+ Seq(Nil, Seq(excludeColumnPruning)).foreach { confs =>
+ withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> operationClass, "columns" ->
columns))
+ }
+ }
+ }
+ }
+ }
+
+ test("column-update: an empty write requirement column name is reported as
<empty>") {
+ checkError(
+ exception =
QueryCompilationErrors.undeclaredWriteRequirementColumnsError(
+ reqAttrsOperationClass, Seq(Seq.empty[String] ->
None)).asInstanceOf[AnalysisException],
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" ->
"[<empty>]"))
+ }
+
+ test("column-update: distribution by an undeclared row ID column is
allowed") {
+ Seq("salary = -1", "salary = -1, pk = pk").foreach { assignments =>
+ withTable(tableNameAsString) {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin,
+ tablePropsWith(reqAttrsTableProps("salary,dep"),
"column-update-cluster-by" -> "pk"))
+
+ sql(s"UPDATE $tableNameAsString SET $assignments WHERE pk = 1")
+
+ checkLastUpdatedColumns("salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 200, "hr") :: Nil)
+ }
+ }
+ }
+
+ test("column-update: an undeclared, unreferenced partition column is not
read") {
+ createAndInitTableWithReqAttrs("pk,id", "pk INT NOT NULL, id INT, dep
STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1", "pk INT,
_partition STRING")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively
without adopting the " +
+ "declared spelling") {
+ createAndInitTableWithReqAttrs("PK,SALARY,dep", "pk INT NOT NULL, salary
INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ checkLastColumnUpdateSchema("pk", "salary", "dep")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: DELETE on a scan that returns unrequested columns is
unaffected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |""".stripMargin,
+ tablePropsWith(extraTableProps,
InMemoryBaseTable.SIMULATE_PARTIAL_COLUMN_PRUNING -> "true"))
+
+ sql(s"DELETE FROM $tableNameAsString WHERE pk < 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(2, 200, 20, "software") :: Nil)
+ }
+
+ test("column-update: case-sensitive analysis accepts columns that differ
only in case") {
+ withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") {
+ createAndInitTableWithReqAttrs("pk,salary,SALARY,dep",
+ "pk INT NOT NULL, salary INT, SALARY INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "SALARY": 200, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateSchema("pk", "salary", "SALARY", "dep")
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, -1, 200,
"hr") :: Nil)
+ }
+ }
+
+ test("column-update: invalid requiredDataAttributes are rejected") {
+ Seq(
+ ("pk,s.c1,dep", "COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE",
+ "nestedAttributes" -> "[`s`.`c1`]"),
+ ("pk,pk,salary", "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ "duplicateAttributes" -> "[`pk`]"),
+ ("salary,PK,SALARY,pk",
"COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ "duplicateAttributes" -> "[`salary`, `PK`]"),
+ ("pk,_partition,salary",
"COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE",
+ "metadataAttributes" -> "[`_partition`]"),
+ ("pk,nonexistent_col,salary",
"COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE",
+ "unknownAttributes" -> "[`nonexistent_col`]"),
+ ("pk,salary",
"COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS",
+ "missingColumns" -> "[`a.b`]")
+ ).foreach { case (reqAttrs, condition, columns) =>
+ withTable(tableNameAsString) {
+ createAndInitTableWithReqAttrs(reqAttrs,
+ "pk INT NOT NULL, salary INT, s STRUCT<c1: INT, c2: INT>, `a.b` INT,
dep STRING",
+ """{ "pk": 1, "salary": 100, "s": { "c1": 1, "c2": 2 }, "a.b": 1,
"dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1, `a.b` = -1 WHERE
pk = 1")
+ },
+ condition = condition,
+ parameters = Map("connector" -> reqAttrsOperationClass, columns))
+ }
+ }
+ }
+
+ test("column-update: a declared column with no name parts is reported as
<empty>") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+ val t = table
+ val operation = new t.DeltaBasedColumnUpdateOperation(UPDATE) {
+ override def requiredDataAttributes(): Array[NamedReference] =
+ Array(FieldReference(Seq.empty[String]), FieldReference("pk"),
+ FieldReference(Seq.empty[String]))
+ }
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ RewriteRowLevelCommand.resolveRequiredDataAttrs(
+ spark.table(tableNameAsString).queryExecution.analyzed, operation)
+ },
+ condition = "COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> operation.getClass.getName, "unknownAttributes" ->
"[<empty>, <empty>]"))
+ }
+
+ test("column-update split: delete and reinsert carry the row ID and
preserved metadata") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ deleteWriteLogEntry(id = 1, metadata = Row("hr", null)),
+ reinsertWriteLogEntry(metadata = Row("hr", null), data = Row(1, "hr",
-1)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "hr") :: Nil)
+ }
+
+ test("column-update split: row-ID reassignment on narrow write is rejected")
{
+ // the declared [pk, dep, salary] leaves out `extra`, so the write is
narrow
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING, extra
STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr", "extra": "x" }
+ |{ "pk": 2, "salary": 200, "dep": "software", "extra": "y" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE
dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT",
+ parameters = Map("connector" -> splitOperationClass, "rowIds" ->
"[`pk`]"))
+ }
+
+ test("column-update split: row-ID reassignment is allowed when
requiredDataAttributes " +
+ "covers every table column") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ splitReqAttrsTableProps("pk,salary,dep"))
+
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep =
'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(2, 200, "software") :: Row(11, -1, "hr") :: Nil)
+ }
+
+ test("column-update split: row-ID reassignment is allowed when
requiredDataAttributes " +
+ "covers every data column of a relation with metadata columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ splitReqAttrsTableProps("pk,salary,dep"))
+
+ // `_partition` and the coerced `salary > 0.5` make the analyzer add the
metadata columns to
Review Comment:
Done.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala:
##########
@@ -65,7 +65,9 @@ object GroupBasedRowLevelOperationScanPlanning extends
Rule[LogicalPlan] with Pr
.mkString(", ")
}
- val (scan, output) = PushDownUtils.pruneColumns(scanBuilder, relation,
relation.output, Nil)
+ val requiredAttrs =
+ if (V2Writes.hasNarrowRows(rd)) readAttrs(rd.query, relation) else
relation.output
Review Comment:
I kept the gate, so connectors without the mix-in see no scan change. As you
note, dropping it would fail `UPDATE t SET dep = 'x'` on a scan that always
reports `dep` as a filter attribute. I think that belongs in a separate change.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/RowLevelOperationTable.scala:
##########
@@ -60,4 +65,14 @@ private[sql] case class RowLevelOperationTable(
override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = {
operation.newWriteBuilder(info)
}
+
+ /**
Review Comment:
Done.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,114 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can
implement this interface to
+ * receive a narrow row containing only the columns declared via {@link
#requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, updated,
copied, and reinserted
+ * rows follow it: group-based operations receive them through
+ * {@link DataWriter#writeColumnUpdate(Object, Object)} or
+ * {@link DataWriter#writeColumnUpdate(Object)}, and operations that mix in
{@link SupportsDelta}
+ * through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert}. Rows
the command inserts
+ * always follow {@link LogicalWriteInfo#schema()} and arrive through
+ * {@link DataWriter#write(Object)} and {@link DeltaWriter#insert}. When
+ * {@link LogicalWriteInfo#columnUpdateSchema()} is absent, every row follows
+ * {@link LogicalWriteInfo#schema()}, as for any other operation. Connectors
must decide which
+ * layout to expect by whether {@link LogicalWriteInfo#columnUpdateSchema()}
is present, not by
+ * the command. A builder that wants full-width rows for some command can check
+ * {@link RowLevelOperationInfo#command()} and build an operation without this
mix-in. Currently
+ * Spark narrows only UPDATE.
+ * <p>
+ * The scan builder returned by {@link #newScanBuilder} should implement
+ * {@link
org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns}, so that
the scan
+ * does not read columns the command neither references nor declares.
Otherwise the scan reads
+ * every column. For operations that do not mix in {@link SupportsDelta}, only
top-level columns
+ * are pruned: a struct column the command references through a nested field
is read in full.
+ *
+ * @since 4.4.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level
operation.
+ * <p>
+ * When Spark narrows the write, the returned columns become
+ * {@link LogicalWriteInfo#columnUpdateSchema()}, in declared order. Spark
matches each declared
+ * column with the session's resolver, and {@link
LogicalWriteInfo#columnUpdateSchema()} uses the
+ * table's spelling of each column, which can differ in case from the
declared reference.
+ * Implementations must include every column they want to receive: every
column reported by
+ * {@link RowLevelOperationInfo#updatedColumns()}, plus any columns needed
for row lookup or
+ * routing. Columns that are not declared are absent from the rows the
connector receives, so
+ * the connector must preserve their values itself. Spark may call this
method more than once
+ * for the same operation, and every call must return the same columns.
+ * <p>
+ * If the returned array is empty, Spark does not narrow the write and the
operation receives
+ * full-width rows. An UPDATE whose assignments are all identity assignments
reports no
+ * {@link RowLevelOperationInfo#updatedColumns()}, so an operation that
declares exactly those
+ * columns takes this path for such an UPDATE. An empty array is not an
opt-out: an UPDATE that
+ * sets a column is still rejected, because the array misses that column.
+ * <p>
+ * An operation that does not mix in {@link SupportsDelta} should locate
original rows through
+ * metadata columns returned by {@link #requiredMetadataAttributes()} that
are preserved on
+ * update ({@link
org.apache.spark.sql.connector.catalog.MetadataColumn#PRESERVE_ON_UPDATE}),
+ * such as a file and a row position. If it locates rows through a data
column, its builder
+ * should build an operation without this mix-in when
+ * {@link RowLevelOperationInfo#updatedColumns()} contains that column.
+ * <p>
+ * Each entry must name a top-level data column of the table. For updates on
nested fields such
Review Comment:
Yes, it's to limit the scope of the PR. The hard part is the analyzer. I
investigated and found these should be addressed.
1. Alignment loses the fields. Spark aligns `SET s.c1 = -1` to one
assignment per top-level column, `s = named_struct('c1', -1, 'c2', s.c2)`. To
report `s.c1`, Spark would have to take that struct apart again and find the
fields whose value differs from the old field.
2. NULL structs. If `s` is NULL, `SET s.c1 = -1` makes `s` non-NULL with
`c2` NULL. The comment on our identity test notes this. A row with only `s.c1`
can't say that `s` itself changed. Spark would have to report `s` whenever the
struct can be NULL, or send the struct's null flag with the row.
3. Write schema and projections. `columnUpdateSchema()` would hold a pruned
struct, `s STRUCT<c1 INT>`, and the update, copy and reinsert projections would
build it. Declarations would need overlap rules, for example for `s` and `s.c1`
together.
4. Scan. Operations that don't mix in `SupportsDelta` prune only top-level
columns, so they would still read all of `s`.
I think this will make this PR more complex. I'm happy to do this
immediately after this PR merges.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -354,3 +364,69 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
}
}
}
+
+object RewriteRowLevelCommand extends SQLConfHelper {
+
+ /**
+ * Whether the operation narrows the write: it mixes in
`SupportsColumnUpdates` and declares at
+ * least one column. An operation that declares none receives full-width
rows.
+ */
+ private[sql] def narrowsWrite(operation: RowLevelOperation): Boolean =
operation match {
+ case op: SupportsColumnUpdates => op.requiredDataAttributes.nonEmpty
+ case _ => false
+ }
+
+ /**
+ * Resolves the connector-declared required data attributes for
column-update writes against
+ * the given relation, in declared order and with the relation's spelling.
Nested fields,
+ * duplicate columns, columns that do not exist in the relation (including a
reference with no
+ * name parts) and metadata columns are rejected with an `AnalysisException`.
+ */
+ private[sql] def resolveRequiredDataAttrs(
+ relation: LogicalPlan,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes.toImmutableArraySeq
+ val nameless = refs.filter(_.fieldNames.isEmpty).map(_ =>
Seq.empty[String])
+ if (nameless.nonEmpty) {
+ throw QueryCompilationErrors.unknownRequiredDataAttributeError(
+ operation.getClass.getName, nameless)
+ }
+ val nested = refs.filter(_.fieldNames.length > 1).map(_.fieldNames.toSeq)
+ if (nested.nonEmpty) {
+ throw QueryCompilationErrors.nestedRequiredDataAttributeError(
+ operation.getClass.getName, nested)
+ }
+ val normalizedNames =
+ refs.map(ref => SchemaUtils.foldName(ref.fieldNames.head,
conf.caseSensitiveAnalysis))
Review Comment:
I left this as is. A case-insensitive session can't use such a table. A
SELECT fails with `COLUMN_ALREADY_EXISTS` from
`V2TableUtil.validateCapturedColumns`, which folds the two names into one. An
UPDATE reaches this check first and fails with
`COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE`, as you traced.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/RowLevelOperationInfoImpl.scala:
##########
@@ -17,9 +17,15 @@
package org.apache.spark.sql.connector.write
+import org.apache.spark.sql.connector.expressions.NamedReference
import org.apache.spark.sql.connector.write.RowLevelOperation.Command
import org.apache.spark.sql.util.CaseInsensitiveStringMap
private[sql] case class RowLevelOperationInfoImpl(
command: Command,
- options: CaseInsensitiveStringMap) extends RowLevelOperationInfo
+ options: CaseInsensitiveStringMap,
+ private val updatedCols: Seq[NamedReference])
Review Comment:
Per Anton, I removed `operationInfo` and `withTable`, and
`RowLevelOperationTable` matches master again. Iceberg's `ResolveBranch` will
handle this on its side.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala:
##########
@@ -111,10 +112,10 @@ class
RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla
// and must transform the runtime filter condition to use correct expr
IDs for each relation
// note this only applies to group-based row-level operations (i.e.
ReplaceData)
// see RewriteUpdateTable for more details
- // the condition may also reference metadata columns resolved from the
metadata output
- // of the target relation, which are not in the table output and must
be mapped too
- val metadataAttrs = (cond.references -- write.table.output).toSeq
- val attrMap = buildTableToScanAttrMap(write.table.output ++
metadataAttrs, relation.output)
+ // the write table may leave out data columns, so map those of the
original table
Review Comment:
Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -237,4 +253,101 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Returns the table attributes that these assignments set to a new value,
skipping identity
+ * assignments.
+ */
+ private def collectUpdatedAttrs(assignments: Seq[Assignment]):
Seq[AttributeReference] = {
+ assignments.collect {
+ case Assignment(key: AttributeReference, value) if
!isIdentityAssignment(key, value) => key
Review Comment:
Done. My reply to item 56 was wrong for tables that accept any schema, which
SPARK-60721 skips. On such a table, a narrow write now also rejects a new value
for a metadata column, because the connector finds rows through those columns.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/ColumnUpdateTableSharedTests.scala:
##########
@@ -0,0 +1,427 @@
+/*
+ * 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
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.{SparkException, SparkRuntimeException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.{InMemoryBaseTable,
InMemoryRowLevelOperationTable}
+import org.apache.spark.sql.connector.write.{RowLevelOperation,
RowLevelOperationTable}
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation,
V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+// Column-update tests that hold for both delta-based and group-based
operations; each suite
+// supplies its column-update table properties through `extraTableProps`.
+trait ColumnUpdateTableSharedTests extends UpdateTableSuiteBase {
+
+ // the properties of a column-update table whose connector declares exactly
`reqAttrs`
+ protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String,
String]
+
+ // the properties of a table of the same kind whose operation does not mix
in column updates
+ protected def fullWidthTableProps: java.util.Map[String, String]
+
+ // the metadata columns the scan of a column-update table reads, in DDL form
+ protected def scanMetadataColumns: String = ", _partition STRING, index INT"
+
+ // the class name of the operation that `reqAttrsTableProps` builds
+ protected def reqAttrsOperationClass: String = if (deltaUpdate) {
+
classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateOperationWithReqAttrs].getName
+ } else {
+
classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperationWithReqAttrs].getName
+ }
+
+ // `base` with the given table properties added; a given key replaces the
same key in `base`
+ protected def tablePropsWith(
+ base: java.util.Map[String, String],
+ props: (String, String)*): java.util.Map[String, String] = {
+ (base.asScala ++ props).asJava
+ }
+
+ // t(pk, id, dep, salary, bonus, extra) with CHECK (extra > 0), partitioned
by dep
+ protected def createAndInitTableWithProps(props: java.util.Map[String,
String]): Unit = {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, salary INT, bonus
INT, extra INT",
+ """{ "pk": 1, "id": 1, "dep": "hr", "salary": 100, "bonus": 10, "extra":
1 }
+ |{ "pk": 2, "id": 2, "dep": "hr", "salary": 200, "bonus": 20, "extra":
1 }
+ |{ "pk": 3, "id": 3, "dep": "software", "salary": 300, "bonus": 30,
"extra": 1 }
+ |""".stripMargin,
+ props)
+ sql(s"ALTER TABLE $tableNameAsString ADD CONSTRAINT positive_extra CHECK
(extra > 0)")
+ }
+
+ test("column-update: undeclared columns are not written and keep their
values") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkLastColumnUpdateSchema("pk", "dep", "salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update: write relation has the declared columns in declared
order") {
+ createAndInitTableWithProps(reqAttrsTableProps("salary,pk"))
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep =
'hr'")
+ }
+
+ assert(!write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "salary", "pk")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update: the plan writes the declared columns and scans the
referenced ones") {
Review Comment:
Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -159,14 +164,24 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
relation: DataSourceV2Relation,
operationTable: RowLevelOperationTable,
assignments: Seq[Assignment],
- cond: Expression): WriteDelta = {
+ cond: Expression,
+ writeAttrs: Seq[AttributeReference],
+ updatedAttrs: Seq[AttributeReference]): WriteDelta = {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
// resolve all needed attrs (e.g. row ID and any required metadata attrs)
- val rowAttrs = relation.output
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
+ val splitUpdates = operation.representUpdateAsDeleteAndInsert
+
+ operation match {
+ case op: SupportsColumnUpdates if splitUpdates &&
RewriteRowLevelCommand.narrowsWrite(op) =>
Review Comment:
Done.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -372,6 +875,66 @@ private class DeltaBufferWriter(schema: StructType)
extends BufferWriter(schema)
override def commit(): WriterCommitMessage = buffer
}
+// Test-only: a writer factory for a connector that never overrides
writeColumnUpdate(), to
+// exercise the DataWriter#writeColumnUpdate default that throws
+// DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED.
+// Declared top-level (no reference to the enclosing table) so it stays
task-serializable.
+private class NoWriteColumnUpdateOverrideWriterFactory extends
DataWriterFactory {
Review Comment:
Done. The test now also fails if the overrides move back to `BufferWriter`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -237,4 +253,101 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Returns the table attributes that these assignments set to a new value,
skipping identity
+ * assignments.
+ */
+ private def collectUpdatedAttrs(assignments: Seq[Assignment]):
Seq[AttributeReference] = {
+ assignments.collect {
+ case Assignment(key: AttributeReference, value) if
!isIdentityAssignment(key, value) => key
+ }
+ }
+
+ /**
+ * Returns the output of the write relation, which is the layout of the data
rows the connector
+ * receives. If the operation opts into column updates, this is the
connector's
+ * `requiredDataAttributes()`. An empty declaration yields the full relation
output, as does any
+ * operation without the mix-in. A mixed-in operation's declaration must
cover every updated
+ * column, even an empty one.
+ */
+ private def resolveWriteAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation,
+ updatedAttrs: Seq[AttributeReference]): Seq[AttributeReference] =
operation match {
+ case op: SupportsColumnUpdates =>
+ val attrs = if (RewriteRowLevelCommand.narrowsWrite(op)) {
Review Comment:
Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -159,14 +164,24 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
relation: DataSourceV2Relation,
operationTable: RowLevelOperationTable,
assignments: Seq[Assignment],
- cond: Expression): WriteDelta = {
+ cond: Expression,
+ writeAttrs: Seq[AttributeReference],
+ updatedAttrs: Seq[AttributeReference]): WriteDelta = {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
// resolve all needed attrs (e.g. row ID and any required metadata attrs)
- val rowAttrs = relation.output
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
+ val splitUpdates = operation.representUpdateAsDeleteAndInsert
+
+ operation match {
Review Comment:
The problem is that the reinsert row is partial. With narrow rows it carries
only the declared columns, but the connector still has to combine it with the
row it replaces. The delete knows the row but not the new value. The reinsert
has the new value but not the row.
With full-width rows this never comes up, because the reinsert row is
complete.
That's why the checks require the reinsert row to carry the row ID, either
as a declared data column or as a metadata column preserved on reinsert.
Are you suggesting the connectors should match each reinsert with its delete
themselves? I see three ways, and none of them feels solid.
- Rely on the Expand order. Spark doesn't guarantee it once a distribution
or ordering applies.
- Have the connector request a distribution and ordering that keep each pair
together. But no column on the reinsert row ties it to its delete.
- Have Spark pass the row ID to `reinsert`. That changes the API.
If you had another way in mind, I'm happy to brainstorm.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,114 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can
implement this interface to
+ * receive a narrow row containing only the columns declared via {@link
#requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, updated,
copied, and reinserted
+ * rows follow it: group-based operations receive them through
+ * {@link DataWriter#writeColumnUpdate(Object, Object)} or
+ * {@link DataWriter#writeColumnUpdate(Object)}, and operations that mix in
{@link SupportsDelta}
+ * through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert}. Rows
the command inserts
+ * always follow {@link LogicalWriteInfo#schema()} and arrive through
+ * {@link DataWriter#write(Object)} and {@link DeltaWriter#insert}. When
+ * {@link LogicalWriteInfo#columnUpdateSchema()} is absent, every row follows
+ * {@link LogicalWriteInfo#schema()}, as for any other operation. Connectors
must decide which
+ * layout to expect by whether {@link LogicalWriteInfo#columnUpdateSchema()}
is present, not by
+ * the command. A builder that wants full-width rows for some command can check
+ * {@link RowLevelOperationInfo#command()} and build an operation without this
mix-in. Currently
+ * Spark narrows only UPDATE.
+ * <p>
+ * The scan builder returned by {@link #newScanBuilder} should implement
Review Comment:
Done.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DistributionAndOrderingUtils.scala:
##########
@@ -38,15 +39,8 @@ object DistributionAndOrderingUtils {
val numPartitions = write.requiredNumPartitions()
val partitionSize = write.advisoryPartitionSizeInBytes()
- val distribution = write.requiredDistribution match {
Review Comment:
It makes the check and `prepareQuery` read the same expressions. The check
rejects a distribution or ordering column that the narrow write doesn't
receive. Then `prepareQuery` turns the same expressions into the repartition
and sort. If the two ever disagreed on a distribution kind, Spark would reject
or resolve an undeclared column depending on whether `ColumnPruning` had
already removed it. If you'd rather keep `prepareQuery` as it was and match
locally in `V2Writes`, I'll revert it.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala:
##########
@@ -109,22 +112,95 @@ object V2Writes extends Rule[LogicalPlan] with
PredicateHelper {
WriteToDataSourceV2(Some(r), microBatchWrite, newQuery, customMetrics)
case rd @ ReplaceData(r: DataSourceV2Relation, _, query, _, projections,
_, None) =>
- val rowSchema = projections.rowProjection.schema
+ validateWriteRelation(rd)
+ val (rowSchema, columnUpdateSchema) =
+ writeSchemas(projections.rowProjection.schema, hasNarrowRows(rd))
val metadataSchema = projections.metadataProjection.map(_.schema)
val writeOptions = mergeOptions(Map.empty,
r.options.asCaseSensitiveMap.asScala.toMap)
- val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema,
metadataSchema)
+ val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema,
metadataSchema,
+ columnUpdateSchema)
val write = writeBuilder.build()
+ validateColumnUpdateWrite(rd, write)
val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query,
r.funCatalog)
rd.copy(write = Some(write), query = newQuery)
case wd @ WriteDelta(r: DataSourceV2Relation, _, query, _, projections, _,
None) =>
+ validateWriteRelation(wd)
val writeOptions = mergeOptions(Map.empty,
r.options.asCaseSensitiveMap.asScala.toMap)
- val deltaWriteBuilder = newDeltaWriteBuilder(r.table, writeOptions,
projections)
+ val deltaWriteBuilder =
+ newDeltaWriteBuilder(r.table, writeOptions, projections,
hasNarrowRows(wd))
val deltaWrite = deltaWriteBuilder.build()
+ validateColumnUpdateWrite(wd, deltaWrite)
val newQuery = DistributionAndOrderingUtils.prepareQuery(deltaWrite,
query, r.funCatalog)
wd.copy(write = Some(deltaWrite), query = newQuery)
}
+ /**
+ * Whether a row-level write delivers narrow rows for updated, copied, and
reinserted records.
+ * This must match the writes whose write relation RewriteUpdateTable
narrows, even when the
+ * declared columns cover the whole table.
+ */
+ private[v2] def hasNarrowRows(write: RowLevelWrite): Boolean = {
+ write.operation.command == UPDATE &&
RewriteRowLevelCommand.narrowsWrite(write.operation)
+ }
+
+ /**
+ * Checks that the write relation matches the current operation: a write
that delivers narrow
+ * rows must have exactly the declared columns, in declared order, and any
other write the full
+ * table columns. A rule that replaces the operation after the rewrite can
break this, and the
+ * connector would then decode the rows with the wrong layout.
+ */
+ private def validateWriteRelation(rowLevelWrite: RowLevelWrite): Unit = {
+ val writeNames = rowLevelWrite.table.output.map(_.name)
+ val expectedNames = rowLevelWrite.operation match {
+ case op: SupportsColumnUpdates if hasNarrowRows(rowLevelWrite) =>
+
RewriteRowLevelCommand.resolveRequiredDataAttrs(rowLevelWrite.originalTable, op)
+ .map(_.name)
+ case _ =>
+ rowLevelWrite.originalTable.output.map(_.name)
+ }
+ if (writeNames != expectedNames) {
+ throw SparkException.internalError(
+ s"The write relation of operation
${rowLevelWrite.operation.getClass.getName} has " +
+ s"columns ${writeNames.mkString("[", ", ", "]")}, but the operation
expects " +
+ s"${expectedNames.mkString("[", ", ", "]")}.")
+ }
+ }
+
+ /**
+ * Checks that the distribution and ordering of a write that delivers narrow
rows reference only
+ * columns passed to the write. Other columns stay in the query only until
column pruning removes
+ * them, so they are rejected here whether or not pruning ran.
+ */
+ private def validateColumnUpdateWrite(rowLevelWrite: RowLevelWrite, write:
Write): Unit = {
+ write match {
+ case w: RequiresDistributionAndOrdering if hasNarrowRows(rowLevelWrite)
=>
+ val distribution =
+
DistributionAndOrderingUtils.distributionExpressions(w.requiredDistribution)
+ val readAttrs = rowLevelWrite.references.toSeq.distinctBy(_.name)
Review Comment:
Done.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
