peter-toth commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3735139305
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -226,4 +415,142 @@ 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)
+ }
+ }
+
+ /**
+ * For connectors that opt into narrow column updates AND represent UPDATE
as delete + insert,
+ * reject reassignment of any row-ID column as the REINSERT path has no
row-ID channel to
+ * reconstruct columns outside `requiredDataAttributes()`.
+ */
+ private def validateNoRowIdReassignment(
Review Comment:
**Finding 15.** This closes the repro I gave for finding 6 — `SET pk = pk +
10, salary = -1` now fails with `SPLIT_UPDATE_ROW_ID_REASSIGNMENT` — but the
guard is narrower than the loss it protects against, and the error text points
at why: it says a reassigned row ID leaves "no row-ID channel to reconstruct
columns outside `requiredDataAttributes()`". That channel exists only because
the fixture happens to put `pk` in `requiredDataAttributes()`
(`InMemoryRowLevelOperationTable.scala:312`, `(pk +: updatedCols)`). Nothing in
the interface requires it, and nothing checks it.
I measured this on `cef886e` with a split fixture whose
`requiredDataAttributes()` returns the assigned columns only, against the
existing `column-update-split` fixture as a control. Table `pk INT NOT NULL,
salary INT, dep STRING` partitioned by `dep`, rows `(1, 100, hr)` and `(2, 200,
software)`, then `UPDATE t SET salary = -1 WHERE dep = 'hr'`:
| `requiredDataAttributes()` | `updateSchema()` | result |
| --- | --- | --- |
| `[pk, salary]` (control) | `STRUCT<pk: INT NOT NULL, salary: INT NOT
NULL>` | `[1,-1,hr], [2,200,software]` |
| `[salary]` | `STRUCT<salary: INT NOT NULL>` | **`[2,200,software]`** |
The matched row is gone — not NULLed, deleted. No error from analysis and
none from the writer: the DELETE half of the split has the row ID it needs, the
REINSERT half gets a one-column row it cannot place, and nothing notices.
Every validation passes on the way there. `validateUpdatedColumnsSubset` is
happy because the only assigned key is `salary`, which is declared.
`validateNoOverlap` is happy. `validateNoRowIdReassignment` is happy because no
row-ID column is assigned. And the projection asymmetry does the rest:
`buildNarrowRelationWithAttrs(..., rowIdAttrs)` puts `pk` in the scan and
`buildNarrowDeletesAndInserts` puts it in the `Expand` output, but
`buildWriteDeltaProjections(rowDeltaPlan, rowAttrs, rowIdAttrs, metadataAttrs,
connectorDataAttrs)` projects the REINSERT payload down to
`connectorDataAttrs`. `dataAttrsResolved` then compares that payload against
`projectedDataAttrs`, which is the same `[salary]`, so resolution is satisfied
by construction.
Since the scan already carries the row ID, widening this validation is the
cheap fix:
```scala
// the REINSERT payload is projected down to requiredDataAttributes(),
so the row ID has to be
// declared there too or the reinserted row cannot be placed at all
val undeclaredRowIds = rowIdAttrs
.filterNot(AttributeSet(connectorDataAttrs).contains)
.map(_.name)
.distinct
```
reusing `SPLIT_UPDATE_ROW_ID_REASSIGNMENT` or adding a sibling condition.
The other direction — projecting `connectorDataAttrs ++ rowIdAttrs` for
REINSERT rows — would make `updateSchema()` differ by operation shape, which
`dataAttrsResolved` deliberately forbids ("same columns, same order"), so
validating looks like the right side to fix this on.
Worth a test next to the one you added: the same split fixture with `pk`
dropped from `requiredDataAttributes()`, asserting the analysis error. Happy to
hand over the fixture variant I used if that saves you a few minutes.
##########
project/MimaExcludes.scala:
##########
@@ -60,7 +60,10 @@ object MimaExcludes {
// [SPARK-57987] Add desc field to the SQL REST API Node case class
ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.status.api.v1.sql.Node.apply"),
ProblemFilters.exclude[DirectMissingMethodProblem]("org.apache.spark.status.api.v1.sql.Node.copy"),
-
ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.status.api.v1.sql.Node$")
+
ProblemFilters.exclude[MissingTypesProblem]("org.apache.spark.status.api.v1.sql.Node$"),
+ // [SPARK-56599] Add scan and write schema narrowing for column-level
UPDATEs in DSv2
Review Comment:
**Finding 18.** Wrong ticket — this PR is SPARK-58111; SPARK-56599 is a
different issue. These comments are how someone later works out why a filter
exists, so it's worth getting right.
```suggestion
// [SPARK-58111] Add scan and write schema narrowing for column-level
UPDATEs in DSv2
```
(Which block this filter belongs in depends on the 4.3-vs-4.4 question, so
I'm leaving that alone.)
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DeltaWriter.java:
##########
@@ -40,6 +40,11 @@ public interface DeltaWriter<T> extends DataWriter<T> {
/**
* Updates a row.
+ * <p>
+ * When the associated {@link RowLevelOperation} mixes in {@link
SupportsColumnUpdates}, the
Review Comment:
**Finding 16.** Finding 11's note landed on the interface as "Currently
honored only for UPDATE. DELETE and MERGE ignore this interface: those
operations receive full-width rows and `LogicalWriteInfo.updateSchema()` is
absent for them." This sentence keys the layout on something different —
whether the operation *mixes in* the interface — and for MERGE the two disagree.
A `SupportsColumnUpdates` operation is reachable from MERGE:
`RowLevelOperationBuilder.build()` only sees `info.command`, and your own
fixture returns the same `DeltaBasedColumnUpdateOperation` for every command
(`InMemoryRowLevelOperationTable.scala:154-160`). `RewriteMergeIntoTable` never
consults `SupportsColumnUpdates`, so `updateRowProjection` is absent,
`updateProj` falls back to `rowProjection`
(`WriteToDataSourceV2Exec.scala:891`), and `update(...)` receives a full-width
row — while this Javadoc tells the implementer to read it as narrow. A
connector that follows the doc mis-parses every MERGE update, and the failure
is a wrong value per column rather than an error.
Keying both sentences on `updateSchema()` instead of on the mix-in removes
the contradiction and stays true for every command:
```suggestion
* When {@link LogicalWriteInfo#updateSchema()} is present, the {@code
row} follows the narrow
* layout it declares rather than the full table schema from {@link
LogicalWriteInfo#schema()}.
* It is present only for UPDATE on an operation that mixes in {@link
SupportsColumnUpdates};
* otherwise {@code row} follows the full table schema.
```
Same wording applies to the `reinsert` copy at `:61`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -226,4 +415,142 @@ 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)
Review Comment:
**Finding 17.** `scanOnlyDataAttributes()` is the right answer to finding 9
and I like that it separates "must be in the scan" from "must reach the writer"
explicitly. The part I'd push back on is dropping the automatic partition-ref
inclusion in the same change: it turns something Spark used to guarantee into
something the connector has to remember, with no check and no clear error.
The code you removed kept partition refs precisely "so downstream rules
(V2ScanPartitioningAndOrdering, GroupBasedRowLevelOperationScanPlanning) can
resolve the table's partitioning expressions against the scan output". A
connector that partitions by `dep`, opts into column updates, and names `dep`
in neither method now gets a narrow scan without it and lands in exactly the
failure finding 9 was about — `Unable to resolve dep given [...]` out of
`DistributionAndOrderingUtils` when the write clusters on it — or, on the read
side, silently reports no partitioning because the refs don't resolve. The
evidence that it's easy to trip is in this PR:
`DeltaBasedColumnUpdateOperation` had to start returning `dep` from
`scanOnlyDataAttributes()` unconditionally
(`InMemoryRowLevelOperationTable.scala:316-319`) for the pre-existing narrowing
tests to keep passing, and two of them had to grow an `extra` column to still
have something to assert exclusion on.
Either check it, next to `validateNoOverlap`, so the connector gets told
what it forgot:
```scala
// Spark no longer adds partition refs implicitly, so a connector that
needs them for
// partitioning resolution or write clustering must declare them in one
of the two methods.
val declared = AttributeSet(connectorDataAttrs ++ scanOnlyDataAttrs)
val undeclaredPartitionAttrs = <partition refs resolved against relation>
.filterNot(declared.contains).map(_.name)
```
or keep the implicit inclusion for partition refs only. That is safe now in
a way it wasn't before: the write payload is projected down to
`connectorDataAttrs` explicitly by `buildWriteDeltaProjections(...,
updateRowAttrs)`, so a partition column present in the scan no longer reaches
the writer. `scanOnlyDataAttributes()` would still earn its keep for a
clustering key that isn't a partition column.
--
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]
