anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4068060968
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -109,6 +120,22 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
relation)
}
+ /**
+ * Resolves the connector-declared required data attributes for
column-update writes against
+ * the given relation. Non-existent columns are rejected with an
`AnalysisException`.
+ */
+ protected def resolveRequiredDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes
+ if (refs.isEmpty) {
Review Comment:
Done. Added a new error type when nested columns are required and added a
test for it.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -226,4 +421,186 @@ 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. Returns `Nil` otherwise.
+ */
+ private def resolveConnectorDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation): Seq[AttributeReference] = operation match
{
+ case scu: SupportsColumnUpdates => resolveRequiredDataAttrs(relation, scu)
+ case _ => Nil
+ }
+
+ /**
+ * Resolves the connector's `scanOnlyDataAttributes()` if the operation opts
into column
+ * updates. Returns `Nil` otherwise.
+ */
+ private def resolveScanOnlyDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation): Seq[AttributeReference] = operation match
{
+ case scu: SupportsColumnUpdates =>
+ V2ExpressionUtils.resolveRefs[AttributeReference](
+ scu.scanOnlyDataAttributes.toImmutableArraySeq, relation)
+ case _ => Nil
+ }
+
+ /**
+ * Computes the narrow set of data columns that must be present in the scan
for a column-update
+ * write: connector-declared attrs (both `requiredDataAttributes()` and
+ * `scanOnlyDataAttributes()`), unioned with any table columns referenced by
non-identity
+ * assignment RHS expressions and the operation condition.
+ */
+ private def computeNarrowReadAttrs(
+ relation: DataSourceV2Relation,
+ connectorDataAttrs: Seq[AttributeReference],
+ scanOnlyDataAttrs: 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)
+ .collect { case a: AttributeReference => a }
+ .filter(relationSet.contains)
+ dedupAttrs(connectorDataAttrs ++ scanOnlyDataAttrs ++ extraRefs)
+ }
+
+ /**
+ * Enforces that every column being assigned (non-identity) is present in
the connector-declared
+ * `requiredDataAttributes()`. Comparison is at root-column granularity
+ */
+ private def validateUpdatedColumnsSubset(
+ operation: RowLevelOperation,
+ assignments: Seq[Assignment],
+ connectorDataAttrs: Seq[AttributeReference]): Unit = {
+ val declaredIds = connectorDataAttrs.map(_.exprId).toSet
+ val missing = assignments.collect {
+ case Assignment(key: AttributeReference, value)
+ if !isIdentityAssignment(key, value) &&
!declaredIds.contains(key.exprId) =>
+ key.name
+ }.distinct
+ if (missing.nonEmpty) {
+ throw
QueryCompilationErrors.requiredDataAttributesMissingUpdatedColumnsError(
+ operation.getClass.getName, missing)
+ }
+ }
+
+ /**
+ * Enforces that `requiredDataAttributes()` and `scanOnlyDataAttributes()`
are disjoint.
+ */
+ private def validateNoOverlap(
+ operation: RowLevelOperation,
+ connectorDataAttrs: Seq[AttributeReference],
+ scanOnlyDataAttrs: Seq[AttributeReference]): Unit = {
+ val requiredIds = connectorDataAttrs.map(_.exprId).toSet
+ val overlapping = scanOnlyDataAttrs.collect {
+ case attr if requiredIds.contains(attr.exprId) => attr.name
+ }.distinct
+ if (overlapping.nonEmpty) {
+ throw
QueryCompilationErrors.requiredDataAttributesOverlapScanOnlyAttributesError(
+ operation.getClass.getName, overlapping)
+ }
+ }
+
+ /**
+ * Enforces that every partition-source column is declared in either
`requiredDataAttributes()`
+ * or `scanOnlyDataAttributes()`. Spark no longer adds partition refs to the
narrow scan
+ * implicitly, so a connector that needs them for partitioning resolution or
write-side
+ * clustering must declare them in one of the two methods.
+ */
+ private def validatePartitionAttrsDeclared(
+ operation: RowLevelOperation,
+ relation: DataSourceV2Relation,
+ connectorDataAttrs: Seq[AttributeReference],
+ scanOnlyDataAttrs: Seq[AttributeReference]): Unit = {
+ val partitionRefNames = relation.table.partitioning().toImmutableArraySeq
+ .flatMap(_.references.toImmutableArraySeq)
+ .map(_.fieldNames.head)
+ .toSet
+ val declared = AttributeSet(connectorDataAttrs ++ scanOnlyDataAttrs)
+ val undeclared = relation.output
+ .filter(a => partitionRefNames.exists(name => conf.resolver(name,
a.name)))
+ .filterNot(declared.contains).map(_.name)
+ if (undeclared.nonEmpty) {
+ throw
QueryCompilationErrors.requiredDataAttributesMissingPartitionColumnsError(
+ operation.getClass.getName, undeclared)
+ }
+ }
+
+ /**
+ * For connectors that opt into narrow column updates AND represent UPDATE
as delete + insert,
+ * reject reassignment of any row-ID column as the REINSERT path has no
row-ID channel to
+ * reconstruct columns outside `requiredDataAttributes()`.
+ */
+ private def validateNoRowIdReassignment(
Review Comment:
This has been fixed in an earlier commit 33b3f071352
##########
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:
Good catch. For constraints I did both read side and the write side changes
to include the constraint columns. They are discarded by the update projects
before the write so it should still pass narrow writes to the connector.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -253,6 +304,382 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends
only the
+ // declared + assigned columns in the row projection instead of the full row
schema. The base
+ // class composes its required-attrs set as `pk` (the row-lookup key) plus
whatever columns
+ // Spark reports as being assigned via
`RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil)
+ extends DeltaBasedOperation(command, CaseInsensitiveStringMap.empty())
Review Comment:
Done.
##########
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:
Thanks for testing this. I like option 1 but I think that change belongs to
it is own PR as it touches code unrelated to this PR. If you agree, I will
create a separate JIRA for this. I would also defer adding the test that skips
adding the partition column since it will fail now.
##########
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:
I took the suggestion and added java docs to reflect this.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,683 @@
+/*
+ * 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 {
Review Comment:
Good call. Just to see what it would break, I made changes.
`DeltaBasedColumnUpdateTableSuite` now extends
`DeltaBasedUpdateTableSuiteBase`, `GroupBasedColumnUpdateTableSuite` now
extends `UpdateTableSuiteBase`.
I found that only 1 test in each suit failed and it was due to the default
values. So I fixed the `doCommit` schema-indexing bug in all three affected
spots. Each now checks `i < base.numFields` and falls back to
`ResolveDefaultColumns.getExistenceDefaultValue(field)` for columns added after
a row was written, matching the exact pattern already used elsewhere in
`InMemoryBaseTable.scala` for the same problem.
I decided to make this small change in this PR.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -109,6 +120,22 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
relation)
}
+ /**
+ * Resolves the connector-declared required data attributes for
column-update writes against
+ * the given relation. Non-existent columns are rejected with an
`AnalysisException`.
+ */
+ protected def resolveRequiredDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes
+ if (refs.isEmpty) {
+ throw
QueryCompilationErrors.emptyRequiredDataAttributesError(operation.getClass.getName)
Review Comment:
I added duplicate check and associated error message.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,68 @@
+/*
+ * 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
Review Comment:
Good points. Addressed both.
##########
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()));
+ }
+
+ /**
+ * Writes one updated, copied, or reinserted record without metadata.
+ * <p>
+ * Equivalent to {@link #writeUpdate(Object, Object)} for writers that do
not require metadata.
+ * 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 record) throws IOException {
Review Comment:
I have added a test `column-update ReplaceData: connector missing
writeUpdate override is rejected` in `GroupBasedColumnUpdateTableSuite.scala`
##########
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:
Done. I fixed all the mentioned places.
--
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]
