aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3786270204
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -226,4 +421,186 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Variant of `buildDeletesAndInserts` for the `SupportsColumnUpdates`
narrow-scan path.
+ * This variant realigns the assignments to one value per surviving rowAttr
padding unassigned
+ * rowAttrs with identity, so the reinsert output arity matches the delete
output arity in
+ * the resulting Expand.
+ */
+ private def buildNarrowDeletesAndInserts(
+ matchedRowsPlan: LogicalPlan,
+ assignments: Seq[Assignment],
+ rowIdAttrs: Seq[Attribute]): Expand = {
+
+ val (metadataAttrs, rowAttrs) = matchedRowsPlan.output.partition { attr =>
+ MetadataAttribute.isValid(attr.metadata)
+ }
+ val assignmentMap = AttributeMap(assignments.collect {
+ case a @ Assignment(key: Attribute, _) => key -> a
+ })
+ val reinsertAssignments = rowAttrs.map { attr =>
+ assignmentMap.get(attr) match {
+ case Some(a) => a
+ case None => Assignment(attr, attr)
+ }
+ }
+ val deleteOutput = deltaDeleteOutput(rowAttrs, rowIdAttrs, metadataAttrs)
+ val insertOutput = deltaReinsertOutput(reinsertAssignments, metadataAttrs)
+ val outputs = Seq(deleteOutput, insertOutput)
+ val operationTypeAttr = AttributeReference(OPERATION_COLUMN, IntegerType,
nullable = false)()
+ val attrs = operationTypeAttr +: matchedRowsPlan.output
+ val expandOutput = generateExpandOutput(attrs, outputs)
+ Expand(outputs, expandOutput, matchedRowsPlan)
+ }
+
+ /**
+ * Resolves the connector's `requiredDataAttributes()` if the operation opts
into column
+ * updates. Returns `Nil` otherwise.
+ */
+ private def resolveConnectorDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation): Seq[AttributeReference] = operation match
{
+ case scu: SupportsColumnUpdates => resolveRequiredDataAttrs(relation, scu)
+ case _ => Nil
+ }
+
+ /**
+ * Resolves the connector's `scanOnlyDataAttributes()` if the operation opts
into column
+ * updates. Returns `Nil` otherwise.
+ */
+ private def resolveScanOnlyDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation): Seq[AttributeReference] = operation match
{
+ case scu: SupportsColumnUpdates =>
+ V2ExpressionUtils.resolveRefs[AttributeReference](
+ scu.scanOnlyDataAttributes.toImmutableArraySeq, relation)
+ case _ => Nil
+ }
+
+ /**
+ * Computes the narrow set of data columns that must be present in the scan
for a column-update
+ * write: connector-declared attrs (both `requiredDataAttributes()` and
+ * `scanOnlyDataAttributes()`), unioned with any table columns referenced by
non-identity
+ * assignment RHS expressions and the operation condition.
+ */
+ private def computeNarrowReadAttrs(
+ relation: DataSourceV2Relation,
+ connectorDataAttrs: Seq[AttributeReference],
+ scanOnlyDataAttrs: Seq[AttributeReference],
+ assignments: Seq[Assignment],
+ cond: Expression): Seq[AttributeReference] = {
+ val relationSet = relation.outputSet
+ val nonIdentityRhsRefs = assignments.iterator
+ .filterNot(a => a.key.isInstanceOf[Attribute] &&
+ isIdentityAssignment(a.key.asInstanceOf[Attribute], a.value))
+ .flatMap(_.value.references.toSeq)
+ .toSeq
+ val extraRefs = (cond.references.toSeq ++ nonIdentityRhsRefs)
+ .collect { case a: AttributeReference => a }
+ .filter(relationSet.contains)
+ dedupAttrs(connectorDataAttrs ++ scanOnlyDataAttrs ++ extraRefs)
+ }
+
+ /**
+ * Enforces that every column being assigned (non-identity) is present in
the connector-declared
+ * `requiredDataAttributes()`. Comparison is at root-column granularity
+ */
+ private def validateUpdatedColumnsSubset(
+ operation: RowLevelOperation,
+ assignments: Seq[Assignment],
+ connectorDataAttrs: Seq[AttributeReference]): Unit = {
+ val declaredIds = connectorDataAttrs.map(_.exprId).toSet
+ val missing = assignments.collect {
+ case Assignment(key: AttributeReference, value)
+ if !isIdentityAssignment(key, value) &&
!declaredIds.contains(key.exprId) =>
+ key.name
+ }.distinct
+ if (missing.nonEmpty) {
+ throw
QueryCompilationErrors.requiredDataAttributesMissingUpdatedColumnsError(
+ operation.getClass.getName, missing)
+ }
+ }
+
+ /**
+ * Enforces that `requiredDataAttributes()` and `scanOnlyDataAttributes()`
are disjoint.
+ */
+ private def validateNoOverlap(
+ operation: RowLevelOperation,
+ connectorDataAttrs: Seq[AttributeReference],
+ scanOnlyDataAttrs: Seq[AttributeReference]): Unit = {
+ val requiredIds = connectorDataAttrs.map(_.exprId).toSet
+ val overlapping = scanOnlyDataAttrs.collect {
+ case attr if requiredIds.contains(attr.exprId) => attr.name
+ }.distinct
+ if (overlapping.nonEmpty) {
+ throw
QueryCompilationErrors.requiredDataAttributesOverlapScanOnlyAttributesError(
+ operation.getClass.getName, overlapping)
+ }
+ }
+
+ /**
+ * Enforces that every partition-source column is declared in either
`requiredDataAttributes()`
+ * or `scanOnlyDataAttributes()`. Spark no longer adds partition refs to the
narrow scan
+ * implicitly, so a connector that needs them for partitioning resolution or
write-side
+ * clustering must declare them in one of the two methods.
+ */
+ private def validatePartitionAttrsDeclared(
Review Comment:
We can't be very restrictive on the Spark side and make assumptions about
what the connectors may want to do.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
