peter-toth commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3735141175


##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -226,4 +389,98 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
     val expandOutput = generateExpandOutput(attrs, outputs)
     Expand(outputs, expandOutput, matchedRowsPlan)
   }
+
+  /**
+   * Variant of `buildDeletesAndInserts` for the `SupportsColumnUpdates` 
narrow-scan path.
+   * This variant realigns the assignments to one value per surviving rowAttr 
padding unassigned
+   * rowAttrs with identity, so the reinsert output arity matches the delete 
output arity in
+   * the resulting Expand.
+   */
+  private def buildNarrowDeletesAndInserts(
+      matchedRowsPlan: LogicalPlan,
+      assignments: Seq[Assignment],
+      rowIdAttrs: Seq[Attribute]): Expand = {
+
+    val (metadataAttrs, rowAttrs) = matchedRowsPlan.output.partition { attr =>
+      MetadataAttribute.isValid(attr.metadata)
+    }
+    val assignmentMap = AttributeMap(assignments.collect {
+      case a @ Assignment(key: Attribute, _) => key -> a
+    })
+    val reinsertAssignments = rowAttrs.map { attr =>
+      assignmentMap.get(attr) match {
+        case Some(a) => a
+        case None => Assignment(attr, attr)
+      }
+    }
+    val deleteOutput = deltaDeleteOutput(rowAttrs, rowIdAttrs, metadataAttrs)
+    val insertOutput = deltaReinsertOutput(reinsertAssignments, metadataAttrs)
+    val outputs = Seq(deleteOutput, insertOutput)
+    val operationTypeAttr = AttributeReference(OPERATION_COLUMN, IntegerType, 
nullable = false)()
+    val attrs = operationTypeAttr +: matchedRowsPlan.output
+    val expandOutput = generateExpandOutput(attrs, outputs)
+    Expand(outputs, expandOutput, matchedRowsPlan)
+  }
+
+  /**
+   * Resolves the connector's `requiredDataAttributes()` if the operation opts 
into column
+   * updates. Returns `Nil` otherwise.
+   */
+  private def resolveConnectorDataAttrs(
+      relation: DataSourceV2Relation,
+      operation: RowLevelOperation): Seq[AttributeReference] = operation match 
{
+    case scu: SupportsColumnUpdates => resolveRequiredDataAttrs(relation, scu)
+    case _ => Nil
+  }
+
+  /**
+   * Computes the narrow set of data columns that must be present in the scan 
for a column-update
+   * write: connector-declared attrs, unioned with any table columns 
referenced by non-identity
+   * assignment RHS expressions, the operation condition, and the table's 
partition expressions.
+   * Partition-column refs are always kept so downstream rules 
(V2ScanPartitioningAndOrdering,
+   * GroupBasedRowLevelOperationScanPlanning) can resolve the table's 
partitioning expressions
+   * against the scan output.
+   */
+  private def computeNarrowReadAttrs(
+      relation: DataSourceV2Relation,
+      connectorDataAttrs: Seq[AttributeReference],
+      assignments: Seq[Assignment],
+      cond: Expression): Seq[AttributeReference] = {
+    val relationSet = relation.outputSet
+    val nonIdentityRhsRefs = assignments.iterator
+      .filterNot(a => a.key.isInstanceOf[Attribute] &&
+        isIdentityAssignment(a.key.asInstanceOf[Attribute], a.value))
+      .flatMap(_.value.references.toSeq)
+      .toSeq
+    val extraRefs = (cond.references.toSeq ++ nonIdentityRhsRefs)
+      .collect { case a: AttributeReference => a }
+      .filter(relationSet.contains)
+    val partitionRefNames = relation.table.partitioning().toImmutableArraySeq

Review Comment:
   **Finding 9.** Option 1 is what I hoped you'd take, and the implementation 
is right: `scanOnlyPassThroughValues` carries the column through the write 
query so `RequiresDistributionAndOrdering` can resolve it, while the payload 
stays projected to `connectorDataAttrs`, so the column never reaches the 
writer. `validateNoOverlap` closes the obvious way to get it wrong. I re-traced 
the MoR repro from my last round and it resolves now.
   
   The test is the part I'd single out — `column-update: scanOnlyDataAttributes 
resolves write-side clustering on a data column without leaking it into the 
write payload` asserts the *absence* from `updateSchema()` as well as the 
successful resolution, which is what makes it a real regression test rather 
than a smoke test. Making the fixture's `clusterColumnRef` overridable so the 
scan-only variant clusters on a data column instead of `_partition` was the 
right way to reach it.
   
   Two follow-ups from this change rather than objections to it: finding 17 on 
the removal of the implicit partition-ref inclusion that came with it, and 
finding 15 on the row-ID gap in the split path, which the same projection 
asymmetry causes.
   



-- 
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]

Reply via email to