anuragmantri commented on code in PR #55518: URL: https://github.com/apache/spark/pull/55518#discussion_r4240160391
########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java: ########## @@ -0,0 +1,47 @@ +/* + * 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. + * + * @since 4.3.0 Review Comment: 4.3.0 has been released, so the APIs are now `@since 4.4.0` and the MiMa exclude is in the 4.4 section. ########## 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) + } + V2ExpressionUtils.resolveRefs[AttributeReference]( Review Comment: Fixed: each resolved attribute is mapped back to the relation's own attribute by expr ID, and the narrow read relation this fed into no longer exists. Covered by the case-insensitive declaration tests (`requiredDataAttributes resolves case-insensitively ...`) and the differently-cased tests, which check the table's spelling in the scan. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala: ########## @@ -146,6 +203,48 @@ object RewriteUpdateTable extends RewriteRowLevelCommand { Project(operationCol +: updatedValues, plan) } + /** + * Variant of `buildReplaceDataUpdateProjection` for the `SupportsColumnUpdates` narrow-scan + * path. + * For narrow attributes, looks up assignments by `ExprId` via `AttributeMap`and passes through Review Comment: This code was removed in the redesign. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala: ########## @@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand { val operation = operationTable.operation.asInstanceOf[SupportsDelta] - // resolve all needed attrs (e.g. row ID and any required metadata attrs) + // resolve all needed attrs (e.g. row ID, any required metadata attrs and optionally connector + // declared attrs) val rowAttrs = relation.output + val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates] + val connectorDataAttrs = if (supportsColumnUpdate) { + resolveConnectorDataAttrs(relation, operation) + } else Nil + val scanOnlyDataAttrs = if (supportsColumnUpdate) { + resolveScanOnlyDataAttrs(relation, operation) + } else Nil + + if (supportsColumnUpdate) { + validateUpdatedColumnsSubset(operation, assignments, connectorDataAttrs) + validateNoOverlap(operation, connectorDataAttrs, scanOnlyDataAttrs) + validatePartitionAttrsDeclared(operation, relation, connectorDataAttrs, scanOnlyDataAttrs) + } + + val rowIdAttrs = resolveRowIdAttrs(relation, operation) val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation) - // construct a read relation and include all required metadata columns - val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs) + if (supportsColumnUpdate && operation.representUpdateAsDeleteAndInsert) { + validateNoRowIdReassignment(operation, assignments, rowIdAttrs) + validateRowIdDeclared(operation, connectorDataAttrs, rowIdAttrs) + } + + val narrowDataAttrs = if (supportsColumnUpdate) { + computeNarrowReadAttrs(relation, connectorDataAttrs, scanOnlyDataAttrs, assignments, cond) + } else { + relation.output + } + + val readRelation = if (supportsColumnUpdate) { + buildNarrowRelationWithAttrs(relation, operationTable, narrowDataAttrs, metadataAttrs, + rowIdAttrs) + } else { + buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs) + } // build a plan for updated records that match the condition val matchedRowsPlan = Filter(cond, readRelation) - val rowDeltaPlan = if (operation.representUpdateAsDeleteAndInsert) { - buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs) + val rowDeltaPlan = if (supportsColumnUpdate) { + if (operation.representUpdateAsDeleteAndInsert) { + buildNarrowDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs) + } else { + buildColumnUpdateProjection( + matchedRowsPlan, assignments, rowIdAttrs, metadataAttrs, connectorDataAttrs, + scanOnlyDataAttrs) + } } else { - buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs) + if (operation.representUpdateAsDeleteAndInsert) { + buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs) + } else { + buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs) + } } // build a plan to write the row delta to the table val writeRelation = relation.copy(table = operationTable) - val projections = buildWriteDeltaProjections(rowDeltaPlan, rowAttrs, rowIdAttrs, metadataAttrs) + val projections = buildWriteDeltaProjections( + rowDeltaPlan, rowAttrs, rowIdAttrs, metadataAttrs, connectorDataAttrs) val groupFilterCond = if (groupFilterEnabled) Some(cond) else None WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond) } + /** + * Builds the WriteDelta projection for the column-update path. + */ + private def buildColumnUpdateProjection( + plan: LogicalPlan, + assignments: Seq[Assignment], + rowIdAttrs: Seq[Attribute], + metadataAttrs: Seq[Attribute], + connectorDataAttrs: Seq[AttributeReference], + scanOnlyDataAttrs: Seq[AttributeReference]): LogicalPlan = { + + val assignedValues = assignments.collect { + case Assignment(key: Attribute, value) if !isIdentityAssignment(key, value) => + Alias(value, key.name)() + } + + // Connector-required columns whose value isn't being changed by the UPDATE: pass through the + // current value so the connector receives a complete write row. Row-ID columns are excluded + // here even when also connector-required (e.g. a primary key used for both row lookup and + // write payload) they are emitted once, below, via rowIdValues. + val assignedKeyIds = assignments.collect { + case Assignment(key: AttributeReference, value) if !isIdentityAssignment(key, value) => + key.exprId + }.toSet + val rowIdAttrSet = AttributeSet(rowIdAttrs) + val connectorPassThroughValues = connectorDataAttrs.filterNot(a => + assignedKeyIds.contains(a.exprId) || rowIdAttrSet.contains(a)) + + // scanOnlyDataAttrs are never assigned carry them through the write query so Review Comment: This code was removed in the redesign. ########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/LogicalWriteInfo.java: ########## @@ -65,4 +65,14 @@ default Optional<StructType> metadataSchema() { throw new SparkUnsupportedOperationException( "DATA_SOURCE_METADATA_SCHEMA_NOT_IMPLEMENTED", Map.of("class", getClass().getName())); } + + /** + * the narrow schema for updates. Present only when the connector mixes in + * {@link SupportsColumnUpdates}. + * + * @since 4.3.0 + */ + default Optional<StructType> updateSchema() { Review Comment: Renamed to `columnUpdateSchema()`, together with `writeColumnUpdate`. ########## common/utils/src/main/resources/error/error-conditions.json: ########## @@ -2385,6 +2421,12 @@ ], "sqlState" : "42K03" }, + "DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED" : { + "message" : [ + "<class> mixes in `SupportsColumnUpdates` but does not implement `writeUpdate`. Connectors that opt into narrow column-level updates must override `writeUpdate` to handle rows in the `LogicalWriteInfo.updateSchema()` layout; the default implementation would forward narrow rows to `write`, which expects the full `LogicalWriteInfo.schema()` layout." Review Comment: Fixed. The message no longer describes a default that forwards to `write`, and `LogicalWriteInfo.schema()` now says that when `columnUpdateSchema()` is present it covers only newly inserted rows, so it is empty for UPDATE. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedColumnUpdateTableSuite.scala: ########## @@ -0,0 +1,320 @@ +/* + * 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, SparkUnsupportedOperationException} +import org.apache.spark.sql.Row +import org.apache.spark.sql.connector.catalog.{CatalogV2Util, TableInfo, WriteUpdate} +import org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity, reference} +import org.apache.spark.sql.connector.expressions.Transform +import org.apache.spark.sql.types.StructType + +class GroupBasedColumnUpdateTableSuite extends UpdateTableSuiteBase { + + override protected lazy val extraTableProps: java.util.Map[String, String] = { + val props = new java.util.HashMap[String, String]() + props.put("column-update-cow", "true") + props + } + + private def createAndInitTableReplaceData(schemaString: String, jsonData: String): Unit = { Review Comment: Removed the duplicate helpers and the `FromInfo` fixture. The suites use one `createAndInitTable(schema, data, tableProps)` overload. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala: ########## @@ -306,13 +308,80 @@ abstract class RowLevelOperationSuiteBase protected def checkLastWriteInfo( expectedRowSchema: StructType = new StructType(), expectedRowIdSchema: Option[StructType] = None, - expectedMetadataSchema: Option[StructType] = None): Unit = { + expectedMetadataSchema: Option[StructType] = None, + expectedUpdateSchema: Option[StructType] = None): Unit = { val info = table.lastWriteInfo assert(info.schema == expectedRowSchema, "row schema must match") val actualRowIdSchema = Option(info.rowIdSchema.orElse(null)) assert(actualRowIdSchema == expectedRowIdSchema, "row ID schema must match") val actualMetadataSchema = Option(info.metadataSchema.orElse(null)) assert(actualMetadataSchema == expectedMetadataSchema, "metadata schema must match") + val actualUpdateSchema = Option(info.updateSchema.orElse(null)) + assert(actualUpdateSchema == expectedUpdateSchema, "update schema must match") + } + + protected def getUpdateSummary(): UpdateSummary = { + catalog.loadTable(ident).asInstanceOf[InMemoryTable] + .commits.last.writeSummary.get + .asInstanceOf[UpdateSummary] + } + + /** + * Asserts the last UPDATE's write summary metrics. `deltaUpdate` controls the expected + * COPY count: MoR connectors emit only deltas (no COPY rows) so `numCopiedRows` is forced + * to 0; CoW connectors emit COPY rows for unchanged rows in matched groups. + */ + protected def checkUpdateMetrics( + numUpdatedRows: Long, + numCopiedRows: Long, + deltaUpdate: Boolean = false): Unit = { + val summary = getUpdateSummary() + assert(summary.numUpdatedRows() === numUpdatedRows, + s"Expected numUpdatedRows=$numUpdatedRows, got ${summary.numUpdatedRows()}") + val expectedCopied = if (deltaUpdate) 0L else numCopiedRows + assert(summary.numCopiedRows() === expectedCopied, + s"Expected numCopiedRows=$expectedCopied, got ${summary.numCopiedRows()}") + } + + /** + * Asserts that the column names in RowLevelOperationInfo.updatedColumns() received by the + * last operation match exactly the expected set. Order is ignored. + */ + protected def checkLastUpdatedColumns(expectedNames: String*): Unit = { + val actual = table.lastUpdatedColumns.map(_.describe()).toSet + val expected = expectedNames.toSet + assert(actual == expected, + s"updatedColumns mismatch: expected ${expected.mkString("[", ", ", "]")} " + + s"but got ${actual.mkString("[", ", ", "]")}") + } + + /** + * Asserts that the last connector scan schema does NOT contain any of the given columns. + * Useful for negative narrowing assertions -- proves that a column outside the analysis-time + * narrow set was pruned by the scan, without over-specifying what else is present (which + * ColumnPruning may further tighten in ways unrelated to the narrowing contract). + */ + protected def checkLastScanExcludes(excludedNames: String*): Unit = { + val schema = Option(table.lastScanSchema).getOrElse(StructType(Nil)) Review Comment: `checkLastScanExcludes` and `checkLastScanIncludes` now fail when no scan schema was recorded. The group-filtering test checks both scan schemas and the replaced partitions. ########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java: ########## @@ -105,4 +105,45 @@ default String description() { default NamedReference[] requiredMetadataAttributes() { return new NamedReference[0]; } + + + /** + * Controls whether to send only the required data columns to the connector rather than the + * full row. + * <p> + * When true, Spark narrows the data column schema ({@link LogicalWriteInfo#schema()}) to only + * the columns declared via {@link #requiredDataAttributes()}. Metadata columns (from + * {@link #requiredMetadataAttributes()}) and row ID columns (from + * {@link SupportsDelta#rowId()}) are unaffected and always projected separately. + * <p> + * If {@link #requiredDataAttributes()} returns a non-empty array, the write schema is exactly + * those columns in declared order. The connector must include all columns it wants to receive, + * including the columns being updated. If {@link #requiredDataAttributes()} returns an empty + * array, Spark sends only the non-identity assigned columns (heuristic path). + * + * @since 4.2.0 + */ + default boolean supportsColumnUpdates() { + return false; + } + + /** + * Returns data column references required to perform this row-level operation. + * <p> + * This method is only consulted by Spark when {@link #supportsColumnUpdates()} returns + * {@code true}. If {@code supportsColumnUpdates()} returns {@code false}, the returned array + * is ignored and the full table row is sent (the default behavior). + * <p> + * When non-empty, the returned columns become the write schema in declared order. + * The connector must declare all columns it wants to receive, including the columns being + * updated. Use {@link RowLevelOperationInfo#updatedColumns()} to learn which columns are being + * assigned, then add any extra columns needed for row lookup or routing (e.g., primary key). + * <p> + * When empty (the default), Spark falls back to sending only the non-identity assigned columns. + * + * @since 4.2.0 Review Comment: The APIs are now `@since 4.4.0`, since 4.3.0 has been released. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala: ########## @@ -106,38 +105,78 @@ object RewriteUpdateTable extends RewriteRowLevelCommand { val remainingRowsPlan = addOperationColumn(COPY_OPERATION, Filter(remainingRowFilter, readRelation)) - // the new state is a union of updated and copied over records - val query = Union(updatedRowsPlan, remainingRowsPlan) + val updatedAndRemainingRowsPlan = Union(updatedRowsPlan, remainingRowsPlan) - // build a plan to replace read groups in the table val writeRelation = relation.copy(table = operationTable) Review Comment: You were right. The PR now narrows only the write relation here, with the existing read relation and plan builders, and validates against `relation.output` as before. To keep the API MERGE-proof, the narrow layout is reported as `columnUpdateSchema()` and `schema()` covers only newly inserted rows, so MERGE can narrow its write relation later without API changes. ########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java: ########## @@ -105,4 +105,45 @@ default String description() { default NamedReference[] requiredMetadataAttributes() { return new NamedReference[0]; } + + + /** + * Controls whether to send only the required data columns to the connector rather than the + * full row. + * <p> + * When true, Spark narrows the data column schema ({@link LogicalWriteInfo#schema()}) to only + * the columns declared via {@link #requiredDataAttributes()}. Metadata columns (from + * {@link #requiredMetadataAttributes()}) and row ID columns (from + * {@link SupportsDelta#rowId()}) are unaffected and always projected separately. + * <p> + * If {@link #requiredDataAttributes()} returns a non-empty array, the write schema is exactly + * those columns in declared order. The connector must include all columns it wants to receive, + * including the columns being updated. If {@link #requiredDataAttributes()} returns an empty + * array, Spark sends only the non-identity assigned columns (heuristic path). + * + * @since 4.2.0 + */ + default boolean supportsColumnUpdates() { Review Comment: Done. The `SupportsColumnUpdates` Javadoc says Spark currently narrows only UPDATE, and `RowLevelOperationInfo.updatedColumns()` says other commands report an empty array. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala: ########## @@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand { val operation = operationTable.operation.asInstanceOf[SupportsDelta] - // resolve all needed attrs (e.g. row ID and any required metadata attrs) + // resolve all needed attrs (e.g. row ID, any required metadata attrs and optionally connector + // declared attrs) val rowAttrs = relation.output + val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates] + val connectorDataAttrs = if (supportsColumnUpdate) { + resolveConnectorDataAttrs(relation, operation) + } else Nil + val scanOnlyDataAttrs = if (supportsColumnUpdate) { + resolveScanOnlyDataAttrs(relation, operation) + } else Nil + + if (supportsColumnUpdate) { + validateUpdatedColumnsSubset(operation, assignments, connectorDataAttrs) + validateNoOverlap(operation, connectorDataAttrs, scanOnlyDataAttrs) + validatePartitionAttrsDeclared(operation, relation, connectorDataAttrs, scanOnlyDataAttrs) + } + + val rowIdAttrs = resolveRowIdAttrs(relation, operation) val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation) - // construct a read relation and include all required metadata columns - val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs) + if (supportsColumnUpdate && operation.representUpdateAsDeleteAndInsert) { + validateNoRowIdReassignment(operation, assignments, rowIdAttrs) + validateRowIdDeclared(operation, connectorDataAttrs, rowIdAttrs) + } + + val narrowDataAttrs = if (supportsColumnUpdate) { + computeNarrowReadAttrs(relation, connectorDataAttrs, scanOnlyDataAttrs, assignments, cond) + } else { + relation.output + } + + val readRelation = if (supportsColumnUpdate) { + buildNarrowRelationWithAttrs(relation, operationTable, narrowDataAttrs, metadataAttrs, + rowIdAttrs) + } else { + buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs) + } // build a plan for updated records that match the condition val matchedRowsPlan = Filter(cond, readRelation) - val rowDeltaPlan = if (operation.representUpdateAsDeleteAndInsert) { - buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs) + val rowDeltaPlan = if (supportsColumnUpdate) { Review Comment: Reworked. The column-update decisions are now in two places: `resolveWriteAttrs` in `RewriteUpdateTable`, a single three-arm match that returns the write relation's columns, and `V2Writes.hasNarrowRows`, which the write builders, the distribution check, copy-on-write scan planning and the exec routing use. The plan builders are the existing ones and have no column-update branches. ########## 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 { Review Comment: Renamed to `writeColumnUpdate(record)` and `writeColumnUpdate(metadata, record)`. The Javadoc says it receives both updated and copied rows of a group-based operation. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala: ########## @@ -226,4 +389,98 @@ 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 + } + + /** + * Computes the narrow set of data columns that must be present in the scan for a column-update + * write: connector-declared attrs, unioned with any table columns referenced by non-identity + * assignment RHS expressions, the operation condition, and the table's partition expressions. + * Partition-column refs are always kept so downstream rules (V2ScanPartitioningAndOrdering, + * GroupBasedRowLevelOperationScanPlanning) can resolve the table's partitioning expressions + * against the scan output. + */ + 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) + .collect { case a: AttributeReference => a } + .filter(relationSet.contains) + val partitionRefNames = relation.table.partitioning().toImmutableArraySeq Review Comment: Following Anton's later comment, `scanOnlyDataAttributes()` was removed. Columns the write's distribution or ordering needs are declared in `requiredDataAttributes()`, and an undeclared one is now rejected with `COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS`. Before, the same write failed with `_LEGACY_ERROR_TEMP_1137` only when column pruning ran, and succeeded with it excluded. Your delta repro is covered by `column-update: undeclared distribution column is rejected`, with and without `ColumnPruning` excluded. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala: ########## @@ -267,13 +268,80 @@ abstract class RowLevelOperationSuiteBase protected def checkLastWriteInfo( expectedRowSchema: StructType = new StructType(), expectedRowIdSchema: Option[StructType] = None, - expectedMetadataSchema: Option[StructType] = None): Unit = { + expectedMetadataSchema: Option[StructType] = None, + expectedUpdateSchema: Option[StructType] = None): Unit = { val info = table.lastWriteInfo assert(info.schema == expectedRowSchema, "row schema must match") val actualRowIdSchema = Option(info.rowIdSchema.orElse(null)) assert(actualRowIdSchema == expectedRowIdSchema, "row ID schema must match") val actualMetadataSchema = Option(info.metadataSchema.orElse(null)) assert(actualMetadataSchema == expectedMetadataSchema, "metadata schema must match") + val actualUpdateSchema = Option(info.updateSchema.orElse(null)) + assert(actualUpdateSchema == expectedUpdateSchema, "update schema must match") + } + + protected def getUpdateSummary(): org.apache.spark.sql.connector.write.UpdateSummary = { Review Comment: The update summary helpers are back in `UpdateTableSuiteBase`, which imports `UpdateSummary`. ########## 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: SPARK-59721 (#58979) is merged, so I added the deferred test: `column-update: partition source column not in requiredDataAttributes` declares `[pk, id]` on a table partitioned by `dep`. It passes, and fails with `Unable to resolve dep given [pk,id,_partition,index]` if the #58979 change to `V2ScanPartitioningAndOrdering` is reverted. The description links the JIRA. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala: ########## @@ -771,6 +777,7 @@ trait WritingSparkTask[W <: DataWriter[InternalRow]] extends Logging with Serial case class DataAndMetadataWritingSparkTask( dataProj: ProjectingInternalRow, + updateDataProj: ProjectingInternalRow, Review Comment: The tasks no longer take a nullable projection; they take a `useWriteColumnUpdate: Boolean` and reuse the row projection. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala: ########## @@ -109,6 +123,43 @@ 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) + } + val nested = refs.filter(_.fieldNames.length != 1).map(_.describe()).toImmutableArraySeq + if (nested.nonEmpty) { + throw QueryCompilationErrors.nestedRequiredDataAttributeError( + operation.getClass.getName, nested) + } + val normalizedNames = refs.map { ref => + val name = ref.fieldNames.head + if (conf.caseSensitiveAnalysis) name else name.toLowerCase(Locale.ROOT) + } + val duplicates = normalizedNames.groupBy(identity).collect { + case (_, occurrences) if occurrences.length > 1 => occurrences.head + }.toSeq Review Comment: The duplicate error reports the spelling the connector declared, as the nested check does. `["PK", "Pk"]` reports `[PK]` (`column-update: case-only duplicate requiredDataAttributes throws AnalysisException`). ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala: ########## @@ -157,29 +228,112 @@ object RewriteUpdateTable extends RewriteRowLevelCommand { val operation = operationTable.operation.asInstanceOf[SupportsDelta] - // resolve all needed attrs (e.g. row ID and any required metadata attrs) + // resolve all needed attrs (e.g. row ID, required metadata attrs, and any connector-declared + // data attrs for column-update writes) val rowAttrs = relation.output + val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates] + val connectorDataAttrs = resolveConnectorDataAttrs(relation, operation, assignments) + val rowIdAttrs = resolveRowIdAttrs(relation, operation) val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation) - // construct a read relation and include all required metadata columns - val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs) + if (supportsColumnUpdate && operation.representUpdateAsDeleteAndInsert) { + validateNoRowIdReassignment(operation, relation, connectorDataAttrs, assignments, rowIdAttrs) + validateRowIdDeclared(operation, connectorDataAttrs, rowIdAttrs) + } + + val narrowDataAttrs = if (supportsColumnUpdate) { + computeNarrowReadAttrs(relation, connectorDataAttrs, assignments, cond) + } else { + relation.output + } + + val readRelation = if (supportsColumnUpdate) { + buildNarrowRelationWithAttrs(relation, operationTable, narrowDataAttrs, metadataAttrs, + rowIdAttrs) + } else { + buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs) + } // build a plan for updated records that match the condition val matchedRowsPlan = Filter(cond, readRelation) - val rowDeltaPlan = if (operation.representUpdateAsDeleteAndInsert) { - buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs) + val rowDeltaPlan = if (supportsColumnUpdate) { + if (operation.representUpdateAsDeleteAndInsert) { + buildNarrowDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs) + } else { + buildColumnUpdateProjection( + matchedRowsPlan, assignments, rowIdAttrs, metadataAttrs, connectorDataAttrs) + } } else { - buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs) + if (operation.representUpdateAsDeleteAndInsert) { + buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs) + } else { + buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs) + } } // build a plan to write the row delta to the table val writeRelation = relation.copy(table = operationTable) - val projections = buildWriteDeltaProjections(rowDeltaPlan, rowAttrs, rowIdAttrs, metadataAttrs) + val projections = buildWriteDeltaProjections( + rowDeltaPlan, rowAttrs, rowIdAttrs, metadataAttrs, connectorDataAttrs) val groupFilterCond = if (groupFilterEnabled) Some(cond) else None WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond) } + /** + * Builds the WriteDelta projection for the column-update path. + */ + private def buildColumnUpdateProjection( + plan: LogicalPlan, + assignments: Seq[Assignment], + rowIdAttrs: Seq[Attribute], + metadataAttrs: Seq[Attribute], + connectorDataAttrs: Seq[AttributeReference]): LogicalPlan = { + + val assignedValues = assignments.collect { + case Assignment(key: Attribute, value) if !isIdentityAssignment(key, value) => + Alias(value, key.name)() + } + + // Connector-required columns whose value isn't being changed by the UPDATE: pass through the + // current value so the connector receives a complete write row. Row-ID columns are excluded + // here even when also connector-required (e.g. a primary key used for both row lookup and + // write payload); they are emitted once, below, via rowIdValues. + val assignedKeyIds = collectUpdatedAttrs(assignments).map(_.exprId).toSet + val rowIdAttrSet = AttributeSet(rowIdAttrs) + val connectorPassThroughValues = connectorDataAttrs.filterNot(a => + assignedKeyIds.contains(a.exprId) || rowIdAttrSet.contains(a)) + + val metadataAttrSet = AttributeSet(metadataAttrs) + val metadataValues = plan.output.filter(metadataAttrSet.contains).map { attr => + if (MetadataAttribute.isPreservedOnUpdate(attr)) { + attr + } else { + Alias(Literal(null, attr.dataType), attr.name)(explicitMetadata = Some(attr.metadata)) + } + } + + val rowIdValues = plan.output.filter(rowIdAttrSet.contains) Review Comment: The rewrite now uses the existing projection builders, which project the original row ID under a separate name, so `pk` appears once. The row-ID reassignment tests with a CHECK constraint on `pk` still pass, for example `column-update: CHECK constraint on a reassigned row-ID column is enforced`. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala: ########## @@ -228,4 +382,169 @@ 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, the operation condition, and the + * table's CHECK constraints. + */ + 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) + .collect { case a: AttributeReference => a } + .filter(relationSet.contains) + val checkConstraintAttrs = resolveCheckConstraintAttrs(relation) + dedupAttrs(connectorDataAttrs ++ extraRefs ++ checkConstraintAttrs) + } + + /** + * Returns the relation attributes referenced by the table's CHECK constraints, so that + * narrowing does not drop a column `ResolveTableConstraints` needs later to re-validate the + * rewritten write query. Falls back to the full relation output whenever a constraint's + * condition cannot be resolved to a `Predicate` (see `Check#predicate()`). + */ + private def resolveCheckConstraintAttrs( + relation: DataSourceV2Relation): Seq[AttributeReference] = { + val checks = Option(relation.table.constraints).getOrElse(Array.empty[Constraint]).collect { + case c: Check => c + } + if (checks.isEmpty) { + Nil + } else if (checks.exists(_.predicate() == null)) { + relation.output + } else { + val refs = checks.flatMap(_.predicate().references()).toSeq + V2ExpressionUtils.resolveRefs[AttributeReference](refs, relation) Review Comment: With the read relation no longer narrowed, CHECK constraints resolve against the full write query, so a nested-field CHECK such as `CHECK (s.c1 > 0)` no longer needs special handling. Covered by `CHECK constraint on a nested field narrows correctly` in both suites. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala: ########## @@ -228,4 +382,169 @@ 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, the operation condition, and the + * table's CHECK constraints. + */ + 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) + .collect { case a: AttributeReference => a } + .filter(relationSet.contains) + val checkConstraintAttrs = resolveCheckConstraintAttrs(relation) + dedupAttrs(connectorDataAttrs ++ extraRefs ++ checkConstraintAttrs) Review Comment: This no longer applies: `RewriteUpdateTable` now keeps the existing read relation and narrows only the write relation, so no attribute in the connector's or the user's spelling can reach a scan. The declared attributes are still mapped back to the relation's own attributes by expr ID. I kept the regression tests for your example on both paths, and they now also check that the scan uses the table's spelling (`undeclared, differently-cased columns in the RHS and WHERE clause are read with the table's own spelling`). ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala: ########## @@ -88,6 +92,16 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] { relation.copy(table = table, output = attrs) } + protected def buildNarrowRelationWithAttrs( + relation: DataSourceV2Relation, + table: RowLevelOperationTable, + dataAttrs: Seq[AttributeReference], + metadataAttrs: Seq[AttributeReference], + rowIdAttrs: Seq[AttributeReference] = Nil): DataSourceV2Relation = { + val attrs = dedupAttrs(dataAttrs ++ rowIdAttrs ++ metadataAttrs) Review Comment: This no longer applies: the read relation keeps every table column, so `PushDownUtils.pruneColumns` and `toOutputAttrs` see the same relation as for any other row-level operation. A scan builder without `SupportsPushDownRequiredColumns` reads every column and the write stays correct; the Javadoc now says the builder should implement it. The partial-pruning case is covered by `scan that returns unrequested columns still writes narrow rows` in both suites. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala: ########## @@ -109,6 +123,43 @@ 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) + } + val nested = refs.filter(_.fieldNames.length != 1).map(_.describe()).toImmutableArraySeq + if (nested.nonEmpty) { + throw QueryCompilationErrors.nestedRequiredDataAttributeError( + operation.getClass.getName, nested) + } + val normalizedNames = refs.map { ref => + val name = ref.fieldNames.head + if (conf.caseSensitiveAnalysis) name else name.toLowerCase(Locale.ROOT) + } + val duplicates = normalizedNames.groupBy(identity).collect { + case (_, occurrences) if occurrences.length > 1 => occurrences.head + }.toSeq + if (duplicates.nonEmpty) { + throw QueryCompilationErrors.duplicateRequiredDataAttributeError( + operation.getClass.getName, duplicates) + } + val resolved = V2ExpressionUtils.resolveRefs[AttributeReference]( Review Comment: A metadata column in `requiredDataAttributes()` is now rejected with `COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE`, and the Javadoc says to return metadata columns from `requiredMetadataAttributes()` instead. Covered on the in-place and split delta paths (`metadata column in requiredDataAttributes throws AnalysisException`). ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala: ########## @@ -513,7 +557,31 @@ case class WriteDelta( case Some(projection) => DataTypeUtils.toAttributes(projection.schema) case None => Nil } - table.skipSchemaResolution || areCompatible(inRowAttrs, outRowAttrs) + val inUpdateAttrs = projections.updateRowProjection match { + case Some(projection) => DataTypeUtils.toAttributes(projection.schema) + case None => Nil + } + // `rowProjection` (INSERT-tagged rows) validates against `outRowAttrs`. For column-update + // connectors, `updateRowProjection` is narrow and validates against `projectedDataAttrs` + // (the connector-declared narrow set) instead. When the connector does not mix in + // `SupportsColumnUpdates`, `updateRowProjection` is absent and `updateResolved` is trivially + // true. + val insertResolved = table.skipSchemaResolution || + (inUpdateAttrs.nonEmpty && inRowAttrs.isEmpty) || + areCompatible(inRowAttrs, outRowAttrs) + val updateResolved = table.skipSchemaResolution || + inUpdateAttrs.isEmpty || dataAttrsResolved(inUpdateAttrs) + insertResolved && updateResolved + } + + /** + * Validates the narrow-write-schema row projection output for a column-update connector. + * The write schema must exactly match the columns declared via + * `SupportsColumnUpdates.requiredDataAttributes()` (same columns, same order). + */ + private def dataAttrsResolved(inRowAttrs: Seq[Attribute]): Boolean = { Review Comment: This code is gone. `ReplaceData` and `WriteDelta` validate against the narrowed write relation with the existing `rowAttrsResolved`. The only addition is `RowLevelWrite.references`, defined once in the trait. ########## sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala: ########## @@ -252,6 +304,413 @@ 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) and `dep` (an + // unconditionally declared base column) plus whatever columns Spark reports as being + // assigned via `RowLevelOperationInfo#updatedColumns()`. + class DeltaBasedColumnUpdateOperation( + command: Command, + updatedCols: Seq[NamedReference] = Nil, + options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty()) + extends DeltaBasedOperation(command, options) + with SupportsColumnUpdates { + override def representUpdateAsDeleteAndInsert(): Boolean = false + override def requiredDataAttributes(): Array[NamedReference] = { + val base = Seq(FieldReference("pk"), FieldReference("dep")) + val baseNames = base.map(_.describe()).toSet + (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray + } + + protected def clusterColumnRef: NamedReference = PARTITION_COLUMN_REF + + override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = { + lastWriteInfo = info + // Capture info into a local val so nested writer/commit closures see a stable schema + // even if a subsequent newWriteBuilder call mutates lastWriteInfo. + val capturedInfo = info + val capturedWriteSchema = if (capturedInfo.updateSchema().isPresent) { + capturedInfo.updateSchema().get() + } else { + capturedInfo.schema() + } + new DeltaWriteBuilder { + override def build(): DeltaWrite = + new DeltaWrite with RequiresDistributionAndOrdering { + + override def requiredDistribution(): Distribution = { + Distributions.clustered(Array(clusterColumnRef)) + } + + override def requiredOrdering(): Array[SortOrder] = { + Array[SortOrder]( + LogicalExpressions.sort( + clusterColumnRef, + SortDirection.ASCENDING, + SortDirection.ASCENDING.defaultNullOrdering()) + ) + } + + override def toBatch: DeltaBatchWrite = + new TestBatchWrite with DeltaBatchWrite { + override def createBatchWriterFactory( + info: PhysicalWriteInfo): DeltaWriterFactory = { + // Use the narrow update schema for UPDATE/COPY/REINSERT rows when present. + new DeltaBufferedRowsWriterFactory(capturedWriteSchema) + } + + // For column-update writes, rows contain only the assigned columns + // (narrow schema from LogicalWriteInfo). We expand each row to the full table + // schema by overlaying write-schema columns on the base row found by pk. + override protected def doCommit(messages: Array[WriterCommitMessage]): Unit = + dataMap.synchronized { + val newData = messages.map(_.asInstanceOf[BufferedRows]) + val writeSchema = capturedWriteSchema + val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap + + val mergedData = newData.map { buf => + val merged = new BufferedRows(buf.key, schema) + val updateOpName = UTF8String.fromString(Update.toString) + val insertOpName = UTF8String.fromString(Insert.toString) + buf.log.foreach { logRow => + val opName = logRow.getUTF8String(0) + if (opName == updateOpName) { + val pk = logRow.getInt(1) + val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow] + val baseRow = dataMap.values.iterator.flatten + .flatMap(_.rows) + .find(r => r.getInt(schema.fieldIndex("pk")) == pk) + val fullRow = new GenericInternalRow(schema.length) + baseRow.foreach { base => + for (i <- schema.fields.indices) { + val field = schema.fields(i) + val value = if (i < base.numFields) { + base.get(i, field.dataType) + } else { + ResolveDefaultColumns.getExistenceDefaultValue(field) + } + fullRow.update(i, value) + } + } + schema.fields.zipWithIndex.foreach { case (field, i) => + writeFieldIdx.get(field.name).foreach { j => + fullRow.update(i, narrowRow.get(j, field.dataType)) + } + } + merged.rows.append(fullRow) + } else if (opName == insertOpName) { + // INSERT rows arrive with the full table schema via writer.insert() + val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow] + merged.rows.append(insertRow.copy()) + } + } + merged + } + + withDeletes(newData) + withData(mergedData) + lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq + } + + override def abort(messages: Array[WriterCommitMessage]): Unit = {} + } + } + } + } + } + + class DeltaBasedColumnUpdateOperationWithReqAttrs( + command: Command, + reqCols: Array[String], + options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty()) + extends DeltaBasedColumnUpdateOperation(command, options = options) { + override def requiredDataAttributes(): Array[NamedReference] = reqCols.map(FieldReference(_)) + } + + class DeltaBasedColumnUpdateOperationFromInfo( + command: Command, + updatedCols: Seq[NamedReference], + options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty()) + extends DeltaBasedColumnUpdateOperation(command, updatedCols, options) { + } + + class DeltaBasedColumnUpdateSplitOperation( + command: Command, + updatedCols: Seq[NamedReference] = Nil, + options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty()) + extends DeltaBasedColumnUpdateOperation(command, updatedCols, options) { + override def representUpdateAsDeleteAndInsert(): Boolean = true + + override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = { + lastWriteInfo = info + // Capture info into a local val so nested writer/commit closures see a stable schema + // even if a subsequent newWriteBuilder call mutates lastWriteInfo. + val capturedInfo = info + val capturedWriteSchema = if (capturedInfo.updateSchema().isPresent) { + capturedInfo.updateSchema().get() + } else { + capturedInfo.schema() + } + new DeltaWriteBuilder { + override def build(): DeltaWrite = + new DeltaWrite with RequiresDistributionAndOrdering { + override def requiredDistribution(): Distribution = + Distributions.clustered(Array(PARTITION_COLUMN_REF)) + override def requiredOrdering(): Array[SortOrder] = Array[SortOrder]( + LogicalExpressions.sort( + PARTITION_COLUMN_REF, + SortDirection.ASCENDING, + SortDirection.ASCENDING.defaultNullOrdering())) + override def toBatch: DeltaBatchWrite = + new TestBatchWrite with DeltaBatchWrite { + override def createBatchWriterFactory( + info: PhysicalWriteInfo): DeltaWriterFactory = { + new DeltaBufferedRowsWriterFactory(capturedWriteSchema) + } + + // For delete+reinsert with narrow writes, the REINSERT row has only the + // connector-declared columns (requiredDataAttributes order). + // pk is the first field in the write schema (declared before updatedCols). + // Reconstruct the full row by overlaying the narrow row onto the original. + override protected def doCommit(messages: Array[WriterCommitMessage]): Unit = + dataMap.synchronized { + val newData = messages.map(_.asInstanceOf[BufferedRows]) + val writeSchema = capturedWriteSchema + val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap + val reinsertOpName = UTF8String.fromString(Reinsert.toString) + val pkIdx = writeFieldIdx("pk") + + val expandedData = newData.map { buf => + val expanded = new BufferedRows(buf.key, schema) + buf.log.foreach { logRow => + val opName = logRow.getUTF8String(0) + if (opName == reinsertOpName) { Review Comment: The split fixture's commit now handles `Insert` as well as `Reinsert`, and MERGE runs on the split path too (`column-update split: MERGE falls back to full-width rows`). ########## sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java: ########## @@ -82,6 +84,54 @@ 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} without also mixing in + * {@link SupportsDelta} (i.e. group-based operations) receive records here in the schema + * declared by {@link LogicalWriteInfo#updateSchema()}. An operation that mixes in both + * instead delivers narrow rows through {@link DeltaWriter#update} / {@link DeltaWriter#reinsert} + * and never calls this method. By default, delegates to {@link #writeUpdate(Object)} for + * connectors that do not need metadata; implementations that do need it should override this + * method directly. + * <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 overrides neither {@code writeUpdate} overload. + * + * @since 4.3.0 + */ + default void writeUpdate(T metadata, T record) throws IOException { + writeUpdate(record); + } + + /** + * Writes one updated, copied, or reinserted record without metadata. + * <p> + * Connectors that mix in {@link SupportsColumnUpdates} without also mixing in + * {@link SupportsDelta} (i.e. group-based operations) receive records here in the schema + * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must override this method, + * or {@link #writeUpdate(Object, Object)}, when mixing in {@link SupportsColumnUpdates} without + * {@link SupportsDelta}. + * <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 overrides neither {@code writeUpdate} overload. + * + * @since 4.3.0 + */ + default void writeUpdate(T record) throws IOException { Review Comment: The overloads are now `writeColumnUpdate(record)` and `writeColumnUpdate(metadata, record)`. The Javadoc says Spark calls the metadata overload when the operation declares `requiredMetadataAttributes()` and the record-only overload otherwise. The record-only default throws and the metadata overload delegates to it, so a writer must override the overload Spark calls. Tests cover both dispatches, a writer that overrides only `writeColumnUpdate(record)`, and a writer that overrides neither. ########## sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala: ########## @@ -82,11 +83,29 @@ class InMemoryRowLevelOperationTable private ( private final val noMetadata = properties.getOrDefault(NO_METADATA, "false") == "true" private final val useCatalystRuntimeFiltering = properties.getOrDefault(USE_CATALYST_RUNTIME_FILTERING, "false") == "true" + private final val COLUMN_UPDATE = "column-update" + private final val COLUMN_UPDATE_REQ_ATTRS = "column-update-req-attrs" + private final val COLUMN_UPDATE_COW = "column-update-cow" + private final val COLUMN_UPDATE_COW_NO_WRITE_UPDATE = "column-update-cow-no-write-update" + private final val COLUMN_UPDATE_FROM_INFO = "column-update-from-info" + private final val COLUMN_UPDATE_SPLIT = "column-update-split" + private final val COLUMN_UPDATE_SPLIT_REQ_ATTRS = "column-update-split-req-attrs" + private final val COLUMN_UPDATE_EMPTY_REQ_ATTRS = "column-update-empty-req-attrs" + private final val COLUMN_UPDATE_SPLIT_MISSING_ROW_ID = "column-update-split-missing-row-id" // used in row-level operation tests to verify replaced partitions var replacedPartitions: Seq[Seq[Any]] = Seq.empty // used in row-level operation tests to verify reported write schema var lastWriteInfo: LogicalWriteInfo = _ + // used in column-update tests to verify that Spark passed the correct updated column list + // to the connector via RowLevelOperationInfo.updatedColumns() + var lastUpdatedColumns: Array[NamedReference] = Array.empty Review Comment: `copy()` now carries `lastUpdatedColumns`, and the transactional table sets it on commit instead of when the operation is built. -- 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]
