peter-toth commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3721001039
##########
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.** I hit the concrete failure this thread is about, so here it
is with numbers, plus my read on the three options.
A `SupportsDelta` + `SupportsColumnUpdates` connector declaring
`requiredDataAttributes() = [pk, id]` on a table partitioned by `dep`, whose
write requires `Distributions.clustered(dep)` and an ordering on `dep`:
```sql
UPDATE t SET id = -1 WHERE pk = 1
```
fails at `V2Writes` with
```
AnalysisException: Unable to resolve dep given
[__row_operation,id,pk,_partition,index,pk]
at
DistributionAndOrderingUtils$.prepareQuery(DistributionAndOrderingUtils.scala:47)
at V2Writes$$anonfun$apply$1.applyOrElse(V2Writes.scala:126)
```
The same connector on the CoW path resolves fine. The asymmetry is that
`buildNarrowReplaceDataUpdateProjection` maps over all of `plan.output`, so the
partition columns stay in the ReplaceData write query, while
`buildColumnUpdateProjection` emits only `assignedValues ++
connectorPassThroughValues ++ metadata ++ rowIds`, so they are gone from the
WriteDelta query. `computeNarrowReadAttrs`' comment says the partition refs are
kept "so downstream rules ... can resolve the table's partitioning expressions
against the scan output" -- that holds for the scan, but the delta write query
then drops them again, which is exactly the hole you describe.
Worth noting for whichever option you pick: `prepareQuery` resolves against
`query.output`, and `updateRowProjection` selects the payload out of that query
*by name*. So a column can sit in the write query without entering the write
payload -- no `updateSchema()` change needed. That means option 1 is
implementable by adding the scan-only attrs to the `Project` list in
`buildColumnUpdateProjection` and leaving `updateRowAttrs = connectorDataAttrs`
alone, which is what the CoW path effectively already does.
On the options:
- **Option 1 (`scanOnlyDataAttributes()`)** -- my preference. It is the only
one that names the actual distinction (needed for planning, not for the write
row), and it keeps `updateSchema()` meaning "the columns being written".
- **Option 2 (fold into `requiredDataAttributes()`)** --
`requiredDataAttributes()` *is* the write-row schema today: `updateSchema()` is
exactly that list, and `ReplaceData`/`WriteDelta.dataAttrsResolved` validates
the projection against it position for position. Folding scan-only columns in
would push them into the payload and make `updateSchema()` stop meaning that,
which also silently re-introduces the reconstruct-or-lose problem for those
columns.
- **Option 3 (widen `requiredMetadataAttributes()`)** -- those are validated
against `projectedMetadataAttrs` and delivered through the metadata row, so
data columns would arrive in the metadata slot. That is a larger contract
change than option 1 for a smaller gain.
Whichever way it goes, the two paths should end up agreeing -- right now CoW
carries the columns and delta does not, and nothing in the tests catches the
difference because every test connector clusters on the `_partition` metadata
column rather than a data column.
Small aside in the same method: `pk` appears twice in that projection
(`connectorPassThroughValues` and `rowIdValues` both emit it, hence the
duplicate in the output list above). It is harmless -- `findColOrdinal` takes
the first and both share an `exprId`, so resolution does not go ambiguous --
but the second one is dead and could be filtered out.
--
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]
