anuragmantri commented on code in PR #55518: URL: https://github.com/apache/spark/pull/55518#discussion_r4341122600
########## 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: Done in 22332b3. The doc now says that an operation without `SupportsDelta` should locate original rows through metadata columns preserved on update, such as a file and a row position, and, if it locates rows by a data column, build an operation without the mix-in when `updatedColumns()` contains that column. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala: ########## @@ -109,22 +111,75 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper { WriteToDataSourceV2(Some(r), microBatchWrite, newQuery, customMetrics) case rd @ ReplaceData(r: DataSourceV2Relation, _, query, _, projections, _, None) => - val rowSchema = projections.rowProjection.schema + val (rowSchema, columnUpdateSchema) = + writeSchemas(projections.rowProjection.schema, hasNarrowRows(rd)) val metadataSchema = projections.metadataProjection.map(_.schema) val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap) - val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema) + val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema, + columnUpdateSchema) val write = writeBuilder.build() + validateColumnUpdateWrite(rd, write) val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, r.funCatalog) rd.copy(write = Some(write), query = newQuery) case wd @ WriteDelta(r: DataSourceV2Relation, _, query, _, projections, _, None) => val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap) - val deltaWriteBuilder = newDeltaWriteBuilder(r.table, writeOptions, projections) + val deltaWriteBuilder = + newDeltaWriteBuilder(r.table, writeOptions, projections, hasNarrowRows(wd)) val deltaWrite = deltaWriteBuilder.build() + validateColumnUpdateWrite(wd, deltaWrite) val newQuery = DistributionAndOrderingUtils.prepareQuery(deltaWrite, query, r.funCatalog) wd.copy(write = Some(deltaWrite), query = newQuery) } + /** + * Whether a row-level write delivers narrow rows for updated, copied, and reinserted records. + * This must match the writes whose write relation RewriteUpdateTable narrows, even when the + * declared columns cover the whole table. + */ + private[v2] def hasNarrowRows(write: RowLevelWrite): Boolean = { + write.operation.isInstanceOf[SupportsColumnUpdates] && write.operation.command == UPDATE Review Comment: Good idea. Done in 22332b3. `V2Writes.validateWriteRelation` runs before building every `ReplaceData` and `WriteDelta` write. The Javadoc of `requiredDataAttributes()` also says it must return the same columns on every call. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala: ########## @@ -109,22 +111,75 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper { WriteToDataSourceV2(Some(r), microBatchWrite, newQuery, customMetrics) case rd @ ReplaceData(r: DataSourceV2Relation, _, query, _, projections, _, None) => - val rowSchema = projections.rowProjection.schema + val (rowSchema, columnUpdateSchema) = + writeSchemas(projections.rowProjection.schema, hasNarrowRows(rd)) val metadataSchema = projections.metadataProjection.map(_.schema) val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap) - val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema) + val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema, + columnUpdateSchema) val write = writeBuilder.build() + validateColumnUpdateWrite(rd, write) val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, r.funCatalog) rd.copy(write = Some(write), query = newQuery) case wd @ WriteDelta(r: DataSourceV2Relation, _, query, _, projections, _, None) => val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap) - val deltaWriteBuilder = newDeltaWriteBuilder(r.table, writeOptions, projections) + val deltaWriteBuilder = + newDeltaWriteBuilder(r.table, writeOptions, projections, hasNarrowRows(wd)) val deltaWrite = deltaWriteBuilder.build() + validateColumnUpdateWrite(wd, deltaWrite) val newQuery = DistributionAndOrderingUtils.prepareQuery(deltaWrite, query, r.funCatalog) wd.copy(write = Some(deltaWrite), query = newQuery) } + /** + * Whether a row-level write delivers narrow rows for updated, copied, and reinserted records. + * This must match the writes whose write relation RewriteUpdateTable narrows, even when the + * declared columns cover the whole table. + */ + private[v2] def hasNarrowRows(write: RowLevelWrite): Boolean = { + write.operation.isInstanceOf[SupportsColumnUpdates] && write.operation.command == UPDATE + } + + /** + * Checks that the distribution and ordering of a write that delivers narrow rows reference only + * columns the write reads. Other columns stay in the query only until column pruning removes + * them, so they are rejected here whether or not pruning ran. + */ + private def validateColumnUpdateWrite(rowLevelWrite: RowLevelWrite, write: Write): Unit = { + write match { + case w: RequiresDistributionAndOrdering if hasNarrowRows(rowLevelWrite) => + val distribution = w.requiredDistribution match { + case d: ClusteredDistribution => d.clustering.toImmutableArraySeq + case d: OrderedDistribution => d.ordering.toImmutableArraySeq + case _ => Seq.empty + } + val readNames = rowLevelWrite.references.toSeq.map(_.name) + val unread = (distribution ++ w.requiredOrdering.toImmutableArraySeq) + .flatMap(_.references.toImmutableArraySeq) Review Comment: Done in 22332b3, as suggested. The helper moved to `V2ExpressionUtils`. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/RowLevelOperationInfoImpl.scala: ########## @@ -17,9 +17,15 @@ package org.apache.spark.sql.connector.write +import org.apache.spark.sql.connector.expressions.NamedReference import org.apache.spark.sql.connector.write.RowLevelOperation.Command import org.apache.spark.sql.util.CaseInsensitiveStringMap private[sql] case class RowLevelOperationInfoImpl( command: Command, - options: CaseInsensitiveStringMap) extends RowLevelOperationInfo + options: CaseInsensitiveStringMap, + private val updatedCols: Seq[NamedReference]) Review Comment: Done in 22332b3. `RowLevelOperationTable` keeps the info in a second parameter list and `withTable(newTable)` rebuilds the operation for another table from the same info. `ResolveBranch` would replace building the info, the operation and the table with `operationTable.withTable(newTable)`. This is a change on Iceberg either way. @aokolnychyi, could you check this one from Iceberg perspective? ########## 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: Thanks for this comment. I learned a new thing here. I documented this in 22332b3. ########## 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: Done in 22332b3. The default now throws `DATA_SOURCE_COLUMN_UPDATE_SCHEMA_NOT_IMPLEMENTED`. I added a test also to cover this. ########## 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: Done in 22332b3. ########## 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: I like this suggestion, it's cleaner than my reordering. I implemented this in 22332b3. This now applies to every row-level write, so DELETE and MERGE queries can also drop columns their projections do not read. To make sure this does not break anything, I ran the DSv2 suites and they all pass. ########## 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: Done in 22332b3. ########## 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: Done in 22332b3. The JVM-global recorder is gone. `executeAndCheckScan` now checks every scan of the target table in the statement's own executed plan and fails if there is none, and it and `executeAndCheckScans` compare field names exactly; `executeAndCheckScanIncludes` and `executeAndCheckScanExcludes` cover the include and exclude checks. `lastUpdatedColumns` starts as null, so `checkLastUpdatedColumns()` fails when no operation was built. ########## 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: Done in 22332b3. All `COLUMN_UPDATE_*` helpers take name parts and format them with `toSQLId` ########## 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 { + case Assignment(key: AttributeReference, value) + if !isIdentityAssignment(key, value) && !declaredIds.contains(key.exprId) => + key.name + }.distinct + if (missing.nonEmpty) { + throw QueryCompilationErrors.requiredDataAttributesMissingUpdatedColumnsError( + operation.getClass.getName, missing) + } + } + + /** + * For connectors that opt into narrow column updates AND represent UPDATE as delete + insert, + * reject reassignment of any row-ID column, unless `requiredDataAttributes()` covers every + * column in the relation. In that case the REINSERT payload already is the full row with the + * new row-ID value, and the DELETE half still carries the original row-ID via + * `newLazyRowIdProjection`, so reassignment is safe. Otherwise the REINSERT path has no + * row-ID channel to reconstruct columns outside `requiredDataAttributes()`. + */ + private def validateNoRowIdReassignment( + operation: RowLevelOperation, + relation: DataSourceV2Relation, + connectorDataAttrs: Seq[AttributeReference], + assignments: Seq[Assignment], + rowIdAttrs: Seq[Attribute]): Unit = { + val declaredIds = connectorDataAttrs.map(_.exprId).toSet + if (relation.output.forall(a => declaredIds.contains(a.exprId))) { Review Comment: Done in 22332b3: the exemption compares only the data columns. ########## 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: Done in 22332b3. ########## 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( + 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 + } + + private val clusterRef: NamedReference = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY)) + .map(FieldReference(_)) + .getOrElse(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.columnUpdateSchema().isPresent) { + capturedInfo.columnUpdateSchema().get() + } else { + capturedInfo.schema() + } + new DeltaWriteBuilder { + override def build(): DeltaWrite = + new DeltaWrite with RequiresDistributionAndOrdering { + + override def requiredDistribution(): Distribution = { + Distributions.clustered(Array(clusterRef)) + } + + override def requiredOrdering(): Array[SortOrder] = { + Array[SortOrder]( + LogicalExpressions.sort( + clusterRef, + SortDirection.ASCENDING, + SortDirection.ASCENDING.defaultNullOrdering()) + ) + } + + override def toBatch: DeltaBatchWrite = + new TestBatchWrite with DeltaBatchWrite { + override def createBatchWriterFactory( + info: PhysicalWriteInfo): DeltaWriterFactory = { + new DeltaBufferedRowsWriterFactory(capturedWriteSchema) + } + + // For column-update writes, updated rows contain only the declared columns + // (columnUpdateSchema from LogicalWriteInfo). We expand each row to the full + // table schema by overlaying them 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 = currentRowByPk(pk) + val fullRow = overlayNarrowRow(baseRow, narrowRow, writeFieldIdx) + 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 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 = { Review Comment: Done in 22332b3. The split operation now reuses the parent's write builder and overrides only the writer factory and the commit, so it honors `column-update-cluster-by`. The dynamic-options copies are dropped, as the inherited SPARK-57681 tests already run on both tables, and `target_deps` is wrapped in `withTempView`. ########## sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala: ########## @@ -92,6 +95,38 @@ object GroupBasedRowLevelOperationScanPlanning extends Rule[LogicalPlan] with Pr } } + /** + * Returns the columns of `relation` that the projections and filters directly above each read + * of the table use, as the rest of the query sees only their output. Only whole columns are + * pruned, as the query is not rewritten for a pruned nested schema. Review Comment: Done in 22332b3: the scan-builder paragraph of `SupportsColumnUpdates` now says that for operations without `SupportsDelta` only top-level columns are pruned, so a struct referenced through a nested field is read in full,. I will file the nested pruning for copy-on-write scans as a follow-up after this PR. ########## sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedColumnUpdateTableSuite.scala: ########## @@ -0,0 +1,754 @@ +/* + * 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.{AnalysisException, Row} +import org.apache.spark.sql.catalyst.optimizer.ColumnPruning +import org.apache.spark.sql.catalyst.plans.logical.Union +import org.apache.spark.sql.connector.catalog.{InMemoryBaseTable, InMemoryRowLevelOperationTable} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{IntegerType, StructField, StructType} + +class GroupBasedColumnUpdateTableSuite extends UpdateTableSuiteBase { Review Comment: Done in 22332b3. "update table with expression-based default values" and "update with current_timestamp default value using DEFAULT keyword" now create their tables with the suite's properties, so they run narrow in both column-update suites and on delta tables in the existing delta suites. The first one needed `pk` to be `NOT NULL`, since delta row IDs cannot be nullable; it had only ever run on copy-on-write tables. The fixtures now declare `pk` and `dep` only when the table has them. "update table with recursive CTE" stays on a plain table, since its table has no row ID column. ########## 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( + 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 + } + + private val clusterRef: NamedReference = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY)) + .map(FieldReference(_)) + .getOrElse(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.columnUpdateSchema().isPresent) { + capturedInfo.columnUpdateSchema().get() + } else { + capturedInfo.schema() + } + new DeltaWriteBuilder { + override def build(): DeltaWrite = + new DeltaWrite with RequiresDistributionAndOrdering { Review Comment: Done in 22332b3. With `no-metadata`, the delta column-update write is a plain `DeltaWrite` without distribution or ordering, like the parent, unless `column-update-cluster-by` is set, and `split-updates` on a column-update table now selects the split operation ("column-update: split-updates sends a narrow UPDATE as delete and reinsert"). The new `DeltaBasedNoMetadataColumnUpdateTableSuite` extends `UpdateTableSuiteBase` with the shared tests, so narrow rows go through `DeltaWritingSparkTask`; most of its tests fail when that task's UPDATE branch stops projecting the row. ########## 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: Done in 22332b3 ########## 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: Done in 22332b3. The split commit now locates the original row by the operation's `rowId()`, taking a data row-ID column from the narrow row and a metadata row-ID column from the preserved metadata, and fails unless exactly one row matches. The delta writer logs a non-integer row ID as null, and the commit resolves such a DELETE from its metadata. ########## 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( + 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 + } + + private val clusterRef: NamedReference = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY)) + .map(FieldReference(_)) + .getOrElse(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.columnUpdateSchema().isPresent) { + capturedInfo.columnUpdateSchema().get() + } else { + capturedInfo.schema() + } + new DeltaWriteBuilder { + override def build(): DeltaWrite = + new DeltaWrite with RequiresDistributionAndOrdering { + + override def requiredDistribution(): Distribution = { + Distributions.clustered(Array(clusterRef)) + } + + override def requiredOrdering(): Array[SortOrder] = { + Array[SortOrder]( + LogicalExpressions.sort( + clusterRef, + SortDirection.ASCENDING, + SortDirection.ASCENDING.defaultNullOrdering()) + ) + } + + override def toBatch: DeltaBatchWrite = + new TestBatchWrite with DeltaBatchWrite { + override def createBatchWriterFactory( + info: PhysicalWriteInfo): DeltaWriterFactory = { + new DeltaBufferedRowsWriterFactory(capturedWriteSchema) + } + + // For column-update writes, updated rows contain only the declared columns + // (columnUpdateSchema from LogicalWriteInfo). We expand each row to the full + // table schema by overlaying them 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 = currentRowByPk(pk) + val fullRow = overlayNarrowRow(baseRow, narrowRow, writeFieldIdx) + 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 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.columnUpdateSchema().isPresent) { + capturedInfo.columnUpdateSchema().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). The base row is + // looked up by `pk`, so `pk` must be declared for the write to commit. + // 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 insertOpName = UTF8String.fromString(Insert.toString) + val pkIdx = writeFieldIdx("pk") Review Comment: Done in 22332b3. ########## sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala: ########## @@ -129,6 +184,24 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] { V2ExpressionUtils.resolveRef[AttributeReference](FieldReference(name), plan) } + protected def isIdentityAssignment(key: Attribute, value: Expression): Boolean = { Review Comment: Done in 22332b3: `buildOriginalRowIdValues` now uses `isIdentityAssignment`. -- 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]
