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..b2004002ef5df --- /dev/null +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedNestedRowIdTableSuite.scala @@ -0,0 +1,171 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connector + +import org.apache.spark.sql.Row +import org.apache.spark.sql.connector.catalog.CatalogV2Util +import org.apache.spark.sql.types.{IntegerType, 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 + } + + // 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))), + 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)