peter-toth commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4053024876
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -228,4 +371,146 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Variant of `buildDeletesAndInserts` for the `SupportsColumnUpdates`
narrow-scan path.
+ * This variant realigns the assignments to one value per surviving rowAttr
padding unassigned
+ * rowAttrs with identity, so the reinsert output arity matches the delete
output arity in
+ * the resulting Expand.
+ */
+ private def buildNarrowDeletesAndInserts(
+ matchedRowsPlan: LogicalPlan,
+ assignments: Seq[Assignment],
+ rowIdAttrs: Seq[Attribute]): Expand = {
+
+ val (metadataAttrs, rowAttrs) = matchedRowsPlan.output.partition { attr =>
+ MetadataAttribute.isValid(attr.metadata)
+ }
+ val assignmentMap = AttributeMap(assignments.collect {
+ case a @ Assignment(key: Attribute, _) => key -> a
+ })
+ val reinsertAssignments = rowAttrs.map { attr =>
+ assignmentMap.get(attr) match {
+ case Some(a) => a
+ case None => Assignment(attr, attr)
+ }
+ }
+ val deleteOutput = deltaDeleteOutput(rowAttrs, rowIdAttrs, metadataAttrs)
+ val insertOutput = deltaReinsertOutput(reinsertAssignments, metadataAttrs)
+ val outputs = Seq(deleteOutput, insertOutput)
+ val operationTypeAttr = AttributeReference(OPERATION_COLUMN, IntegerType,
nullable = false)()
+ val attrs = operationTypeAttr +: matchedRowsPlan.output
+ val expandOutput = generateExpandOutput(attrs, outputs)
+ Expand(outputs, expandOutput, matchedRowsPlan)
+ }
+
+ /**
+ * Resolves the connector's `requiredDataAttributes()` if the operation opts
into column
+ * updates, validating that every assigned column is covered. Returns `Nil`
otherwise.
+ */
+ private def resolveConnectorDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation,
+ assignments: Seq[Assignment]): Seq[AttributeReference] = operation match
{
+ case scu: SupportsColumnUpdates =>
+ val attrs = resolveRequiredDataAttrs(relation, scu)
+ validateUpdatedColumnsSubset(scu, assignments, attrs)
+ attrs
+ case _ => Nil
+ }
+
+ /**
+ * Computes the narrow set of data columns that must be present in the scan
for a column-update
+ * write: connector-declared attrs (`requiredDataAttributes()`), unioned
with any table columns
+ * referenced by non-identity assignment RHS expressions and the operation
condition.
+ */
+ private def computeNarrowReadAttrs(
+ relation: DataSourceV2Relation,
+ connectorDataAttrs: Seq[AttributeReference],
+ assignments: Seq[Assignment],
+ cond: Expression): Seq[AttributeReference] = {
+ val relationSet = relation.outputSet
+ val nonIdentityRhsRefs = assignments.iterator
+ .filterNot(a => a.key.isInstanceOf[Attribute] &&
+ isIdentityAssignment(a.key.asInstanceOf[Attribute], a.value))
+ .flatMap(_.value.references.toSeq)
+ .toSeq
+ val extraRefs = (cond.references.toSeq ++ nonIdentityRhsRefs)
Review Comment:
**Finding 27.** A table carrying a `CHECK` constraint cannot be updated
through this path at all.
`ResolveTableConstraints` wraps the rewritten query in
`Filter(CheckInvariant(...), query)` and resolves the constraint's column
references against `query.output` (`ResolveTableConstraints.scala:60-61`).
`RowLevelOperationTable.constraints()` delegates to the real table, so the rule
fires for the operation table too. Narrowing takes those columns out of both
the scan and the write query, and nothing here accounts for it.
Measured on `72ec505`. Same table in every arm, `pk INT NOT NULL, id INT,
dep STRING, extra INT` partitioned by `identity(dep)`, same constraint `ALTER
TABLE ... ADD CONSTRAINT positive_extra CHECK (extra > 0)`, same statement
`UPDATE t SET id = -1 WHERE pk = 1`:
| connector | result |
|---|---|
| wide `supports-deltas` (MoR) | works |
| wide default fixture (CoW) | works |
| `column-update` (MoR) | ``[UNRESOLVED_COLUMN.WITH_SUGGESTION] A column,
variable, or function parameter with name `extra` cannot be resolved. Did you
mean one of the following? [`dep`, `id`, `pk`, `index`, `_partition`]`` |
| `column-update-cow` (CoW) | same |
So it is a narrowing regression, not pre-existing. The error names a column
the statement never mentions and says nothing about `requiredDataAttributes()`.
The second arm is worse. When the constraint's column *is* in the narrow
scan, because the condition or an assignment RHS references it, the delta path
still drops it from the write query in the `Project` below (`:320-323`), and
the Resolution batch never converges:
| connector | statement | result |
|---|---|---|
| `column-update` (MoR) | `SET id = -1 WHERE extra > 3` | `RuntimeException:
Max iterations (100) reached for batch Resolution` |
| `column-update` (MoR) | `SET id = extra + 1 WHERE pk = 1` | same |
| `column-update-cow` (CoW) | `SET id = -1 WHERE extra > 3` | works |
The CoW row is the control for that pair:
`buildNarrowReplaceDataUpdateProjection` maps all of `plan.output`, so the
column survives into `query.output` there, and only the delta builder drops it.
I applied the matching fix locally and it holds. Carrying the remaining
narrow scan columns through the `Project`:
```scala
val connectorIds = connectorDataAttrs.map(_.exprId).toSet
val carryAlong = plan.output.filterNot { a =>
MetadataAttribute.isValid(a.metadata) || rowIdAttrSet.contains(a) ||
assignedKeyIds.contains(a.exprId) || connectorIds.contains(a.exprId)
}
Project(
Seq(operationType) ++ assignedValues ++ connectorPassThroughValues ++
metadataValues ++ rowIdValues ++ originalRowIdValues ++ carryAlong,
plan)
```
turns both looping arms green and keeps the payload narrow: `updateSchema`
stays `struct<pk, dep, id>` and `extra` does not leak into it, because
`updateRowProjection` picks `connectorDataAttrs` out of the `Project` by name.
`DeltaBasedColumnUpdateTableSuite` + `GroupBasedColumnUpdateTableSuite` stay
43/43. And the constraint is genuinely enforced again afterwards: with `CHECK
(id > 0)` instead, `SET id = -1 WHERE extra > 3` gives
`CHECK_CONSTRAINT_VIOLATION`.
That still leaves the first arm, where the constraint column is not in the
narrow scan at all. Reading it is the only way Spark can re-validate the
constraint on the new row, so `computeNarrowReadAttrs` would have to include
the `Check` constraints' references as well. The alternative is to refuse to
narrow a table that carries `Check` constraints, with a clear error. I would
widen rather than refuse, and note in the javadoc that a table with CHECK
constraints narrows less.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,57 @@
+/*
+ * 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>
+ * Currently honored only for UPDATE. DELETE and MERGE ignore this interface:
those operations
+ * receive full-width rows and {@link LogicalWriteInfo#updateSchema()} is
absent for them.
+ *
+ * @since 4.3.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level
operation.
+ * <p>
+ * The returned columns become the schema of updated, copied, and reinserted
rows, in declared
+ * order. Implementations must include every column they want to receive
(typically the columns
+ * reported by {@link RowLevelOperationInfo#updatedColumns()} plus any
columns needed for row
+ * lookup or routing, e.g. a primary key).
+ * <p>
+ * If any of the columns from {@link RowLevelOperationInfo#updatedColumns()}
are
+ * missing, an analysis exception is thrown.
+ * <p>
+ * For updates on nested fields such as {@code SET t.s.c1 = -1} the
connector should declare the
+ * root struct column {@code s} rather than any nested field.
+ * <p>
+ * This also covers columns needed only for planning, e.g. resolving the
table's partitioning
Review Comment:
**Finding 17 (regressed).** Dropping `validatePartitionAttrsDeclared` is a
fair policy call, and I agree Spark should not decide what a connector wants.
But the thing that breaks when a partition source column is missing is not the
connector's write requirement. It is Spark's own SPJ reporting rule, and it
throws instead of degrading.
This paragraph reads as advisory, "also covers columns needed only for
planning". For any partitioned table whose scan reports
`KeyGroupedPartitioning` it is mandatory. Measured on `72ec505`,
`requiredDataAttributes() = [pk, id]`, table `pk INT NOT NULL, id INT, dep
STRING` partitioned by `identity(dep)`, `UPDATE t SET id = -1 WHERE pk = 1`:
```
org.apache.spark.sql.AnalysisException: Unable to resolve dep given
[pk,id,_partition,index].
at
QueryCompilationErrors$.cannotResolveAttributeError(QueryCompilationErrors.scala:1988)
at V2ExpressionUtils$.resolveRef(V2ExpressionUtils.scala:56)
at V2ExpressionUtils$.toCatalystTransformOpt(V2ExpressionUtils.scala:147)
at V2ExpressionUtils$.toCatalystOpt(V2ExpressionUtils.scala:128)
at V2ScanPartitioningAndOrdering$$anonfun$partitioning$2 ...
(V2ScanPartitioningAndOrdering.scala:51)
```
The same connector against an unpartitioned table passes, so the
partitioning is the trigger.
This also corrects what I wrote at round 5
([r3735139324](https://github.com/apache/spark/pull/55518#discussion_r3735139324)),
where I said the read side would "silently report no partitioning". It does
not. `V2ScanPartitioningAndOrdering.partitioning` resolves the keys against
`scanRelation.relation` (`ExtractV2ScanInfo`,
`DataSourceV2Relation.scala:457-461`), and that relation is exactly what this
PR narrows. `toCatalystTransformOpt`'s `IdentityTransform` case then calls the
throwing `resolveRef`. Both the `Opt` naming and the comment two lines below
the call site ("Keep the partitioning when at least one of its keys is still in
the scan output") say the rule means to tolerate a pruned key. It never gets
the chance. The write-side site I measured at round 4
([r3735141175](https://github.com/apache/spark/pull/55518#discussion_r3735141175))
is still reachable too, for a connector that clusters on a data column.
Either fix closes the read side:
- resolve the keys with a non-throwing lookup and drop the partitioning when
a key is absent. This asks nothing of connectors and is what the rule already
intends;
- or have `computeNarrowReadAttrs` add the table's partition-transform
references to the *scan*. This is not restrictive the way the removed guard
was: the payload is projected down to `connectorDataAttrs` by name, so a
partition column sitting in the scan never reaches the writer. I confirmed that
projection behaviour while measuring finding 27, where a carried-along column
stayed out of `updateSchema()`.
Whichever you take, this sentence should say the declaration is required
rather than that it "covers" the case, and there should be a test for a
column-update UPDATE on a partitioned table whose partition column is not
declared. Every fixture in the suite declares `dep`, which is why 43 tests pass
over this.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +84,51 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated, copied, or reinserted record with metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} receive records here
in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must
override this
+ * method when mixing in {@link SupportsColumnUpdates}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be
called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like
writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but does not override this method.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T metadata, T record) throws IOException {
+ throw new SparkUnsupportedOperationException(
+ "DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED",
+ Map.of("class", getClass().getName()));
+ }
Review Comment:
**Finding 29.** A connector that implements only `writeUpdate(record)` still
fails here.
Which overload Spark calls depends on whether the operation declares
required metadata attributes (`ReplaceDataExec.writingTask`,
`WriteToDataSourceV2Exec.scala:385-397`), which is not visible from the
interface. The single-argument form's doc says it is "Equivalent to
`writeUpdate(Object, Object)` for writers that do not require metadata", so
implementing that one alone reads as sufficient, and it is not.
The sibling pair a few lines up already solves this: `write(T metadata, T
record)` defaults to `write(record)` (`:83-85`). Mirroring it keeps the two
channels consistent:
```suggestion
writeUpdate(record);
}
```
A connector that overrides neither form still gets
`DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED`, from the single-argument default.
The trade-off is the one `write` already makes: a writer that needs the
metadata must override the two-argument form.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,684 @@
+/*
+ * 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 org.apache.spark.sql.Row
+import org.apache.spark.sql.connector.catalog.{CatalogV2Util, TableInfo}
+import
org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity,
reference}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField,
StructType}
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, Spark narrows the update-row
projection
+ * (LogicalWriteInfo.updateSchema()) to contain only the declared columns
rather than
+ * the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends RowLevelOperationSuiteBase {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] =
{
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ test("column-update: rowSchema contains only the single assigned column") {
+ 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)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ // info.schema() is empty: UPDATE-only column-update writes have no
INSERT-shaped rows,
+ // so nothing flows through the writer's write() path. The narrow
update-row layout is
+ // carried by info.updateSchema().
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: rowSchema contains multiple assigned columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'engineering' WHERE pk
= 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false),
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: rowSchema is empty for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ // All assignments are identity, so updatedColumns is empty -- the
connector still declares
+ // pk (row lookup) and dep (write-side clustering key), so the narrow
update schema is
+ // just [pk, dep].
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Array(PK_FIELD, DEP_FIELD))))
+ }
+
+ test("column-update: row filter condition is orthogonal to column
narrowing") {
+ 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)
+
+ sql(s"UPDATE $tableNameAsString SET dep = 'engineering' WHERE pk IN (1,
3)")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: update all rows (no WHERE clause)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("salary", IntegerType, nullable = true)
+ ))))
+ }
+
+ test("column-update: rowSchema excludes identity assignments in a mixed
UPDATE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk
= 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: cross-column assignment is not treated as identity") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = dep, id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: nested struct field update narrows to the root struct
column") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep
STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("s"),
+ s"expected [s] in updatedColumns (root struct) but got: $updatedNames")
+
+ // info.updateSchema() carries the narrow row layout; info.schema() is the
full table.
+ // `dep` is present because the connector also declares it as the
write-side clustering key.
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("s"),
+ s"s must be in update schema: $updateSchema")
+ assert(updateSchema.fieldNames.contains("dep"),
+ s"dep must be in update schema: $updateSchema")
+ }
+
+ test("column-update: nested field identity update reports root struct as
updated") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep
STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = s.c1 WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("s"),
+ s"nested identity is reported as an update at root-column granularity:
$updatedNames")
+
+ // Data correctness: the struct is rewritten but with equal values, so
rows are unchanged.
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, Row(1, 2), "hr") :: Nil)
+ }
+
+ test("column-update: updatedColumns contains non-identity assigned columns")
{
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'eng' WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("id", "dep"),
+ s"expected [id, dep] in updatedColumns but got: $updatedNames")
+ }
+
+ test("column-update: updatedColumns excludes identity assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = dep WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("id"),
+ s"expected only [id] in updatedColumns but got: $updatedNames")
+ }
+
+ test("column-update: updatedColumns is empty for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ assert(table.lastUpdatedColumns.isEmpty,
+ s"expected empty updatedColumns but got:
${table.lastUpdatedColumns.mkString(", ")}")
+ }
+
+ test("column-update: updatedColumns is empty for DELETE (Javadoc contract)")
{
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"DELETE FROM $tableNameAsString WHERE dep = 'hr'")
+
+ assert(table.lastUpdatedColumns.isEmpty,
+ s"DELETE must pass empty updatedColumns but got:
${table.lastUpdatedColumns.mkString(", ")}")
+ }
+
+ test("column-update: data correctness -- single column update") {
+ 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)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ // 1 row matched WHERE pk = 1 -> 1 update. MoR emits no COPY rows.
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0, deltaUpdate =
true)
+ }
+
+ test("column-update: data correctness -- update all rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 200, "hr") :: Row(2, 400, "software") :: Row(3, 600, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 3, numCopiedRows = 0, deltaUpdate =
true)
+ }
+
+ test("column-update: data correctness -- mixed identity and real
assignments") {
+ 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)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk
= 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "engineering") :: Row(2, 2, "software") :: Row(3, 3, "hr") ::
Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0, deltaUpdate =
true)
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-req-attrs", reqAttrs)
+ val columns =
CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update: requiredDataAttributes - data correctness") {
+ createAndInitTableWithReqAttrs("dep,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)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ 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 on
the row-ID " +
+ "column without adopting the declared spelling") {
+ // The connector declares `PK` (not `pk`); resolution must match
case-insensitively but the
+ // narrow scan/write must still use the table's own column name, `pk`, not
the connector's
+ // declared spelling -- otherwise column pruning can't find `PK` in the
physical scan.
+ 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")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("pk"),
+ s"update schema must use the table's own spelling `pk`, not `PK`:
$updateSchema")
+ assert(!updateSchema.fieldNames.contains("PK"),
+ s"update schema must not adopt the connector's declared spelling `PK`:
$updateSchema")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively on
an assigned " +
+ "column without adopting the declared spelling") {
+ // Same as above, but the case mismatch is on `salary` (the assigned
column) rather than
+ // the row-ID column.
+ 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")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("salary"),
+ s"update schema must use the table's own spelling `salary`, not
`SALARY`: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("SALARY"),
+ s"update schema must not adopt the connector's declared spelling
`SALARY`: $updateSchema")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: empty requiredDataAttributes throws AnalysisException")
{
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-empty-req-attrs", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(
+ StructType.fromDDL("pk INT NOT NULL, id INT, dep STRING"))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }""".stripMargin)
+
+ // A connector that mixes in SupportsColumnUpdates but returns an empty
+ // requiredDataAttributes() violates the mix-in contract -- analysis must
reject the
+ // operation rather than silently falling through to the wide write path.
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ }
+ assert(ex.getCondition == "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ s"expected COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES but got:
${ex.getCondition}")
+ }
+
+ test("column-update: requiredDataAttributes throws AnalysisException for
invalid column") {
+ createAndInitTableWithReqAttrs("nonexistent_col", "pk INT NOT NULL, id
INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ }
+ assert(ex.getMessage.contains("nonexistent_col"),
+ s"Expected error about unresolvable column but got: ${ex.getMessage}")
+ }
+
+ private def createAndInitTableFromInfo(schemaString: String, jsonData:
String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-from-info", "true")
+ val columns =
CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update from-info: write schema is updatedColumns + pk/dep
pass-through") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, id INT, dep
STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("salary"),
+ s"salary must be in update schema: $updateSchema")
+ assert(updateSchema.fieldNames.contains("pk"),
+ s"pk must be in update schema: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("id"),
+ s"id must not be in update schema: $updateSchema")
+ assert(updateSchema.fieldNames.contains("dep"),
+ s"dep must be in update schema: $updateSchema")
+ }
+
+ test("column-update from-info: pk already in updatedColumns is not
duplicated") {
+ createAndInitTableFromInfo("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 pk = pk + 10, salary = -1 WHERE dep =
'hr'")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ val pkCount = updateSchema.fieldNames.count(_ == "pk")
+ assert(pkCount == 1, s"pk must appear exactly once in update schema:
$updateSchema")
+ }
+
+ test("column-update from-info: data correctness") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, id INT, dep
STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "id": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ 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)
+ }
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String):
Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-split", "true")
+ val columns =
CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update split: write schema is narrow (assigned + pk/dep
pass-through)") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ // For representUpdateAsDeleteAndInsert connectors, reinsert-tagged rows
still route through
+ // writeUpdate() (info.updateSchema() = narrow), not write().
info.schema() stays empty.
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD,
INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update split: data correctness") {
+ createAndInitTableSplit("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)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, -1, "hr") :: Nil)
+ }
+
+ test("column-update split: row-ID reassignment on narrow write is rejected")
{
+ // `extra` is neither declared, updated, nor referenced, so
requiredDataAttributes()
+ // ([pk, dep, salary]) doesn't cover every table column -- the write stays
genuinely narrow
+ // and reassigning the row ID must be rejected.
+ 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)
+
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep
= 'hr'")
+ }
+ assert(ex.getCondition == "COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT",
+ s"expected COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT but got:
${ex.getCondition}")
+ }
+
+ test("column-update split: row-ID reassignment is allowed when
requiredDataAttributes " +
+ "covers every table column") {
+ // Finding 22: the SPLIT_UPDATE_ROW_ID_REASSIGNMENT message advertises a
second remedy --
+ // "include every table column in requiredDataAttributes()" -- but the
guard never checked
+ // for it. When the declaration covers the whole relation, the REINSERT
payload already is
+ // the full row with the new row-ID value, and the DELETE half still
carries the original
+ // row-ID via newLazyRowIdProjection, so reassignment is safe and must be
allowed.
+ createAndInitTableSplitWithReqAttrs("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 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)
+ }
+
+ private def createAndInitTableSplitWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-split-req-attrs", reqAttrs)
+ val columns =
CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update split: undeclared row-ID column on narrow write is
rejected") {
+ // The connector's requiredDataAttributes() excludes `pk` (the row ID).
The REINSERT
+ // payload for the split-update path is projected down to
requiredDataAttributes(), so
+ // without `pk` there the reinserted row would have no identity for the
connector to
+ // place it by. This must be rejected at analysis time, not discovered at
write time.
+ createAndInitTableSplitMissingRowId("pk INT NOT NULL, salary INT, dep
STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+ }
+ assert(ex.getCondition == "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED",
+ s"expected COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED but got:
${ex.getCondition}")
+ assert(ex.getMessage.contains("pk"),
+ s"error message must name the undeclared row ID column `pk`:
${ex.getMessage}")
+ }
+
+ private def createAndInitTableSplitMissingRowId(
+ schemaString: String, jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-split-missing-row-id", "true")
+ val columns =
CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ //
---------------------------------------------------------------------------
+ // Scan-narrowing tests: verify that when the connector opts into column
updates,
+ // the physical scan reads only the columns actually needed by the operation.
+ //
+ // Assertions use `checkLastScanExcludes` (proves narrowing) and
`checkLastScanIncludes`
+ // (proves transparent widening) rather than exact-set matches --
ColumnPruning is
+ // free to further tighten the scan beyond our analysis-time narrow set
(e.g. when an
+ // assignment RHS is a literal, the assigned column doesn't need to be read).
+ //
---------------------------------------------------------------------------
+
+ test("column-update: scan excludes columns outside connector-declared + cond
+ RHS refs") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, extra STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr", "extra": "x" }
+ |{ "pk": 2, "id": 2, "dep": "software", "extra": "y" }
+ |""".stripMargin)
+
+ // requiredDataAttributes = [pk, id]; cond refs [pk] only; RHS is a
literal.
+ // `dep` is scan-only declared (needed for scan-side partitioning
resolution) so it must
Review Comment:
**Finding 30.** Wording left over from the removed
`scanOnlyDataAttributes()`. There is no scan-only category any more: `dep` is
in `requiredDataAttributes()` like every other declared column.
Here, and again at `:654-657` ("`dep` is scan-only declared (needed for
scan-side partitioning resolution) so it also appears; `extra` is neither
declared, scan-only, nor referenced"). Also
`InMemoryRowLevelOperationTable.scala:303`, which calls `dep` "the write-side
clustering key, see `clusterColumnRef`" while `clusterColumnRef` is
`PARTITION_COLUMN_REF` (`_partition`), a metadata column, not `dep`.
Two test names in this file no longer match what they assert, same cause.
"rowSchema contains only the single assigned column" (`:42`) asserts `[pk, dep,
id]`, and "rowSchema is empty for a full identity update" (`:82`) asserts `[pk,
dep]`. "rowSchema" is also the pre-redesign name for what is now
`updateSchema()`.
--
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]
