dongjoon-hyun commented on code in PR #55518: URL: https://github.com/apache/spark/pull/55518#discussion_r4315467514
########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java: ########## @@ -0,0 +1,94 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector.write; + +import org.apache.spark.annotation.Experimental; +import org.apache.spark.sql.connector.expressions.NamedReference; + +/** + * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to + * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()} + * for updated, copied, and reinserted records, instead of the full table row. + * <p> + * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, updated, copied, and reinserted + * rows follow it: group-based operations receive them through + * {@link DataWriter#writeColumnUpdate(Object, Object)} or + * {@link DataWriter#writeColumnUpdate(Object)}, and operations that mix in {@link SupportsDelta} + * through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert}. Inserted rows always follow + * {@link LogicalWriteInfo#schema()} and arrive through {@link DataWriter#write(Object)} and + * {@link DeltaWriter#insert}. When {@link LogicalWriteInfo#columnUpdateSchema()} is absent, every + * row follows {@link LogicalWriteInfo#schema()}, as for any other operation. Connectors must + * decide which layout to expect by whether {@link LogicalWriteInfo#columnUpdateSchema()} is + * present, not by the command. A builder that wants full-width rows for some command can check + * {@link RowLevelOperationInfo#command()} and build an operation without this mix-in. Currently + * Spark narrows only UPDATE. + * <p> + * The scan builder returned by {@link #newScanBuilder} should implement + * {@link org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns}, so that the scan + * does not read columns the command neither references nor declares. Otherwise the scan reads + * every column. + * + * @since 4.4.0 + */ +@Experimental +public interface SupportsColumnUpdates extends RowLevelOperation { + /** + * Returns the data column references required to perform this row-level operation. + * <p> + * When Spark narrows the write, the returned columns become + * {@link LogicalWriteInfo#columnUpdateSchema()}, in declared order. Implementations must include + * every column they want to receive: every column reported by + * {@link RowLevelOperationInfo#updatedColumns()}, plus any columns needed for row lookup or + * routing, e.g. a primary key. Columns that are not declared are absent from the rows the Review Comment: **5. A copy-on-write column update that assigns the connector's lookup column leaves the connector no way to find the original row.** For a group-based operation, an updated row carries only the new values of the declared columns, and metadata columns that are not preserved on update are nulled (`RewriteUpdateTable.scala:146-150`), so the original value of an assigned column reaches the connector through no channel. This paragraph still suggests "a primary key" for row lookup, and Spark guards the analogous case only for split delta operations (`validateNoRowIdReassignment`). The in-repo fixture shows the result. On a `column-update-cow` table `pk INT NOT NULL, id INT, dep STRING` with rows (1, 1, hr), (2, 2, software) and (3, 3, hr), `UPDATE t SET pk = 10 WHERE id = 1` delivers `[pk=10, dep=hr]` with metadata `[hr, null]`. `PartitionBasedNarrowReplaceData.doCommit` looks the base row up by the new `pk` (`InMemoryRowLevelOperationTable.scala:736-738`), finds none, and commits (10, null, hr) instead of (10, 1, hr). `SET pk = 3 WHERE id = 1` silently copies `id` from the other row and commits (3, 3, hr). The delta path is fine thanks to `__original_row_id_pk`, but `GroupBasedColumnUpdateTableSuite` does not inherit "update with assignments to row ID" from `DeltaBasedUpdateTableSuiteBase`, so no test covers this. Could you document that a group-based connector has to locate rows through metadata preserved on update, or build an operation without this mix-in when `updatedColumns()` contains its lookup column? Could the fixture also fail loudly when the lookup misses, and could the copy-on-write suite get a port of that test? ########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/LogicalWriteInfo.java: ########## @@ -65,4 +69,18 @@ default Optional<StructType> metadataSchema() { throw new SparkUnsupportedOperationException( "DATA_SOURCE_METADATA_SCHEMA_NOT_IMPLEMENTED", Map.of("class", getClass().getName())); } + + /** + * the schema of updated, copied, and reinserted rows from Spark to data source in a + * column-level update. Present when the operation mixes in {@link SupportsColumnUpdates} and + * Spark delivers these rows with the columns of + * {@link SupportsColumnUpdates#requiredDataAttributes()}, which currently happens only for + * UPDATE. It covers every table column if every column is declared. When present, + * {@link #schema()} covers only newly inserted rows. + * + * @since 4.4.0 + */ + default Optional<StructType> columnUpdateSchema() { Review Comment: **10.** Minor: `columnUpdateSchema()` silently defaults to `Optional.empty()`, while its siblings `rowIdSchema()` and `metadataSchema()` (L60-71) throw `DATA_SOURCE_*_SCHEMA_NOT_IMPLEMENTED`. For a narrow UPDATE, Spark sends an empty `schema()`. A connector-side wrapper that rebuilds the info, the way `FileTable.mergedWriteInfo` does (this PR had to add the field there), and forwards the other methods but not this one makes the connector see "absent". By the documented rule it then decodes the narrow rows of `DeltaWriter.update` and `reinsert` with the empty `schema()` layout, silently. Spark reads this method only on infos it builds itself, so a throwing default would not affect any Spark path. Could the default throw like its siblings? ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala: ########## @@ -253,6 +326,24 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] { WriteDeltaProjections(rowProjection, rowIdProjection, metadataProjection) } + /** + * Moves the query columns that no projection reads after the ones that are read, and returns + * the reordered query with a map from each old ordinal to the new one. The optimizer prunes the + * query columns after the last one the write reads, which then never changes an ordinal. + */ + protected def orderReadColumnsFirst( Review Comment: **11.** Minor (design, optional): keeping the ordinal-based projections valid under pruning takes three cooperating pieces: this analysis-time reorder plus `remapOrdinals` at three call sites in `RewriteUpdateTable`, the prefix `references` in `RowLevelWrite` (`v2Commands.scala:367-370`), and the new `ColumnPruning` case. All of them rely on "read columns come first", which only `RewriteUpdateTable` establishes, and only for the plan shapes `reorderOutputs` handles. `extractOutputs` handles `MergeRows` (L351), but `reorderOutputs` throws for it (L364), which matters for the MERGE follow-up in the PR description. A path that skips the reorder does not fail either; the prefix `references` silently keeps every column up to the last read ordinal. Could `references` be exactly the attributes at the read ordinals (plus ordinal 0), with the `ColumnPruning` case remapping the projections when it prunes, as it already does for `Generate.unrequiredChildIndex` (`Optimizer.scala:1543-1549`)? The optimized plan would be the same, without the analysis-time reorder and its shape restriction. ########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java: ########## @@ -37,4 +38,16 @@ public interface RowLevelOperationInfo { * Returns the row-level SQL command (e.g. DELETE, UPDATE, MERGE). */ Command command(); + + /** + * Returns the columns being updated by this operation. Currently only UPDATE populates it; + * other commands report an empty array. + * <p> + * A column is reported only if it is assigned a new value, so identity assignments such as Review Comment: **8. An explicit `SET v = v` on a CHAR/VARCHAR column is reported here.** With the default CHAR/VARCHAR handling, assignment alignment wraps the value in a write-side length check, `Alias(StaticInvoke(varcharTypeWriteSideCheck, Cast(v, string), 10))` (`TableOutputResolver.checkUpdate`), and `isIdentityAssignment` strips only one `Alias` before `semanticEquals`, so it does not recognize the identity. For `t(pk INT NOT NULL, salary INT, v VARCHAR(10))` and a connector that declares `[pk, salary]`, `UPDATE t SET salary = 1, v = v` fails with `COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS` and `[v]`, although `v` cannot change. CHAR(n) behaves the same, and on a split delta operation `SET rid = rid` on a VARCHAR row-ID column is rejected with `COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT`. Unassigned CHAR/VARCHAR columns stay bare attributes, so only explicit self-assignments are affected, and both column-update fixtures declare whatever Spark reports, so no test notices. Could this Javadoc say that an assignment is excluded only when its aligned value is the column itself, and that `SET v = v` on CHAR/VARCHAR (like `SET s.c1 = s.c1`) may still be reported? Seeing through the wrappers would skip the padding and length checks, so documenting it looks safer. ########## sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala: ########## @@ -373,6 +904,21 @@ private class DeltaBufferWriter(schema: StructType) extends BufferWriter(schema) } object InMemoryRowLevelOperationTable { + // Global holder for the scan schema of the last row-level operation. See the class-level + // `lastScanSchema` def for why this is a companion-object field rather than a per-instance one. + // Tests that assert on scan pruning should reset this before running the operation under test + // (see `RowLevelOperationSuiteBase.beforeEach`). + private[catalog] val lastScanSchemaRef = new AtomicReference[StructType]() Review Comment: **15.** Minor: this JVM-global, last-writer-wins `AtomicReference`, hooked into the existing `PartitionBasedOperation` and `DeltaBasedOperation` scan builders and reset in `before {}`, duplicates `executeAndCheckScan(s)` in `RowLevelOperationSuiteBase`, which asserts the same `scan.readSchema()` from the statement's own executed plan; `GroupBasedColumnUpdateTableSuite.scala:109`, `:285` and `:344` already use it. It also keeps the previous statement's value: a statement that plans no scan, e.g. a metadata-only `DELETE ... WHERE dep = 'hr'`, leaves the last UPDATE's schema behind, and `GroupBasedColumnUpdateTableSuite.scala:430-457` checks UPDATE, DELETE and MERGE against one expected string. This is the same vacuous-pass class as round 4 ([r4074793084](https://github.com/apache/spark/pull/55518#discussion_r4074793084)), which was fixed only for `null`. Two more one-sided checks: `checkLastScanSchema` and `executeAndCheckScan` compare with `sameType`, which ignores field-name case by default, so the "table's own spelling" tests (`DeltaBasedColumnUpdateTableSuite.scala:395-408`, `GroupBasedColumnUpdateTableSuite.scala:660-675`) cannot detect a spelling regression, and an expected-empty `checkLastUpdatedColumns()` matches the initial `Array.empty`. Could the suites use `executeAndCheckScan(s)` and drop this global? ########## common/utils/src/main/resources/error/error-conditions.json: ########## @@ -1253,6 +1253,60 @@ ], "sqlState" : "22003" }, + "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE" : { + "message" : [ + "Connector <connector> mixes in `SupportsColumnUpdates` but declared duplicate column(s) <duplicateAttributes> in `requiredDataAttributes()`. Each column must be declared at most once." + ], + "sqlState" : "42000" + }, + "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES" : { + "message" : [ + "Connector <connector> mixes in `SupportsColumnUpdates` but returned an empty array from `requiredDataAttributes()`. Connectors that opt into column-level updates must declare at least one required data attribute." + ], + "sqlState" : "42000" + }, + "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE" : { + "message" : [ + "Connector <connector> mixes in `SupportsColumnUpdates` but declared metadata column(s) <metadataAttributes> in `requiredDataAttributes()`. Declare metadata columns in `requiredMetadataAttributes()` instead." + ], + "sqlState" : "42000" + }, + "COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE" : { + "message" : [ + "Connector <connector> mixes in `SupportsColumnUpdates` but declared nested field(s) <nestedAttributes> in `requiredDataAttributes()`. Declare the root struct column instead of a nested field." + ], + "sqlState" : "42000" + }, + "COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS" : { + "message" : [ + "Connector <connector> mixes in `SupportsColumnUpdates` but its `requiredDataAttributes()` does not cover every column being updated. Missing columns: <missingColumns>. Connectors must include every column reported by `RowLevelOperationInfo.updatedColumns()` in `requiredDataAttributes()`." + ], + "sqlState" : "42000" + }, + "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED" : { + "message" : [ + "Connector <connector> mixes in `SupportsColumnUpdates` and represents UPDATE as delete and insert, but row ID column(s) <rowIds> do not reach the reinserted row, which then has no identity for the connector to place it by. Declare each data row ID column in `requiredDataAttributes()`, and return each metadata row ID column from `requiredMetadataAttributes()` with `MetadataColumn.PRESERVE_ON_REINSERT` set." + ], + "sqlState" : "42000" + }, + "COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT" : { + "message" : [ + "Connector <connector> mixes in `SupportsColumnUpdates` and represents UPDATE as delete and insert, so UPDATE cannot assign row ID column(s) <rowIds>. The reinserted row carries only the columns in `requiredDataAttributes()`, and with a new row ID it cannot be matched to the row whose other columns the connector must preserve. Do not assign row ID columns, or include every table column in `requiredDataAttributes()`." + ], + "sqlState" : "42000" + }, + "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS" : { + "message" : [ + "Connector <connector> mixes in `SupportsColumnUpdates`, but its write requires a distribution or ordering by column(s) <columns> that the column-level UPDATE does not read. Declare each data column in `requiredDataAttributes()` and each metadata column in `requiredMetadataAttributes()`." Review Comment: **12.** Minor: "that the column-level UPDATE does not read" is not accurate. The check compares against the columns the write receives (`rowLevelWrite.references`), so it also fires for an undeclared column that the UPDATE reads: `GroupBasedColumnUpdateTableSuite.scala:401-409` expects this error with `[dep]` for `... WHERE dep = 'hr'`, where the scan reads `dep`. The Javadoc in `SupportsColumnUpdates.java:76-84` already describes it correctly. Could the message say something like "by column(s) <columns> that are not passed to the write"? ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala: ########## @@ -228,4 +240,114 @@ object RewriteUpdateTable extends RewriteRowLevelCommand { val expandOutput = generateExpandOutput(attrs, outputs) Expand(outputs, expandOutput, matchedRowsPlan) } + + /** + * Returns the output of the write relation, which is the layout of the data rows the connector + * receives. If the operation opts into column updates, this is the connector's + * `requiredDataAttributes()`, validated against the assignments. Otherwise it is the full + * relation output. + */ + private def resolveWriteAttrs( + relation: DataSourceV2Relation, + operation: RowLevelOperation, + assignments: Seq[Assignment]): Seq[AttributeReference] = operation match { + case op: SupportsColumnUpdates with SupportsDelta if op.representUpdateAsDeleteAndInsert => + val attrs = resolveDeclaredAttrs(relation, op, assignments) + val rowIdAttrs = resolveRowIdAttrs(relation, op) + validateNoRowIdReassignment(op, relation, attrs, assignments, rowIdAttrs) + validateRowIdDeclared(op, attrs, rowIdAttrs, resolveRequiredMetadataAttrs(relation, op)) + attrs + case op: SupportsColumnUpdates => + resolveDeclaredAttrs(relation, op, assignments) + case _ => + relation.output + } + + /** + * Resolves the connector's `requiredDataAttributes()` and checks that it covers every assigned + * column. + */ + private def resolveDeclaredAttrs( + relation: DataSourceV2Relation, + operation: SupportsColumnUpdates, + assignments: Seq[Assignment]): Seq[AttributeReference] = { + val attrs = resolveRequiredDataAttrs(relation, operation) + validateUpdatedColumnsSubset(operation, assignments, attrs) + attrs + } + + /** + * 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 { Review Comment: **16.** Minor: `apply` already computes `updatedAttrs = collectUpdatedAttrs(assignments)` (L47), the list reported as `updatedColumns()` and named by this error, but passes raw `assignments` to `resolveWriteAttrs` (L50). So this method and `validateNoRowIdReassignment` (L318-322) re-derive the non-identity predicate; aokolnychyi's earlier suggestion ([r3786035399](https://github.com/apache/spark/pull/55518#discussion_r3786035399)) already noted that `collectUpdatedAttrs` could be reused in a few other places. With `updatedAttrs`, these become `updatedAttrs.filterNot(...)` and `updatedAttrs.filter(rowIdAttrSet.contains)`, and both `.distinct` calls can go, since aligned assignments have one entry per column. Similarly, the split arm resolves the row-ID and metadata attributes (L256, L258) that `buildWriteDeltaPlan` resolves again (L171-172), the same caller-side recomputation as in round 1 ([r3192416031](https://github.com/apache/spark/pull/55518#discussion_r3192416031)). In `resol veRequiredDataAttrs`, the `getOrElse(a.exprId, a)` fallback (`RewriteRowLevelCommand.scala:164`) is unreachable now that metadata columns are rejected, the cast there is a no-op, and the duplicate list follows `groupBy`'s hash order instead of the declared order. Could you pass `updatedAttrs` and resolve those attributes once? ########## sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala: ########## @@ -0,0 +1,999 @@ +/* + * 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.SparkRuntimeException +import org.apache.spark.sql.{AnalysisException, Row} +import org.apache.spark.sql.catalyst.optimizer.ColumnPruning +import org.apache.spark.sql.catalyst.plans.logical.WriteDelta +import org.apache.spark.sql.connector.catalog.{Delete, InMemoryBaseTable, InMemoryRowLevelOperationTable, Reinsert} +import org.apache.spark.sql.internal.SQLConf +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, updated rows contain only the declared columns + * (LogicalWriteInfo.columnUpdateSchema()) rather than the full table row. + */ +class DeltaBasedColumnUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase { + + override protected lazy val extraTableProps: java.util.Map[String, String] = { + val props = new java.util.HashMap[String, String]() + props.put("column-update", "true") + props + } + + private val reqAttrsOperationClass = + classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateOperationWithReqAttrs].getName + + private val splitOperationClass = + classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperation].getName + + private val splitReqAttrsOperationClass = + classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithReqAttrs] + .getName + + private val splitWithRowIdOperationClass = + classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithRowId].getName + + test("column-update: column update schema contains a 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") + + checkLastWriteInfo( + expectedRowIdSchema = Some(StructType(Array(PK_FIELD))), + expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))), + expectedColumnUpdateSchema = Some(StructType(Seq( + PK_FIELD, + DEP_FIELD, + StructField("id", IntegerType, nullable = false) + )))) + } + + test("column-update: column update schema 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))), + expectedColumnUpdateSchema = Some(StructType(Seq( + PK_FIELD, + StructField("dep", StringType, nullable = false), + StructField("id", IntegerType, nullable = false) + )))) + } + + test("column-update: column update schema is pk and dep 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; pk and dep remain because the + // connector unconditionally declares them as base columns. + checkLastWriteInfo( + expectedRowIdSchema = Some(StructType(Array(PK_FIELD))), + expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))), + expectedColumnUpdateSchema = 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))), + expectedColumnUpdateSchema = 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))), + expectedColumnUpdateSchema = Some(StructType(Seq( + PK_FIELD, + DEP_FIELD, + StructField("salary", IntegerType, nullable = true) + )))) + } + + test("column-update: column update schema 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))), + expectedColumnUpdateSchema = Some(StructType(Seq( + PK_FIELD, + StructField("dep", StringType, 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") + + checkLastUpdatedColumns("s") + + // `dep` is present because the connector unconditionally declares it as a base column. + assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq == Seq("pk", "dep", "s")) + } + + 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") + + checkLastUpdatedColumns("s") + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(1, Row(1, 2), "hr") :: Nil) + } + + test("column-update: updatedColumns is empty for DELETE") { + 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'") + + checkLastUpdatedColumns() + } + + test("column-update: row-level DELETE falls back to full-width rows") { + createAndInitTable("pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |{ "pk": 2, "salary": 200, "dep": "hr" } + |""".stripMargin) + + sql(s"DELETE FROM $tableNameAsString WHERE pk < 2") + + assert(!table.lastWriteInfo.columnUpdateSchema().isPresent, + s"DELETE must not report a column update schema: " + + s"${table.lastWriteInfo.columnUpdateSchema()}") + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(2, 200, "hr") :: Nil) + } + + test("column-update: MERGE falls back to full-width rows") { + createAndInitTable("pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |{ "pk": 2, "salary": 200, "dep": "software" } + |""".stripMargin) + checkMergeFallsBackToFullWidthRows() + } + + test("column-update split: MERGE falls back to full-width rows") { + createAndInitTable("pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |{ "pk": 2, "salary": 200, "dep": "software" } + |""".stripMargin, + java.util.Map.of("column-update-split", "true")) + checkMergeFallsBackToFullWidthRows() + } + + private def checkMergeFallsBackToFullWidthRows(): Unit = { + withTempView("source") { + import testImplicits._ + Seq((1, 150, "hr"), (3, 300, "software")).toDF("pk", "salary", "dep") + .createOrReplaceTempView("source") + + sql( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.pk = s.pk + |WHEN MATCHED THEN UPDATE SET salary = s.salary + |WHEN NOT MATCHED THEN INSERT * + |""".stripMargin) + + assert(!table.lastWriteInfo.columnUpdateSchema().isPresent, + s"MERGE must not report a column update schema: " + + s"${table.lastWriteInfo.columnUpdateSchema()}") + checkLastUpdatedColumns() + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"), + Row(1, 150, "hr") :: Row(2, 200, "software") :: Row(3, 300, "software") :: Nil) + } + } + + 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) + checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0) + } + + 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) + } + + 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) + } + + private def createAndInitTableWithReqAttrs( + reqAttrs: String, + schemaString: String, + jsonData: String): Unit = { + createAndInitTable(schemaString, jsonData, + java.util.Map.of("column-update-req-attrs", reqAttrs)) + } + + 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: undeclared distribution column is rejected") { + createAndInitTable("pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |""".stripMargin, + java.util.Map.of( + "column-update-req-attrs", "pk,salary", + "column-update-cluster-by", "dep")) + + val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName + Seq(Nil, Seq(excludeColumnPruning)).foreach { confs => + withSQLConf(confs: _*) { + checkError( + exception = intercept[AnalysisException] { + sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1") + }, + condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS", + parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[dep]")) + } + } + } + + test("column-update: distribution by an undeclared row ID column is allowed") { + createAndInitTable("pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |{ "pk": 2, "salary": 200, "dep": "hr" } + |""".stripMargin, + java.util.Map.of( + "column-update-req-attrs", "salary,dep", + "column-update-cluster-by", "pk")) + + sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1") + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"), + Row(1, -1, "hr") :: Row(2, 200, "hr") :: Nil) + } + + test("column-update: partition source column not in requiredDataAttributes") { + createAndInitTableWithReqAttrs("pk,id", "pk INT NOT NULL, id INT, dep STRING", + """{ "pk": 1, "id": 1, "dep": "hr" } + |{ "pk": 2, "id": 2, "dep": "software" } + |{ "pk": 3, "id": 3, "dep": "hr" } + |""".stripMargin) + + sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1") + + checkLastScanExcludes("dep") + 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") { + 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") + + assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq == + Seq("pk", "salary", "dep")) + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"), + Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil) + } + + + test("column-update: undeclared, differently-cased columns in the RHS and WHERE clause " + + "are read with the table's own spelling") { + createAndInitTableWithReqAttrs("pk,dep,salary", + "pk INT NOT NULL, salary INT, bonus INT, extra INT, dep STRING", + """{ "pk": 1, "salary": 100, "bonus": 10, "extra": 5, "dep": "hr" } + |""".stripMargin) + + sql(s"UPDATE $tableNameAsString SET salary = salary + BONUS WHERE pk = 1 AND EXTRA > 3") + + checkLastScanSchema("pk INT, salary INT, bonus INT, extra INT, dep STRING, _partition STRING") + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(1, 110, 10, 5, "hr") :: Nil) + } + + test("column-update: scan that returns unrequested columns still writes narrow rows") { + createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING", + """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" } + |""".stripMargin, + java.util.Map.of( + "column-update", "true", + InMemoryBaseTable.SIMULATE_PARTIAL_COLUMN_PRUNING, "true")) + + sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1") + + checkLastScanIncludes("bonus") + val columnUpdateSchema = table.lastWriteInfo.columnUpdateSchema().get() + assert(!columnUpdateSchema.fieldNames.contains("bonus"), + s"bonus must not be in column update schema: $columnUpdateSchema") + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(1, -1, 10, "hr") :: Nil) + } + + test("column-update: DELETE on a scan that returns unrequested columns is unaffected") { + createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING", + """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" } + |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" } + |""".stripMargin, + java.util.Map.of( + "column-update", "true", + InMemoryBaseTable.SIMULATE_PARTIAL_COLUMN_PRUNING, "true")) + + sql(s"DELETE FROM $tableNameAsString WHERE pk < 2") + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString"), + Row(2, 200, 20, "software") :: Nil) + } + + test("column-update: requiredDataAttributes resolves case-insensitively on an assigned " + + "column without adopting the declared spelling") { + createAndInitTableWithReqAttrs("pk,SALARY,dep", "pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |{ "pk": 2, "salary": 200, "dep": "software" } + |""".stripMargin) + + sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1") + + assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq == + Seq("pk", "salary", "dep")) + + checkAnswer( + sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"), + Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil) + } + + test("column-update: case-sensitive analysis accepts columns that differ only in case") { + withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") { + createAndInitTableWithReqAttrs("pk,salary,SALARY,dep", + "pk INT NOT NULL, salary INT, SALARY INT, dep STRING", + """{ "pk": 1, "salary": 100, "SALARY": 200, "dep": "hr" } + |""".stripMargin) + + sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1") + + assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq == + Seq("pk", "salary", "SALARY", "dep")) + checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, -1, 200, "hr") :: Nil) + } + } + + test("column-update: empty requiredDataAttributes throws AnalysisException") { + createAndInitTable("pk INT NOT NULL, id INT, dep STRING", + """{ "pk": 1, "id": 1, "dep": "hr" }""".stripMargin, + java.util.Map.of("column-update-empty-req-attrs", "true")) + + checkError( + exception = intercept[AnalysisException] { + sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1") + }, + condition = "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES", + parameters = Map("connector" -> reqAttrsOperationClass)) + } + + test("column-update: nested requiredDataAttributes throws AnalysisException") { + createAndInitTableWithReqAttrs("pk,s.c1,dep", + "pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING", + """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" } + |""".stripMargin) + + checkError( + exception = intercept[AnalysisException] { + sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1") + }, + condition = "COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE", + parameters = Map("connector" -> reqAttrsOperationClass, "nestedAttributes" -> "[s.c1]")) + } + + test("column-update: duplicate requiredDataAttributes throws AnalysisException") { + createAndInitTableWithReqAttrs("pk,pk,salary", "pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |""".stripMargin) + + checkError( + exception = intercept[AnalysisException] { + sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1") + }, + condition = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE", + parameters = Map("connector" -> reqAttrsOperationClass, "duplicateAttributes" -> "[pk]")) + } + + test("column-update: case-only duplicate requiredDataAttributes throws AnalysisException") { + createAndInitTableWithReqAttrs("PK,pk,salary", "pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |""".stripMargin) + + checkError( + exception = intercept[AnalysisException] { + sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1") + }, + condition = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE", + parameters = Map("connector" -> reqAttrsOperationClass, "duplicateAttributes" -> "[PK]")) + } + + test("column-update: metadata column in requiredDataAttributes throws AnalysisException") { + createAndInitTableWithReqAttrs("pk,_partition,dep", "pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |""".stripMargin) + + checkError( + exception = intercept[AnalysisException] { + sql(s"UPDATE $tableNameAsString SET dep = 'x' WHERE pk = 1") + }, + condition = "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE", + parameters = Map( + "connector" -> reqAttrsOperationClass, + "metadataAttributes" -> "[_partition]")) + } + + test("column-update split: metadata column in requiredDataAttributes throws " + + "AnalysisException") { + createAndInitTable("pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |""".stripMargin, + java.util.Map.of("column-update-split-req-attrs", "pk,_partition,dep")) + + checkError( + exception = intercept[AnalysisException] { + sql(s"UPDATE $tableNameAsString SET dep = 'x' WHERE pk = 1") + }, + condition = "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE", + parameters = Map( + "connector" -> splitReqAttrsOperationClass, + "metadataAttributes" -> "[_partition]")) + } + + test("column-update: requiredDataAttributes rejects a column that does not exist") { + createAndInitTableWithReqAttrs("pk,nonexistent_col,id", "pk INT NOT NULL, id INT, dep STRING", + """{ "pk": 1, "id": 1, "dep": "hr" } + |""".stripMargin) + + checkError( + exception = intercept[AnalysisException] { + sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1") + }, + condition = "COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE", + parameters = Map( + "connector" -> reqAttrsOperationClass, + "unknownAttributes" -> "[nonexistent_col]")) + } + + test("column-update: column update schema excludes undeclared columns") { + createAndInitTable("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") + + assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq == + Seq("pk", "dep", "salary")) + } + + test("column-update: data correctness") { + createAndInitTable("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 = { + createAndInitTable(schemaString, jsonData, + java.util.Map.of("column-update-split", "true")) + } + + test("column-update split: column update schema has the declared columns") { + 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") + + checkLastWriteInfo( + expectedRowIdSchema = Some(StructType(Array(PK_FIELD))), + expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))), + expectedColumnUpdateSchema = Some(StructType(Seq( + PK_FIELD, + DEP_FIELD, + StructField("id", IntegerType, nullable = false) + )))) + } + + test("column-update 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) + + checkError( + exception = intercept[AnalysisException] { + sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'") + }, + condition = "COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT", + parameters = Map("connector" -> splitOperationClass, "rowIds" -> "[pk]")) + } + + test("column-update split: row-ID reassignment is allowed when requiredDataAttributes " + + "covers every table column") { + 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 = { + createAndInitTable(schemaString, jsonData, + java.util.Map.of("column-update-split-req-attrs", reqAttrs)) + } + + test("column-update split: undeclared row-ID column on narrow write is rejected") { + createAndInitTable("pk INT NOT NULL, salary INT, dep STRING", + """{ "pk": 1, "salary": 100, "dep": "hr" } + |{ "pk": 2, "salary": 200, "dep": "software" } + |""".stripMargin, + java.util.Map.of("column-update-split-row-id", "pk")) + + checkError( + exception = intercept[AnalysisException] { + sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'") + }, + condition = "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED", + parameters = Map("connector" -> splitWithRowIdOperationClass, "rowIds" -> "[pk]")) + } + + // analyzes without executing, as this fixture can't commit a write that omits `pk` + private def analyzeWriteDelta(sqlText: String): WriteDelta = { + val parsed = spark.sessionState.sqlParser.parsePlan(sqlText) + spark.sessionState.executePlan(parsed).analyzed.collectFirst { + case w: WriteDelta => w + }.getOrElse(fail("couldn't find WriteDelta in analyzed plan")) + } + + test("column-update split: metadata row ID preserved on reinsert need not be declared") { Review Comment: **14. The split configuration that `validateRowIdDeclared` explicitly accepts is never executed.** This test (an undeclared metadata row ID preserved on reinsert) only analyzes the plan (L697), because the fixture cannot commit it: the split commit needs `pk` in the narrow row (`InMemoryRowLevelOperationTable.scala:507`), and `DeltaBufferWriter.delete` calls `id.getInt(0)` on the STRING `_partition` (`InMemoryRowLevelOperationTable.scala:870`). More broadly, no narrow delta test checks the row-ID or metadata values in the write log. So a regression where the reinserted metadata is shifted by `orderReadColumnsFirst` or pruning, or arrives as null instead of the preserved value, would go unnoticed; L716 only catches `_partition` disappearing from the schema. Could you assert the write log of a default split UPDATE with `checkLastColumnUpdateWriteLog(...)`, including the reinsert metadata, and make the fixture able to execute this test? ########## sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala: ########## @@ -258,6 +354,407 @@ class InMemoryRowLevelOperationTable private ( } } + // A delta-based operation that supports column-level updates: Spark sends only the declared + // columns in the row projection instead of the full row schema. It declares `pk` (the + // row-lookup key) and `dep` (an unconditionally declared base column) plus whatever columns + // Spark reports as being assigned via `RowLevelOperationInfo#updatedColumns()`. + class DeltaBasedColumnUpdateOperation( Review Comment: **19.** Nit: this block (L357-756) sits between `DeltaBasedOperation` (L303-355) and `TestDeltaBatchWrite` (L758-771), which only `DeltaBasedOperation` uses, while the class otherwise keeps each operation next to its batch write (`PartitionBasedOperation` at L236, `PartitionBasedReplaceData` at L288). AGENTS.md asks to "put the new member with the code it belongs with". A few more nits: - `MimaExcludes.scala:50` now reads "Write schema narrowing ...": moving the filter to `v44excludes` dropped the "scan and write" wording that matches the PR title. - L117 refers to `InMemoryRowLevelOperationTable.lastScanSchema`, but the companion member is `lastScanSchemaRef` (L911), and L910 refers to `RowLevelOperationSuiteBase.beforeEach`, but the reset is in `before { ... }`. - Redundant double blank lines at L549-550, `RowLevelOperationSuiteBase.scala:357-358`, `DeltaBasedColumnUpdateTableSuite.scala:161-162` and `:393-394`, and `GroupBasedColumnUpdateTableSuite.scala:560-561` and `:658-659`. Could you move the block below `TestDeltaBatchWrite` and fix these? ########## sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala: ########## @@ -4638,6 +4638,92 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat messageParameters = Map("nullableRowIdAttrs" -> nullableRowIdAttrs.mkString(", "))) } + def emptyRequiredDataAttributesError(connectorClass: String): Throwable = { + new AnalysisException( + errorClass = "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES", + messageParameters = Map("connector" -> connectorClass)) + } + + def duplicateRequiredDataAttributeError( + connectorClass: String, + duplicateAttributes: Seq[String]): Throwable = { + new AnalysisException( + errorClass = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE", + messageParameters = Map( + "connector" -> connectorClass, + "duplicateAttributes" -> duplicateAttributes.mkString("[", ", ", "]"))) Review Comment: **17.** Minor: the new `COLUMN_UPDATE_*` helpers format identifier lists by hand with `mkString("[", ", ", "]")` instead of `toSQLId`, which `QueryErrorsBase` asks for ("Elements of identifiers shall be wrapped by backticks by using `toSQLId()`"), and their inputs mix `ref.describe()` (`RewriteRowLevelCommand.scala:128`, `:138`, `:146`, `:154` and `V2Writes.scala:161`) with raw names (`RewriteUpdateTable.scala:291`, `:321`, `:347`). For a top-level column named `a.b`, the UNKNOWN error shows "[`a.b`]" while MISSING_UPDATED_COLUMNS shows "[a.b]", which reads like the nested field in the NESTED error's "[s.c1]". In the same suite, `dep` appears as "`dep`" (`GroupBasedColumnUpdateTableSuite.scala:375`) and as "[dep]" (`:409`). These parameters become part of the `checkError` contract once released. Could you use `toSQLId` for all of them? ########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java: ########## @@ -82,6 +84,60 @@ default void write(T metadata, T record) throws IOException { write(record); } + /** + * Writes one updated or copied record with metadata in the column update layout. + * <p> + * When {@link LogicalWriteInfo#columnUpdateSchema()} is present for a row-level operation that + * does not mix in {@link SupportsDelta}, Spark passes updated and copied records to this method + * instead of {@link #write(Object, Object)} if the operation returns a non-empty Review Comment: **13.** Minor: the Javadoc of the existing `write(T metadata, T record)` (L72-74) still says "This method is used by group-based row-level operations to pass back metadata for records that are updated or copied", but for a narrow UPDATE those records come here instead (`WriteToDataSourceV2Exec.scala:786-787`). A connector author who reads only `write(T, T)` gets `DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED`. This PR adds such a pointer to `DeltaWriter.update`, `DeltaWriter.reinsert` and `LogicalWriteInfo.schema()`, but not there. Could you add a sentence to L72-74, e.g. "When {@link LogicalWriteInfo#columnUpdateSchema()} is present for an operation that does not mix in {@link SupportsDelta}, Spark passes these records to {@link #writeColumnUpdate(Object, Object)} instead."? -- 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]
