dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4428138885
##########
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:
**66. `updatedColumns()` reports the user's spelling when the SET clause is
already aligned, so the check that `SupportsColumnUpdates.java:74-76` asks
builders to make can miss the lookup column.** This line takes the reported
name from the assignment key, and a key resolves to `a.withName(<typed
spelling>)` (`expressions/package.scala:197`). Alignment is the only step that
restores the table's attribute (`AssignmentUtils.scala:91`), and it runs only
for `!u.aligned` (`ResolveRowLevelCommandAssignments.scala:51`).
`AssignmentUtils.aligned` compares names with the resolver
(`AssignmentUtils.scala:402`), so in a case-insensitive session a SET clause
that lists every column in table order skips it, and "Column names match the
case used in the table schema" (`RowLevelOperationInfo.java:53`) no longer
holds. For example:
- With `t(pk INT NOT NULL, salary INT, bonus INT, dep STRING)` and the rows
(1, 100, 10, hr) and (2, 200, 20, hr) on `column-update-cow-req-attrs =
pk,salary`, `UPDATE t SET PK = 2, SALARY = 0, BONUS = BONUS, DEP = DEP WHERE PK
= 1` reports `[PK, SALARY]`. `updatesPk`
(`InMemoryRowLevelOperationTable.scala:161`) is then false, so the narrow
operation is built instead of the full-width one, and the commit finds the base
row by the new `pk = 2`. The updated row becomes (2, 0, 20, hr) instead of (2,
0, 10, hr), with no error. The same statement in lower case falls back and gets
it right. A real connector that follows `SupportsColumnUpdates.java:74-76` with
an exact comparison would look rows up the same way.
- With `t(pk INT NOT NULL, id INT, dep STRING)` on the default
`column-update` fixture, `UPDATE t SET PK = PK, ID = -1, DEP = 'eng' WHERE pk =
1` declares `[pk, dep, ID, DEP]`
(`InMemoryRowLevelOperationTable.scala:392-396`) and fails with
`COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE`, while `SET id = -1, dep =
'eng'` works.
`("ID = -1", Seq("id"), ...)` at
`DeltaBasedColumnUpdateTableSuite.scala:169` does not cover this, because
assigning one of three columns goes through alignment. Since `u.aligned`
guarantees one assignment per column in table order, could this use the
relation's attributes, e.g. `relation.output.zip(assignments).collect { case
(attr, Assignment(key: Attribute, value)) if !isIdentityAssignment(key, value)
=> attr }`, the same idea as the exprId map-back in `resolveRequiredDataAttrs`,
and could that test get a row with upper-case keys?
##########
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:
**67. A metadata key can still reach `updatedColumns()` on a table that
accepts any schema, so item 56's dead end remains there.** a6e22b24587 removed
the metadata filter on these keys, and [the reply to item
56](https://github.com/apache/spark/pull/55518#discussion_r4426725837) says
SPARK-60721 handles it. That holds for other tables, but SPARK-60721's check
and SPARK-60667's target exclusion both skip `skipSchemaResolution` tables
(`ResolveRowLevelCommandAssignments.scala:98` and `Analyzer.scala:1198`). On a
`column-update=true, accept-any-schema=true` table `t(pk INT NOT NULL, salary
INT, dep STRING)`, `UPDATE t SET pk = pk, salary = -1, dep = dep, index = 5,
_partition = 'x'` needs no injected rule: `AddMetadataColumns` widens the
target to `[pk, salary, dep, index, _partition]`, the assignments are then
aligned, and this reports `[salary, index, _partition]`. Declaring those fails
with `COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE`, leaving them out fails
with `COLUMN_UPDAT
E_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS`, and both blame the
connector for the user's assignment. Before a6e22b24587 this reported
`[salary]`. It is a rare shape, since the SET clause has to list every data and
metadata column in output order, and only an error-quality problem. Could this
line skip metadata keys again with `!MetadataAttribute.isValid(key.metadata)`,
or could `RowLevelOperationInfo.updatedColumns()` mention the case?
##########
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:
**69. The PR description says "`LogicalWriteInfoImpl.apply` takes a required
`columnUpdateSchema`", but only the `Option` overload at L43 does.** The
`Optional` `apply` keeps this default, so master's five-argument shape, as in
`FileTable.mergedWriteInfo` before this PR, still compiles. An extension under
`org.apache.spark.sql` that rebuilds the info that way hands the connector an
empty `schema()` and no `columnUpdateSchema()`, while
`DataSourceV2Strategy.scala:651` still sends narrow rows. Item 10's default
throw does not help here, because this class overrides the method. Spark's own
paths all forward it, and `LogicalWriteInfo.java:83-84` already says a wrapper
must forward it, so only the description needs a change. Could the
compatibility paragraph say that the compiler catches the `Option` overload but
not this one?
##########
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:
**73.** Nit: some of the new tests overlap or only exercise the fixture.
- L89 and L105 run the same schema, statement and rows. L89 adds the
declared order, the original table and the empty `schema()`, and L105 adds the
exact scan and the node name, so one test with `"salary,pk"` could check all of
them. The regex at L118 repeats the first assertion of `checkColumnUpdateWrite`
(`RowLevelOperationSuiteBase.scala:229`).
- L141 and L292 make the same assertion from `"pk,salary"` and `""`, so
`Seq(...).foreach` could keep both inputs.
- `DeltaBasedColumnUpdateTableSuite.scala:504` asserts only the text of the
fixture's own exception. DELETE is never narrow (`V2Writes.scala:144`), so it
covers no Spark behavior.
- `DeltaBasedColumnUpdateTableSuite.scala:572` cannot fail on its CHECK,
since `extra > 0` holds for every value it writes. It shows that the column
resolves rather than that the constraint is enforced.
##########
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:
**71. This comment is no longer true since SPARK-60667, so the test no
longer covers what its name says.** With the merge of master,
`AddMetadataColumns` skips the target of a row-level command
(`Analyzer.scala:1144` and `Analyzer.scala:1197-1198`), so the relation stays
`[pk, salary, dep]` with no metadata columns in it. The statement should also
run now: the delete and reinsert outputs of `buildDeletesAndInserts` have the
same width, 6, and SPARK-60667's commit message lists this `ExpandExec` failure
as fixed. `DeltaBasedUpdateAsDeleteAndInsertTableSuite.scala:98-128` already
runs the same shape. As a result:
- The only assertion, at L453, checks the declaration, and the relation is
now the same as in the test at L420. Removing the `filterNot` at
`RewriteUpdateTable.scala:319`, the fix for [item
21](https://github.com/apache/spark/pull/55518#discussion_r4316536100), fails
no test. The filter now matters only on a table that accepts any schema, where
it swaps `COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT` for an `ExpandExec` failure.
- No narrow suite executes a later-iteration UPDATE that references a
metadata column, because the SPARK-60667 tests are in concrete suites that the
narrow suites do not extend.
- `executePlan(parsed).analyzed` at L449 leaves a transaction open, because
the default mode begins one (`QueryExecution.scala:145-151`) and nothing
commits or aborts it. It is harmless today only because `after` resets the
catalog manager.
a6e22b24587 removed two other tests that carried this premise, and this one
stayed. Could this test run the statement with `sql(...)` and check `Row(2,
200, "software") :: Row(11, -1, "hr") :: Nil`, with the name and the comment
updated? `dataAttrs` at `RewriteUpdateTable.scala:319` could then be
`relation.output`, which keeps the row-ID error on such a table.
##########
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:
**78.** Nit: since item 58 moved the `writeColumnUpdate` overrides off
`BufferWriter`, this writer and its factory duplicate the plain writer that
`TestBatchWrite` already creates (`InMemoryBaseTable.scala:1009-1011`), which
ends in the same throw at `DataWriter.java:138-142`. L730 could be
`super.createBatchWriterFactory(info)`, with L878-893 gone and
`GroupBasedColumnUpdateTableSuite.scala:167` expecting
`org.apache.spark.sql.connector.catalog.BufferWriter`. That test would then
also fail if the overrides came back to `BufferWriter`, which no test catches
today.
##########
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:
**77.** Nit, and fine to leave as is or to do separately: this gate ties a
scan rule to the write layout, although `readAttrs` gives the same columns for
wide writes unless the condition folds to true. Copied rows of DELETE and MERGE
reference every read column, and for UPDATE only `If(TrueLiteral, trueValue, _)
=> trueValue` drops one (`optimizer/expressions.scala:713`), in which case no
COPY row needs it. So a wide `UPDATE t SET big_col = NULL` with no WHERE clause
reads `big_col` for every row, while delta and narrow copy-on-write scans do
not. `val requiredAttrs = readAttrs(rd.query, relation)` would remove that
difference and match `ColumnPruning`, which applies to every write since item
11. It is not free of behavior change, though: a scan with unconditional
`filterAttributes()` would then fail `UPDATE t SET dep = 'x'` with
`DATA_SOURCE_INVALID_RUNTIME_FILTER_ATTRIBUTE.CANNOT_RESOLVE`
(`DataSourceV2Strategy.scala:207`), where master runs it.
##########
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:
**68. The duplicate check treats `foldName` as column identity, which its
Scaladoc says it is not.** `SchemaUtils.scala:541-549` says "The fold alone is
therefore not an identity" and "Callers must narrow a multi-candidate match
with the resolver". In a case-insensitive session, a column named U+0130 and a
column named `i` plus U+0307 fold to the same string while `equalsIgnoreCase`
keeps them apart. Each reference resolves to its own column at L410, yet
declaring both fails here with
`COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE`, whose message says "Each
column must be declared at most once". This is rare and it fails loudly, so it
is low priority. Moving the check after L418 and counting `resolved` by exprId,
filtering `refs.zip(resolved)` in order, would use the resolver's answer and
keep `` [`salary`, `PK`] ``. Two messages would change, and no test pins
either: `x,x` would become `COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE`, and
case-insensitive `salary,SALARY` on a ta
ble with both columns would become `AMBIGUOUS_REFERENCE` from L410.
##########
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:
**70. The PR description's exception list for identity assignments leaves
out nested self-assignments.** It says "Identity assignments such as `SET a =
a` are excluded, except self-assignments on CHAR or VARCHAR columns with the
default CHAR and VARCHAR handling", which reads as a closed list, but `SET s.c1
= s.c1` is reported as `s`, as this Javadoc and
`ColumnUpdateTableSharedTests.scala:362-371` say. An operation that declares
`[pk, dep]` therefore rejects that statement with
`COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS`. Could the
description mention it? Reporting `s` looks right to me:
`applyFieldAssignments` (`AssignmentUtils.scala:222`) has no null guard, so a
NULL `s` becomes `{null, null}`, as "update fields inside NULL structs"
(`UpdateTableSuiteBase.scala:456-467`) implies. For the same reason, "although
its values don't change" at `ColumnUpdateTableSharedTests.scala:367-368` holds
only for a non-NULL `s`.
##########
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:
**72. No test covers `&& RewriteRowLevelCommand.narrowsWrite(op)`.** Every
test of an empty declaration (`reqAttrsTableProps("")` at
`ColumnUpdateTableSharedTests.scala:237` and
`ColumnUpdateTableSharedTests.scala:293`) is non-split, and every split test
declares at least one column, so reverting this line to `case op:
SupportsColumnUpdates if splitUpdates =>` fails nothing. The difference shows
on `t(pk INT NOT NULL, salary INT, dep STRING)` with
`column-update-split-row-id = _partition, no-metadata = true`: `UPDATE t SET
salary = salary WHERE pk = 1` analyzes to a full-width `WriteDelta` with the
guard, and fails with `COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED` without it.
This is the `DeltaBasedColumnUpdateSplitOperationWithRowId` shape from item 55.
The fixture's commit throws `IllegalStateException` for a matching row
(`InMemoryRowLevelOperationTable.scala:561-564`), so a test would have to check
the analyzed plan.
##########
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:
**75.** Nit: `distinctBy(_.name)` and the `sortBy` at L184 do nothing this
check needs. Every read ordinal comes from `indexWhere` with the resolver
(`RewriteRowLevelCommand.scala:332`), so two attributes in `references` cannot
share a name, and `LocalRelation(rowLevelWrite.references.toSeq)`, as in item
7, is enough. If such a path ever appeared, `distinctBy` would let the column
pass here and fail with `AMBIGUOUS_REFERENCE` in `prepareQuery`.
`unresolvableReferences` already walks the expressions in order and ends with
`.distinct`, so the list is deterministic without the sort, and the sort makes
this the only `COLUMN_UPDATE_*` list not in the connector's or the assignments'
order: `column-update-cluster-by = id,dep` reports `` [`dep`, `id`] ``.
Dropping it changes the expected value at
`DeltaBasedColumnUpdateTableSuite.scala:228`.
##########
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:
**76.** Nit: this branch is no longer needed. Since 8d9856dd55b removed the
`refs.isEmpty` throw, `resolveRequiredDataAttrs` returns an empty `Seq` for an
empty declaration, so L279-283 can be `val attrs =
RewriteRowLevelCommand.resolveRequiredDataAttrs(relation, op)`, with L284-285
unchanged. Similarly, the `operation match` at L178-184 can be `if
(splitUpdates && RewriteRowLevelCommand.narrowsWrite(operation))`, because
`narrowsWrite` takes a `RowLevelOperation` and both validators use the
operation only for its class name. `writeRelation.output` at L88, L127 and L203
is `writeAttrs`.
##########
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:
**74.** Nit: since SPARK-60664 this hunk no longer changes the result, and
it dropped master's comment. `attrMap` only rewrites attributes in `cond`
(L119-121), and master's `cond.references -- write.table.output` already
includes an undeclared column such as `id`, so
`GroupBasedColumnUpdateTableSuite.scala:193` should behave the same without
this change. Master's "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" is gone although L117 still does that, and
[the reply to item
57](https://github.com/apache/spark/pull/55518#discussion_r4426726027) covers
the code but not the comment. Could this be `val attrMap =
buildTableToScanAttrMap(cond.references.toSeq, relation.output)`, which needs
neither `tableAttrs`, `metadataAttrs` nor either comment? Keeping these lines
and restoring the comment works too.
--
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]
