From 29dc5154d17a68369d1b321ee46705ba031e760f Mon Sep 17 00:00:00 2001 From: Eames Trinh Date: Fri, 24 Jul 2026 10:43:00 +0000 Subject: [PATCH 1/4] [SPARK-56942][SQL] Support nested field references as row-level operation row IDs Row-level DELETE/UPDATE/MERGE resolve each SupportsDelta.rowId() reference to an attribute. A nested reference such as _metadata.row_index resolves to an Alias(GetStructField(...)), not an AttributeReference, so it previously failed with a ClassCastException and connectors could only use top-level columns as row IDs. Resolve row-id references as NamedExpression and materialize any nested field as a top-level extraction column on the read plan. Expand and MergeRows regenerate expression IDs, so carry the extraction as a trailing output column and rebind the row-id projection to it by output position. Resolve a reference whose head names a metadata column by its logical name so a user data column of the same physical name cannot shadow it, and share that resolution between the rewrite and WriteDelta's resolution check. --- .../analysis/RewriteDeleteFromTable.scala | 20 +- .../analysis/RewriteMergeIntoTable.scala | 98 ++++++---- .../analysis/RewriteRowLevelCommand.scala | 126 +++++++++++-- .../analysis/RewriteUpdateTable.scala | 73 +++++--- .../expressions/V2ExpressionUtils.scala | 21 +++ .../catalyst/plans/logical/v2Commands.scala | 5 +- .../expressions/V2ExpressionUtilsSuite.scala | 23 ++- .../connector/catalog/InMemoryBaseTable.scala | 6 +- .../InMemoryRowLevelOperationTable.scala | 19 +- .../DeltaBasedNestedRowIdTableSuite.scala | 174 ++++++++++++++++++ 10 files changed, 477 insertions(+), 88 deletions(-) create mode 100644 sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteDeleteFromTable.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteDeleteFromTable.scala index f8881e2077103..c7555698abc28 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteDeleteFromTable.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteDeleteFromTable.scala @@ -73,7 +73,7 @@ object RewriteDeleteFromTable extends RewriteRowLevelCommand { val metadataAttrs = resolveRequiredMetadataAttrs(relation, operationTable.operation) // construct a read relation and include all required metadata columns - val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs) + val readPlan = buildRelationWithAttrs(relation, operationTable, metadataAttrs) // construct a plan that contains unmatched rows in matched groups that must be carried over // such rows do not match the condition but have to be copied over as the source can replace @@ -82,7 +82,7 @@ object RewriteDeleteFromTable extends RewriteRowLevelCommand { // it is safe to negate the condition here as the predicate pushdown for group-based row-level // operations is handled in a special way val remainingRowsFilter = Not(EqualNullSafe(cond, TrueLiteral)) - val remainingRowsPlan = Filter(remainingRowsFilter, readRelation) + val remainingRowsPlan = Filter(remainingRowsFilter, readPlan) // build a plan to replace read groups in the table val writeRelation = relation.copy(table = operationTable) @@ -100,16 +100,22 @@ object RewriteDeleteFromTable extends RewriteRowLevelCommand { // resolve all needed attrs (e.g. row ID and any required metadata attrs) val operation = operationTable.operation.asInstanceOf[SupportsDelta] - val rowIdAttrs = resolveRowIdAttrs(relation, operation) + val rowIdRefs = resolveRowIdRefs(relation, operation) + val rowIdAttrs = rowIdRefs.rowIdAttrs val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation) - // construct a read relation and include all required metadata columns - val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs) + // construct a read plan with all required row ID and metadata columns + val readRelation = buildRelationWithAttrs( + relation, operationTable, metadataAttrs, rowIdRefs.scanAttrs) + val readPlan = withExtractedRowIds(readRelation, rowIdRefs.extractionAliases) // construct a plan that only contains records to delete - val deletedRowsPlan = Filter(cond, readRelation) + val deletedRowsPlan = Filter(cond, readPlan) val operationType = Alias(Literal(DELETE_OPERATION), OPERATION_COLUMN)() - val requiredWriteAttrs = nullifyMetadataOnDelete(dedupAttrs(rowIdAttrs ++ metadataAttrs)) + // preserve row IDs even when they are also required metadata attributes + val rowIdExprIds = rowIdAttrs.map(_.exprId).toSet + val metadataOnlyAttrs = metadataAttrs.filterNot(attr => rowIdExprIds.contains(attr.exprId)) + val requiredWriteAttrs = rowIdAttrs ++ nullifyMetadataOnDelete(metadataOnlyAttrs) val project = Project(operationType +: requiredWriteAttrs, deletedRowsPlan) // build a plan to write deletes to the table diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala index 714f815161aeb..41143e2afb5ee 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteMergeIntoTable.scala @@ -17,6 +17,7 @@ package org.apache.spark.sql.catalyst.analysis +import org.apache.spark.SparkException import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.expressions.{Alias, And, Attribute, AttributeReference, Exists, Expression, IsNotNull, Literal, MetadataAttribute, MonotonicallyIncreasingID, OuterReference, PredicateHelper, SubqueryExpression} import org.apache.spark.sql.catalyst.expressions.Literal.{FalseLiteral, TrueLiteral} @@ -270,18 +271,20 @@ object RewriteMergeIntoTable extends RewriteRowLevelCommand with PredicateHelper // resolve all needed attrs (e.g. row ID and any required metadata attrs) val rowAttrs = relation.output - val rowIdAttrs = resolveRowIdAttrs(relation, operation) + val rowIdRefs = resolveRowIdRefs(relation, operation) val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation) - // construct a read relation and include all required metadata columns - val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs) + // construct a read plan with all required row ID and metadata columns + val readRelation = buildRelationWithAttrs( + relation, operationTable, metadataAttrs, rowIdRefs.scanAttrs) + val readPlan = withExtractedRowIds(readRelation, rowIdRefs.extractionAliases) // if there is no NOT MATCHED BY SOURCE clause, predicates of the ON condition that // reference only the target table can be pushed down val (filteredReadRelation, joinCond) = if (notMatchedBySourceActions.isEmpty) { - pushDownTargetPredicates(readRelation, cond) + pushDownTargetPredicates(readPlan, cond) } else { - (readRelation, cond) + (readPlan, cond) } val checkCardinality = shouldCheckCardinality(matchedActions) @@ -289,20 +292,21 @@ object RewriteMergeIntoTable extends RewriteRowLevelCommand with PredicateHelper val joinType = chooseWriteDeltaJoinType(notMatchedActions, notMatchedBySourceActions) val joinPlan = join(filteredReadRelation, source, joinType, joinCond, checkCardinality) - val mergeRowsPlan = buildWriteDeltaMergeRowsPlan( - readRelation, joinPlan, matchedActions, notMatchedActions, - notMatchedBySourceActions, rowIdAttrs, checkCardinality, + val rowDelta = buildWriteDeltaMergeRowsPlan( + readPlan, joinPlan, matchedActions, notMatchedActions, + notMatchedBySourceActions, rowIdRefs, checkCardinality, operation.representUpdateAsDeleteAndInsert) // build a plan to write the row delta to the table val writeRelation = relation.copy(table = operationTable) - val projections = buildWriteDeltaProjections(mergeRowsPlan, rowAttrs, rowIdAttrs, metadataAttrs) + val projections = buildWriteDeltaProjections( + rowDelta.plan, rowAttrs, rowDelta.rowIdAttrs, metadataAttrs) val groupFilterCond = if (notMatchedBySourceActions.isEmpty && groupFilterEnabled) { Some(toGroupFilterCondition(relation, source, cond)) } else { None } - WriteDelta(writeRelation, cond, mergeRowsPlan, relation, projections, groupFilterCond) + WriteDelta(writeRelation, cond, rowDelta.plan, relation, projections, groupFilterCond) } private def chooseWriteDeltaJoinType( @@ -324,18 +328,27 @@ object RewriteMergeIntoTable extends RewriteRowLevelCommand with PredicateHelper } private def buildWriteDeltaMergeRowsPlan( - targetTable: DataSourceV2Relation, + targetTable: LogicalPlan, joinPlan: LogicalPlan, matchedActions: Seq[MergeAction], notMatchedActions: Seq[MergeAction], notMatchedBySourceActions: Seq[MergeAction], - rowIdAttrs: Seq[Attribute], + rowIdRefs: ResolvedRowIdRefs, checkCardinality: Boolean, - splitUpdates: Boolean): MergeRows = { + splitUpdates: Boolean): RowDeltaPlan = { - val (metadataAttrs, rowAttrs) = targetTable.output.partition { attr => + val extractionExprIds = rowIdRefs.extractionAttrs.map(_.exprId).toSet + val (extractionAttrs, nonExtractionAttrs) = targetTable.output.partition { attr => + extractionExprIds.contains(attr.exprId) + } + val (metadataAttrs, rowAttrs) = nonExtractionAttrs.partition { attr => MetadataAttribute.isValid(attr.metadata) } + // buildRowDeltaPlan rebinds extracted row IDs by position + if (extractionAttrs.map(_.exprId) != rowIdRefs.extractionAttrs.map(_.exprId)) { + throw SparkException.internalError( + "Extracted row ID attributes are missing or out of order") + } val originalRowIdValues = if (splitUpdates) { Seq.empty @@ -346,19 +359,25 @@ object RewriteMergeIntoTable extends RewriteRowLevelCommand with PredicateHelper case UpdateAction(_, assignments, _) => assignments case _ => Nil } - buildOriginalRowIdValues(rowIdAttrs, updateAssignments) + buildOriginalRowIdValues(rowIdRefs.rowIdAttrs, updateAssignments) } val matchedInstructions = matchedActions.map { action => - toInstruction(action, rowAttrs, rowIdAttrs, metadataAttrs, originalRowIdValues, splitUpdates) + toInstruction( + action, rowAttrs, rowIdRefs.rowIdAttrs, metadataAttrs, originalRowIdValues, + extractionAttrs, splitUpdates) } val notMatchedInstructions = notMatchedActions.map { action => - toInstruction(action, rowAttrs, rowIdAttrs, metadataAttrs, originalRowIdValues, splitUpdates) + toInstruction( + action, rowAttrs, rowIdRefs.rowIdAttrs, metadataAttrs, originalRowIdValues, + extractionAttrs, splitUpdates) } val notMatchedBySourceInstructions = notMatchedBySourceActions.map { action => - toInstruction(action, rowAttrs, rowIdAttrs, metadataAttrs, originalRowIdValues, splitUpdates) + toInstruction( + action, rowAttrs, rowIdRefs.rowIdAttrs, metadataAttrs, originalRowIdValues, + extractionAttrs, splitUpdates) } val rowFromSourceAttr = resolveAttrRef(ROW_FROM_SOURCE, joinPlan) @@ -370,17 +389,19 @@ object RewriteMergeIntoTable extends RewriteRowLevelCommand with PredicateHelper val operationTypeAttr = AttributeReference(OPERATION_COLUMN, IntegerType, nullable = false)() val originalRowIdAttrs = originalRowIdValues.map(_.toAttribute) - val attrs = Seq(operationTypeAttr) ++ targetTable.output ++ originalRowIdAttrs - - MergeRows( - isSourceRowPresent = IsNotNull(rowFromSourceAttr), - isTargetRowPresent = IsNotNull(rowFromTargetAttr), - matchedInstructions = matchedInstructions, - notMatchedInstructions = notMatchedInstructions, - notMatchedBySourceInstructions = notMatchedBySourceInstructions, - checkCardinality = checkCardinality, - output = generateExpandOutput(attrs, outputs), - joinPlan) + val baseAttrs = Seq(operationTypeAttr) ++ rowAttrs ++ metadataAttrs ++ originalRowIdAttrs + + buildRowDeltaPlan(baseAttrs, outputs, rowIdRefs) { output => + MergeRows( + isSourceRowPresent = IsNotNull(rowFromSourceAttr), + isTargetRowPresent = IsNotNull(rowFromTargetAttr), + matchedInstructions = matchedInstructions, + notMatchedInstructions = notMatchedInstructions, + notMatchedBySourceInstructions = notMatchedBySourceInstructions, + checkCardinality = checkCardinality, + output = output, + joinPlan) + } } private def pushDownTargetPredicates( @@ -472,24 +493,33 @@ object RewriteMergeIntoTable extends RewriteRowLevelCommand with PredicateHelper rowIdAttrs: Seq[Attribute], metadataAttrs: Seq[Attribute], originalRowIdValues: Seq[Alias], + extractionAttrs: Seq[Attribute], splitUpdates: Boolean): Instruction = { + // inserts and reinserts have no existing row ID, so use null placeholders + val nullExtractionValues = extractionAttrs.map(attr => Literal(null, attr.dataType)) + action match { case UpdateAction(cond, assignments, _) if splitUpdates => - val output = deltaDeleteOutput(rowAttrs, rowIdAttrs, metadataAttrs, originalRowIdValues) - val otherOutput = deltaReinsertOutput(assignments, metadataAttrs, originalRowIdValues) + val output = deltaDeleteOutput( + rowAttrs, rowIdAttrs, metadataAttrs, originalRowIdValues) ++ extractionAttrs + val otherOutput = deltaReinsertOutput( + assignments, metadataAttrs, originalRowIdValues) ++ nullExtractionValues Split(cond.getOrElse(TrueLiteral), output, otherOutput) case UpdateAction(cond, assignments, _) => - val output = deltaUpdateOutput(assignments, metadataAttrs, originalRowIdValues) + val output = deltaUpdateOutput( + assignments, metadataAttrs, originalRowIdValues) ++ extractionAttrs Keep(Update, cond.getOrElse(TrueLiteral), output) case DeleteAction(cond) => - val output = deltaDeleteOutput(rowAttrs, rowIdAttrs, metadataAttrs, originalRowIdValues) + val output = deltaDeleteOutput( + rowAttrs, rowIdAttrs, metadataAttrs, originalRowIdValues) ++ extractionAttrs Keep(Delete, cond.getOrElse(TrueLiteral), output) case InsertAction(cond, assignments) => - val output = deltaInsertOutput(assignments, metadataAttrs, originalRowIdValues) + val output = deltaInsertOutput( + assignments, metadataAttrs, originalRowIdValues) ++ nullExtractionValues Keep(Insert, cond.getOrElse(TrueLiteral), output) case other => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala index 60143ae46c3c9..fac436d4ce833 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala @@ -28,7 +28,7 @@ import org.apache.spark.sql.catalyst.util.{GeneratedColumn, ReplaceDataProjectio WriteDeltaProjections} import org.apache.spark.sql.catalyst.util.RowDeltaUtils._ import org.apache.spark.sql.connector.catalog.SupportsRowLevelOperations -import org.apache.spark.sql.connector.expressions.FieldReference +import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference} import org.apache.spark.sql.connector.write.{RowLevelOperation, RowLevelOperationInfoImpl, RowLevelOperationTable, SupportsDelta} import org.apache.spark.sql.connector.write.RowLevelOperation.Command import org.apache.spark.sql.errors.QueryCompilationErrors @@ -100,29 +100,108 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] { } } + protected case class ResolvedRowIdRefs( + // row ID attributes used by the write projections + rowIdAttrs: Seq[AttributeReference], + // aliases that materialize nested row IDs as top-level attributes + extractionAliases: Seq[Alias], + // scan attributes needed to read or extract the row IDs + scanAttrs: Seq[AttributeReference]) { + + lazy val extractionAttrs: Seq[AttributeReference] = extractionAliases.map { alias => + rowIdAttrs.find(_.exprId == alias.exprId).getOrElse { + throw SparkException.internalError(s"Cannot find extracted row ID attribute: $alias") + } + } + + // rebind nested row IDs after Expand or MergeRows generates fresh output attributes + // top-level row IDs retain their original attributes for update handling + def rebindExtractions(reboundAttrs: Seq[Attribute]): Seq[Attribute] = { + if (extractionAttrs.size != reboundAttrs.size) { + throw SparkException.internalError( + s"Expected ${extractionAttrs.size} extracted row ID attributes, " + + s"but found ${reboundAttrs.size}") + } + val reboundByExprId = extractionAttrs.zip(reboundAttrs).map { case (attr, reboundAttr) => + attr.exprId -> reboundAttr + }.toMap + rowIdAttrs.map(attr => reboundByExprId.getOrElse(attr.exprId, attr)) + } + } + + protected case class RowDeltaPlan(plan: LogicalPlan, rowIdAttrs: Seq[Attribute]) + + // generate fresh output attributes and rebind nested row IDs to the trailing extracted fields + protected def buildRowDeltaPlan( + baseAttrs: Seq[Attribute], + outputs: Seq[Seq[Expression]], + rowIdRefs: ResolvedRowIdRefs)(buildPlan: Seq[Attribute] => LogicalPlan): RowDeltaPlan = { + val attrs = baseAttrs ++ rowIdRefs.extractionAttrs + outputs.find(_.size != attrs.size).foreach { output => + throw SparkException.internalError( + s"Expected ${attrs.size} row delta values, but found ${output.size}") + } + val output = generateExpandOutput(attrs, outputs) + // nested row ID attributes are the output suffix + val reboundRowIdAttrs = rowIdRefs.rebindExtractions( + output.takeRight(rowIdRefs.extractionAttrs.size)) + RowDeltaPlan(buildPlan(output), reboundRowIdAttrs) + } + + private def resolveRefs( + relation: DataSourceV2Relation, + refs: Seq[NamedReference]): ResolvedRowIdRefs = { + val rowIdAttrs = mutable.ArrayBuffer.empty[AttributeReference] + val extractionAliases = mutable.ArrayBuffer.empty[Alias] + val scanAttrs = mutable.ArrayBuffer.empty[AttributeReference] + refs.foreach { ref => + V2ExpressionUtils.resolveRowIdRef(ref, relation) match { + case attr: AttributeReference => + rowIdAttrs += attr + scanAttrs += attr + case alias: Alias => + extractionAliases += alias + scanAttrs ++= alias.references.collect { case ref: AttributeReference => ref } + alias.toAttribute match { + case attr: AttributeReference => rowIdAttrs += attr + case other => + throw SparkException.internalError(s"Row ID reference did not resolve: $other") + } + case other => + throw SparkException.internalError("Unexpected resolved row-level reference: " + other) + } + } + ResolvedRowIdRefs( + rowIdAttrs.toSeq, extractionAliases.toSeq, dedupAttrs(scanAttrs.toSeq)) + } + protected def resolveRequiredMetadataAttrs( relation: DataSourceV2Relation, operation: RowLevelOperation): Seq[AttributeReference] = { - V2ExpressionUtils.resolveRefs[AttributeReference]( - operation.requiredMetadataAttributes.toImmutableArraySeq, - relation) + operation.requiredMetadataAttributes.toImmutableArraySeq, relation) } - protected def resolveRowIdAttrs( + protected def resolveRowIdRefs( relation: DataSourceV2Relation, - operation: SupportsDelta): Seq[AttributeReference] = { - - val rowIdAttrs = V2ExpressionUtils.resolveRefs[AttributeReference]( - operation.rowId.toImmutableArraySeq, - relation) - - val nullableRowIdAttrs = rowIdAttrs.filter(_.nullable) + operation: SupportsDelta): ResolvedRowIdRefs = { + val resolved = resolveRefs(relation, operation.rowId.toImmutableArraySeq) + val nullableRowIdAttrs = resolved.rowIdAttrs.filter(_.nullable) if (nullableRowIdAttrs.nonEmpty) { throw QueryCompilationErrors.nullableRowIdError(nullableRowIdAttrs) } + resolved + } - rowIdAttrs + // materialize nested row IDs as top-level attributes + protected def withExtractedRowIds( + plan: LogicalPlan, + extractionAliases: Seq[Alias]): LogicalPlan = { + if (extractionAliases.isEmpty) { + plan + } else { + Project(plan.output ++ extractionAliases, plan) + } } protected def resolveAttrRef(name: String, plan: LogicalPlan): AttributeReference = { @@ -279,7 +358,11 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] { plan: LogicalPlan, outputs: Seq[Seq[Expression]], attrs: Seq[Attribute]): ProjectingInternalRow = { - val colOrdinals = attrs.map(attr => findColOrdinal(plan, attr.name)) + val colOrdinals = attrs.map { attr => + val byExprId = findColOrdinalByExprId(plan, attr.exprId) + // generated outputs have fresh exprIds, so fall back to the column name + if (byExprId != -1) byExprId else findColOrdinal(plan, attr.name) + } createProjectingInternalRow(outputs, colOrdinals, attrs) } @@ -290,8 +373,15 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] { outputs: Seq[Seq[Expression]], rowIdAttrs: Seq[Attribute]): ProjectingInternalRow = { val colOrdinals = rowIdAttrs.map { attr => - val originalValueIndex = findColOrdinal(plan, ORIGINAL_ROW_ID_VALUE_PREFIX + attr.name) - if (originalValueIndex != -1) originalValueIndex else findColOrdinal(plan, attr.name) + // prefer exprId to avoid binding nested row IDs to same-named data columns + // updated top-level row IDs fall back to their saved original values + val byExprId = findColOrdinalByExprId(plan, attr.exprId) + if (byExprId != -1) { + byExprId + } else { + val originalValueIndex = findColOrdinal(plan, ORIGINAL_ROW_ID_VALUE_PREFIX + attr.name) + if (originalValueIndex != -1) originalValueIndex else findColOrdinal(plan, attr.name) + } } createProjectingInternalRow(outputs, colOrdinals, rowIdAttrs) } @@ -311,6 +401,10 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] { plan.output.indexWhere(attr => conf.resolver(attr.name, name)) } + private def findColOrdinalByExprId(plan: LogicalPlan, exprId: ExprId): Int = { + plan.output.indexWhere(_.exprId == exprId) + } + protected def buildOriginalRowIdValues( rowIdAttrs: Seq[Attribute], assignments: Seq[Assignment]): Seq[Alias] = { diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala index 207b3e7217040..692ddb798f392 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala @@ -17,6 +17,7 @@ package org.apache.spark.sql.catalyst.analysis +import org.apache.spark.SparkException import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, EqualNullSafe, Expression, If, Literal, MetadataAttribute, Not, SubqueryExpression} import org.apache.spark.sql.catalyst.expressions.Literal.TrueLiteral import org.apache.spark.sql.catalyst.plans.logical.{Assignment, Expand, Filter, LogicalPlan, Project, ReplaceData, Union, UpdateTable, WriteDelta} @@ -69,10 +70,10 @@ object RewriteUpdateTable extends RewriteRowLevelCommand { val metadataAttrs = resolveRequiredMetadataAttrs(relation, operationTable.operation) // construct a read relation and include all required metadata columns - val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs) + val readPlan = buildRelationWithAttrs(relation, operationTable, metadataAttrs) // build a plan with updated and copied over records - val query = buildReplaceDataUpdateProjection(readRelation, assignments, cond) + val query = buildReplaceDataUpdateProjection(readPlan, assignments, cond) // build a plan to replace read groups in the table val writeRelation = relation.copy(table = operationTable) @@ -95,16 +96,16 @@ object RewriteUpdateTable extends RewriteRowLevelCommand { // construct a read relation and include all required metadata columns // the same read relation will be used to read records that must be updated and copied over // the analyzer will take care of duplicated attr IDs - val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs) + val readPlan = buildRelationWithAttrs(relation, operationTable, metadataAttrs) // build a plan for updated records that match the condition - val matchedRowsPlan = Filter(cond, readRelation) + val matchedRowsPlan = Filter(cond, readPlan) val updatedRowsPlan = buildReplaceDataUpdateProjection(matchedRowsPlan, assignments) // build a plan that contains unmatched rows in matched groups that must be copied over val remainingRowFilter = Not(EqualNullSafe(cond, Literal.TrueLiteral)) val remainingRowsPlan = addOperationColumn(COPY_OPERATION, - Filter(remainingRowFilter, readRelation)) + Filter(remainingRowFilter, readPlan)) // the new state is a union of updated and copied over records val query = Union(updatedRowsPlan, remainingRowsPlan) @@ -157,40 +158,50 @@ object RewriteUpdateTable extends RewriteRowLevelCommand { // resolve all needed attrs (e.g. row ID and any required metadata attrs) val rowAttrs = relation.output - val rowIdAttrs = resolveRowIdAttrs(relation, operation) + val rowIdRefs = resolveRowIdRefs(relation, operation) val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation) + val rowIdAttrs = rowIdRefs.rowIdAttrs - // construct a read relation and include all required metadata columns - val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs) + // construct a read plan with all required row ID and metadata columns + val readRelation = buildRelationWithAttrs( + relation, operationTable, metadataAttrs, rowIdRefs.scanAttrs) + val readPlan = withExtractedRowIds(readRelation, rowIdRefs.extractionAliases) // 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 matchedRowsPlan = Filter(cond, readPlan) + val rowDelta = if (operation.representUpdateAsDeleteAndInsert) { + buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdRefs) } else { - buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs) + val updatePlan = buildWriteDeltaUpdateProjection( + matchedRowsPlan, assignments, rowIdAttrs, rowIdRefs.extractionAttrs) + RowDeltaPlan(updatePlan, 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( + rowDelta.plan, rowAttrs, rowDelta.rowIdAttrs, metadataAttrs) val groupFilterCond = if (groupFilterEnabled) Some(cond) else None - WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond) + WriteDelta(writeRelation, cond, rowDelta.plan, relation, projections, groupFilterCond) } // this method assumes the assignments have been already aligned before private def buildWriteDeltaUpdateProjection( plan: LogicalPlan, assignments: Seq[Assignment], - rowIdAttrs: Seq[Attribute]): LogicalPlan = { + rowIdAttrs: Seq[Attribute], + rowIdExtractionAttrs: Seq[Attribute]): LogicalPlan = { - // the plan output may include immutable metadata columns at the end - // that's why the number of assignments may not match the number of plan output columns + // the plan output may include immutable metadata columns and extracted row IDs at the end val assignedValues = assignments.map(_.value) + val extractionExprIds = rowIdExtractionAttrs.map(_.exprId).toSet val updatedValues = plan.output.zipWithIndex.map { case (attr, index) => if (index < assignments.size) { val assignedExpr = assignedValues(index) Alias(assignedExpr, attr.name)() + } else if (extractionExprIds.contains(attr.exprId)) { + // preserve the original nested row ID for update encoding + attr } else { assert(MetadataAttribute.isValid(attr.metadata)) if (MetadataAttribute.isPreservedOnUpdate(attr)) { @@ -213,17 +224,31 @@ object RewriteUpdateTable extends RewriteRowLevelCommand { private def buildDeletesAndInserts( matchedRowsPlan: LogicalPlan, assignments: Seq[Assignment], - rowIdAttrs: Seq[Attribute]): Expand = { + rowIdRefs: ResolvedRowIdRefs): RowDeltaPlan = { - val (metadataAttrs, rowAttrs) = matchedRowsPlan.output.partition { attr => + val extractionExprIds = rowIdRefs.extractionAttrs.map(_.exprId).toSet + val (extractionAttrs, nonExtractionAttrs) = matchedRowsPlan.output.partition { attr => + extractionExprIds.contains(attr.exprId) + } + val (metadataAttrs, rowAttrs) = nonExtractionAttrs.partition { attr => MetadataAttribute.isValid(attr.metadata) } - val deleteOutput = deltaDeleteOutput(rowAttrs, rowIdAttrs, metadataAttrs) - val insertOutput = deltaReinsertOutput(assignments, metadataAttrs) + // buildRowDeltaPlan rebinds extracted row IDs by position + if (extractionAttrs.map(_.exprId) != rowIdRefs.extractionAttrs.map(_.exprId)) { + throw SparkException.internalError( + "Extracted row ID attributes are missing or out of order") + } + + // only the delete half needs the old row ID + val deleteOutput = + deltaDeleteOutput(rowAttrs, rowIdRefs.rowIdAttrs, metadataAttrs) ++ extractionAttrs + val insertOutput = deltaReinsertOutput(assignments, metadataAttrs) ++ + extractionAttrs.map(attr => Literal(null, attr.dataType)) 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) + val baseAttrs = operationTypeAttr +: (rowAttrs ++ metadataAttrs) + buildRowDeltaPlan(baseAttrs, outputs, rowIdRefs) { output => + Expand(outputs, output, matchedRowsPlan) + } } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala index 702255e075743..4956224b94eef 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala @@ -60,6 +60,27 @@ object V2ExpressionUtils extends SQLConfHelper with Logging { refs.map(ref => resolveRef[T](ref, plan)) } + // resolve metadata-rooted row IDs by logical name to avoid collisions with data columns + def resolveRowIdRef(ref: NamedReference, plan: LogicalPlan): NamedExpression = { + val parts = ref.fieldNames.toImmutableArraySeq + plan.getMetadataAttributeByNameOpt(parts.head) match { + case Some(struct) if parts.length == 1 => + struct + case Some(struct) => + val extracted = parts.tail.foldLeft(struct: Expression) { (base, field) => + ExtractValue.applyOrNull(base, Literal(field), conf.resolver) + } + Alias(extracted, parts.last)() + case None => + resolveRef[NamedExpression](ref, plan) + } + } + + def resolveRowIdRefs( + refs: Seq[NamedReference], plan: LogicalPlan): Seq[NamedExpression] = { + refs.map(ref => resolveRowIdRef(ref, plan)) + } + /** * Resolves [[NamedReference]]s against the given output and returns them as an [[AttributeSet]]. */ diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala index e82a65c8f99e7..197135c29f536 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala @@ -514,9 +514,10 @@ case class WriteDelta( // validates row ID projection output is compatible with row ID attributes private def rowIdAttrsResolved: Boolean = { - val outRowIdAttrs = V2ExpressionUtils.resolveRefs[AttributeReference]( + // must match the rewrite so a data column can't shadow a metadata-rooted row ID + val outRowIdAttrs = V2ExpressionUtils.resolveRowIdRefs( operation.rowId.toImmutableArraySeq, - originalTable) + originalTable).map(_.toAttribute) val inRowIdAttrs = DataTypeUtils.toAttributes(projections.rowIdProjection.schema) areCompatible(inRowIdAttrs, outRowIdAttrs) } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtilsSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtilsSuite.scala index d1c23d68555af..6efbebe849516 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtilsSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtilsSuite.scala @@ -20,8 +20,9 @@ package org.apache.spark.sql.catalyst.expressions import org.apache.spark.SparkFunSuite import org.apache.spark.sql.AnalysisException import org.apache.spark.sql.catalyst.plans.logical.LocalRelation +import org.apache.spark.sql.catalyst.util.METADATA_COL_ATTR_KEY import org.apache.spark.sql.connector.expressions._ -import org.apache.spark.sql.types.StringType +import org.apache.spark.sql.types.{IntegerType, MetadataBuilder, StringType, StructType} class V2ExpressionUtilsSuite extends SparkFunSuite { @@ -37,4 +38,24 @@ class V2ExpressionUtilsSuite extends SparkFunSuite { } assert(exc.message.contains("v2Fun(a) ASC NULLS FIRST is not currently supported")) } + + test("resolveRowIdRef resolves a metadata-rooted ref past a colliding data column") { + // simulate a metadata column renamed after colliding with a user `_metadata` column + val userColumn = AttributeReference("_metadata", IntegerType)() + val metadataStruct = AttributeReference( + "__metadata", + new StructType().add("row_index", IntegerType), + nullable = false, + metadata = new MetadataBuilder().putString(METADATA_COL_ATTR_KEY, "_metadata").build())() + val dataColumn = AttributeReference("id", IntegerType)() + val plan = LocalRelation(Seq(userColumn, metadataStruct, dataColumn)) + + val metadataRowId = V2ExpressionUtils.resolveRowIdRef( + FieldReference(Seq("_metadata", "row_index")), plan) + assert(metadataRowId.references.contains(metadataStruct)) + assert(!metadataRowId.references.contains(userColumn)) + + val dataRowId = V2ExpressionUtils.resolveRowIdRef(FieldReference("id"), plan) + assert(dataRowId == dataColumn) + } } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala index 78775d80cfa6d..bafdd14166003 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala @@ -376,12 +376,14 @@ abstract class InMemoryBaseTable( dataMap.update(key, Seq(new BufferedRows(key, schema))) } - def withDeletes(data: Array[BufferedRows]): InMemoryBaseTable = { + def withDeletes( + data: Array[BufferedRows], + rowId: InternalRow => Int = _.getInt(0)): InMemoryBaseTable = { data.foreach { p => dataMap ++= dataMap.map { case (key, currentSplits) => val newSplits = currentSplits.map { currentRows => val newRows = new BufferedRows(currentRows.key, currentRows.schema) - newRows.rows ++= currentRows.rows.filter(r => !p.deletes.contains(r.getInt(0))) + newRows.rows ++= currentRows.rows.filter(r => !p.deletes.contains(rowId(r))) newRows } key -> newSplits diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala index f62f1b0dce25e..e8a7f1519af57 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala @@ -76,8 +76,10 @@ class InMemoryRowLevelOperationTable private ( private final val INDEX_COLUMN_REF = FieldReference(IndexColumn.name) private final val SUPPORTS_DELTAS = "supports-deltas" private final val SPLIT_UPDATES = "split-updates" + private final val NESTED_ROW_ID = "nested-row-id" private final val NO_METADATA = "no-metadata" private final val noMetadata = properties.getOrDefault(NO_METADATA, "false") == "true" + private final val nestedRowId = properties.getOrDefault(NESTED_ROW_ID, "false") == "true" // used in row-level operation tests to verify replaced partitions var replacedPartitions: Seq[Seq[Any]] = Seq.empty @@ -195,6 +197,7 @@ class InMemoryRowLevelOperationTable private ( case class DeltaBasedOperation(command: Command, options: CaseInsensitiveStringMap) extends RowLevelOperation with SupportsDelta with RowLevelOperationWithOptions { private final val PK_COLUMN_REF = FieldReference("pk") + private final val NESTED_PK_COLUMN_REF = FieldReference(Seq("nested", "pk")) override def requiredMetadataAttributes(): Array[NamedReference] = { if (noMetadata) { @@ -204,7 +207,8 @@ class InMemoryRowLevelOperationTable private ( } } - override def rowId(): Array[NamedReference] = Array(PK_COLUMN_REF) + override def rowId(): Array[NamedReference] = + if (nestedRowId) Array(NESTED_PK_COLUMN_REF) else Array(PK_COLUMN_REF) override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { new InMemoryScanBuilder(schema, options) @@ -251,7 +255,18 @@ class InMemoryRowLevelOperationTable private ( override protected def doCommit(messages: Array[WriterCommitMessage]): Unit = { val newData = messages.map(_.asInstanceOf[BufferedRows]) - withDeletes(newData) + // delete rows using the configured row ID + val tableSchema = schema() + val rowId: InternalRow => Int = if (nestedRowId) { + val nestedOrdinal = tableSchema.fieldIndex("nested") + val nestedType = tableSchema(nestedOrdinal).dataType.asInstanceOf[StructType] + val pkOrdinal = nestedType.fieldIndex("pk") + row => row.getStruct(nestedOrdinal, nestedType.length).getInt(pkOrdinal) + } else { + val pkOrdinal = tableSchema.fieldIndex("pk") + row => row.getInt(pkOrdinal) + } + withDeletes(newData, rowId) withData(newData, columns()) lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala new file mode 100644 index 0000000000000..fdffbf7bfc16f --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala @@ -0,0 +1,174 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector + +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.util.METADATA_COL_ATTR_KEY +import org.apache.spark.sql.connector.catalog.CatalogV2Util +import org.apache.spark.sql.types.{IntegerType, MetadataBuilder, StringType, StructField, StructType} + +abstract class DeltaBasedNestedRowIdTableSuiteBase(splitUpdates: Boolean) + extends RowLevelOperationSuiteBase { + + import testImplicits._ + + override protected def extraTableProps: java.util.Map[String, String] = { + val props = new java.util.HashMap[String, String]() + props.put("supports-deltas", "true") + props.put("nested-row-id", "true") + if (splitUpdates) props.put("split-updates", "true") + props + } + + // mirror the metadata marker on file-source row IDs + private val nestedPkMetadata = + new MetadataBuilder().putString(METADATA_COL_ATTR_KEY, "row_index").build() + + private val tableSchema = StructType(Seq( + StructField("pk", IntegerType, nullable = false), + StructField("nested", StructType(Seq( + StructField("pk", IntegerType, nullable = false, metadata = nestedPkMetadata))), + nullable = false), + StructField("id", IntegerType), + StructField("dep", StringType))) + + // use different top-level and nested PK values to expose name-based binding + private val initialRows = + """{ "pk": 10, "nested": { "pk": 1 }, "id": 1, "dep": "hr" } + |{ "pk": 20, "nested": { "pk": 2 }, "id": 2, "dep": "software" } + |{ "pk": 30, "nested": { "pk": 3 }, "id": 3, "dep": "hr" } + |""".stripMargin + + private def createNestedRowIdTable(): Unit = { + createTable(CatalogV2Util.structTypeToV2Columns(tableSchema)) + append(tableSchema.toDDL, initialRows) + } + + private def checkTable(expected: Seq[Row]): Unit = { + checkAnswer( + sql(s"SELECT pk, nested.pk, id, dep FROM $tableNameAsString ORDER BY pk"), + expected) + } + + test("delete with nested row id") { + createNestedRowIdTable() + // use a multi-value IN to force the row-level delta path instead of metadata-only deleteWhere + sql(s"DELETE FROM $tableNameAsString WHERE id IN (1, 100)") + checkTable(Seq( + Row(20, 2, 2, "software"), + Row(30, 3, 3, "hr"))) + } + + test("update with nested row id") { + createNestedRowIdTable() + sql(s"UPDATE $tableNameAsString SET dep = 'it' WHERE id = 1") + checkTable(Seq( + Row(10, 1, 1, "it"), + Row(20, 2, 2, "software"), + Row(30, 3, 3, "hr"))) + } + + test("update replacing the struct that holds the nested row id") { + createNestedRowIdTable() + sql( + s"""UPDATE $tableNameAsString + |SET nested = named_struct('pk', 11) + |WHERE id = 1 + |""".stripMargin) + checkTable(Seq( + Row(10, 11, 1, "hr"), + Row(20, 2, 2, "software"), + Row(30, 3, 3, "hr"))) + } + + test("merge update with nested row id") { + withTempView("source") { + createNestedRowIdTable() + Seq((1, "it")).toDF("id", "dep").createOrReplaceTempView("source") + sql( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.id = s.id + |WHEN MATCHED THEN UPDATE SET t.dep = s.dep + |""".stripMargin) + checkTable(Seq( + Row(10, 1, 1, "it"), + Row(20, 2, 2, "software"), + Row(30, 3, 3, "hr"))) + } + } + + test("merge delete with nested row id") { + withTempView("source") { + createNestedRowIdTable() + Seq(2).toDF("id").createOrReplaceTempView("source") + sql( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.id = s.id + |WHEN MATCHED THEN DELETE + |""".stripMargin) + checkTable(Seq( + Row(10, 1, 1, "hr"), + Row(30, 3, 3, "hr"))) + } + } + + test("merge not matched by source update with nested row id") { + withTempView("source") { + createNestedRowIdTable() + Seq(1).toDF("id").createOrReplaceTempView("source") + sql( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.id = s.id + |WHEN NOT MATCHED BY SOURCE THEN UPDATE SET t.dep = 'gone' + |""".stripMargin) + checkTable(Seq( + Row(10, 1, 1, "hr"), + Row(20, 2, 2, "gone"), + Row(30, 3, 3, "gone"))) + } + } + + test("merge with matched update and not-matched insert with nested row id") { + withTempView("source") { + createNestedRowIdTable() + Seq((1, "it"), (4, "new")).toDF("id", "dep").createOrReplaceTempView("source") + sql( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.id = s.id + |WHEN MATCHED THEN UPDATE SET t.dep = s.dep + |WHEN NOT MATCHED THEN + |INSERT (pk, nested, id, dep) VALUES (s.id * 10, named_struct('pk', s.id), s.id, s.dep) + |""".stripMargin) + checkTable(Seq( + Row(10, 1, 1, "it"), + Row(20, 2, 2, "software"), + Row(30, 3, 3, "hr"), + Row(40, 4, 4, "new"))) + } + } +} + +class DeltaBasedNestedRowIdTableSuite + extends DeltaBasedNestedRowIdTableSuiteBase(splitUpdates = false) + +class DeltaBasedNestedRowIdUpdateAsDeleteAndInsertTableSuite + extends DeltaBasedNestedRowIdTableSuiteBase(splitUpdates = true) From cf815a57c1b8200832d48301e2382e941ecc056f Mon Sep 17 00:00:00 2001 From: Eames Trinh Date: Sat, 25 Jul 2026 10:41:13 +0000 Subject: [PATCH 2/4] [SPARK-56942][SQL] Remove unread metadata marker from nested row ID test The leaf StructField metadata marker was never read: getMetadataAttributeByNameOpt and MetadataAttributeWithLogicalName inspect only top-level AttributeReference metadata, and the test's row id (nested.pk) resolves through the plain reference path. The suite's real coverage is the leaf-name collision between the nested row id and the top-level pk data column, so drop the marker and correct the comment. Metadata-name collision remains covered by V2ExpressionUtilsSuite. --- .../connector/DeltaBasedNestedRowIdTableSuite.scala | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala index fdffbf7bfc16f..b2004002ef5df 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala @@ -18,9 +18,8 @@ package org.apache.spark.sql.connector import org.apache.spark.sql.Row -import org.apache.spark.sql.catalyst.util.METADATA_COL_ATTR_KEY import org.apache.spark.sql.connector.catalog.CatalogV2Util -import org.apache.spark.sql.types.{IntegerType, MetadataBuilder, StringType, StructField, StructType} +import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType} abstract class DeltaBasedNestedRowIdTableSuiteBase(splitUpdates: Boolean) extends RowLevelOperationSuiteBase { @@ -35,14 +34,12 @@ abstract class DeltaBasedNestedRowIdTableSuiteBase(splitUpdates: Boolean) props } - // mirror the metadata marker on file-source row IDs - private val nestedPkMetadata = - new MetadataBuilder().putString(METADATA_COL_ATTR_KEY, "row_index").build() - + // the nested row ID pk shares its leaf name with the top-level pk data column, so a name-based + // bind would confuse the two; the suite checks that the nested field is used private val tableSchema = StructType(Seq( StructField("pk", IntegerType, nullable = false), StructField("nested", StructType(Seq( - StructField("pk", IntegerType, nullable = false, metadata = nestedPkMetadata))), + StructField("pk", IntegerType, nullable = false))), nullable = false), StructField("id", IntegerType), StructField("dep", StringType))) From 5db6e7f7428d83bf193e8cd7e124bcdaa7b5ae74 Mon Sep 17 00:00:00 2001 From: Eames Trinh Date: Mon, 27 Jul 2026 22:37:28 +0000 Subject: [PATCH 3/4] [SPARK-56942][SQL] Address review: rename helper, revert churn, add collision tests - Rename the private resolveRefs helper to resolveRowIdRefsToAttrs and un-shadow the alias-reference variable. - Restore resolveRequiredMetadataAttrs formatting to reduce diff noise. - Assert the write log / write info in the nested delete and update tests so the nested row id (not the same-named top-level column) is shown to be used. - Add DELETE/UPDATE/MERGE tests where a data column shares a metadata column's name, covering metadata-vs-data resolution. --- .../analysis/RewriteRowLevelCommand.scala | 10 ++- .../DeltaBasedNestedRowIdTableSuite.scala | 73 +++++++++++++++++++ 2 files changed, 79 insertions(+), 4 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala index fac436d4ce833..87152f986cdcd 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala @@ -148,7 +148,7 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] { RowDeltaPlan(buildPlan(output), reboundRowIdAttrs) } - private def resolveRefs( + private def resolveRowIdRefsToAttrs( relation: DataSourceV2Relation, refs: Seq[NamedReference]): ResolvedRowIdRefs = { val rowIdAttrs = mutable.ArrayBuffer.empty[AttributeReference] @@ -161,7 +161,7 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] { scanAttrs += attr case alias: Alias => extractionAliases += alias - scanAttrs ++= alias.references.collect { case ref: AttributeReference => ref } + scanAttrs ++= alias.references.collect { case scanRef: AttributeReference => scanRef } alias.toAttribute match { case attr: AttributeReference => rowIdAttrs += attr case other => @@ -178,14 +178,16 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] { protected def resolveRequiredMetadataAttrs( relation: DataSourceV2Relation, operation: RowLevelOperation): Seq[AttributeReference] = { + V2ExpressionUtils.resolveRefs[AttributeReference]( - operation.requiredMetadataAttributes.toImmutableArraySeq, relation) + operation.requiredMetadataAttributes.toImmutableArraySeq, + relation) } protected def resolveRowIdRefs( relation: DataSourceV2Relation, operation: SupportsDelta): ResolvedRowIdRefs = { - val resolved = resolveRefs(relation, operation.rowId.toImmutableArraySeq) + val resolved = resolveRowIdRefsToAttrs(relation, operation.rowId.toImmutableArraySeq) val nullableRowIdAttrs = resolved.rowIdAttrs.filter(_.nullable) if (nullableRowIdAttrs.nonEmpty) { throw QueryCompilationErrors.nullableRowIdError(nullableRowIdAttrs) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala index b2004002ef5df..ced3344a51741 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala @@ -69,6 +69,11 @@ abstract class DeltaBasedNestedRowIdTableSuiteBase(splitUpdates: Boolean) checkTable(Seq( Row(20, 2, 2, "software"), Row(30, 3, 3, "hr"))) + // the row id written to the delta is the nested pk (1), not the top-level pk (10) + checkLastWriteInfo( + expectedRowIdSchema = Some(StructType(Array(PK_FIELD))), + expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE)))) + checkLastWriteLog(deleteWriteLogEntry(id = 1, metadata = Row("hr", null))) } test("update with nested row id") { @@ -78,6 +83,15 @@ abstract class DeltaBasedNestedRowIdTableSuiteBase(splitUpdates: Boolean) Row(10, 1, 1, "it"), Row(20, 2, 2, "software"), Row(30, 3, 3, "hr"))) + // the row id written to the delta is the nested pk (1), not the top-level pk (10) + if (splitUpdates) { + checkLastWriteLog( + deleteWriteLogEntry(id = 1, metadata = Row("hr", null)), + reinsertWriteLogEntry(metadata = Row("hr", null), data = Row(10, Row(1), 1, "it"))) + } else { + checkLastWriteLog( + updateWriteLogEntry(id = 1, metadata = Row("hr", null), data = Row(10, Row(1), 1, "it"))) + } } test("update replacing the struct that holds the nested row id") { @@ -162,6 +176,65 @@ abstract class DeltaBasedNestedRowIdTableSuiteBase(splitUpdates: Boolean) Row(40, 4, 4, "new"))) } } + + // A data column `index` collides with the in-memory table's `index` metadata column: the row + // level rewrite must still resolve the nested row id and the metadata column correctly. + private def createDataMetadataConflictTable(): Unit = { + val schema = StructType(Seq( + StructField("pk", IntegerType, nullable = false), + StructField("nested", StructType(Seq( + StructField("pk", IntegerType, nullable = false))), + nullable = false), + StructField("id", IntegerType), + StructField("index", IntegerType), + StructField("dep", StringType))) + createTable(CatalogV2Util.structTypeToV2Columns(schema)) + append(schema.toDDL, + """{ "pk": 10, "nested": { "pk": 1 }, "id": 1, "index": 100, "dep": "hr" } + |{ "pk": 20, "nested": { "pk": 2 }, "id": 2, "index": 200, "dep": "software" } + |{ "pk": 30, "nested": { "pk": 3 }, "id": 3, "index": 300, "dep": "hr" } + |""".stripMargin) + } + + private def checkConflictTable(expected: Seq[Row]): Unit = { + checkAnswer( + sql(s"SELECT pk, nested.pk, id, index, dep FROM $tableNameAsString ORDER BY pk"), + expected) + } + + test("delete with a data column named like a metadata column") { + createDataMetadataConflictTable() + sql(s"DELETE FROM $tableNameAsString WHERE id IN (1, 100)") + checkConflictTable(Seq( + Row(20, 2, 2, 200, "software"), + Row(30, 3, 3, 300, "hr"))) + } + + test("update with a data column named like a metadata column") { + createDataMetadataConflictTable() + sql(s"UPDATE $tableNameAsString SET dep = 'it' WHERE id = 1") + checkConflictTable(Seq( + Row(10, 1, 1, 100, "it"), + Row(20, 2, 2, 200, "software"), + Row(30, 3, 3, 300, "hr"))) + } + + test("merge with a data column named like a metadata column") { + withTempView("source") { + createDataMetadataConflictTable() + Seq((1, "it")).toDF("id", "dep").createOrReplaceTempView("source") + sql( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.id = s.id + |WHEN MATCHED THEN UPDATE SET t.dep = s.dep + |""".stripMargin) + checkConflictTable(Seq( + Row(10, 1, 1, 100, "it"), + Row(20, 2, 2, 200, "software"), + Row(30, 3, 3, 300, "hr"))) + } + } } class DeltaBasedNestedRowIdTableSuite From 74fdd3c877e60e217ab99ebdd84e8647deee6eb5 Mon Sep 17 00:00:00 2001 From: Eames Trinh Date: Tue, 28 Jul 2026 20:22:21 +0000 Subject: [PATCH 4/4] [SPARK-56942][SQL] Null metadata columns colliding with a nested row id When a nested row id (e.g. a.b) is materialized under its leaf name and a required metadata column shares that name, nullifyMetadata built the null alias with a fresh expression id. The metadata projection then failed to bind by id and fell back to name, picking the row id's same-named column, so a PRESERVE_ON_DELETE=false metadata column was written with the row id value instead of null. Preserve the original expression id on the null alias so the projection binds it by id. Add DELETE/UPDATE/MERGE tests asserting the write log for a nested row id whose leaf collides with a metadata column. --- .../analysis/RewriteRowLevelCommand.scala | 4 +- .../InMemoryRowLevelOperationTable.scala | 21 ++++-- .../DeltaBasedNestedRowIdTableSuite.scala | 72 +++++++++++++++++++ 3 files changed, 91 insertions(+), 6 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala index 87152f986cdcd..b5680755af5d9 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala @@ -237,7 +237,9 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] { shouldPreserve: Attribute => Boolean): Seq[NamedExpression] = { attrs.map { case MetadataAttribute(attr) if !shouldPreserve(attr) => - Alias(Literal(null, attr.dataType), attr.name)(explicitMetadata = Some(attr.metadata)) + // keep the exprId so the projection binds this by id, not by a name a row id may share + Alias(Literal(null, attr.dataType), attr.name)( + exprId = attr.exprId, explicitMetadata = Some(attr.metadata)) case attr => attr } diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala index e8a7f1519af57..848c551ef8c70 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala @@ -77,9 +77,13 @@ class InMemoryRowLevelOperationTable private ( private final val SUPPORTS_DELTAS = "supports-deltas" private final val SPLIT_UPDATES = "split-updates" private final val NESTED_ROW_ID = "nested-row-id" + // row id is nested.index, whose leaf collides with the `index` metadata column + private final val NESTED_METADATA_NAME_ROW_ID = "nested-metadata-name-row-id" private final val NO_METADATA = "no-metadata" private final val noMetadata = properties.getOrDefault(NO_METADATA, "false") == "true" private final val nestedRowId = properties.getOrDefault(NESTED_ROW_ID, "false") == "true" + private final val nestedMetadataNameRowId = + properties.getOrDefault(NESTED_METADATA_NAME_ROW_ID, "false") == "true" // used in row-level operation tests to verify replaced partitions var replacedPartitions: Seq[Seq[Any]] = Seq.empty @@ -198,6 +202,7 @@ class InMemoryRowLevelOperationTable private ( extends RowLevelOperation with SupportsDelta with RowLevelOperationWithOptions { private final val PK_COLUMN_REF = FieldReference("pk") private final val NESTED_PK_COLUMN_REF = FieldReference(Seq("nested", "pk")) + private final val NESTED_INDEX_COLUMN_REF = FieldReference(Seq("nested", "index")) override def requiredMetadataAttributes(): Array[NamedReference] = { if (noMetadata) { @@ -207,8 +212,13 @@ class InMemoryRowLevelOperationTable private ( } } - override def rowId(): Array[NamedReference] = - if (nestedRowId) Array(NESTED_PK_COLUMN_REF) else Array(PK_COLUMN_REF) + override def rowId(): Array[NamedReference] = if (nestedMetadataNameRowId) { + Array(NESTED_INDEX_COLUMN_REF) + } else if (nestedRowId) { + Array(NESTED_PK_COLUMN_REF) + } else { + Array(PK_COLUMN_REF) + } override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { new InMemoryScanBuilder(schema, options) @@ -257,11 +267,12 @@ class InMemoryRowLevelOperationTable private ( val newData = messages.map(_.asInstanceOf[BufferedRows]) // delete rows using the configured row ID val tableSchema = schema() - val rowId: InternalRow => Int = if (nestedRowId) { + val rowId: InternalRow => Int = if (nestedRowId || nestedMetadataNameRowId) { + val field = if (nestedMetadataNameRowId) "index" else "pk" val nestedOrdinal = tableSchema.fieldIndex("nested") val nestedType = tableSchema(nestedOrdinal).dataType.asInstanceOf[StructType] - val pkOrdinal = nestedType.fieldIndex("pk") - row => row.getStruct(nestedOrdinal, nestedType.length).getInt(pkOrdinal) + val fieldOrdinal = nestedType.fieldIndex(field) + row => row.getStruct(nestedOrdinal, nestedType.length).getInt(fieldOrdinal) } else { val pkOrdinal = tableSchema.fieldIndex("pk") row => row.getInt(pkOrdinal) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala index ced3344a51741..4306d234d91a3 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala @@ -242,3 +242,75 @@ class DeltaBasedNestedRowIdTableSuite class DeltaBasedNestedRowIdUpdateAsDeleteAndInsertTableSuite extends DeltaBasedNestedRowIdTableSuiteBase(splitUpdates = true) + +// Row id is nested.index, whose leaf name collides with the `index` metadata column +// (PRESERVE_ON_DELETE = false). The metadata must be nulled on write, not bound to the row id +// value it shares a name with. +abstract class DeltaBasedNestedRowIdMetadataCollisionSuiteBase(splitUpdates: Boolean) + extends RowLevelOperationSuiteBase { + + import testImplicits._ + + override protected def extraTableProps: java.util.Map[String, String] = { + val props = new java.util.HashMap[String, String]() + props.put("supports-deltas", "true") + props.put("nested-metadata-name-row-id", "true") + if (splitUpdates) props.put("split-updates", "true") + props + } + + private def createCollisionTable(): Unit = { + val schema = StructType(Seq( + StructField("pk", IntegerType, nullable = false), + StructField("nested", StructType(Seq( + StructField("index", IntegerType, nullable = false))), + nullable = false), + StructField("id", IntegerType), + StructField("dep", StringType))) + createTable(CatalogV2Util.structTypeToV2Columns(schema)) + append(schema.toDDL, + """{ "pk": 10, "nested": { "index": 1 }, "id": 1, "dep": "hr" } + |{ "pk": 20, "nested": { "index": 2 }, "id": 2, "dep": "software" } + |{ "pk": 30, "nested": { "index": 3 }, "id": 3, "dep": "hr" } + |""".stripMargin) + } + + test("delete nulls a metadata column colliding with the nested row id") { + createCollisionTable() + sql(s"DELETE FROM $tableNameAsString WHERE id IN (1, 100)") + checkLastWriteLog(deleteWriteLogEntry(id = 1, metadata = Row("hr", null))) + } + + test("update nulls a metadata column colliding with the nested row id") { + createCollisionTable() + sql(s"UPDATE $tableNameAsString SET dep = 'it' WHERE id = 1") + if (splitUpdates) { + checkLastWriteLog( + deleteWriteLogEntry(id = 1, metadata = Row("hr", null)), + reinsertWriteLogEntry(metadata = Row("hr", null), data = Row(10, Row(1), 1, "it"))) + } else { + checkLastWriteLog( + updateWriteLogEntry(id = 1, metadata = Row("hr", null), data = Row(10, Row(1), 1, "it"))) + } + } + + test("merge delete nulls a metadata column colliding with the nested row id") { + withTempView("source") { + createCollisionTable() + Seq(1).toDF("id").createOrReplaceTempView("source") + sql( + s"""MERGE INTO $tableNameAsString t + |USING source s + |ON t.id = s.id + |WHEN MATCHED THEN DELETE + |""".stripMargin) + checkLastWriteLog(deleteWriteLogEntry(id = 1, metadata = Row("hr", null))) + } + } +} + +class DeltaBasedNestedRowIdMetadataCollisionSuite + extends DeltaBasedNestedRowIdMetadataCollisionSuiteBase(splitUpdates = false) + +class DeltaBasedNestedRowIdMetadataCollisionUpdateAsDeleteAndInsertSuite + extends DeltaBasedNestedRowIdMetadataCollisionSuiteBase(splitUpdates = true)