You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "anuragmantri (via GitHub)" <gi...@apache.org> on 2026/04/23 18:30:06 UTC
[PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
anuragmantri opened a new pull request, #55518:
URL: https://github.com/apache/spark/pull/55518
**What changes were proposed in this pull request?**
For SPIP: [SPARK-56599](https://issues.apache.org/jira/browse/SPARK-56599)
This PR adds three new default methods to the DSv2 connector API to enable scan and write-schema narrowing for column-level UPDATEs:
- `updatedColumns()` on RowLevelOperationInfo — Spark informs the connector which columns are being assigned (non-identity only) before the operation is
built.
- `requiredDataAttributes()` on RowLevelOperation — the connector declares the exact set of data columns it needs in the write schema, symmetric with
`requiredMetadataAttributes()`.
- `supportsColumnUpdates()` on RowLevelOperation — explicit opt-in for receiving a partial row instead of the full table row.
When a connector opts in, Spark removes identity assignments from the write plan's Project node, unblocking ColumnPruning to narrow the physical scan automatically (MOR path). For CoW, scan narrowing is done at analysis time via `buildRelationWithAttrs()` since GroupBasedRowLevelOperationScanPlanning reads DataSourceV2Relation.output before ColumnPruning fires.
All three methods have default implementations that preserve today's full-row behavior. No existing connector is affected.
**Why are the changes needed?**
Today, Spark's analyzer generates identity assignments for every column during [UPDATE alignment](https://github.com/apache/spark/blob/master/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/AssignmentUtils.scala#L62). These are used to build a Project that references [all columns ](https://github.com/apache/spark/blob/master/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala#L179), preventing [Optimizer](https://github.com/apache/spark/blob/master/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala#L1050) from narrowing the scan. The cost scales as O(table width) regardless of how many columns are being updated.
This is especially wasteful for columnar formats like Parquet/Iceberg and is a blocker for efficient column-level update implementations in connectors (see the [Efficient Column Updates Proposal](https://docs.google.com/document/d/1Bd7JVzgajA8-DozzeEE24mID_GLuz6iwj0g4TlcVJcs/edit?pli=1&tab=t.0) in Iceberg).
**Does this PR introduce any user-facing change?**
Yes. Three new default methods are added to the public DSv2 connector API:
- `RowLevelOperation.supportsColumnUpdates()`
- `RowLevelOperation.requiredDataAttributes()`
- `RowLevelOperationInfo.updatedColumns()`
All are opt-in with backward-compatible defaults. Existing connectors see no change.
**How was this patch tested?**
- 31 new tests in DeltaBasedColumnUpdateTableSuite covering scan narrowing, write-schema narrowing, data correctness, identity assignment filtering, updatedColumns behavior, and requiredDataAttributes across MOR (delta), CoW (ReplaceData), and delete-then-reinsert paths.
- 6 new updatedColumns tests in DeltaBasedUpdateTableSuiteBase.
**Was this patch authored or co-authored using generative AI tooling?**
Generated-by: Claude Sonnet 4.6
I used Claude Code to generate code and tests and manually reviewed the generated code.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3575345200
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -154,30 +193,114 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
cond: Expression): WriteDelta = {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
+ val supportsColumnUpdate = operation.supportsColumnUpdates()
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
- val rowAttrs = relation.output
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveRequiredDataAttrs(relation, operation)
+ } else Nil
+
val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ // Connector-declared attrs not being assigned are passed through so ColumnPruning
+ // keeps them in the scan and the connector receives their current values.
+ val assignedAttrs = if (supportsColumnUpdate) computeAssignedAttrs(assignments)
+ else relation.output
+ val connectorExtraAttrs: Seq[AttributeReference] = if (connectorDataAttrs.nonEmpty) {
+ val assignedAttrSet = AttributeSet(assignedAttrs)
+ connectorDataAttrs.filterNot(assignedAttrSet.contains)
+ } else Nil
+
// build a plan for updated records that match the condition
val matchedRowsPlan = Filter(cond, readRelation)
val rowDeltaPlan = if (operation.representUpdateAsDeleteAndInsert) {
buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else if (supportsColumnUpdate) {
+ buildColumnUpdateProjection(
+ matchedRowsPlan, assignments, rowIdAttrs, metadataAttrs, connectorExtraAttrs)
} else {
buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs)
}
+ val effectiveRowAttrs = if (supportsColumnUpdate && connectorDataAttrs.nonEmpty) {
+ connectorDataAttrs
+ } else if (supportsColumnUpdate) {
+ assignedAttrs
+ } else {
+ relation.output
+ }
+
// 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(
+ rowDeltaPlan, effectiveRowAttrs, rowIdAttrs, metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond)
}
+ /**
+ * Builds the WriteDelta projection for the column update path. The resulting Project
+ * references only the columns needed for the write, so ColumnPruning narrows the scan
+ * to match.
+ */
+ private def buildColumnUpdateProjection(
+ plan: LogicalPlan,
+ assignments: Seq[Assignment],
+ rowIdAttrs: Seq[Attribute],
+ metadataAttrs: Seq[Attribute],
+ connectorExtraAttrs: Seq[AttributeReference] = Nil): LogicalPlan = {
+
+ val assignedValues = assignments.collect {
+ case Assignment(key: Attribute, value) if !isIdentityAssignment(key, value) =>
+ Alias(value, key.name)()
+ }
+
+ val connectorExtraAttrSet = AttributeSet(connectorExtraAttrs)
+ val connectorPassThroughValues = plan.output.filter { a =>
+ connectorExtraAttrSet.contains(a) && !MetadataAttribute.isValid(a.metadata)
+ }
+
+ val metadataAttrSet = AttributeSet(metadataAttrs)
+ val metadataValues = plan.output.filter(metadataAttrSet.contains).map { attr =>
+ if (MetadataAttribute.isPreservedOnUpdate(attr)) {
+ attr
+ } else {
+ Alias(Literal(null, attr.dataType), attr.name)(explicitMetadata = Some(attr.metadata))
+ }
+ }
+
+ val rowIdAttrSet = AttributeSet(rowIdAttrs)
+ val rowIdValues = plan.output.filter(rowIdAttrSet.contains)
+
+ val originalRowIdValues = buildOriginalRowIdValues(rowIdAttrs, assignments)
+ val operationType = Alias(Literal(UPDATE_OPERATION), OPERATION_COLUMN)()
+
+ Project(
+ Seq(operationType) ++ assignedValues ++ connectorPassThroughValues ++
+ metadataValues ++ rowIdValues ++ originalRowIdValues,
+ plan)
+ }
+
+ // Returns the table attributes that are genuinely updated (non-identity) in this UPDATE.
+ private def computeAssignedAttrs(assignments: Seq[Assignment]): Seq[AttributeReference] = {
+ assignments.collect {
+ case Assignment(key: AttributeReference, value) if !isIdentityAssignment(key, value) => key
+ }
+ }
+
+ private def isIdentityAssignment(key: Attribute, value: Expression): Boolean = {
Review Comment:
Yes, I added a test for this.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-4383981429
> Could you resolve the conflicts, @anuragmantri ?
Thanks. I rebased and fixed the conflicts.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192409086
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -65,18 +72,15 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
assignments: Seq[Assignment],
cond: Expression): ReplaceData = {
- // resolve all required metadata attrs that may be used for grouping data on write
- val metadataAttrs = resolveRequiredMetadataAttrs(relation, operationTable.operation)
-
- // construct a read relation and include all required metadata columns
- val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs)
+ val (readRelation, rowAttrs) = buildCoWReadSetup(relation, operationTable, assignments, cond)
- // build a plan with updated and copied over records
- val query = buildReplaceDataUpdateProjection(readRelation, assignments, cond)
+ val updatedAndRemainingRowsPlan = buildReplaceDataUpdateProjection(
+ readRelation, assignments, cond)
- // build a plan to replace read groups in the table
val writeRelation = relation.copy(table = operationTable)
- val projections = buildReplaceDataProjections(query, relation.output, metadataAttrs)
+ val query = updatedAndRemainingRowsPlan
Review Comment:
This looks like duplications: Let's use one variable instead of mixing two variables, `updatedAndRemainingRowsPlan` and `query`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-6444705076
Gentle ping, @anuragmantri .
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub) via reviews" <re...@spark.apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4428900269
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/LogicalWriteInfo.java:
##########
@@ -47,6 +47,10 @@ public interface LogicalWriteInfo {
/**
* the schema of the input data from Spark to data source.
+ * <p>
Review Comment:
This makes sense. We already report this empty for DELETE operations that go via SupportsDelta.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub) via reviews" <re...@spark.apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4428971790
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,114 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, updated, copied, and reinserted
+ * rows follow it: group-based operations receive them through
+ * {@link DataWriter#writeColumnUpdate(Object, Object)} or
+ * {@link DataWriter#writeColumnUpdate(Object)}, and operations that mix in {@link SupportsDelta}
+ * through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert}. Rows the command inserts
+ * always follow {@link LogicalWriteInfo#schema()} and arrive through
+ * {@link DataWriter#write(Object)} and {@link DeltaWriter#insert}. When
+ * {@link LogicalWriteInfo#columnUpdateSchema()} is absent, every row follows
+ * {@link LogicalWriteInfo#schema()}, as for any other operation. Connectors must decide which
+ * layout to expect by whether {@link LogicalWriteInfo#columnUpdateSchema()} is present, not by
+ * the command. A builder that wants full-width rows for some command can check
+ * {@link RowLevelOperationInfo#command()} and build an operation without this mix-in. Currently
+ * Spark narrows only UPDATE.
+ * <p>
+ * The scan builder returned by {@link #newScanBuilder} should implement
+ * {@link org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns}, so that the scan
+ * does not read columns the command neither references nor declares. Otherwise the scan reads
+ * every column. For operations that do not mix in {@link SupportsDelta}, only top-level columns
+ * are pruned: a struct column the command references through a nested field is read in full.
+ *
+ * @since 4.4.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * When Spark narrows the write, the returned columns become
+ * {@link LogicalWriteInfo#columnUpdateSchema()}, in declared order. Spark matches each declared
+ * column with the session's resolver, and {@link LogicalWriteInfo#columnUpdateSchema()} uses the
+ * table's spelling of each column, which can differ in case from the declared reference.
+ * Implementations must include every column they want to receive: every column reported by
+ * {@link RowLevelOperationInfo#updatedColumns()}, plus any columns needed for row lookup or
+ * routing. Columns that are not declared are absent from the rows the connector receives, so
+ * the connector must preserve their values itself. Spark may call this method more than once
+ * for the same operation, and every call must return the same columns.
+ * <p>
+ * If the returned array is empty, Spark does not narrow the write and the operation receives
Review Comment:
Not sure this paragraph makes sense. If there are no updated columns, what's the point of executing it?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3423730403
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java:
##########
@@ -37,4 +38,20 @@ public interface RowLevelOperationInfo {
* Returns the row-level SQL command (e.g. DELETE, UPDATE, MERGE).
*/
Command command();
+
+ /**
+ * Returns the columns being updated in an UPDATE statement, as non-identity assignments.
+ *
+ * <p>For DELETE and MERGE, returns an empty array.
+ *
+ * <p>Connectors can use this to decide what {@link RowLevelOperation#requiredDataAttributes()}
+ * to declare. For instance, a connector that needs its primary key for row lookup can check
+ * whether pk is already in the updated columns list and, if not, add it to
+ * requiredDataAttributes().
+ *
+ * @since 4.3.0
+ */
+ default NamedReference[] updatedColumns() {
+ return new NamedReference[0];
Review Comment:
Do we handle nested column update correctly? The doc must be clear and implementation consistent.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3423871912
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -50,20 +51,36 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
protected def buildOperationTable(
table: SupportsRowLevelOperations,
command: Command,
- options: CaseInsensitiveStringMap): RowLevelOperationTable = {
- val info = RowLevelOperationInfoImpl(command, options)
+ options: CaseInsensitiveStringMap,
+ updatedColumns: Seq[NamedReference] = Nil): RowLevelOperationTable = {
+ val info = RowLevelOperationInfoImpl(command, options, updatedColumns)
val operation = table.newRowLevelOperationBuilder(info).build()
RowLevelOperationTable(table, operation)
}
+ /**
+ * Builds a DataSourceV2Relation for a row-level operation, optionally narrowing its output.
+ *
+ * When dataAttrs is non-empty, the relation output is narrowed to include only columns
+ * required for a column-update write. When dataAttrs is empty, the full relation.output is
+ * preserved.
+ */
protected def buildRelationWithAttrs(
Review Comment:
Not a big fan of this change, it makes the method fragile. I would consider adding another one instead.
Something like `buildNarrowRelationWithAttrs()` or something similar.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192395270
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -154,30 +206,145 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
cond: Expression): WriteDelta = {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
+ // Column-update support applies to the standard delta path and the delete+reinsert path.
+ // When representUpdateAsDeleteAndInsert is true, the REINSERT leg of the Expand already
+ // uses only assigned values, so the narrow effectiveRowAttrs applies correctly.
+ val supportsColumnUpdate = operation.supportsColumnUpdates()
// resolve all needed attrs (e.g. row ID and any required metadata attrs)
- val rowAttrs = relation.output
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
+ // Connector-declared data attrs used to determine pass-through columns in the write plan.
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveRequiredDataAttrs(relation, operation)
+ } else Nil
+
+ // MOR uses a full-schema scan; ColumnPruning narrows it via Project references.
val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ // Connector-required attrs that are NOT being assigned are added as pass-throughs in the
+ // plan so that ColumnPruning keeps them in the physical scan AND the connector receives
+ // their current values via DeltaWriter.update's row argument.
+ val assignedAttrs = if (supportsColumnUpdate) computeAssignedAttrs(assignments)
+ else relation.output
+ val connectorExtraAttrs: Seq[AttributeReference] = if (connectorDataAttrs.nonEmpty) {
+ val assignedAttrSet = AttributeSet(assignedAttrs)
+ connectorDataAttrs.filterNot(assignedAttrSet.contains)
+ } else Nil
+
// build a plan for updated records that match the condition
val matchedRowsPlan = Filter(cond, readRelation)
val rowDeltaPlan = if (operation.representUpdateAsDeleteAndInsert) {
buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else if (supportsColumnUpdate) {
+ buildColumnUpdateProjection(
+ matchedRowsPlan, assignments, rowIdAttrs, metadataAttrs, connectorExtraAttrs)
} else {
buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs)
}
+ // Effective row write schema:
+ // - Narrow path (connectorDataAttrs declared): exactly connector-declared cols in declared
+ // order. The connector must declare ALL columns it wants to receive (including updated
+ // ones). This mirrors the metadata pattern and enables strict areCompatible validation.
+ // - Heuristic path (connectorDataAttrs empty): only the assigned (changed) columns.
+ // - Full path (no column-update support): full table output.
+ val effectiveRowAttrs = if (supportsColumnUpdate && connectorDataAttrs.nonEmpty) {
+ connectorDataAttrs
+ } else if (supportsColumnUpdate) {
+ assignedAttrs
+ } else {
+ relation.output
+ }
+
// 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(
+ rowDeltaPlan, effectiveRowAttrs, rowIdAttrs, metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond)
}
+ // Builds the row delta projection for the column update path.
+ //
+ // The resulting Project references only:
+ // - assigned column values (new values being written)
+ // - connector pass-through values (connector declared but not assigned)
+ // - metadata columns (nulled or preserved)
+ // - row ID columns (for delta identification)
+ // - original row ID values (only when a row ID column is being reassigned)
+ //
+ // ColumnPruning observes exactly these references and narrows the physical scan accordingly.
+ // Connectors that need additional columns in the scan (e.g., partition columns for
+ // distribution) should declare them in requiredDataAttributes().
Review Comment:
IIUC, for the correctness, we need to throw `AnalysisException` if `requiredDataAttributes` is invalid.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192416031
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -106,38 +104,92 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val remainingRowsPlan = addOperationColumn(COPY_OPERATION,
Filter(remainingRowFilter, readRelation))
- // the new state is a union of updated and copied over records
- val query = Union(updatedRowsPlan, remainingRowsPlan)
+ val updatedAndRemainingRowsPlan = Union(updatedRowsPlan, remainingRowsPlan)
- // build a plan to replace read groups in the table
val writeRelation = relation.copy(table = operationTable)
- val projections = buildReplaceDataProjections(query, relation.output, metadataAttrs)
+ val query = updatedAndRemainingRowsPlan
+ val metadataAttrs = resolveRequiredMetadataAttrs(relation, operationTable.operation)
+ val projections = buildReplaceDataProjections(query, rowAttrs, metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
ReplaceData(writeRelation, cond, query, relation, projections, groupFilterCond)
}
+ // Common read-relation setup shared by both CoW plan builders.
+ //
+ // When the connector supports column updates and declares required data attributes,
+ // the read relation is narrowed at analysis time so that
+ // GroupBasedRowLevelOperationScanPlanning uses only the needed columns for the scan.
+ // Otherwise the full relation output is used.
+ private def buildCoWReadSetup(
+ relation: DataSourceV2Relation,
+ operationTable: RowLevelOperationTable,
+ assignments: Seq[Assignment],
+ cond: Expression): (DataSourceV2Relation, Seq[Attribute]) = {
+
+ val operation = operationTable.operation
+ val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
+ val connectorDataAttrs = resolveRequiredDataAttrs(relation, operation)
+ val isNarrow = operation.supportsColumnUpdates() && connectorDataAttrs.nonEmpty
+
+ // CoW scan narrowing must be done manually at analysis time.
+ // GroupBasedRowLevelOperationScanPlanning (an optimizer rule that fires after analysis)
+ // always reads relation.output directly when building the physical scan -- it does not
+ // observe Project nodes above the relation, so optimizer-driven column pruning has no
+ // effect on CoW scans. We narrow DataSourceV2Relation.output here so that rule picks
+ // up the narrow set.
+ val readRelation = if (isNarrow) {
+ val allRequired = (connectorDataAttrs ++ computeAssignedAttrs(assignments)).distinct
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs, dataAttrs = allRequired,
+ cond = cond)
+ } else {
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs)
+ }
+
+ // CoW write schema (two paths only, no heuristic for CoW):
+ // - Narrow path (connectorDataAttrs declared): exactly connector-declared cols in declared
+ // order. The connector must declare ALL columns it wants to receive.
+ // - Full path (connectorDataAttrs empty OR supportsColumnUpdates=false): full table output.
+ // Unlike MOR, CoW does not have a heuristic assigned-only path because
+ // GroupBasedRowLevelOperationScanPlanning needs explicit column declarations to narrow.
+ val rowAttrs: Seq[Attribute] = if (isNarrow) connectorDataAttrs else relation.output
+
+ (readRelation, rowAttrs)
Review Comment:
Please return `metadataAttrs` too to avoid the following recomputation in the caller-side.
> val metadataAttrs = resolveRequiredMetadataAttrs(relation, operationTable.operation)
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3633962122
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -210,11 +247,25 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
protected def buildReplaceDataProjections(
plan: LogicalPlan,
rowAttrs: Seq[Attribute],
- metadataAttrs: Seq[Attribute]): ReplaceDataProjections = {
+ metadataAttrs: Seq[Attribute],
+ updateRowAttrs: Seq[Attribute] = Nil): ReplaceDataProjections = {
val outputs = extractOutputs(plan)
- val outputsWithRow = filterOutputs(outputs, OPERATIONS_WITH_ROW)
- val rowProjection = newLazyProjection(plan, outputsWithRow, rowAttrs)
+ val rowProjection = if (updateRowAttrs.nonEmpty) {
+ val outputsForInsert = filterOutputs(outputs,
+ OPERATIONS_WITH_ROW -- Set(UPDATE_OPERATION, COPY_OPERATION))
+ newLazyProjection(plan, outputsForInsert, rowAttrs)
Review Comment:
+1 for the above comment.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3635902078
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala:
##########
@@ -139,17 +138,13 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla
tableAttrs: Seq[Attribute],
scanAttrs: Seq[Attribute]): AttributeMap[Attribute] = {
- val attrMapping = tableAttrs.map { tableAttr =>
+ // The scan may be narrowed to exclude columns not needed by the connector.
+ // Attributes absent from the scan are skipped here; the caller must ensure
+ // that any attribute referenced in the condition is present in the scan.
+ val attrMapping = tableAttrs.flatMap { tableAttr =>
scanAttrs
.find(scanAttr => conf.resolver(scanAttr.name, tableAttr.name))
.map(scanAttr => tableAttr -> scanAttr)
- .getOrElse {
- throw new AnalysisException(
- errorClass = "_LEGACY_ERROR_TEMP_3075",
- messageParameters = Map(
- "tableAttr" -> tableAttr.toString,
- "scanAttrs" -> scanAttrs.mkString(",")))
- }
}
Review Comment:
Makes sense. Similar to other changes for column updates paths, I created a conditional method `buildNarrowTableToScanAttrMap()` which is called only during column updates and throws when any of the condition references are missing. My rationale is that the runtime filtering applies to only the filters so it is sufficient remap the filters only. Let me know if this understanding is incorrect.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,47 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ *
+ * @since 4.3.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * The returned columns become the schema of updated, copied, and reinserted rows, in declared
+ * order. Implementations must include every column they want to receive (typically the columns
+ * reported by {@link RowLevelOperationInfo#updatedColumns()} plus any columns needed for row
+ * lookup or routing, e.g. a primary key).
+ * <p>
+ * If any of the columns from {@link RowLevelOperationInfo#updatedColumns()}) is
Review Comment:
Done.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java:
##########
@@ -37,4 +38,13 @@ public interface RowLevelOperationInfo {
* Returns the row-level SQL command (e.g. DELETE, UPDATE, MERGE).
*/
Command command();
+
+ /**
+ * Returns the columns being updated by this operation.
Review Comment:
Done.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,540 @@
+/*
+ * 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, TableInfo}
+import org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity, reference}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+
+/**
+ * Tests for UPDATE statements targeting connectors that return true from
+ * [[org.apache.spark.sql.connector.write.RowLevelOperation#supportsColumnUpdates]].
+ *
+ * When a connector supports column updates, Spark narrows the row projection
+ * (LogicalWriteInfo.schema()) to contain only the assigned/changed columns rather than
+ * the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends RowLevelOperationSuiteBase {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ test("column-update: rowSchema contains only the single assigned column") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ // info.schema() is empty: UPDATE-only column-update writes have no INSERT-shaped rows,
+ // so nothing flows through the writer's write() path. The narrow update-row layout is
+ // carried by info.updateSchema().
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: rowSchema contains multiple assigned columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("id", IntegerType, nullable = false),
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: rowSchema is empty for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ // All assignments are identity, so updatedColumns is empty -- the connector still declares
+ // pk for row lookup, so the narrow update schema is just [pk].
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Array(PK_FIELD))))
+ }
+
+ test("column-update: row filter condition is orthogonal to column narrowing") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = 'engineering' WHERE pk IN (1, 3)")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: update all rows (no WHERE clause)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("salary", IntegerType, nullable = true)
+ ))))
+ }
+
+ test("column-update: rowSchema excludes identity assignments in a mixed UPDATE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: cross-column assignment is not treated as identity") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = dep, id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: nested struct field update narrows to the root struct column") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("s"),
+ s"expected [s] in updatedColumns (root struct) but got: $updatedNames")
+
+ // info.updateSchema() carries the narrow row layout; info.schema() is the full table.
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("s"),
+ s"s must be in update schema: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("dep"),
+ s"dep must not be in update schema: $updateSchema")
+ }
+
+ test("column-update: nested field identity update reports root struct as updated") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = s.c1 WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("s"),
+ s"nested identity is reported as an update at root-column granularity: $updatedNames")
+
+ // Data correctness: the struct is rewritten but with equal values, so rows are unchanged.
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, Row(1, 2), "hr") :: Nil)
+ }
+
+ test("column-update: updatedColumns contains non-identity assigned columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'eng' WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("id", "dep"),
+ s"expected [id, dep] in updatedColumns but got: $updatedNames")
+ }
+
+ test("column-update: updatedColumns excludes identity assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = dep WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("id"),
+ s"expected only [id] in updatedColumns but got: $updatedNames")
+ }
+
+ test("column-update: updatedColumns is empty for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ assert(table.lastUpdatedColumns.isEmpty,
+ s"expected empty updatedColumns but got: ${table.lastUpdatedColumns.mkString(", ")}")
+ }
+
+ test("column-update: updatedColumns is empty for DELETE (Javadoc contract)") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"DELETE FROM $tableNameAsString WHERE dep = 'hr'")
+
+ assert(table.lastUpdatedColumns.isEmpty,
+ s"DELETE must pass empty updatedColumns but got: ${table.lastUpdatedColumns.mkString(", ")}")
+ }
+
+ test("column-update: data correctness -- single column update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ // 1 row matched WHERE pk = 1 -> 1 update. MoR emits no COPY rows.
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0, deltaUpdate = true)
+ }
+
+ test("column-update: data correctness -- update all rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 200, "hr") :: Row(2, 400, "software") :: Row(3, 600, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 3, numCopiedRows = 0, deltaUpdate = true)
+ }
+
+ test("column-update: data correctness -- mixed identity and real assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "engineering") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0, deltaUpdate = true)
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-req-attrs", reqAttrs)
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update: requiredDataAttributes - data correctness") {
+ createAndInitTableWithReqAttrs("dep,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: empty requiredDataAttributes throws AnalysisException") {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-empty-req-attrs", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(
+ StructType.fromDDL("pk INT NOT NULL, id INT, dep STRING"))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }""".stripMargin)
+
+ // A connector that mixes in SupportsColumnUpdates but returns an empty
+ // requiredDataAttributes() violates the mix-in contract -- analysis must reject the
+ // operation rather than silently falling through to the wide write path.
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ }
+ assert(ex.getCondition == "EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ s"expected EMPTY_REQUIRED_DATA_ATTRIBUTES but got: ${ex.getCondition}")
+ }
+
+ test("column-update: requiredDataAttributes throws AnalysisException for invalid column") {
+ createAndInitTableWithReqAttrs("nonexistent_col", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ }
+ assert(ex.getMessage.contains("nonexistent_col"),
+ s"Expected error about unresolvable column but got: ${ex.getMessage}")
+ }
+
+ private def createAndInitTableFromInfo(schemaString: String, jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-from-info", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update from-info: write schema is updatedColumns + pk pass-through") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("salary"),
+ s"salary must be in update schema: $updateSchema")
+ assert(updateSchema.fieldNames.contains("pk"),
+ s"pk must be in update schema: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("id"),
+ s"id must not be in update schema: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("dep"),
+ s"dep must not be in update schema: $updateSchema")
+ }
+
+ test("column-update from-info: pk already in updatedColumns is not duplicated") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ val pkCount = updateSchema.fieldNames.count(_ == "pk")
+ assert(pkCount == 1, s"pk must appear exactly once in update schema: $updateSchema")
+ }
+
+ test("column-update from-info: data correctness") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "id": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ }
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-split", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update split: write schema is narrow (assigned + pk pass-through)") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ // For representUpdateAsDeleteAndInsert connectors, reinsert-tagged rows still route through
+ // writeUpdate() (info.updateSchema() = narrow), not write(). info.schema() stays empty.
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update split: data correctness") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, -1, "hr") :: Nil)
+ }
+
+ // ---------------------------------------------------------------------------
+ // Scan-narrowing tests: verify that when the connector opts into column updates,
+ // the physical scan reads only the columns actually needed by the operation.
+ //
+ // Assertions use `checkLastScanExcludes` (proves narrowing) and `checkLastScanIncludes`
+ // (proves transparent widening) rather than exact-set matches -- ColumnPruning is
+ // free to further tighten the scan beyond our analysis-time narrow set (e.g. when an
+ // assignment RHS is a literal, the assigned column doesn't need to be read).
+ // ---------------------------------------------------------------------------
+
+ test("column-update: scan excludes columns outside connector-declared + cond + RHS refs") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ // requiredDataAttributes = [pk, id]; cond refs [pk] only; RHS is a literal.
+ // `dep` is neither declared nor referenced -- it must not appear in the scan.
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastScanExcludes("dep")
+ }
+
+ test("column-update: scan transparently widens for cond-referenced columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ // requiredDataAttributes = [pk, id]; cond refs [dep].
+ // `dep` must appear even though the connector didn't declare it.
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE dep = 'hr'")
+
+ checkLastScanIncludes("dep")
+ }
+
+ test("column-update: scan transparently widens for assignment RHS references") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin)
+
+ // requiredDataAttributes = [pk, salary]; RHS `salary + bonus` references `bonus`.
+ // `bonus` must appear in the scan even though it's not declared as required.
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE pk = 1")
+
+ checkLastScanIncludes("bonus", "salary")
+ checkLastScanExcludes("dep")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 110, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update: analysis fails when assignment key is outside requiredDataAttributes") {
+ // Connector declares only [pk] but the user assigns to `id`. We enforce
+ // updatedColumns ⊆ requiredDataAttributes at analysis time (root-column granularity).
Review Comment:
Yes, this was auto generated. I rewrote it as per your suggestion.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -210,11 +247,25 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
protected def buildReplaceDataProjections(
plan: LogicalPlan,
rowAttrs: Seq[Attribute],
- metadataAttrs: Seq[Attribute]): ReplaceDataProjections = {
+ metadataAttrs: Seq[Attribute],
+ updateRowAttrs: Seq[Attribute] = Nil): ReplaceDataProjections = {
val outputs = extractOutputs(plan)
- val outputsWithRow = filterOutputs(outputs, OPERATIONS_WITH_ROW)
- val rowProjection = newLazyProjection(plan, outputsWithRow, rowAttrs)
+ val rowProjection = if (updateRowAttrs.nonEmpty) {
+ val outputsForInsert = filterOutputs(outputs,
+ OPERATIONS_WITH_ROW -- Set(UPDATE_OPERATION, COPY_OPERATION))
+ newLazyProjection(plan, outputsForInsert, rowAttrs)
Review Comment:
Thanks for catching this. I mirrored the guard in `buildWriteDeltaProjections`
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -83,11 +84,20 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
table: RowLevelOperationTable,
metadataAttrs: Seq[AttributeReference],
rowIdAttrs: Seq[AttributeReference] = Nil): DataSourceV2Relation = {
-
Review Comment:
Done
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -41,7 +44,11 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
EliminateSubqueryAliases(aliasedTable) match {
case r @ ExtractV2Table(tbl: SupportsRowLevelOperations) =>
checkNoGeneratedColumns(r, UPDATE)
- val table = buildOperationTable(tbl, UPDATE, r.options)
+ val updatedCols = assignments.collect {
+ case Assignment(key: AttributeReference, value) if !isIdentityAssignment(key, value) =>
+ FieldReference(key.name)
Review Comment:
Thanks for the pointer. Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -223,19 +274,39 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
None
}
- ReplaceDataProjections(rowProjection, metadataProjection)
+ ReplaceDataProjections(rowProjection, updateRowProjection, metadataProjection)
}
protected def buildWriteDeltaProjections(
plan: LogicalPlan,
rowAttrs: Seq[Attribute],
rowIdAttrs: Seq[Attribute],
- metadataAttrs: Seq[Attribute]): WriteDeltaProjections = {
+ metadataAttrs: Seq[Attribute],
+ updateRowAttrs: Seq[Attribute] = Nil): WriteDeltaProjections = {
val outputs = extractOutputs(plan)
- val rowProjection = if (rowAttrs.nonEmpty) {
+ // Always produce Some(rowProjection) even for empty rowAttrs (identity-only column updates).
+ // When updateRowAttrs is non-empty, the row projection covers INSERT-shaped rows only and a
+ // separate updateRowProjection handles UPDATE/COPY/REINSERT-shaped rows.
+ val rowProjection = if (updateRowAttrs.nonEmpty) {
+ val outputsForInsert = filterOutputs(outputs,
+ OPERATIONS_WITH_ROW -- Set(UPDATE_OPERATION, COPY_OPERATION, REINSERT_OPERATION))
+ if (outputsForInsert.isEmpty) {
+ Some(ProjectingInternalRow(StructType(Nil), Nil))
+ } else {
+ Some(newLazyProjection(plan, outputsForInsert, rowAttrs))
+ }
+ } else if (rowAttrs.nonEmpty) {
val outputsWithRow = filterOutputs(outputs, OPERATIONS_WITH_ROW)
Some(newLazyProjection(plan, outputsWithRow, rowAttrs))
+ } else {
+ Some(ProjectingInternalRow(StructType(Nil), Nil))
Review Comment:
Thanks for the explanation. I moved it back to `None` and updated the comment.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -244,6 +281,327 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends only the
+ // declared + assigned columns in the row projection instead of the full row schema. The base
+ // class composes its required-attrs set as `pk` (the row-lookup key) plus whatever columns
+ // Spark reports as being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil)
+ extends DeltaBasedOperation(command, CaseInsensitiveStringMap.empty())
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = {
+ val pk: NamedReference = FieldReference("pk")
+ val updatedNames = updatedCols.map(_.describe()).toSet
+ if (updatedNames.contains("pk")) updatedCols.toArray
+ else (pk +: updatedCols).toArray
+ }
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.updateSchema().isPresent) {
+ capturedInfo.updateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(Array(PARTITION_COLUMN_REF))
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ Array[SortOrder](
+ LogicalExpressions.sort(
+ PARTITION_COLUMN_REF,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering())
+ )
+ }
+
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ // Use the narrow update schema for UPDATE/COPY/REINSERT rows when present.
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For column-update writes, rows contain only the assigned columns
+ // (narrow schema from LogicalWriteInfo). We expand each row to the full table
+ // schema by overlaying write-schema columns on the base row found by pk.
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+
+ val mergedData = newData.map { buf =>
+ val merged = new BufferedRows(buf.key, schema)
+ val updateOpName = UTF8String.fromString(Update.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == updateOpName) {
+ val pk = logRow.getInt(1)
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val baseRow = dataMap.values.iterator.flatten
+ .flatMap(_.rows)
+ .find(r => r.getInt(schema.fieldIndex("pk")) == pk)
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ fullRow.update(i, base.get(i, schema.fields(i).dataType))
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ writeFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ merged.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ merged.rows.append(insertRow.copy())
+ }
+ }
+ merged
+ }
+
+ withDeletes(newData)
+ withData(mergedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ }
+ }
+ }
+ }
+
+ class DeltaBasedColumnUpdateOperationWithReqAttrs(command: Command, reqCols: Array[String])
+ extends DeltaBasedColumnUpdateOperation(command) {
+ override def requiredDataAttributes(): Array[NamedReference] = reqCols.map(FieldReference(_))
+ }
+
+ class DeltaBasedColumnUpdateOperationFromInfo(
+ command: Command,
+ updatedCols: Seq[NamedReference])
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols) {
+ }
+
+ class DeltaBasedColumnUpdateSplitOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil)
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols) {
+ override def representUpdateAsDeleteAndInsert(): Boolean = true
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.updateSchema().isPresent) {
+ capturedInfo.updateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+ override def requiredDistribution(): Distribution =
+ Distributions.clustered(Array(PARTITION_COLUMN_REF))
+ override def requiredOrdering(): Array[SortOrder] = Array[SortOrder](
+ LogicalExpressions.sort(
+ PARTITION_COLUMN_REF,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()))
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For delete+reinsert with narrow writes, the REINSERT row has only the
+ // connector-declared columns (requiredDataAttributes order).
+ // pk is the first field in the write schema (declared before updatedCols).
+ // Reconstruct the full row by overlaying the narrow row onto the original.
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+ val reinsertOpName = UTF8String.fromString(Reinsert.toString)
+ val pkIdx = writeFieldIdx("pk")
+
+ val expandedData = newData.map { buf =>
+ val expanded = new BufferedRows(buf.key, schema)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == reinsertOpName) {
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val pk = narrowRow.getInt(pkIdx)
+ val baseRow = dataMap.values.iterator.flatten
+ .flatMap(_.rows)
+ .find(r => r.getInt(schema.fieldIndex("pk")) == pk)
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ fullRow.update(i, base.get(i, schema.fields(i).dataType))
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ writeFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ expanded.rows.append(fullRow)
+ }
+ }
+ expanded
+ }
+
+ withDeletes(newData)
+ withData(expandedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ }
+ }
+ }
+ }
+
+ class PartitionBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil)
+ extends RowLevelOperation with SupportsColumnUpdates {
+ var configuredScan: InMemoryBatchScan = _
+
+ override def command(): Command = command
+
+ override def requiredDataAttributes(): Array[NamedReference] = {
+ val base = Seq(FieldReference("pk"), FieldReference("dep"))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ override def requiredMetadataAttributes(): Array[NamedReference] =
+ Array(PARTITION_COLUMN_REF, INDEX_COLUMN_REF)
+
+ override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = {
+ new InMemoryScanBuilder(schema, options) {
+ override def build(): Scan = {
+ val scan = super.build()
+ InMemoryRowLevelOperationTable.recordLastScanSchema(scan.readSchema())
+ configuredScan = scan.asInstanceOf[InMemoryBatchScan]
+ scan
+ }
+ }
+ }
+
+ override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = {
+ lastWriteInfo = info
+ new WriteBuilder {
+ override def build(): Write = new Write with RequiresDistributionAndOrdering {
+ override def requiredDistribution: Distribution =
+ Distributions.clustered(Array(PARTITION_COLUMN_REF))
+
+ override def requiredOrdering: Array[SortOrder] = Array[SortOrder](
+ LogicalExpressions.sort(
+ PARTITION_COLUMN_REF,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()))
+
+ override def toBatch: BatchWrite = {
+ val narrowSchema = if (info.updateSchema().isPresent) {
+ info.updateSchema().get()
+ } else {
+ info.schema()
+ }
+ PartitionBasedNarrowReplaceData(configuredScan, narrowSchema, info.schema())
Review Comment:
Done
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3786136982
##########
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:
Why do we need this? Like the comment says, if the connector needs, it adds them. Why should Spark be concerned about it?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3779998895
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +84,51 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated, copied, or reinserted record with metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must override this
+ * method when mixing in {@link SupportsColumnUpdates}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but does not override this method.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T metadata, T record) throws IOException {
Review Comment:
I wonder whether even something like `writeColumnUpdate` would be more descriptive.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3575251797
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -106,38 +105,78 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val remainingRowsPlan = addOperationColumn(COPY_OPERATION,
Filter(remainingRowFilter, readRelation))
- // the new state is a union of updated and copied over records
- val query = Union(updatedRowsPlan, remainingRowsPlan)
+ val updatedAndRemainingRowsPlan = Union(updatedRowsPlan, remainingRowsPlan)
- // build a plan to replace read groups in the table
val writeRelation = relation.copy(table = operationTable)
Review Comment:
I thought about narrowing the write relation here. I preferred not to because it works for UPDATE but when we implement MERGE later, this would not work and we would have to have the full write relation because INSERT rows will need it.
Also, if we narrow it here, the writeSchema() and updateSchema() would look the same making updateSchema() redundant for UPDATEs. I would like to have writeSchema() always represent full width and updateSchema() be the narrow width to be consistent.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-5210103684
@peter-toth, this is ready for another review
- Finding 6 - 14 (except 9) https://github.com/apache/spark/pull/55518/commits/2b07cec6c49b27c28a69e448aeb4c6dc450d51f6
- Finding 9 [cef886e](https://github.com/apache/spark/pull/55518/commits/cef886e16b7fb400b90841caa8975355f4c7ff42)
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192341324
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java:
##########
@@ -105,4 +105,45 @@ default String description() {
default NamedReference[] requiredMetadataAttributes() {
return new NamedReference[0];
}
+
+
Review Comment:
nit. Remove redundant empty line.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192365368
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala:
##########
@@ -378,7 +386,35 @@ case class ReplaceData(
// validates row projection output is compatible with table attributes
private def rowAttrsResolved: Boolean = {
val inRowAttrs = DataTypeUtils.toAttributes(projections.rowProjection.schema)
- table.skipSchemaResolution || areCompatible(inRowAttrs, table.output)
+ table.skipSchemaResolution ||
+ areCompatible(inRowAttrs, table.output) ||
+ dataAttrsResolved(inRowAttrs)
+ }
+
+ // Validates the narrow-write-schema row projection output.
Review Comment:
For function description, please follow the community style like the other code path.
```
/**
* ...
*/
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3633944202
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -83,11 +84,20 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
table: RowLevelOperationTable,
metadataAttrs: Seq[AttributeReference],
rowIdAttrs: Seq[AttributeReference] = Nil): DataSourceV2Relation = {
-
Review Comment:
Please recover this empty line removal. This removal seems that a leftover when you recover the original method. @aokolnychyi 's request was to keep the original method unchanged completely.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-5039922301
CC: @dongjoon-hyun and @peter-toth, since you last reviewed, I changed the following
- SupportsColumnUpdate is now a mixin for RowLevelOperations
- LogicalWriteInfo.updateSchema(): Optional — narrow row schema for writeUpdate()
- DataWriter.writeUpdate(record) and writeUpdate(metadata, record) for writing column update records with narrow schema.
Please take another look. It will be great to have this in by feature freeze for Spark 4.3.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3784965468
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +84,51 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated, copied, or reinserted record with metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must override this
+ * method when mixing in {@link SupportsColumnUpdates}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but does not override this method.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T metadata, T record) throws IOException {
Review Comment:
Need to think more, no need to update.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3763603723
##########
common/utils/src/main/resources/error/error-conditions.json:
##########
@@ -2461,6 +2467,12 @@
],
"sqlState" : "KD009"
},
+ "EMPTY_REQUIRED_DATA_ATTRIBUTES" : {
Review Comment:
Good idea, I have prefixed all the error conditions with `COLUMN_UPDATE_*`.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java:
##########
@@ -37,4 +38,15 @@ public interface RowLevelOperationInfo {
* Returns the row-level SQL command (e.g. DELETE, UPDATE, MERGE).
*/
Command command();
+
+ /**
+ * Returns the columns being updated by this operation. Currently populated only for UPDATE;
+ * DELETE and MERGE report an empty array.
+ * <p>
+ * Nested struct field updates are reported at root-column granularity
Review Comment:
This is from `AssignmentUtils.alignUpdateAssignments`, which runs during analysis before this rule ever sees the assignments. A nested-field `SET s.c1 = -1` is rewritten there into a full-struct-replacement assignment on the root attribute `s`.
In this PR I did not want to change the `AssignmentUtils`
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,68 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * Currently honored only for UPDATE. DELETE and MERGE ignore this interface: those operations
+ * receive full-width rows and {@link LogicalWriteInfo#updateSchema()} is absent for them.
+ *
+ * @since 4.3.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * The returned columns become the schema of updated, copied, and reinserted rows, in declared
+ * order. Implementations must include every column they want to receive (typically the columns
+ * reported by {@link RowLevelOperationInfo#updatedColumns()} plus any columns needed for row
+ * lookup or routing, e.g. a primary key).
+ * <p>
+ * If any of the columns from {@link RowLevelOperationInfo#updatedColumns()} are
+ * missing, an analysis exception is thrown.
+ * <p>
+ * For updates on nested fields such as {@code SET t.s.c1 = -1} the connector should declare the
+ * root struct column {@code s} rather than any nested field.
+ */
+ NamedReference[] requiredDataAttributes();
+
+ /**
+ * Returns additional data column references that must be present in the narrow scan but are
+ * NOT delivered to the writer.
+ * <p>
+ * Use this for columns needed only for planning -- e.g. resolving the table's partitioning
+ * expressions against the scan output, or as clustering keys returned from
+ * {@link RequiresDistributionAndOrdering#requiredDistribution()} -- where the column itself
+ * should not appear in the row passed to
+ * {@link DataWriter#writeUpdate(Object, Object)} / {@link DeltaWriter#update} /
+ * {@link DeltaWriter#reinsert}. Columns returned here are not reflected in
+ * {@link LogicalWriteInfo#updateSchema()}.
+ * <p>
+ * Must not overlap with {@link #requiredDataAttributes()}; defaults to an empty array.
+ */
+ default NamedReference[] scanOnlyDataAttributes() {
Review Comment:
Sounds good. I reverted this API.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -41,7 +44,11 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
EliminateSubqueryAliases(aliasedTable) match {
case r @ ExtractV2Table(tbl: SupportsRowLevelOperations) =>
checkNoGeneratedColumns(r, UPDATE)
- val table = buildOperationTable(tbl, UPDATE, r.options)
+ val updatedCols = assignments.collect {
Review Comment:
Good idea. Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, any required metadata attrs and optionally connector
Review Comment:
I went with this
```
// resolve all needed attrs (e.g. row ID, required metadata attrs, and any connector-declared
// data attrs for column-update writes)
```
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, any required metadata attrs and optionally connector
+ // declared attrs)
val rowAttrs = relation.output
+ val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates]
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveConnectorDataAttrs(relation, operation)
+ } else Nil
+ val scanOnlyDataAttrs = if (supportsColumnUpdate) {
+ resolveScanOnlyDataAttrs(relation, operation)
+ } else Nil
+
+ if (supportsColumnUpdate) {
+ validateUpdatedColumnsSubset(operation, assignments, connectorDataAttrs)
+ validateNoOverlap(operation, connectorDataAttrs, scanOnlyDataAttrs)
+ validatePartitionAttrsDeclared(operation, relation, connectorDataAttrs, scanOnlyDataAttrs)
+ }
+
+
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
- val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ if (supportsColumnUpdate && operation.representUpdateAsDeleteAndInsert) {
+ validateNoRowIdReassignment(operation, assignments, rowIdAttrs)
Review Comment:
Good catch. This was a bug. I fixed it by exempting metadata columns from the check:
```scala
val undeclared = rowIdAttrs
.filterNot(a => MetadataAttribute.isValid(a.metadata) || declaredIds.contains(a.exprId))
.map(_.name).distinct
```
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, any required metadata attrs and optionally connector
+ // declared attrs)
val rowAttrs = relation.output
+ val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates]
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveConnectorDataAttrs(relation, operation)
+ } else Nil
+ val scanOnlyDataAttrs = if (supportsColumnUpdate) {
+ resolveScanOnlyDataAttrs(relation, operation)
+ } else Nil
+
+ if (supportsColumnUpdate) {
+ validateUpdatedColumnsSubset(operation, assignments, connectorDataAttrs)
Review Comment:
Done, moved it into `resolveConnectorDataAttrs()`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -129,6 +156,17 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
V2ExpressionUtils.resolveRef[AttributeReference](FieldReference(name), plan)
}
+ protected def isIdentityAssignment(key: Attribute, value: Expression): Boolean = {
+ val unwrapped = value match {
+ case Alias(child, _) => child
+ case other => other
+ }
+ unwrapped match {
+ case attr: Attribute => AttributeSet(Seq(key)).contains(attr)
Review Comment:
Thanks for the suggestion. Done.
##########
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:
Got it. I removed this.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3718167431
##########
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:
After speaking to @aokolnychyi, the current unconditional inclusion of partition source columns in the narrow scan is incorrect. It should be connector-driven, not implicit on Spark's side (e.g. only requested when the connector actually needs it for V2ScanPartitioningAndOrdering or write-side clustering).
The tension is that connectors may need a column visible in the scan (for planning/clustering resolution) but not in the write payload. I'm considering three options:
1. Add `scanOnlyDataAttributes()` to the mixin for columns the connector wants in the narrow scan but not in the write payload.
2. Use `requiredDataAttributes()` for everything. Connectors that need scan-side resolution declare it here, but by contract it then also flows into the write row, putting the burden on the connector to ignore/filter it when it's not being updated.
3. Use existing `requiredMetadataAttributes()` (we would need to expand its scope to allow certain data attributes). These don't flow into the narrow write data row, but they do reach the writer via the metadata slot, which conflicts with requiredMetadataAttributes()'s original scope of synthetic columns only.
I'm leaning towards option 1 as it seems cleanest but I'm open to any other suggestions.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-5781515065
Summarizing the blocking items from my [inline review](https://github.com/apache/spark/pull/55518#pullrequestreview-5281801632). All four are correctness issues specific to the narrow path; the same statements work on the wide path.
1. **User-spelled attributes leak into the narrow scan output** ([`RewriteUpdateTable.scala:454`](https://github.com/apache/spark/pull/55518#discussion_r4074793007)): under the default case-insensitive analysis, referencing an undeclared column with different casing (e.g. `SET salary = salary + BONUS`) puts a `BONUS`-named attribute into the scan, so the UPDATE fails at execution with `INTERNAL_ERROR_ATTRIBUTE_NOT_FOUND`.
2. **Row-ID reassignment emits `pk` twice** ([`RewriteUpdateTable.scala:316`](https://github.com/apache/spark/pull/55518#discussion_r4074793021)): `buildColumnUpdateProjection` outputs both the new `Alias(pk + 10, "pk")` and the original `pk`, so a `CHECK (pk > 0)` constraint or a write clustered on `pk` fails with `AMBIGUOUS_REFERENCE`.
3. **Narrow relation breaks `PushDownUtils.pruneColumns` assumptions** ([`RewriteRowLevelCommand.scala:101`](https://github.com/apache/spark/pull/55518#discussion_r4074793028)): a connector that partially prunes (allowed by `SupportsPushDownRequiredColumns`) hits `NoSuchElementException` in `toOutputAttrs`, and one without `SupportsPushDownRequiredColumns` silently binds full-width rows to the narrow output, writing misaligned values.
4. **Metadata column accepted in `requiredDataAttributes()`** ([`RewriteRowLevelCommand.scala:153`](https://github.com/apache/spark/pull/55518#discussion_r4074793040)): `LogicalPlan.resolve` falls back to metadata columns, so declaring `_partition` puts it among the data attrs; on the split path the `Expand` projections and output attrs then disagree slot by slot, silently swapping `dep` and `_partition` values.
The remaining inline comments (the `writeUpdate` overload contract, the empty `schema()` documentation, and the test-infra items) are non-blocking.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4071667376
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -109,6 +123,43 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
relation)
}
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation. Non-existent columns are rejected with an `AnalysisException`.
+ */
+ protected def resolveRequiredDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes
+ if (refs.isEmpty) {
+ throw QueryCompilationErrors.emptyRequiredDataAttributesError(operation.getClass.getName)
+ }
+ val nested = refs.filter(_.fieldNames.length != 1).map(_.describe()).toImmutableArraySeq
+ if (nested.nonEmpty) {
+ throw QueryCompilationErrors.nestedRequiredDataAttributeError(
+ operation.getClass.getName, nested)
+ }
+ val normalizedNames = refs.map { ref =>
+ val name = ref.fieldNames.head
+ if (conf.caseSensitiveAnalysis) name else name.toLowerCase(Locale.ROOT)
+ }
+ val duplicates = normalizedNames.groupBy(identity).collect {
+ case (_, occurrences) if occurrences.length > 1 => occurrences.head
+ }.toSeq
Review Comment:
**Finding 32.** `occurrences.head` is the normalized name, not the declared one, so with `spark.sql.caseSensitive` off a declaration of `["PK", "Pk"]` reports `[pk]` — a spelling the connector never wrote. The nested-attribute check nine lines up reports `_.describe()`, so the two messages disagree about which spelling the connector author gets to see.
Keeping the first declared spelling:
```suggestion
val duplicates = refs.zip(normalizedNames).groupBy(_._2).collect {
case (_, occurrences) if occurrences.length > 1 => occurrences.head._1.describe()
}.toSeq
```
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -228,4 +382,169 @@ 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, validating that every assigned column is covered. Returns `Nil` otherwise.
+ */
+ private def resolveConnectorDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation,
+ assignments: Seq[Assignment]): Seq[AttributeReference] = operation match {
+ case scu: SupportsColumnUpdates =>
+ val attrs = resolveRequiredDataAttrs(relation, scu)
+ validateUpdatedColumnsSubset(scu, assignments, attrs)
+ attrs
+ case _ => Nil
+ }
+
+ /**
+ * Computes the narrow set of data columns that must be present in the scan for a column-update
+ * write: connector-declared attrs (`requiredDataAttributes()`), unioned with any table columns
+ * referenced by non-identity assignment RHS expressions, the operation condition, and the
+ * table's CHECK constraints.
+ */
+ 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 checkConstraintAttrs = resolveCheckConstraintAttrs(relation)
+ dedupAttrs(connectorDataAttrs ++ extraRefs ++ checkConstraintAttrs)
+ }
+
+ /**
+ * Returns the relation attributes referenced by the table's CHECK constraints, so that
+ * narrowing does not drop a column `ResolveTableConstraints` needs later to re-validate the
+ * rewritten write query. Falls back to the full relation output whenever a constraint's
+ * condition cannot be resolved to a `Predicate` (see `Check#predicate()`).
+ */
+ private def resolveCheckConstraintAttrs(
+ relation: DataSourceV2Relation): Seq[AttributeReference] = {
+ val checks = Option(relation.table.constraints).getOrElse(Array.empty[Constraint]).collect {
+ case c: Check => c
+ }
+ if (checks.isEmpty) {
+ Nil
+ } else if (checks.exists(_.predicate() == null)) {
+ relation.output
+ } else {
+ val refs = checks.flatMap(_.predicate().references()).toSeq
+ V2ExpressionUtils.resolveRefs[AttributeReference](refs, relation)
Review Comment:
**Finding 31.** `Check#predicate().references()` returns references at full depth, so a constraint on a nested field gives a two-part `FieldReference`. `V2ExpressionUtils.resolveRef` puts that through `LogicalPlan.resolve`, which returns an `Alias(GetStructField(...))` for a nested path, and its `asInstanceOf[T]` is erased — so the `Alias` rides along inside a `Seq[AttributeReference]` and blows up at the first frame that needs the real type.
Measured on `9932f88`, delta MoR (`column-update`) and CoW (`column-update-cow`), table `pk INT NOT NULL, id INT, dep STRING, s STRUCT<c1: INT, c2: INT>` partitioned by `dep`:
```sql
ALTER TABLE t ADD CONSTRAINT positive_c1 CHECK (s.c1 > 0);
UPDATE t SET id = -1 WHERE pk = 1;
```
```
java.lang.ClassCastException: class org.apache.spark.sql.catalyst.expressions.Alias cannot be cast to
class org.apache.spark.sql.catalyst.expressions.AttributeReference
at RewriteRowLevelCommand.dedupAttrs(RewriteRowLevelCommand.scala:107)
at RewriteUpdateTable$.computeNarrowReadAttrs(RewriteUpdateTable.scala:454)
```
Control: the same table, constraint and statement on the wide `supports-deltas` connector passes, so this is specific to narrowing. A `CHECK` on a *top-level* undeclared column is fine on all three of delta MoR, CoW and the split path — I ran that too.
Narrowing is at root-column granularity anyway, so truncating each reference to its root column is enough:
```suggestion
val refs = checks.flatMap(_.predicate().references()).toSeq
.map(ref => FieldReference(Seq(ref.fieldNames.head)))
V2ExpressionUtils.resolveRefs[AttributeReference](refs, relation)
```
Needs `import org.apache.spark.sql.connector.expressions.FieldReference`. With that applied both probes pass and the two suites stay at 155. Worth a nested-field-constraint test on both narrow paths.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4316536111
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala:
##########
@@ -92,6 +95,38 @@ object GroupBasedRowLevelOperationScanPlanning extends Rule[LogicalPlan] with Pr
}
}
+ /**
+ * Returns the columns of `relation` that the projections and filters directly above each read
+ * of the table use, as the rest of the query sees only their output. Only whole columns are
+ * pruned, as the query is not rewritten for a pruned nested schema.
Review Comment:
**22.** Minor: copy-on-write narrow scans are pruned only at the top level, which this Scaladoc states but the PR description and the public Javadoc do not. `readAttrs` returns bare relation attributes, so `PushDownUtils.pruneColumns` (L71) keeps every field of a struct that the UPDATE references only through a nested field, while delta-based scans prune nested fields through `V2ScanRelationPushDown` (`V2ScanRelationPushDown.scala:1003-1016`). For `pk INT NOT NULL, salary INT, s STRUCT<c1: INT, c2: INT>, dep STRING` and `UPDATE t SET salary = s.c1 WHERE pk = 1`, the copy-on-write scan is pruned to `s STRUCT<c1: INT, c2: INT>` and the delta scan to `s STRUCT<c1: INT>`. The nested `CHECK` test (`GroupBasedColumnUpdateTableSuite.scala:639-646`) only checks that `s` is read. This is not a regression, since the base reads every column, but on a wide struct it is a large part of the I/O this feature is meant to avoid. Could the PR description and the scan-builder paragraph of `SupportsC
olumnUpdates` (`SupportsColumnUpdates.java:41-44`) state that copy-on-write scans are pruned only at the top level, and leave nested pruning to a separate JIRA?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -129,6 +184,24 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
V2ExpressionUtils.resolveRef[AttributeReference](FieldReference(name), plan)
}
+ protected def isIdentityAssignment(key: Attribute, value: Expression): Boolean = {
Review Comment:
**20. An explicit `SET pk = pk` on an undeclared row-ID column makes the distribution check reject a row-ID column, which the Javadoc allows.** This method strips one `Alias` before `semanticEquals`, so it treats `SET pk = pk`, which assignment alignment turns into `Alias(pk, "pk")` (`TableOutputResolver.scala:331`), as an identity: `pk` is not reported in `updatedColumns()` and does not have to be declared. `buildOriginalRowIdValues` (L416-429) still uses `!key.semanticEquals(value)`, so for the same assignment it adds `__original_row_id_pk` (`RewriteUpdateTable.scala:220`). `newLazyRowIdProjection` (L390-399) then reads that column instead of `pk`, no projection reads `pk`, `orderReadColumnsFirst` moves it after the last read ordinal, and `validateColumnUpdateWrite` (`V2Writes.scala:157-166`) rejects a write that clusters by `pk`.
With the setup of "distribution by an undeclared row ID column is allowed" (`DeltaBasedColumnUpdateTableSuite.scala:346`: `column-update-req-attrs` = `salary,dep` and `column-update-cluster-by` = `pk`), `UPDATE t SET salary = -1 WHERE pk = 1` succeeds, but `UPDATE t SET salary = -1, pk = pk WHERE pk = 1` fails with `COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS` and `[pk]`, although `SupportsColumnUpdates.java:81-84` exempts row-ID columns for delta operations. On the wide path the extra column is harmless, because the row projection still reads `pk`.
Could `buildOriginalRowIdValues` use `isIdentityAssignment` as well, and could that test get a `pk = pk` variant? Its other caller, `RewriteMergeIntoTable.scala:351`, gets the same result, since an identity assignment does not change the row ID.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -258,6 +354,407 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep` (an unconditionally declared base column) plus whatever columns
+ // Spark reports as being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = {
+ val base = Seq(FieldReference("pk"), FieldReference("dep"))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ private val clusterRef: NamedReference = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY))
+ .map(FieldReference(_))
+ .getOrElse(PARTITION_COLUMN_REF)
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.columnUpdateSchema().isPresent) {
+ capturedInfo.columnUpdateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
Review Comment:
**25.** Minor: the delta column-update fixtures ignore `no-metadata`, so no test sends narrow rows through `DeltaWritingSparkTask`. With `no-metadata`, the inherited `requiredMetadataAttributes()` is empty (L307-313), but this write still requires clustering and ordering by `clusterRef`, which defaults to `_partition` (L374-376), and the split write hard-codes `_partition` (item 18). Every narrow UPDATE then fails with `COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS` and `[_partition]`, whereas the parent returns a plain `DeltaWrite` (L326-329) and the copy-on-write operation clusters by `dep` (L587-588). `WriteDeltaExec` uses `DeltaWritingSparkTask` exactly when there is no metadata projection (`WriteToDataSourceV2Exec.scala:435-440`), so that task never receives narrow rows in a test. The round-4 comment ([r4074793050](https://github.com/apache/spark/pull/55518#discussion_r4074793050)) raised the same gap for copy-on-write, which now has no-metadata tests (`GroupBasedColumnU
pdateTableSuite.scala:145`, `:181`). Similarly, L367 overrides `representUpdateAsDeleteAndInsert()` to `false`, so `split-updates` is silently ignored on these tables. Could the delta builders honor `no-metadata` like the parent, and could a no-metadata variant of `DeltaBasedColumnUpdateTableSuite` run, as `DeltaBasedNoMetadataUpdateTableSuite` does for full rows?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -258,6 +354,407 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep` (an unconditionally declared base column) plus whatever columns
+ // Spark reports as being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = {
+ val base = Seq(FieldReference("pk"), FieldReference("dep"))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ private val clusterRef: NamedReference = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY))
+ .map(FieldReference(_))
+ .getOrElse(PARTITION_COLUMN_REF)
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.columnUpdateSchema().isPresent) {
+ capturedInfo.columnUpdateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(Array(clusterRef))
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ Array[SortOrder](
+ LogicalExpressions.sort(
+ clusterRef,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering())
+ )
+ }
+
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For column-update writes, updated rows contain only the declared columns
+ // (columnUpdateSchema from LogicalWriteInfo). We expand each row to the full
+ // table schema by overlaying them on the base row found by pk.
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+
+ val mergedData = newData.map { buf =>
+ val merged = new BufferedRows(buf.key, schema)
+ val updateOpName = UTF8String.fromString(Update.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == updateOpName) {
+ val pk = logRow.getInt(1)
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val baseRow = currentRowByPk(pk)
+ val fullRow = overlayNarrowRow(baseRow, narrowRow, writeFieldIdx)
+ merged.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ merged.rows.append(insertRow.copy())
+ }
+ }
+ merged
+ }
+
+ withDeletes(newData)
+ withData(mergedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ }
+ }
+ }
+ }
+
+ class DeltaBasedColumnUpdateOperationWithReqAttrs(
+ command: Command,
+ reqCols: Array[String],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, options = options) {
+ override def requiredDataAttributes(): Array[NamedReference] = reqCols.map(FieldReference(_))
+ }
+
+ class DeltaBasedColumnUpdateSplitOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols, options) {
+ override def representUpdateAsDeleteAndInsert(): Boolean = true
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.columnUpdateSchema().isPresent) {
+ capturedInfo.columnUpdateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+ override def requiredDistribution(): Distribution =
+ Distributions.clustered(Array(PARTITION_COLUMN_REF))
+ override def requiredOrdering(): Array[SortOrder] = Array[SortOrder](
+ LogicalExpressions.sort(
+ PARTITION_COLUMN_REF,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()))
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For delete+reinsert with narrow writes, the REINSERT row has only the
+ // connector-declared columns (requiredDataAttributes order). The base row is
+ // looked up by `pk`, so `pk` must be declared for the write to commit.
+ // Reconstruct the full row by overlaying the narrow row onto the original.
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+ val reinsertOpName = UTF8String.fromString(Reinsert.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+ val pkIdx = writeFieldIdx("pk")
Review Comment:
**23. Every row-level DELETE on a split column-update table fails at commit, so the DELETE fallback is not tested on that operation.** Unlike the parent, which takes `pk` from the row ID inside its `Update` branch (L427-428), this lookup runs before the log loop. A delta DELETE has no row projection, so `schema()` is empty and `columnUpdateSchema()` is absent (`V2Writes.scala:240-241`), `capturedWriteSchema` (L474-478) is empty, and this line throws `NoSuchElementException: key not found: pk`. Any DELETE that is not metadata-only, e.g. `DELETE FROM t WHERE pk < 2`, fails this way on `column-update-split`, `column-update-split-req-attrs` and `column-update-split-row-id` tables. This is the same line as item 14, but a different case: item 14 is a split UPDATE that cannot commit.
The round-4 comment ([r4074793068](https://github.com/apache/spark/pull/55518#discussion_r4074793068)) asked for the DELETE and MERGE fallback to be tested on both connector types, and the PR description lists "DELETE and MERGE falling back to full-width rows" among what the suites cover. The split operation got "column-update split: MERGE falls back to full-width rows" (`DeltaBasedColumnUpdateTableSuite.scala:225`), but no DELETE test. Could the `pk` lookup move into the `Reinsert` branch, and could a "column-update split: row-level DELETE falls back to full-width rows" test be added?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,754 @@
+/*
+ * 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.{SparkRuntimeException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.{InMemoryBaseTable, InMemoryRowLevelOperationTable}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+
+class GroupBasedColumnUpdateTableSuite extends UpdateTableSuiteBase {
Review Comment:
**24. Three inherited UPDATE tests do not run against the narrow fixture.** Since finding 25 ([r3741376793](https://github.com/apache/spark/pull/55518#discussion_r3741376793)), this suite and `DeltaBasedColumnUpdateTableSuite` extend the existing UPDATE bases, and the PR description says "so the existing UPDATE tests also run against narrow writes". But `extraTableProps` reaches a table only through the `createTable` helper (`RowLevelOperationSuiteBase.scala:136`), and three inherited tests create their table another way, so both suites run them on a plain table:
- "update table with expression-based default values" (`catalog.createTable` with a `TableInfo` without properties, `UpdateTableSuiteBase.scala:100-101`)
- "update with current_timestamp default value using DEFAULT keyword" (`CREATE TABLE`, `UpdateTableSuiteBase.scala:849-850`)
- "update table with recursive CTE" (`CREATE TABLE`, `UpdateTableSuiteBase.scala:916-917`)
The first one, `SET value = DEFAULT` on a column with an existence default, is worth running narrow. The fixtures also declare `pk` and `dep` unconditionally (`InMemoryRowLevelOperationTable.scala:369`, `:574`), so the CTE table `(val INT)` would fail with `COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE` if the properties did reach it. Could these tables pick up `extraTableProps`, with the fixtures declaring only existing columns, or could the description say which inherited tests stay on the wide path?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -228,4 +240,114 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Returns the output of the write relation, which is the layout of the data rows the connector
+ * receives. If the operation opts into column updates, this is the connector's
+ * `requiredDataAttributes()`, validated against the assignments. Otherwise it is the full
+ * relation output.
+ */
+ private def resolveWriteAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation,
+ assignments: Seq[Assignment]): Seq[AttributeReference] = operation match {
+ case op: SupportsColumnUpdates with SupportsDelta if op.representUpdateAsDeleteAndInsert =>
+ val attrs = resolveDeclaredAttrs(relation, op, assignments)
+ val rowIdAttrs = resolveRowIdAttrs(relation, op)
+ validateNoRowIdReassignment(op, relation, attrs, assignments, rowIdAttrs)
+ validateRowIdDeclared(op, attrs, rowIdAttrs, resolveRequiredMetadataAttrs(relation, op))
+ attrs
+ case op: SupportsColumnUpdates =>
+ resolveDeclaredAttrs(relation, op, assignments)
+ case _ =>
+ relation.output
+ }
+
+ /**
+ * Resolves the connector's `requiredDataAttributes()` and checks that it covers every assigned
+ * column.
+ */
+ private def resolveDeclaredAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates,
+ assignments: Seq[Assignment]): Seq[AttributeReference] = {
+ val attrs = resolveRequiredDataAttrs(relation, operation)
+ validateUpdatedColumnsSubset(operation, assignments, attrs)
+ attrs
+ }
+
+ /**
+ * 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)
+ }
+ }
+
+ /**
+ * For connectors that opt into narrow column updates AND represent UPDATE as delete + insert,
+ * reject reassignment of any row-ID column, unless `requiredDataAttributes()` covers every
+ * column in the relation. In that case the REINSERT payload already is the full row with the
+ * new row-ID value, and the DELETE half still carries the original row-ID via
+ * `newLazyRowIdProjection`, so reassignment is safe. Otherwise the REINSERT path has no
+ * row-ID channel to reconstruct columns outside `requiredDataAttributes()`.
+ */
+ private def validateNoRowIdReassignment(
+ operation: RowLevelOperation,
+ relation: DataSourceV2Relation,
+ connectorDataAttrs: Seq[AttributeReference],
+ assignments: Seq[Assignment],
+ rowIdAttrs: Seq[Attribute]): Unit = {
+ val declaredIds = connectorDataAttrs.map(_.exprId).toSet
+ if (relation.output.forall(a => declaredIds.contains(a.exprId))) {
Review Comment:
**21. This exemption never applies when the target relation carries metadata columns.** `relation.output` is not always the data columns. When an UPDATE is not resolved in the first iteration of the Resolution batch, for example because a comparison needs type coercion, `AddMetadataColumns` (`Analyzer.scala:691`, `:1145`) runs on the resolved command in the next iteration and calls `withMetadataColumns()` on the target relation (`Analyzer.scala:1194`), which appends every metadata column to its output (`DataSourceV2Relation.scala:137-141`). Metadata columns cannot be declared (`RewriteRowLevelCommand.scala:157`), so this `forall` is then false even when every data column is declared.
With the setup of "row-ID reassignment is allowed when requiredDataAttributes covers every table column" (`DeltaBasedColumnUpdateTableSuite.scala:659`), `UPDATE t SET pk = pk + 10, salary = -1 WHERE dep = 'hr' AND _partition = 'hr' AND salary > 0.5` fails with `COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT`, while the same statement with `salary > 0`, or without `_partition = 'hr'`, is exempted. This is the exemption added for finding 22 ([r3741376785](https://github.com/apache/spark/pull/55518#discussion_r3741376785)).
Could this compare only the data columns, e.g. `relation.output.filterNot(a => MetadataAttribute.isValid(a.metadata))`?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3575292426
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -106,38 +105,78 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val remainingRowsPlan = addOperationColumn(COPY_OPERATION,
Filter(remainingRowFilter, readRelation))
- // the new state is a union of updated and copied over records
- val query = Union(updatedRowsPlan, remainingRowsPlan)
+ val updatedAndRemainingRowsPlan = Union(updatedRowsPlan, remainingRowsPlan)
- // build a plan to replace read groups in the table
val writeRelation = relation.copy(table = operationTable)
- val projections = buildReplaceDataProjections(query, relation.output, metadataAttrs)
+ val projections = buildReplaceDataProjections(updatedAndRemainingRowsPlan, rowAttrs,
+ metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
- ReplaceData(writeRelation, cond, query, relation, projections, groupFilterCond)
+ ReplaceData(writeRelation, cond, updatedAndRemainingRowsPlan, relation, projections,
+ groupFilterCond)
}
- // this method assumes the assignments have been already aligned before
+ /**
+ * When the connector supports column updates and declares required data attributes,
+ * the read relation is narrowed at analysis time so that GroupBasedRowLevelOperationScanPlanning
+ * uses only the needed columns for the scan. Otherwise, the full relation output is used.
+ */
+ private def buildReplaceDataReadRelation(
+ relation: DataSourceV2Relation,
+ operationTable: RowLevelOperationTable,
+ assignments: Seq[Assignment],
+ cond: Expression): (DataSourceV2Relation, Seq[Attribute], Seq[AttributeReference]) = {
+
+ val operation = operationTable.operation
+ val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
+ val connectorDataAttrs = resolveRequiredDataAttrs(relation, operation)
+ val isNarrow = operation.supportsColumnUpdates() && connectorDataAttrs.nonEmpty
+
+ val readRelation = if (isNarrow) {
+ val allRequired = (connectorDataAttrs ++ computeAssignedAttrs(assignments)).distinct
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs, dataAttrs = allRequired,
+ cond = cond)
+ } else {
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs)
+ }
+
+ val rowAttrs: Seq[Attribute] = if (isNarrow) connectorDataAttrs else relation.output
+
+ (readRelation, rowAttrs, metadataAttrs)
+ }
+
+ /**
+ * Builds the update projection for ReplaceData plans. Assumes assignments are already aligned.
+ *
+ * plan.output may be narrowed by buildReplaceDataReadRelation, so only columns present in the
+ * plan are projected.
+ */
private def buildReplaceDataUpdateProjection(
Review Comment:
Added a new method for narrow projection
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192354146
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -41,7 +42,13 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
EliminateSubqueryAliases(aliasedTable) match {
case r @ ExtractV2Table(tbl: SupportsRowLevelOperations) =>
- val table = buildOperationTable(tbl, UPDATE, CaseInsensitiveStringMap.empty())
+ val updatedCols = assignments.collect {
+ case Assignment(key: AttributeReference, value)
+ if !isIdentityAssignment(key, value) =>
Review Comment:
One liner doesn't violate the line-length rule, does it?
```scala
- case Assignment(key: AttributeReference, value)
- if !isIdentityAssignment(key, value) =>
+ case Assignment(key: AttributeReference, value) if !isIdentityAssignment(key, value) =>
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192416031
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -106,38 +104,92 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val remainingRowsPlan = addOperationColumn(COPY_OPERATION,
Filter(remainingRowFilter, readRelation))
- // the new state is a union of updated and copied over records
- val query = Union(updatedRowsPlan, remainingRowsPlan)
+ val updatedAndRemainingRowsPlan = Union(updatedRowsPlan, remainingRowsPlan)
- // build a plan to replace read groups in the table
val writeRelation = relation.copy(table = operationTable)
- val projections = buildReplaceDataProjections(query, relation.output, metadataAttrs)
+ val query = updatedAndRemainingRowsPlan
+ val metadataAttrs = resolveRequiredMetadataAttrs(relation, operationTable.operation)
+ val projections = buildReplaceDataProjections(query, rowAttrs, metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
ReplaceData(writeRelation, cond, query, relation, projections, groupFilterCond)
}
+ // Common read-relation setup shared by both CoW plan builders.
+ //
+ // When the connector supports column updates and declares required data attributes,
+ // the read relation is narrowed at analysis time so that
+ // GroupBasedRowLevelOperationScanPlanning uses only the needed columns for the scan.
+ // Otherwise the full relation output is used.
+ private def buildCoWReadSetup(
+ relation: DataSourceV2Relation,
+ operationTable: RowLevelOperationTable,
+ assignments: Seq[Assignment],
+ cond: Expression): (DataSourceV2Relation, Seq[Attribute]) = {
+
+ val operation = operationTable.operation
+ val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
+ val connectorDataAttrs = resolveRequiredDataAttrs(relation, operation)
+ val isNarrow = operation.supportsColumnUpdates() && connectorDataAttrs.nonEmpty
+
+ // CoW scan narrowing must be done manually at analysis time.
+ // GroupBasedRowLevelOperationScanPlanning (an optimizer rule that fires after analysis)
+ // always reads relation.output directly when building the physical scan -- it does not
+ // observe Project nodes above the relation, so optimizer-driven column pruning has no
+ // effect on CoW scans. We narrow DataSourceV2Relation.output here so that rule picks
+ // up the narrow set.
+ val readRelation = if (isNarrow) {
+ val allRequired = (connectorDataAttrs ++ computeAssignedAttrs(assignments)).distinct
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs, dataAttrs = allRequired,
+ cond = cond)
+ } else {
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs)
+ }
+
+ // CoW write schema (two paths only, no heuristic for CoW):
+ // - Narrow path (connectorDataAttrs declared): exactly connector-declared cols in declared
+ // order. The connector must declare ALL columns it wants to receive.
+ // - Full path (connectorDataAttrs empty OR supportsColumnUpdates=false): full table output.
+ // Unlike MOR, CoW does not have a heuristic assigned-only path because
+ // GroupBasedRowLevelOperationScanPlanning needs explicit column declarations to narrow.
+ val rowAttrs: Seq[Attribute] = if (isNarrow) connectorDataAttrs else relation.output
+
+ (readRelation, rowAttrs)
Review Comment:
Please return `metadataAttrs` too to avoid recompilation in the caller-side.
> val metadataAttrs = resolveRequiredMetadataAttrs(relation, operationTable.operation)
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3720993832
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,47 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ *
+ * @since 4.3.0
Review Comment:
**Finding 5.** `4.3.0` is no longer a version this can ship in. `branch-4.3` was cut on 2026-05-01 and `branch-4.x` was bumped to `4.4.0-SNAPSHOT` on 2026-08-03 by `[SPARK-58534][4.X][BUILD] Change version in branch-4.x to 4.4.0` -- one day before this head. `dev/next_version_candidates.py` on this checkout prints:
```
master 5.0.0
branch-4.x 4.4.0
```
So this should be `4.4.0` (the version it first ships in via `branch-4.x`), or `5.0.0` if the team decides the abstract-method addition on `RowLevelOperationInfo` makes it master-only. Same fix needed in three more places:
- `RowLevelOperationInfo.java:49`
- `LogicalWriteInfo.java:73`
- `DataWriter.java:98` and `:115`
Two related nits in `project/MimaExcludes.scala:64-66`: the entry sits in `v43excludes`, whose header says "Exclude rules for 4.3.x from 4.2.0" -- 4.3.0 does not have `updatedColumns()`, so the filter belongs in the list for the version that does (`branch-4.x` has a `v44excludes` for exactly this; master currently has no `v50excludes` list of its own, so worth confirming the intended home with a committer). And the comment cites `[SPARK-56599]`, the SPIP umbrella, while every neighbouring entry cites the implementing ticket -- that should be `[SPARK-58111]`.
##########
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)
Review Comment:
**Finding 6.** On this path the connector cannot reconstruct the columns it did not declare once the row ID is itself reassigned, and the result is silent NULLs rather than an error.
Repro -- a `representUpdateAsDeleteAndInsert` + `SupportsColumnUpdates` connector declaring `requiredDataAttributes() = [pk, salary]`, table `pk INT NOT NULL, salary INT, dep STRING` partitioned by `dep`, rows `(1, 100, 'hr')` and `(2, 200, 'software')`:
```sql
UPDATE t SET pk = pk + 10, salary = -1 WHERE dep = 'hr'
```
gives `[11, -1, null]` -- `dep` is gone. The same query against the same table with narrowing off (`supports-deltas` + `split-updates`) gives `[11, -1, 'hr']`, so this is specific to the narrow path.
Why it can't be fixed connector-side: `buildNarrowDeletesAndInserts` does build the REINSERT output over every narrow-scan row attr (so `dep` is in the `Expand`), but `updateRowProjection` then selects only `connectorDataAttrs` by name, so the row reaching `writer.reinsert(metadata, row)` is `[pk_new, salary_new]`. Unlike `writer.update(metadata, rowId, row)`, which gets the pre-update value through `buildOriginalRowIdValues`' `__original_row_id_pk` column, `reinsert` has no row-ID channel at all -- and `deltaDeleteOutput` nullifies every non-row-ID column on the paired DELETE. So nothing Spark hands the connector identifies the row being replaced, and no lookup key can recover `dep`.
Two ways out, smallest first:
```scala
// in buildWriteDeltaPlan, before choosing buildNarrowDeletesAndInserts:
// reject the combination that is provably unreconstructible
if (operation.representUpdateAsDeleteAndInsert &&
assignments.exists(a => rowIdAttrSet.contains(a.key.asInstanceOf[Attribute]) &&
!isIdentityAssignment(a.key.asInstanceOf[Attribute], a.value))) {
throw ... // "column-level updates cannot reassign a row ID column when updates are
// represented as delete + insert"
}
```
or thread `buildOriginalRowIdValues(rowIdAttrs, assignments)` into `deltaReinsertOutput` so the original row ID travels with the REINSERT row. The first is contained enough for this PR; the second needs the reinsert contract to grow. Either way this needs a test -- there is currently none for split updates plus a row-ID assignment (`column-update split: data correctness` only reassigns `id`).
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala:
##########
@@ -418,7 +432,29 @@ case class ReplaceData(
// validates row projection output is compatible with table attributes
private def rowAttrsResolved: Boolean = {
val inRowAttrs = DataTypeUtils.toAttributes(projections.rowProjection.schema)
- table.skipSchemaResolution || areCompatible(inRowAttrs, table.output)
+ val inUpdateAttrs = projections.updateRowProjection match {
+ case Some(projection) => DataTypeUtils.toAttributes(projection.schema)
+ case None => Nil
+ }
+ // `rowProjection` (INSERT-tagged rows) validates against `table.output` -- the full table
+ // shape. `updateRowProjection` (UPDATE/COPY-tagged rows) is narrow for column-update
+ // connectors, so it validates against `projectedDataAttrs` (the connector-declared narrow
+ // set) instead. When the connector does not mix in `SupportsColumnUpdates`,
+ // `updateRowProjection` is absent and `updateResolved` is trivially true.
+ val insertResolved = table.skipSchemaResolution || inRowAttrs.isEmpty ||
Review Comment:
**Finding 10.** The `inRowAttrs.isEmpty` term is needed for the column-update path, where `rowProjection` is deliberately `ProjectingInternalRow(StructType(Nil), Nil)`, but it is written unconditionally, so it also switches off the `areCompatible(inRowAttrs, table.output)` check for connectors that never opt in. That check is the only thing validating the INSERT-shaped row projection against the table shape, and silently skipping it on an empty projection is a wider behaviour change than this PR needs -- the same class of thing @dongjoon-hyun pushed back on for the `Some(empty)` -> `None` restore in `buildWriteDeltaProjections`.
Gating it keeps every non-column-update plan byte-identical to master:
```scala
val insertResolved = table.skipSchemaResolution ||
(inUpdateAttrs.nonEmpty && inRowAttrs.isEmpty) ||
areCompatible(inRowAttrs, table.output)
```
Same at `:558` in `WriteDelta.rowAttrsResolved`.
While you are in here: `updateResolved` does not honour `table.skipSchemaResolution`, so an `ACCEPT_ANY_SCHEMA` connector that opts into column updates now gets strict validation where master validated nothing. Probably worth folding into the same `skipSchemaResolution` short-circuit.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -267,13 +268,97 @@ abstract class RowLevelOperationSuiteBase
protected def checkLastWriteInfo(
expectedRowSchema: StructType = new StructType(),
expectedRowIdSchema: Option[StructType] = None,
- expectedMetadataSchema: Option[StructType] = None): Unit = {
+ expectedMetadataSchema: Option[StructType] = None,
+ expectedUpdateSchema: Option[StructType] = None): Unit = {
val info = table.lastWriteInfo
assert(info.schema == expectedRowSchema, "row schema must match")
val actualRowIdSchema = Option(info.rowIdSchema.orElse(null))
assert(actualRowIdSchema == expectedRowIdSchema, "row ID schema must match")
val actualMetadataSchema = Option(info.metadataSchema.orElse(null))
assert(actualMetadataSchema == expectedMetadataSchema, "metadata schema must match")
+ val actualUpdateSchema = Option(info.updateSchema.orElse(null))
+ assert(actualUpdateSchema == expectedUpdateSchema, "update schema must match")
+ }
+
+ protected def getUpdateSummary(): org.apache.spark.sql.connector.write.UpdateSummary = {
+ catalog.loadTable(ident).asInstanceOf[InMemoryTable]
+ .commits.last.writeSummary.get
+ .asInstanceOf[org.apache.spark.sql.connector.write.UpdateSummary]
+ }
+
+ /**
+ * Asserts the last UPDATE's write summary metrics. `deltaUpdate` controls the expected
+ * COPY count: MoR connectors emit only deltas (no COPY rows) so `numCopiedRows` is forced
+ * to 0; CoW connectors emit COPY rows for unchanged rows in matched groups.
+ */
+ protected def checkUpdateMetrics(
+ numUpdatedRows: Long,
+ numCopiedRows: Long,
+ deltaUpdate: Boolean = false): Unit = {
+ val summary = getUpdateSummary()
+ assert(summary.numUpdatedRows() === numUpdatedRows,
+ s"Expected numUpdatedRows=$numUpdatedRows, got ${summary.numUpdatedRows()}")
+ val expectedCopied = if (deltaUpdate) 0L else numCopiedRows
+ assert(summary.numCopiedRows() === expectedCopied,
+ s"Expected numCopiedRows=$expectedCopied, got ${summary.numCopiedRows()}")
+ }
+
+ /**
+ * Asserts that the column names in RowLevelOperationInfo.updatedColumns() received by the
+ * last operation match exactly the expected set. Order is ignored.
+ */
+ protected def checkLastUpdatedColumns(expectedNames: String*): Unit = {
+ val actual = table.lastUpdatedColumns.map(_.describe()).toSet
+ val expected = expectedNames.toSet
+ assert(actual == expected,
+ s"updatedColumns mismatch: expected ${expected.mkString("[", ", ", "]")} " +
+ s"but got ${actual.mkString("[", ", ", "]")}")
+ }
+
+ /**
+ * Asserts the set of top-level column names present in the last connector scan schema.
+ * Metadata columns are filtered out so tests can focus on data-column narrowing without
+ * having to enumerate `_partition` / `index` on every assertion.
+ */
+ protected def checkLastScanDataColumns(expectedNames: String*): Unit = {
Review Comment:
**Finding 13.** Nothing calls this -- both new suites use `checkLastScanExcludes` / `checkLastScanIncludes`, and the comment block above the scan-narrowing tests explains why exact-set matching was avoided. Worth dropping so a later reader does not take it as the intended assertion helper.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,47 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
Review Comment:
**Finding 11.** The interface doc reads as if the mix-in applies to any `RowLevelOperation`, but only `RewriteUpdateTable` honours it: `RewriteDeleteFromTable` and `RewriteMergeIntoTable` call the three-argument `buildReplaceDataProjections` / four-argument `buildWriteDeltaProjections`, so `updateRowProjection` stays `None`, `updateSchema()` stays absent, and `requiredDataAttributes()` is never read. I confirmed a MERGE against a `SupportsColumnUpdates` connector runs correctly with full-width rows -- so this is not a correctness problem, but a connector reading the doc would reasonably expect narrow rows there.
Finding 2 fixed the equivalent note on `updatedColumns()`; this is the same point on the mix-in itself, which is what @dongjoon-hyun originally asked for ("shall we mention that `DELETE and MERGE ignores this method`?" on the old `RowLevelOperation.java`).
```suggestion
* for updated, copied, and reinserted records, instead of the full table row.
* <p>
* Currently honored only for UPDATE. DELETE and MERGE ignore this interface: those operations
* receive full-width rows and {@link LogicalWriteInfo#updateSchema()} is absent for them.
```
Worth a cross-reference from `DeltaWriter#update` and `DeltaWriter#reinsert` too -- their `row` parameter is documented as "a row with updated values" with no hint that it may follow `updateSchema()` rather than `schema()`.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,541 @@
+/*
+ * 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, TableInfo}
+import org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity, reference}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+
+/**
+ * Tests for UPDATE statements targeting connectors that return true from
+ * [[org.apache.spark.sql.connector.write.RowLevelOperation#supportsColumnUpdates]].
Review Comment:
**Finding 12.** Two stale references from before the mix-in redesign: `RowLevelOperation#supportsColumnUpdates` no longer exists, and the narrowed schema is reported through `updateSchema()`, not `schema()` (which this suite itself asserts is empty).
```suggestion
* Tests for UPDATE statements targeting connectors that mix in
* [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
*
* When a connector supports column updates, Spark narrows the update-row projection
* (LogicalWriteInfo.updateSchema()) to contain only the declared columns rather than
* the full table row.
```
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +82,42 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated, copied, or reinserted record with metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must
+ * override this method when mixing in {@link SupportsColumnUpdates}; the default delegates
+ * to {@link #write(Object, Object)} so existing connectors are unaffected.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T metadata, T record) throws IOException {
+ write(metadata, record);
+ }
+
+ /**
+ * Writes one updated, copied, or reinserted record without metadata.
+ * <p>
+ * Equivalent to {@link #writeUpdate(Object, Object)} for writers that do not require metadata.
+ * The default delegates to {@link #write(Object)}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T record) throws IOException {
Review Comment:
**Finding 8.** The default is only right in the one case where the feature does nothing.
`writeUpdate` is reached from exactly two places -- `DataAndMetadataWritingSparkTask` and `DataWithProjectionWritingSparkTask`, both built by `ReplaceDataExec.writingTask` -- and both gate on `useWriteUpdate = updateDataProj != null`. `updateRowProjection` is `Some` only when `updateRowAttrs.nonEmpty`, which only happens for a `SupportsColumnUpdates` operation. So every call to `writeUpdate` passes the *narrow* projection; the default then forwards that narrow row to `write(...)`, which by its own contract takes a row in `LogicalWriteInfo.schema()` shape. It happens to be correct only if the connector declared every table column, i.e. if narrowing is a no-op.
A connector that mixes in `SupportsColumnUpdates` but forgets to override the writer therefore gets positionally misaligned rows with no error. The javadoc says "Implementations must override this method", but nothing enforces it, and `LogicalWriteInfo` right next door already uses the loud default for exactly this situation (`metadataSchema()` throws `DATA_SOURCE_METADATA_SCHEMA_NOT_IMPLEMENTED`). Since no existing connector can receive this call, making it throw is backward compatible:
```java
default void writeUpdate(T record) throws IOException {
throw new SparkUnsupportedOperationException(
"DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED", Map.of("class", getClass().getName()));
}
```
Same for `writeUpdate(T metadata, T record)` at `:100`.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedUpdateTableSuite.scala:
##########
@@ -45,11 +45,13 @@ class DeltaBasedUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase {
sql(s"SELECT * FROM $tableNameAsString"),
Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ // info.schema() reflects the row projection: `id` is non-nullable because the assignment
+ // supplies a non-null literal, matching master's projection-derived write schema semantics.
Review Comment:
**Finding 14.** These two hunks are not needed. I restored the original `StructType(table.schema.map { case attr if attr.name == "id" => attr.copy(nullable = false); case attr => attr })` form on this head (both call sites) and the whole suite still passes -- 62 tests, 0 failures. `PK_FIELD` and the hardcoded fields are the same values `table.schema` already produces here, so the rewrite is churn in a suite this PR does not otherwise touch, and it drags a widened `org.apache.spark.sql.types` import along with it.
The comment is the part I would definitely drop either way: "matching master's projection-derived write schema semantics" reads as if the non-opted-in path changed behaviour, and it didn't.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3741376788
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -109,6 +120,22 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
relation)
}
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation. Non-existent columns are rejected with an `AnalysisException`.
+ */
+ protected def resolveRequiredDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes
+ if (refs.isEmpty) {
Review Comment:
**Finding 23.** `resolveRefs[AttributeReference]` casts the result of `plan.resolve`, which is an `Alias` — not an `AttributeReference` — for any nested path, so a nested declaration escapes as an unhandled `ClassCastException`.
Measured on `774412c`: `column-update-req-attrs = "pk,s.c1"` on `pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING` with `UPDATE t SET s.c1 = 9 WHERE pk = 1` gives
```
java.lang.ClassCastException: class org.apache.spark.sql.catalyst.expressions.Alias
cannot be cast to class org.apache.spark.sql.catalyst.expressions.AttributeReference
```
`SupportsColumnUpdates`' javadoc does say to declare the root column at root-column granularity, so this is a connector mistake — but it is an easy one, because `FieldReference(String)` parses dots, so `FieldReference("s.c1")` produces the two-part path without the author intending it. A one-line guard next to the existing empty check turns it into a real error:
```scala
val nested = refs.filter(_.fieldNames.length != 1).map(_.describe())
if (nested.nonEmpty) {
throw QueryCompilationErrors.<nestedRequiredDataAttributeError>(
operation.getClass.getName, nested)
}
```
Same treatment for `scanOnlyDataAttributes()` in `RewriteUpdateTable.scala:472`.
##########
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(
+ operation: RowLevelOperation,
+ relation: DataSourceV2Relation,
+ connectorDataAttrs: Seq[AttributeReference],
+ scanOnlyDataAttrs: Seq[AttributeReference]): Unit = {
+ val partitionRefNames = relation.table.partitioning().toImmutableArraySeq
+ .flatMap(_.references.toImmutableArraySeq)
+ .map(_.fieldNames.head)
+ .toSet
+ val declared = AttributeSet(connectorDataAttrs ++ scanOnlyDataAttrs)
+ val undeclared = relation.output
+ .filter(a => partitionRefNames.exists(name => conf.resolver(name, a.name)))
+ .filterNot(declared.contains).map(_.name)
+ if (undeclared.nonEmpty) {
+ throw QueryCompilationErrors.requiredDataAttributesMissingPartitionColumnsError(
+ operation.getClass.getName, undeclared)
+ }
+ }
+
+ /**
+ * 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 22.** The error message tells the connector two ways out, and the guard only implements one.
`SPLIT_UPDATE_ROW_ID_REASSIGNMENT` reads: "Either avoid reassigning row ID columns, **or include every table column in `requiredDataAttributes()`**." But `validateNoRowIdReassignment` takes `(operation, assignments, rowIdAttrs)` and throws on any non-identity assignment to a row-ID column — `connectorDataAttrs` is not a parameter, so the second remedy cannot possibly be honoured. A connector author who follows it declares every column and still gets rejected.
Two ways to make them agree, and I think the first is right:
- Implement the remedy — skip the check when the declaration already covers the relation, since then the REINSERT payload from `buildNarrowDeletesAndInserts` really is the full row with the new row ID, and the DELETE branch still carries the original row ID via `newLazyRowIdProjection`:
```scala
val declared = connectorDataAttrs.map(_.exprId).toSet
if (!relation.output.forall(a => declared.contains(a.exprId))) {
// existing reassignment check
}
```
That also matters because a plain `supports-deltas` + `split-updates` connector *can* do `SET pk = ...` today, so opting into the mix-in currently costs a capability.
- Or, if the intent is to reject it unconditionally for now, drop the second clause so the message stops promising it.
I have not measured that the full-declaration case is actually safe end to end, so treat the first option as a proposal rather than a verified fix — but the message/code mismatch stands on its own either way.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,683 @@
+/*
+ * 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, TableInfo}
+import org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity, reference}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, Spark narrows the update-row projection
+ * (LogicalWriteInfo.updateSchema()) to contain only the declared columns rather than
+ * the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends RowLevelOperationSuiteBase {
Review Comment:
**Finding 25.** This suite extends `RowLevelOperationSuiteBase` directly, and so does `GroupBasedColumnUpdateTableSuite` — but the existing UPDATE battery lives one level down, in `UpdateTableSuiteBase` (`:31`) and `DeltaBasedUpdateTableSuiteBase` (`:22`). So none of it runs against the narrow path.
That matters more here than it would for a smaller change, because `SupportsColumnUpdates` is not a tweak to the existing rewrite — it is a second set of builders alongside it: `buildNarrowReplaceDataUpdateProjection`, `buildColumnUpdateProjection`, `buildNarrowDeletesAndInserts`, `buildNarrowRelationWithAttrs`, `computeNarrowReadAttrs`, plus `dataAttrsResolved` / `updateResolved` in `v2Commands.scala`. Everything the battery covers for the wide path — aliases, `DEFAULT` values, char/varchar padding, `EXISTS` / `NOT IN` / correlated subquery conditions, transactions, explain, nested-field assignment — is unexercised against those builders. The two new suites test the narrowing *mechanism* thoroughly and the UPDATE *semantics* barely.
The cheap version is to make these two suites extend the corresponding battery base with the fixture property set, so the existing tests run twice — once wide, once narrow. Wiring that up locally suggests the large majority pass as-is, with the failures tracing to fixture limitations rather than the rewrite (finding 26 is one of them; another is that the fixture's `doCommit` reconstructs rows by indexing the current schema, which breaks for rows written before an `ALTER TABLE ADD COLUMN`). Those are worth fixing on the fixture regardless, since they are what currently makes the battery unreusable.
If that is too much for this PR, it is a reasonable follow-up — but it should be a named follow-up rather than left implicit, because the narrow builders are where the next regression will land.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,68 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * Currently honored only for UPDATE. DELETE and MERGE ignore this interface: those operations
+ * receive full-width rows and {@link LogicalWriteInfo#updateSchema()} is absent for them.
+ *
+ * @since 4.3.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * The returned columns become the schema of updated, copied, and reinserted rows, in declared
+ * order. Implementations must include every column they want to receive (typically the columns
Review Comment:
**Finding 20.** This says "Implementations must include every column they want to receive", which reads as permissive — but `RewriteUpdateTable` now enforces four hard requirements, and none of them are stated on the two methods that carry them:
1. every column reported by `RowLevelOperationInfo.updatedColumns()` (`REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS`);
2. every row-ID column, when `representUpdateAsDeleteAndInsert()` is `true` (`SPLIT_UPDATE_ROW_ID_NOT_DECLARED`);
3. no overlap with `scanOnlyDataAttributes()` (`REQUIRED_DATA_ATTRIBUTES_OVERLAP_SCAN_ONLY_ATTRIBUTES`);
4. every partition-source column, across the two methods (`REQUIRED_DATA_ATTRIBUTES_MISSING_PARTITION_COLUMNS`).
A connector author currently discovers each one by being rejected. Requirement 4 in particular is worth a sentence of its own, because it bounds what the feature can deliver: on a partitioned table the narrow scan always has to carry the partition columns, so narrowing can never get below them. I checked that this is a real constraint rather than defensive over-strictness — with the three `validatePartitionAttrsDeclared` call sites disabled locally, the same UPDATE fails at write resolution instead, with `AnalysisException: Unable to resolve dep given [pk,id,_partition,index]`. So the guard is converting an obscure downstream failure into a clear one, which is the right call; it just needs to be predictable from the javadoc.
One related doc point while you are here: `DataWriter#writeUpdate`'s javadoc (`DataWriter.java:113`, `:136`) and `DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED` both say a connector mixing in `SupportsColumnUpdates` must override `writeUpdate`. That holds for group-based operations only — `writeUpdate` is reached from `ReplaceDataExec`'s writing tasks, so a `SupportsDelta` connector receives narrow rows through `DeltaWriter.update` / `reinsert` and never sees `writeUpdate` at all. Worth qualifying both to group-based.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -253,6 +304,382 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends only the
+ // declared + assigned columns in the row projection instead of the full row schema. The base
+ // class composes its required-attrs set as `pk` (the row-lookup key) plus whatever columns
+ // Spark reports as being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil)
+ extends DeltaBasedOperation(command, CaseInsensitiveStringMap.empty())
Review Comment:
**Finding 26.** `DeltaBasedColumnUpdateOperation` hardcodes `CaseInsensitiveStringMap.empty()` instead of forwarding `info.options`, so every column-update fixture loses the row-level operation options.
The consequence is invisible in this PR's own suites because they never assert on options — but it means no column-update test can observe whether the dynamic-options plumbing from SPARK-57681 still reaches a connector that opts into narrowing. Pointing the existing dynamic-options tests at this fixture fails them on a null option value, which is the same blind spot as finding 25 and has the same one-line cause:
```scala
extends DeltaBasedOperation(command, info.options)
```
That needs `info` threaded in (or the options passed at the construction sites in `newRowLevelOperationBuilder`), which is a small change to the fixture rather than to the feature.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -109,6 +120,22 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
relation)
}
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation. Non-existent columns are rejected with an `AnalysisException`.
+ */
+ protected def resolveRequiredDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes
+ if (refs.isEmpty) {
+ throw QueryCompilationErrors.emptyRequiredDataAttributesError(operation.getClass.getName)
+ }
+ V2ExpressionUtils.resolveRefs[AttributeReference](
Review Comment:
**Finding 21.** A connector's declaration is resolved case-insensitively but then keeps the *declared* spelling, and that spelling becomes the narrow scan's output names.
`AttributeSeq.resolve` returns `a.withName(name)` for the requested name, preserving `exprId` (`expressions/package.scala:308`). So `resolveRequiredDataAttrs` can hand back an attribute named `PK` for a table column `pk`. `buildNarrowRelationWithAttrs` (`:98`) then builds the read relation as `dedupAttrs(dataAttrs ++ rowIdAttrs ++ metadataAttrs)` with the renamed attrs **first** — and `dedupAttrs` keys on `exprId`, so the renamed copy wins and the table's own attribute is discarded. Column pruning pushes those names to the scan, which does not have them.
Measured on `774412c` with the `column-update-req-attrs` fixture, table `pk INT NOT NULL, salary INT, dep STRING, id INT`, query `UPDATE t SET salary = salary + 1 WHERE pk = 1` — only the declared letter case differs between arms:
| declared | result |
| --- | --- |
| `pk,salary,dep` | passes; rows `[1,101,hr,10], [2,200,hr,20]`; `updateSchema = struct<pk,salary,dep>` |
| `PK,salary,dep` | `[INTERNAL_ERROR_ATTRIBUTE_NOT_FOUND] Could not find PK#162 in [salary#163,dep#164,_partition#168]` |
| `pk,SALARY,dep` | `[INTERNAL_ERROR_ATTRIBUTE_NOT_FOUND] Could not find SALARY#273 in [pk#272,dep#274,_partition#278]` |
An internal error is the wrong response to a connector's declaration either way, and case-insensitive matching is the behaviour a `NamedReference`-based API should have — `validatePartitionAttrsDeclared` already uses `conf.resolver` for exactly this reason.
The pre-existing `buildRelationWithAttrs` (`:88`) is immune only by accident: it puts `relation.output` first, so `dedupAttrs` drops the renamed copies.
Fix: after resolving, map each attribute back to the relation's own, preserving declared order (order matters — `updateSchema()` and `areCompatible` both depend on it):
```scala
val resolved = V2ExpressionUtils.resolveRefs[AttributeReference](refs, relation)
val byId = relation.output.map(a => a.exprId -> a).toMap
resolved.map(a => byId.getOrElse(a.exprId, a).asInstanceOf[AttributeReference])
```
Worth applying at all three resolution sites rather than just this one — `resolveScanOnlyDataAttrs` (`RewriteUpdateTable.scala:472`) and `RowLevelWrite.projectedDataAttrs` (`v2Commands.scala:375`) resolve the same declarations, and `dataAttrsResolved`'s `areCompatible` compares `inAttr.name == outAttr.name`, so fixing only the read relation would leave the two sides disagreeing on the name.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -109,6 +120,22 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
relation)
}
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation. Non-existent columns are rejected with an `AnalysisException`.
+ */
+ protected def resolveRequiredDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes
+ if (refs.isEmpty) {
+ throw QueryCompilationErrors.emptyRequiredDataAttributesError(operation.getClass.getName)
Review Comment:
**Finding 24.** This checks for empty but not for duplicates, and a duplicate survives all the way into the public `updateSchema()`.
Measured on `774412c`: `column-update-req-attrs = "pk,pk,salary"` gives
```
updateSchema = Optional[StructType(StructField(pk,IntegerType,false),
StructField(pk,IntegerType,false),
StructField(salary,IntegerType,true))]
```
and the UPDATE succeeds with correct data — the row happens to be written correctly because the projection resolves ordinals positionally. The hazard is on the connector side: `updateSchema().fieldNames.zipWithIndex.toMap`, or anything else keyed by field name, silently collapses two columns into one and writes the wrong ordinal. A `StructType` with a repeated field name is not something a connector should have to defend against.
Rejecting it belongs next to the empty check, since both are "the declaration itself is malformed":
```scala
val dupes = refs.groupBy(_.describe()).filter(_._2.size > 1).keys.toSeq
if (dupes.nonEmpty) {
throw QueryCompilationErrors.<duplicateRequiredDataAttributesError>(
operation.getClass.getName, dupes)
}
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3786105284
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, any required metadata attrs and optionally connector
Review Comment:
It is unclear what `optionally connector declared attrs` means in this case. I'd opt to simplify.
```
// resolve all needed attrs (e.g. row ID and required data / metadata attrs)
// resolve all needed attrs (e.g. row ID or any required metadata attrs)
```
Or similar
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3785080596
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java:
##########
@@ -37,4 +38,15 @@ public interface RowLevelOperationInfo {
* Returns the row-level SQL command (e.g. DELETE, UPDATE, MERGE).
*/
Command command();
+
+ /**
+ * Returns the columns being updated by this operation. Currently populated only for UPDATE;
+ * DELETE and MERGE report an empty array.
+ * <p>
+ * Nested struct field updates are reported at root-column granularity
Review Comment:
Is this a temp limitation to simplify implementation?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3634003715
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -223,19 +274,39 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
None
}
- ReplaceDataProjections(rowProjection, metadataProjection)
+ ReplaceDataProjections(rowProjection, updateRowProjection, metadataProjection)
}
protected def buildWriteDeltaProjections(
plan: LogicalPlan,
rowAttrs: Seq[Attribute],
rowIdAttrs: Seq[Attribute],
- metadataAttrs: Seq[Attribute]): WriteDeltaProjections = {
+ metadataAttrs: Seq[Attribute],
+ updateRowAttrs: Seq[Attribute] = Nil): WriteDeltaProjections = {
val outputs = extractOutputs(plan)
- val rowProjection = if (rowAttrs.nonEmpty) {
+ // Always produce Some(rowProjection) even for empty rowAttrs (identity-only column updates).
+ // When updateRowAttrs is non-empty, the row projection covers INSERT-shaped rows only and a
+ // separate updateRowProjection handles UPDATE/COPY/REINSERT-shaped rows.
+ val rowProjection = if (updateRowAttrs.nonEmpty) {
+ val outputsForInsert = filterOutputs(outputs,
+ OPERATIONS_WITH_ROW -- Set(UPDATE_OPERATION, COPY_OPERATION, REINSERT_OPERATION))
+ if (outputsForInsert.isEmpty) {
+ Some(ProjectingInternalRow(StructType(Nil), Nil))
+ } else {
+ Some(newLazyProjection(plan, outputsForInsert, rowAttrs))
+ }
+ } else if (rowAttrs.nonEmpty) {
val outputsWithRow = filterOutputs(outputs, OPERATIONS_WITH_ROW)
Some(newLazyProjection(plan, outputsWithRow, rowAttrs))
+ } else {
+ Some(ProjectingInternalRow(StructType(Nil), Nil))
Review Comment:
Could we keep the last branch as `None` and adjust the comment?
```suggestion
None
```
This `else` branch is unreachable for the column-update path, so the change here is an unnecessary behavior change to DELETE:
- The comment says the `Some(...)` is needed for identity-only column updates, but that
case never reaches this branch: an identity-only UPDATE on a `SupportsColumnUpdates`
connector still has non-empty `connectorDataAttrs` (an empty `requiredDataAttributes()`
is already rejected with `EMPTY_REQUIRED_DATA_ATTRIBUTES`), so it always takes the first
branch (`updateRowAttrs.nonEmpty`).
- The only caller that reaches this branch is `RewriteDeleteFromTable`, which passes
`rowAttrs = Nil` and no `updateRowAttrs`. For DELETE, this changes `rowProjection` from
`None` to `Some(empty projection)`. The current consumers happen to be equivalent
(`V2Writes` uses `.map(_.schema).getOrElse(StructType(Nil))` and `DeltaWritingSparkTask`
uses `.orNull`), so nothing breaks today, but it silently alters the DELETE plan shape
for no benefit and widens the blast radius of this PR.
Reverting this branch to `None` keeps the DELETE path byte-for-byte identical to master.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4068060968
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -109,6 +120,22 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
relation)
}
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation. Non-existent columns are rejected with an `AnalysisException`.
+ */
+ protected def resolveRequiredDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes
+ if (refs.isEmpty) {
Review Comment:
Done. Added a new error type when nested columns are required and added a test for it.
##########
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(
+ operation: RowLevelOperation,
+ relation: DataSourceV2Relation,
+ connectorDataAttrs: Seq[AttributeReference],
+ scanOnlyDataAttrs: Seq[AttributeReference]): Unit = {
+ val partitionRefNames = relation.table.partitioning().toImmutableArraySeq
+ .flatMap(_.references.toImmutableArraySeq)
+ .map(_.fieldNames.head)
+ .toSet
+ val declared = AttributeSet(connectorDataAttrs ++ scanOnlyDataAttrs)
+ val undeclared = relation.output
+ .filter(a => partitionRefNames.exists(name => conf.resolver(name, a.name)))
+ .filterNot(declared.contains).map(_.name)
+ if (undeclared.nonEmpty) {
+ throw QueryCompilationErrors.requiredDataAttributesMissingPartitionColumnsError(
+ operation.getClass.getName, undeclared)
+ }
+ }
+
+ /**
+ * 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:
This has been fixed in an earlier commit 33b3f071352
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -228,4 +371,146 @@ 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, validating that every assigned column is covered. Returns `Nil` otherwise.
+ */
+ private def resolveConnectorDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation,
+ assignments: Seq[Assignment]): Seq[AttributeReference] = operation match {
+ case scu: SupportsColumnUpdates =>
+ val attrs = resolveRequiredDataAttrs(relation, scu)
+ validateUpdatedColumnsSubset(scu, assignments, attrs)
+ attrs
+ case _ => Nil
+ }
+
+ /**
+ * Computes the narrow set of data columns that must be present in the scan for a column-update
+ * write: connector-declared attrs (`requiredDataAttributes()`), unioned with any table columns
+ * referenced by non-identity assignment RHS expressions and the operation condition.
+ */
+ 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)
Review Comment:
Good catch. For constraints I did both read side and the write side changes to include the constraint columns. They are discarded by the update projects before the write so it should still pass narrow writes to the connector.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -253,6 +304,382 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends only the
+ // declared + assigned columns in the row projection instead of the full row schema. The base
+ // class composes its required-attrs set as `pk` (the row-lookup key) plus whatever columns
+ // Spark reports as being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil)
+ extends DeltaBasedOperation(command, CaseInsensitiveStringMap.empty())
Review Comment:
Done.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,57 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * Currently honored only for UPDATE. DELETE and MERGE ignore this interface: those operations
+ * receive full-width rows and {@link LogicalWriteInfo#updateSchema()} is absent for them.
+ *
+ * @since 4.3.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * The returned columns become the schema of updated, copied, and reinserted rows, in declared
+ * order. Implementations must include every column they want to receive (typically the columns
+ * reported by {@link RowLevelOperationInfo#updatedColumns()} plus any columns needed for row
+ * lookup or routing, e.g. a primary key).
+ * <p>
+ * If any of the columns from {@link RowLevelOperationInfo#updatedColumns()} are
+ * missing, an analysis exception is thrown.
+ * <p>
+ * For updates on nested fields such as {@code SET t.s.c1 = -1} the connector should declare the
+ * root struct column {@code s} rather than any nested field.
+ * <p>
+ * This also covers columns needed only for planning, e.g. resolving the table's partitioning
Review Comment:
Thanks for testing this. I like option 1 but I think that change belongs to it is own PR as it touches code unrelated to this PR. If you agree, I will create a separate JIRA for this. I would also defer adding the test that skips adding the partition column since it will fail now.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +84,51 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated, copied, or reinserted record with metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must override this
+ * method when mixing in {@link SupportsColumnUpdates}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but does not override this method.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T metadata, T record) throws IOException {
+ throw new SparkUnsupportedOperationException(
+ "DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED",
+ Map.of("class", getClass().getName()));
+ }
Review Comment:
I took the suggestion and added java docs to reflect this.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,683 @@
+/*
+ * 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, TableInfo}
+import org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity, reference}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, Spark narrows the update-row projection
+ * (LogicalWriteInfo.updateSchema()) to contain only the declared columns rather than
+ * the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends RowLevelOperationSuiteBase {
Review Comment:
Good call. Just to see what it would break, I made changes. `DeltaBasedColumnUpdateTableSuite` now extends `DeltaBasedUpdateTableSuiteBase`, `GroupBasedColumnUpdateTableSuite` now extends `UpdateTableSuiteBase`.
I found that only 1 test in each suit failed and it was due to the default values. So I fixed the `doCommit` schema-indexing bug in all three affected spots. Each now checks `i < base.numFields` and falls back to `ResolveDefaultColumns.getExistenceDefaultValue(field)` for columns added after a row was written, matching the exact pattern already used elsewhere in `InMemoryBaseTable.scala` for the same problem.
I decided to make this small change in this PR.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -109,6 +120,22 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
relation)
}
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation. Non-existent columns are rejected with an `AnalysisException`.
+ */
+ protected def resolveRequiredDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes
+ if (refs.isEmpty) {
+ throw QueryCompilationErrors.emptyRequiredDataAttributesError(operation.getClass.getName)
Review Comment:
I added duplicate check and associated error message.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,68 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * Currently honored only for UPDATE. DELETE and MERGE ignore this interface: those operations
+ * receive full-width rows and {@link LogicalWriteInfo#updateSchema()} is absent for them.
+ *
+ * @since 4.3.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * The returned columns become the schema of updated, copied, and reinserted rows, in declared
+ * order. Implementations must include every column they want to receive (typically the columns
Review Comment:
Good points. Addressed both.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +84,51 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated, copied, or reinserted record with metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must override this
+ * method when mixing in {@link SupportsColumnUpdates}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but does not override this method.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T metadata, T record) throws IOException {
+ throw new SparkUnsupportedOperationException(
+ "DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED",
+ Map.of("class", getClass().getName()));
+ }
+
+ /**
+ * Writes one updated, copied, or reinserted record without metadata.
+ * <p>
+ * Equivalent to {@link #writeUpdate(Object, Object)} for writers that do not require metadata.
+ * Connectors that mix in {@link SupportsColumnUpdates} receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must override this
+ * method when mixing in {@link SupportsColumnUpdates}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but does not override this method.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T record) throws IOException {
Review Comment:
I have added a test `column-update ReplaceData: connector missing writeUpdate override is rejected` in `GroupBasedColumnUpdateTableSuite.scala`
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,684 @@
+/*
+ * 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, TableInfo}
+import org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity, reference}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, Spark narrows the update-row projection
+ * (LogicalWriteInfo.updateSchema()) to contain only the declared columns rather than
+ * the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends RowLevelOperationSuiteBase {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ test("column-update: rowSchema contains only the single assigned column") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ // info.schema() is empty: UPDATE-only column-update writes have no INSERT-shaped rows,
+ // so nothing flows through the writer's write() path. The narrow update-row layout is
+ // carried by info.updateSchema().
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: rowSchema contains multiple assigned columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false),
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: rowSchema is empty for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ // All assignments are identity, so updatedColumns is empty -- the connector still declares
+ // pk (row lookup) and dep (write-side clustering key), so the narrow update schema is
+ // just [pk, dep].
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Array(PK_FIELD, DEP_FIELD))))
+ }
+
+ test("column-update: row filter condition is orthogonal to column narrowing") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = 'engineering' WHERE pk IN (1, 3)")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: update all rows (no WHERE clause)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("salary", IntegerType, nullable = true)
+ ))))
+ }
+
+ test("column-update: rowSchema excludes identity assignments in a mixed UPDATE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: cross-column assignment is not treated as identity") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = dep, id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: nested struct field update narrows to the root struct column") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("s"),
+ s"expected [s] in updatedColumns (root struct) but got: $updatedNames")
+
+ // info.updateSchema() carries the narrow row layout; info.schema() is the full table.
+ // `dep` is present because the connector also declares it as the write-side clustering key.
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("s"),
+ s"s must be in update schema: $updateSchema")
+ assert(updateSchema.fieldNames.contains("dep"),
+ s"dep must be in update schema: $updateSchema")
+ }
+
+ test("column-update: nested field identity update reports root struct as updated") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = s.c1 WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("s"),
+ s"nested identity is reported as an update at root-column granularity: $updatedNames")
+
+ // Data correctness: the struct is rewritten but with equal values, so rows are unchanged.
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, Row(1, 2), "hr") :: Nil)
+ }
+
+ test("column-update: updatedColumns contains non-identity assigned columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'eng' WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("id", "dep"),
+ s"expected [id, dep] in updatedColumns but got: $updatedNames")
+ }
+
+ test("column-update: updatedColumns excludes identity assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = dep WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("id"),
+ s"expected only [id] in updatedColumns but got: $updatedNames")
+ }
+
+ test("column-update: updatedColumns is empty for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ assert(table.lastUpdatedColumns.isEmpty,
+ s"expected empty updatedColumns but got: ${table.lastUpdatedColumns.mkString(", ")}")
+ }
+
+ test("column-update: updatedColumns is empty for DELETE (Javadoc contract)") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"DELETE FROM $tableNameAsString WHERE dep = 'hr'")
+
+ assert(table.lastUpdatedColumns.isEmpty,
+ s"DELETE must pass empty updatedColumns but got: ${table.lastUpdatedColumns.mkString(", ")}")
+ }
+
+ test("column-update: data correctness -- single column update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ // 1 row matched WHERE pk = 1 -> 1 update. MoR emits no COPY rows.
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0, deltaUpdate = true)
+ }
+
+ test("column-update: data correctness -- update all rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 200, "hr") :: Row(2, 400, "software") :: Row(3, 600, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 3, numCopiedRows = 0, deltaUpdate = true)
+ }
+
+ test("column-update: data correctness -- mixed identity and real assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "engineering") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0, deltaUpdate = true)
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-req-attrs", reqAttrs)
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update: requiredDataAttributes - data correctness") {
+ createAndInitTableWithReqAttrs("dep,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively on the row-ID " +
+ "column without adopting the declared spelling") {
+ // The connector declares `PK` (not `pk`); resolution must match case-insensitively but the
+ // narrow scan/write must still use the table's own column name, `pk`, not the connector's
+ // declared spelling -- otherwise column pruning can't find `PK` in the physical scan.
+ createAndInitTableWithReqAttrs("PK,salary,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("pk"),
+ s"update schema must use the table's own spelling `pk`, not `PK`: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("PK"),
+ s"update schema must not adopt the connector's declared spelling `PK`: $updateSchema")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively on an assigned " +
+ "column without adopting the declared spelling") {
+ // Same as above, but the case mismatch is on `salary` (the assigned column) rather than
+ // the row-ID column.
+ createAndInitTableWithReqAttrs("pk,SALARY,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("salary"),
+ s"update schema must use the table's own spelling `salary`, not `SALARY`: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("SALARY"),
+ s"update schema must not adopt the connector's declared spelling `SALARY`: $updateSchema")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: empty requiredDataAttributes throws AnalysisException") {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-empty-req-attrs", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(
+ StructType.fromDDL("pk INT NOT NULL, id INT, dep STRING"))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }""".stripMargin)
+
+ // A connector that mixes in SupportsColumnUpdates but returns an empty
+ // requiredDataAttributes() violates the mix-in contract -- analysis must reject the
+ // operation rather than silently falling through to the wide write path.
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ }
+ assert(ex.getCondition == "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ s"expected COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES but got: ${ex.getCondition}")
+ }
+
+ test("column-update: requiredDataAttributes throws AnalysisException for invalid column") {
+ createAndInitTableWithReqAttrs("nonexistent_col", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ }
+ assert(ex.getMessage.contains("nonexistent_col"),
+ s"Expected error about unresolvable column but got: ${ex.getMessage}")
+ }
+
+ private def createAndInitTableFromInfo(schemaString: String, jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-from-info", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update from-info: write schema is updatedColumns + pk/dep pass-through") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("salary"),
+ s"salary must be in update schema: $updateSchema")
+ assert(updateSchema.fieldNames.contains("pk"),
+ s"pk must be in update schema: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("id"),
+ s"id must not be in update schema: $updateSchema")
+ assert(updateSchema.fieldNames.contains("dep"),
+ s"dep must be in update schema: $updateSchema")
+ }
+
+ test("column-update from-info: pk already in updatedColumns is not duplicated") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ val pkCount = updateSchema.fieldNames.count(_ == "pk")
+ assert(pkCount == 1, s"pk must appear exactly once in update schema: $updateSchema")
+ }
+
+ test("column-update from-info: data correctness") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "id": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ }
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-split", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update split: write schema is narrow (assigned + pk/dep pass-through)") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ // For representUpdateAsDeleteAndInsert connectors, reinsert-tagged rows still route through
+ // writeUpdate() (info.updateSchema() = narrow), not write(). info.schema() stays empty.
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update split: data correctness") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, -1, "hr") :: Nil)
+ }
+
+ test("column-update split: row-ID reassignment on narrow write is rejected") {
+ // `extra` is neither declared, updated, nor referenced, so requiredDataAttributes()
+ // ([pk, dep, salary]) doesn't cover every table column -- the write stays genuinely narrow
+ // and reassigning the row ID must be rejected.
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING, extra STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr", "extra": "x" }
+ |{ "pk": 2, "salary": 200, "dep": "software", "extra": "y" }
+ |""".stripMargin)
+
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+ }
+ assert(ex.getCondition == "COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT",
+ s"expected COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT but got: ${ex.getCondition}")
+ }
+
+ test("column-update split: row-ID reassignment is allowed when requiredDataAttributes " +
+ "covers every table column") {
+ // Finding 22: the SPLIT_UPDATE_ROW_ID_REASSIGNMENT message advertises a second remedy --
+ // "include every table column in requiredDataAttributes()" -- but the guard never checked
+ // for it. When the declaration covers the whole relation, the REINSERT payload already is
+ // the full row with the new row-ID value, and the DELETE half still carries the original
+ // row-ID via newLazyRowIdProjection, so reassignment is safe and must be allowed.
+ createAndInitTableSplitWithReqAttrs("pk,salary,dep",
+ "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(2, 200, "software") :: Row(11, -1, "hr") :: Nil)
+ }
+
+ private def createAndInitTableSplitWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-split-req-attrs", reqAttrs)
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update split: undeclared row-ID column on narrow write is rejected") {
+ // The connector's requiredDataAttributes() excludes `pk` (the row ID). The REINSERT
+ // payload for the split-update path is projected down to requiredDataAttributes(), so
+ // without `pk` there the reinserted row would have no identity for the connector to
+ // place it by. This must be rejected at analysis time, not discovered at write time.
+ createAndInitTableSplitMissingRowId("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+ }
+ assert(ex.getCondition == "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED",
+ s"expected COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED but got: ${ex.getCondition}")
+ assert(ex.getMessage.contains("pk"),
+ s"error message must name the undeclared row ID column `pk`: ${ex.getMessage}")
+ }
+
+ private def createAndInitTableSplitMissingRowId(
+ schemaString: String, jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-split-missing-row-id", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ // ---------------------------------------------------------------------------
+ // Scan-narrowing tests: verify that when the connector opts into column updates,
+ // the physical scan reads only the columns actually needed by the operation.
+ //
+ // Assertions use `checkLastScanExcludes` (proves narrowing) and `checkLastScanIncludes`
+ // (proves transparent widening) rather than exact-set matches -- ColumnPruning is
+ // free to further tighten the scan beyond our analysis-time narrow set (e.g. when an
+ // assignment RHS is a literal, the assigned column doesn't need to be read).
+ // ---------------------------------------------------------------------------
+
+ test("column-update: scan excludes columns outside connector-declared + cond + RHS refs") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, extra STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr", "extra": "x" }
+ |{ "pk": 2, "id": 2, "dep": "software", "extra": "y" }
+ |""".stripMargin)
+
+ // requiredDataAttributes = [pk, id]; cond refs [pk] only; RHS is a literal.
+ // `dep` is scan-only declared (needed for scan-side partitioning resolution) so it must
Review Comment:
Done. I fixed all the mentioned places.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-5737269368
I rebased it with latest master.
@szehon-ho would you like to take a look as well since you reviewed the SPIP?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-5999317768
I merged the following. Do you think you can rebase and update this PR, @anuragmantri ?
- https://github.com/apache/spark/pull/58979
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4341122600
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,94 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, updated, copied, and reinserted
+ * rows follow it: group-based operations receive them through
+ * {@link DataWriter#writeColumnUpdate(Object, Object)} or
+ * {@link DataWriter#writeColumnUpdate(Object)}, and operations that mix in {@link SupportsDelta}
+ * through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert}. Inserted rows always follow
+ * {@link LogicalWriteInfo#schema()} and arrive through {@link DataWriter#write(Object)} and
+ * {@link DeltaWriter#insert}. When {@link LogicalWriteInfo#columnUpdateSchema()} is absent, every
+ * row follows {@link LogicalWriteInfo#schema()}, as for any other operation. Connectors must
+ * decide which layout to expect by whether {@link LogicalWriteInfo#columnUpdateSchema()} is
+ * present, not by the command. A builder that wants full-width rows for some command can check
+ * {@link RowLevelOperationInfo#command()} and build an operation without this mix-in. Currently
+ * Spark narrows only UPDATE.
+ * <p>
+ * The scan builder returned by {@link #newScanBuilder} should implement
+ * {@link org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns}, so that the scan
+ * does not read columns the command neither references nor declares. Otherwise the scan reads
+ * every column.
+ *
+ * @since 4.4.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * When Spark narrows the write, the returned columns become
+ * {@link LogicalWriteInfo#columnUpdateSchema()}, in declared order. Implementations must include
+ * every column they want to receive: every column reported by
+ * {@link RowLevelOperationInfo#updatedColumns()}, plus any columns needed for row lookup or
+ * routing, e.g. a primary key. Columns that are not declared are absent from the rows the
Review Comment:
Done in 22332b3. The doc now says that an operation without `SupportsDelta` should locate original rows through metadata columns preserved on update, such as a file and a row position, and, if it locates rows by a data column, build an operation without the mix-in when `updatedColumns()` contains that column.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala:
##########
@@ -109,22 +111,75 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper {
WriteToDataSourceV2(Some(r), microBatchWrite, newQuery, customMetrics)
case rd @ ReplaceData(r: DataSourceV2Relation, _, query, _, projections, _, None) =>
- val rowSchema = projections.rowProjection.schema
+ val (rowSchema, columnUpdateSchema) =
+ writeSchemas(projections.rowProjection.schema, hasNarrowRows(rd))
val metadataSchema = projections.metadataProjection.map(_.schema)
val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap)
- val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema)
+ val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema,
+ columnUpdateSchema)
val write = writeBuilder.build()
+ validateColumnUpdateWrite(rd, write)
val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, r.funCatalog)
rd.copy(write = Some(write), query = newQuery)
case wd @ WriteDelta(r: DataSourceV2Relation, _, query, _, projections, _, None) =>
val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap)
- val deltaWriteBuilder = newDeltaWriteBuilder(r.table, writeOptions, projections)
+ val deltaWriteBuilder =
+ newDeltaWriteBuilder(r.table, writeOptions, projections, hasNarrowRows(wd))
val deltaWrite = deltaWriteBuilder.build()
+ validateColumnUpdateWrite(wd, deltaWrite)
val newQuery = DistributionAndOrderingUtils.prepareQuery(deltaWrite, query, r.funCatalog)
wd.copy(write = Some(deltaWrite), query = newQuery)
}
+ /**
+ * Whether a row-level write delivers narrow rows for updated, copied, and reinserted records.
+ * This must match the writes whose write relation RewriteUpdateTable narrows, even when the
+ * declared columns cover the whole table.
+ */
+ private[v2] def hasNarrowRows(write: RowLevelWrite): Boolean = {
+ write.operation.isInstanceOf[SupportsColumnUpdates] && write.operation.command == UPDATE
Review Comment:
Good idea. Done in 22332b3. `V2Writes.validateWriteRelation` runs before building every `ReplaceData` and `WriteDelta` write. The Javadoc of `requiredDataAttributes()` also says it must return the same columns on every call.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala:
##########
@@ -109,22 +111,75 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper {
WriteToDataSourceV2(Some(r), microBatchWrite, newQuery, customMetrics)
case rd @ ReplaceData(r: DataSourceV2Relation, _, query, _, projections, _, None) =>
- val rowSchema = projections.rowProjection.schema
+ val (rowSchema, columnUpdateSchema) =
+ writeSchemas(projections.rowProjection.schema, hasNarrowRows(rd))
val metadataSchema = projections.metadataProjection.map(_.schema)
val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap)
- val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema)
+ val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema,
+ columnUpdateSchema)
val write = writeBuilder.build()
+ validateColumnUpdateWrite(rd, write)
val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, r.funCatalog)
rd.copy(write = Some(write), query = newQuery)
case wd @ WriteDelta(r: DataSourceV2Relation, _, query, _, projections, _, None) =>
val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap)
- val deltaWriteBuilder = newDeltaWriteBuilder(r.table, writeOptions, projections)
+ val deltaWriteBuilder =
+ newDeltaWriteBuilder(r.table, writeOptions, projections, hasNarrowRows(wd))
val deltaWrite = deltaWriteBuilder.build()
+ validateColumnUpdateWrite(wd, deltaWrite)
val newQuery = DistributionAndOrderingUtils.prepareQuery(deltaWrite, query, r.funCatalog)
wd.copy(write = Some(deltaWrite), query = newQuery)
}
+ /**
+ * Whether a row-level write delivers narrow rows for updated, copied, and reinserted records.
+ * This must match the writes whose write relation RewriteUpdateTable narrows, even when the
+ * declared columns cover the whole table.
+ */
+ private[v2] def hasNarrowRows(write: RowLevelWrite): Boolean = {
+ write.operation.isInstanceOf[SupportsColumnUpdates] && write.operation.command == UPDATE
+ }
+
+ /**
+ * Checks that the distribution and ordering of a write that delivers narrow rows reference only
+ * columns the write reads. Other columns stay in the query only until column pruning removes
+ * them, so they are rejected here whether or not pruning ran.
+ */
+ private def validateColumnUpdateWrite(rowLevelWrite: RowLevelWrite, write: Write): Unit = {
+ write match {
+ case w: RequiresDistributionAndOrdering if hasNarrowRows(rowLevelWrite) =>
+ val distribution = w.requiredDistribution match {
+ case d: ClusteredDistribution => d.clustering.toImmutableArraySeq
+ case d: OrderedDistribution => d.ordering.toImmutableArraySeq
+ case _ => Seq.empty
+ }
+ val readNames = rowLevelWrite.references.toSeq.map(_.name)
+ val unread = (distribution ++ w.requiredOrdering.toImmutableArraySeq)
+ .flatMap(_.references.toImmutableArraySeq)
Review Comment:
Done in 22332b3, as suggested. The helper moved to `V2ExpressionUtils`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/RowLevelOperationInfoImpl.scala:
##########
@@ -17,9 +17,15 @@
package org.apache.spark.sql.connector.write
+import org.apache.spark.sql.connector.expressions.NamedReference
import org.apache.spark.sql.connector.write.RowLevelOperation.Command
import org.apache.spark.sql.util.CaseInsensitiveStringMap
private[sql] case class RowLevelOperationInfoImpl(
command: Command,
- options: CaseInsensitiveStringMap) extends RowLevelOperationInfo
+ options: CaseInsensitiveStringMap,
+ private val updatedCols: Seq[NamedReference])
Review Comment:
Done in 22332b3. `RowLevelOperationTable` keeps the info in a second parameter list and `withTable(newTable)` rebuilds the operation for another table from the same info. `ResolveBranch` would replace building the info, the operation and the table with `operationTable.withTable(newTable)`.
This is a change on Iceberg either way. @aokolnychyi, could you check this one from Iceberg perspective?
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java:
##########
@@ -37,4 +38,16 @@ public interface RowLevelOperationInfo {
* Returns the row-level SQL command (e.g. DELETE, UPDATE, MERGE).
*/
Command command();
+
+ /**
+ * Returns the columns being updated by this operation. Currently only UPDATE populates it;
+ * other commands report an empty array.
+ * <p>
+ * A column is reported only if it is assigned a new value, so identity assignments such as
Review Comment:
Thanks for this comment. I learned a new thing here. I documented this in 22332b3.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/LogicalWriteInfo.java:
##########
@@ -65,4 +69,18 @@ default Optional<StructType> metadataSchema() {
throw new SparkUnsupportedOperationException(
"DATA_SOURCE_METADATA_SCHEMA_NOT_IMPLEMENTED", Map.of("class", getClass().getName()));
}
+
+ /**
+ * the schema of updated, copied, and reinserted rows from Spark to data source in a
+ * column-level update. Present when the operation mixes in {@link SupportsColumnUpdates} and
+ * Spark delivers these rows with the columns of
+ * {@link SupportsColumnUpdates#requiredDataAttributes()}, which currently happens only for
+ * UPDATE. It covers every table column if every column is declared. When present,
+ * {@link #schema()} covers only newly inserted rows.
+ *
+ * @since 4.4.0
+ */
+ default Optional<StructType> columnUpdateSchema() {
Review Comment:
Done in 22332b3. The default now throws `DATA_SOURCE_COLUMN_UPDATE_SCHEMA_NOT_IMPLEMENTED`. I added a test also to cover this.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +84,60 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated or copied record with metadata in the column update layout.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present for a row-level operation that
+ * does not mix in {@link SupportsDelta}, Spark passes updated and copied records to this method
+ * instead of {@link #write(Object, Object)} if the operation returns a non-empty
Review Comment:
Done in 22332b3.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -253,6 +326,24 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
WriteDeltaProjections(rowProjection, rowIdProjection, metadataProjection)
}
+ /**
+ * Moves the query columns that no projection reads after the ones that are read, and returns
+ * the reordered query with a map from each old ordinal to the new one. The optimizer prunes the
+ * query columns after the last one the write reads, which then never changes an ordinal.
+ */
+ protected def orderReadColumnsFirst(
Review Comment:
I like this suggestion, it's cleaner than my reordering. I implemented this in 22332b3. This now applies to every row-level write, so DELETE and MERGE queries can also drop columns their projections do not read. To make sure this does not break anything, I ran the DSv2 suites and they all pass.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -228,4 +240,114 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Returns the output of the write relation, which is the layout of the data rows the connector
+ * receives. If the operation opts into column updates, this is the connector's
+ * `requiredDataAttributes()`, validated against the assignments. Otherwise it is the full
+ * relation output.
+ */
+ private def resolveWriteAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation,
+ assignments: Seq[Assignment]): Seq[AttributeReference] = operation match {
+ case op: SupportsColumnUpdates with SupportsDelta if op.representUpdateAsDeleteAndInsert =>
+ val attrs = resolveDeclaredAttrs(relation, op, assignments)
+ val rowIdAttrs = resolveRowIdAttrs(relation, op)
+ validateNoRowIdReassignment(op, relation, attrs, assignments, rowIdAttrs)
+ validateRowIdDeclared(op, attrs, rowIdAttrs, resolveRequiredMetadataAttrs(relation, op))
+ attrs
+ case op: SupportsColumnUpdates =>
+ resolveDeclaredAttrs(relation, op, assignments)
+ case _ =>
+ relation.output
+ }
+
+ /**
+ * Resolves the connector's `requiredDataAttributes()` and checks that it covers every assigned
+ * column.
+ */
+ private def resolveDeclaredAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates,
+ assignments: Seq[Assignment]): Seq[AttributeReference] = {
+ val attrs = resolveRequiredDataAttrs(relation, operation)
+ validateUpdatedColumnsSubset(operation, assignments, attrs)
+ attrs
+ }
+
+ /**
+ * 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 {
Review Comment:
Done in 22332b3.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -373,6 +904,21 @@ private class DeltaBufferWriter(schema: StructType) extends BufferWriter(schema)
}
object InMemoryRowLevelOperationTable {
+ // Global holder for the scan schema of the last row-level operation. See the class-level
+ // `lastScanSchema` def for why this is a companion-object field rather than a per-instance one.
+ // Tests that assert on scan pruning should reset this before running the operation under test
+ // (see `RowLevelOperationSuiteBase.beforeEach`).
+ private[catalog] val lastScanSchemaRef = new AtomicReference[StructType]()
Review Comment:
Done in 22332b3. The JVM-global recorder is gone. `executeAndCheckScan` now checks every scan of the target table in the statement's own executed plan and fails if there is none, and it and `executeAndCheckScans` compare field names exactly; `executeAndCheckScanIncludes` and `executeAndCheckScanExcludes` cover the include and exclude checks. `lastUpdatedColumns` starts as null, so `checkLastUpdatedColumns()` fails when no operation was built.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala:
##########
@@ -4638,6 +4638,92 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat
messageParameters = Map("nullableRowIdAttrs" -> nullableRowIdAttrs.mkString(", ")))
}
+ def emptyRequiredDataAttributesError(connectorClass: String): Throwable = {
+ new AnalysisException(
+ errorClass = "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ messageParameters = Map("connector" -> connectorClass))
+ }
+
+ def duplicateRequiredDataAttributeError(
+ connectorClass: String,
+ duplicateAttributes: Seq[String]): Throwable = {
+ new AnalysisException(
+ errorClass = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ messageParameters = Map(
+ "connector" -> connectorClass,
+ "duplicateAttributes" -> duplicateAttributes.mkString("[", ", ", "]")))
Review Comment:
Done in 22332b3. All `COLUMN_UPDATE_*` helpers take name parts and format them with `toSQLId`
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -228,4 +240,114 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Returns the output of the write relation, which is the layout of the data rows the connector
+ * receives. If the operation opts into column updates, this is the connector's
+ * `requiredDataAttributes()`, validated against the assignments. Otherwise it is the full
+ * relation output.
+ */
+ private def resolveWriteAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation,
+ assignments: Seq[Assignment]): Seq[AttributeReference] = operation match {
+ case op: SupportsColumnUpdates with SupportsDelta if op.representUpdateAsDeleteAndInsert =>
+ val attrs = resolveDeclaredAttrs(relation, op, assignments)
+ val rowIdAttrs = resolveRowIdAttrs(relation, op)
+ validateNoRowIdReassignment(op, relation, attrs, assignments, rowIdAttrs)
+ validateRowIdDeclared(op, attrs, rowIdAttrs, resolveRequiredMetadataAttrs(relation, op))
+ attrs
+ case op: SupportsColumnUpdates =>
+ resolveDeclaredAttrs(relation, op, assignments)
+ case _ =>
+ relation.output
+ }
+
+ /**
+ * Resolves the connector's `requiredDataAttributes()` and checks that it covers every assigned
+ * column.
+ */
+ private def resolveDeclaredAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates,
+ assignments: Seq[Assignment]): Seq[AttributeReference] = {
+ val attrs = resolveRequiredDataAttrs(relation, operation)
+ validateUpdatedColumnsSubset(operation, assignments, attrs)
+ attrs
+ }
+
+ /**
+ * 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)
+ }
+ }
+
+ /**
+ * For connectors that opt into narrow column updates AND represent UPDATE as delete + insert,
+ * reject reassignment of any row-ID column, unless `requiredDataAttributes()` covers every
+ * column in the relation. In that case the REINSERT payload already is the full row with the
+ * new row-ID value, and the DELETE half still carries the original row-ID via
+ * `newLazyRowIdProjection`, so reassignment is safe. Otherwise the REINSERT path has no
+ * row-ID channel to reconstruct columns outside `requiredDataAttributes()`.
+ */
+ private def validateNoRowIdReassignment(
+ operation: RowLevelOperation,
+ relation: DataSourceV2Relation,
+ connectorDataAttrs: Seq[AttributeReference],
+ assignments: Seq[Assignment],
+ rowIdAttrs: Seq[Attribute]): Unit = {
+ val declaredIds = connectorDataAttrs.map(_.exprId).toSet
+ if (relation.output.forall(a => declaredIds.contains(a.exprId))) {
Review Comment:
Done in 22332b3: the exemption compares only the data columns.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -258,6 +354,407 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep` (an unconditionally declared base column) plus whatever columns
+ // Spark reports as being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
Review Comment:
Done in 22332b3.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -258,6 +354,407 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep` (an unconditionally declared base column) plus whatever columns
+ // Spark reports as being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = {
+ val base = Seq(FieldReference("pk"), FieldReference("dep"))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ private val clusterRef: NamedReference = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY))
+ .map(FieldReference(_))
+ .getOrElse(PARTITION_COLUMN_REF)
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.columnUpdateSchema().isPresent) {
+ capturedInfo.columnUpdateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(Array(clusterRef))
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ Array[SortOrder](
+ LogicalExpressions.sort(
+ clusterRef,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering())
+ )
+ }
+
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For column-update writes, updated rows contain only the declared columns
+ // (columnUpdateSchema from LogicalWriteInfo). We expand each row to the full
+ // table schema by overlaying them on the base row found by pk.
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+
+ val mergedData = newData.map { buf =>
+ val merged = new BufferedRows(buf.key, schema)
+ val updateOpName = UTF8String.fromString(Update.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == updateOpName) {
+ val pk = logRow.getInt(1)
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val baseRow = currentRowByPk(pk)
+ val fullRow = overlayNarrowRow(baseRow, narrowRow, writeFieldIdx)
+ merged.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ merged.rows.append(insertRow.copy())
+ }
+ }
+ merged
+ }
+
+ withDeletes(newData)
+ withData(mergedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ }
+ }
+ }
+ }
+
+ class DeltaBasedColumnUpdateOperationWithReqAttrs(
+ command: Command,
+ reqCols: Array[String],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, options = options) {
+ override def requiredDataAttributes(): Array[NamedReference] = reqCols.map(FieldReference(_))
+ }
+
+ class DeltaBasedColumnUpdateSplitOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols, options) {
+ override def representUpdateAsDeleteAndInsert(): Boolean = true
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
Review Comment:
Done in 22332b3. The split operation now reuses the parent's write builder and overrides only the writer factory and the commit, so it honors `column-update-cluster-by`. The dynamic-options copies are dropped, as the inherited SPARK-57681 tests already run on both tables, and `target_deps` is wrapped in `withTempView`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala:
##########
@@ -92,6 +95,38 @@ object GroupBasedRowLevelOperationScanPlanning extends Rule[LogicalPlan] with Pr
}
}
+ /**
+ * Returns the columns of `relation` that the projections and filters directly above each read
+ * of the table use, as the rest of the query sees only their output. Only whole columns are
+ * pruned, as the query is not rewritten for a pruned nested schema.
Review Comment:
Done in 22332b3: the scan-builder paragraph of `SupportsColumnUpdates` now says that for operations without `SupportsDelta` only top-level columns are pruned, so a struct referenced through a nested field is read in full,. I will file the nested pruning for copy-on-write scans as a follow-up after this PR.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,754 @@
+/*
+ * 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.{SparkRuntimeException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.{InMemoryBaseTable, InMemoryRowLevelOperationTable}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+
+class GroupBasedColumnUpdateTableSuite extends UpdateTableSuiteBase {
Review Comment:
Done in 22332b3. "update table with expression-based default values" and "update with current_timestamp default value using DEFAULT keyword" now create their tables with the suite's properties, so they run narrow in both column-update suites and on delta tables in the existing delta suites. The first one needed `pk` to be `NOT NULL`, since delta row IDs cannot be nullable; it had only ever run on copy-on-write tables. The fixtures now declare `pk` and `dep` only when the table has them. "update table with recursive CTE" stays on a plain table, since its table has no row ID column.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -258,6 +354,407 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep` (an unconditionally declared base column) plus whatever columns
+ // Spark reports as being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = {
+ val base = Seq(FieldReference("pk"), FieldReference("dep"))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ private val clusterRef: NamedReference = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY))
+ .map(FieldReference(_))
+ .getOrElse(PARTITION_COLUMN_REF)
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.columnUpdateSchema().isPresent) {
+ capturedInfo.columnUpdateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
Review Comment:
Done in 22332b3. With `no-metadata`, the delta column-update write is a plain `DeltaWrite` without distribution or ordering, like the parent, unless `column-update-cluster-by` is set, and `split-updates` on a column-update table now selects the split operation ("column-update: split-updates sends a narrow UPDATE as delete and reinsert"). The new `DeltaBasedNoMetadataColumnUpdateTableSuite` extends `UpdateTableSuiteBase` with the shared tests, so narrow rows go through `DeltaWritingSparkTask`; most of its tests fail when that task's UPDATE branch stops projecting the row.
##########
common/utils/src/main/resources/error/error-conditions.json:
##########
@@ -1253,6 +1253,60 @@
],
"sqlState" : "22003"
},
+ "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates` but declared duplicate column(s) <duplicateAttributes> in `requiredDataAttributes()`. Each column must be declared at most once."
+ ],
+ "sqlState" : "42000"
+ },
+ "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates` but returned an empty array from `requiredDataAttributes()`. Connectors that opt into column-level updates must declare at least one required data attribute."
+ ],
+ "sqlState" : "42000"
+ },
+ "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates` but declared metadata column(s) <metadataAttributes> in `requiredDataAttributes()`. Declare metadata columns in `requiredMetadataAttributes()` instead."
+ ],
+ "sqlState" : "42000"
+ },
+ "COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates` but declared nested field(s) <nestedAttributes> in `requiredDataAttributes()`. Declare the root struct column instead of a nested field."
+ ],
+ "sqlState" : "42000"
+ },
+ "COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates` but its `requiredDataAttributes()` does not cover every column being updated. Missing columns: <missingColumns>. Connectors must include every column reported by `RowLevelOperationInfo.updatedColumns()` in `requiredDataAttributes()`."
+ ],
+ "sqlState" : "42000"
+ },
+ "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates` and represents UPDATE as delete and insert, but row ID column(s) <rowIds> do not reach the reinserted row, which then has no identity for the connector to place it by. Declare each data row ID column in `requiredDataAttributes()`, and return each metadata row ID column from `requiredMetadataAttributes()` with `MetadataColumn.PRESERVE_ON_REINSERT` set."
+ ],
+ "sqlState" : "42000"
+ },
+ "COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates` and represents UPDATE as delete and insert, so UPDATE cannot assign row ID column(s) <rowIds>. The reinserted row carries only the columns in `requiredDataAttributes()`, and with a new row ID it cannot be matched to the row whose other columns the connector must preserve. Do not assign row ID columns, or include every table column in `requiredDataAttributes()`."
+ ],
+ "sqlState" : "42000"
+ },
+ "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates`, but its write requires a distribution or ordering by column(s) <columns> that the column-level UPDATE does not read. Declare each data column in `requiredDataAttributes()` and each metadata column in `requiredMetadataAttributes()`."
Review Comment:
Done in 22332b3
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,999 @@
+/*
+ * 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.SparkRuntimeException
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+import org.apache.spark.sql.connector.catalog.{Delete, InMemoryBaseTable, InMemoryRowLevelOperationTable, Reinsert}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, updated rows contain only the declared columns
+ * (LogicalWriteInfo.columnUpdateSchema()) rather than the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ private val reqAttrsOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateOperationWithReqAttrs].getName
+
+ private val splitOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperation].getName
+
+ private val splitReqAttrsOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithReqAttrs]
+ .getName
+
+ private val splitWithRowIdOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithRowId].getName
+
+ test("column-update: column update schema contains a single assigned column") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: column update schema contains multiple assigned columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false),
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: column update schema is pk and dep for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ // All assignments are identity, so updatedColumns is empty; pk and dep remain because the
+ // connector unconditionally declares them as base columns.
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Array(PK_FIELD, DEP_FIELD))))
+ }
+
+ test("column-update: row filter condition is orthogonal to column narrowing") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = 'engineering' WHERE pk IN (1, 3)")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: update all rows (no WHERE clause)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("salary", IntegerType, nullable = true)
+ ))))
+ }
+
+ test("column-update: column update schema excludes identity assignments in a mixed UPDATE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+
+ test("column-update: nested struct field update narrows to the root struct column") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+
+ checkLastUpdatedColumns("s")
+
+ // `dep` is present because the connector unconditionally declares it as a base column.
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq == Seq("pk", "dep", "s"))
+ }
+
+ test("column-update: nested field identity update reports root struct as updated") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = s.c1 WHERE pk = 1")
+
+ checkLastUpdatedColumns("s")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, Row(1, 2), "hr") :: Nil)
+ }
+
+ test("column-update: updatedColumns is empty for DELETE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"DELETE FROM $tableNameAsString WHERE dep = 'hr'")
+
+ checkLastUpdatedColumns()
+ }
+
+ test("column-update: row-level DELETE falls back to full-width rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"DELETE FROM $tableNameAsString WHERE pk < 2")
+
+ assert(!table.lastWriteInfo.columnUpdateSchema().isPresent,
+ s"DELETE must not report a column update schema: " +
+ s"${table.lastWriteInfo.columnUpdateSchema()}")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(2, 200, "hr") :: Nil)
+ }
+
+ test("column-update: MERGE falls back to full-width rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+ checkMergeFallsBackToFullWidthRows()
+ }
+
+ test("column-update split: MERGE falls back to full-width rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split", "true"))
+ checkMergeFallsBackToFullWidthRows()
+ }
+
+ private def checkMergeFallsBackToFullWidthRows(): Unit = {
+ withTempView("source") {
+ import testImplicits._
+ Seq((1, 150, "hr"), (3, 300, "software")).toDF("pk", "salary", "dep")
+ .createOrReplaceTempView("source")
+
+ sql(
+ s"""MERGE INTO $tableNameAsString t
+ |USING source s
+ |ON t.pk = s.pk
+ |WHEN MATCHED THEN UPDATE SET salary = s.salary
+ |WHEN NOT MATCHED THEN INSERT *
+ |""".stripMargin)
+
+ assert(!table.lastWriteInfo.columnUpdateSchema().isPresent,
+ s"MERGE must not report a column update schema: " +
+ s"${table.lastWriteInfo.columnUpdateSchema()}")
+ checkLastUpdatedColumns()
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 150, "hr") :: Row(2, 200, "software") :: Row(3, 300, "software") :: Nil)
+ }
+ }
+
+ test("column-update: data correctness -- single column update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0)
+ }
+
+ test("column-update: data correctness -- update all rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 200, "hr") :: Row(2, 400, "software") :: Row(3, 600, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 3, numCopiedRows = 0)
+ }
+
+ test("column-update: data correctness -- mixed identity and real assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "engineering") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0)
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-req-attrs", reqAttrs))
+ }
+
+ test("column-update: requiredDataAttributes - data correctness") {
+ createAndInitTableWithReqAttrs("dep,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: undeclared distribution column is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "pk,salary",
+ "column-update-cluster-by", "dep"))
+
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName
+ Seq(Nil, Seq(excludeColumnPruning)).foreach { confs =>
+ withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[dep]"))
+ }
+ }
+ }
+
+ test("column-update: distribution by an undeclared row ID column is allowed") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "salary,dep",
+ "column-update-cluster-by", "pk"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 200, "hr") :: Nil)
+ }
+
+ test("column-update: partition source column not in requiredDataAttributes") {
+ createAndInitTableWithReqAttrs("pk,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastScanExcludes("dep")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively on the row-ID " +
+ "column without adopting the declared spelling") {
+ createAndInitTableWithReqAttrs("PK,salary,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "dep"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+
+ test("column-update: undeclared, differently-cased columns in the RHS and WHERE clause " +
+ "are read with the table's own spelling") {
+ createAndInitTableWithReqAttrs("pk,dep,salary",
+ "pk INT NOT NULL, salary INT, bonus INT, extra INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "extra": 5, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + BONUS WHERE pk = 1 AND EXTRA > 3")
+
+ checkLastScanSchema("pk INT, salary INT, bonus INT, extra INT, dep STRING, _partition STRING")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, 110, 10, 5, "hr") :: Nil)
+ }
+
+ test("column-update: scan that returns unrequested columns still writes narrow rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update", "true",
+ InMemoryBaseTable.SIMULATE_PARTIAL_COLUMN_PRUNING, "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastScanIncludes("bonus")
+ val columnUpdateSchema = table.lastWriteInfo.columnUpdateSchema().get()
+ assert(!columnUpdateSchema.fieldNames.contains("bonus"),
+ s"bonus must not be in column update schema: $columnUpdateSchema")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, -1, 10, "hr") :: Nil)
+ }
+
+ test("column-update: DELETE on a scan that returns unrequested columns is unaffected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update", "true",
+ InMemoryBaseTable.SIMULATE_PARTIAL_COLUMN_PRUNING, "true"))
+
+ sql(s"DELETE FROM $tableNameAsString WHERE pk < 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(2, 200, 20, "software") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively on an assigned " +
+ "column without adopting the declared spelling") {
+ createAndInitTableWithReqAttrs("pk,SALARY,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "dep"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: case-sensitive analysis accepts columns that differ only in case") {
+ withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") {
+ createAndInitTableWithReqAttrs("pk,salary,SALARY,dep",
+ "pk INT NOT NULL, salary INT, SALARY INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "SALARY": 200, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "SALARY", "dep"))
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, -1, 200, "hr") :: Nil)
+ }
+ }
+
+ test("column-update: empty requiredDataAttributes throws AnalysisException") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }""".stripMargin,
+ java.util.Map.of("column-update-empty-req-attrs", "true"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ parameters = Map("connector" -> reqAttrsOperationClass))
+ }
+
+ test("column-update: nested requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("pk,s.c1,dep",
+ "pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map("connector" -> reqAttrsOperationClass, "nestedAttributes" -> "[s.c1]"))
+ }
+
+ test("column-update: duplicate requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("pk,pk,salary", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map("connector" -> reqAttrsOperationClass, "duplicateAttributes" -> "[pk]"))
+ }
+
+ test("column-update: case-only duplicate requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("PK,pk,salary", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map("connector" -> reqAttrsOperationClass, "duplicateAttributes" -> "[PK]"))
+ }
+
+ test("column-update: metadata column in requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("pk,_partition,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET dep = 'x' WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "metadataAttributes" -> "[_partition]"))
+ }
+
+ test("column-update split: metadata column in requiredDataAttributes throws " +
+ "AnalysisException") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-req-attrs", "pk,_partition,dep"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET dep = 'x' WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> splitReqAttrsOperationClass,
+ "metadataAttributes" -> "[_partition]"))
+ }
+
+ test("column-update: requiredDataAttributes rejects a column that does not exist") {
+ createAndInitTableWithReqAttrs("pk,nonexistent_col,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "unknownAttributes" -> "[nonexistent_col]"))
+ }
+
+ test("column-update: column update schema excludes undeclared columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "dep", "salary"))
+ }
+
+ test("column-update: data correctness") {
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "id": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ }
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-split", "true"))
+ }
+
+ test("column-update split: column update schema has the declared columns") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update split: data correctness") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, -1, "hr") :: Nil)
+ }
+
+ test("column-update split: row-ID reassignment on narrow write is rejected") {
+ // `extra` is neither declared, updated, nor referenced, so requiredDataAttributes()
+ // ([pk, dep, salary]) doesn't cover every table column -- the write stays genuinely narrow
+ // and reassigning the row ID must be rejected.
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING, extra STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr", "extra": "x" }
+ |{ "pk": 2, "salary": 200, "dep": "software", "extra": "y" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT",
+ parameters = Map("connector" -> splitOperationClass, "rowIds" -> "[pk]"))
+ }
+
+ test("column-update split: row-ID reassignment is allowed when requiredDataAttributes " +
+ "covers every table column") {
+ createAndInitTableSplitWithReqAttrs("pk,salary,dep",
+ "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(2, 200, "software") :: Row(11, -1, "hr") :: Nil)
+ }
+
+ private def createAndInitTableSplitWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-split-req-attrs", reqAttrs))
+ }
+
+ test("column-update split: undeclared row-ID column on narrow write is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "pk"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED",
+ parameters = Map("connector" -> splitWithRowIdOperationClass, "rowIds" -> "[pk]"))
+ }
+
+ // analyzes without executing, as this fixture can't commit a write that omits `pk`
+ private def analyzeWriteDelta(sqlText: String): WriteDelta = {
+ val parsed = spark.sessionState.sqlParser.parsePlan(sqlText)
+ spark.sessionState.executePlan(parsed).analyzed.collectFirst {
+ case w: WriteDelta => w
+ }.getOrElse(fail("couldn't find WriteDelta in analyzed plan"))
+ }
+
+ test("column-update split: metadata row ID preserved on reinsert need not be declared") {
Review Comment:
Done in 22332b3. The split commit now locates the original row by the operation's `rowId()`, taking a data row-ID column from the narrow row and a metadata row-ID column from the preserved metadata, and fails unless exactly one row matches. The delta writer logs a non-integer row ID as null, and the commit resolves such a DELETE from its metadata.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -258,6 +354,407 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep` (an unconditionally declared base column) plus whatever columns
+ // Spark reports as being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = {
+ val base = Seq(FieldReference("pk"), FieldReference("dep"))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ private val clusterRef: NamedReference = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY))
+ .map(FieldReference(_))
+ .getOrElse(PARTITION_COLUMN_REF)
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.columnUpdateSchema().isPresent) {
+ capturedInfo.columnUpdateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(Array(clusterRef))
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ Array[SortOrder](
+ LogicalExpressions.sort(
+ clusterRef,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering())
+ )
+ }
+
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For column-update writes, updated rows contain only the declared columns
+ // (columnUpdateSchema from LogicalWriteInfo). We expand each row to the full
+ // table schema by overlaying them on the base row found by pk.
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+
+ val mergedData = newData.map { buf =>
+ val merged = new BufferedRows(buf.key, schema)
+ val updateOpName = UTF8String.fromString(Update.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == updateOpName) {
+ val pk = logRow.getInt(1)
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val baseRow = currentRowByPk(pk)
+ val fullRow = overlayNarrowRow(baseRow, narrowRow, writeFieldIdx)
+ merged.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ merged.rows.append(insertRow.copy())
+ }
+ }
+ merged
+ }
+
+ withDeletes(newData)
+ withData(mergedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ }
+ }
+ }
+ }
+
+ class DeltaBasedColumnUpdateOperationWithReqAttrs(
+ command: Command,
+ reqCols: Array[String],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, options = options) {
+ override def requiredDataAttributes(): Array[NamedReference] = reqCols.map(FieldReference(_))
+ }
+
+ class DeltaBasedColumnUpdateSplitOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols, options) {
+ override def representUpdateAsDeleteAndInsert(): Boolean = true
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.columnUpdateSchema().isPresent) {
+ capturedInfo.columnUpdateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+ override def requiredDistribution(): Distribution =
+ Distributions.clustered(Array(PARTITION_COLUMN_REF))
+ override def requiredOrdering(): Array[SortOrder] = Array[SortOrder](
+ LogicalExpressions.sort(
+ PARTITION_COLUMN_REF,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()))
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For delete+reinsert with narrow writes, the REINSERT row has only the
+ // connector-declared columns (requiredDataAttributes order). The base row is
+ // looked up by `pk`, so `pk` must be declared for the write to commit.
+ // Reconstruct the full row by overlaying the narrow row onto the original.
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+ val reinsertOpName = UTF8String.fromString(Reinsert.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+ val pkIdx = writeFieldIdx("pk")
Review Comment:
Done in 22332b3.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -129,6 +184,24 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
V2ExpressionUtils.resolveRef[AttributeReference](FieldReference(name), plan)
}
+ protected def isIdentityAssignment(key: Attribute, value: Expression): Boolean = {
Review Comment:
Done in 22332b3: `buildOriginalRowIdValues` now uses `isIdentityAssignment`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub) via reviews" <re...@spark.apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4429130835
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -159,14 +164,24 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
relation: DataSourceV2Relation,
operationTable: RowLevelOperationTable,
assignments: Seq[Assignment],
- cond: Expression): WriteDelta = {
+ cond: Expression,
+ writeAttrs: Seq[AttributeReference],
+ updatedAttrs: Seq[AttributeReference]): WriteDelta = {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
// resolve all needed attrs (e.g. row ID and any required metadata attrs)
- val rowAttrs = relation.output
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
+ val splitUpdates = operation.representUpdateAsDeleteAndInsert
+
+ operation match {
Review Comment:
I'll need to think about this a bit more. Specifically, what's the correct behavior if we opt in as representing update as delete and insert with narrow column updates. Keep in mind that Iceberg splits updates but uses synthetic row IDs... Seems like that would be incompatible with this arrangement then?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3634079860
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -41,7 +44,11 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
EliminateSubqueryAliases(aliasedTable) match {
case r @ ExtractV2Table(tbl: SupportsRowLevelOperations) =>
checkNoGeneratedColumns(r, UPDATE)
- val table = buildOperationTable(tbl, UPDATE, r.options)
+ val updatedCols = assignments.collect {
+ case Assignment(key: AttributeReference, value) if !isIdentityAssignment(key, value) =>
+ FieldReference(key.name)
Review Comment:
`FieldReference(key.name)` goes through the parser — `FieldReference.apply(String)` is
`LogicalExpressions.parseReference(column)` — rather than wrapping the name as-is. So a
top-level column whose name contains a dot (e.g. ``CREATE TABLE t (`a.b` INT, ...)``) gets
reported to the connector as a nested reference `a`.`b` instead of a single column named
`a.b`, and names containing backticks can even fail to parse.
Since `key` is an already-resolved `AttributeReference` at this point, there is no reason
to re-parse its name. This matters because the value flows into
`RowLevelOperationInfo.updatedColumns()`, which connectors use to size
`requiredDataAttributes()`, and it also feeds `validateUpdatedColumnsSubset` — a
mis-parsed reference would surface as a spurious validation failure or a wrong narrow
schema.
```suggestion
FieldReference(Seq(key.name))
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3632495355
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -210,11 +247,25 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
protected def buildReplaceDataProjections(
plan: LogicalPlan,
rowAttrs: Seq[Attribute],
- metadataAttrs: Seq[Attribute]): ReplaceDataProjections = {
+ metadataAttrs: Seq[Attribute],
+ updateRowAttrs: Seq[Attribute] = Nil): ReplaceDataProjections = {
val outputs = extractOutputs(plan)
- val outputsWithRow = filterOutputs(outputs, OPERATIONS_WITH_ROW)
- val rowProjection = newLazyProjection(plan, outputsWithRow, rowAttrs)
+ val rowProjection = if (updateRowAttrs.nonEmpty) {
+ val outputsForInsert = filterOutputs(outputs,
+ OPERATIONS_WITH_ROW -- Set(UPDATE_OPERATION, COPY_OPERATION))
+ newLazyProjection(plan, outputsForInsert, rowAttrs)
Review Comment:
This narrow branch doesn't guard an empty `outputsForInsert`, unlike `buildWriteDeltaProjections` right below (lines 294-296). For a column-update UPDATE there are no INSERT/REINSERT rows, so `outputsForInsert` is empty and `rowAttrs` is the full `relation.output`; `newLazyProjection` then resolves ordinals against the *narrow* query and returns `-1` for every column absent from the narrow scan, and `info.schema()` ends up as the full table shape -- whereas the delta path reports an empty schema for the same UPDATE.
It's harmless today because this projection is never applied for UPDATE, but the `-1` ordinals become a real `row.get(-1)` once MERGE/INSERT rows flow through this path. Mirroring the WriteDelta guard fixes both the inconsistency and the latent case:
```suggestion
if (outputsForInsert.isEmpty) {
ProjectingInternalRow(StructType(Nil), Nil)
} else {
newLazyProjection(plan, outputsForInsert, rowAttrs)
}
```
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java:
##########
@@ -37,4 +38,13 @@ public interface RowLevelOperationInfo {
* Returns the row-level SQL command (e.g. DELETE, UPDATE, MERGE).
*/
Command command();
+
+ /**
+ * Returns the columns being updated by this operation.
Review Comment:
`updatedColumns()` is only populated by `RewriteUpdateTable`; DELETE and MERGE fall through to the default `Nil`, so a connector running a MERGE sees an empty array even though columns are being updated. Worth stating the UPDATE-only scope on the public method (the note @dongjoon-hyun asked for earlier):
```suggestion
* Returns the columns being updated by this operation. Currently populated only for UPDATE;
* DELETE and MERGE report an empty array.
```
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,47 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ *
+ * @since 4.3.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * The returned columns become the schema of updated, copied, and reinserted rows, in declared
+ * order. Implementations must include every column they want to receive (typically the columns
+ * reported by {@link RowLevelOperationInfo#updatedColumns()} plus any columns needed for row
+ * lookup or routing, e.g. a primary key).
+ * <p>
+ * If any of the columns from {@link RowLevelOperationInfo#updatedColumns()}) is
Review Comment:
Stray `)` after `updatedColumns()}` and `is` -> `are`:
```suggestion
* If any of the columns from {@link RowLevelOperationInfo#updatedColumns()} are
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3633977690
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -244,6 +281,327 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends only the
+ // declared + assigned columns in the row projection instead of the full row schema. The base
+ // class composes its required-attrs set as `pk` (the row-lookup key) plus whatever columns
+ // Spark reports as being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil)
+ extends DeltaBasedOperation(command, CaseInsensitiveStringMap.empty())
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = {
+ val pk: NamedReference = FieldReference("pk")
+ val updatedNames = updatedCols.map(_.describe()).toSet
+ if (updatedNames.contains("pk")) updatedCols.toArray
+ else (pk +: updatedCols).toArray
+ }
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.updateSchema().isPresent) {
+ capturedInfo.updateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(Array(PARTITION_COLUMN_REF))
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ Array[SortOrder](
+ LogicalExpressions.sort(
+ PARTITION_COLUMN_REF,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering())
+ )
+ }
+
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ // Use the narrow update schema for UPDATE/COPY/REINSERT rows when present.
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For column-update writes, rows contain only the assigned columns
+ // (narrow schema from LogicalWriteInfo). We expand each row to the full table
+ // schema by overlaying write-schema columns on the base row found by pk.
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+
+ val mergedData = newData.map { buf =>
+ val merged = new BufferedRows(buf.key, schema)
+ val updateOpName = UTF8String.fromString(Update.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == updateOpName) {
+ val pk = logRow.getInt(1)
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val baseRow = dataMap.values.iterator.flatten
+ .flatMap(_.rows)
+ .find(r => r.getInt(schema.fieldIndex("pk")) == pk)
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ fullRow.update(i, base.get(i, schema.fields(i).dataType))
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ writeFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ merged.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ merged.rows.append(insertRow.copy())
+ }
+ }
+ merged
+ }
+
+ withDeletes(newData)
+ withData(mergedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ }
+ }
+ }
+ }
+
+ class DeltaBasedColumnUpdateOperationWithReqAttrs(command: Command, reqCols: Array[String])
+ extends DeltaBasedColumnUpdateOperation(command) {
+ override def requiredDataAttributes(): Array[NamedReference] = reqCols.map(FieldReference(_))
+ }
+
+ class DeltaBasedColumnUpdateOperationFromInfo(
+ command: Command,
+ updatedCols: Seq[NamedReference])
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols) {
+ }
+
+ class DeltaBasedColumnUpdateSplitOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil)
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols) {
+ override def representUpdateAsDeleteAndInsert(): Boolean = true
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.updateSchema().isPresent) {
+ capturedInfo.updateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+ override def requiredDistribution(): Distribution =
+ Distributions.clustered(Array(PARTITION_COLUMN_REF))
+ override def requiredOrdering(): Array[SortOrder] = Array[SortOrder](
+ LogicalExpressions.sort(
+ PARTITION_COLUMN_REF,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()))
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For delete+reinsert with narrow writes, the REINSERT row has only the
+ // connector-declared columns (requiredDataAttributes order).
+ // pk is the first field in the write schema (declared before updatedCols).
+ // Reconstruct the full row by overlaying the narrow row onto the original.
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+ val reinsertOpName = UTF8String.fromString(Reinsert.toString)
+ val pkIdx = writeFieldIdx("pk")
+
+ val expandedData = newData.map { buf =>
+ val expanded = new BufferedRows(buf.key, schema)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == reinsertOpName) {
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val pk = narrowRow.getInt(pkIdx)
+ val baseRow = dataMap.values.iterator.flatten
+ .flatMap(_.rows)
+ .find(r => r.getInt(schema.fieldIndex("pk")) == pk)
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ fullRow.update(i, base.get(i, schema.fields(i).dataType))
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ writeFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ expanded.rows.append(fullRow)
+ }
+ }
+ expanded
+ }
+
+ withDeletes(newData)
+ withData(expandedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ }
+ }
+ }
+ }
+
+ class PartitionBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil)
+ extends RowLevelOperation with SupportsColumnUpdates {
+ var configuredScan: InMemoryBatchScan = _
+
+ override def command(): Command = command
+
+ override def requiredDataAttributes(): Array[NamedReference] = {
+ val base = Seq(FieldReference("pk"), FieldReference("dep"))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ override def requiredMetadataAttributes(): Array[NamedReference] =
+ Array(PARTITION_COLUMN_REF, INDEX_COLUMN_REF)
+
+ override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = {
+ new InMemoryScanBuilder(schema, options) {
+ override def build(): Scan = {
+ val scan = super.build()
+ InMemoryRowLevelOperationTable.recordLastScanSchema(scan.readSchema())
+ configuredScan = scan.asInstanceOf[InMemoryBatchScan]
+ scan
+ }
+ }
+ }
+
+ override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = {
+ lastWriteInfo = info
+ new WriteBuilder {
+ override def build(): Write = new Write with RequiresDistributionAndOrdering {
+ override def requiredDistribution: Distribution =
+ Distributions.clustered(Array(PARTITION_COLUMN_REF))
+
+ override def requiredOrdering: Array[SortOrder] = Array[SortOrder](
+ LogicalExpressions.sort(
+ PARTITION_COLUMN_REF,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()))
+
+ override def toBatch: BatchWrite = {
+ val narrowSchema = if (info.updateSchema().isPresent) {
+ info.updateSchema().get()
+ } else {
+ info.schema()
+ }
+ PartitionBasedNarrowReplaceData(configuredScan, narrowSchema, info.schema())
Review Comment:
```suggestion
// info.schema() is empty for column-update UPDATE writes (no INSERT rows);
// use the table schema to align any INSERT-tagged rows.
PartitionBasedNarrowReplaceData(configuredScan, narrowSchema, schema)
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-4964392329
@aokolnychyi - This is ready for another review.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-4406789733
Thanks for the review @dongjoon-hyun. For the SPIP, we are waiting for a few more maintainers to also review the design as well as the PR before going for a vote.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192166920
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -50,20 +51,35 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
protected def buildOperationTable(
table: SupportsRowLevelOperations,
command: Command,
- options: CaseInsensitiveStringMap): RowLevelOperationTable = {
- val info = RowLevelOperationInfoImpl(command, options)
+ options: CaseInsensitiveStringMap,
+ updatedColumns: Seq[NamedReference] = Nil): RowLevelOperationTable = {
+ val info = RowLevelOperationInfoImpl(command, options, updatedColumns)
val operation = table.newRowLevelOperationBuilder(info).build()
RowLevelOperationTable(table, operation)
}
+ // Builds a DataSourceV2Relation for a row-level operation, optionally narrowing its output.
+ //
+ // When dataAttrs is non-empty, the relation output is narrowed to include only columns
+ // required for a column-update write. When dataAttrs is empty, the full relation.output is
+ // preserved.
protected def buildRelationWithAttrs(
relation: DataSourceV2Relation,
table: RowLevelOperationTable,
metadataAttrs: Seq[AttributeReference],
- rowIdAttrs: Seq[AttributeReference] = Nil): DataSourceV2Relation = {
-
- val attrs = dedupAttrs(relation.output ++ rowIdAttrs ++ metadataAttrs)
- relation.copy(table = table, output = attrs)
+ rowIdAttrs: Seq[AttributeReference] = Nil,
+ dataAttrs: Seq[AttributeReference] = Nil,
+ cond: Expression = TrueLiteral): DataSourceV2Relation = {
+
+ if (dataAttrs.nonEmpty) {
+ val required =
+ AttributeSet(dataAttrs) ++ AttributeSet(Seq(cond)) ++ AttributeSet(rowIdAttrs)
+ val narrowOutput = relation.output.filter(required.contains)
+ relation.copy(table = table, output = dedupAttrs(narrowOutput ++ rowIdAttrs ++ metadataAttrs))
Review Comment:
> Can an attribute in required be missing from relation.output?
No. dataAttrs come from the connector's `requiredDataAttributes()` which are resolved against relation (via `V2ExpressionUtils.resolveRefs`), so they're guaranteed to be present. The condition's referenced columns are also table columns from the user's `WHERE` clause. `rowIdAttrs` and `metadataAttrs` can be absent from relation.output (they're resolved separately), but they're not part of the filter. They're appended unconditionally afterward via `dedupAttrs(narrowOutput ++ rowIdAttrs ++ metadataAttrs)`
> rowIdAttrs seems to be added 2 times. If we already have dedupAttrs() then probably doesn't make sense to build AttributeSets.
Agreed. I fixed it.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3785101803
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/LogicalWriteInfo.java:
##########
@@ -65,4 +65,14 @@ default Optional<StructType> metadataSchema() {
throw new SparkUnsupportedOperationException(
"DATA_SOURCE_METADATA_SCHEMA_NOT_IMPLEMENTED", Map.of("class", getClass().getName()));
}
+
+ /**
+ * the narrow schema for updates. Present only when the connector mixes in
+ * {@link SupportsColumnUpdates}.
+ *
+ * @since 4.3.0
+ */
+ default Optional<StructType> updateSchema() {
Review Comment:
This could technically also become `columnUpdateSchema()`, if we decide to do the rename. It does indicate the schema of copied rows in CoW and reinserted rows.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3785929879
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -129,6 +156,17 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
V2ExpressionUtils.resolveRef[AttributeReference](FieldReference(name), plan)
}
+ protected def isIdentityAssignment(key: Attribute, value: Expression): Boolean = {
+ val unwrapped = value match {
+ case Alias(child, _) => child
+ case other => other
+ }
+ unwrapped match {
+ case attr: Attribute => AttributeSet(Seq(key)).contains(attr)
Review Comment:
It is a bit odd to see AttributeSet being used here? Can we do this instead?
```
protected def isIdentityAssignment(key: Attribute, value: Expression): Boolean = {
val valueWithoutAlias = value match {
case Alias(child, _) => child
case other => other
}
key.semanticEquals(valueWithoutAlias)
}
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4392772308
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala:
##########
@@ -380,14 +380,17 @@ case class ReplaceDataExec(
write: Write,
rowLevelCommand: RowLevelOperation.Command,
tableName: String,
+ useWriteColumnUpdate: Boolean,
Review Comment:
Done.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DeltaWriter.java:
##########
@@ -52,6 +55,9 @@ public interface DeltaWriter<T> extends DataWriter<T> {
* Reinserts a row with metadata.
* <p>
* This method handles the insert portion of updated rows split into deletes and inserts.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, the {@code row} follows it;
+ * otherwise it follows {@link LogicalWriteInfo#schema()}.
Review Comment:
Done.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,623 @@
+/*
+ * 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 scala.jdk.CollectionConverters._
+
+import org.apache.spark.{SparkRuntimeException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.analysis.RewriteRowLevelCommand
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+import org.apache.spark.sql.catalyst.util.RowDeltaUtils.ORIGINAL_ROW_ID_VALUE_PREFIX
+import org.apache.spark.sql.connector.catalog.{Delete, InMemoryBaseTable, InMemoryRowLevelOperationTable, Reinsert}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.write.LogicalWriteInfo
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.errors.QueryCompilationErrors
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, updated rows contain only the declared columns
+ * (LogicalWriteInfo.columnUpdateSchema()) rather than the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase
Review Comment:
Done.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala:
##########
@@ -92,6 +95,38 @@ object GroupBasedRowLevelOperationScanPlanning extends Rule[LogicalPlan] with Pr
}
}
+ /**
+ * Returns the columns of `relation` that the projections and filters directly above each read
+ * of the table use, as the rest of the query sees only their output. Only whole columns are
+ * pruned, as the query is not rewritten for a pruned nested schema.
+ *
+ * An UPDATE with a subquery in its condition reads the table twice, once per Union child, and
+ * the analyzer gives the second read new expr IDs. Both reads share one scan, so the columns
+ * are matched to `relation` by name.
+ */
+ private def readAttrs(
+ plan: LogicalPlan,
+ table: RowLevelOperationTable,
+ relation: DataSourceV2Relation): Seq[AttributeReference] = {
+ val readNames = tableReads(plan, table).flatMap { case (projects, filters, read) =>
+ val references = AttributeSet((projects ++ filters).flatMap(_.references))
+ read.output.filter(references.contains).map(_.name)
+ }.toSet
+ relation.output.filter(attr => readNames.contains(attr.name))
+ }
+
+ private def tableReads(
Review Comment:
Done.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,105 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, updated, copied, and reinserted
+ * rows follow it: group-based operations receive them through
+ * {@link DataWriter#writeColumnUpdate(Object, Object)} or
+ * {@link DataWriter#writeColumnUpdate(Object)}, and operations that mix in {@link SupportsDelta}
+ * through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert}. Inserted rows always follow
+ * {@link LogicalWriteInfo#schema()} and arrive through {@link DataWriter#write(Object)} and
+ * {@link DeltaWriter#insert}. When {@link LogicalWriteInfo#columnUpdateSchema()} is absent, every
+ * row follows {@link LogicalWriteInfo#schema()}, as for any other operation. Connectors must
+ * decide which layout to expect by whether {@link LogicalWriteInfo#columnUpdateSchema()} is
+ * present, not by the command. A builder that wants full-width rows for some command can check
+ * {@link RowLevelOperationInfo#command()} and build an operation without this mix-in. Currently
+ * Spark narrows only UPDATE.
+ * <p>
+ * The scan builder returned by {@link #newScanBuilder} should implement
+ * {@link org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns}, so that the scan
+ * does not read columns the command neither references nor declares. Otherwise the scan reads
+ * every column. For operations that do not mix in {@link SupportsDelta}, only top-level columns
+ * are pruned: a struct column the command references through a nested field is read in full.
+ *
+ * @since 4.4.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * When Spark narrows the write, the returned columns become
+ * {@link LogicalWriteInfo#columnUpdateSchema()}, in declared order. Implementations must include
Review Comment:
Done.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -273,6 +339,423 @@ class InMemoryRowLevelOperationTable private (
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
}
+ // Finds the one current row with the given row-ID values, matching a data column by value and
+ // `_partition` by the partition key.
+ private def currentRowByRowId(values: Seq[(String, Any)]): InternalRow = {
+ val matches = dataMap.values.iterator.flatten.flatMap { buf =>
+ buf.rows.filter { row =>
+ values.forall {
+ case (name, value) if name == PartitionKeyColumn.name => buf.keyString() == value.toString
+ case (name, value) =>
+ val i = schema.fieldIndex(name)
+ row.get(i, schema.fields(i).dataType) == value
+ }
+ }
+ }.toSeq
+ matches match {
+ case Seq(row) => row
+ case _ => throw new IllegalStateException(
+ s"Expected one row with row ID $values, found ${matches.size}")
+ }
+ }
+
+ // Rebuilds a full table row for a narrow column-update row: the base row's values (or the
+ // existence default for a column added after the base row was written), overlaid with the
+ // narrow row's columns.
+ private def overlayNarrowRow(
+ baseRow: Option[InternalRow],
+ narrowRow: InternalRow,
+ narrowFieldIdx: Map[String, Int]): InternalRow = {
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ val field = schema.fields(i)
+ val value = if (i < base.numFields) {
+ base.get(i, field.dataType)
+ } else {
+ ResolveDefaultColumns.getExistenceDefaultValue(field)
+ }
+ fullRow.update(i, value)
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ narrowFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ fullRow
+ }
+
+ // The base columns `pk` and `dep`, where the table has them, followed by the updated columns.
+ private def declaredColumns(updatedCols: Seq[NamedReference]): Array[NamedReference] = {
+ val base = Seq("pk", "dep").filter(schema.fieldNames.contains).map(FieldReference(_))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep`, where the table has them, plus whatever columns Spark reports as
+ // being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = declaredColumns(updatedCols)
+
+ private val clusterExprs: Seq[Expression] = {
+ val refs = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY))
+ .map(_.split(",").map(name => FieldReference(name.trim)).toSeq)
+ .getOrElse(Seq(PARTITION_COLUMN_REF))
+ if (properties.getOrDefault(COLUMN_UPDATE_CLUSTER_BY_NESTED_TRANSFORM, "false") == "true") {
+ refs.map(ref => LogicalExpressions.apply("f", LogicalExpressions.apply("g", ref)))
+ } else {
+ refs
+ }
+ }
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ val writeSchema = info.columnUpdateSchema().orElse(info.schema())
+ val batchWrite = new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ physicalInfo: PhysicalWriteInfo): DeltaWriterFactory = {
+ newWriterFactory(info, writeSchema)
+ }
+
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ commitColumnUpdate(info, writeSchema, newData)
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ if (noMetadata && !properties.containsKey(COLUMN_UPDATE_CLUSTER_BY)) {
+ new DeltaWrite {
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ } else {
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(clusterExprs.toArray)
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ clusterExprs.map { expr =>
+ LogicalExpressions.sort(
+ expr,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()): SortOrder
+ }.toArray
+ }
+
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ }
+ }
+ }
+
+ protected def newWriterFactory(
+ info: LogicalWriteInfo,
+ writeSchema: StructType): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(writeSchema)
+ }
+
+ // For column-update writes, updated rows contain only the declared columns
+ // (columnUpdateSchema from LogicalWriteInfo). We expand each row to the full
+ // table schema by overlaying them on the base row found by pk.
+ protected def commitColumnUpdate(
+ info: LogicalWriteInfo,
+ writeSchema: StructType,
+ newData: Array[BufferedRows]): Unit = {
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+ val updateOpName = UTF8String.fromString(Update.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+
+ val mergedData = newData.map { buf =>
+ val merged = new BufferedRows(buf.key, schema)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == updateOpName) {
+ val pk = logRow.getInt(1)
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val baseRow = Some(currentRowByRowId(Seq("pk" -> pk)))
+ val fullRow = overlayNarrowRow(baseRow, narrowRow, writeFieldIdx)
+ merged.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ merged.rows.append(insertRow.copy())
+ }
+ }
+ merged
+ }
+
+ withDeletes(newData)
+ withData(mergedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+ }
+
+ class DeltaBasedColumnUpdateOperationWithReqAttrs(
+ command: Command,
+ reqCols: Array[String],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, options = options) {
+ override def requiredDataAttributes(): Array[NamedReference] = reqCols.map(FieldReference(_))
+ }
+
+ class DeltaBasedColumnUpdateSplitOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols, options) {
+ override def representUpdateAsDeleteAndInsert(): Boolean = true
+
+ override protected def newWriterFactory(
+ info: LogicalWriteInfo,
+ writeSchema: StructType): DeltaWriterFactory = {
+ val pkRowId = rowId().map(_.describe()).toSeq == Seq("pk") &&
+ info.rowIdSchema().get().fields.map(_.dataType).toSeq == Seq(IntegerType)
+ new DeltaBufferedRowsWriterFactory(writeSchema, pkRowId)
+ }
+
+ // For delete+reinsert with narrow writes, the REINSERT row has only the connector-declared
+ // columns (requiredDataAttributes order). The original row is located by the row ID, taking a
+ // data row-ID column from the narrow row and a metadata row-ID column from the preserved
+ // metadata, and the full row is rebuilt by overlaying the narrow row onto it. A DELETE logged
+ // without a `pk` row ID is resolved the same way from its metadata, so such a row ID must
+ // consist of metadata columns.
+ override protected def commitColumnUpdate(
+ info: LogicalWriteInfo,
+ writeSchema: StructType,
+ newData: Array[BufferedRows]): Unit = {
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+ val metadataSchema = info.metadataSchema().orElse(new StructType())
+ val coversTable = schema.fieldNames.forall(writeFieldIdx.contains)
+ val deleteOpName = UTF8String.fromString(Delete.toString)
+ val reinsertOpName = UTF8String.fromString(Reinsert.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+
+ def rowIdValues(narrowRow: Option[InternalRow], metadata: InternalRow): Seq[(String, Any)] = {
+ rowId().toSeq.map(_.describe()).map { name =>
+ (writeFieldIdx.get(name), narrowRow) match {
+ case (Some(i), Some(row)) => name -> row.get(i, writeSchema(i).dataType)
+ case _ =>
+ if (metadata == null || !metadataSchema.fieldNames.contains(name)) {
+ throw new IllegalStateException(s"The split fixture can only locate a row by pk " +
+ s"or by metadata row-ID columns, but row-ID column $name is in neither")
+ }
+ val i = metadataSchema.fieldIndex(name)
+ name -> metadata.get(i, metadataSchema(i).dataType)
+ }
+ }
+ }
+
+ def metadataOf(logRow: InternalRow): InternalRow =
+ if (logRow.isNullAt(2)) null else logRow.getStruct(2, metadataSchema.length)
+
+ val expandedData = newData.map { buf =>
+ val expanded = new BufferedRows(buf.key, schema)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == deleteOpName && logRow.isNullAt(1)) {
+ val values = rowIdValues(None, metadataOf(logRow))
+ buf.deletes += currentRowByRowId(values).getInt(0)
+ } else if (opName == reinsertOpName) {
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val baseRow = if (coversTable) {
+ None
+ } else {
+ Some(currentRowByRowId(rowIdValues(Some(narrowRow), metadataOf(logRow))))
+ }
+ val fullRow = overlayNarrowRow(baseRow, narrowRow, writeFieldIdx)
+ expanded.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ expanded.rows.append(insertRow.copy())
+ }
+ }
+ expanded
+ }
+
+ withDeletes(newData)
+ withData(expandedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+ }
+
+ // Test-only: a split-update fixture whose requiredDataAttributes() is the given column list.
+ class DeltaBasedColumnUpdateSplitOperationWithReqAttrs(
+ command: Command,
+ reqCols: Array[String],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateSplitOperation(command, options = options) {
+ override def requiredDataAttributes(): Array[NamedReference] = reqCols.map(FieldReference(_))
+ }
+
+ // Test-only: a split-update fixture whose row ID is the given column and whose
+ // requiredDataAttributes() lists only the updated columns, to exercise which undeclared
+ // row-ID columns still reach the REINSERT row.
+ class DeltaBasedColumnUpdateSplitOperationWithRowId(
+ command: Command,
+ rowIdCol: String,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateSplitOperation(command, updatedCols, options) {
+ override def rowId(): Array[NamedReference] = Array(FieldReference(rowIdCol))
+ override def requiredDataAttributes(): Array[NamedReference] = updatedCols.toArray
+ }
+
+ class PartitionBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ override val options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends RowLevelOperation with SupportsColumnUpdates with RowLevelOperationWithOptions {
+ var configuredScan: BatchScanBaseClass = _
+
+ override def command(): Command = command
+
+ override def requiredDataAttributes(): Array[NamedReference] = declaredColumns(updatedCols)
+
+ override def requiredMetadataAttributes(): Array[NamedReference] = {
+ if (noMetadata) {
+ Array.empty
+ } else {
+ Array(PARTITION_COLUMN_REF, INDEX_COLUMN_REF)
+ }
+ }
+
+ private def clusterColumnRef: NamedReference =
+ if (noMetadata) FieldReference("dep") else PARTITION_COLUMN_REF
+
+ override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = {
+ val onBuild: BatchScanBaseClass => Unit = { scan =>
+ configuredScan = scan
+ }
+ if (properties.getOrDefault(COLUMN_UPDATE_COW_UNREAD_FILTER_ATTRS, "false") == "true") {
Review Comment:
I went with the fallback. The unread-filter branch now honors `useCatalystRuntimeFiltering`, and its test runs with the flag off and on and checks the scan type. Sharing the other parts would mean changing master's `PartitionBasedOperation`, `PartitionBasedReplaceData` and `newRowLevelScanBuilder`, so I left them as they are.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -321,25 +804,63 @@ class InMemoryRowLevelOperationTable private (
}
}
-private class DeltaBufferedRowsWriterFactory(schema: StructType) extends DeltaWriterFactory {
+private class DeltaBufferedRowsWriterFactory(
+ schema: StructType,
+ pkRowId: Boolean = true) extends DeltaWriterFactory {
override def createWriter(partitionId: Int, taskId: Long): DeltaWriter[InternalRow] = {
- new DeltaBufferWriter(schema)
+ new DeltaBufferWriter(schema, pkRowId)
+ }
+}
+
+// Test-only: a writer factory for a connector that never overrides writeColumnUpdate(), to
Review Comment:
Done.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/ColumnUpdateTableSharedTests.scala:
##########
@@ -0,0 +1,355 @@
+/*
+ * 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.{SparkException, SparkRuntimeException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.{InMemoryBaseTable, InMemoryRowLevelOperationTable}
+import org.apache.spark.sql.connector.write.{RowLevelOperation, RowLevelOperationTable}
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+// Column-update tests that hold for both delta-based and group-based operations; each suite
+// supplies its column-update table properties through `extraTableProps`.
+trait ColumnUpdateTableSharedTests extends UpdateTableSuiteBase {
+
+ // the properties of a column-update table whose connector declares exactly `reqAttrs`
+ protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String]
+
+ // the properties of a table of the same kind whose operation does not mix in column updates
+ protected def fullWidthTableProps: java.util.Map[String, String]
+
+ // the metadata columns the scan of a column-update table reads, in DDL form
+ protected def scanMetadataColumns: String = ", _partition STRING, index INT"
+
+ // the class name of the operation that `reqAttrsTableProps` builds
+ protected def reqAttrsOperationClass: String = if (deltaUpdate) {
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateOperationWithReqAttrs].getName
+ } else {
+ classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperationWithReqAttrs].getName
+ }
+
+ // t(pk, id, dep, salary, bonus, extra) with CHECK (extra > 0), partitioned by dep
+ protected def createAndInitTableWithProps(props: java.util.Map[String, String]): Unit = {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, salary INT, bonus INT, extra INT",
+ """{ "pk": 1, "id": 1, "dep": "hr", "salary": 100, "bonus": 10, "extra": 1 }
+ |{ "pk": 2, "id": 2, "dep": "hr", "salary": 200, "bonus": 20, "extra": 1 }
+ |{ "pk": 3, "id": 3, "dep": "software", "salary": 300, "bonus": 30, "extra": 1 }
+ |""".stripMargin,
+ props)
+ sql(s"ALTER TABLE $tableNameAsString ADD CONSTRAINT positive_extra CHECK (extra > 0)")
+ }
+
+ test("column-update: undeclared columns are not written and keep their values") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
Review Comment:
Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala:
##########
@@ -56,6 +56,56 @@ object V2ExpressionUtils extends SQLConfHelper with Logging {
plan.resolve(ref.fieldNames.toImmutableArraySeq, conf.resolver)
}
+ /**
+ * Returns the names of the columns referenced by the given expressions that cannot be resolved
+ * against the plan, at any depth, including inside transform arguments.
+ */
+ private[sql] def unresolvableColumns(
+ exprs: Seq[V2Expression],
+ relation: LogicalPlan): Seq[String] = {
+ unresolvableReferences(exprs, relation).map {
+ case (Seq(), _) => "<empty>"
+ case (parts, None) => parts.quoted
+ case (parts, Some(condition)) => s"${parts.quoted} ($condition)"
+ }.distinct
Review Comment:
Done.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -273,6 +339,423 @@ class InMemoryRowLevelOperationTable private (
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
}
+ // Finds the one current row with the given row-ID values, matching a data column by value and
+ // `_partition` by the partition key.
+ private def currentRowByRowId(values: Seq[(String, Any)]): InternalRow = {
+ val matches = dataMap.values.iterator.flatten.flatMap { buf =>
+ buf.rows.filter { row =>
+ values.forall {
+ case (name, value) if name == PartitionKeyColumn.name => buf.keyString() == value.toString
+ case (name, value) =>
+ val i = schema.fieldIndex(name)
+ row.get(i, schema.fields(i).dataType) == value
+ }
+ }
+ }.toSeq
+ matches match {
+ case Seq(row) => row
+ case _ => throw new IllegalStateException(
+ s"Expected one row with row ID $values, found ${matches.size}")
+ }
+ }
+
+ // Rebuilds a full table row for a narrow column-update row: the base row's values (or the
+ // existence default for a column added after the base row was written), overlaid with the
+ // narrow row's columns.
+ private def overlayNarrowRow(
+ baseRow: Option[InternalRow],
+ narrowRow: InternalRow,
+ narrowFieldIdx: Map[String, Int]): InternalRow = {
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ val field = schema.fields(i)
+ val value = if (i < base.numFields) {
+ base.get(i, field.dataType)
+ } else {
+ ResolveDefaultColumns.getExistenceDefaultValue(field)
+ }
+ fullRow.update(i, value)
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ narrowFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ fullRow
+ }
+
+ // The base columns `pk` and `dep`, where the table has them, followed by the updated columns.
+ private def declaredColumns(updatedCols: Seq[NamedReference]): Array[NamedReference] = {
+ val base = Seq("pk", "dep").filter(schema.fieldNames.contains).map(FieldReference(_))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep`, where the table has them, plus whatever columns Spark reports as
+ // being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
Review Comment:
Done. I also overrode `canEqual`, because with only `equals` and `hashCode` the wide operation's case-class `equals` still matched the narrow one. The swap test now builds its replacements from the plan's table and asserts each one differs from the original operation. Before, `table` returned a different instance, so the pair never compared equal there.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,362 @@
+/*
+ * 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 java.util.regex.Pattern
+
+import org.apache.spark.SparkUnsupportedOperationException
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTable
+import org.apache.spark.sql.connector.write.RowLevelOperationTable
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+
+class GroupBasedColumnUpdateTableSuite extends UpdateTableSuiteBase
+ with ColumnUpdateTableSharedTests {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-cow", "true")
+ props
+ }
+
+ override protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String] = {
+ java.util.Map.of("column-update-cow-req-attrs", reqAttrs)
+ }
+
+ override protected def fullWidthTableProps: java.util.Map[String, String] = java.util.Map.of()
+
+ test("column-update ReplaceData: runtime group filtering reads only the referenced columns") {
+ Seq(true, false).foreach { dppEnabled =>
+ Seq(true, false).foreach { aqeEnabled =>
+ withSQLConf(
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> dppEnabled.toString,
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString) {
+ withTable(tableNameAsString) {
+ withTempView("matched_pk") {
+ createAndInitTable("pk INT, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ import testImplicits._
+ val matchedPkDF = Seq(Some(1), None).toDF()
+ matchedPkDF.createOrReplaceTempView("matched_pk")
+
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = -1 " +
+ s"WHERE pk IN (SELECT * FROM matched_pk)",
+ primaryScanSchema =
+ "pk INT, salary INT, dep STRING, _partition STRING, index INT",
+ groupFilterScanSchema = Some("pk INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, 300, 30, "hr") :: Nil)
+
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("column-update ReplaceData: updated and copied rows reach writeColumnUpdate with " +
+ "metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", null), data = Row(1, "hr", -1)),
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", 1), data = Row(2, "hr", 200)))
+ }
+
+ test("column-update ReplaceData: without metadata, rows reach writeColumnUpdate(record)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "no-metadata", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD, DEP_FIELD, StructField("salary", IntegerType, nullable = true)))))
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Row(3, 300, 30, "software") :: Nil)
+ }
+
+ test("column-update ReplaceData: a writer overriding only writeColumnUpdate(record) " +
+ "receives rows with metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "column-update-cow-record-only-writer", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: connector missing writeColumnUpdate override is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-cow", "true", "column-update-cow-no-write-column-update", "true"))
+
+ checkError(
+ exception = intercept[SparkUnsupportedOperationException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED",
+ parameters = Map(
+ "class" -> "org.apache.spark.sql.connector.catalog.NoWriteColumnUpdateOverrideWriter"))
+ }
+
+ test("column-update ReplaceData: the subquery plan narrows the write relation and the scan") {
+ withTempView("deps") {
+ createAndInitTableWithProps(reqAttrsTableProps("pk,salary"))
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ val write = executeAndKeepAnalyzedWrite {
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE dep IN (SELECT dep FROM deps)",
+ "pk INT, dep STRING, salary INT, bonus INT, extra INT, _partition STRING, index INT")
+ }
+
+ assert(write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "pk", "salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+ }
+
+ test("column-update ReplaceData: runtime group filtering on an undeclared condition column") {
+ withTempView("matched_ids") {
+ // `dep` is declared because the scan filters groups at runtime by the partition column
+ createAndInitTableWithProps(reqAttrsTableProps("pk,dep,salary"))
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `id` is neither declared nor in the write relation; the group filter must still map it
+ // onto the second read of the table that the subquery rewrite adds
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, dep STRING, salary INT, bonus INT, extra INT, " +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = Some("id INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+
+ test("column-update ReplaceData: unread partition column skips runtime group filtering") {
+ withTempView("matched_ids") {
+ createAndInitTableWithProps(reqAttrsTableProps("pk,salary"))
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `dep` is neither declared nor referenced, so the scan doesn't read it and doesn't
+ // report it from filterAttributes(); every group is read and replaced
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, salary INT, bonus INT, extra INT, " +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = None)
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr", "software"))
+ }
+ }
+
+ test("column-update ReplaceData: runtime filter attribute the scan doesn't read is rejected") {
+ withTempView("matched_ids") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, salary INT",
+ """{ "pk": 1, "id": 1, "dep": "hr", "salary": 100 }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-cow-req-attrs", "pk,salary",
+ "column-update-cow-unread-filter-attrs", "true"))
+ import testImplicits._
+ Seq(1).toDF().createOrReplaceTempView("matched_ids")
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE id IN (SELECT * FROM matched_ids)")
+ },
+ condition = "DATA_SOURCE_INVALID_RUNTIME_FILTER_ATTRIBUTE.CANNOT_RESOLVE",
+ sqlState = Some("KD000"),
+ parameters = Map(
+ "attribute" -> Pattern.quote("`dep`"),
+ "method" -> Pattern.quote("filterAttributes()"),
+ "scanClass" -> (Pattern.quote(
+ classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperation].getName) +
+ """\$\$anon\$\d+\$\$anon\$\d+"""),
+ "relationOutput" -> Pattern.quote("\"STRUCT<pk: INT NOT NULL, id: INT, salary: INT, " +
+ "_partition: STRING NOT NULL COMMENT 'Partition key used to store the row', " +
+ "index: INT NOT NULL COMMENT 'Metadata column used to conflict with a data " +
+ "column'>\"")),
+ matchPVals = true)
+ }
+ }
+
+ test("column-update ReplaceData: undeclared distribution column is rejected") {
+ // without metadata, the write clusters by `dep`
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow-req-attrs", "pk,salary", "no-metadata", "true"))
+
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName
+ for {
+ cond <- Seq("pk = 1", "dep = 'hr'")
+ confs <- Seq(Nil, Seq(excludeColumnPruning))
+ } withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE $cond")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[`dep`]"))
+ }
+ }
+
+ test("column-update ReplaceData: UPDATE without column updates, DELETE and MERGE read every " +
+ "column") {
+ val fullScanSchema = "pk INT, salary INT, bonus INT, dep STRING, _partition STRING, index INT"
+ Seq(fullWidthTableProps, extraTableProps).foreach { props =>
+ withTable(tableNameAsString) {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ props)
+
+ if (props == fullWidthTableProps) {
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1", fullScanSchema)
+ }
+ executeAndCheckScan(s"DELETE FROM $tableNameAsString WHERE pk > 1", fullScanSchema)
+ withTempView("source") {
+ import testImplicits._
+ Seq((1, 150)).toDF("pk", "salary").createOrReplaceTempView("source")
+ executeAndCheckScan(
+ s"""MERGE INTO $tableNameAsString t
+ |USING source s
+ |ON t.pk = s.pk
+ |WHEN MATCHED THEN UPDATE SET salary = s.salary
+ |""".stripMargin,
+ fullScanSchema)
+ }
+
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, 150, 10, "hr") :: Nil)
+ }
+ }
+ }
+
+ test("column-update ReplaceData: an operation rebuilt with withTable keeps the updated columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ }
+ val relation = write.table.asInstanceOf[DataSourceV2Relation]
+ val operationTable = relation.table.asInstanceOf[RowLevelOperationTable]
+ val t = table
+ t.lastUpdatedColumns = null
+
+ val rebuilt = operationTable.withTable(t)
+
+ checkLastUpdatedColumns("salary")
+ assert(rebuilt.operationInfo eq operationTable.operationInfo)
+ val rebuiltQuery = write.query.transform {
+ case r: DataSourceV2Relation if r.table eq operationTable => r.copy(table = rebuilt)
+ }
+ assert(rebuiltQuery.exists {
+ case r: DataSourceV2Relation => r.table eq rebuilt
+ case _ => false
+ })
+ V2Writes(write.withNewTable(relation.copy(table = rebuilt)).withNewQuery(rebuiltQuery))
+ }
+
+ test("column-update ReplaceData: column update schema nullability covers UPDATE and COPY " +
+ "rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT NOT NULL, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": null, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow-req-attrs", "pk,salary,bonus"))
Review Comment:
Done.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,623 @@
+/*
+ * 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 scala.jdk.CollectionConverters._
+
+import org.apache.spark.{SparkRuntimeException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.analysis.RewriteRowLevelCommand
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+import org.apache.spark.sql.catalyst.util.RowDeltaUtils.ORIGINAL_ROW_ID_VALUE_PREFIX
+import org.apache.spark.sql.connector.catalog.{Delete, InMemoryBaseTable, InMemoryRowLevelOperationTable, Reinsert}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.write.LogicalWriteInfo
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.errors.QueryCompilationErrors
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, updated rows contain only the declared columns
+ * (LogicalWriteInfo.columnUpdateSchema()) rather than the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase
+ with ColumnUpdateTableSharedTests {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ override protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String] = {
+ java.util.Map.of("column-update-req-attrs", reqAttrs)
+ }
+
+ override protected def fullWidthTableProps: java.util.Map[String, String] = {
+ java.util.Map.of("supports-deltas", "true")
+ }
+
+ override protected def scanMetadataColumns: String = ", _partition STRING"
+
+ private val splitOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperation].getName
+
+ private val splitWithRowIdOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithRowId].getName
+
+ private def splitReqAttrsTableProps(reqAttrs: String): java.util.Map[String, String] = {
+ java.util.Map.of("column-update-req-attrs", reqAttrs, "split-updates", "true")
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData, reqAttrsTableProps(reqAttrs))
+ }
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update", "true", "split-updates", "true"))
+ }
+
+ test("column-update: column update schema contains a single assigned column") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: CHAR and VARCHAR self-assignments are reported as updated") {
+ createTable("pk INT NOT NULL, id INT, c CHAR(3), v VARCHAR(10), dep STRING")
+ sql(s"INSERT INTO $tableNameAsString VALUES (1, 1, 'a', 'x', 'hr')")
+
+ sql(s"UPDATE $tableNameAsString SET id = id, c = c, v = v WHERE pk = 1")
+
+ checkLastUpdatedColumns("c", "v")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, 1, "a ", "x", "hr") :: Nil)
+ }
+
+ test("column-update: a LogicalWriteInfo that does not implement columnUpdateSchema throws") {
+ val info = new LogicalWriteInfo {
+ override def options(): CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty()
+ override def queryId(): String = "query"
+ override def schema(): StructType = new StructType()
+ }
+
+ checkError(
+ exception = intercept[SparkUnsupportedOperationException](info.columnUpdateSchema()),
+ condition = "DATA_SOURCE_COLUMN_UPDATE_SCHEMA_NOT_IMPLEMENTED",
+ parameters = Map("class" -> info.getClass.getName))
+ }
+
+ test("column-update: updatedColumns and the column update schema follow the non-identity " +
+ "assignments") {
+ Seq(
+ ("id = -1", Seq("id"), Seq("pk", "dep", "id"), Row(1, -1, "hr")),
+ ("id = -1, dep = 'eng'", Seq("id", "dep"), Seq("pk", "dep", "id"), Row(1, -1, "eng")),
+ ("id = -1, dep = dep", Seq("id"), Seq("pk", "dep", "id"), Row(1, -1, "hr")),
+ ("id = id, dep = 'eng'", Seq("dep"), Seq("pk", "dep"), Row(1, 10, "eng")),
+ ("id = id, dep = dep", Seq.empty, Seq("pk", "dep"), Row(1, 10, "hr")),
+ ("id = pk", Seq("id"), Seq("pk", "dep", "id"), Row(1, 1, "hr"))
+ ).foreach { case (assignments, updatedColumns, columnUpdateNames, updatedRow) =>
+ withTable(tableNameAsString) {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 10, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET $assignments WHERE pk = 1")
+
+ checkLastUpdatedColumns(updatedColumns: _*)
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq == columnUpdateNames)
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), updatedRow :: Nil)
+ }
+ }
+ }
+
+ test("column-update split: row-level DELETE falls back to full-width rows") {
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin)
+ checkDeleteFallsBackToFullWidthRows(Row(2, 200, "hr"))
+ }
+
+ test("column-update split: row-level DELETE by a metadata row ID falls back to full-width " +
+ "rows") {
+ // Each partition holds one row, so `_partition` identifies it.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "_partition"))
+ checkDeleteFallsBackToFullWidthRows(Row(2, 200, "software"))
+ }
+
+ test("column-update split: MERGE falls back to full-width rows") {
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+ checkMergeFallsBackToFullWidthRows()
+ }
+
+ test("column-update: undeclared write requirement columns are rejected") {
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName
+ val splitReqAttrsOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithReqAttrs]
+ .getName
+ Seq(
+ ("pk,salary", "dep", Map.empty[String, String], reqAttrsOperationClass, "[`dep`]"),
+ ("pk,salary", "dep", Map("split-updates" -> "true"), splitReqAttrsOperationClass,
+ "[`dep`]"),
+ ("pk,salary", "dep", Map("column-update-cluster-by-nested-transform" -> "true"),
+ reqAttrsOperationClass, "[`dep`]"),
+ ("pk,salary", "id,dep", Map.empty[String, String], reqAttrsOperationClass,
+ "[`dep`, `id`]"),
+ ("pk,salary,s", "s.missing", Map.empty[String, String], reqAttrsOperationClass,
+ "[`s`.`missing` (FIELD_NOT_FOUND)]")
+ ).foreach { case (reqAttrs, clusterBy, props, operationClass, columns) =>
+ withTable(tableNameAsString) {
+ val tableProps = props ++ Map(
+ "column-update-req-attrs" -> reqAttrs, "column-update-cluster-by" -> clusterBy)
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, s STRUCT<c1: INT>, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 1, "s": { "c1": 1 }, "dep": "hr" }
+ |""".stripMargin,
+ tableProps.asJava)
+
+ Seq(Nil, Seq(excludeColumnPruning)).foreach { confs =>
+ withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> operationClass, "columns" -> columns))
+ }
+ }
+ }
+ }
+ }
+
+ test("column-update: an empty write requirement column name is reported as <empty>") {
+ checkError(
+ exception = QueryCompilationErrors.undeclaredWriteRequirementColumnsError(
+ reqAttrsOperationClass, Seq(Seq.empty[String] -> None)).asInstanceOf[AnalysisException],
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[<empty>]"))
+ }
+
+ test("column-update: distribution by an undeclared row ID column is allowed") {
+ Seq("salary = -1", "salary = -1, pk = pk").foreach { assignments =>
+ withTable(tableNameAsString) {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "salary,dep",
+ "column-update-cluster-by", "pk"))
+
+ sql(s"UPDATE $tableNameAsString SET $assignments WHERE pk = 1")
+
+ checkLastUpdatedColumns("salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 200, "hr") :: Nil)
+ }
+ }
+ }
+
+ test("column-update: an undeclared, unreferenced partition column is not read") {
+ createAndInitTableWithReqAttrs("pk,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ executeAndCheckScanExcludes(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1", "dep")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively without adopting the " +
+ "declared spelling") {
+ createAndInitTableWithReqAttrs("PK,SALARY,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "dep"))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: DELETE on a scan that returns unrequested columns is unaffected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update", "true",
+ InMemoryBaseTable.SIMULATE_PARTIAL_COLUMN_PRUNING, "true"))
+
+ sql(s"DELETE FROM $tableNameAsString WHERE pk < 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(2, 200, 20, "software") :: Nil)
+ }
+
+ test("column-update: case-sensitive analysis accepts columns that differ only in case") {
+ withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") {
+ createAndInitTableWithReqAttrs("pk,salary,SALARY,dep",
+ "pk INT NOT NULL, salary INT, SALARY INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "SALARY": 200, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "SALARY", "dep"))
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, -1, 200, "hr") :: Nil)
+ }
+ }
+
+ test("column-update: empty requiredDataAttributes throws AnalysisException") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }""".stripMargin,
+ java.util.Map.of("column-update-empty-req-attrs", "true"))
+
+ // `id = id` assigns nothing, so only the non-empty check can reject the declaration
+ Seq("id = -1", "id = id").foreach { assignment =>
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET $assignment WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ parameters = Map("connector" -> reqAttrsOperationClass))
+ }
+ }
+
+ test("column-update: invalid requiredDataAttributes are rejected") {
+ Seq(
+ ("pk,s.c1,dep", "COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE",
+ "nestedAttributes" -> "[`s`.`c1`]"),
+ ("pk,pk,salary", "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ "duplicateAttributes" -> "[`pk`]"),
+ ("salary,PK,SALARY,pk", "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ "duplicateAttributes" -> "[`salary`, `PK`]"),
+ ("pk,_partition,salary", "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE",
+ "metadataAttributes" -> "[`_partition`]"),
+ ("pk,nonexistent_col,salary", "COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE",
+ "unknownAttributes" -> "[`nonexistent_col`]"),
+ ("pk,salary", "COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS",
+ "missingColumns" -> "[`a.b`]")
+ ).foreach { case (reqAttrs, condition, columns) =>
+ withTable(tableNameAsString) {
+ createAndInitTableWithReqAttrs(reqAttrs,
+ "pk INT NOT NULL, salary INT, s STRUCT<c1: INT, c2: INT>, `a.b` INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "s": { "c1": 1, "c2": 2 }, "a.b": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1, `a.b` = -1 WHERE pk = 1")
+ },
+ condition = condition,
+ parameters = Map("connector" -> reqAttrsOperationClass, columns))
+ }
+ }
+ }
+
+ test("column-update: a declared column with no name parts is reported as <empty>") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+ val t = table
+ val operation = new t.DeltaBasedColumnUpdateOperation(UPDATE) {
+ override def requiredDataAttributes(): Array[NamedReference] =
+ Array(FieldReference(Seq.empty[String]), FieldReference("pk"),
+ FieldReference(Seq.empty[String]))
+ }
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ RewriteRowLevelCommand.resolveRequiredDataAttrs(
+ spark.table(tableNameAsString).queryExecution.analyzed, operation)
+ },
+ condition = "COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> operation.getClass.getName, "unknownAttributes" -> "[<empty>, <empty>]"))
+ }
+
+ test("column-update split: column update schema has the declared columns") {
Review Comment:
Done.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -246,18 +290,55 @@ abstract class RowLevelOperationSuiteBase
assert(scans.isEmpty, "plan must not contain scan relations")
}
+ // the schemas of the scans of the target table, directly or through a row-level operation, in
+ // the plan of the query
+ private def executeAndCollectScanSchemas(query: String): Seq[StructType] = {
+ val executedPlan = executeAndKeepPlan {
+ sql(query)
+ }
+ val targetName = table.name
+ def scansTarget(scan: BatchScanExec): Boolean = scan.table match {
Review Comment:
Done.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -273,6 +339,423 @@ class InMemoryRowLevelOperationTable private (
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
}
+ // Finds the one current row with the given row-ID values, matching a data column by value and
+ // `_partition` by the partition key.
+ private def currentRowByRowId(values: Seq[(String, Any)]): InternalRow = {
+ val matches = dataMap.values.iterator.flatten.flatMap { buf =>
+ buf.rows.filter { row =>
+ values.forall {
+ case (name, value) if name == PartitionKeyColumn.name => buf.keyString() == value.toString
+ case (name, value) =>
+ val i = schema.fieldIndex(name)
+ row.get(i, schema.fields(i).dataType) == value
+ }
+ }
+ }.toSeq
+ matches match {
+ case Seq(row) => row
+ case _ => throw new IllegalStateException(
+ s"Expected one row with row ID $values, found ${matches.size}")
+ }
+ }
+
+ // Rebuilds a full table row for a narrow column-update row: the base row's values (or the
+ // existence default for a column added after the base row was written), overlaid with the
+ // narrow row's columns.
+ private def overlayNarrowRow(
+ baseRow: Option[InternalRow],
+ narrowRow: InternalRow,
+ narrowFieldIdx: Map[String, Int]): InternalRow = {
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ val field = schema.fields(i)
+ val value = if (i < base.numFields) {
+ base.get(i, field.dataType)
+ } else {
+ ResolveDefaultColumns.getExistenceDefaultValue(field)
+ }
+ fullRow.update(i, value)
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ narrowFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ fullRow
+ }
+
+ // The base columns `pk` and `dep`, where the table has them, followed by the updated columns.
+ private def declaredColumns(updatedCols: Seq[NamedReference]): Array[NamedReference] = {
+ val base = Seq("pk", "dep").filter(schema.fieldNames.contains).map(FieldReference(_))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep`, where the table has them, plus whatever columns Spark reports as
+ // being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = declaredColumns(updatedCols)
+
+ private val clusterExprs: Seq[Expression] = {
+ val refs = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY))
+ .map(_.split(",").map(name => FieldReference(name.trim)).toSeq)
+ .getOrElse(Seq(PARTITION_COLUMN_REF))
+ if (properties.getOrDefault(COLUMN_UPDATE_CLUSTER_BY_NESTED_TRANSFORM, "false") == "true") {
+ refs.map(ref => LogicalExpressions.apply("f", LogicalExpressions.apply("g", ref)))
+ } else {
+ refs
+ }
+ }
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ val writeSchema = info.columnUpdateSchema().orElse(info.schema())
+ val batchWrite = new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ physicalInfo: PhysicalWriteInfo): DeltaWriterFactory = {
+ newWriterFactory(info, writeSchema)
+ }
+
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ commitColumnUpdate(info, writeSchema, newData)
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ if (noMetadata && !properties.containsKey(COLUMN_UPDATE_CLUSTER_BY)) {
+ new DeltaWrite {
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ } else {
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(clusterExprs.toArray)
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ clusterExprs.map { expr =>
+ LogicalExpressions.sort(
+ expr,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()): SortOrder
+ }.toArray
+ }
+
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ }
+ }
+ }
+
+ protected def newWriterFactory(
+ info: LogicalWriteInfo,
+ writeSchema: StructType): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(writeSchema)
+ }
+
+ // For column-update writes, updated rows contain only the declared columns
+ // (columnUpdateSchema from LogicalWriteInfo). We expand each row to the full
+ // table schema by overlaying them on the base row found by pk.
+ protected def commitColumnUpdate(
+ info: LogicalWriteInfo,
+ writeSchema: StructType,
+ newData: Array[BufferedRows]): Unit = {
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+ val updateOpName = UTF8String.fromString(Update.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+
+ val mergedData = newData.map { buf =>
+ val merged = new BufferedRows(buf.key, schema)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == updateOpName) {
+ val pk = logRow.getInt(1)
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val baseRow = Some(currentRowByRowId(Seq("pk" -> pk)))
+ val fullRow = overlayNarrowRow(baseRow, narrowRow, writeFieldIdx)
+ merged.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ merged.rows.append(insertRow.copy())
+ }
+ }
+ merged
+ }
+
+ withDeletes(newData)
+ withData(mergedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+ }
+
+ class DeltaBasedColumnUpdateOperationWithReqAttrs(
+ command: Command,
+ reqCols: Array[String],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
Review Comment:
Done.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4402822957
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,111 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, updated, copied, and reinserted
+ * rows follow it: group-based operations receive them through
+ * {@link DataWriter#writeColumnUpdate(Object, Object)} or
+ * {@link DataWriter#writeColumnUpdate(Object)}, and operations that mix in {@link SupportsDelta}
+ * through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert}. Rows the command inserts
+ * always follow {@link LogicalWriteInfo#schema()} and arrive through
+ * {@link DataWriter#write(Object)} and {@link DeltaWriter#insert}. When
+ * {@link LogicalWriteInfo#columnUpdateSchema()} is absent, every row follows
+ * {@link LogicalWriteInfo#schema()}, as for any other operation. Connectors must decide which
+ * layout to expect by whether {@link LogicalWriteInfo#columnUpdateSchema()} is present, not by
+ * the command. A builder that wants full-width rows for some command can check
+ * {@link RowLevelOperationInfo#command()} and build an operation without this mix-in. Currently
+ * Spark narrows only UPDATE.
+ * <p>
+ * The scan builder returned by {@link #newScanBuilder} should implement
+ * {@link org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns}, so that the scan
+ * does not read columns the command neither references nor declares. Otherwise the scan reads
+ * every column. For operations that do not mix in {@link SupportsDelta}, only top-level columns
+ * are pruned: a struct column the command references through a nested field is read in full.
+ *
+ * @since 4.4.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * When Spark narrows the write, the returned columns become
+ * {@link LogicalWriteInfo#columnUpdateSchema()}, in declared order. Spark matches each declared
+ * column with the session's resolver, and {@link LogicalWriteInfo#columnUpdateSchema()} uses the
+ * table's spelling of each column, which can differ in case from the declared reference.
+ * Implementations must include every column they want to receive: every column reported by
+ * {@link RowLevelOperationInfo#updatedColumns()}, plus any columns needed for row lookup or
+ * routing. Columns that are not declared are absent from the rows the connector receives, so
+ * the connector must preserve their values itself. Spark may call this method more than once
+ * for the same operation, and every call must return the same columns.
+ * <p>
+ * If {@link RowLevelOperationInfo#updatedColumns()} is empty, the builder should build an
+ * operation without this mix-in.
Review Comment:
Superseded by Dongjoon's item 55. An empty declaration now means Spark does not narrow, so builders no longer drop the mix-in. The Javadoc says so.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -343,3 +353,63 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
}
}
}
+
+object RewriteRowLevelCommand extends SQLConfHelper {
+
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation, in declared order and with the relation's spelling. An empty array,
+ * nested fields, duplicate columns, columns that do not exist in the relation (including a
+ * reference with no name parts) and metadata columns are rejected with an `AnalysisException`.
+ */
+ private[sql] def resolveRequiredDataAttrs(
+ relation: LogicalPlan,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes.toImmutableArraySeq
+ if (refs.isEmpty) {
Review Comment:
Done.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala:
##########
@@ -110,8 +111,10 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla
// this rule assigns runtime filters to both scan relations (will be shared at runtime)
// and must transform the runtime filter condition to use correct expr IDs for each relation
// note this only applies to group-based row-level operations (i.e. ReplaceData)
+ // the map is built from the original table as the write table may be narrower than the
Review Comment:
Done. Master added the metadata mapping on this line, and the merge keeps it.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala:
##########
@@ -1500,6 +1501,21 @@ private class BufferWriter(schema: StructType) extends DataWriter[InternalRow] {
buffer.log.append(logEntry)
}
+ // Tag UPDATE/COPY rows distinctly so tests can verify that Spark dispatches
+ // through writeColumnUpdate (the SupportsColumnUpdates path) rather than write.
+ override def writeColumnUpdate(metadata: InternalRow, row: InternalRow): Unit = {
Review Comment:
Done, both options. Mutating the strategy or `DataWithProjectionWritingSparkTask` now fails tests.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/ColumnUpdateTableSharedTests.scala:
##########
@@ -0,0 +1,387 @@
+/*
+ * 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 scala.jdk.CollectionConverters._
+
+import org.apache.spark.{SparkException, SparkRuntimeException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.{InMemoryBaseTable, InMemoryRowLevelOperationTable}
+import org.apache.spark.sql.connector.write.{RowLevelOperation, RowLevelOperationTable}
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+// Column-update tests that hold for both delta-based and group-based operations; each suite
+// supplies its column-update table properties through `extraTableProps`.
+trait ColumnUpdateTableSharedTests extends UpdateTableSuiteBase {
+
+ // the properties of a column-update table whose connector declares exactly `reqAttrs`
+ protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String]
+
+ // the properties of a table of the same kind whose operation does not mix in column updates
+ protected def fullWidthTableProps: java.util.Map[String, String]
+
+ // the metadata columns the scan of a column-update table reads, in DDL form
+ protected def scanMetadataColumns: String = ", _partition STRING, index INT"
+
+ // the class name of the operation that `reqAttrsTableProps` builds
+ protected def reqAttrsOperationClass: String = if (deltaUpdate) {
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateOperationWithReqAttrs].getName
+ } else {
+ classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperationWithReqAttrs].getName
+ }
+
+ // `base` with the given table properties added; a given key replaces the same key in `base`
+ protected def tablePropsWith(
+ base: java.util.Map[String, String],
+ props: (String, String)*): java.util.Map[String, String] = {
+ (base.asScala ++ props).asJava
+ }
+
+ // t(pk, id, dep, salary, bonus, extra) with CHECK (extra > 0), partitioned by dep
+ protected def createAndInitTableWithProps(props: java.util.Map[String, String]): Unit = {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, salary INT, bonus INT, extra INT",
+ """{ "pk": 1, "id": 1, "dep": "hr", "salary": 100, "bonus": 10, "extra": 1 }
+ |{ "pk": 2, "id": 2, "dep": "hr", "salary": 200, "bonus": 20, "extra": 1 }
+ |{ "pk": 3, "id": 3, "dep": "software", "salary": 300, "bonus": 30, "extra": 1 }
+ |""".stripMargin,
+ props)
+ sql(s"ALTER TABLE $tableNameAsString ADD CONSTRAINT positive_extra CHECK (extra > 0)")
+ }
+
+ test("column-update: undeclared columns are not written and keep their values") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkLastColumnUpdateSchema("pk", "dep", "salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update: write relation has the declared columns in declared order") {
+ createAndInitTableWithProps(reqAttrsTableProps("salary,pk"))
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'")
+ }
+
+ assert(!write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "salary", "pk")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update: the plan writes the declared columns and scans the referenced ones") {
+ createAndInitTableWithProps(reqAttrsTableProps("pk,salary"))
+
+ // `id` is neither declared nor referenced; `dep`, `bonus` and `extra` are read for the
+ // condition, the assignment and the CHECK constraint
+ val write = executeAndKeepAnalyzedWrite {
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'",
+ s"pk INT, dep STRING, salary INT, bonus INT, extra INT$scanMetadataColumns")
+ }
+
+ val writeNodeName = if (deltaUpdate) "WriteDelta" else "ReplaceData"
+ val header = write.simpleString(Int.MaxValue)
+ assert(header.matches(s"""$writeNodeName RelationV2\\[pk#\\d+, salary#\\d+\\] .*"""), header)
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update: undeclared, differently-cased columns in the RHS and WHERE clause are " +
+ "read with the table's own spelling") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, extra INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "extra": 5, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "extra": 5, "dep": "hr" }
+ |""".stripMargin)
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + BONUS WHERE pk = 1 AND EXTRA > 3",
+ s"pk INT, salary INT, bonus INT, extra INT, dep STRING$scanMetadataColumns")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 110, 10, 5, "hr") :: Row(2, 200, 20, 5, "hr") :: Nil)
+ }
+
+ test("column-update: assigned column outside requiredDataAttributes is rejected") {
+ createAndInitTableWithProps(reqAttrsTableProps("pk,salary"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET bonus = 0 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "missingColumns" -> "[`bonus`]"))
+ }
+
+ test("column-update: an operation replaced after the rewrite is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ }
+ val relation = write.table.asInstanceOf[DataSourceV2Relation]
+ val operationTable = relation.table.asInstanceOf[RowLevelOperationTable]
+ val t = operationTable.table.asInstanceOf[InMemoryRowLevelOperationTable]
+ val replacements: Seq[(RowLevelOperation, String)] = if (deltaUpdate) {
+ Seq(
+ t.DeltaBasedOperation(UPDATE, CaseInsensitiveStringMap.empty()) -> "[pk, salary, dep]",
+ new t.DeltaBasedColumnUpdateOperation(UPDATE) -> "[pk, dep]")
+ } else {
+ Seq(
+ t.PartitionBasedOperation(UPDATE, CaseInsensitiveStringMap.empty()) -> "[pk, salary, dep]",
+ new t.PartitionBasedColumnUpdateOperation(UPDATE) -> "[pk, dep]")
+ }
+ replacements.foreach { case (operation, expectedNames) =>
+ assert(operation != operationTable.operation)
+ val replaced = write.withNewTable(
+ relation.copy(table = operationTable.copy(operation = operation)(
+ operationTable.operationInfo)))
+ checkError(
+ exception = intercept[SparkException](V2Writes(replaced)),
+ condition = "INTERNAL_ERROR",
+ parameters = Map("message" ->
+ (s"The write relation of operation ${operation.getClass.getName} has columns " +
+ s"[pk, dep, salary], but the operation expects $expectedNames.")))
+ }
+ }
+
+ test("column-update: nested struct field update narrows to the root struct column") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |{ "pk": 2, "s": { "c1": 3, "c2": 4 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+
+ checkLastUpdatedColumns("s")
+ checkLastColumnUpdateSchema("pk", "dep", "s")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, Row(-1, 2), "hr") :: Row(2, Row(3, 4), "hr") :: Nil)
+ }
+
+ test("column-update: excluding ColumnPruning keeps the write narrow") {
+ createAndInitTableWithProps(reqAttrsTableProps("pk,salary"))
+
+ withSQLConf(SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName) {
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'",
+ s"pk INT, id INT, dep STRING, salary INT, bonus INT, extra INT$scanMetadataColumns")
+ }
+
+ checkLastColumnUpdateSchema("pk", "salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update: reassigning a declared row-ID column keeps the row's other columns") {
+ createAndInitTableWithProps(reqAttrsTableProps("pk,salary"))
+
+ sql(s"UPDATE $tableNameAsString SET pk = 2, salary = 0 WHERE pk = 1")
+
+ // the copy-on-write fixture finds original rows by `pk`, so it falls back to full-width rows
+ assert(table.lastWriteInfo.columnUpdateSchema().isPresent == deltaUpdate)
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk, id"),
+ Row(2, 1, "hr", 0, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update: an UPDATE whose assignments are all identity assignments is narrowed " +
+ "only if the operation keeps the mix-in") {
+ Seq(
+ extraTableProps -> None,
+ reqAttrsTableProps("pk,salary") -> Some(Seq("pk", "salary"))
+ ).foreach { case (props, columnUpdateNames) =>
+ withTable(tableNameAsString) {
+ createAndInitTableWithProps(props)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary, dep = dep WHERE pk = 1")
+
+ checkLastUpdatedColumns()
+ columnUpdateNames match {
+ case Some(names) => checkLastColumnUpdateSchema(names: _*)
+ case None =>
+ assert(!table.lastWriteInfo.columnUpdateSchema().isPresent,
Review Comment:
Done.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -128,12 +151,61 @@ class InMemoryRowLevelOperationTable private (
copied.replacedPartitions = replacedPartitions
copied.lastWriteInfo = lastWriteInfo
copied.lastWriteLog = lastWriteLog
+ copied.lastUpdatedColumns = lastUpdatedColumns
copied
}
override def newRowLevelOperationBuilder(
info: RowLevelOperationInfo): RowLevelOperationBuilder = {
- if (properties.getOrDefault(SUPPORTS_DELTAS, "false") == "true") {
+ lastUpdatedColumns = info.updatedColumns()
+ val splitUpdates = properties.getOrDefault(SPLIT_UPDATES, "false").toBoolean
+ // the copy-on-write commit finds the original row by `pk`
+ val updatesPk = info.updatedColumns().exists(_.describe() == "pk")
+ // The default column-update fixtures drop the mix-in for an UPDATE with no updated columns,
+ // as SupportsColumnUpdates asks. The fixed-declaration fixtures, the split row-ID fixture and
+ // all fixtures for DELETE and MERGE keep it.
+ val updatesNothing = info.command == Command.UPDATE && info.updatedColumns().isEmpty
+ if (properties.getOrDefault(COLUMN_UPDATE, "false") == "true") {
+ if (updatesNothing) {
+ () => DeltaBasedOperation(info.command, info.options)
+ } else if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ }
+ } else if (properties.containsKey(COLUMN_UPDATE_REQ_ATTRS)) {
+ val reqCols = properties.get(COLUMN_UPDATE_REQ_ATTRS).split(",").map(_.trim)
+ if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperationWithReqAttrs(
+ info.command, reqCols, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperationWithReqAttrs(info.command, reqCols, info.options)
+ }
+ } else if (properties.getOrDefault(COLUMN_UPDATE_EMPTY_REQ_ATTRS, "false") == "true") {
Review Comment:
Done through 55.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala:
##########
@@ -109,22 +113,98 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper {
WriteToDataSourceV2(Some(r), microBatchWrite, newQuery, customMetrics)
case rd @ ReplaceData(r: DataSourceV2Relation, _, query, _, projections, _, None) =>
- val rowSchema = projections.rowProjection.schema
+ validateWriteRelation(rd)
+ val (rowSchema, columnUpdateSchema) =
+ writeSchemas(projections.rowProjection.schema, hasNarrowRows(rd))
val metadataSchema = projections.metadataProjection.map(_.schema)
val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap)
- val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema)
+ val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema,
+ columnUpdateSchema)
val write = writeBuilder.build()
+ validateColumnUpdateWrite(rd, write)
val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, r.funCatalog)
rd.copy(write = Some(write), query = newQuery)
case wd @ WriteDelta(r: DataSourceV2Relation, _, query, _, projections, _, None) =>
+ validateWriteRelation(wd)
val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap)
- val deltaWriteBuilder = newDeltaWriteBuilder(r.table, writeOptions, projections)
+ val deltaWriteBuilder =
+ newDeltaWriteBuilder(r.table, writeOptions, projections, hasNarrowRows(wd))
val deltaWrite = deltaWriteBuilder.build()
+ validateColumnUpdateWrite(wd, deltaWrite)
val newQuery = DistributionAndOrderingUtils.prepareQuery(deltaWrite, query, r.funCatalog)
wd.copy(write = Some(deltaWrite), query = newQuery)
}
+ /**
+ * Whether a row-level write delivers narrow rows for updated, copied, and reinserted records.
+ * This must match the writes whose write relation RewriteUpdateTable narrows, even when the
+ * declared columns cover the whole table.
+ */
+ private[v2] def hasNarrowRows(write: RowLevelWrite): Boolean = {
+ write.operation.isInstanceOf[SupportsColumnUpdates] && write.operation.command == UPDATE
+ }
+
+ /**
+ * Checks that the write relation matches the current operation: a write that delivers narrow
+ * rows must have exactly the declared columns, in declared order, and any other write the full
+ * table columns. A rule that replaces the operation after the rewrite can break this, and the
+ * connector would then decode the rows with the wrong layout.
+ */
+ private def validateWriteRelation(rowLevelWrite: RowLevelWrite): Unit = {
+ val writeNames = rowLevelWrite.table.output.map(_.name)
+ val expectedNames = rowLevelWrite.operation match {
+ case op: SupportsColumnUpdates if hasNarrowRows(rowLevelWrite) =>
+ RewriteRowLevelCommand.resolveRequiredDataAttrs(rowLevelWrite.originalTable, op)
+ .map(_.name)
+ case _ =>
+ rowLevelWrite.originalTable.output.map(_.name)
+ }
+ if (writeNames != expectedNames) {
+ throw SparkException.internalError(
+ s"The write relation of operation ${rowLevelWrite.operation.getClass.getName} has " +
+ s"columns ${writeNames.mkString("[", ", ", "]")}, but the operation expects " +
+ s"${expectedNames.mkString("[", ", ", "]")}.")
+ }
+ }
+
+ /**
+ * Checks that the distribution and ordering of a write that delivers narrow rows reference only
+ * columns passed to the write. Other columns stay in the query only until column pruning removes
+ * them, so they are rejected here whether or not pruning ran.
+ */
+ private def validateColumnUpdateWrite(rowLevelWrite: RowLevelWrite, write: Write): Unit = {
+ write match {
+ case w: RequiresDistributionAndOrdering if hasNarrowRows(rowLevelWrite) =>
+ val distribution = w.requiredDistribution match {
+ case d: ClusteredDistribution => d.clustering.toImmutableArraySeq
+ case d: OrderedDistribution => d.ordering.toImmutableArraySeq
+ case _ => Seq.empty
Review Comment:
Done.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RowLevelWriteColumnPruningSuite.scala:
##########
@@ -0,0 +1,278 @@
+/*
+ * 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.catalyst.optimizer
+
+import java.util
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.ProjectingInternalRow
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, EqualNullSafe, Expression, If, Literal, Not}
+import org.apache.spark.sql.catalyst.plans.{LeftOuter, PlanTest}
+import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.plans.logical.MergeRows.{Copy, Keep, Update}
+import org.apache.spark.sql.catalyst.rules.RuleExecutor
+import org.apache.spark.sql.catalyst.util.{ReplaceDataProjections, WriteDeltaProjections}
+import org.apache.spark.sql.catalyst.util.RowDeltaUtils._
+import org.apache.spark.sql.connector.catalog.{Column, Table, TableCapability}
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Tests the [[ColumnPruning]] case that prunes the query columns of a [[RowLevelWrite]] that its
+ * projections don't read and moves the projections to the ordinals of the pruned query.
+ */
+class RowLevelWriteColumnPruningSuite extends PlanTest {
+
+ object Optimize extends RuleExecutor[LogicalPlan] {
+ val batches = Batch("Column Pruning", FixedPoint(10),
+ ColumnPruning, CollapseProject, RemoveNoopOperators) :: Nil
+ }
+
+ // t(pk, id, dep, salary, bonus, extra) with the metadata column _partition
+ private val read = LocalRelation(
+ $"pk".int.notNull, $"id".int, $"dep".string, $"salary".int, $"bonus".int, $"extra".int,
+ $"_partition".string.notNull)
+
+ // the write and original table relations; the pruning never inspects them
+ private val tableRelation: DataSourceV2Relation = {
Review Comment:
Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -228,4 +244,96 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Returns the table attributes that these assignments set to a new value, skipping identity
+ * assignments.
+ */
+ private def collectUpdatedAttrs(assignments: Seq[Assignment]): Seq[AttributeReference] = {
+ assignments.collect {
+ case Assignment(key: AttributeReference, value) if !isIdentityAssignment(key, value) => key
Review Comment:
Handled by SPARK-60721 on master, which rejects assignments to metadata columns at analysis. Nothing left to do here.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -273,6 +345,440 @@ class InMemoryRowLevelOperationTable private (
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
}
+ // Finds the one current row with the given row-ID values, matching a data column by value and
+ // `_partition` by the partition key.
+ private def currentRowByRowId(values: Seq[(String, Any)]): InternalRow = {
+ val matches = dataMap.values.iterator.flatten.flatMap { buf =>
+ buf.rows.filter { row =>
+ values.forall {
+ case (name, value) if name == PartitionKeyColumn.name => buf.keyString() == value.toString
+ case (name, value) =>
+ val i = schema.fieldIndex(name)
+ row.get(i, schema.fields(i).dataType) == value
+ }
+ }
+ }.toSeq
+ matches match {
+ case Seq(row) => row
+ case _ => throw new IllegalStateException(
+ s"Expected one row with row ID $values, found ${matches.size}")
+ }
+ }
+
+ // Rebuilds a full table row for a narrow column-update row: the base row's values (or the
+ // existence default for a column added after the base row was written), overlaid with the
+ // narrow row's columns.
+ private def overlayNarrowRow(
+ baseRow: Option[InternalRow],
+ narrowRow: InternalRow,
+ narrowFieldIdx: Map[String, Int]): InternalRow = {
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ val field = schema.fields(i)
+ val value = if (i < base.numFields) {
+ base.get(i, field.dataType)
Review Comment:
Done, except moving `pk`. That case fails on the full-width fixture too, because `InMemoryBaseTable.withDeletes` reads the key at stored ordinal 0.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -72,10 +72,12 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
protected def buildOperationTable(
table: SupportsRowLevelOperations,
command: Command,
- options: CaseInsensitiveStringMap): RowLevelOperationTable = {
- val info = RowLevelOperationInfoImpl(command, options)
+ options: CaseInsensitiveStringMap,
+ updatedAttrs: Seq[AttributeReference] = Nil): RowLevelOperationTable = {
+ val updatedColumns = updatedAttrs.map(attr => FieldReference(Seq(attr.name)))
Review Comment:
Done.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -273,6 +345,440 @@ class InMemoryRowLevelOperationTable private (
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
}
+ // Finds the one current row with the given row-ID values, matching a data column by value and
+ // `_partition` by the partition key.
+ private def currentRowByRowId(values: Seq[(String, Any)]): InternalRow = {
+ val matches = dataMap.values.iterator.flatten.flatMap { buf =>
+ buf.rows.filter { row =>
+ values.forall {
+ case (name, value) if name == PartitionKeyColumn.name => buf.keyString() == value.toString
+ case (name, value) =>
+ val i = schema.fieldIndex(name)
+ row.get(i, schema.fields(i).dataType) == value
+ }
+ }
+ }.toSeq
+ matches match {
+ case Seq(row) => row
+ case _ => throw new IllegalStateException(
+ s"Expected one row with row ID $values, found ${matches.size}")
+ }
+ }
+
+ // Rebuilds a full table row for a narrow column-update row: the base row's values (or the
+ // existence default for a column added after the base row was written), overlaid with the
+ // narrow row's columns.
+ private def overlayNarrowRow(
+ baseRow: Option[InternalRow],
+ narrowRow: InternalRow,
+ narrowFieldIdx: Map[String, Int]): InternalRow = {
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ val field = schema.fields(i)
+ val value = if (i < base.numFields) {
+ base.get(i, field.dataType)
+ } else {
+ ResolveDefaultColumns.getExistenceDefaultValue(field)
+ }
+ fullRow.update(i, value)
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ narrowFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ fullRow
+ }
+
+ // The base columns `pk` and `dep`, where the table has them, followed by the updated columns.
+ private def declaredColumns(updatedCols: Seq[NamedReference]): Array[NamedReference] = {
+ val base = Seq("pk", "dep").filter(schema.fieldNames.contains).map(FieldReference(_))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep`, where the table has them, plus whatever columns Spark reports as
+ // being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = declaredColumns(updatedCols)
+
+ // The case-class equality of DeltaBasedOperation would make this operation equal to a wide
+ // one, or to another column-update operation, with the same command and options.
+ override def equals(other: Any): Boolean = other match {
+ case ref: AnyRef => this eq ref
+ case _ => false
+ }
+ override def hashCode(): Int = System.identityHashCode(this)
+ override def canEqual(other: Any): Boolean = false
+
+ private val clusterExprs: Seq[Expression] = {
+ val refs = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY))
+ .map(_.split(",").map(name => FieldReference(name.trim)).toSeq)
+ .getOrElse(Seq(PARTITION_COLUMN_REF))
+ if (properties.getOrDefault(COLUMN_UPDATE_CLUSTER_BY_NESTED_TRANSFORM, "false") == "true") {
+ refs.map(ref => LogicalExpressions.apply("f", LogicalExpressions.apply("g", ref)))
+ } else {
+ refs
+ }
+ }
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ val writeSchema = info.columnUpdateSchema().orElse(info.schema())
+ val batchWrite = new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ physicalInfo: PhysicalWriteInfo): DeltaWriterFactory = {
+ newWriterFactory(info, writeSchema)
+ }
+
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ commitColumnUpdate(info, writeSchema, newData)
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ if (noMetadata && !properties.containsKey(COLUMN_UPDATE_CLUSTER_BY)) {
+ new DeltaWrite {
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ } else {
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(clusterExprs.toArray)
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ clusterExprs.map { expr =>
+ LogicalExpressions.sort(
+ expr,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()): SortOrder
+ }.toArray
+ }
+
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ }
+ }
+ }
+
+ protected def newWriterFactory(
+ info: LogicalWriteInfo,
+ writeSchema: StructType): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(writeSchema)
+ }
+
+ // For column-update writes, updated rows contain only the declared columns
+ // (columnUpdateSchema from LogicalWriteInfo). We expand each row to the full
+ // table schema by overlaying them on the base row found by pk.
+ protected def commitColumnUpdate(
+ info: LogicalWriteInfo,
+ writeSchema: StructType,
+ newData: Array[BufferedRows]): Unit = {
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
Review Comment:
Done.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub) via reviews" <re...@spark.apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4428927820
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,114 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, updated, copied, and reinserted
+ * rows follow it: group-based operations receive them through
+ * {@link DataWriter#writeColumnUpdate(Object, Object)} or
+ * {@link DataWriter#writeColumnUpdate(Object)}, and operations that mix in {@link SupportsDelta}
+ * through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert}. Rows the command inserts
+ * always follow {@link LogicalWriteInfo#schema()} and arrive through
+ * {@link DataWriter#write(Object)} and {@link DeltaWriter#insert}. When
+ * {@link LogicalWriteInfo#columnUpdateSchema()} is absent, every row follows
+ * {@link LogicalWriteInfo#schema()}, as for any other operation. Connectors must decide which
+ * layout to expect by whether {@link LogicalWriteInfo#columnUpdateSchema()} is present, not by
+ * the command. A builder that wants full-width rows for some command can check
+ * {@link RowLevelOperationInfo#command()} and build an operation without this mix-in. Currently
+ * Spark narrows only UPDATE.
+ * <p>
+ * The scan builder returned by {@link #newScanBuilder} should implement
Review Comment:
In this particular sentence, I am not sure referencing a particular scan API is the right call. We may evolve it in the future. I'd just say that scans must support schema pruning.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub) via reviews" <re...@spark.apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4429185545
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -159,14 +164,24 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
relation: DataSourceV2Relation,
operationTable: RowLevelOperationTable,
assignments: Seq[Assignment],
- cond: Expression): WriteDelta = {
+ cond: Expression,
+ writeAttrs: Seq[AttributeReference],
+ updatedAttrs: Seq[AttributeReference]): WriteDelta = {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
// resolve all needed attrs (e.g. row ID and any required metadata attrs)
- val rowAttrs = relation.output
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
+ val splitUpdates = operation.representUpdateAsDeleteAndInsert
+
+ operation match {
Review Comment:
I think these checks give reinsert a completely different meaning. DeltaWriter explicitly doesn't pass row ID for reinserts so it must be OK for the connector. We will pass the updated row ID to DeltaWriter$delete. That should be enough for a connector that represents updates as deletes and inserts.
Am I missing anything?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192427427
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java:
##########
@@ -105,4 +105,45 @@ default String description() {
default NamedReference[] requiredMetadataAttributes() {
return new NamedReference[0];
}
+
+
+ /**
+ * Controls whether to send only the required data columns to the connector rather than the
+ * full row.
+ * <p>
+ * When true, Spark narrows the data column schema ({@link LogicalWriteInfo#schema()}) to only
+ * the columns declared via {@link #requiredDataAttributes()}. Metadata columns (from
+ * {@link #requiredMetadataAttributes()}) and row ID columns (from
+ * {@link SupportsDelta#rowId()}) are unaffected and always projected separately.
+ * <p>
+ * If {@link #requiredDataAttributes()} returns a non-empty array, the write schema is exactly
+ * those columns in declared order. The connector must include all columns it wants to receive,
+ * including the columns being updated. If {@link #requiredDataAttributes()} returns an empty
+ * array, Spark sends only the non-identity assigned columns (heuristic path).
+ *
+ * @since 4.2.0
+ */
+ default boolean supportsColumnUpdates() {
Review Comment:
Given the scope of this PR, shall we mention that `DELETE and MERGE ignores this method`?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192402487
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java:
##########
@@ -37,4 +38,20 @@ public interface RowLevelOperationInfo {
* Returns the row-level SQL command (e.g. DELETE, UPDATE, MERGE).
*/
Command command();
+
+ /**
+ * Returns the columns being updated in an UPDATE statement, as non-identity assignments.
+ *
+ * <p>For DELETE and MERGE, returns an empty array.
+ *
+ * <p>Connectors can use this to decide what {@link RowLevelOperation#requiredDataAttributes()}
+ * to declare. For instance, a connector that needs its primary key for row lookup can check
+ * whether pk is already in the updated columns list and, if not, add it to
+ * requiredDataAttributes().
+ *
+ * @since 4.2.0
Review Comment:
4.3.0
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-4384373728
I finished the first round review, @anuragmantri .
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3633929809
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,540 @@
+/*
+ * 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, TableInfo}
+import org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity, reference}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+
+/**
+ * Tests for UPDATE statements targeting connectors that return true from
+ * [[org.apache.spark.sql.connector.write.RowLevelOperation#supportsColumnUpdates]].
+ *
+ * When a connector supports column updates, Spark narrows the row projection
+ * (LogicalWriteInfo.schema()) to contain only the assigned/changed columns rather than
+ * the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends RowLevelOperationSuiteBase {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ test("column-update: rowSchema contains only the single assigned column") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ // info.schema() is empty: UPDATE-only column-update writes have no INSERT-shaped rows,
+ // so nothing flows through the writer's write() path. The narrow update-row layout is
+ // carried by info.updateSchema().
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: rowSchema contains multiple assigned columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("id", IntegerType, nullable = false),
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: rowSchema is empty for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ // All assignments are identity, so updatedColumns is empty -- the connector still declares
+ // pk for row lookup, so the narrow update schema is just [pk].
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Array(PK_FIELD))))
+ }
+
+ test("column-update: row filter condition is orthogonal to column narrowing") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = 'engineering' WHERE pk IN (1, 3)")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: update all rows (no WHERE clause)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("salary", IntegerType, nullable = true)
+ ))))
+ }
+
+ test("column-update: rowSchema excludes identity assignments in a mixed UPDATE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: cross-column assignment is not treated as identity") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = dep, id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: nested struct field update narrows to the root struct column") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("s"),
+ s"expected [s] in updatedColumns (root struct) but got: $updatedNames")
+
+ // info.updateSchema() carries the narrow row layout; info.schema() is the full table.
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("s"),
+ s"s must be in update schema: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("dep"),
+ s"dep must not be in update schema: $updateSchema")
+ }
+
+ test("column-update: nested field identity update reports root struct as updated") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = s.c1 WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("s"),
+ s"nested identity is reported as an update at root-column granularity: $updatedNames")
+
+ // Data correctness: the struct is rewritten but with equal values, so rows are unchanged.
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, Row(1, 2), "hr") :: Nil)
+ }
+
+ test("column-update: updatedColumns contains non-identity assigned columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'eng' WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("id", "dep"),
+ s"expected [id, dep] in updatedColumns but got: $updatedNames")
+ }
+
+ test("column-update: updatedColumns excludes identity assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = dep WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("id"),
+ s"expected only [id] in updatedColumns but got: $updatedNames")
+ }
+
+ test("column-update: updatedColumns is empty for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ assert(table.lastUpdatedColumns.isEmpty,
+ s"expected empty updatedColumns but got: ${table.lastUpdatedColumns.mkString(", ")}")
+ }
+
+ test("column-update: updatedColumns is empty for DELETE (Javadoc contract)") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"DELETE FROM $tableNameAsString WHERE dep = 'hr'")
+
+ assert(table.lastUpdatedColumns.isEmpty,
+ s"DELETE must pass empty updatedColumns but got: ${table.lastUpdatedColumns.mkString(", ")}")
+ }
+
+ test("column-update: data correctness -- single column update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ // 1 row matched WHERE pk = 1 -> 1 update. MoR emits no COPY rows.
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0, deltaUpdate = true)
+ }
+
+ test("column-update: data correctness -- update all rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 200, "hr") :: Row(2, 400, "software") :: Row(3, 600, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 3, numCopiedRows = 0, deltaUpdate = true)
+ }
+
+ test("column-update: data correctness -- mixed identity and real assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "engineering") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0, deltaUpdate = true)
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-req-attrs", reqAttrs)
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update: requiredDataAttributes - data correctness") {
+ createAndInitTableWithReqAttrs("dep,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: empty requiredDataAttributes throws AnalysisException") {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-empty-req-attrs", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(
+ StructType.fromDDL("pk INT NOT NULL, id INT, dep STRING"))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }""".stripMargin)
+
+ // A connector that mixes in SupportsColumnUpdates but returns an empty
+ // requiredDataAttributes() violates the mix-in contract -- analysis must reject the
+ // operation rather than silently falling through to the wide write path.
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ }
+ assert(ex.getCondition == "EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ s"expected EMPTY_REQUIRED_DATA_ATTRIBUTES but got: ${ex.getCondition}")
+ }
+
+ test("column-update: requiredDataAttributes throws AnalysisException for invalid column") {
+ createAndInitTableWithReqAttrs("nonexistent_col", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ }
+ assert(ex.getMessage.contains("nonexistent_col"),
+ s"Expected error about unresolvable column but got: ${ex.getMessage}")
+ }
+
+ private def createAndInitTableFromInfo(schemaString: String, jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-from-info", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update from-info: write schema is updatedColumns + pk pass-through") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("salary"),
+ s"salary must be in update schema: $updateSchema")
+ assert(updateSchema.fieldNames.contains("pk"),
+ s"pk must be in update schema: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("id"),
+ s"id must not be in update schema: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("dep"),
+ s"dep must not be in update schema: $updateSchema")
+ }
+
+ test("column-update from-info: pk already in updatedColumns is not duplicated") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ val pkCount = updateSchema.fieldNames.count(_ == "pk")
+ assert(pkCount == 1, s"pk must appear exactly once in update schema: $updateSchema")
+ }
+
+ test("column-update from-info: data correctness") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "id": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ }
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-split", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update split: write schema is narrow (assigned + pk pass-through)") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ // For representUpdateAsDeleteAndInsert connectors, reinsert-tagged rows still route through
+ // writeUpdate() (info.updateSchema() = narrow), not write(). info.schema() stays empty.
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update split: data correctness") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, -1, "hr") :: Nil)
+ }
+
+ // ---------------------------------------------------------------------------
+ // Scan-narrowing tests: verify that when the connector opts into column updates,
+ // the physical scan reads only the columns actually needed by the operation.
+ //
+ // Assertions use `checkLastScanExcludes` (proves narrowing) and `checkLastScanIncludes`
+ // (proves transparent widening) rather than exact-set matches -- ColumnPruning is
+ // free to further tighten the scan beyond our analysis-time narrow set (e.g. when an
+ // assignment RHS is a literal, the assigned column doesn't need to be read).
+ // ---------------------------------------------------------------------------
+
+ test("column-update: scan excludes columns outside connector-declared + cond + RHS refs") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ // requiredDataAttributes = [pk, id]; cond refs [pk] only; RHS is a literal.
+ // `dep` is neither declared nor referenced -- it must not appear in the scan.
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastScanExcludes("dep")
+ }
+
+ test("column-update: scan transparently widens for cond-referenced columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ // requiredDataAttributes = [pk, id]; cond refs [dep].
+ // `dep` must appear even though the connector didn't declare it.
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE dep = 'hr'")
+
+ checkLastScanIncludes("dep")
+ }
+
+ test("column-update: scan transparently widens for assignment RHS references") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin)
+
+ // requiredDataAttributes = [pk, salary]; RHS `salary + bonus` references `bonus`.
+ // `bonus` must appear in the scan even though it's not declared as required.
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE pk = 1")
+
+ checkLastScanIncludes("bonus", "salary")
+ checkLastScanExcludes("dep")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 110, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update: analysis fails when assignment key is outside requiredDataAttributes") {
+ // Connector declares only [pk] but the user assigns to `id`. We enforce
+ // updatedColumns ⊆ requiredDataAttributes at analysis time (root-column granularity).
Review Comment:
Is this auto-generated sentence, @anuragmantri ? AFAIK, Apache Spark CI will fail for non-ASCII characters like this.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-5225523711
@anuragmantri, I believe you forgot to push the new commit.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3742227035
##########
common/utils/src/main/resources/error/error-conditions.json:
##########
@@ -2461,6 +2467,12 @@
],
"sqlState" : "KD009"
},
+ "EMPTY_REQUIRED_DATA_ATTRIBUTES" : {
Review Comment:
nit. The five new error conditions for the same `SupportsColumnUpdates` contract use three different prefixes (`EMPTY_REQUIRED_DATA_ATTRIBUTES`, `REQUIRED_DATA_ATTRIBUTES_*`, `SPLIT_UPDATE_*`). Could you consider a common prefix or a single parent condition with sub-conditions so they group together and are easier to discover?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, any required metadata attrs and optionally connector
+ // declared attrs)
val rowAttrs = relation.output
+ val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates]
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveConnectorDataAttrs(relation, operation)
+ } else Nil
+ val scanOnlyDataAttrs = if (supportsColumnUpdate) {
+ resolveScanOnlyDataAttrs(relation, operation)
+ } else Nil
+
+ if (supportsColumnUpdate) {
+ validateUpdatedColumnsSubset(operation, assignments, connectorDataAttrs)
+ validateNoOverlap(operation, connectorDataAttrs, scanOnlyDataAttrs)
+ validatePartitionAttrsDeclared(operation, relation, connectorDataAttrs, scanOnlyDataAttrs)
+ }
+
+
Review Comment:
nit. Redundant empty line.
```suggestion
```
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, any required metadata attrs and optionally connector
+ // declared attrs)
val rowAttrs = relation.output
+ val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates]
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveConnectorDataAttrs(relation, operation)
+ } else Nil
+ val scanOnlyDataAttrs = if (supportsColumnUpdate) {
+ resolveScanOnlyDataAttrs(relation, operation)
+ } else Nil
+
+ if (supportsColumnUpdate) {
+ validateUpdatedColumnsSubset(operation, assignments, connectorDataAttrs)
+ validateNoOverlap(operation, connectorDataAttrs, scanOnlyDataAttrs)
+ validatePartitionAttrsDeclared(operation, relation, connectorDataAttrs, scanOnlyDataAttrs)
+ }
+
+
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
- val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ if (supportsColumnUpdate && operation.representUpdateAsDeleteAndInsert) {
+ validateNoRowIdReassignment(operation, assignments, rowIdAttrs)
+ validateRowIdDeclared(operation, connectorDataAttrs, rowIdAttrs)
+ }
+
+ val narrowDataAttrs = if (supportsColumnUpdate) {
+ computeNarrowReadAttrs(relation, connectorDataAttrs, scanOnlyDataAttrs, assignments, cond)
+ } else {
+ relation.output
+ }
+
+ val readRelation = if (supportsColumnUpdate) {
+ buildNarrowRelationWithAttrs(relation, operationTable, narrowDataAttrs, metadataAttrs,
+ rowIdAttrs)
+ } else {
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ }
// 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 rowDeltaPlan = if (supportsColumnUpdate) {
+ if (operation.representUpdateAsDeleteAndInsert) {
+ buildNarrowDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else {
+ buildColumnUpdateProjection(
+ matchedRowsPlan, assignments, rowIdAttrs, metadataAttrs, connectorDataAttrs,
+ scanOnlyDataAttrs)
+ }
} else {
- buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs)
+ if (operation.representUpdateAsDeleteAndInsert) {
+ buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else {
+ buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, 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(
+ rowDeltaPlan, rowAttrs, rowIdAttrs, metadataAttrs, connectorDataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond)
}
+ /**
+ * Builds the WriteDelta projection for the column-update path.
+ */
+ private def buildColumnUpdateProjection(
+ plan: LogicalPlan,
+ assignments: Seq[Assignment],
+ rowIdAttrs: Seq[Attribute],
+ metadataAttrs: Seq[Attribute],
+ connectorDataAttrs: Seq[AttributeReference],
+ scanOnlyDataAttrs: Seq[AttributeReference]): LogicalPlan = {
+
+ val assignedValues = assignments.collect {
+ case Assignment(key: Attribute, value) if !isIdentityAssignment(key, value) =>
+ Alias(value, key.name)()
+ }
+
+ // Connector-required columns whose value isn't being changed by the UPDATE: pass through the
+ // current value so the connector receives a complete write row. Row-ID columns are excluded
+ // here even when also connector-required (e.g. a primary key used for both row lookup and
+ // write payload) they are emitted once, below, via rowIdValues.
+ val assignedKeyIds = assignments.collect {
+ case Assignment(key: AttributeReference, value) if !isIdentityAssignment(key, value) =>
+ key.exprId
+ }.toSet
+ val rowIdAttrSet = AttributeSet(rowIdAttrs)
+ val connectorPassThroughValues = connectorDataAttrs.filterNot(a =>
+ assignedKeyIds.contains(a.exprId) || rowIdAttrSet.contains(a))
+
+ // scanOnlyDataAttrs are never assigned carry them through the write query so
Review Comment:
nit. Missing punctuation here and at line 341 (`write payload) they are emitted once` -> `write payload); they are emitted once`).
```suggestion
// scanOnlyDataAttrs are never assigned; carry them through the write query so
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala:
##########
@@ -771,6 +777,7 @@ trait WritingSparkTask[W <: DataWriter[InternalRow]] extends Logging with Serial
case class DataAndMetadataWritingSparkTask(
dataProj: ProjectingInternalRow,
+ updateDataProj: ProjectingInternalRow,
Review Comment:
nit (optional). A nullable `updateDataProj: ProjectingInternalRow` constructor parameter (here and in `DataWithProjectionWritingSparkTask`) is consistent with the existing `orNull` usage in the delta tasks, but `Option[ProjectingInternalRow]` would be more idiomatic for a case class constructor. Feel free to keep as is if you prefer the symmetry.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -146,6 +203,48 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
Project(operationCol +: updatedValues, plan)
}
+ /**
+ * Variant of `buildReplaceDataUpdateProjection` for the `SupportsColumnUpdates` narrow-scan
+ * path.
+ * For narrow attributes, looks up assignments by `ExprId` via `AttributeMap`and passes through
Review Comment:
nit. Missing space.
```suggestion
* For narrow attributes, looks up assignments by `ExprId` via `AttributeMap` and passes through
```
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +84,51 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated, copied, or reinserted record with metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must override this
+ * method when mixing in {@link SupportsColumnUpdates}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but does not override this method.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T metadata, T record) throws IOException {
+ throw new SparkUnsupportedOperationException(
+ "DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED",
+ Map.of("class", getClass().getName()));
+ }
+
+ /**
+ * Writes one updated, copied, or reinserted record without metadata.
+ * <p>
+ * Equivalent to {@link #writeUpdate(Object, Object)} for writers that do not require metadata.
+ * Connectors that mix in {@link SupportsColumnUpdates} receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must override this
+ * method when mixing in {@link SupportsColumnUpdates}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but does not override this method.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T record) throws IOException {
Review Comment:
`DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED` is defined and thrown here (and at `:103`), but no test exercises this path. Could you add a test where a connector mixes in `SupportsColumnUpdates` without overriding `writeUpdate`, and assert this error condition is raised? That keeps the new error condition covered and guards the dispatch contract.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -267,13 +268,80 @@ abstract class RowLevelOperationSuiteBase
protected def checkLastWriteInfo(
expectedRowSchema: StructType = new StructType(),
expectedRowIdSchema: Option[StructType] = None,
- expectedMetadataSchema: Option[StructType] = None): Unit = {
+ expectedMetadataSchema: Option[StructType] = None,
+ expectedUpdateSchema: Option[StructType] = None): Unit = {
val info = table.lastWriteInfo
assert(info.schema == expectedRowSchema, "row schema must match")
val actualRowIdSchema = Option(info.rowIdSchema.orElse(null))
assert(actualRowIdSchema == expectedRowIdSchema, "row ID schema must match")
val actualMetadataSchema = Option(info.metadataSchema.orElse(null))
assert(actualMetadataSchema == expectedMetadataSchema, "metadata schema must match")
+ val actualUpdateSchema = Option(info.updateSchema.orElse(null))
+ assert(actualUpdateSchema == expectedUpdateSchema, "update schema must match")
+ }
+
+ protected def getUpdateSummary(): org.apache.spark.sql.connector.write.UpdateSummary = {
Review Comment:
nit. Please import `org.apache.spark.sql.connector.write.UpdateSummary` instead of using the fully-qualified name inline (here and at line 286).
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3787082425
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, any required metadata attrs and optionally connector
+ // declared attrs)
val rowAttrs = relation.output
+ val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates]
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveConnectorDataAttrs(relation, operation)
+ } else Nil
+ val scanOnlyDataAttrs = if (supportsColumnUpdate) {
+ resolveScanOnlyDataAttrs(relation, operation)
+ } else Nil
+
+ if (supportsColumnUpdate) {
+ validateUpdatedColumnsSubset(operation, assignments, connectorDataAttrs)
+ validateNoOverlap(operation, connectorDataAttrs, scanOnlyDataAttrs)
+ validatePartitionAttrsDeclared(operation, relation, connectorDataAttrs, scanOnlyDataAttrs)
+ }
+
+
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
- val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ if (supportsColumnUpdate && operation.representUpdateAsDeleteAndInsert) {
+ validateNoRowIdReassignment(operation, assignments, rowIdAttrs)
Review Comment:
I am not sure I understand this validation. Specifically, validateRowIdDeclared seems wrong. What if my row id is a metadata 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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4074793021
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -157,29 +228,112 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, required metadata attrs, and any connector-declared
+ // data attrs for column-update writes)
val rowAttrs = relation.output
+ val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates]
+ val connectorDataAttrs = resolveConnectorDataAttrs(relation, operation, assignments)
+
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
- val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ if (supportsColumnUpdate && operation.representUpdateAsDeleteAndInsert) {
+ validateNoRowIdReassignment(operation, relation, connectorDataAttrs, assignments, rowIdAttrs)
+ validateRowIdDeclared(operation, connectorDataAttrs, rowIdAttrs)
+ }
+
+ val narrowDataAttrs = if (supportsColumnUpdate) {
+ computeNarrowReadAttrs(relation, connectorDataAttrs, assignments, cond)
+ } else {
+ relation.output
+ }
+
+ val readRelation = if (supportsColumnUpdate) {
+ buildNarrowRelationWithAttrs(relation, operationTable, narrowDataAttrs, metadataAttrs,
+ rowIdAttrs)
+ } else {
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ }
// 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 rowDeltaPlan = if (supportsColumnUpdate) {
+ if (operation.representUpdateAsDeleteAndInsert) {
+ buildNarrowDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else {
+ buildColumnUpdateProjection(
+ matchedRowsPlan, assignments, rowIdAttrs, metadataAttrs, connectorDataAttrs)
+ }
} else {
- buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs)
+ if (operation.representUpdateAsDeleteAndInsert) {
+ buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else {
+ buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, 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(
+ rowDeltaPlan, rowAttrs, rowIdAttrs, metadataAttrs, connectorDataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond)
}
+ /**
+ * Builds the WriteDelta projection for the column-update path.
+ */
+ private def buildColumnUpdateProjection(
+ plan: LogicalPlan,
+ assignments: Seq[Assignment],
+ rowIdAttrs: Seq[Attribute],
+ metadataAttrs: Seq[Attribute],
+ connectorDataAttrs: Seq[AttributeReference]): LogicalPlan = {
+
+ val assignedValues = assignments.collect {
+ case Assignment(key: Attribute, value) if !isIdentityAssignment(key, value) =>
+ Alias(value, key.name)()
+ }
+
+ // Connector-required columns whose value isn't being changed by the UPDATE: pass through the
+ // current value so the connector receives a complete write row. Row-ID columns are excluded
+ // here even when also connector-required (e.g. a primary key used for both row lookup and
+ // write payload); they are emitted once, below, via rowIdValues.
+ val assignedKeyIds = collectUpdatedAttrs(assignments).map(_.exprId).toSet
+ val rowIdAttrSet = AttributeSet(rowIdAttrs)
+ val connectorPassThroughValues = connectorDataAttrs.filterNot(a =>
+ assignedKeyIds.contains(a.exprId) || rowIdAttrSet.contains(a))
+
+ val metadataAttrSet = AttributeSet(metadataAttrs)
+ val metadataValues = plan.output.filter(metadataAttrSet.contains).map { attr =>
+ if (MetadataAttribute.isPreservedOnUpdate(attr)) {
+ attr
+ } else {
+ Alias(Literal(null, attr.dataType), attr.name)(explicitMetadata = Some(attr.metadata))
+ }
+ }
+
+ val rowIdValues = plan.output.filter(rowIdAttrSet.contains)
Review Comment:
When a row-ID column is reassigned, this projection emits `pk` twice with different exprIds: `assignedValues` contributes `Alias(pk + 10, "pk")` (fresh exprId) and `rowIdValues` contributes the original `pk#1`. The write projections survive only because `findColOrdinal` takes the first name match, but name-based resolution over `query.output` does not:
- `ResolveTableConstraints` wraps the query in `Filter(CheckInvariant(UnresolvedAttribute("pk") > 0), query)` for a table with `CHECK (pk > 0)` -> `AttributeSeq.resolve` sees two candidates -> `AMBIGUOUS_REFERENCE`.
- A connector whose `Write` clusters or orders by `pk` fails the same way in `V2Writes` via `DistributionAndOrderingUtils.prepareQuery`.
Repro on the non-split `column-update` fixture: `ALTER TABLE t ADD CONSTRAINT positive_pk CHECK (pk > 0)` then `UPDATE t SET pk = pk + 10, salary = -1 WHERE dep = 'hr'`. The existing test `pk already in updatedColumns is not duplicated` only inspects `updateSchema()` (already de-duplicated by name lookup), so it cannot catch this. The wide `buildWriteDeltaUpdateProjection` replaces the value in place and never has the duplicate.
Building the output in a single pass over `plan.output` (the way `buildNarrowReplaceDataUpdateProjection` does) -- assigned attrs replaced in place, metadata preserved/nullified, everything else passed through -- followed by `originalRowIdValues` removes the duplicate and the first-match ordering dependency at the same time.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -109,6 +123,43 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
relation)
}
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation. Non-existent columns are rejected with an `AnalysisException`.
+ */
+ protected def resolveRequiredDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes
+ if (refs.isEmpty) {
+ throw QueryCompilationErrors.emptyRequiredDataAttributesError(operation.getClass.getName)
+ }
+ val nested = refs.filter(_.fieldNames.length != 1).map(_.describe()).toImmutableArraySeq
+ if (nested.nonEmpty) {
+ throw QueryCompilationErrors.nestedRequiredDataAttributeError(
+ operation.getClass.getName, nested)
+ }
+ val normalizedNames = refs.map { ref =>
+ val name = ref.fieldNames.head
+ if (conf.caseSensitiveAnalysis) name else name.toLowerCase(Locale.ROOT)
+ }
+ val duplicates = normalizedNames.groupBy(identity).collect {
+ case (_, occurrences) if occurrences.length > 1 => occurrences.head
+ }.toSeq
+ if (duplicates.nonEmpty) {
+ throw QueryCompilationErrors.duplicateRequiredDataAttributeError(
+ operation.getClass.getName, duplicates)
+ }
+ val resolved = V2ExpressionUtils.resolveRefs[AttributeReference](
Review Comment:
`LogicalPlan.resolve` falls back to `metadataOutput`, so a metadata column such as `_partition` is accepted here as a *data* attribute (the map-back `byExprId.getOrElse(a.exprId, a)` passes it through untouched). The javadoc actively steers connectors this way ("Partition columns ... must be declared here too"), and in our fixture the partition column is exactly a metadata column.
On the split path this corrupts the row silently: `buildNarrowRelationWithAttrs` leaves the metadata attr in the middle of the data attrs (declared order, `dedupAttrs` first-wins), `buildNarrowDeletesAndInserts` partitions `output` into `rowAttrs ++ metadataAttrs` for the DELETE/REINSERT projections, but the `Expand` output attrs keep the interleaved order, so projection slot `i` and output attribute `i` disagree. With `column-update-split-req-attrs = "pk,_partition,dep"` and `UPDATE t SET dep = 'x' WHERE pk = 1`, `dep`'s value lands in the `_partition` slot and vice versa (both STRING, so nothing fails); with a type mismatch it reads garbage. The plan stays resolved because `areCompatible` compares names/types only. The non-split and CoW paths are self-consistent (name/exprId based) but still put the metadata column into `updateSchema()`.
Suggest rejecting `MetadataAttribute.isValid(a.metadata)` attrs here with a dedicated error (or resolving against `LocalRelation(relation.output)` so the fallback never applies), and clarifying in the javadoc that metadata partition columns travel through `requiredMetadataAttributes()`.
##########
common/utils/src/main/resources/error/error-conditions.json:
##########
@@ -2385,6 +2421,12 @@
],
"sqlState" : "42K03"
},
+ "DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED" : {
+ "message" : [
+ "<class> mixes in `SupportsColumnUpdates` but does not implement `writeUpdate`. Connectors that opt into narrow column-level updates must override `writeUpdate` to handle rows in the `LogicalWriteInfo.updateSchema()` layout; the default implementation would forward narrow rows to `write`, which expects the full `LogicalWriteInfo.schema()` layout."
Review Comment:
For every narrow UPDATE the row projection is the zero-column placeholder, so `LogicalWriteInfo.schema()` is an empty `StructType` here -- the suites assert exactly that (`checkLastWriteInfo` with the default `new StructType()`). This message says the opposite ("`write`, which expects the full `LogicalWriteInfo.schema()` layout"). The PR description was fixed for this, but the two API javadocs were not: `LogicalWriteInfo.schema()` still reads "the schema of the input data from Spark to data source" and `updateSchema()` does not mention that `schema()` is empty (or INSERT-only) when it is present. This is new for group-based writes -- before this PR `ReplaceData` never handed a connector an empty `schema()` -- and a `newWriteBuilder` that validates `info.schema()` against the table will reject every narrow UPDATE.
Suggest: drop "full" here, and add one sentence to `LogicalWriteInfo.schema()`/`updateSchema()` stating that when `updateSchema()` is present, `schema()` covers only INSERT-shaped rows and is empty for UPDATE.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -306,13 +308,80 @@ abstract class RowLevelOperationSuiteBase
protected def checkLastWriteInfo(
expectedRowSchema: StructType = new StructType(),
expectedRowIdSchema: Option[StructType] = None,
- expectedMetadataSchema: Option[StructType] = None): Unit = {
+ expectedMetadataSchema: Option[StructType] = None,
+ expectedUpdateSchema: Option[StructType] = None): Unit = {
val info = table.lastWriteInfo
assert(info.schema == expectedRowSchema, "row schema must match")
val actualRowIdSchema = Option(info.rowIdSchema.orElse(null))
assert(actualRowIdSchema == expectedRowIdSchema, "row ID schema must match")
val actualMetadataSchema = Option(info.metadataSchema.orElse(null))
assert(actualMetadataSchema == expectedMetadataSchema, "metadata schema must match")
+ val actualUpdateSchema = Option(info.updateSchema.orElse(null))
+ assert(actualUpdateSchema == expectedUpdateSchema, "update schema must match")
+ }
+
+ protected def getUpdateSummary(): UpdateSummary = {
+ catalog.loadTable(ident).asInstanceOf[InMemoryTable]
+ .commits.last.writeSummary.get
+ .asInstanceOf[UpdateSummary]
+ }
+
+ /**
+ * Asserts the last UPDATE's write summary metrics. `deltaUpdate` controls the expected
+ * COPY count: MoR connectors emit only deltas (no COPY rows) so `numCopiedRows` is forced
+ * to 0; CoW connectors emit COPY rows for unchanged rows in matched groups.
+ */
+ protected def checkUpdateMetrics(
+ numUpdatedRows: Long,
+ numCopiedRows: Long,
+ deltaUpdate: Boolean = false): Unit = {
+ val summary = getUpdateSummary()
+ assert(summary.numUpdatedRows() === numUpdatedRows,
+ s"Expected numUpdatedRows=$numUpdatedRows, got ${summary.numUpdatedRows()}")
+ val expectedCopied = if (deltaUpdate) 0L else numCopiedRows
+ assert(summary.numCopiedRows() === expectedCopied,
+ s"Expected numCopiedRows=$expectedCopied, got ${summary.numCopiedRows()}")
+ }
+
+ /**
+ * Asserts that the column names in RowLevelOperationInfo.updatedColumns() received by the
+ * last operation match exactly the expected set. Order is ignored.
+ */
+ protected def checkLastUpdatedColumns(expectedNames: String*): Unit = {
+ val actual = table.lastUpdatedColumns.map(_.describe()).toSet
+ val expected = expectedNames.toSet
+ assert(actual == expected,
+ s"updatedColumns mismatch: expected ${expected.mkString("[", ", ", "]")} " +
+ s"but got ${actual.mkString("[", ", ", "]")}")
+ }
+
+ /**
+ * Asserts that the last connector scan schema does NOT contain any of the given columns.
+ * Useful for negative narrowing assertions -- proves that a column outside the analysis-time
+ * narrow set was pruned by the scan, without over-specifying what else is present (which
+ * ColumnPruning may further tighten in ways unrelated to the narrowing contract).
+ */
+ protected def checkLastScanExcludes(excludedNames: String*): Unit = {
+ val schema = Option(table.lastScanSchema).getOrElse(StructType(Nil))
Review Comment:
With `lastScanSchema == null` (the state `before {}` resets to) this substitutes an empty schema, so the negative assertion passes vacuously. `GroupBasedColumnUpdateTableSuite`'s "scan excludes columns outside required + cond" relies solely on this helper, so removing `recordLastScanSchema` from `PartitionBasedColumnUpdateOperation`'s scan builder would leave it green. `assert(table.lastScanSchema != null, ...)` instead of `getOrElse(StructType(Nil))` closes it.
Related: the group-path "runtime group filtering data correctness" test only calls `checkAnswer`, so it cannot tell whether a group-filter subquery was injected or which partitions were rewritten. `DeltaBasedUpdateTableSuite.checkUpdateRuntimeGroupFiltering` shows the shape (`executeAndCheckScans(..., groupFilterScanSchema = Some(...))` + `checkReplacedPartitions(Seq("hr"))`); the narrow variant should assert the same.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,320 @@
+/*
+ * 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.{SparkRuntimeException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.connector.catalog.{CatalogV2Util, TableInfo, WriteUpdate}
+import org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity, reference}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.StructType
+
+class GroupBasedColumnUpdateTableSuite extends UpdateTableSuiteBase {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-cow", "true")
+ props
+ }
+
+ private def createAndInitTableReplaceData(schemaString: String, jsonData: String): Unit = {
Review Comment:
nit. This helper is functionally identical to the inherited `createAndInitTable` (same `identity(dep)` partitioning, and `column-update-cow` is exactly this suite's `extraTableProps`), and `DeltaBasedColumnUpdateTableSuite.createAndInitTableFromInfo` is the same story because `DeltaBasedColumnUpdateOperationFromInfo` is an empty subclass of `DeltaBasedColumnUpdateOperation` (the `column-update-from-info` flag and class can go). The other six `createAndInitTableXxx` copies differ only in one `props.put`; a `props` parameter on the base helper (replacing, not merging, `extraTableProps`, since `column-update` is matched first in `newRowLevelOperationBuilder`) would absorb all of them.
While here: the ten hand-written `table.lastUpdatedColumns.map(_.describe()).toSet == ...` assertions in the two suites bypass the `checkLastUpdatedColumns` helper this PR adds, three of the Delta suite's `updatedColumns` tests duplicate the parent-suite tests added by the same PR (same DDL, SQL and expectation), `writeUpdateLogEntry`/`writeUpdateWithMetadataLogEntry` in `RowLevelOperationSuiteBase` have no callers, and the three ~17-line "overlay narrow row on base row" blocks in `InMemoryRowLevelOperationTable` could share one helper.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -88,6 +92,16 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
relation.copy(table = table, output = attrs)
}
+ protected def buildNarrowRelationWithAttrs(
+ relation: DataSourceV2Relation,
+ table: RowLevelOperationTable,
+ dataAttrs: Seq[AttributeReference],
+ metadataAttrs: Seq[AttributeReference],
+ rowIdAttrs: Seq[AttributeReference] = Nil): DataSourceV2Relation = {
+ val attrs = dedupAttrs(dataAttrs ++ rowIdAttrs ++ metadataAttrs)
Review Comment:
This is the first time a leaf `DataSourceV2Relation.output` is a strict subset of the table's columns, and `PushDownUtils.pruneColumns` still assumes the old invariant in two places:
1. `toOutputAttrs` maps every field in `scan.readSchema()` back to `relation.output` by exact name (`nameToAttr(a.name)`). A connector that reports more columns than requested -- which `SupportsPushDownRequiredColumns` explicitly allows ("it's also OK to do the pruning partially") and which our own `SIMULATE_PARTIAL_COLUMN_PRUNING` fixture models -- now throws `NoSuchElementException: key not found: salary` from the optimizer. Same call in `GroupBasedRowLevelOperationScanPlanning`.
2. The `case _ => scanBuilder.build() -> relation.output` fallback for a `ScanBuilder` without `SupportsPushDownRequiredColumns` pairs full-width reader rows with the narrow `output`, so `BatchScanExec` binds the wrong ordinals and misaligned values flow into `writeUpdate`/`update` with no error.
Nothing validates or documents that a `SupportsColumnUpdates` scan builder must implement `SupportsPushDownRequiredColumns` and must not over-report. A sub-suite with `SIMULATE_PARTIAL_COLUMN_PRUNING -> "true"` in `extraTableProps` would reproduce (1) immediately. I would either make `toOutputAttrs` tolerate extra fields (they are unreferenced above the scan) or reject/document both shapes.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +84,54 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated, copied, or reinserted record with metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} without also mixing in
+ * {@link SupportsDelta} (i.e. group-based operations) receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. An operation that mixes in both
+ * instead delivers narrow rows through {@link DeltaWriter#update} / {@link DeltaWriter#reinsert}
+ * and never calls this method. By default, delegates to {@link #writeUpdate(Object)} for
+ * connectors that do not need metadata; implementations that do need it should override this
+ * method directly.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but overrides neither {@code writeUpdate} overload.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T metadata, T record) throws IOException {
+ writeUpdate(record);
+ }
+
+ /**
+ * Writes one updated, copied, or reinserted record without metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} without also mixing in
+ * {@link SupportsDelta} (i.e. group-based operations) receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must override this method,
+ * or {@link #writeUpdate(Object, Object)}, when mixing in {@link SupportsColumnUpdates} without
+ * {@link SupportsDelta}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but overrides neither {@code writeUpdate} overload.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T record) throws IOException {
Review Comment:
After the delegation change, the javadoc on both overloads says overriding *either* one is enough ("Implementations must override this method, or `writeUpdate(Object, Object)`", and both `@throws` clauses say "overrides neither"). That is not what the dispatch does: `ReplaceDataExec.writingTask` picks `DataWithProjectionWritingSparkTask` whenever `requiredMetadataAttributes()` is empty (the default), and that task calls the 1-arg overload only. So a group-based connector that overrides just `writeUpdate(metadata, record)` -- the overload `SupportsColumnUpdates`' javadoc links to -- gets `DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED` ("does not implement `writeUpdate`", which is false) on every UPDATE. The reverse shape works, so the contract is asymmetric in exactly the direction the docs invite.
The PR description already states the stricter rule ("must override `writeUpdate(record)`"). I would make the javadoc and the error text say that, and drop the "or" clause. Also, no test goes through the 1-arg dispatch at all: `PartitionBasedColumnUpdateOperation` always declares `[_partition, index]` and ignores `no-metadata`, so every CoW column-update test uses `DataAndMetadataWritingSparkTask`. A no-metadata CoW fixture variant would cover it.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -82,11 +83,29 @@ class InMemoryRowLevelOperationTable private (
private final val noMetadata = properties.getOrDefault(NO_METADATA, "false") == "true"
private final val useCatalystRuntimeFiltering =
properties.getOrDefault(USE_CATALYST_RUNTIME_FILTERING, "false") == "true"
+ private final val COLUMN_UPDATE = "column-update"
+ private final val COLUMN_UPDATE_REQ_ATTRS = "column-update-req-attrs"
+ private final val COLUMN_UPDATE_COW = "column-update-cow"
+ private final val COLUMN_UPDATE_COW_NO_WRITE_UPDATE = "column-update-cow-no-write-update"
+ private final val COLUMN_UPDATE_FROM_INFO = "column-update-from-info"
+ private final val COLUMN_UPDATE_SPLIT = "column-update-split"
+ private final val COLUMN_UPDATE_SPLIT_REQ_ATTRS = "column-update-split-req-attrs"
+ private final val COLUMN_UPDATE_EMPTY_REQ_ATTRS = "column-update-empty-req-attrs"
+ private final val COLUMN_UPDATE_SPLIT_MISSING_ROW_ID = "column-update-split-missing-row-id"
// used in row-level operation tests to verify replaced partitions
var replacedPartitions: Seq[Seq[Any]] = Seq.empty
// used in row-level operation tests to verify reported write schema
var lastWriteInfo: LogicalWriteInfo = _
+ // used in column-update tests to verify that Spark passed the correct updated column list
+ // to the connector via RowLevelOperationInfo.updatedColumns()
+ var lastUpdatedColumns: Array[NamedReference] = Array.empty
Review Comment:
Two bookkeeping gaps for this new field:
- `copy()` below copies `replacedPartitions`, `lastWriteInfo` and `lastWriteLog` but not `lastUpdatedColumns`, so `InMemoryRowLevelOperationTableCatalog.loadTable` snapshots always report an empty array. The suites happen to read `liveTable`, so nothing fails today, but the next test that follows the `lastWriteInfo`/`loadTable` pattern will.
- `TxnTable.newRowLevelOperationBuilder` (`txns.scala:146`) writes it to the live delegate immediately, whereas the sibling fields are staged and copied in `commit()` -- which already copies `lastUpdatedColumns` too, so the immediate write is redundant. After an UPDATE that fails analysis, `table.lastUpdatedColumns` reflects the failed statement while `lastWriteInfo`/`lastWriteLog` reflect the previous commit.
Add `copied.lastUpdatedColumns = lastUpdatedColumns` and delete the immediate delegate write in `txns.scala`.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -252,6 +304,413 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends only the
+ // declared + assigned columns in the row projection instead of the full row schema. The base
+ // class composes its required-attrs set as `pk` (the row-lookup key) and `dep` (an
+ // unconditionally declared base column) plus whatever columns Spark reports as being
+ // assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = {
+ val base = Seq(FieldReference("pk"), FieldReference("dep"))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ protected def clusterColumnRef: NamedReference = PARTITION_COLUMN_REF
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.updateSchema().isPresent) {
+ capturedInfo.updateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(Array(clusterColumnRef))
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ Array[SortOrder](
+ LogicalExpressions.sort(
+ clusterColumnRef,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering())
+ )
+ }
+
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ // Use the narrow update schema for UPDATE/COPY/REINSERT rows when present.
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For column-update writes, rows contain only the assigned columns
+ // (narrow schema from LogicalWriteInfo). We expand each row to the full table
+ // schema by overlaying write-schema columns on the base row found by pk.
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+
+ val mergedData = newData.map { buf =>
+ val merged = new BufferedRows(buf.key, schema)
+ val updateOpName = UTF8String.fromString(Update.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == updateOpName) {
+ val pk = logRow.getInt(1)
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val baseRow = dataMap.values.iterator.flatten
+ .flatMap(_.rows)
+ .find(r => r.getInt(schema.fieldIndex("pk")) == pk)
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ val field = schema.fields(i)
+ val value = if (i < base.numFields) {
+ base.get(i, field.dataType)
+ } else {
+ ResolveDefaultColumns.getExistenceDefaultValue(field)
+ }
+ fullRow.update(i, value)
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ writeFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ merged.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ merged.rows.append(insertRow.copy())
+ }
+ }
+ merged
+ }
+
+ withDeletes(newData)
+ withData(mergedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ }
+ }
+ }
+ }
+
+ class DeltaBasedColumnUpdateOperationWithReqAttrs(
+ command: Command,
+ reqCols: Array[String],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, options = options) {
+ override def requiredDataAttributes(): Array[NamedReference] = reqCols.map(FieldReference(_))
+ }
+
+ class DeltaBasedColumnUpdateOperationFromInfo(
+ command: Command,
+ updatedCols: Seq[NamedReference],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols, options) {
+ }
+
+ class DeltaBasedColumnUpdateSplitOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols, options) {
+ override def representUpdateAsDeleteAndInsert(): Boolean = true
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.updateSchema().isPresent) {
+ capturedInfo.updateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+ override def requiredDistribution(): Distribution =
+ Distributions.clustered(Array(PARTITION_COLUMN_REF))
+ override def requiredOrdering(): Array[SortOrder] = Array[SortOrder](
+ LogicalExpressions.sort(
+ PARTITION_COLUMN_REF,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()))
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For delete+reinsert with narrow writes, the REINSERT row has only the
+ // connector-declared columns (requiredDataAttributes order).
+ // pk is the first field in the write schema (declared before updatedCols).
+ // Reconstruct the full row by overlaying the narrow row onto the original.
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+ val reinsertOpName = UTF8String.fromString(Reinsert.toString)
+ val pkIdx = writeFieldIdx("pk")
+
+ val expandedData = newData.map { buf =>
+ val expanded = new BufferedRows(buf.key, schema)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == reinsertOpName) {
Review Comment:
This `doCommit` only materializes `Reinsert` entries; `Insert` (and `Update`) entries are dropped, while the parent `DeltaBasedColumnUpdateOperation.doCommit` handles both. The operation is returned for every command, so a `MERGE ... WHEN NOT MATCHED THEN INSERT` on a `column-update-split` table would commit nothing for the inserted rows and no test would notice: neither column-update suite runs `MERGE INTO`, and the only `DELETE` test (delta) asserts just `updatedColumns.isEmpty`; there is no CoW `DELETE` test at all. Since the mix-in's documented contract is precisely that DELETE and MERGE fall back to full-width rows, that fallback deserves a test on both connector types. Mirroring the parent's `insertOpName` branch here is a one-liner.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -228,4 +382,169 @@ 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, validating that every assigned column is covered. Returns `Nil` otherwise.
+ */
+ private def resolveConnectorDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation,
+ assignments: Seq[Assignment]): Seq[AttributeReference] = operation match {
+ case scu: SupportsColumnUpdates =>
+ val attrs = resolveRequiredDataAttrs(relation, scu)
+ validateUpdatedColumnsSubset(scu, assignments, attrs)
+ attrs
+ case _ => Nil
+ }
+
+ /**
+ * Computes the narrow set of data columns that must be present in the scan for a column-update
+ * write: connector-declared attrs (`requiredDataAttributes()`), unioned with any table columns
+ * referenced by non-identity assignment RHS expressions, the operation condition, and the
+ * table's CHECK constraints.
+ */
+ 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 checkConstraintAttrs = resolveCheckConstraintAttrs(relation)
+ dedupAttrs(connectorDataAttrs ++ extraRefs ++ checkConstraintAttrs)
Review Comment:
The exprId map-back added for the connector declaration (`resolveRequiredDataAttrs`) is missing for the other sources that feed the narrow relation output: `cond.references`, the assignment RHS references, `resolveCheckConstraintAttrs`, and `rowIdAttrs`. Those instances come out of `AttributeSeq.resolve` as `a.withName(<requested spelling>)`, and since `dedupAttrs` keeps the first instance per exprId, a user-spelled attribute becomes the scan's column name whenever the column is not also declared.
With the default `spark.sql.caseSensitive=false`, `column-update` table `pk INT NOT NULL, salary INT, bonus INT, dep STRING` (declared `[pk, dep, salary]`):
```sql
UPDATE t SET salary = salary + BONUS WHERE pk = 1
```
narrow output = `[pk#1, dep#4, salary#2, BONUS#3, ...]` -> `pruneColumns` asks the connector for `BONUS` -> `InMemoryScanBuilder` drops it (case-sensitive name match) -> the scan relation has no `BONUS#3` -> `INTERNAL_ERROR_ATTRIBUTE_NOT_FOUND` at execution. A connector that echoes the table spelling instead hits `NoSuchElementException` in `PushDownUtils.toOutputAttrs`. `WHERE EXTRA > 3` and `CHECK (EXTRA > 0)` take the same path. The wide path is immune because it always starts from `relation.output`.
Simplest fix is to build the narrow output from the relation's own attributes so both spelling and column order are preserved:
```scala
val narrowSet = AttributeSet(connectorDataAttrs ++ extraRefs ++ checkConstraintAttrs)
relation.output.filter(narrowSet.contains)
```
(and the same `relation.output.filter(...)` for `rowIdAttrs` in `buildNarrowRelationWithAttrs`). Worth a case test with an undeclared, differently-cased column in the RHS and in the WHERE clause.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala:
##########
@@ -513,7 +557,31 @@ case class WriteDelta(
case Some(projection) => DataTypeUtils.toAttributes(projection.schema)
case None => Nil
}
- table.skipSchemaResolution || areCompatible(inRowAttrs, outRowAttrs)
+ val inUpdateAttrs = projections.updateRowProjection match {
+ case Some(projection) => DataTypeUtils.toAttributes(projection.schema)
+ case None => Nil
+ }
+ // `rowProjection` (INSERT-tagged rows) validates against `outRowAttrs`. For column-update
+ // connectors, `updateRowProjection` is narrow and validates against `projectedDataAttrs`
+ // (the connector-declared narrow set) instead. When the connector does not mix in
+ // `SupportsColumnUpdates`, `updateRowProjection` is absent and `updateResolved` is trivially
+ // true.
+ val insertResolved = table.skipSchemaResolution ||
+ (inUpdateAttrs.nonEmpty && inRowAttrs.isEmpty) ||
+ areCompatible(inRowAttrs, outRowAttrs)
+ val updateResolved = table.skipSchemaResolution ||
+ inUpdateAttrs.isEmpty || dataAttrsResolved(inUpdateAttrs)
+ insertResolved && updateResolved
+ }
+
+ /**
+ * Validates the narrow-write-schema row projection output for a column-update connector.
+ * The write schema must exactly match the columns declared via
+ * `SupportsColumnUpdates.requiredDataAttributes()` (same columns, same order).
+ */
+ private def dataAttrsResolved(inRowAttrs: Seq[Attribute]): Boolean = {
Review Comment:
nit. This is byte-identical to `ReplaceData.dataAttrsResolved`, and the `inUpdateAttrs` / `insertResolved` / `updateResolved` blocks above differ from `ReplaceData`'s only in `outRowAttrs`. Everything they touch (`operation`, `projectedDataAttrs`, `areCompatible`) already lives on `RowLevelWrite`, next to `projectedMetadataAttrs`, so both could move there (with an abstract `updateRowProjection` accessor, since the two projections types differ). The `operation.isInstanceOf[SupportsColumnUpdates]` guard is also redundant: `projectedDataAttrs` returns `Nil` otherwise and `areCompatible` fails on the size mismatch.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4240160391
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,47 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ *
+ * @since 4.3.0
Review Comment:
4.3.0 has been released, so the APIs are now `@since 4.4.0` and the MiMa exclude is in the 4.4 section.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -109,6 +120,22 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
relation)
}
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation. Non-existent columns are rejected with an `AnalysisException`.
+ */
+ protected def resolveRequiredDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes
+ if (refs.isEmpty) {
+ throw QueryCompilationErrors.emptyRequiredDataAttributesError(operation.getClass.getName)
+ }
+ V2ExpressionUtils.resolveRefs[AttributeReference](
Review Comment:
Fixed: each resolved attribute is mapped back to the relation's own attribute by expr ID, and the narrow read relation this fed into no longer exists. Covered by the case-insensitive declaration tests (`requiredDataAttributes resolves case-insensitively ...`) and the differently-cased tests, which check the table's spelling in the scan.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -146,6 +203,48 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
Project(operationCol +: updatedValues, plan)
}
+ /**
+ * Variant of `buildReplaceDataUpdateProjection` for the `SupportsColumnUpdates` narrow-scan
+ * path.
+ * For narrow attributes, looks up assignments by `ExprId` via `AttributeMap`and passes through
Review Comment:
This code was removed in the redesign.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, any required metadata attrs and optionally connector
+ // declared attrs)
val rowAttrs = relation.output
+ val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates]
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveConnectorDataAttrs(relation, operation)
+ } else Nil
+ val scanOnlyDataAttrs = if (supportsColumnUpdate) {
+ resolveScanOnlyDataAttrs(relation, operation)
+ } else Nil
+
+ if (supportsColumnUpdate) {
+ validateUpdatedColumnsSubset(operation, assignments, connectorDataAttrs)
+ validateNoOverlap(operation, connectorDataAttrs, scanOnlyDataAttrs)
+ validatePartitionAttrsDeclared(operation, relation, connectorDataAttrs, scanOnlyDataAttrs)
+ }
+
+
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
- val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ if (supportsColumnUpdate && operation.representUpdateAsDeleteAndInsert) {
+ validateNoRowIdReassignment(operation, assignments, rowIdAttrs)
+ validateRowIdDeclared(operation, connectorDataAttrs, rowIdAttrs)
+ }
+
+ val narrowDataAttrs = if (supportsColumnUpdate) {
+ computeNarrowReadAttrs(relation, connectorDataAttrs, scanOnlyDataAttrs, assignments, cond)
+ } else {
+ relation.output
+ }
+
+ val readRelation = if (supportsColumnUpdate) {
+ buildNarrowRelationWithAttrs(relation, operationTable, narrowDataAttrs, metadataAttrs,
+ rowIdAttrs)
+ } else {
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ }
// 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 rowDeltaPlan = if (supportsColumnUpdate) {
+ if (operation.representUpdateAsDeleteAndInsert) {
+ buildNarrowDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else {
+ buildColumnUpdateProjection(
+ matchedRowsPlan, assignments, rowIdAttrs, metadataAttrs, connectorDataAttrs,
+ scanOnlyDataAttrs)
+ }
} else {
- buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs)
+ if (operation.representUpdateAsDeleteAndInsert) {
+ buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else {
+ buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, 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(
+ rowDeltaPlan, rowAttrs, rowIdAttrs, metadataAttrs, connectorDataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond)
}
+ /**
+ * Builds the WriteDelta projection for the column-update path.
+ */
+ private def buildColumnUpdateProjection(
+ plan: LogicalPlan,
+ assignments: Seq[Assignment],
+ rowIdAttrs: Seq[Attribute],
+ metadataAttrs: Seq[Attribute],
+ connectorDataAttrs: Seq[AttributeReference],
+ scanOnlyDataAttrs: Seq[AttributeReference]): LogicalPlan = {
+
+ val assignedValues = assignments.collect {
+ case Assignment(key: Attribute, value) if !isIdentityAssignment(key, value) =>
+ Alias(value, key.name)()
+ }
+
+ // Connector-required columns whose value isn't being changed by the UPDATE: pass through the
+ // current value so the connector receives a complete write row. Row-ID columns are excluded
+ // here even when also connector-required (e.g. a primary key used for both row lookup and
+ // write payload) they are emitted once, below, via rowIdValues.
+ val assignedKeyIds = assignments.collect {
+ case Assignment(key: AttributeReference, value) if !isIdentityAssignment(key, value) =>
+ key.exprId
+ }.toSet
+ val rowIdAttrSet = AttributeSet(rowIdAttrs)
+ val connectorPassThroughValues = connectorDataAttrs.filterNot(a =>
+ assignedKeyIds.contains(a.exprId) || rowIdAttrSet.contains(a))
+
+ // scanOnlyDataAttrs are never assigned carry them through the write query so
Review Comment:
This code was removed in the redesign.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/LogicalWriteInfo.java:
##########
@@ -65,4 +65,14 @@ default Optional<StructType> metadataSchema() {
throw new SparkUnsupportedOperationException(
"DATA_SOURCE_METADATA_SCHEMA_NOT_IMPLEMENTED", Map.of("class", getClass().getName()));
}
+
+ /**
+ * the narrow schema for updates. Present only when the connector mixes in
+ * {@link SupportsColumnUpdates}.
+ *
+ * @since 4.3.0
+ */
+ default Optional<StructType> updateSchema() {
Review Comment:
Renamed to `columnUpdateSchema()`, together with `writeColumnUpdate`.
##########
common/utils/src/main/resources/error/error-conditions.json:
##########
@@ -2385,6 +2421,12 @@
],
"sqlState" : "42K03"
},
+ "DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED" : {
+ "message" : [
+ "<class> mixes in `SupportsColumnUpdates` but does not implement `writeUpdate`. Connectors that opt into narrow column-level updates must override `writeUpdate` to handle rows in the `LogicalWriteInfo.updateSchema()` layout; the default implementation would forward narrow rows to `write`, which expects the full `LogicalWriteInfo.schema()` layout."
Review Comment:
Fixed. The message no longer describes a default that forwards to `write`, and `LogicalWriteInfo.schema()` now says that when `columnUpdateSchema()` is present it covers only newly inserted rows, so it is empty for UPDATE.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,320 @@
+/*
+ * 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.{SparkRuntimeException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.connector.catalog.{CatalogV2Util, TableInfo, WriteUpdate}
+import org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity, reference}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.StructType
+
+class GroupBasedColumnUpdateTableSuite extends UpdateTableSuiteBase {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-cow", "true")
+ props
+ }
+
+ private def createAndInitTableReplaceData(schemaString: String, jsonData: String): Unit = {
Review Comment:
Removed the duplicate helpers and the `FromInfo` fixture. The suites use one `createAndInitTable(schema, data, tableProps)` overload.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -306,13 +308,80 @@ abstract class RowLevelOperationSuiteBase
protected def checkLastWriteInfo(
expectedRowSchema: StructType = new StructType(),
expectedRowIdSchema: Option[StructType] = None,
- expectedMetadataSchema: Option[StructType] = None): Unit = {
+ expectedMetadataSchema: Option[StructType] = None,
+ expectedUpdateSchema: Option[StructType] = None): Unit = {
val info = table.lastWriteInfo
assert(info.schema == expectedRowSchema, "row schema must match")
val actualRowIdSchema = Option(info.rowIdSchema.orElse(null))
assert(actualRowIdSchema == expectedRowIdSchema, "row ID schema must match")
val actualMetadataSchema = Option(info.metadataSchema.orElse(null))
assert(actualMetadataSchema == expectedMetadataSchema, "metadata schema must match")
+ val actualUpdateSchema = Option(info.updateSchema.orElse(null))
+ assert(actualUpdateSchema == expectedUpdateSchema, "update schema must match")
+ }
+
+ protected def getUpdateSummary(): UpdateSummary = {
+ catalog.loadTable(ident).asInstanceOf[InMemoryTable]
+ .commits.last.writeSummary.get
+ .asInstanceOf[UpdateSummary]
+ }
+
+ /**
+ * Asserts the last UPDATE's write summary metrics. `deltaUpdate` controls the expected
+ * COPY count: MoR connectors emit only deltas (no COPY rows) so `numCopiedRows` is forced
+ * to 0; CoW connectors emit COPY rows for unchanged rows in matched groups.
+ */
+ protected def checkUpdateMetrics(
+ numUpdatedRows: Long,
+ numCopiedRows: Long,
+ deltaUpdate: Boolean = false): Unit = {
+ val summary = getUpdateSummary()
+ assert(summary.numUpdatedRows() === numUpdatedRows,
+ s"Expected numUpdatedRows=$numUpdatedRows, got ${summary.numUpdatedRows()}")
+ val expectedCopied = if (deltaUpdate) 0L else numCopiedRows
+ assert(summary.numCopiedRows() === expectedCopied,
+ s"Expected numCopiedRows=$expectedCopied, got ${summary.numCopiedRows()}")
+ }
+
+ /**
+ * Asserts that the column names in RowLevelOperationInfo.updatedColumns() received by the
+ * last operation match exactly the expected set. Order is ignored.
+ */
+ protected def checkLastUpdatedColumns(expectedNames: String*): Unit = {
+ val actual = table.lastUpdatedColumns.map(_.describe()).toSet
+ val expected = expectedNames.toSet
+ assert(actual == expected,
+ s"updatedColumns mismatch: expected ${expected.mkString("[", ", ", "]")} " +
+ s"but got ${actual.mkString("[", ", ", "]")}")
+ }
+
+ /**
+ * Asserts that the last connector scan schema does NOT contain any of the given columns.
+ * Useful for negative narrowing assertions -- proves that a column outside the analysis-time
+ * narrow set was pruned by the scan, without over-specifying what else is present (which
+ * ColumnPruning may further tighten in ways unrelated to the narrowing contract).
+ */
+ protected def checkLastScanExcludes(excludedNames: String*): Unit = {
+ val schema = Option(table.lastScanSchema).getOrElse(StructType(Nil))
Review Comment:
`checkLastScanExcludes` and `checkLastScanIncludes` now fail when no scan schema was recorded. The group-filtering test checks both scan schemas and the replaced partitions.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java:
##########
@@ -105,4 +105,45 @@ default String description() {
default NamedReference[] requiredMetadataAttributes() {
return new NamedReference[0];
}
+
+
+ /**
+ * Controls whether to send only the required data columns to the connector rather than the
+ * full row.
+ * <p>
+ * When true, Spark narrows the data column schema ({@link LogicalWriteInfo#schema()}) to only
+ * the columns declared via {@link #requiredDataAttributes()}. Metadata columns (from
+ * {@link #requiredMetadataAttributes()}) and row ID columns (from
+ * {@link SupportsDelta#rowId()}) are unaffected and always projected separately.
+ * <p>
+ * If {@link #requiredDataAttributes()} returns a non-empty array, the write schema is exactly
+ * those columns in declared order. The connector must include all columns it wants to receive,
+ * including the columns being updated. If {@link #requiredDataAttributes()} returns an empty
+ * array, Spark sends only the non-identity assigned columns (heuristic path).
+ *
+ * @since 4.2.0
+ */
+ default boolean supportsColumnUpdates() {
+ return false;
+ }
+
+ /**
+ * Returns data column references required to perform this row-level operation.
+ * <p>
+ * This method is only consulted by Spark when {@link #supportsColumnUpdates()} returns
+ * {@code true}. If {@code supportsColumnUpdates()} returns {@code false}, the returned array
+ * is ignored and the full table row is sent (the default behavior).
+ * <p>
+ * When non-empty, the returned columns become the write schema in declared order.
+ * The connector must declare all columns it wants to receive, including the columns being
+ * updated. Use {@link RowLevelOperationInfo#updatedColumns()} to learn which columns are being
+ * assigned, then add any extra columns needed for row lookup or routing (e.g., primary key).
+ * <p>
+ * When empty (the default), Spark falls back to sending only the non-identity assigned columns.
+ *
+ * @since 4.2.0
Review Comment:
The APIs are now `@since 4.4.0`, since 4.3.0 has been released.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -106,38 +105,78 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val remainingRowsPlan = addOperationColumn(COPY_OPERATION,
Filter(remainingRowFilter, readRelation))
- // the new state is a union of updated and copied over records
- val query = Union(updatedRowsPlan, remainingRowsPlan)
+ val updatedAndRemainingRowsPlan = Union(updatedRowsPlan, remainingRowsPlan)
- // build a plan to replace read groups in the table
val writeRelation = relation.copy(table = operationTable)
Review Comment:
You were right. The PR now narrows only the write relation here, with the existing read relation and plan builders, and validates against `relation.output` as before. To keep the API MERGE-proof, the narrow layout is reported as `columnUpdateSchema()` and `schema()` covers only newly inserted rows, so MERGE can narrow its write relation later without API changes.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java:
##########
@@ -105,4 +105,45 @@ default String description() {
default NamedReference[] requiredMetadataAttributes() {
return new NamedReference[0];
}
+
+
+ /**
+ * Controls whether to send only the required data columns to the connector rather than the
+ * full row.
+ * <p>
+ * When true, Spark narrows the data column schema ({@link LogicalWriteInfo#schema()}) to only
+ * the columns declared via {@link #requiredDataAttributes()}. Metadata columns (from
+ * {@link #requiredMetadataAttributes()}) and row ID columns (from
+ * {@link SupportsDelta#rowId()}) are unaffected and always projected separately.
+ * <p>
+ * If {@link #requiredDataAttributes()} returns a non-empty array, the write schema is exactly
+ * those columns in declared order. The connector must include all columns it wants to receive,
+ * including the columns being updated. If {@link #requiredDataAttributes()} returns an empty
+ * array, Spark sends only the non-identity assigned columns (heuristic path).
+ *
+ * @since 4.2.0
+ */
+ default boolean supportsColumnUpdates() {
Review Comment:
Done. The `SupportsColumnUpdates` Javadoc says Spark currently narrows only UPDATE, and `RowLevelOperationInfo.updatedColumns()` says other commands report an empty array.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, any required metadata attrs and optionally connector
+ // declared attrs)
val rowAttrs = relation.output
+ val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates]
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveConnectorDataAttrs(relation, operation)
+ } else Nil
+ val scanOnlyDataAttrs = if (supportsColumnUpdate) {
+ resolveScanOnlyDataAttrs(relation, operation)
+ } else Nil
+
+ if (supportsColumnUpdate) {
+ validateUpdatedColumnsSubset(operation, assignments, connectorDataAttrs)
+ validateNoOverlap(operation, connectorDataAttrs, scanOnlyDataAttrs)
+ validatePartitionAttrsDeclared(operation, relation, connectorDataAttrs, scanOnlyDataAttrs)
+ }
+
+
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
- val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ if (supportsColumnUpdate && operation.representUpdateAsDeleteAndInsert) {
+ validateNoRowIdReassignment(operation, assignments, rowIdAttrs)
+ validateRowIdDeclared(operation, connectorDataAttrs, rowIdAttrs)
+ }
+
+ val narrowDataAttrs = if (supportsColumnUpdate) {
+ computeNarrowReadAttrs(relation, connectorDataAttrs, scanOnlyDataAttrs, assignments, cond)
+ } else {
+ relation.output
+ }
+
+ val readRelation = if (supportsColumnUpdate) {
+ buildNarrowRelationWithAttrs(relation, operationTable, narrowDataAttrs, metadataAttrs,
+ rowIdAttrs)
+ } else {
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ }
// 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 rowDeltaPlan = if (supportsColumnUpdate) {
Review Comment:
Reworked. The column-update decisions are now in two places: `resolveWriteAttrs` in `RewriteUpdateTable`, a single three-arm match that returns the write relation's columns, and `V2Writes.hasNarrowRows`, which the write builders, the distribution check, copy-on-write scan planning and the exec routing use. The plan builders are the existing ones and have no column-update branches.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +84,51 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated, copied, or reinserted record with metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must override this
+ * method when mixing in {@link SupportsColumnUpdates}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but does not override this method.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T metadata, T record) throws IOException {
Review Comment:
Renamed to `writeColumnUpdate(record)` and `writeColumnUpdate(metadata, record)`. The Javadoc says it receives both updated and copied rows of a group-based operation.
##########
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:
Following Anton's later comment, `scanOnlyDataAttributes()` was removed. Columns the write's distribution or ordering needs are declared in `requiredDataAttributes()`, and an undeclared one is now rejected with `COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS`. Before, the same write failed with `_LEGACY_ERROR_TEMP_1137` only when column pruning ran, and succeeded with it excluded. Your delta repro is covered by `column-update: undeclared distribution column is rejected`, with and without `ColumnPruning` excluded.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -267,13 +268,80 @@ abstract class RowLevelOperationSuiteBase
protected def checkLastWriteInfo(
expectedRowSchema: StructType = new StructType(),
expectedRowIdSchema: Option[StructType] = None,
- expectedMetadataSchema: Option[StructType] = None): Unit = {
+ expectedMetadataSchema: Option[StructType] = None,
+ expectedUpdateSchema: Option[StructType] = None): Unit = {
val info = table.lastWriteInfo
assert(info.schema == expectedRowSchema, "row schema must match")
val actualRowIdSchema = Option(info.rowIdSchema.orElse(null))
assert(actualRowIdSchema == expectedRowIdSchema, "row ID schema must match")
val actualMetadataSchema = Option(info.metadataSchema.orElse(null))
assert(actualMetadataSchema == expectedMetadataSchema, "metadata schema must match")
+ val actualUpdateSchema = Option(info.updateSchema.orElse(null))
+ assert(actualUpdateSchema == expectedUpdateSchema, "update schema must match")
+ }
+
+ protected def getUpdateSummary(): org.apache.spark.sql.connector.write.UpdateSummary = {
Review Comment:
The update summary helpers are back in `UpdateTableSuiteBase`, which imports `UpdateSummary`.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,57 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * Currently honored only for UPDATE. DELETE and MERGE ignore this interface: those operations
+ * receive full-width rows and {@link LogicalWriteInfo#updateSchema()} is absent for them.
+ *
+ * @since 4.3.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * The returned columns become the schema of updated, copied, and reinserted rows, in declared
+ * order. Implementations must include every column they want to receive (typically the columns
+ * reported by {@link RowLevelOperationInfo#updatedColumns()} plus any columns needed for row
+ * lookup or routing, e.g. a primary key).
+ * <p>
+ * If any of the columns from {@link RowLevelOperationInfo#updatedColumns()} are
+ * missing, an analysis exception is thrown.
+ * <p>
+ * For updates on nested fields such as {@code SET t.s.c1 = -1} the connector should declare the
+ * root struct column {@code s} rather than any nested field.
+ * <p>
+ * This also covers columns needed only for planning, e.g. resolving the table's partitioning
Review Comment:
SPARK-59721 (#58979) is merged, so I added the deferred test: `column-update: partition source column not in requiredDataAttributes` declares `[pk, id]` on a table partitioned by `dep`. It passes, and fails with `Unable to resolve dep given [pk,id,_partition,index]` if the #58979 change to `V2ScanPartitioningAndOrdering` is reverted. The description links the JIRA.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala:
##########
@@ -771,6 +777,7 @@ trait WritingSparkTask[W <: DataWriter[InternalRow]] extends Logging with Serial
case class DataAndMetadataWritingSparkTask(
dataProj: ProjectingInternalRow,
+ updateDataProj: ProjectingInternalRow,
Review Comment:
The tasks no longer take a nullable projection; they take a `useWriteColumnUpdate: Boolean` and reuse the row projection.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -109,6 +123,43 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
relation)
}
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation. Non-existent columns are rejected with an `AnalysisException`.
+ */
+ protected def resolveRequiredDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes
+ if (refs.isEmpty) {
+ throw QueryCompilationErrors.emptyRequiredDataAttributesError(operation.getClass.getName)
+ }
+ val nested = refs.filter(_.fieldNames.length != 1).map(_.describe()).toImmutableArraySeq
+ if (nested.nonEmpty) {
+ throw QueryCompilationErrors.nestedRequiredDataAttributeError(
+ operation.getClass.getName, nested)
+ }
+ val normalizedNames = refs.map { ref =>
+ val name = ref.fieldNames.head
+ if (conf.caseSensitiveAnalysis) name else name.toLowerCase(Locale.ROOT)
+ }
+ val duplicates = normalizedNames.groupBy(identity).collect {
+ case (_, occurrences) if occurrences.length > 1 => occurrences.head
+ }.toSeq
Review Comment:
The duplicate error reports the spelling the connector declared, as the nested check does. `["PK", "Pk"]` reports `[PK]` (`column-update: case-only duplicate requiredDataAttributes throws AnalysisException`).
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -157,29 +228,112 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, required metadata attrs, and any connector-declared
+ // data attrs for column-update writes)
val rowAttrs = relation.output
+ val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates]
+ val connectorDataAttrs = resolveConnectorDataAttrs(relation, operation, assignments)
+
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
- val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ if (supportsColumnUpdate && operation.representUpdateAsDeleteAndInsert) {
+ validateNoRowIdReassignment(operation, relation, connectorDataAttrs, assignments, rowIdAttrs)
+ validateRowIdDeclared(operation, connectorDataAttrs, rowIdAttrs)
+ }
+
+ val narrowDataAttrs = if (supportsColumnUpdate) {
+ computeNarrowReadAttrs(relation, connectorDataAttrs, assignments, cond)
+ } else {
+ relation.output
+ }
+
+ val readRelation = if (supportsColumnUpdate) {
+ buildNarrowRelationWithAttrs(relation, operationTable, narrowDataAttrs, metadataAttrs,
+ rowIdAttrs)
+ } else {
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ }
// 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 rowDeltaPlan = if (supportsColumnUpdate) {
+ if (operation.representUpdateAsDeleteAndInsert) {
+ buildNarrowDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else {
+ buildColumnUpdateProjection(
+ matchedRowsPlan, assignments, rowIdAttrs, metadataAttrs, connectorDataAttrs)
+ }
} else {
- buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs)
+ if (operation.representUpdateAsDeleteAndInsert) {
+ buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else {
+ buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, 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(
+ rowDeltaPlan, rowAttrs, rowIdAttrs, metadataAttrs, connectorDataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond)
}
+ /**
+ * Builds the WriteDelta projection for the column-update path.
+ */
+ private def buildColumnUpdateProjection(
+ plan: LogicalPlan,
+ assignments: Seq[Assignment],
+ rowIdAttrs: Seq[Attribute],
+ metadataAttrs: Seq[Attribute],
+ connectorDataAttrs: Seq[AttributeReference]): LogicalPlan = {
+
+ val assignedValues = assignments.collect {
+ case Assignment(key: Attribute, value) if !isIdentityAssignment(key, value) =>
+ Alias(value, key.name)()
+ }
+
+ // Connector-required columns whose value isn't being changed by the UPDATE: pass through the
+ // current value so the connector receives a complete write row. Row-ID columns are excluded
+ // here even when also connector-required (e.g. a primary key used for both row lookup and
+ // write payload); they are emitted once, below, via rowIdValues.
+ val assignedKeyIds = collectUpdatedAttrs(assignments).map(_.exprId).toSet
+ val rowIdAttrSet = AttributeSet(rowIdAttrs)
+ val connectorPassThroughValues = connectorDataAttrs.filterNot(a =>
+ assignedKeyIds.contains(a.exprId) || rowIdAttrSet.contains(a))
+
+ val metadataAttrSet = AttributeSet(metadataAttrs)
+ val metadataValues = plan.output.filter(metadataAttrSet.contains).map { attr =>
+ if (MetadataAttribute.isPreservedOnUpdate(attr)) {
+ attr
+ } else {
+ Alias(Literal(null, attr.dataType), attr.name)(explicitMetadata = Some(attr.metadata))
+ }
+ }
+
+ val rowIdValues = plan.output.filter(rowIdAttrSet.contains)
Review Comment:
The rewrite now uses the existing projection builders, which project the original row ID under a separate name, so `pk` appears once. The row-ID reassignment tests with a CHECK constraint on `pk` still pass, for example `column-update: CHECK constraint on a reassigned row-ID column is enforced`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -228,4 +382,169 @@ 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, validating that every assigned column is covered. Returns `Nil` otherwise.
+ */
+ private def resolveConnectorDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation,
+ assignments: Seq[Assignment]): Seq[AttributeReference] = operation match {
+ case scu: SupportsColumnUpdates =>
+ val attrs = resolveRequiredDataAttrs(relation, scu)
+ validateUpdatedColumnsSubset(scu, assignments, attrs)
+ attrs
+ case _ => Nil
+ }
+
+ /**
+ * Computes the narrow set of data columns that must be present in the scan for a column-update
+ * write: connector-declared attrs (`requiredDataAttributes()`), unioned with any table columns
+ * referenced by non-identity assignment RHS expressions, the operation condition, and the
+ * table's CHECK constraints.
+ */
+ 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 checkConstraintAttrs = resolveCheckConstraintAttrs(relation)
+ dedupAttrs(connectorDataAttrs ++ extraRefs ++ checkConstraintAttrs)
+ }
+
+ /**
+ * Returns the relation attributes referenced by the table's CHECK constraints, so that
+ * narrowing does not drop a column `ResolveTableConstraints` needs later to re-validate the
+ * rewritten write query. Falls back to the full relation output whenever a constraint's
+ * condition cannot be resolved to a `Predicate` (see `Check#predicate()`).
+ */
+ private def resolveCheckConstraintAttrs(
+ relation: DataSourceV2Relation): Seq[AttributeReference] = {
+ val checks = Option(relation.table.constraints).getOrElse(Array.empty[Constraint]).collect {
+ case c: Check => c
+ }
+ if (checks.isEmpty) {
+ Nil
+ } else if (checks.exists(_.predicate() == null)) {
+ relation.output
+ } else {
+ val refs = checks.flatMap(_.predicate().references()).toSeq
+ V2ExpressionUtils.resolveRefs[AttributeReference](refs, relation)
Review Comment:
With the read relation no longer narrowed, CHECK constraints resolve against the full write query, so a nested-field CHECK such as `CHECK (s.c1 > 0)` no longer needs special handling. Covered by `CHECK constraint on a nested field narrows correctly` in both suites.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -228,4 +382,169 @@ 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, validating that every assigned column is covered. Returns `Nil` otherwise.
+ */
+ private def resolveConnectorDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation,
+ assignments: Seq[Assignment]): Seq[AttributeReference] = operation match {
+ case scu: SupportsColumnUpdates =>
+ val attrs = resolveRequiredDataAttrs(relation, scu)
+ validateUpdatedColumnsSubset(scu, assignments, attrs)
+ attrs
+ case _ => Nil
+ }
+
+ /**
+ * Computes the narrow set of data columns that must be present in the scan for a column-update
+ * write: connector-declared attrs (`requiredDataAttributes()`), unioned with any table columns
+ * referenced by non-identity assignment RHS expressions, the operation condition, and the
+ * table's CHECK constraints.
+ */
+ 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 checkConstraintAttrs = resolveCheckConstraintAttrs(relation)
+ dedupAttrs(connectorDataAttrs ++ extraRefs ++ checkConstraintAttrs)
Review Comment:
This no longer applies: `RewriteUpdateTable` now keeps the existing read relation and narrows only the write relation, so no attribute in the connector's or the user's spelling can reach a scan. The declared attributes are still mapped back to the relation's own attributes by expr ID. I kept the regression tests for your example on both paths, and they now also check that the scan uses the table's spelling (`undeclared, differently-cased columns in the RHS and WHERE clause are read with the table's own spelling`).
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -88,6 +92,16 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
relation.copy(table = table, output = attrs)
}
+ protected def buildNarrowRelationWithAttrs(
+ relation: DataSourceV2Relation,
+ table: RowLevelOperationTable,
+ dataAttrs: Seq[AttributeReference],
+ metadataAttrs: Seq[AttributeReference],
+ rowIdAttrs: Seq[AttributeReference] = Nil): DataSourceV2Relation = {
+ val attrs = dedupAttrs(dataAttrs ++ rowIdAttrs ++ metadataAttrs)
Review Comment:
This no longer applies: the read relation keeps every table column, so `PushDownUtils.pruneColumns` and `toOutputAttrs` see the same relation as for any other row-level operation. A scan builder without `SupportsPushDownRequiredColumns` reads every column and the write stays correct; the Javadoc now says the builder should implement it. The partial-pruning case is covered by `scan that returns unrequested columns still writes narrow rows` in both suites.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -109,6 +123,43 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
relation)
}
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation. Non-existent columns are rejected with an `AnalysisException`.
+ */
+ protected def resolveRequiredDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes
+ if (refs.isEmpty) {
+ throw QueryCompilationErrors.emptyRequiredDataAttributesError(operation.getClass.getName)
+ }
+ val nested = refs.filter(_.fieldNames.length != 1).map(_.describe()).toImmutableArraySeq
+ if (nested.nonEmpty) {
+ throw QueryCompilationErrors.nestedRequiredDataAttributeError(
+ operation.getClass.getName, nested)
+ }
+ val normalizedNames = refs.map { ref =>
+ val name = ref.fieldNames.head
+ if (conf.caseSensitiveAnalysis) name else name.toLowerCase(Locale.ROOT)
+ }
+ val duplicates = normalizedNames.groupBy(identity).collect {
+ case (_, occurrences) if occurrences.length > 1 => occurrences.head
+ }.toSeq
+ if (duplicates.nonEmpty) {
+ throw QueryCompilationErrors.duplicateRequiredDataAttributeError(
+ operation.getClass.getName, duplicates)
+ }
+ val resolved = V2ExpressionUtils.resolveRefs[AttributeReference](
Review Comment:
A metadata column in `requiredDataAttributes()` is now rejected with `COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE`, and the Javadoc says to return metadata columns from `requiredMetadataAttributes()` instead. Covered on the in-place and split delta paths (`metadata column in requiredDataAttributes throws AnalysisException`).
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala:
##########
@@ -513,7 +557,31 @@ case class WriteDelta(
case Some(projection) => DataTypeUtils.toAttributes(projection.schema)
case None => Nil
}
- table.skipSchemaResolution || areCompatible(inRowAttrs, outRowAttrs)
+ val inUpdateAttrs = projections.updateRowProjection match {
+ case Some(projection) => DataTypeUtils.toAttributes(projection.schema)
+ case None => Nil
+ }
+ // `rowProjection` (INSERT-tagged rows) validates against `outRowAttrs`. For column-update
+ // connectors, `updateRowProjection` is narrow and validates against `projectedDataAttrs`
+ // (the connector-declared narrow set) instead. When the connector does not mix in
+ // `SupportsColumnUpdates`, `updateRowProjection` is absent and `updateResolved` is trivially
+ // true.
+ val insertResolved = table.skipSchemaResolution ||
+ (inUpdateAttrs.nonEmpty && inRowAttrs.isEmpty) ||
+ areCompatible(inRowAttrs, outRowAttrs)
+ val updateResolved = table.skipSchemaResolution ||
+ inUpdateAttrs.isEmpty || dataAttrsResolved(inUpdateAttrs)
+ insertResolved && updateResolved
+ }
+
+ /**
+ * Validates the narrow-write-schema row projection output for a column-update connector.
+ * The write schema must exactly match the columns declared via
+ * `SupportsColumnUpdates.requiredDataAttributes()` (same columns, same order).
+ */
+ private def dataAttrsResolved(inRowAttrs: Seq[Attribute]): Boolean = {
Review Comment:
This code is gone. `ReplaceData` and `WriteDelta` validate against the narrowed write relation with the existing `rowAttrsResolved`. The only addition is `RowLevelWrite.references`, defined once in the trait.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -252,6 +304,413 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends only the
+ // declared + assigned columns in the row projection instead of the full row schema. The base
+ // class composes its required-attrs set as `pk` (the row-lookup key) and `dep` (an
+ // unconditionally declared base column) plus whatever columns Spark reports as being
+ // assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = {
+ val base = Seq(FieldReference("pk"), FieldReference("dep"))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ protected def clusterColumnRef: NamedReference = PARTITION_COLUMN_REF
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.updateSchema().isPresent) {
+ capturedInfo.updateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(Array(clusterColumnRef))
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ Array[SortOrder](
+ LogicalExpressions.sort(
+ clusterColumnRef,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering())
+ )
+ }
+
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ // Use the narrow update schema for UPDATE/COPY/REINSERT rows when present.
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For column-update writes, rows contain only the assigned columns
+ // (narrow schema from LogicalWriteInfo). We expand each row to the full table
+ // schema by overlaying write-schema columns on the base row found by pk.
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+
+ val mergedData = newData.map { buf =>
+ val merged = new BufferedRows(buf.key, schema)
+ val updateOpName = UTF8String.fromString(Update.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == updateOpName) {
+ val pk = logRow.getInt(1)
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val baseRow = dataMap.values.iterator.flatten
+ .flatMap(_.rows)
+ .find(r => r.getInt(schema.fieldIndex("pk")) == pk)
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ val field = schema.fields(i)
+ val value = if (i < base.numFields) {
+ base.get(i, field.dataType)
+ } else {
+ ResolveDefaultColumns.getExistenceDefaultValue(field)
+ }
+ fullRow.update(i, value)
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ writeFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ merged.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ merged.rows.append(insertRow.copy())
+ }
+ }
+ merged
+ }
+
+ withDeletes(newData)
+ withData(mergedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ }
+ }
+ }
+ }
+
+ class DeltaBasedColumnUpdateOperationWithReqAttrs(
+ command: Command,
+ reqCols: Array[String],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, options = options) {
+ override def requiredDataAttributes(): Array[NamedReference] = reqCols.map(FieldReference(_))
+ }
+
+ class DeltaBasedColumnUpdateOperationFromInfo(
+ command: Command,
+ updatedCols: Seq[NamedReference],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols, options) {
+ }
+
+ class DeltaBasedColumnUpdateSplitOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols, options) {
+ override def representUpdateAsDeleteAndInsert(): Boolean = true
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.updateSchema().isPresent) {
+ capturedInfo.updateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+ override def requiredDistribution(): Distribution =
+ Distributions.clustered(Array(PARTITION_COLUMN_REF))
+ override def requiredOrdering(): Array[SortOrder] = Array[SortOrder](
+ LogicalExpressions.sort(
+ PARTITION_COLUMN_REF,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()))
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For delete+reinsert with narrow writes, the REINSERT row has only the
+ // connector-declared columns (requiredDataAttributes order).
+ // pk is the first field in the write schema (declared before updatedCols).
+ // Reconstruct the full row by overlaying the narrow row onto the original.
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+ val reinsertOpName = UTF8String.fromString(Reinsert.toString)
+ val pkIdx = writeFieldIdx("pk")
+
+ val expandedData = newData.map { buf =>
+ val expanded = new BufferedRows(buf.key, schema)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == reinsertOpName) {
Review Comment:
The split fixture's commit now handles `Insert` as well as `Reinsert`, and MERGE runs on the split path too (`column-update split: MERGE falls back to full-width rows`).
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +84,54 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated, copied, or reinserted record with metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} without also mixing in
+ * {@link SupportsDelta} (i.e. group-based operations) receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. An operation that mixes in both
+ * instead delivers narrow rows through {@link DeltaWriter#update} / {@link DeltaWriter#reinsert}
+ * and never calls this method. By default, delegates to {@link #writeUpdate(Object)} for
+ * connectors that do not need metadata; implementations that do need it should override this
+ * method directly.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but overrides neither {@code writeUpdate} overload.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T metadata, T record) throws IOException {
+ writeUpdate(record);
+ }
+
+ /**
+ * Writes one updated, copied, or reinserted record without metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} without also mixing in
+ * {@link SupportsDelta} (i.e. group-based operations) receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must override this method,
+ * or {@link #writeUpdate(Object, Object)}, when mixing in {@link SupportsColumnUpdates} without
+ * {@link SupportsDelta}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but overrides neither {@code writeUpdate} overload.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T record) throws IOException {
Review Comment:
The overloads are now `writeColumnUpdate(record)` and `writeColumnUpdate(metadata, record)`. The Javadoc says Spark calls the metadata overload when the operation declares `requiredMetadataAttributes()` and the record-only overload otherwise. The record-only default throws and the metadata overload delegates to it, so a writer must override the overload Spark calls. Tests cover both dispatches, a writer that overrides only `writeColumnUpdate(record)`, and a writer that overrides neither.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -82,11 +83,29 @@ class InMemoryRowLevelOperationTable private (
private final val noMetadata = properties.getOrDefault(NO_METADATA, "false") == "true"
private final val useCatalystRuntimeFiltering =
properties.getOrDefault(USE_CATALYST_RUNTIME_FILTERING, "false") == "true"
+ private final val COLUMN_UPDATE = "column-update"
+ private final val COLUMN_UPDATE_REQ_ATTRS = "column-update-req-attrs"
+ private final val COLUMN_UPDATE_COW = "column-update-cow"
+ private final val COLUMN_UPDATE_COW_NO_WRITE_UPDATE = "column-update-cow-no-write-update"
+ private final val COLUMN_UPDATE_FROM_INFO = "column-update-from-info"
+ private final val COLUMN_UPDATE_SPLIT = "column-update-split"
+ private final val COLUMN_UPDATE_SPLIT_REQ_ATTRS = "column-update-split-req-attrs"
+ private final val COLUMN_UPDATE_EMPTY_REQ_ATTRS = "column-update-empty-req-attrs"
+ private final val COLUMN_UPDATE_SPLIT_MISSING_ROW_ID = "column-update-split-missing-row-id"
// used in row-level operation tests to verify replaced partitions
var replacedPartitions: Seq[Seq[Any]] = Seq.empty
// used in row-level operation tests to verify reported write schema
var lastWriteInfo: LogicalWriteInfo = _
+ // used in column-update tests to verify that Spark passed the correct updated column list
+ // to the connector via RowLevelOperationInfo.updatedColumns()
+ var lastUpdatedColumns: Array[NamedReference] = Array.empty
Review Comment:
`copy()` now carries `lastUpdatedColumns`, and the transactional table sets it on commit instead of when the operation is built.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub) via reviews" <re...@spark.apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4428867391
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,9 +87,67 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated or copied record with metadata in the column update layout.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present for a row-level operation that
+ * does not mix in {@link SupportsDelta}, Spark passes updated and copied records to this method
+ * instead of {@link #write(Object, Object)} if the operation returns a non-empty
+ * {@link RowLevelOperation#requiredMetadataAttributes()}, and to
+ * {@link #writeColumnUpdate(Object)} otherwise. The record follows
+ * {@link LogicalWriteInfo#columnUpdateSchema()} and the metadata follows
+ * {@link LogicalWriteInfo#metadataSchema()}. Operations that mix in {@link SupportsDelta} receive
+ * such rows through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert} instead.
+ * <p>
+ * By default, delegates to {@link #writeColumnUpdate(Object)} and drops the metadata.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if neither this method nor
+ * {@link #writeColumnUpdate(Object)} is overridden.
+ *
+ * @since 4.4.0
+ */
+ default void writeColumnUpdate(T metadata, T record) throws IOException {
+ writeColumnUpdate(record);
+ }
+
+ /**
+ * Writes one updated or copied record without metadata in the column update layout.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present for a row-level operation that
+ * does not mix in {@link SupportsDelta}, Spark passes updated and copied records to this method
+ * instead of {@link #write(Object)} if the operation returns no
+ * {@link RowLevelOperation#requiredMetadataAttributes()}. The record follows
+ * {@link LogicalWriteInfo#columnUpdateSchema()}.
+ * <p>
+ * A writer for such an operation must override this method, unless the operation returns a
+ * non-empty {@link RowLevelOperation#requiredMetadataAttributes()} and the writer overrides
+ * {@link #writeColumnUpdate(Object, Object)}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if this method is not overridden.
+ *
+ * @since 4.4.0
+ */
+ default void writeColumnUpdate(T record) throws IOException {
Review Comment:
This also makes sense.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub) via reviews" <re...@spark.apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4428856251
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -68,8 +70,11 @@ public interface DataWriter<T> extends Closeable {
* Writes one record with metadata.
* <p>
* This method is used by group-based row-level operations to pass back metadata for records
- * that are updated or copied. New records added during a MERGE operation are written using
- * {@link #write(Object)} as there is no metadata associated with those records.
+ * that are updated or copied. When {@link LogicalWriteInfo#columnUpdateSchema()} is present for
Review Comment:
This makes sense.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4372866939
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala:
##########
@@ -380,14 +380,17 @@ case class ReplaceDataExec(
write: Write,
rowLevelCommand: RowLevelOperation.Command,
tableName: String,
+ useWriteColumnUpdate: Boolean,
Review Comment:
**42.** Minor: after item 27, this flag is the only record of whether UPDATE and COPY rows follow `columnUpdateSchema()`, and it cannot be derived from `write` or `projections`. Plugins that replace `ReplaceDataExec` rebuild their node from the other fields, so they compile unchanged and drop it:
- NVIDIA/cudf-spark builds `GpuReplaceDataExec(childPlan, cpuExec.refreshCache, cpuExec.projections, gpuWrite)` ([ReplaceDataExecShim.scala L41](https://github.com/NVIDIA/cudf-spark/blob/a97083ad4c47cd35d894c7b5d91a5899fce748b5/sql-plugin/src/main/spark400/scala/com/nvidia/spark/rapids/shims/ReplaceDataExecShim.scala#L41), whose shim list includes 420 and 500), and its task calls `writer.write(metadata, projected)` or `writer.write(projected)` ([GpuReplaceDataExec.scala L87-L89](https://github.com/NVIDIA/cudf-spark/blob/a97083ad4c47cd35d894c7b5d91a5899fce748b5/sql-plugin/src/main/spark400/scala/org/apache/spark/sql/execution/datasources/v2/GpuReplaceDataExec.scala#L87-L89)).
- Gluten builds `VeloxIcebergReplaceDataExec(original.query, original.refreshCache, original.write)` ([VeloxIcebergReplaceDataExec.scala L31-L37](https://github.com/apache/gluten/blob/ffb8af3fedfc5c9c74244c83a4c33adb87caccd8/backends-velox/src-iceberg/main/scala/org/apache/gluten/execution/VeloxIcebergReplaceDataExec.scala#L31-L37)).
Nothing is affected today, but once a connector they offload, such as Iceberg copy-on-write, mixes in `SupportsColumnUpdates`, narrow UPDATE and COPY rows would reach `write(...)` instead of `writeColumnUpdate(...)`, without any error. Could the Scaladoc of this parameter say that a node replacing this one must send UPDATE and COPY rows to `writeColumnUpdate` when it is set, or not replace it, and could the compatibility paragraph of the PR description say the same?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -273,6 +339,423 @@ class InMemoryRowLevelOperationTable private (
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
}
+ // Finds the one current row with the given row-ID values, matching a data column by value and
+ // `_partition` by the partition key.
+ private def currentRowByRowId(values: Seq[(String, Any)]): InternalRow = {
+ val matches = dataMap.values.iterator.flatten.flatMap { buf =>
+ buf.rows.filter { row =>
+ values.forall {
+ case (name, value) if name == PartitionKeyColumn.name => buf.keyString() == value.toString
+ case (name, value) =>
+ val i = schema.fieldIndex(name)
+ row.get(i, schema.fields(i).dataType) == value
+ }
+ }
+ }.toSeq
+ matches match {
+ case Seq(row) => row
+ case _ => throw new IllegalStateException(
+ s"Expected one row with row ID $values, found ${matches.size}")
+ }
+ }
+
+ // Rebuilds a full table row for a narrow column-update row: the base row's values (or the
+ // existence default for a column added after the base row was written), overlaid with the
+ // narrow row's columns.
+ private def overlayNarrowRow(
+ baseRow: Option[InternalRow],
+ narrowRow: InternalRow,
+ narrowFieldIdx: Map[String, Int]): InternalRow = {
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ val field = schema.fields(i)
+ val value = if (i < base.numFields) {
+ base.get(i, field.dataType)
+ } else {
+ ResolveDefaultColumns.getExistenceDefaultValue(field)
+ }
+ fullRow.update(i, value)
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ narrowFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ fullRow
+ }
+
+ // The base columns `pk` and `dep`, where the table has them, followed by the updated columns.
+ private def declaredColumns(updatedCols: Seq[NamedReference]): Array[NamedReference] = {
+ val base = Seq("pk", "dep").filter(schema.fieldNames.contains).map(FieldReference(_))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep`, where the table has them, plus whatever columns Spark reports as
+ // being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
Review Comment:
**46.** Nit: this extends the case class `DeltaBasedOperation` (L275), so this class and its subclasses (L507, L515, L597, L608) inherit `equals`, `hashCode` and `copy` on `(command, options)`. For the same table, `new DeltaBasedColumnUpdateOperation(UPDATE)` equals `DeltaBasedOperation(UPDATE, CaseInsensitiveStringMap.empty())`, which is exactly the pair at `ColumnUpdateTableSharedTests.scala:157-158`, and a split and a non-split operation, or two with different `requiredDataAttributes()`, are equal too. `RowLevelOperationTable` and the relation then compare equal as well, so a future test that swaps one of these operations in with `plan.transform` would keep the old node (`fastEquals`), and a `Set` or `Map` of them, unlike the `Seq` at `ColumnUpdateTableSharedTests.scala:155`, would collapse them. No current test is affected. Could this class override `equals` and `hashCode` with reference equality?
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,105 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, updated, copied, and reinserted
+ * rows follow it: group-based operations receive them through
+ * {@link DataWriter#writeColumnUpdate(Object, Object)} or
+ * {@link DataWriter#writeColumnUpdate(Object)}, and operations that mix in {@link SupportsDelta}
+ * through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert}. Inserted rows always follow
+ * {@link LogicalWriteInfo#schema()} and arrive through {@link DataWriter#write(Object)} and
+ * {@link DeltaWriter#insert}. When {@link LogicalWriteInfo#columnUpdateSchema()} is absent, every
+ * row follows {@link LogicalWriteInfo#schema()}, as for any other operation. Connectors must
+ * decide which layout to expect by whether {@link LogicalWriteInfo#columnUpdateSchema()} is
+ * present, not by the command. A builder that wants full-width rows for some command can check
+ * {@link RowLevelOperationInfo#command()} and build an operation without this mix-in. Currently
+ * Spark narrows only UPDATE.
+ * <p>
+ * The scan builder returned by {@link #newScanBuilder} should implement
+ * {@link org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns}, so that the scan
+ * does not read columns the command neither references nor declares. Otherwise the scan reads
+ * every column. For operations that do not mix in {@link SupportsDelta}, only top-level columns
+ * are pruned: a struct column the command references through a nested field is read in full.
+ *
+ * @since 4.4.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * When Spark narrows the write, the returned columns become
+ * {@link LogicalWriteInfo#columnUpdateSchema()}, in declared order. Implementations must include
Review Comment:
**43.** Minor: this says that the declared columns become `columnUpdateSchema()` "in declared order", but not that they take the table's spelling. Declarations resolve case-insensitively, and Spark maps each one back to the relation's own attribute (`RewriteRowLevelCommand.scala:409-413`), as "requiredDataAttributes resolves case-insensitively without adopting the declared spelling" checks (`DeltaBasedColumnUpdateTableSuite.scala:260-274`: declaring `PK,SALARY,dep` gives `pk, salary, dep`). `updatedColumns()` uses the table's spelling as well (`RewriteRowLevelCommand.scala:77`). So a connector that declares `FieldReference("PK")` for a column `pk` and follows L62-67 with `Arrays.asList(info.updatedColumns()).contains(FieldReference("PK"))` gets `false` for `UPDATE t SET pk = 2`, because `FieldReference` equality is case-sensitive. It then builds the narrow operation and looks rows up by the new `pk`, as in item 5, and its `columnUpdateSchema().get().fieldIndex("PK")` throws as wel
l. Could this paragraph and `RowLevelOperationInfo.updatedColumns()` say that the names follow the table's spelling?
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DeltaWriter.java:
##########
@@ -52,6 +55,9 @@ public interface DeltaWriter<T> extends DataWriter<T> {
* Reinserts a row with metadata.
* <p>
* This method handles the insert portion of updated rows split into deletes and inserts.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, the {@code row} follows it;
+ * otherwise it follows {@link LogicalWriteInfo#schema()}.
Review Comment:
**41. The 4.0 default `reinsert(metadata, row)` still forwards to `insert(row)` (L68-70), so a split-update `SupportsColumnUpdates` writer that keeps it gets narrow rows, without their metadata, in `insert()`.** For a narrow UPDATE on an operation that represents updates as deletes and inserts, `WriteToDataSourceV2Exec.scala:903` and `:957` pass the narrow row projection to `reinsert`, so this default hands a `columnUpdateSchema()` row to `insert()`. However, `SupportsColumnUpdates.java:32-34` says that inserted rows always follow `schema()` and arrive through `DeltaWriter#insert`, and `schema()` is empty for UPDATE (`LogicalWriteInfo.java:51-53`). The default also drops `metadata`, which is the only identity of the reinserted row when the row ID is a metadata column preserved on reinsert, a setup that `SupportsColumnUpdates.java:78-81` and `validateRowIdDeclared` (`RewriteUpdateTable.scala:323-337`) accept.
With the setup of "metadata row ID preserved on reinsert need not be declared" (`DeltaBasedColumnUpdateTableSuite.scala:486-504`: `_partition` as the row ID, `columnUpdateSchema()` = `[salary]`) and a writer that implements `delete`, `update` and `insert` but keeps this default, `UPDATE t SET salary = -1 WHERE dep = 'hr'` deletes the matched row and then calls `insert([-1])`. Depending on the writer, the row is lost or written without `pk` and `dep`, and a writer that decodes `insert` rows by the table layout fails with an `IndexOutOfBoundsException` from `ProjectingInternalRow`. The group-based path is guarded against the same mistake, because `DataWriter.writeColumnUpdate(T)` throws by default (`DataWriter.java:138-142`), but nothing guards this one, and no test reaches it, since the only split writer, `DeltaBufferWriter`, overrides `reinsert` (`InMemoryRowLevelOperationTable.scala:876-881`). Because this default has been public since 4.0.0, changing it would affect existing wri
ters. Could this Javadoc say that a writer for an operation that mixes in `SupportsColumnUpdates` and represents updates as deletes and inserts must override `reinsert`, and could `SupportsColumnUpdates.java:32-34` be qualified accordingly?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,362 @@
+/*
+ * 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 java.util.regex.Pattern
+
+import org.apache.spark.SparkUnsupportedOperationException
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTable
+import org.apache.spark.sql.connector.write.RowLevelOperationTable
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+
+class GroupBasedColumnUpdateTableSuite extends UpdateTableSuiteBase
+ with ColumnUpdateTableSharedTests {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-cow", "true")
+ props
+ }
+
+ override protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String] = {
+ java.util.Map.of("column-update-cow-req-attrs", reqAttrs)
+ }
+
+ override protected def fullWidthTableProps: java.util.Map[String, String] = java.util.Map.of()
+
+ test("column-update ReplaceData: runtime group filtering reads only the referenced columns") {
+ Seq(true, false).foreach { dppEnabled =>
+ Seq(true, false).foreach { aqeEnabled =>
+ withSQLConf(
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> dppEnabled.toString,
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString) {
+ withTable(tableNameAsString) {
+ withTempView("matched_pk") {
+ createAndInitTable("pk INT, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ import testImplicits._
+ val matchedPkDF = Seq(Some(1), None).toDF()
+ matchedPkDF.createOrReplaceTempView("matched_pk")
+
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = -1 " +
+ s"WHERE pk IN (SELECT * FROM matched_pk)",
+ primaryScanSchema =
+ "pk INT, salary INT, dep STRING, _partition STRING, index INT",
+ groupFilterScanSchema = Some("pk INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, 300, 30, "hr") :: Nil)
+
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("column-update ReplaceData: updated and copied rows reach writeColumnUpdate with " +
+ "metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", null), data = Row(1, "hr", -1)),
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", 1), data = Row(2, "hr", 200)))
+ }
+
+ test("column-update ReplaceData: without metadata, rows reach writeColumnUpdate(record)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "no-metadata", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD, DEP_FIELD, StructField("salary", IntegerType, nullable = true)))))
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Row(3, 300, 30, "software") :: Nil)
+ }
+
+ test("column-update ReplaceData: a writer overriding only writeColumnUpdate(record) " +
+ "receives rows with metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "column-update-cow-record-only-writer", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: connector missing writeColumnUpdate override is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-cow", "true", "column-update-cow-no-write-column-update", "true"))
+
+ checkError(
+ exception = intercept[SparkUnsupportedOperationException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED",
+ parameters = Map(
+ "class" -> "org.apache.spark.sql.connector.catalog.NoWriteColumnUpdateOverrideWriter"))
+ }
+
+ test("column-update ReplaceData: the subquery plan narrows the write relation and the scan") {
+ withTempView("deps") {
+ createAndInitTableWithProps(reqAttrsTableProps("pk,salary"))
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ val write = executeAndKeepAnalyzedWrite {
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE dep IN (SELECT dep FROM deps)",
+ "pk INT, dep STRING, salary INT, bonus INT, extra INT, _partition STRING, index INT")
+ }
+
+ assert(write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "pk", "salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+ }
+
+ test("column-update ReplaceData: runtime group filtering on an undeclared condition column") {
+ withTempView("matched_ids") {
+ // `dep` is declared because the scan filters groups at runtime by the partition column
+ createAndInitTableWithProps(reqAttrsTableProps("pk,dep,salary"))
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `id` is neither declared nor in the write relation; the group filter must still map it
+ // onto the second read of the table that the subquery rewrite adds
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, dep STRING, salary INT, bonus INT, extra INT, " +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = Some("id INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+
+ test("column-update ReplaceData: unread partition column skips runtime group filtering") {
+ withTempView("matched_ids") {
+ createAndInitTableWithProps(reqAttrsTableProps("pk,salary"))
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `dep` is neither declared nor referenced, so the scan doesn't read it and doesn't
+ // report it from filterAttributes(); every group is read and replaced
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, salary INT, bonus INT, extra INT, " +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = None)
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr", "software"))
+ }
+ }
+
+ test("column-update ReplaceData: runtime filter attribute the scan doesn't read is rejected") {
+ withTempView("matched_ids") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, salary INT",
+ """{ "pk": 1, "id": 1, "dep": "hr", "salary": 100 }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-cow-req-attrs", "pk,salary",
+ "column-update-cow-unread-filter-attrs", "true"))
+ import testImplicits._
+ Seq(1).toDF().createOrReplaceTempView("matched_ids")
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE id IN (SELECT * FROM matched_ids)")
+ },
+ condition = "DATA_SOURCE_INVALID_RUNTIME_FILTER_ATTRIBUTE.CANNOT_RESOLVE",
+ sqlState = Some("KD000"),
+ parameters = Map(
+ "attribute" -> Pattern.quote("`dep`"),
+ "method" -> Pattern.quote("filterAttributes()"),
+ "scanClass" -> (Pattern.quote(
+ classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperation].getName) +
+ """\$\$anon\$\d+\$\$anon\$\d+"""),
+ "relationOutput" -> Pattern.quote("\"STRUCT<pk: INT NOT NULL, id: INT, salary: INT, " +
+ "_partition: STRING NOT NULL COMMENT 'Partition key used to store the row', " +
+ "index: INT NOT NULL COMMENT 'Metadata column used to conflict with a data " +
+ "column'>\"")),
+ matchPVals = true)
+ }
+ }
+
+ test("column-update ReplaceData: undeclared distribution column is rejected") {
+ // without metadata, the write clusters by `dep`
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow-req-attrs", "pk,salary", "no-metadata", "true"))
+
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName
+ for {
+ cond <- Seq("pk = 1", "dep = 'hr'")
+ confs <- Seq(Nil, Seq(excludeColumnPruning))
+ } withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE $cond")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[`dep`]"))
+ }
+ }
+
+ test("column-update ReplaceData: UPDATE without column updates, DELETE and MERGE read every " +
+ "column") {
+ val fullScanSchema = "pk INT, salary INT, bonus INT, dep STRING, _partition STRING, index INT"
+ Seq(fullWidthTableProps, extraTableProps).foreach { props =>
+ withTable(tableNameAsString) {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ props)
+
+ if (props == fullWidthTableProps) {
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1", fullScanSchema)
+ }
+ executeAndCheckScan(s"DELETE FROM $tableNameAsString WHERE pk > 1", fullScanSchema)
+ withTempView("source") {
+ import testImplicits._
+ Seq((1, 150)).toDF("pk", "salary").createOrReplaceTempView("source")
+ executeAndCheckScan(
+ s"""MERGE INTO $tableNameAsString t
+ |USING source s
+ |ON t.pk = s.pk
+ |WHEN MATCHED THEN UPDATE SET salary = s.salary
+ |""".stripMargin,
+ fullScanSchema)
+ }
+
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, 150, 10, "hr") :: Nil)
+ }
+ }
+ }
+
+ test("column-update ReplaceData: an operation rebuilt with withTable keeps the updated columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ }
+ val relation = write.table.asInstanceOf[DataSourceV2Relation]
+ val operationTable = relation.table.asInstanceOf[RowLevelOperationTable]
+ val t = table
+ t.lastUpdatedColumns = null
+
+ val rebuilt = operationTable.withTable(t)
+
+ checkLastUpdatedColumns("salary")
+ assert(rebuilt.operationInfo eq operationTable.operationInfo)
+ val rebuiltQuery = write.query.transform {
+ case r: DataSourceV2Relation if r.table eq operationTable => r.copy(table = rebuilt)
+ }
+ assert(rebuiltQuery.exists {
+ case r: DataSourceV2Relation => r.table eq rebuilt
+ case _ => false
+ })
+ V2Writes(write.withNewTable(relation.copy(table = rebuilt)).withNewQuery(rebuiltQuery))
+ }
+
+ test("column-update ReplaceData: column update schema nullability covers UPDATE and COPY " +
+ "rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT NOT NULL, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": null, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow-req-attrs", "pk,salary,bonus"))
Review Comment:
**52.** Nit: this map is `reqAttrsTableProps("pk,salary,bonus")` (L41-43), which L158, L182 and L206 use. Item 36 asked the delta suite to use its own hook instead of such literals, and this one is left in this suite, so it would not follow an override of the hook. Could it be `reqAttrsTableProps("pk,salary,bonus")`?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -273,6 +339,423 @@ class InMemoryRowLevelOperationTable private (
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
}
+ // Finds the one current row with the given row-ID values, matching a data column by value and
+ // `_partition` by the partition key.
+ private def currentRowByRowId(values: Seq[(String, Any)]): InternalRow = {
+ val matches = dataMap.values.iterator.flatten.flatMap { buf =>
+ buf.rows.filter { row =>
+ values.forall {
+ case (name, value) if name == PartitionKeyColumn.name => buf.keyString() == value.toString
+ case (name, value) =>
+ val i = schema.fieldIndex(name)
+ row.get(i, schema.fields(i).dataType) == value
+ }
+ }
+ }.toSeq
+ matches match {
+ case Seq(row) => row
+ case _ => throw new IllegalStateException(
+ s"Expected one row with row ID $values, found ${matches.size}")
+ }
+ }
+
+ // Rebuilds a full table row for a narrow column-update row: the base row's values (or the
+ // existence default for a column added after the base row was written), overlaid with the
+ // narrow row's columns.
+ private def overlayNarrowRow(
+ baseRow: Option[InternalRow],
+ narrowRow: InternalRow,
+ narrowFieldIdx: Map[String, Int]): InternalRow = {
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ val field = schema.fields(i)
+ val value = if (i < base.numFields) {
+ base.get(i, field.dataType)
+ } else {
+ ResolveDefaultColumns.getExistenceDefaultValue(field)
+ }
+ fullRow.update(i, value)
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ narrowFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ fullRow
+ }
+
+ // The base columns `pk` and `dep`, where the table has them, followed by the updated columns.
+ private def declaredColumns(updatedCols: Seq[NamedReference]): Array[NamedReference] = {
+ val base = Seq("pk", "dep").filter(schema.fieldNames.contains).map(FieldReference(_))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep`, where the table has them, plus whatever columns Spark reports as
+ // being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = declaredColumns(updatedCols)
+
+ private val clusterExprs: Seq[Expression] = {
+ val refs = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY))
+ .map(_.split(",").map(name => FieldReference(name.trim)).toSeq)
+ .getOrElse(Seq(PARTITION_COLUMN_REF))
+ if (properties.getOrDefault(COLUMN_UPDATE_CLUSTER_BY_NESTED_TRANSFORM, "false") == "true") {
+ refs.map(ref => LogicalExpressions.apply("f", LogicalExpressions.apply("g", ref)))
+ } else {
+ refs
+ }
+ }
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ val writeSchema = info.columnUpdateSchema().orElse(info.schema())
+ val batchWrite = new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ physicalInfo: PhysicalWriteInfo): DeltaWriterFactory = {
+ newWriterFactory(info, writeSchema)
+ }
+
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ commitColumnUpdate(info, writeSchema, newData)
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ if (noMetadata && !properties.containsKey(COLUMN_UPDATE_CLUSTER_BY)) {
+ new DeltaWrite {
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ } else {
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(clusterExprs.toArray)
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ clusterExprs.map { expr =>
+ LogicalExpressions.sort(
+ expr,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()): SortOrder
+ }.toArray
+ }
+
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ }
+ }
+ }
+
+ protected def newWriterFactory(
+ info: LogicalWriteInfo,
+ writeSchema: StructType): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(writeSchema)
+ }
+
+ // For column-update writes, updated rows contain only the declared columns
+ // (columnUpdateSchema from LogicalWriteInfo). We expand each row to the full
+ // table schema by overlaying them on the base row found by pk.
+ protected def commitColumnUpdate(
+ info: LogicalWriteInfo,
+ writeSchema: StructType,
+ newData: Array[BufferedRows]): Unit = {
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+ val updateOpName = UTF8String.fromString(Update.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+
+ val mergedData = newData.map { buf =>
+ val merged = new BufferedRows(buf.key, schema)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == updateOpName) {
+ val pk = logRow.getInt(1)
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val baseRow = Some(currentRowByRowId(Seq("pk" -> pk)))
+ val fullRow = overlayNarrowRow(baseRow, narrowRow, writeFieldIdx)
+ merged.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ merged.rows.append(insertRow.copy())
+ }
+ }
+ merged
+ }
+
+ withDeletes(newData)
+ withData(mergedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+ }
+
+ class DeltaBasedColumnUpdateOperationWithReqAttrs(
+ command: Command,
+ reqCols: Array[String],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
Review Comment:
**51.** Nit: the dispatcher always passes `info.options` and the updated columns (L164-201), so this `options` default and those at L518, L600, L612 and L697, as well as `updatedCols = Nil` at L611, are never used, and neither is `pkRowId = true` of `DeltaBufferWriter` (L851), whose only construction passes it (L811). The SPARK-57681 options tests (`UpdateTableSuiteBase.scala:1274`, `:1292` and `:1349`) run only with each suite's `extraTableProps`, so if a req-attrs, split or row-ID branch stopped passing `info.options`, it would still compile, silently use empty options, and no test would notice. Could these defaults go?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/ColumnUpdateTableSharedTests.scala:
##########
@@ -0,0 +1,355 @@
+/*
+ * 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.{SparkException, SparkRuntimeException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.{InMemoryBaseTable, InMemoryRowLevelOperationTable}
+import org.apache.spark.sql.connector.write.{RowLevelOperation, RowLevelOperationTable}
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+// Column-update tests that hold for both delta-based and group-based operations; each suite
+// supplies its column-update table properties through `extraTableProps`.
+trait ColumnUpdateTableSharedTests extends UpdateTableSuiteBase {
+
+ // the properties of a column-update table whose connector declares exactly `reqAttrs`
+ protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String]
+
+ // the properties of a table of the same kind whose operation does not mix in column updates
+ protected def fullWidthTableProps: java.util.Map[String, String]
+
+ // the metadata columns the scan of a column-update table reads, in DDL form
+ protected def scanMetadataColumns: String = ", _partition STRING, index INT"
+
+ // the class name of the operation that `reqAttrsTableProps` builds
+ protected def reqAttrsOperationClass: String = if (deltaUpdate) {
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateOperationWithReqAttrs].getName
+ } else {
+ classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperationWithReqAttrs].getName
+ }
+
+ // t(pk, id, dep, salary, bonus, extra) with CHECK (extra > 0), partitioned by dep
+ protected def createAndInitTableWithProps(props: java.util.Map[String, String]): Unit = {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, salary INT, bonus INT, extra INT",
+ """{ "pk": 1, "id": 1, "dep": "hr", "salary": 100, "bonus": 10, "extra": 1 }
+ |{ "pk": 2, "id": 2, "dep": "hr", "salary": 200, "bonus": 20, "extra": 1 }
+ |{ "pk": 3, "id": 3, "dep": "software", "salary": 300, "bonus": 30, "extra": 1 }
+ |""".stripMargin,
+ props)
+ sql(s"ALTER TABLE $tableNameAsString ADD CONSTRAINT positive_extra CHECK (extra > 0)")
+ }
+
+ test("column-update: undeclared columns are not written and keep their values") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
Review Comment:
**45.** Minor: this reads `columnUpdateSchema.get` without a guard, as do L186, L201 and L348 here, `DeltaBasedColumnUpdateTableSuite.scala:147`, `:269`, `:301` and `:498`, and `checkLastColumnUpdateWriteLog` (`RowLevelOperationSuiteBase.scala:424`), which seven tests use. If a regression stops narrowing, they all fail with `java.util.NoSuchElementException: No value present`, without the expected names or the actual write info. Like `checkLastUpdatedColumns` (`RowLevelOperationSuiteBase.scala:409-416`, item 33), could a `checkLastColumnUpdateSchema(expectedNames: String*)` helper first assert that the schema is present, with a message, and could these sites and `checkColumnUpdateWrite` use it?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala:
##########
@@ -92,6 +95,38 @@ object GroupBasedRowLevelOperationScanPlanning extends Rule[LogicalPlan] with Pr
}
}
+ /**
+ * Returns the columns of `relation` that the projections and filters directly above each read
+ * of the table use, as the rest of the query sees only their output. Only whole columns are
+ * pruned, as the query is not rewritten for a pruned nested schema.
+ *
+ * An UPDATE with a subquery in its condition reads the table twice, once per Union child, and
+ * the analyzer gives the second read new expr IDs. Both reads share one scan, so the columns
+ * are matched to `relation` by name.
+ */
+ private def readAttrs(
+ plan: LogicalPlan,
+ table: RowLevelOperationTable,
+ relation: DataSourceV2Relation): Seq[AttributeReference] = {
+ val readNames = tableReads(plan, table).flatMap { case (projects, filters, read) =>
+ val references = AttributeSet((projects ++ filters).flatMap(_.references))
+ read.output.filter(references.contains).map(_.name)
+ }.toSet
+ relation.output.filter(attr => readNames.contains(attr.name))
+ }
+
+ private def tableReads(
Review Comment:
**48.** Nit: `table` is always `relation.table` (L52 only casts it), and this method only returns a tuple that `readAttrs` takes apart at L111; the `NamedExpression` (L21) and `RowLevelOperationTable` (L30) imports exist only for these signatures. Could `readAttrs` take only the relation and use a local recursive function? For example:
```scala
private def readAttrs(
plan: LogicalPlan,
relation: DataSourceV2Relation): Seq[AttributeReference] = {
def readNames(node: LogicalPlan): Seq[String] = node match {
case PhysicalOperation(projects, filters, r: DataSourceV2Relation)
if r.table eq relation.table =>
val references = AttributeSet((projects ++ filters).flatMap(_.references))
r.output.filter(references.contains).map(_.name)
case _ =>
node.children.flatMap(readNames)
}
val names = readNames(plan).toSet
relation.output.filter(attr => names.contains(attr.name))
}
```
L70 would then call `readAttrs(rd.query, relation)`. It has to keep stopping at the first match, because `plan.collect` would also match the `Filter` under a `Project`, with `projects = r.output`.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -273,6 +339,423 @@ class InMemoryRowLevelOperationTable private (
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
}
+ // Finds the one current row with the given row-ID values, matching a data column by value and
+ // `_partition` by the partition key.
+ private def currentRowByRowId(values: Seq[(String, Any)]): InternalRow = {
+ val matches = dataMap.values.iterator.flatten.flatMap { buf =>
+ buf.rows.filter { row =>
+ values.forall {
+ case (name, value) if name == PartitionKeyColumn.name => buf.keyString() == value.toString
+ case (name, value) =>
+ val i = schema.fieldIndex(name)
+ row.get(i, schema.fields(i).dataType) == value
+ }
+ }
+ }.toSeq
+ matches match {
+ case Seq(row) => row
+ case _ => throw new IllegalStateException(
+ s"Expected one row with row ID $values, found ${matches.size}")
+ }
+ }
+
+ // Rebuilds a full table row for a narrow column-update row: the base row's values (or the
+ // existence default for a column added after the base row was written), overlaid with the
+ // narrow row's columns.
+ private def overlayNarrowRow(
+ baseRow: Option[InternalRow],
+ narrowRow: InternalRow,
+ narrowFieldIdx: Map[String, Int]): InternalRow = {
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ val field = schema.fields(i)
+ val value = if (i < base.numFields) {
+ base.get(i, field.dataType)
+ } else {
+ ResolveDefaultColumns.getExistenceDefaultValue(field)
+ }
+ fullRow.update(i, value)
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ narrowFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ fullRow
+ }
+
+ // The base columns `pk` and `dep`, where the table has them, followed by the updated columns.
+ private def declaredColumns(updatedCols: Seq[NamedReference]): Array[NamedReference] = {
+ val base = Seq("pk", "dep").filter(schema.fieldNames.contains).map(FieldReference(_))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep`, where the table has them, plus whatever columns Spark reports as
+ // being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = declaredColumns(updatedCols)
+
+ private val clusterExprs: Seq[Expression] = {
+ val refs = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY))
+ .map(_.split(",").map(name => FieldReference(name.trim)).toSeq)
+ .getOrElse(Seq(PARTITION_COLUMN_REF))
+ if (properties.getOrDefault(COLUMN_UPDATE_CLUSTER_BY_NESTED_TRANSFORM, "false") == "true") {
+ refs.map(ref => LogicalExpressions.apply("f", LogicalExpressions.apply("g", ref)))
+ } else {
+ refs
+ }
+ }
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ val writeSchema = info.columnUpdateSchema().orElse(info.schema())
+ val batchWrite = new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ physicalInfo: PhysicalWriteInfo): DeltaWriterFactory = {
+ newWriterFactory(info, writeSchema)
+ }
+
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ commitColumnUpdate(info, writeSchema, newData)
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ if (noMetadata && !properties.containsKey(COLUMN_UPDATE_CLUSTER_BY)) {
+ new DeltaWrite {
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ } else {
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(clusterExprs.toArray)
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ clusterExprs.map { expr =>
+ LogicalExpressions.sort(
+ expr,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()): SortOrder
+ }.toArray
+ }
+
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ }
+ }
+ }
+
+ protected def newWriterFactory(
+ info: LogicalWriteInfo,
+ writeSchema: StructType): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(writeSchema)
+ }
+
+ // For column-update writes, updated rows contain only the declared columns
+ // (columnUpdateSchema from LogicalWriteInfo). We expand each row to the full
+ // table schema by overlaying them on the base row found by pk.
+ protected def commitColumnUpdate(
+ info: LogicalWriteInfo,
+ writeSchema: StructType,
+ newData: Array[BufferedRows]): Unit = {
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+ val updateOpName = UTF8String.fromString(Update.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+
+ val mergedData = newData.map { buf =>
+ val merged = new BufferedRows(buf.key, schema)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == updateOpName) {
+ val pk = logRow.getInt(1)
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val baseRow = Some(currentRowByRowId(Seq("pk" -> pk)))
+ val fullRow = overlayNarrowRow(baseRow, narrowRow, writeFieldIdx)
+ merged.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ merged.rows.append(insertRow.copy())
+ }
+ }
+ merged
+ }
+
+ withDeletes(newData)
+ withData(mergedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+ }
+
+ class DeltaBasedColumnUpdateOperationWithReqAttrs(
+ command: Command,
+ reqCols: Array[String],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, options = options) {
+ override def requiredDataAttributes(): Array[NamedReference] = reqCols.map(FieldReference(_))
+ }
+
+ class DeltaBasedColumnUpdateSplitOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols, options) {
+ override def representUpdateAsDeleteAndInsert(): Boolean = true
+
+ override protected def newWriterFactory(
+ info: LogicalWriteInfo,
+ writeSchema: StructType): DeltaWriterFactory = {
+ val pkRowId = rowId().map(_.describe()).toSeq == Seq("pk") &&
+ info.rowIdSchema().get().fields.map(_.dataType).toSeq == Seq(IntegerType)
+ new DeltaBufferedRowsWriterFactory(writeSchema, pkRowId)
+ }
+
+ // For delete+reinsert with narrow writes, the REINSERT row has only the connector-declared
+ // columns (requiredDataAttributes order). The original row is located by the row ID, taking a
+ // data row-ID column from the narrow row and a metadata row-ID column from the preserved
+ // metadata, and the full row is rebuilt by overlaying the narrow row onto it. A DELETE logged
+ // without a `pk` row ID is resolved the same way from its metadata, so such a row ID must
+ // consist of metadata columns.
+ override protected def commitColumnUpdate(
+ info: LogicalWriteInfo,
+ writeSchema: StructType,
+ newData: Array[BufferedRows]): Unit = {
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+ val metadataSchema = info.metadataSchema().orElse(new StructType())
+ val coversTable = schema.fieldNames.forall(writeFieldIdx.contains)
+ val deleteOpName = UTF8String.fromString(Delete.toString)
+ val reinsertOpName = UTF8String.fromString(Reinsert.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+
+ def rowIdValues(narrowRow: Option[InternalRow], metadata: InternalRow): Seq[(String, Any)] = {
+ rowId().toSeq.map(_.describe()).map { name =>
+ (writeFieldIdx.get(name), narrowRow) match {
+ case (Some(i), Some(row)) => name -> row.get(i, writeSchema(i).dataType)
+ case _ =>
+ if (metadata == null || !metadataSchema.fieldNames.contains(name)) {
+ throw new IllegalStateException(s"The split fixture can only locate a row by pk " +
+ s"or by metadata row-ID columns, but row-ID column $name is in neither")
+ }
+ val i = metadataSchema.fieldIndex(name)
+ name -> metadata.get(i, metadataSchema(i).dataType)
+ }
+ }
+ }
+
+ def metadataOf(logRow: InternalRow): InternalRow =
+ if (logRow.isNullAt(2)) null else logRow.getStruct(2, metadataSchema.length)
+
+ val expandedData = newData.map { buf =>
+ val expanded = new BufferedRows(buf.key, schema)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == deleteOpName && logRow.isNullAt(1)) {
+ val values = rowIdValues(None, metadataOf(logRow))
+ buf.deletes += currentRowByRowId(values).getInt(0)
+ } else if (opName == reinsertOpName) {
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val baseRow = if (coversTable) {
+ None
+ } else {
+ Some(currentRowByRowId(rowIdValues(Some(narrowRow), metadataOf(logRow))))
+ }
+ val fullRow = overlayNarrowRow(baseRow, narrowRow, writeFieldIdx)
+ expanded.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ expanded.rows.append(insertRow.copy())
+ }
+ }
+ expanded
+ }
+
+ withDeletes(newData)
+ withData(expandedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+ }
+
+ // Test-only: a split-update fixture whose requiredDataAttributes() is the given column list.
+ class DeltaBasedColumnUpdateSplitOperationWithReqAttrs(
+ command: Command,
+ reqCols: Array[String],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateSplitOperation(command, options = options) {
+ override def requiredDataAttributes(): Array[NamedReference] = reqCols.map(FieldReference(_))
+ }
+
+ // Test-only: a split-update fixture whose row ID is the given column and whose
+ // requiredDataAttributes() lists only the updated columns, to exercise which undeclared
+ // row-ID columns still reach the REINSERT row.
+ class DeltaBasedColumnUpdateSplitOperationWithRowId(
+ command: Command,
+ rowIdCol: String,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateSplitOperation(command, updatedCols, options) {
+ override def rowId(): Array[NamedReference] = Array(FieldReference(rowIdCol))
+ override def requiredDataAttributes(): Array[NamedReference] = updatedCols.toArray
+ }
+
+ class PartitionBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ override val options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends RowLevelOperation with SupportsColumnUpdates with RowLevelOperationWithOptions {
+ var configuredScan: BatchScanBaseClass = _
+
+ override def command(): Command = command
+
+ override def requiredDataAttributes(): Array[NamedReference] = declaredColumns(updatedCols)
+
+ override def requiredMetadataAttributes(): Array[NamedReference] = {
+ if (noMetadata) {
+ Array.empty
+ } else {
+ Array(PARTITION_COLUMN_REF, INDEX_COLUMN_REF)
+ }
+ }
+
+ private def clusterColumnRef: NamedReference =
+ if (noMetadata) FieldReference("dep") else PARTITION_COLUMN_REF
+
+ override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = {
+ val onBuild: BatchScanBaseClass => Unit = { scan =>
+ configuredScan = scan
+ }
+ if (properties.getOrDefault(COLUMN_UPDATE_COW_UNREAD_FILTER_ATTRS, "false") == "true") {
Review Comment:
**47.** Nit: `PartitionBasedColumnUpdateOperation` and `PartitionBasedNarrowReplaceData` copy parts of `PartitionBasedOperation` and `PartitionBasedReplaceData` instead of sharing them: `requiredMetadataAttributes()` (L629-635 and L213-219), the default scan builder (L663 and L222-224), the commit bookkeeping (L722-726 and L265-269) and the `build()` override here (L656-660 and L780-784). This branch has already drifted: it always builds a plain `InMemoryBatchScan`, so it skips the `useCatalystRuntimeFiltering` switch of `newRowLevelScanBuilder` (L773-777), and a table with both `use-catalyst-runtime-filtering` and `column-update-cow-unread-filter-attrs` would silently plan the plain scan. No test combines them today. Could the shared parts go through helpers, or could this branch at least honor `useCatalystRuntimeFiltering`?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,623 @@
+/*
+ * 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 scala.jdk.CollectionConverters._
+
+import org.apache.spark.{SparkRuntimeException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.analysis.RewriteRowLevelCommand
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+import org.apache.spark.sql.catalyst.util.RowDeltaUtils.ORIGINAL_ROW_ID_VALUE_PREFIX
+import org.apache.spark.sql.connector.catalog.{Delete, InMemoryBaseTable, InMemoryRowLevelOperationTable, Reinsert}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.write.LogicalWriteInfo
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.errors.QueryCompilationErrors
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, updated rows contain only the declared columns
+ * (LogicalWriteInfo.columnUpdateSchema()) rather than the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase
Review Comment:
**44. No test checks what the connector receives for a narrow delta UPDATE with metadata columns, or for a narrow copy-on-write UPDATE with a subquery condition.** The tests that assert the metadata, row IDs and rows in the write log on these paths, "update handles metadata columns correctly" and "update with subquery handles metadata columns correctly", live in `DeltaBasedUpdateTableSuite.scala:35` and `:60` and in `GroupBasedUpdateTableSuite.scala:31` and `:53`, which these suites do not extend. The narrow write-log checks cover only delete and reinsert (L401, L499), the no-metadata delta writer (`DeltaBasedNoMetadataColumnUpdateTableSuite.scala:54` and `:69`) and copy-on-write without a subquery (`GroupBasedColumnUpdateTableSuite.scala:96`, `:113` and `:131`). The subquery test there (`GroupBasedColumnUpdateTableSuite.scala:156-177`) checks the scan, the write relation and the result, and L85-102 here checks only the schemas. So nothing asserts the `update(metadata, id, row)` e
ntries of `DeltaWithMetadataWritingSparkTask` (`_partition` kept, `index` nulled, the row ID and the narrow row), or the `writeColumnUpdate(metadata, record)` entries of the `Union` plan. The fixture commits do not read metadata on these paths (`InMemoryRowLevelOperationTable.scala:486-491` and `:736-745`), so a metadata projection remapped to the wrong ordinal in either plan shape, or an `index` that is no longer nulled, would pass every suite. The PR description also says that the suites cover "the rows, row IDs and metadata the connector receives". Could each path get a `checkLastColumnUpdateWriteLog(...)` assertion, e.g. by porting the two metadata tests?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -246,18 +290,55 @@ abstract class RowLevelOperationSuiteBase
assert(scans.isEmpty, "plan must not contain scan relations")
}
+ // the schemas of the scans of the target table, directly or through a row-level operation, in
+ // the plan of the query
+ private def executeAndCollectScanSchemas(query: String): Seq[StructType] = {
+ val executedPlan = executeAndKeepPlan {
+ sql(query)
+ }
+ val targetName = table.name
+ def scansTarget(scan: BatchScanExec): Boolean = scan.table match {
Review Comment:
**50.** Nit: both cases compute `t.name == targetName`, because `RowLevelOperationTable.name` already returns `table.name` (`RowLevelOperationTable.scala:47`), so this can be `scan.table.name == targetName`. Also, `executeAndCheckScanExcludes` (L334-342) has a single caller, `DeltaBasedColumnUpdateTableSuite.scala:254`, where the exact check fits: that scan should read exactly `pk INT, _partition STRING`, which the inherited "update with schema pruning" (`DeltaBasedUpdateTableSuiteBase.scala:79-90`) also expects in that suite. Could L254 use `executeAndCheckScan(..., "pk INT, _partition STRING")`, which would also catch a scan that starts reading `id` or `index`, and could the Excludes helper go?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,623 @@
+/*
+ * 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 scala.jdk.CollectionConverters._
+
+import org.apache.spark.{SparkRuntimeException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.analysis.RewriteRowLevelCommand
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+import org.apache.spark.sql.catalyst.util.RowDeltaUtils.ORIGINAL_ROW_ID_VALUE_PREFIX
+import org.apache.spark.sql.connector.catalog.{Delete, InMemoryBaseTable, InMemoryRowLevelOperationTable, Reinsert}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.write.LogicalWriteInfo
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.errors.QueryCompilationErrors
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, updated rows contain only the declared columns
+ * (LogicalWriteInfo.columnUpdateSchema()) rather than the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase
+ with ColumnUpdateTableSharedTests {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ override protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String] = {
+ java.util.Map.of("column-update-req-attrs", reqAttrs)
+ }
+
+ override protected def fullWidthTableProps: java.util.Map[String, String] = {
+ java.util.Map.of("supports-deltas", "true")
+ }
+
+ override protected def scanMetadataColumns: String = ", _partition STRING"
+
+ private val splitOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperation].getName
+
+ private val splitWithRowIdOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithRowId].getName
+
+ private def splitReqAttrsTableProps(reqAttrs: String): java.util.Map[String, String] = {
+ java.util.Map.of("column-update-req-attrs", reqAttrs, "split-updates", "true")
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData, reqAttrsTableProps(reqAttrs))
+ }
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update", "true", "split-updates", "true"))
+ }
+
+ test("column-update: column update schema contains a single assigned column") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: CHAR and VARCHAR self-assignments are reported as updated") {
+ createTable("pk INT NOT NULL, id INT, c CHAR(3), v VARCHAR(10), dep STRING")
+ sql(s"INSERT INTO $tableNameAsString VALUES (1, 1, 'a', 'x', 'hr')")
+
+ sql(s"UPDATE $tableNameAsString SET id = id, c = c, v = v WHERE pk = 1")
+
+ checkLastUpdatedColumns("c", "v")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, 1, "a ", "x", "hr") :: Nil)
+ }
+
+ test("column-update: a LogicalWriteInfo that does not implement columnUpdateSchema throws") {
+ val info = new LogicalWriteInfo {
+ override def options(): CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty()
+ override def queryId(): String = "query"
+ override def schema(): StructType = new StructType()
+ }
+
+ checkError(
+ exception = intercept[SparkUnsupportedOperationException](info.columnUpdateSchema()),
+ condition = "DATA_SOURCE_COLUMN_UPDATE_SCHEMA_NOT_IMPLEMENTED",
+ parameters = Map("class" -> info.getClass.getName))
+ }
+
+ test("column-update: updatedColumns and the column update schema follow the non-identity " +
+ "assignments") {
+ Seq(
+ ("id = -1", Seq("id"), Seq("pk", "dep", "id"), Row(1, -1, "hr")),
+ ("id = -1, dep = 'eng'", Seq("id", "dep"), Seq("pk", "dep", "id"), Row(1, -1, "eng")),
+ ("id = -1, dep = dep", Seq("id"), Seq("pk", "dep", "id"), Row(1, -1, "hr")),
+ ("id = id, dep = 'eng'", Seq("dep"), Seq("pk", "dep"), Row(1, 10, "eng")),
+ ("id = id, dep = dep", Seq.empty, Seq("pk", "dep"), Row(1, 10, "hr")),
+ ("id = pk", Seq("id"), Seq("pk", "dep", "id"), Row(1, 1, "hr"))
+ ).foreach { case (assignments, updatedColumns, columnUpdateNames, updatedRow) =>
+ withTable(tableNameAsString) {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 10, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET $assignments WHERE pk = 1")
+
+ checkLastUpdatedColumns(updatedColumns: _*)
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq == columnUpdateNames)
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), updatedRow :: Nil)
+ }
+ }
+ }
+
+ test("column-update split: row-level DELETE falls back to full-width rows") {
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin)
+ checkDeleteFallsBackToFullWidthRows(Row(2, 200, "hr"))
+ }
+
+ test("column-update split: row-level DELETE by a metadata row ID falls back to full-width " +
+ "rows") {
+ // Each partition holds one row, so `_partition` identifies it.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "_partition"))
+ checkDeleteFallsBackToFullWidthRows(Row(2, 200, "software"))
+ }
+
+ test("column-update split: MERGE falls back to full-width rows") {
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+ checkMergeFallsBackToFullWidthRows()
+ }
+
+ test("column-update: undeclared write requirement columns are rejected") {
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName
+ val splitReqAttrsOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithReqAttrs]
+ .getName
+ Seq(
+ ("pk,salary", "dep", Map.empty[String, String], reqAttrsOperationClass, "[`dep`]"),
+ ("pk,salary", "dep", Map("split-updates" -> "true"), splitReqAttrsOperationClass,
+ "[`dep`]"),
+ ("pk,salary", "dep", Map("column-update-cluster-by-nested-transform" -> "true"),
+ reqAttrsOperationClass, "[`dep`]"),
+ ("pk,salary", "id,dep", Map.empty[String, String], reqAttrsOperationClass,
+ "[`dep`, `id`]"),
+ ("pk,salary,s", "s.missing", Map.empty[String, String], reqAttrsOperationClass,
+ "[`s`.`missing` (FIELD_NOT_FOUND)]")
+ ).foreach { case (reqAttrs, clusterBy, props, operationClass, columns) =>
+ withTable(tableNameAsString) {
+ val tableProps = props ++ Map(
+ "column-update-req-attrs" -> reqAttrs, "column-update-cluster-by" -> clusterBy)
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, s STRUCT<c1: INT>, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 1, "s": { "c1": 1 }, "dep": "hr" }
+ |""".stripMargin,
+ tableProps.asJava)
+
+ Seq(Nil, Seq(excludeColumnPruning)).foreach { confs =>
+ withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> operationClass, "columns" -> columns))
+ }
+ }
+ }
+ }
+ }
+
+ test("column-update: an empty write requirement column name is reported as <empty>") {
+ checkError(
+ exception = QueryCompilationErrors.undeclaredWriteRequirementColumnsError(
+ reqAttrsOperationClass, Seq(Seq.empty[String] -> None)).asInstanceOf[AnalysisException],
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[<empty>]"))
+ }
+
+ test("column-update: distribution by an undeclared row ID column is allowed") {
+ Seq("salary = -1", "salary = -1, pk = pk").foreach { assignments =>
+ withTable(tableNameAsString) {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "salary,dep",
+ "column-update-cluster-by", "pk"))
+
+ sql(s"UPDATE $tableNameAsString SET $assignments WHERE pk = 1")
+
+ checkLastUpdatedColumns("salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 200, "hr") :: Nil)
+ }
+ }
+ }
+
+ test("column-update: an undeclared, unreferenced partition column is not read") {
+ createAndInitTableWithReqAttrs("pk,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ executeAndCheckScanExcludes(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1", "dep")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively without adopting the " +
+ "declared spelling") {
+ createAndInitTableWithReqAttrs("PK,SALARY,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "dep"))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: DELETE on a scan that returns unrequested columns is unaffected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update", "true",
+ InMemoryBaseTable.SIMULATE_PARTIAL_COLUMN_PRUNING, "true"))
+
+ sql(s"DELETE FROM $tableNameAsString WHERE pk < 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(2, 200, 20, "software") :: Nil)
+ }
+
+ test("column-update: case-sensitive analysis accepts columns that differ only in case") {
+ withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") {
+ createAndInitTableWithReqAttrs("pk,salary,SALARY,dep",
+ "pk INT NOT NULL, salary INT, SALARY INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "SALARY": 200, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "SALARY", "dep"))
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, -1, 200, "hr") :: Nil)
+ }
+ }
+
+ test("column-update: empty requiredDataAttributes throws AnalysisException") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }""".stripMargin,
+ java.util.Map.of("column-update-empty-req-attrs", "true"))
+
+ // `id = id` assigns nothing, so only the non-empty check can reject the declaration
+ Seq("id = -1", "id = id").foreach { assignment =>
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET $assignment WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ parameters = Map("connector" -> reqAttrsOperationClass))
+ }
+ }
+
+ test("column-update: invalid requiredDataAttributes are rejected") {
+ Seq(
+ ("pk,s.c1,dep", "COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE",
+ "nestedAttributes" -> "[`s`.`c1`]"),
+ ("pk,pk,salary", "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ "duplicateAttributes" -> "[`pk`]"),
+ ("salary,PK,SALARY,pk", "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ "duplicateAttributes" -> "[`salary`, `PK`]"),
+ ("pk,_partition,salary", "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE",
+ "metadataAttributes" -> "[`_partition`]"),
+ ("pk,nonexistent_col,salary", "COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE",
+ "unknownAttributes" -> "[`nonexistent_col`]"),
+ ("pk,salary", "COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS",
+ "missingColumns" -> "[`a.b`]")
+ ).foreach { case (reqAttrs, condition, columns) =>
+ withTable(tableNameAsString) {
+ createAndInitTableWithReqAttrs(reqAttrs,
+ "pk INT NOT NULL, salary INT, s STRUCT<c1: INT, c2: INT>, `a.b` INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "s": { "c1": 1, "c2": 2 }, "a.b": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1, `a.b` = -1 WHERE pk = 1")
+ },
+ condition = condition,
+ parameters = Map("connector" -> reqAttrsOperationClass, columns))
+ }
+ }
+ }
+
+ test("column-update: a declared column with no name parts is reported as <empty>") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+ val t = table
+ val operation = new t.DeltaBasedColumnUpdateOperation(UPDATE) {
+ override def requiredDataAttributes(): Array[NamedReference] =
+ Array(FieldReference(Seq.empty[String]), FieldReference("pk"),
+ FieldReference(Seq.empty[String]))
+ }
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ RewriteRowLevelCommand.resolveRequiredDataAttrs(
+ spark.table(tableNameAsString).queryExecution.analyzed, operation)
+ },
+ condition = "COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> operation.getClass.getName, "unknownAttributes" -> "[<empty>, <empty>]"))
+ }
+
+ test("column-update split: column update schema has the declared columns") {
Review Comment:
**49.** Nit: this test repeats "column-update: column update schema contains a single assigned column" (L85-102): the same DDL, the same `UPDATE ... SET id = -1 WHERE pk = 1` and the same eight-line `checkLastWriteInfo(...)` (L383-390 and L94-101), with only the table properties and one data row differing. The two configurations take different rewrite paths, so both are worth running, but the expected write info is maintained twice. Could they become one test that loops over the two property sets, like L129-151, with the split properties shared with `createAndInitTableSplit` (L80-83)?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala:
##########
@@ -56,6 +56,56 @@ object V2ExpressionUtils extends SQLConfHelper with Logging {
plan.resolve(ref.fieldNames.toImmutableArraySeq, conf.resolver)
}
+ /**
+ * Returns the names of the columns referenced by the given expressions that cannot be resolved
+ * against the plan, at any depth, including inside transform arguments.
+ */
+ private[sql] def unresolvableColumns(
+ exprs: Seq[V2Expression],
+ relation: LogicalPlan): Seq[String] = {
+ unresolvableReferences(exprs, relation).map {
+ case (Seq(), _) => "<empty>"
+ case (parts, None) => parts.quoted
+ case (parts, Some(condition)) => s"${parts.quoted} ($condition)"
+ }.distinct
Review Comment:
**54.** Nit: this `.distinct` became a no-op when `unresolvableReferences` got its own `.distinct` (L106): distinct pairs map to distinct strings, since `quoted` is injective, an entry with a condition ends with `)` while a quoted name cannot, and `<empty>` is not a quoted name. Could it go?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -321,25 +804,63 @@ class InMemoryRowLevelOperationTable private (
}
}
-private class DeltaBufferedRowsWriterFactory(schema: StructType) extends DeltaWriterFactory {
+private class DeltaBufferedRowsWriterFactory(
+ schema: StructType,
+ pkRowId: Boolean = true) extends DeltaWriterFactory {
override def createWriter(partitionId: Int, taskId: Long): DeltaWriter[InternalRow] = {
- new DeltaBufferWriter(schema)
+ new DeltaBufferWriter(schema, pkRowId)
+ }
+}
+
+// Test-only: a writer factory for a connector that never overrides writeColumnUpdate(), to
Review Comment:
**53.** Nit: these four copy-on-write writer classes (L815-847) sit between `DeltaBufferedRowsWriterFactory` (L807-813) and the `DeltaBufferWriter` it creates (L811, L849-894), which were adjacent at the merge-base, so the factory's `pkRowId` (L809) is explained only at the writer (L849-850), 40 lines and four unrelated classes away. As in item 19, AGENTS.md asks to "put the new member with the code it belongs with". Could they move below `DeltaBufferWriter`?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3198406083
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala:
##########
@@ -139,17 +138,13 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla
tableAttrs: Seq[Attribute],
scanAttrs: Seq[Attribute]): AttributeMap[Attribute] = {
- val attrMapping = tableAttrs.map { tableAttr =>
+ // The scan may be narrowed to exclude columns not needed by the connector.
+ // Attributes absent from the scan are skipped here; the caller must ensure
+ // that any attribute referenced in the condition is present in the scan.
+ val attrMapping = tableAttrs.flatMap { tableAttr =>
scanAttrs
.find(scanAttr => conf.resolver(scanAttr.name, tableAttr.name))
.map(scanAttr => tableAttr -> scanAttr)
- .getOrElse {
- throw new AnalysisException(
- errorClass = "_LEGACY_ERROR_TEMP_3075",
- messageParameters = Map(
- "tableAttr" -> tableAttr.toString,
- "scanAttrs" -> scanAttrs.mkString(",")))
- }
}
Review Comment:
I believe this is safe because condition-referenced columns are guaranteed to be in the scan. Please correct me if I'm wrong.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192418867
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -106,38 +104,92 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val remainingRowsPlan = addOperationColumn(COPY_OPERATION,
Filter(remainingRowFilter, readRelation))
- // the new state is a union of updated and copied over records
- val query = Union(updatedRowsPlan, remainingRowsPlan)
+ val updatedAndRemainingRowsPlan = Union(updatedRowsPlan, remainingRowsPlan)
- // build a plan to replace read groups in the table
val writeRelation = relation.copy(table = operationTable)
- val projections = buildReplaceDataProjections(query, relation.output, metadataAttrs)
+ val query = updatedAndRemainingRowsPlan
+ val metadataAttrs = resolveRequiredMetadataAttrs(relation, operationTable.operation)
+ val projections = buildReplaceDataProjections(query, rowAttrs, metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
ReplaceData(writeRelation, cond, query, relation, projections, groupFilterCond)
}
+ // Common read-relation setup shared by both CoW plan builders.
+ //
+ // When the connector supports column updates and declares required data attributes,
+ // the read relation is narrowed at analysis time so that
+ // GroupBasedRowLevelOperationScanPlanning uses only the needed columns for the scan.
+ // Otherwise the full relation output is used.
+ private def buildCoWReadSetup(
+ relation: DataSourceV2Relation,
+ operationTable: RowLevelOperationTable,
+ assignments: Seq[Assignment],
+ cond: Expression): (DataSourceV2Relation, Seq[Attribute]) = {
+
+ val operation = operationTable.operation
+ val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
+ val connectorDataAttrs = resolveRequiredDataAttrs(relation, operation)
+ val isNarrow = operation.supportsColumnUpdates() && connectorDataAttrs.nonEmpty
+
+ // CoW scan narrowing must be done manually at analysis time.
+ // GroupBasedRowLevelOperationScanPlanning (an optimizer rule that fires after analysis)
+ // always reads relation.output directly when building the physical scan -- it does not
+ // observe Project nodes above the relation, so optimizer-driven column pruning has no
+ // effect on CoW scans. We narrow DataSourceV2Relation.output here so that rule picks
+ // up the narrow set.
+ val readRelation = if (isNarrow) {
+ val allRequired = (connectorDataAttrs ++ computeAssignedAttrs(assignments)).distinct
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs, dataAttrs = allRequired,
+ cond = cond)
+ } else {
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs)
+ }
+
+ // CoW write schema (two paths only, no heuristic for CoW):
+ // - Narrow path (connectorDataAttrs declared): exactly connector-declared cols in declared
+ // order. The connector must declare ALL columns it wants to receive.
+ // - Full path (connectorDataAttrs empty OR supportsColumnUpdates=false): full table output.
+ // Unlike MOR, CoW does not have a heuristic assigned-only path because
+ // GroupBasedRowLevelOperationScanPlanning needs explicit column declarations to narrow.
+ val rowAttrs: Seq[Attribute] = if (isNarrow) connectorDataAttrs else relation.output
+
+ (readRelation, rowAttrs)
+ }
+
// this method assumes the assignments have been already aligned before
+ //
+ // Works for both the full-scan and narrow-scan CoW paths. In the narrow case,
+ // readRelation.output is already restricted by buildCoWReadSetup, so projecting
+ // all plan.output gives the correct narrow write schema.
Review Comment:
Use function description style.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3196568267
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java:
##########
@@ -105,4 +105,45 @@ default String description() {
default NamedReference[] requiredMetadataAttributes() {
return new NamedReference[0];
}
+
+
+ /**
+ * Controls whether to send only the required data columns to the connector rather than the
+ * full row.
+ * <p>
+ * When true, Spark narrows the data column schema ({@link LogicalWriteInfo#schema()}) to only
+ * the columns declared via {@link #requiredDataAttributes()}. Metadata columns (from
+ * {@link #requiredMetadataAttributes()}) and row ID columns (from
+ * {@link SupportsDelta#rowId()}) are unaffected and always projected separately.
+ * <p>
+ * If {@link #requiredDataAttributes()} returns a non-empty array, the write schema is exactly
+ * those columns in declared order. The connector must include all columns it wants to receive,
+ * including the columns being updated. If {@link #requiredDataAttributes()} returns an empty
+ * array, Spark sends only the non-identity assigned columns (heuristic path).
+ *
+ * @since 4.2.0
Review Comment:
Done
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -50,20 +51,34 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
protected def buildOperationTable(
table: SupportsRowLevelOperations,
command: Command,
- options: CaseInsensitiveStringMap): RowLevelOperationTable = {
- val info = RowLevelOperationInfoImpl(command, options)
+ options: CaseInsensitiveStringMap,
+ updatedColumns: Seq[NamedReference] = Nil): RowLevelOperationTable = {
+ val info = RowLevelOperationInfoImpl(command, options, updatedColumns)
val operation = table.newRowLevelOperationBuilder(info).build()
RowLevelOperationTable(table, operation)
}
+ // Builds a DataSourceV2Relation for a row-level operation, optionally narrowing its output.
+ //
+ // When dataAttrs is non-empty, the relation output is narrowed to include only columns
+ // required for a column-update write. When dataAttrs is empty, the full relation.output is
+ // preserved.
Review Comment:
Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -65,18 +72,15 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
assignments: Seq[Assignment],
cond: Expression): ReplaceData = {
- // resolve all required metadata attrs that may be used for grouping data on write
- val metadataAttrs = resolveRequiredMetadataAttrs(relation, operationTable.operation)
-
- // construct a read relation and include all required metadata columns
- val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs)
+ val (readRelation, rowAttrs) = buildCoWReadSetup(relation, operationTable, assignments, cond)
- // build a plan with updated and copied over records
- val query = buildReplaceDataUpdateProjection(readRelation, assignments, cond)
+ val updatedAndRemainingRowsPlan = buildReplaceDataUpdateProjection(
+ readRelation, assignments, cond)
- // build a plan to replace read groups in the table
val writeRelation = relation.copy(table = operationTable)
- val projections = buildReplaceDataProjections(query, relation.output, metadataAttrs)
+ val query = updatedAndRemainingRowsPlan
Review Comment:
Done, used a single variable
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java:
##########
@@ -105,4 +105,45 @@ default String description() {
default NamedReference[] requiredMetadataAttributes() {
return new NamedReference[0];
}
+
+
Review Comment:
Done
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java:
##########
@@ -105,4 +105,45 @@ default String description() {
default NamedReference[] requiredMetadataAttributes() {
return new NamedReference[0];
}
+
+
+ /**
+ * Controls whether to send only the required data columns to the connector rather than the
+ * full row.
+ * <p>
+ * When true, Spark narrows the data column schema ({@link LogicalWriteInfo#schema()}) to only
+ * the columns declared via {@link #requiredDataAttributes()}. Metadata columns (from
+ * {@link #requiredMetadataAttributes()}) and row ID columns (from
+ * {@link SupportsDelta#rowId()}) are unaffected and always projected separately.
+ * <p>
+ * If {@link #requiredDataAttributes()} returns a non-empty array, the write schema is exactly
+ * those columns in declared order. The connector must include all columns it wants to receive,
+ * including the columns being updated. If {@link #requiredDataAttributes()} returns an empty
+ * array, Spark sends only the non-identity assigned columns (heuristic path).
+ *
+ * @since 4.2.0
+ */
+ default boolean supportsColumnUpdates() {
+ return false;
+ }
+
+ /**
+ * Returns data column references required to perform this row-level operation.
+ * <p>
+ * This method is only consulted by Spark when {@link #supportsColumnUpdates()} returns
+ * {@code true}. If {@code supportsColumnUpdates()} returns {@code false}, the returned array
+ * is ignored and the full table row is sent (the default behavior).
+ * <p>
+ * When non-empty, the returned columns become the write schema in declared order.
+ * The connector must declare all columns it wants to receive, including the columns being
+ * updated. Use {@link RowLevelOperationInfo#updatedColumns()} to learn which columns are being
+ * assigned, then add any extra columns needed for row lookup or routing (e.g., primary key).
+ * <p>
+ * When empty (the default), Spark falls back to sending only the non-identity assigned columns.
+ *
+ * @since 4.2.0
+ */
+ default NamedReference[] requiredDataAttributes() {
Review Comment:
Even though the scope of this PR is UPDATE only, we'd like this API to work for MERGE as well (DELETE doesn't benefit since it doesn't write data columns). I'm still assessing what it takes and will add a section in the SPIP on how it could be implemented.
Happy to add a "currently only consulted for UPDATE" note in the Javadoc for now and remove it when MERGE support lands.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -41,7 +42,13 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
EliminateSubqueryAliases(aliasedTable) match {
case r @ ExtractV2Table(tbl: SupportsRowLevelOperations) =>
- val table = buildOperationTable(tbl, UPDATE, CaseInsensitiveStringMap.empty())
+ val updatedCols = assignments.collect {
+ case Assignment(key: AttributeReference, value)
+ if !isIdentityAssignment(key, value) =>
Review Comment:
Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -106,38 +104,92 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val remainingRowsPlan = addOperationColumn(COPY_OPERATION,
Filter(remainingRowFilter, readRelation))
- // the new state is a union of updated and copied over records
- val query = Union(updatedRowsPlan, remainingRowsPlan)
+ val updatedAndRemainingRowsPlan = Union(updatedRowsPlan, remainingRowsPlan)
- // build a plan to replace read groups in the table
val writeRelation = relation.copy(table = operationTable)
- val projections = buildReplaceDataProjections(query, relation.output, metadataAttrs)
+ val query = updatedAndRemainingRowsPlan
+ val metadataAttrs = resolveRequiredMetadataAttrs(relation, operationTable.operation)
+ val projections = buildReplaceDataProjections(query, rowAttrs, metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
ReplaceData(writeRelation, cond, query, relation, projections, groupFilterCond)
}
+ // Common read-relation setup shared by both CoW plan builders.
+ //
+ // When the connector supports column updates and declares required data attributes,
+ // the read relation is narrowed at analysis time so that
+ // GroupBasedRowLevelOperationScanPlanning uses only the needed columns for the scan.
+ // Otherwise the full relation output is used.
+ private def buildCoWReadSetup(
+ relation: DataSourceV2Relation,
+ operationTable: RowLevelOperationTable,
+ assignments: Seq[Assignment],
+ cond: Expression): (DataSourceV2Relation, Seq[Attribute]) = {
+
+ val operation = operationTable.operation
+ val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
+ val connectorDataAttrs = resolveRequiredDataAttrs(relation, operation)
+ val isNarrow = operation.supportsColumnUpdates() && connectorDataAttrs.nonEmpty
+
+ // CoW scan narrowing must be done manually at analysis time.
+ // GroupBasedRowLevelOperationScanPlanning (an optimizer rule that fires after analysis)
+ // always reads relation.output directly when building the physical scan -- it does not
+ // observe Project nodes above the relation, so optimizer-driven column pruning has no
+ // effect on CoW scans. We narrow DataSourceV2Relation.output here so that rule picks
+ // up the narrow set.
+ val readRelation = if (isNarrow) {
+ val allRequired = (connectorDataAttrs ++ computeAssignedAttrs(assignments)).distinct
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs, dataAttrs = allRequired,
+ cond = cond)
+ } else {
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs)
+ }
+
+ // CoW write schema (two paths only, no heuristic for CoW):
+ // - Narrow path (connectorDataAttrs declared): exactly connector-declared cols in declared
+ // order. The connector must declare ALL columns it wants to receive.
+ // - Full path (connectorDataAttrs empty OR supportsColumnUpdates=false): full table output.
+ // Unlike MOR, CoW does not have a heuristic assigned-only path because
+ // GroupBasedRowLevelOperationScanPlanning needs explicit column declarations to narrow.
+ val rowAttrs: Seq[Attribute] = if (isNarrow) connectorDataAttrs else relation.output
+
+ (readRelation, rowAttrs)
Review Comment:
I changed this to return metadataAttrs too.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -106,38 +104,92 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val remainingRowsPlan = addOperationColumn(COPY_OPERATION,
Filter(remainingRowFilter, readRelation))
- // the new state is a union of updated and copied over records
- val query = Union(updatedRowsPlan, remainingRowsPlan)
+ val updatedAndRemainingRowsPlan = Union(updatedRowsPlan, remainingRowsPlan)
- // build a plan to replace read groups in the table
val writeRelation = relation.copy(table = operationTable)
- val projections = buildReplaceDataProjections(query, relation.output, metadataAttrs)
+ val query = updatedAndRemainingRowsPlan
+ val metadataAttrs = resolveRequiredMetadataAttrs(relation, operationTable.operation)
+ val projections = buildReplaceDataProjections(query, rowAttrs, metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
ReplaceData(writeRelation, cond, query, relation, projections, groupFilterCond)
}
+ // Common read-relation setup shared by both CoW plan builders.
+ //
+ // When the connector supports column updates and declares required data attributes,
+ // the read relation is narrowed at analysis time so that
+ // GroupBasedRowLevelOperationScanPlanning uses only the needed columns for the scan.
+ // Otherwise the full relation output is used.
Review Comment:
Done
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -106,38 +104,92 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val remainingRowsPlan = addOperationColumn(COPY_OPERATION,
Filter(remainingRowFilter, readRelation))
- // the new state is a union of updated and copied over records
- val query = Union(updatedRowsPlan, remainingRowsPlan)
+ val updatedAndRemainingRowsPlan = Union(updatedRowsPlan, remainingRowsPlan)
- // build a plan to replace read groups in the table
val writeRelation = relation.copy(table = operationTable)
- val projections = buildReplaceDataProjections(query, relation.output, metadataAttrs)
+ val query = updatedAndRemainingRowsPlan
+ val metadataAttrs = resolveRequiredMetadataAttrs(relation, operationTable.operation)
+ val projections = buildReplaceDataProjections(query, rowAttrs, metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
ReplaceData(writeRelation, cond, query, relation, projections, groupFilterCond)
}
+ // Common read-relation setup shared by both CoW plan builders.
+ //
+ // When the connector supports column updates and declares required data attributes,
+ // the read relation is narrowed at analysis time so that
+ // GroupBasedRowLevelOperationScanPlanning uses only the needed columns for the scan.
+ // Otherwise the full relation output is used.
+ private def buildCoWReadSetup(
+ relation: DataSourceV2Relation,
+ operationTable: RowLevelOperationTable,
+ assignments: Seq[Assignment],
+ cond: Expression): (DataSourceV2Relation, Seq[Attribute]) = {
+
+ val operation = operationTable.operation
+ val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
+ val connectorDataAttrs = resolveRequiredDataAttrs(relation, operation)
+ val isNarrow = operation.supportsColumnUpdates() && connectorDataAttrs.nonEmpty
+
+ // CoW scan narrowing must be done manually at analysis time.
+ // GroupBasedRowLevelOperationScanPlanning (an optimizer rule that fires after analysis)
+ // always reads relation.output directly when building the physical scan -- it does not
+ // observe Project nodes above the relation, so optimizer-driven column pruning has no
+ // effect on CoW scans. We narrow DataSourceV2Relation.output here so that rule picks
+ // up the narrow set.
+ val readRelation = if (isNarrow) {
+ val allRequired = (connectorDataAttrs ++ computeAssignedAttrs(assignments)).distinct
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs, dataAttrs = allRequired,
+ cond = cond)
+ } else {
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs)
+ }
+
+ // CoW write schema (two paths only, no heuristic for CoW):
+ // - Narrow path (connectorDataAttrs declared): exactly connector-declared cols in declared
+ // order. The connector must declare ALL columns it wants to receive.
+ // - Full path (connectorDataAttrs empty OR supportsColumnUpdates=false): full table output.
+ // Unlike MOR, CoW does not have a heuristic assigned-only path because
+ // GroupBasedRowLevelOperationScanPlanning needs explicit column declarations to narrow.
+ val rowAttrs: Seq[Attribute] = if (isNarrow) connectorDataAttrs else relation.output
+
+ (readRelation, rowAttrs)
+ }
+
// this method assumes the assignments have been already aligned before
+ //
+ // Works for both the full-scan and narrow-scan CoW paths. In the narrow case,
+ // readRelation.output is already restricted by buildCoWReadSetup, so projecting
+ // all plan.output gives the correct narrow write schema.
Review Comment:
Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -154,30 +206,145 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
cond: Expression): WriteDelta = {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
+ // Column-update support applies to the standard delta path and the delete+reinsert path.
+ // When representUpdateAsDeleteAndInsert is true, the REINSERT leg of the Expand already
+ // uses only assigned values, so the narrow effectiveRowAttrs applies correctly.
+ val supportsColumnUpdate = operation.supportsColumnUpdates()
// resolve all needed attrs (e.g. row ID and any required metadata attrs)
- val rowAttrs = relation.output
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
+ // Connector-declared data attrs used to determine pass-through columns in the write plan.
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveRequiredDataAttrs(relation, operation)
+ } else Nil
+
+ // MOR uses a full-schema scan; ColumnPruning narrows it via Project references.
val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ // Connector-required attrs that are NOT being assigned are added as pass-throughs in the
+ // plan so that ColumnPruning keeps them in the physical scan AND the connector receives
+ // their current values via DeltaWriter.update's row argument.
+ val assignedAttrs = if (supportsColumnUpdate) computeAssignedAttrs(assignments)
+ else relation.output
+ val connectorExtraAttrs: Seq[AttributeReference] = if (connectorDataAttrs.nonEmpty) {
+ val assignedAttrSet = AttributeSet(assignedAttrs)
+ connectorDataAttrs.filterNot(assignedAttrSet.contains)
+ } else Nil
+
// build a plan for updated records that match the condition
val matchedRowsPlan = Filter(cond, readRelation)
val rowDeltaPlan = if (operation.representUpdateAsDeleteAndInsert) {
buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else if (supportsColumnUpdate) {
+ buildColumnUpdateProjection(
+ matchedRowsPlan, assignments, rowIdAttrs, metadataAttrs, connectorExtraAttrs)
} else {
buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs)
}
+ // Effective row write schema:
+ // - Narrow path (connectorDataAttrs declared): exactly connector-declared cols in declared
+ // order. The connector must declare ALL columns it wants to receive (including updated
+ // ones). This mirrors the metadata pattern and enables strict areCompatible validation.
+ // - Heuristic path (connectorDataAttrs empty): only the assigned (changed) columns.
+ // - Full path (no column-update support): full table output.
+ val effectiveRowAttrs = if (supportsColumnUpdate && connectorDataAttrs.nonEmpty) {
+ connectorDataAttrs
+ } else if (supportsColumnUpdate) {
+ assignedAttrs
+ } else {
+ relation.output
+ }
+
// 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(
+ rowDeltaPlan, effectiveRowAttrs, rowIdAttrs, metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond)
}
+ // Builds the row delta projection for the column update path.
+ //
+ // The resulting Project references only:
+ // - assigned column values (new values being written)
+ // - connector pass-through values (connector declared but not assigned)
+ // - metadata columns (nulled or preserved)
+ // - row ID columns (for delta identification)
+ // - original row ID values (only when a row ID column is being reassigned)
+ //
+ // ColumnPruning observes exactly these references and narrows the physical scan accordingly.
+ // Connectors that need additional columns in the scan (e.g., partition columns for
+ // distribution) should declare them in requiredDataAttributes().
+ //
+ // Note: AlignUpdateAssignments guarantees all assignment keys are top-level
Review Comment:
I added a new test `test("column-update: nested struct field update narrows to the root struct column")` that updates an inner field in a struct, the `AlignUpdateAssignment` returns only the root key.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java:
##########
@@ -37,4 +38,20 @@ public interface RowLevelOperationInfo {
* Returns the row-level SQL command (e.g. DELETE, UPDATE, MERGE).
*/
Command command();
+
+ /**
+ * Returns the columns being updated in an UPDATE statement, as non-identity assignments.
+ *
+ * <p>For DELETE and MERGE, returns an empty array.
+ *
+ * <p>Connectors can use this to decide what {@link RowLevelOperation#requiredDataAttributes()}
+ * to declare. For instance, a connector that needs its primary key for row lookup can check
+ * whether pk is already in the updated columns list and, if not, add it to
+ * requiredDataAttributes().
+ *
+ * @since 4.2.0
Review Comment:
Done
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala:
##########
@@ -470,7 +506,28 @@ case class WriteDelta(
case Some(projection) => DataTypeUtils.toAttributes(projection.schema)
case None => Nil
}
- table.skipSchemaResolution || areCompatible(inRowAttrs, outRowAttrs)
+ table.skipSchemaResolution ||
+ areCompatible(inRowAttrs, outRowAttrs) ||
+ dataAttrsResolved(inRowAttrs)
Review Comment:
Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -154,30 +206,145 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
cond: Expression): WriteDelta = {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
+ // Column-update support applies to the standard delta path and the delete+reinsert path.
+ // When representUpdateAsDeleteAndInsert is true, the REINSERT leg of the Expand already
+ // uses only assigned values, so the narrow effectiveRowAttrs applies correctly.
+ val supportsColumnUpdate = operation.supportsColumnUpdates()
// resolve all needed attrs (e.g. row ID and any required metadata attrs)
- val rowAttrs = relation.output
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
+ // Connector-declared data attrs used to determine pass-through columns in the write plan.
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveRequiredDataAttrs(relation, operation)
+ } else Nil
+
+ // MOR uses a full-schema scan; ColumnPruning narrows it via Project references.
val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ // Connector-required attrs that are NOT being assigned are added as pass-throughs in the
+ // plan so that ColumnPruning keeps them in the physical scan AND the connector receives
+ // their current values via DeltaWriter.update's row argument.
+ val assignedAttrs = if (supportsColumnUpdate) computeAssignedAttrs(assignments)
+ else relation.output
+ val connectorExtraAttrs: Seq[AttributeReference] = if (connectorDataAttrs.nonEmpty) {
+ val assignedAttrSet = AttributeSet(assignedAttrs)
+ connectorDataAttrs.filterNot(assignedAttrSet.contains)
+ } else Nil
+
// build a plan for updated records that match the condition
val matchedRowsPlan = Filter(cond, readRelation)
val rowDeltaPlan = if (operation.representUpdateAsDeleteAndInsert) {
buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else if (supportsColumnUpdate) {
+ buildColumnUpdateProjection(
+ matchedRowsPlan, assignments, rowIdAttrs, metadataAttrs, connectorExtraAttrs)
} else {
buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs)
}
+ // Effective row write schema:
+ // - Narrow path (connectorDataAttrs declared): exactly connector-declared cols in declared
+ // order. The connector must declare ALL columns it wants to receive (including updated
+ // ones). This mirrors the metadata pattern and enables strict areCompatible validation.
+ // - Heuristic path (connectorDataAttrs empty): only the assigned (changed) columns.
+ // - Full path (no column-update support): full table output.
+ val effectiveRowAttrs = if (supportsColumnUpdate && connectorDataAttrs.nonEmpty) {
+ connectorDataAttrs
+ } else if (supportsColumnUpdate) {
+ assignedAttrs
+ } else {
+ relation.output
+ }
+
// 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(
+ rowDeltaPlan, effectiveRowAttrs, rowIdAttrs, metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond)
}
+ // Builds the row delta projection for the column update path.
Review Comment:
Done.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala:
##########
@@ -378,7 +386,35 @@ case class ReplaceData(
// validates row projection output is compatible with table attributes
private def rowAttrsResolved: Boolean = {
val inRowAttrs = DataTypeUtils.toAttributes(projections.rowProjection.schema)
- table.skipSchemaResolution || areCompatible(inRowAttrs, table.output)
+ table.skipSchemaResolution ||
+ areCompatible(inRowAttrs, table.output) ||
+ dataAttrsResolved(inRowAttrs)
+ }
+
+ // Validates the narrow-write-schema row projection output.
Review Comment:
Done.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java:
##########
@@ -105,4 +105,45 @@ default String description() {
default NamedReference[] requiredMetadataAttributes() {
return new NamedReference[0];
}
+
+
+ /**
+ * Controls whether to send only the required data columns to the connector rather than the
+ * full row.
+ * <p>
+ * When true, Spark narrows the data column schema ({@link LogicalWriteInfo#schema()}) to only
+ * the columns declared via {@link #requiredDataAttributes()}. Metadata columns (from
+ * {@link #requiredMetadataAttributes()}) and row ID columns (from
+ * {@link SupportsDelta#rowId()}) are unaffected and always projected separately.
+ * <p>
+ * If {@link #requiredDataAttributes()} returns a non-empty array, the write schema is exactly
+ * those columns in declared order. The connector must include all columns it wants to receive,
+ * including the columns being updated. If {@link #requiredDataAttributes()} returns an empty
+ * array, Spark sends only the non-identity assigned columns (heuristic path).
+ *
+ * @since 4.2.0
+ */
+ default boolean supportsColumnUpdates() {
+ return false;
+ }
+
+ /**
+ * Returns data column references required to perform this row-level operation.
+ * <p>
+ * This method is only consulted by Spark when {@link #supportsColumnUpdates()} returns
+ * {@code true}. If {@code supportsColumnUpdates()} returns {@code false}, the returned array
+ * is ignored and the full table row is sent (the default behavior).
+ * <p>
+ * When non-empty, the returned columns become the write schema in declared order.
+ * The connector must declare all columns it wants to receive, including the columns being
Review Comment:
Each column the connector returns passes through `V2ExpressionUtils.resolveRefs` which throws `AnalysisException` if the column is non existent.
I added a test `test("column-update: requiredDataAttributes throws AnalysisException for invalid column")`
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -154,30 +206,145 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
cond: Expression): WriteDelta = {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
+ // Column-update support applies to the standard delta path and the delete+reinsert path.
+ // When representUpdateAsDeleteAndInsert is true, the REINSERT leg of the Expand already
+ // uses only assigned values, so the narrow effectiveRowAttrs applies correctly.
+ val supportsColumnUpdate = operation.supportsColumnUpdates()
// resolve all needed attrs (e.g. row ID and any required metadata attrs)
- val rowAttrs = relation.output
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
+ // Connector-declared data attrs used to determine pass-through columns in the write plan.
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveRequiredDataAttrs(relation, operation)
+ } else Nil
+
+ // MOR uses a full-schema scan; ColumnPruning narrows it via Project references.
val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ // Connector-required attrs that are NOT being assigned are added as pass-throughs in the
+ // plan so that ColumnPruning keeps them in the physical scan AND the connector receives
+ // their current values via DeltaWriter.update's row argument.
+ val assignedAttrs = if (supportsColumnUpdate) computeAssignedAttrs(assignments)
+ else relation.output
+ val connectorExtraAttrs: Seq[AttributeReference] = if (connectorDataAttrs.nonEmpty) {
+ val assignedAttrSet = AttributeSet(assignedAttrs)
+ connectorDataAttrs.filterNot(assignedAttrSet.contains)
+ } else Nil
+
// build a plan for updated records that match the condition
val matchedRowsPlan = Filter(cond, readRelation)
val rowDeltaPlan = if (operation.representUpdateAsDeleteAndInsert) {
buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else if (supportsColumnUpdate) {
+ buildColumnUpdateProjection(
+ matchedRowsPlan, assignments, rowIdAttrs, metadataAttrs, connectorExtraAttrs)
} else {
buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs)
}
+ // Effective row write schema:
+ // - Narrow path (connectorDataAttrs declared): exactly connector-declared cols in declared
+ // order. The connector must declare ALL columns it wants to receive (including updated
+ // ones). This mirrors the metadata pattern and enables strict areCompatible validation.
+ // - Heuristic path (connectorDataAttrs empty): only the assigned (changed) columns.
+ // - Full path (no column-update support): full table output.
+ val effectiveRowAttrs = if (supportsColumnUpdate && connectorDataAttrs.nonEmpty) {
+ connectorDataAttrs
+ } else if (supportsColumnUpdate) {
+ assignedAttrs
+ } else {
+ relation.output
+ }
+
// 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(
+ rowDeltaPlan, effectiveRowAttrs, rowIdAttrs, metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond)
}
+ // Builds the row delta projection for the column update path.
+ //
+ // The resulting Project references only:
+ // - assigned column values (new values being written)
+ // - connector pass-through values (connector declared but not assigned)
+ // - metadata columns (nulled or preserved)
+ // - row ID columns (for delta identification)
+ // - original row ID values (only when a row ID column is being reassigned)
+ //
+ // ColumnPruning observes exactly these references and narrows the physical scan accordingly.
+ // Connectors that need additional columns in the scan (e.g., partition columns for
+ // distribution) should declare them in requiredDataAttributes().
Review Comment:
Each column the connector returns passes through V2ExpressionUtils.resolveRefs which throws AnalysisException if the column is non existent.
I added a test test("column-update: requiredDataAttributes throws AnalysisException for invalid 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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3423871912
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -50,20 +51,36 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
protected def buildOperationTable(
table: SupportsRowLevelOperations,
command: Command,
- options: CaseInsensitiveStringMap): RowLevelOperationTable = {
- val info = RowLevelOperationInfoImpl(command, options)
+ options: CaseInsensitiveStringMap,
+ updatedColumns: Seq[NamedReference] = Nil): RowLevelOperationTable = {
+ val info = RowLevelOperationInfoImpl(command, options, updatedColumns)
val operation = table.newRowLevelOperationBuilder(info).build()
RowLevelOperationTable(table, operation)
}
+ /**
+ * Builds a DataSourceV2Relation for a row-level operation, optionally narrowing its output.
+ *
+ * When dataAttrs is non-empty, the relation output is narrowed to include only columns
+ * required for a column-update write. When dataAttrs is empty, the full relation.output is
+ * preserved.
+ */
protected def buildRelationWithAttrs(
Review Comment:
Not a big fan of this change, it makes the method fragile. I would consider adding another one instead.
Something like `buildNarrowRelationWithAttrs()` or something.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3423698058
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java:
##########
@@ -105,4 +105,48 @@ default String description() {
default NamedReference[] requiredMetadataAttributes() {
return new NamedReference[0];
}
+
+ /**
+ * Controls whether to send only the required data columns to the connector rather than the
+ * full row.
+ * <p>
+ * When true, Spark narrows the data column schema ({@link LogicalWriteInfo#schema()}) to only
+ * the columns declared via {@link #requiredDataAttributes()}. Metadata columns (from
+ * {@link #requiredMetadataAttributes()}) and row ID columns (from
+ * {@link SupportsDelta#rowId()}) are unaffected and always projected separately.
+ * <p>
+ * If {@link #requiredDataAttributes()} returns a non-empty array, the write schema is exactly
+ * those columns in declared order. The connector must include all columns it wants to receive,
+ * including the columns being updated. If {@link #requiredDataAttributes()} returns an empty
+ * array, Spark sends only the non-identity assigned columns (heuristic path).
+ * <p>
+ * Currently only consulted for UPDATE operations.
+ *
+ * @since 4.3.0
+ */
+ default boolean supportsColumnUpdates() {
+ return false;
+ }
+
+ /**
+ * Returns data column references required to perform this row-level operation.
+ * <p>
+ * This method is only consulted by Spark when {@link #supportsColumnUpdates()} returns
+ * {@code true}. If {@code supportsColumnUpdates()} returns {@code false}, the returned array
+ * is ignored and the full table row is sent (the default behavior).
+ * <p>
+ * When non-empty, the returned columns become the write schema in declared order.
+ * The connector must declare all columns it wants to receive, including the columns being
+ * updated. Use {@link RowLevelOperationInfo#updatedColumns()} to learn which columns are being
+ * assigned, then add any extra columns needed for row lookup or routing (e.g., primary key).
+ * <p>
+ * When empty (the default), Spark falls back to sending only the non-identity assigned columns.
+ * <p>
+ * Currently only consulted for UPDATE operations.
+ *
+ * @since 4.3.0
+ */
+ default NamedReference[] requiredDataAttributes() {
+ return new NamedReference[0];
Review Comment:
As I commented in the SPIP doc, this is not a good default.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3575235190
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -50,20 +51,36 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
protected def buildOperationTable(
table: SupportsRowLevelOperations,
command: Command,
- options: CaseInsensitiveStringMap): RowLevelOperationTable = {
- val info = RowLevelOperationInfoImpl(command, options)
+ options: CaseInsensitiveStringMap,
+ updatedColumns: Seq[NamedReference] = Nil): RowLevelOperationTable = {
+ val info = RowLevelOperationInfoImpl(command, options, updatedColumns)
val operation = table.newRowLevelOperationBuilder(info).build()
RowLevelOperationTable(table, operation)
}
+ /**
+ * Builds a DataSourceV2Relation for a row-level operation, optionally narrowing its output.
+ *
+ * When dataAttrs is non-empty, the relation output is narrowed to include only columns
+ * required for a column-update write. When dataAttrs is empty, the full relation.output is
+ * preserved.
+ */
protected def buildRelationWithAttrs(
Review Comment:
I added a `buildNarrowRelationWithAttrs()` version and kept the original unchanged.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3785118090
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,47 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ *
+ * @since 4.3.0
Review Comment:
I'd be open to target 4.3 if the rest of the community is OK with it.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3786035399
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -41,7 +44,11 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
EliminateSubqueryAliases(aliasedTable) match {
case r @ ExtractV2Table(tbl: SupportsRowLevelOperations) =>
checkNoGeneratedColumns(r, UPDATE)
- val table = buildOperationTable(tbl, UPDATE, r.options)
+ val updatedCols = assignments.collect {
Review Comment:
Can we add a helper like below?
```
private def collectUpdatedAttrs(assignments: Seq[Assignment]): Seq[AttributeReference] = {
assignments.collect {
case Assignment(key: AttributeReference, value) if !isIdentityAssignment(key, value) => key
}
}
```
And then pass AttributeReference to buildOperationTable?
```
protected def buildOperationTable(
table: SupportsRowLevelOperations,
command: Command,
options: CaseInsensitiveStringMap,
updatedAttrs: Seq[AttributeReference] = Nil): RowLevelOperationTable = {
val updatedColumns = updatedAttrs.map(attr => FieldReference(Seq(attr.name)))
...
}
```
So that RewriteUpdateTable becomes a bit easier. I also see you may use `collectUpdatedAttrs` in a few other places.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-5844401384
Gentle ping, @anuragmantri .
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4315467526
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala:
##########
@@ -109,22 +111,75 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper {
WriteToDataSourceV2(Some(r), microBatchWrite, newQuery, customMetrics)
case rd @ ReplaceData(r: DataSourceV2Relation, _, query, _, projections, _, None) =>
- val rowSchema = projections.rowProjection.schema
+ val (rowSchema, columnUpdateSchema) =
+ writeSchemas(projections.rowProjection.schema, hasNarrowRows(rd))
val metadataSchema = projections.metadataProjection.map(_.schema)
val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap)
- val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema)
+ val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema,
+ columnUpdateSchema)
val write = writeBuilder.build()
+ validateColumnUpdateWrite(rd, write)
val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, r.funCatalog)
rd.copy(write = Some(write), query = newQuery)
case wd @ WriteDelta(r: DataSourceV2Relation, _, query, _, projections, _, None) =>
val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap)
- val deltaWriteBuilder = newDeltaWriteBuilder(r.table, writeOptions, projections)
+ val deltaWriteBuilder =
+ newDeltaWriteBuilder(r.table, writeOptions, projections, hasNarrowRows(wd))
val deltaWrite = deltaWriteBuilder.build()
+ validateColumnUpdateWrite(wd, deltaWrite)
val newQuery = DistributionAndOrderingUtils.prepareQuery(deltaWrite, query, r.funCatalog)
wd.copy(write = Some(deltaWrite), query = newQuery)
}
+ /**
+ * Whether a row-level write delivers narrow rows for updated, copied, and reinserted records.
+ * This must match the writes whose write relation RewriteUpdateTable narrows, even when the
+ * declared columns cover the whole table.
+ */
+ private[v2] def hasNarrowRows(write: RowLevelWrite): Boolean = {
+ write.operation.isInstanceOf[SupportsColumnUpdates] && write.operation.command == UPDATE
Review Comment:
**6. Whether a write delivers narrow rows is decided twice and recorded nowhere in the plan.** `RewriteUpdateTable.resolveWriteAttrs` narrows the write relation for any `SupportsColumnUpdates` operation, while `hasNarrowRows` re-derives the decision later from whatever operation sits in the write relation at that point, plus `operation.command == UPDATE`, at five call sites (L115, L128, L151, `GroupBasedRowLevelOperationScanPlanning.scala:70` and `DataSourceV2Strategy.scala:594`). Nothing re-validates the narrowed write relation against the current operation either: `metadataAttrsResolved` re-derives its expected columns from `operation.requiredMetadataAttributes`, but `rowAttrsResolved` compares the narrow projection with the narrow `table.output`.
If the two decisions disagree, for example after a post-hoc rule swaps in a rebuilt operation (Iceberg's `ResolveBranch` does this after `RewriteUpdateTable` and keeps the narrow output and projections, see item 9), this rule reports the narrow row schema as `schema()` with no `columnUpdateSchema()`, and `ReplaceDataExec` calls `write()`. A copy-on-write connector that follows the documented contract then rewrites groups with only the declared columns, and nothing fails at analysis or execution.
Could you fail fast here, e.g. with an internal error when `!hasNarrowRows(w) && w.table.output != w.originalTable.output`? Every write Spark builds without narrowing keeps `table.output == originalTable.output`, so existing plans are not affected.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub) via reviews" <re...@spark.apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4428907236
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/LogicalWriteInfo.java:
##########
@@ -65,4 +69,24 @@ default Optional<StructType> metadataSchema() {
throw new SparkUnsupportedOperationException(
"DATA_SOURCE_METADATA_SCHEMA_NOT_IMPLEMENTED", Map.of("class", getClass().getName()));
}
+
+ /**
+ * the schema of updated, copied, and reinserted rows from Spark to data source in a
Review Comment:
This makes sense too.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub) via reviews" <re...@spark.apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4429353236
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/RowLevelOperationInfoImpl.scala:
##########
@@ -17,9 +17,15 @@
package org.apache.spark.sql.connector.write
+import org.apache.spark.sql.connector.expressions.NamedReference
import org.apache.spark.sql.connector.write.RowLevelOperation.Command
import org.apache.spark.sql.util.CaseInsensitiveStringMap
private[sql] case class RowLevelOperationInfoImpl(
command: Command,
- options: CaseInsensitiveStringMap) extends RowLevelOperationInfo
+ options: CaseInsensitiveStringMap,
+ private val updatedCols: Seq[NamedReference])
Review Comment:
It is OK.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub) via reviews" <re...@spark.apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4429371574
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DistributionAndOrderingUtils.scala:
##########
@@ -38,15 +39,8 @@ object DistributionAndOrderingUtils {
val numPartitions = write.requiredNumPartitions()
val partitionSize = write.advisoryPartitionSizeInBytes()
- val distribution = write.requiredDistribution match {
Review Comment:
Is this change needed?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-6499096229
Please resolve the conflicts, @anuragmantri .
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192364870
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -154,30 +206,145 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
cond: Expression): WriteDelta = {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
+ // Column-update support applies to the standard delta path and the delete+reinsert path.
+ // When representUpdateAsDeleteAndInsert is true, the REINSERT leg of the Expand already
+ // uses only assigned values, so the narrow effectiveRowAttrs applies correctly.
+ val supportsColumnUpdate = operation.supportsColumnUpdates()
// resolve all needed attrs (e.g. row ID and any required metadata attrs)
- val rowAttrs = relation.output
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
+ // Connector-declared data attrs used to determine pass-through columns in the write plan.
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveRequiredDataAttrs(relation, operation)
+ } else Nil
+
+ // MOR uses a full-schema scan; ColumnPruning narrows it via Project references.
val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ // Connector-required attrs that are NOT being assigned are added as pass-throughs in the
+ // plan so that ColumnPruning keeps them in the physical scan AND the connector receives
+ // their current values via DeltaWriter.update's row argument.
+ val assignedAttrs = if (supportsColumnUpdate) computeAssignedAttrs(assignments)
+ else relation.output
+ val connectorExtraAttrs: Seq[AttributeReference] = if (connectorDataAttrs.nonEmpty) {
+ val assignedAttrSet = AttributeSet(assignedAttrs)
+ connectorDataAttrs.filterNot(assignedAttrSet.contains)
+ } else Nil
+
// build a plan for updated records that match the condition
val matchedRowsPlan = Filter(cond, readRelation)
val rowDeltaPlan = if (operation.representUpdateAsDeleteAndInsert) {
buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else if (supportsColumnUpdate) {
+ buildColumnUpdateProjection(
+ matchedRowsPlan, assignments, rowIdAttrs, metadataAttrs, connectorExtraAttrs)
} else {
buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs)
}
+ // Effective row write schema:
+ // - Narrow path (connectorDataAttrs declared): exactly connector-declared cols in declared
+ // order. The connector must declare ALL columns it wants to receive (including updated
+ // ones). This mirrors the metadata pattern and enables strict areCompatible validation.
+ // - Heuristic path (connectorDataAttrs empty): only the assigned (changed) columns.
+ // - Full path (no column-update support): full table output.
+ val effectiveRowAttrs = if (supportsColumnUpdate && connectorDataAttrs.nonEmpty) {
+ connectorDataAttrs
+ } else if (supportsColumnUpdate) {
+ assignedAttrs
+ } else {
+ relation.output
+ }
+
// 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(
+ rowDeltaPlan, effectiveRowAttrs, rowIdAttrs, metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond)
}
+ // Builds the row delta projection for the column update path.
Review Comment:
For function description, please follow the community style like the other code path.
```
/**
* ...
*/
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192364149
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -50,20 +51,34 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
protected def buildOperationTable(
table: SupportsRowLevelOperations,
command: Command,
- options: CaseInsensitiveStringMap): RowLevelOperationTable = {
- val info = RowLevelOperationInfoImpl(command, options)
+ options: CaseInsensitiveStringMap,
+ updatedColumns: Seq[NamedReference] = Nil): RowLevelOperationTable = {
+ val info = RowLevelOperationInfoImpl(command, options, updatedColumns)
val operation = table.newRowLevelOperationBuilder(info).build()
RowLevelOperationTable(table, operation)
}
+ // Builds a DataSourceV2Relation for a row-level operation, optionally narrowing its output.
+ //
+ // When dataAttrs is non-empty, the relation output is narrowed to include only columns
+ // required for a column-update write. When dataAttrs is empty, the full relation.output is
+ // preserved.
Review Comment:
For function description, please follow the community style like the other code path.
```
/**
* ...
*/
```
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -106,38 +104,92 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val remainingRowsPlan = addOperationColumn(COPY_OPERATION,
Filter(remainingRowFilter, readRelation))
- // the new state is a union of updated and copied over records
- val query = Union(updatedRowsPlan, remainingRowsPlan)
+ val updatedAndRemainingRowsPlan = Union(updatedRowsPlan, remainingRowsPlan)
- // build a plan to replace read groups in the table
val writeRelation = relation.copy(table = operationTable)
- val projections = buildReplaceDataProjections(query, relation.output, metadataAttrs)
+ val query = updatedAndRemainingRowsPlan
+ val metadataAttrs = resolveRequiredMetadataAttrs(relation, operationTable.operation)
+ val projections = buildReplaceDataProjections(query, rowAttrs, metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
ReplaceData(writeRelation, cond, query, relation, projections, groupFilterCond)
}
+ // Common read-relation setup shared by both CoW plan builders.
+ //
+ // When the connector supports column updates and declares required data attributes,
+ // the read relation is narrowed at analysis time so that
+ // GroupBasedRowLevelOperationScanPlanning uses only the needed columns for the scan.
+ // Otherwise the full relation output is used.
Review Comment:
For function description, please follow the community style like the other code path.
```
/**
* ...
*/
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3424004293
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -106,38 +105,78 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val remainingRowsPlan = addOperationColumn(COPY_OPERATION,
Filter(remainingRowFilter, readRelation))
- // the new state is a union of updated and copied over records
- val query = Union(updatedRowsPlan, remainingRowsPlan)
+ val updatedAndRemainingRowsPlan = Union(updatedRowsPlan, remainingRowsPlan)
- // build a plan to replace read groups in the table
val writeRelation = relation.copy(table = operationTable)
Review Comment:
Instead of changing logic in ReplaceData that validates against `relation.output`?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3423724504
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java:
##########
@@ -37,4 +38,20 @@ public interface RowLevelOperationInfo {
* Returns the row-level SQL command (e.g. DELETE, UPDATE, MERGE).
*/
Command command();
+
+ /**
+ * Returns the columns being updated in an UPDATE statement, as non-identity assignments.
+ *
+ * <p>For DELETE and MERGE, returns an empty array.
+ *
+ * <p>Connectors can use this to decide what {@link RowLevelOperation#requiredDataAttributes()}
+ * to declare. For instance, a connector that needs its primary key for row lookup can check
+ * whether pk is already in the updated columns list and, if not, add it to
+ * requiredDataAttributes().
+ *
+ * @since 4.3.0
+ */
+ default NamedReference[] updatedColumns() {
+ return new NamedReference[0];
Review Comment:
Hm... Defaulting to an empty array is fragile. Instances of this type are never constructed by connectors. Can Spark always fill the right value so that we don't have to default here?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4315467582
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/RowLevelOperationInfoImpl.scala:
##########
@@ -17,9 +17,15 @@
package org.apache.spark.sql.connector.write
+import org.apache.spark.sql.connector.expressions.NamedReference
import org.apache.spark.sql.connector.write.RowLevelOperation.Command
import org.apache.spark.sql.util.CaseInsensitiveStringMap
private[sql] case class RowLevelOperationInfoImpl(
command: Command,
- options: CaseInsensitiveStringMap) extends RowLevelOperationInfo
+ options: CaseInsensitiveStringMap,
+ private val updatedCols: Seq[NamedReference])
Review Comment:
**9. This required parameter breaks Iceberg's `ResolveBranch`, and the original info cannot be recovered from the plan.** Iceberg's spark-extensions call `RowLevelOperationInfoImpl(operation.command, options)` in `ResolveBranch` ([spark/v4.2 L54](https://github.com/apache/iceberg/blob/09c80c42e40190774ccaabb2045784df9c43f16a/spark/v4.2/spark-extensions/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveBranch.scala#L54), the same in spark/v4.1). It is a post-hoc rule that rebuilds the operation for a branch write after `RewriteUpdateTable` and swaps it into the write relation and the reads. Iceberg's current modules target Spark 4.1.3 and 4.2.0, but a future module copied from v4.2 for a Spark release with this change will not compile. Once adapted, it can only pass an empty list, because the plan keeps only `RowLevelOperationTable(table, operation)`. The rebuilt operation then sees `updatedColumns() == []` and may declare different `requiredDataAttributes()` than the on
es Spark validated and narrowed against (see item 6).
I am not suggesting a default for `updatedColumns()`. Could the info be kept on `RowLevelOperationTable` as a body member instead, so a rule that rebuilds the operation can pass the same info? A new constructor field would break Iceberg's two-argument `RowLevelOperationTable(_, _)` extractor.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -258,6 +354,407 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep` (an unconditionally declared base column) plus whatever columns
+ // Spark reports as being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = {
+ val base = Seq(FieldReference("pk"), FieldReference("dep"))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ private val clusterRef: NamedReference = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY))
+ .map(FieldReference(_))
+ .getOrElse(PARTITION_COLUMN_REF)
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ // Capture info into a local val so nested writer/commit closures see a stable schema
+ // even if a subsequent newWriteBuilder call mutates lastWriteInfo.
+ val capturedInfo = info
+ val capturedWriteSchema = if (capturedInfo.columnUpdateSchema().isPresent) {
+ capturedInfo.columnUpdateSchema().get()
+ } else {
+ capturedInfo.schema()
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(Array(clusterRef))
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ Array[SortOrder](
+ LogicalExpressions.sort(
+ clusterRef,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering())
+ )
+ }
+
+ override def toBatch: DeltaBatchWrite =
+ new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ info: PhysicalWriteInfo): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(capturedWriteSchema)
+ }
+
+ // For column-update writes, updated rows contain only the declared columns
+ // (columnUpdateSchema from LogicalWriteInfo). We expand each row to the full
+ // table schema by overlaying them on the base row found by pk.
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ val writeSchema = capturedWriteSchema
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
+
+ val mergedData = newData.map { buf =>
+ val merged = new BufferedRows(buf.key, schema)
+ val updateOpName = UTF8String.fromString(Update.toString)
+ val insertOpName = UTF8String.fromString(Insert.toString)
+ buf.log.foreach { logRow =>
+ val opName = logRow.getUTF8String(0)
+ if (opName == updateOpName) {
+ val pk = logRow.getInt(1)
+ val narrowRow = logRow.get(3, writeSchema).asInstanceOf[InternalRow]
+ val baseRow = currentRowByPk(pk)
+ val fullRow = overlayNarrowRow(baseRow, narrowRow, writeFieldIdx)
+ merged.rows.append(fullRow)
+ } else if (opName == insertOpName) {
+ // INSERT rows arrive with the full table schema via writer.insert()
+ val insertRow = logRow.get(3, schema).asInstanceOf[InternalRow]
+ merged.rows.append(insertRow.copy())
+ }
+ }
+ merged
+ }
+
+ withDeletes(newData)
+ withData(mergedData)
+ lastWriteLog = newData.flatMap(buffer => buffer.log).toIndexedSeq
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ }
+ }
+ }
+ }
+
+ class DeltaBasedColumnUpdateOperationWithReqAttrs(
+ command: Command,
+ reqCols: Array[String],
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, options = options) {
+ override def requiredDataAttributes(): Array[NamedReference] = reqCols.map(FieldReference(_))
+ }
+
+ class DeltaBasedColumnUpdateSplitOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedColumnUpdateOperation(command, updatedCols, options) {
+ override def representUpdateAsDeleteAndInsert(): Boolean = true
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
Review Comment:
**18.** Minor: `DeltaBasedColumnUpdateSplitOperation.newWriteBuilder` copies about 70 lines of the parent's `newWriteBuilder` (L378-451); only the operation tag and the `pk` source differ. The copies have drifted twice: round 4 found the dropped `Insert` rows, fixed by mirroring the code, and this copy hard-codes `PARTITION_COLUMN_REF` (L483, L486), so `column-update-cluster-by` (L374-376) is silently ignored on split tables. Smaller items in the fixture: `info.columnUpdateSchema().orElse(info.schema())` would replace the three `isPresent` blocks (L383-387, L474-478, L632-636), `val capturedInfo = info` has no effect, `PartitionBasedColumnUpdateOperationNoWriteColumnUpdate` could be a flag in `PartitionBasedNarrowReplaceData.createBatchWriterFactory` like `COLUMN_UPDATE_COW_RECORD_ONLY_WRITER`, and the `fullSchema` realignment (L742-747) is an identity copy. In the suites, about seven tests are copied between `DeltaBasedColumnUpdateTableSuite` and `GroupBasedColumnUpdateTableSuite
` (DELETE and MERGE fallback, CHECK enforcement, nested identity update, nested CHECK, partial pruning, and dynamic options, which `UpdateTableSuiteBase` already covers), and the `target_deps` view (`GroupBasedColumnUpdateTableSuite.scala:75`) is not wrapped in `withTempView`. Could the split operation reuse the parent's builder, and could the shared tests live in one trait, like `MergeIntoSchemaEvolutionExtraSQLTests`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala:
##########
@@ -109,22 +111,75 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper {
WriteToDataSourceV2(Some(r), microBatchWrite, newQuery, customMetrics)
case rd @ ReplaceData(r: DataSourceV2Relation, _, query, _, projections, _, None) =>
- val rowSchema = projections.rowProjection.schema
+ val (rowSchema, columnUpdateSchema) =
+ writeSchemas(projections.rowProjection.schema, hasNarrowRows(rd))
val metadataSchema = projections.metadataProjection.map(_.schema)
val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap)
- val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema)
+ val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema,
+ columnUpdateSchema)
val write = writeBuilder.build()
+ validateColumnUpdateWrite(rd, write)
val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, r.funCatalog)
rd.copy(write = Some(write), query = newQuery)
case wd @ WriteDelta(r: DataSourceV2Relation, _, query, _, projections, _, None) =>
val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap)
- val deltaWriteBuilder = newDeltaWriteBuilder(r.table, writeOptions, projections)
+ val deltaWriteBuilder =
+ newDeltaWriteBuilder(r.table, writeOptions, projections, hasNarrowRows(wd))
val deltaWrite = deltaWriteBuilder.build()
+ validateColumnUpdateWrite(wd, deltaWrite)
val newQuery = DistributionAndOrderingUtils.prepareQuery(deltaWrite, query, r.funCatalog)
wd.copy(write = Some(deltaWrite), query = newQuery)
}
+ /**
+ * Whether a row-level write delivers narrow rows for updated, copied, and reinserted records.
+ * This must match the writes whose write relation RewriteUpdateTable narrows, even when the
+ * declared columns cover the whole table.
+ */
+ private[v2] def hasNarrowRows(write: RowLevelWrite): Boolean = {
+ write.operation.isInstanceOf[SupportsColumnUpdates] && write.operation.command == UPDATE
+ }
+
+ /**
+ * Checks that the distribution and ordering of a write that delivers narrow rows reference only
+ * columns the write reads. Other columns stay in the query only until column pruning removes
+ * them, so they are rejected here whether or not pruning ran.
+ */
+ private def validateColumnUpdateWrite(rowLevelWrite: RowLevelWrite, write: Write): Unit = {
+ write match {
+ case w: RequiresDistributionAndOrdering if hasNarrowRows(rowLevelWrite) =>
+ val distribution = w.requiredDistribution match {
+ case d: ClusteredDistribution => d.clustering.toImmutableArraySeq
+ case d: OrderedDistribution => d.ordering.toImmutableArraySeq
+ case _ => Seq.empty
+ }
+ val readNames = rowLevelWrite.references.toSeq.map(_.name)
+ val unread = (distribution ++ w.requiredOrdering.toImmutableArraySeq)
+ .flatMap(_.references.toImmutableArraySeq)
Review Comment:
**7. `validateColumnUpdateWrite` misses columns nested inside transform arguments.** `.flatMap(_.references.toImmutableArraySeq)` relies on `V2Expression.references()`, and `ApplyTransform.references` returns only top-level `NamedReference` arguments (`expressions.scala:220-222`). For a write that requires `clustered(Expressions.apply("f", Expressions.apply("g", Expressions.column("dep"))))` with `requiredDataAttributes()` = `[pk, salary]`, `UPDATE t SET salary = -1 WHERE pk = 1` passes this check. With the default optimizer, `prepareQuery` then fails with the generic `_LEGACY_ERROR_TEMP_1137` while resolving the inner `dep`, and with `ColumnPruning` excluded the write is accepted. That is the pruning-dependent outcome that the Scaladoc above (L144-148) says this check removes.
`V2ScanPartitioningAndOrdering.unresolvableColumns`, added by SPARK-59721 (#58979) while reviewing this PR, walks `Transform.arguments()` and `SortOrder.expression()` for exactly this reason. Could you share it (e.g. make it `private[v2]` or move it to `V2ExpressionUtils`) and call it against `LocalRelation(rowLevelWrite.references.toSeq)` here?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-5771914538
Thanks for the new round @peter-toth. I addressed them in [9932f88](https://github.com/apache/spark/pull/55518/commits/9932f880077cb4332dc8bc291efaacea63d39854).
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-6091300595
Gentle ping, @anuragmantri . Please rebase your PR in order to keep healthy in sync with the Apache Spark `master` branch.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4364254863
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala:
##########
@@ -380,14 +380,17 @@ case class ReplaceDataExec(
write: Write,
rowLevelCommand: RowLevelOperation.Command,
tableName: String,
+ useWriteColumnUpdate: Boolean = false,
Review Comment:
Done in d8a4641: `useWriteColumnUpdate` has no default again. The only construction, in `DataSourceV2Strategy`, passes it by name.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala:
##########
@@ -356,6 +357,32 @@ trait RowLevelWrite extends V2WriteCommand with SupportsSubquery {
def condition: Expression
def originalTable: NamedRelation
+ /** The projections that read query columns by ordinal. */
+ protected def queryProjections: Seq[ProjectingInternalRow]
+
+ /**
+ * The write reads the operation column at ordinal 0 and other query columns only through
+ * projections, by ordinal. Any other query column can be pruned with [[withPrunedQuery]].
+ */
+ override lazy val references: AttributeSet = {
+ val output = query.output
Review Comment:
Done in d8a4641.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -321,23 +810,61 @@ class InMemoryRowLevelOperationTable private (
}
}
-private class DeltaBufferedRowsWriterFactory(schema: StructType) extends DeltaWriterFactory {
+private class DeltaBufferedRowsWriterFactory(
+ schema: StructType,
+ intRowId: Boolean = true) extends DeltaWriterFactory {
override def createWriter(partitionId: Int, taskId: Long): DeltaWriter[InternalRow] = {
- new DeltaBufferWriter(schema)
+ new DeltaBufferWriter(schema, intRowId)
+ }
+}
+
+// Test-only: a writer factory for a connector that never overrides writeColumnUpdate(), to
+// exercise the DataWriter#writeColumnUpdate default that throws
+// DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED.
+// Declared top-level (no reference to the enclosing table) so it stays task-serializable.
+private class NoWriteColumnUpdateOverrideWriterFactory extends DataWriterFactory {
+ override def createWriter(partitionId: Int, taskId: Long): DataWriter[InternalRow] = {
+ new NoWriteColumnUpdateOverrideWriter
}
}
-private class DeltaBufferWriter(schema: StructType) extends BufferWriter(schema)
- with DeltaWriter[InternalRow] {
+private class NoWriteColumnUpdateOverrideWriter extends DataWriter[InternalRow] {
+ override def write(record: InternalRow): Unit = {}
+ override def commit(): WriterCommitMessage = new BufferedRows(Seq.empty, StructType(Nil))
+ override def abort(): Unit = {}
+ override def close(): Unit = {}
+}
+
+// Test-only: a writer that overrides writeColumnUpdate(record) but not
+// writeColumnUpdate(metadata, record), so rows with metadata reach it through the default.
+private class RecordOnlyColumnUpdateWriterFactory(schema: StructType) extends DataWriterFactory {
+ override def createWriter(partitionId: Int, taskId: Long): DataWriter[InternalRow] = {
+ new RecordOnlyColumnUpdateWriter(schema)
+ }
+}
+
+private class RecordOnlyColumnUpdateWriter(schema: StructType) extends DataWriter[InternalRow] {
+ private val delegate = new BufferWriter(schema)
+ override def write(record: InternalRow): Unit = delegate.write(record)
+ override def writeColumnUpdate(record: InternalRow): Unit = delegate.writeColumnUpdate(record)
+ override def commit(): WriterCommitMessage = delegate.commit()
+ override def abort(): Unit = delegate.abort()
+ override def close(): Unit = delegate.close()
+}
+
+// Logs the row ID of a DELETE as an integer when `intRowId` is set and as null otherwise; the
+// commit then resolves such a DELETE from its metadata.
+private class DeltaBufferWriter(schema: StructType, intRowId: Boolean = true)
+ extends BufferWriter(schema) with DeltaWriter[InternalRow] {
private final val DELETE = UTF8String.fromString(Delete.toString)
private final val UPDATE = UTF8String.fromString(Update.toString)
private final val REINSERT = UTF8String.fromString(Reinsert.toString)
private final val INSERT = UTF8String.fromString(Insert.toString)
override def delete(meta: InternalRow, id: InternalRow): Unit = {
- val pk = id.getInt(0)
- buffer.deletes += pk
+ val pk: Any = if (intRowId) id.getInt(0) else null
Review Comment:
Done in d8a4641
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala:
##########
@@ -356,6 +357,32 @@ trait RowLevelWrite extends V2WriteCommand with SupportsSubquery {
def condition: Expression
def originalTable: NamedRelation
+ /** The projections that read query columns by ordinal. */
+ protected def queryProjections: Seq[ProjectingInternalRow]
Review Comment:
Done: the compatibility paragraph now says "`RowLevelWrite` has new abstract `queryProjections` and `withPrunedQuery`".
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -128,12 +153,61 @@ class InMemoryRowLevelOperationTable private (
copied.replacedPartitions = replacedPartitions
copied.lastWriteInfo = lastWriteInfo
copied.lastWriteLog = lastWriteLog
+ copied.lastUpdatedColumns = lastUpdatedColumns
copied
}
override def newRowLevelOperationBuilder(
info: RowLevelOperationInfo): RowLevelOperationBuilder = {
- if (properties.getOrDefault(SUPPORTS_DELTAS, "false") == "true") {
+ lastUpdatedColumns = info.updatedColumns()
+ val splitUpdates = properties.getOrDefault(SPLIT_UPDATES, "false").toBoolean
+ if (properties.getOrDefault(COLUMN_UPDATE, "false") == "true") {
+ if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ }
+ } else if (properties.containsKey(COLUMN_UPDATE_REQ_ATTRS)) {
+ val reqCols = properties.get(COLUMN_UPDATE_REQ_ATTRS).split(",").map(_.trim)
+ if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperationWithReqAttrs(
+ info.command, reqCols, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperationWithReqAttrs(info.command, reqCols, info.options)
+ }
+ } else if (properties.getOrDefault(COLUMN_UPDATE_EMPTY_REQ_ATTRS, "false") == "true") {
+ // Test-only: returns an empty requiredDataAttributes() so we can verify Spark rejects it.
+ () => new DeltaBasedColumnUpdateOperationWithReqAttrs(
+ info.command, Array.empty, info.options)
+ } else if (properties.getOrDefault(COLUMN_UPDATE_COW, "false") == "true") {
+ if (info.updatedColumns().exists(_.describe() == "pk")) {
+ () => PartitionBasedOperation(info.command, info.options)
+ } else {
+ () => new PartitionBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ }
+ } else if (properties.containsKey(COLUMN_UPDATE_COW_REQ_ATTRS)) {
+ val reqCols = properties.get(COLUMN_UPDATE_COW_REQ_ATTRS).split(",").map(_.trim)
+ () => new PartitionBasedColumnUpdateOperationWithReqAttrs(
+ info.command, reqCols, info.options)
+ } else if (
+ properties.getOrDefault(COLUMN_UPDATE_COW_NO_WRITE_COLUMN_UPDATE, "false") == "true") {
+ () => new PartitionBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ } else if (properties.getOrDefault(COLUMN_UPDATE_SPLIT, "false") == "true") {
Review Comment:
Done in d8a4641.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -343,3 +365,60 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
}
}
}
+
+object RewriteRowLevelCommand extends SQLConfHelper {
+
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation, in declared order and with the relation's spelling. An empty declaration,
+ * nested fields, duplicate columns, columns that do not exist in the relation and metadata
+ * columns are rejected with an `AnalysisException`.
+ */
+ private[sql] def resolveRequiredDataAttrs(
+ relation: LogicalPlan,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes.toImmutableArraySeq
+ if (refs.isEmpty) {
+ throw QueryCompilationErrors.emptyRequiredDataAttributesError(operation.getClass.getName)
+ }
+ val nested = refs.filter(_.fieldNames.length != 1).map(_.fieldNames.toSeq)
Review Comment:
Done in d8a4641.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,637 @@
+/*
+ * 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.{SparkException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTable
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.connector.write.RowLevelOperationTable
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+class GroupBasedColumnUpdateTableSuite extends UpdateTableSuiteBase
+ with ColumnUpdateTableSharedTests {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-cow", "true")
+ props
+ }
+
+ override protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String] = {
+ java.util.Map.of("column-update-cow-req-attrs", reqAttrs)
+ }
+
+ override protected def fullWidthTableProps: java.util.Map[String, String] = java.util.Map.of()
+
+ test("column-update ReplaceData: column update schema has only the declared columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "dep", "salary"))
+ }
+
+ test("column-update ReplaceData: data correctness -- bonus preserved, salary updated") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ // CoW: pk=1 and pk=3 match (2 updates); the 'hr' partition has no other rows so no copies.
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update ReplaceData: subquery WHERE condition data correctness") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ withTempView("target_deps") {
+ import testImplicits._
+ val subqueryDF = Seq("hr").toDF()
+ subqueryDF.createOrReplaceTempView("target_deps")
+
+ sql(
+ s"""UPDATE $tableNameAsString
+ |SET salary = -1
+ |WHERE dep IN (SELECT * FROM target_deps)
+ |""".stripMargin)
+ }
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update ReplaceData: runtime group filtering data correctness") {
+ Seq(true, false).foreach { dppEnabled =>
+ Seq(true, false).foreach { aqeEnabled =>
+ withSQLConf(
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> dppEnabled.toString,
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString) {
+ withTable(tableNameAsString) {
+ withTempView("matched_pk") {
+ createAndInitTable("pk INT, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ import testImplicits._
+ val matchedPkDF = Seq(Some(1), None).toDF()
+ matchedPkDF.createOrReplaceTempView("matched_pk")
+
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = -1 " +
+ s"WHERE pk IN (SELECT * FROM matched_pk)",
+ primaryScanSchema =
+ "pk INT, salary INT, dep STRING, _partition STRING, index INT",
+ groupFilterScanSchema = Some("pk INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, 300, 30, "hr") :: Nil)
+
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("column-update ReplaceData: updated and copied rows reach writeColumnUpdate with " +
+ "metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", null), data = Row(1, "hr", -1)),
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", 1), data = Row(2, "hr", 200)))
+ }
+
+ test("column-update ReplaceData: without metadata, rows reach writeColumnUpdate(record)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "no-metadata", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Row(3, 300, 30, "software") :: Nil)
+ }
+
+ test("column-update ReplaceData: a writer overriding only writeColumnUpdate(record) " +
+ "receives rows with metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "column-update-cow-record-only-writer", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: without metadata, column update schema is the declared " +
+ "layout") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "no-metadata", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedColumnUpdateSchema = Some(StructType(Seq(PK_FIELD, DEP_FIELD,
+ StructField("salary", IntegerType, nullable = true)))))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 110, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: connector missing writeColumnUpdate override is rejected") {
+ createAndInitTableReplaceDataNoWriteColumnUpdate("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[SparkUnsupportedOperationException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED",
+ parameters = Map(
+ "class" -> "org.apache.spark.sql.connector.catalog.NoWriteColumnUpdateOverrideWriter"))
+ }
+
+ private def createAndInitTableReplaceDataNoWriteColumnUpdate(
+ schemaString: String, jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-cow-no-write-column-update", "true"))
+ }
+
+ // ---------------------------------------------------------------------------
+ // Write relation tests: verify that only the write relation is narrowed to the
+ // declared columns, on both copy-on-write plan shapes.
+ // ---------------------------------------------------------------------------
+
+ private val reqAttrsOperationClass =
+ classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperationWithReqAttrs].getName
+
+ private def createAndInitTableWithReqAttrs(reqAttrs: String): Unit = {
+ createAndInitTableWithProps(reqAttrsTableProps(reqAttrs))
+ }
+
+ test("column-update ReplaceData: write relation has the declared columns in declared order") {
+ createAndInitTableWithReqAttrs("salary,pk")
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'")
+ }
+
+ assert(!write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "salary", "pk")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update ReplaceData: subquery write relation has the declared columns") {
+ withTempView("deps") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE dep IN (SELECT dep FROM deps)")
+ }
+
+ assert(write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "pk", "salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+ }
+
+ test("column-update ReplaceData: runtime group filtering on an undeclared condition column") {
+ withTempView("matched_ids") {
+ // `dep` is declared because the scan filters groups at runtime by the partition column
+ createAndInitTableWithReqAttrs("pk,dep,salary")
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `id` is neither declared nor in the write relation; the group filter must still map it
+ // onto the second read of the table that the subquery rewrite adds
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, dep STRING, salary INT, bonus INT, extra INT, " +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = Some("id INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+
+ test("column-update ReplaceData: scan reads only the referenced columns") {
+ createAndInitTableWithReqAttrs("pk,salary")
+
+ // `id` is neither declared nor referenced; `dep`, `bonus` and `extra` are read for the
+ // condition, the assignment and the CHECK constraint
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'",
+ "pk INT, dep STRING, salary INT, bonus INT, extra INT, _partition STRING, index INT")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update ReplaceData: subquery scan reads only the referenced columns") {
+ withTempView("deps") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE dep IN (SELECT dep FROM deps)",
+ "pk INT, dep STRING, salary INT, bonus INT, extra INT, _partition STRING, index INT")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+ }
+
+ test("column-update ReplaceData: unread partition column skips runtime group filtering") {
+ withTempView("matched_ids") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `dep` is neither declared nor referenced, so the scan doesn't read it and doesn't
+ // report it from filterAttributes(); every group is read and replaced
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, salary INT, bonus INT, extra INT, " +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = None)
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr", "software"))
+ }
+ }
+
+ test("column-update ReplaceData: runtime filter attribute the scan doesn't read is rejected") {
+ withTempView("matched_ids") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, salary INT",
+ """{ "pk": 1, "id": 1, "dep": "hr", "salary": 100 }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-cow-req-attrs", "pk,salary",
+ "column-update-cow-unread-filter-attrs", "true"))
+ import testImplicits._
+ Seq(1).toDF().createOrReplaceTempView("matched_ids")
+
+ val ex = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE id IN (SELECT * FROM matched_ids)")
+ }
+ assert(ex.getCondition == "DATA_SOURCE_INVALID_RUNTIME_FILTER_ATTRIBUTE.CANNOT_RESOLVE")
Review Comment:
Done in d8a4641.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,637 @@
+/*
+ * 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.{SparkException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTable
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.connector.write.RowLevelOperationTable
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+class GroupBasedColumnUpdateTableSuite extends UpdateTableSuiteBase
+ with ColumnUpdateTableSharedTests {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-cow", "true")
+ props
+ }
+
+ override protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String] = {
+ java.util.Map.of("column-update-cow-req-attrs", reqAttrs)
+ }
+
+ override protected def fullWidthTableProps: java.util.Map[String, String] = java.util.Map.of()
+
+ test("column-update ReplaceData: column update schema has only the declared columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "dep", "salary"))
+ }
+
+ test("column-update ReplaceData: data correctness -- bonus preserved, salary updated") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ // CoW: pk=1 and pk=3 match (2 updates); the 'hr' partition has no other rows so no copies.
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update ReplaceData: subquery WHERE condition data correctness") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ withTempView("target_deps") {
+ import testImplicits._
+ val subqueryDF = Seq("hr").toDF()
+ subqueryDF.createOrReplaceTempView("target_deps")
+
+ sql(
+ s"""UPDATE $tableNameAsString
+ |SET salary = -1
+ |WHERE dep IN (SELECT * FROM target_deps)
+ |""".stripMargin)
+ }
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update ReplaceData: runtime group filtering data correctness") {
+ Seq(true, false).foreach { dppEnabled =>
+ Seq(true, false).foreach { aqeEnabled =>
+ withSQLConf(
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> dppEnabled.toString,
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString) {
+ withTable(tableNameAsString) {
+ withTempView("matched_pk") {
+ createAndInitTable("pk INT, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ import testImplicits._
+ val matchedPkDF = Seq(Some(1), None).toDF()
+ matchedPkDF.createOrReplaceTempView("matched_pk")
+
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = -1 " +
+ s"WHERE pk IN (SELECT * FROM matched_pk)",
+ primaryScanSchema =
+ "pk INT, salary INT, dep STRING, _partition STRING, index INT",
+ groupFilterScanSchema = Some("pk INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, 300, 30, "hr") :: Nil)
+
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("column-update ReplaceData: updated and copied rows reach writeColumnUpdate with " +
+ "metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", null), data = Row(1, "hr", -1)),
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", 1), data = Row(2, "hr", 200)))
+ }
+
+ test("column-update ReplaceData: without metadata, rows reach writeColumnUpdate(record)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "no-metadata", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Row(3, 300, 30, "software") :: Nil)
+ }
+
+ test("column-update ReplaceData: a writer overriding only writeColumnUpdate(record) " +
+ "receives rows with metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "column-update-cow-record-only-writer", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: without metadata, column update schema is the declared " +
+ "layout") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "no-metadata", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedColumnUpdateSchema = Some(StructType(Seq(PK_FIELD, DEP_FIELD,
+ StructField("salary", IntegerType, nullable = true)))))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 110, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: connector missing writeColumnUpdate override is rejected") {
+ createAndInitTableReplaceDataNoWriteColumnUpdate("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[SparkUnsupportedOperationException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED",
+ parameters = Map(
+ "class" -> "org.apache.spark.sql.connector.catalog.NoWriteColumnUpdateOverrideWriter"))
+ }
+
+ private def createAndInitTableReplaceDataNoWriteColumnUpdate(
+ schemaString: String, jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-cow-no-write-column-update", "true"))
+ }
+
+ // ---------------------------------------------------------------------------
+ // Write relation tests: verify that only the write relation is narrowed to the
+ // declared columns, on both copy-on-write plan shapes.
+ // ---------------------------------------------------------------------------
+
+ private val reqAttrsOperationClass =
+ classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperationWithReqAttrs].getName
+
+ private def createAndInitTableWithReqAttrs(reqAttrs: String): Unit = {
+ createAndInitTableWithProps(reqAttrsTableProps(reqAttrs))
+ }
+
+ test("column-update ReplaceData: write relation has the declared columns in declared order") {
+ createAndInitTableWithReqAttrs("salary,pk")
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'")
+ }
+
+ assert(!write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "salary", "pk")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update ReplaceData: subquery write relation has the declared columns") {
+ withTempView("deps") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE dep IN (SELECT dep FROM deps)")
+ }
+
+ assert(write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "pk", "salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+ }
+
+ test("column-update ReplaceData: runtime group filtering on an undeclared condition column") {
+ withTempView("matched_ids") {
+ // `dep` is declared because the scan filters groups at runtime by the partition column
+ createAndInitTableWithReqAttrs("pk,dep,salary")
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `id` is neither declared nor in the write relation; the group filter must still map it
+ // onto the second read of the table that the subquery rewrite adds
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, dep STRING, salary INT, bonus INT, extra INT, " +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = Some("id INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+
+ test("column-update ReplaceData: scan reads only the referenced columns") {
+ createAndInitTableWithReqAttrs("pk,salary")
+
+ // `id` is neither declared nor referenced; `dep`, `bonus` and `extra` are read for the
+ // condition, the assignment and the CHECK constraint
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'",
+ "pk INT, dep STRING, salary INT, bonus INT, extra INT, _partition STRING, index INT")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update ReplaceData: subquery scan reads only the referenced columns") {
+ withTempView("deps") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE dep IN (SELECT dep FROM deps)",
+ "pk INT, dep STRING, salary INT, bonus INT, extra INT, _partition STRING, index INT")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+ }
+
+ test("column-update ReplaceData: unread partition column skips runtime group filtering") {
+ withTempView("matched_ids") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `dep` is neither declared nor referenced, so the scan doesn't read it and doesn't
+ // report it from filterAttributes(); every group is read and replaced
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, salary INT, bonus INT, extra INT, " +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = None)
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr", "software"))
+ }
+ }
+
+ test("column-update ReplaceData: runtime filter attribute the scan doesn't read is rejected") {
+ withTempView("matched_ids") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, salary INT",
+ """{ "pk": 1, "id": 1, "dep": "hr", "salary": 100 }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-cow-req-attrs", "pk,salary",
+ "column-update-cow-unread-filter-attrs", "true"))
+ import testImplicits._
+ Seq(1).toDF().createOrReplaceTempView("matched_ids")
+
+ val ex = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE id IN (SELECT * FROM matched_ids)")
+ }
+ assert(ex.getCondition == "DATA_SOURCE_INVALID_RUNTIME_FILTER_ATTRIBUTE.CANNOT_RESOLVE")
+ assert(ex.getMessageParameters.get("attribute") == "`dep`")
+ assert(ex.getMessageParameters.get("method") == "filterAttributes()")
+ }
+ }
+
+ test("column-update ReplaceData: a connector without column updates reads every column") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of[String, String]())
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = 0",
+ "pk INT, salary INT, dep STRING, _partition STRING, index INT")
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, 0, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: undeclared distribution column is rejected") {
+ // without metadata, the write clusters by `dep`
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow-req-attrs", "pk,salary", "no-metadata", "true"))
+
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName
+ for {
+ cond <- Seq("pk = 1", "dep = 'hr'")
+ confs <- Seq(Nil, Seq(excludeColumnPruning))
+ } withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE $cond")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[`dep`]"))
+ }
+ }
+
+ test("column-update ReplaceData: other operations read every column") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of())
+ val fullScanSchema = "pk INT, salary INT, bonus INT, dep STRING, _partition STRING, index INT"
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1", fullScanSchema)
+
+ executeAndCheckScan(s"DELETE FROM $tableNameAsString WHERE pk > 1", fullScanSchema)
+
+ withTempView("source") {
+ import testImplicits._
+ Seq((1, 150)).toDF("pk", "salary").createOrReplaceTempView("source")
+ executeAndCheckScan(
+ s"""MERGE INTO $tableNameAsString t
+ |USING source s
+ |ON t.pk = s.pk
+ |WHEN MATCHED THEN UPDATE SET salary = s.salary
+ |""".stripMargin,
+ fullScanSchema)
+ }
+
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, 150, 10, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: DELETE and MERGE on a column-update table read every column") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin)
+ val fullScanSchema = "pk INT, salary INT, bonus INT, dep STRING, _partition STRING, index INT"
+
+ executeAndCheckScan(s"DELETE FROM $tableNameAsString WHERE pk > 1", fullScanSchema)
+
+ withTempView("source") {
+ import testImplicits._
+ Seq((1, 150)).toDF("pk", "salary").createOrReplaceTempView("source")
+ executeAndCheckScan(
+ s"""MERGE INTO $tableNameAsString t
+ |USING source s
+ |ON t.pk = s.pk
+ |WHEN MATCHED THEN UPDATE SET salary = s.salary
+ |""".stripMargin,
+ fullScanSchema)
+ }
+
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, 150, 10, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: assigned column outside requiredDataAttributes is rejected") {
+ createAndInitTableWithReqAttrs("pk,salary")
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET bonus = 0 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "missingColumns" -> "[`bonus`]"))
+ }
+
+ test("column-update ReplaceData: assigned column outside requiredDataAttributes is rejected " +
+ "on the subquery path") {
+ withTempView("deps") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET bonus = 0 WHERE dep IN (SELECT dep FROM deps)")
+ },
+ condition = "COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "missingColumns" -> "[`bonus`]"))
+ }
+ }
+
+ test("column-update ReplaceData: analyzed plan shows the declared columns as the write target") {
+ createAndInitTableWithReqAttrs("pk,salary")
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'")
+ }
+
+ val header = write.simpleString(Int.MaxValue)
+ assert(header.matches("""ReplaceData RelationV2\[pk#\d+, salary#\d+\] .*"""), header)
+ }
+
+ test("column-update ReplaceData: an operation replaced after the rewrite is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ }
+ val relation = write.table.asInstanceOf[DataSourceV2Relation]
+ val operationTable = relation.table.asInstanceOf[RowLevelOperationTable]
+ val t = table
+ Seq(
+ t.PartitionBasedOperation(UPDATE, CaseInsensitiveStringMap.empty()) -> "[pk, salary, dep]",
+ new t.PartitionBasedColumnUpdateOperation(UPDATE) -> "[pk, dep]"
+ ).foreach { case (operation, expectedNames) =>
+ val replaced = write.withNewTable(
+ relation.copy(table = operationTable.copy(operation = operation)(
+ operationTable.operationInfo)))
+ checkError(
+ exception = intercept[SparkException](V2Writes(replaced)),
+ condition = "INTERNAL_ERROR",
+ parameters = Map("message" ->
+ (s"The write relation of operation ${operation.getClass.getName} has columns " +
+ s"[pk, dep, salary], but the operation expects $expectedNames.")))
+ }
+ }
+
+ test("column-update ReplaceData: an operation rebuilt with withTable keeps the updated columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ }
+ val relation = write.table.asInstanceOf[DataSourceV2Relation]
+ val operationTable = relation.table.asInstanceOf[RowLevelOperationTable]
+ val t = table
+ t.lastUpdatedColumns = null
+
+ val rebuilt = operationTable.withTable(t)
+
+ assert(t.lastUpdatedColumns.map(_.describe()).toSeq == Seq("salary"))
Review Comment:
Done in d8a4641: it is `checkLastUpdatedColumns("salary")` now.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala:
##########
@@ -53,6 +53,68 @@ abstract class UpdateTableSuiteBase extends RowLevelOperationSuiteBase {
s"Expected numCopiedRows=$expectedCopied, got ${summary.numCopiedRows()}")
}
+ test("SPARK-58111: updatedColumns: single non-identity assignment") {
Review Comment:
Done in d8a4641.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/RowLevelOperationTable.scala:
##########
@@ -58,4 +63,13 @@ private[sql] case class RowLevelOperationTable(
override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = {
operation.newWriteBuilder(info)
}
+
+ /**
+ * Builds the operation for another table from the same info, for a rule that redirects the
+ * write after the rewrite, such as to a branch of the table.
+ */
+ def withTable(newTable: Table with SupportsRowLevelOperations): RowLevelOperationTable = {
Review Comment:
Done in d8a4641.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -128,12 +153,61 @@ class InMemoryRowLevelOperationTable private (
copied.replacedPartitions = replacedPartitions
copied.lastWriteInfo = lastWriteInfo
copied.lastWriteLog = lastWriteLog
+ copied.lastUpdatedColumns = lastUpdatedColumns
copied
}
override def newRowLevelOperationBuilder(
info: RowLevelOperationInfo): RowLevelOperationBuilder = {
- if (properties.getOrDefault(SUPPORTS_DELTAS, "false") == "true") {
+ lastUpdatedColumns = info.updatedColumns()
+ val splitUpdates = properties.getOrDefault(SPLIT_UPDATES, "false").toBoolean
+ if (properties.getOrDefault(COLUMN_UPDATE, "false") == "true") {
+ if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ }
+ } else if (properties.containsKey(COLUMN_UPDATE_REQ_ATTRS)) {
+ val reqCols = properties.get(COLUMN_UPDATE_REQ_ATTRS).split(",").map(_.trim)
+ if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperationWithReqAttrs(
+ info.command, reqCols, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperationWithReqAttrs(info.command, reqCols, info.options)
+ }
+ } else if (properties.getOrDefault(COLUMN_UPDATE_EMPTY_REQ_ATTRS, "false") == "true") {
+ // Test-only: returns an empty requiredDataAttributes() so we can verify Spark rejects it.
+ () => new DeltaBasedColumnUpdateOperationWithReqAttrs(
+ info.command, Array.empty, info.options)
+ } else if (properties.getOrDefault(COLUMN_UPDATE_COW, "false") == "true") {
+ if (info.updatedColumns().exists(_.describe() == "pk")) {
+ () => PartitionBasedOperation(info.command, info.options)
+ } else {
+ () => new PartitionBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ }
+ } else if (properties.containsKey(COLUMN_UPDATE_COW_REQ_ATTRS)) {
Review Comment:
Done in d8a4641.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -343,3 +365,60 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
}
}
}
+
+object RewriteRowLevelCommand extends SQLConfHelper {
+
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation, in declared order and with the relation's spelling. An empty declaration,
+ * nested fields, duplicate columns, columns that do not exist in the relation and metadata
+ * columns are rejected with an `AnalysisException`.
+ */
+ private[sql] def resolveRequiredDataAttrs(
+ relation: LogicalPlan,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes.toImmutableArraySeq
+ if (refs.isEmpty) {
+ throw QueryCompilationErrors.emptyRequiredDataAttributesError(operation.getClass.getName)
+ }
+ val nested = refs.filter(_.fieldNames.length != 1).map(_.fieldNames.toSeq)
+ if (nested.nonEmpty) {
+ throw QueryCompilationErrors.nestedRequiredDataAttributeError(
+ operation.getClass.getName, nested)
+ }
+ val normalizedNames = refs.map { ref =>
Review Comment:
Done in d8a4641: it is `SchemaUtils.foldName(ref.fieldNames.head, conf.caseSensitiveAnalysis)` now, and the `Locale` import is gone.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -141,6 +153,19 @@ abstract class RowLevelOperationSuiteBase
append(schemaString, jsonData)
}
+ protected def createAndInitTable(
Review Comment:
Done in d8a4641.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,1114 @@
+/*
+ * 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.{SparkException, SparkRuntimeException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+import org.apache.spark.sql.catalyst.util.RowDeltaUtils.ORIGINAL_ROW_ID_VALUE_PREFIX
+import org.apache.spark.sql.connector.catalog.{Delete, InMemoryBaseTable, InMemoryRowLevelOperationTable, Reinsert}
+import org.apache.spark.sql.connector.write.{LogicalWriteInfo, RowLevelOperationTable}
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.errors.QueryCompilationErrors
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, updated rows contain only the declared columns
+ * (LogicalWriteInfo.columnUpdateSchema()) rather than the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase
+ with ColumnUpdateTableSharedTests {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ override protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String] = {
+ java.util.Map.of("column-update-req-attrs", reqAttrs)
+ }
+
+ override protected def fullWidthTableProps: java.util.Map[String, String] = {
+ java.util.Map.of("supports-deltas", "true")
+ }
+
+ override protected def scanMetadataColumns: String = ", _partition STRING"
+
+ private val reqAttrsOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateOperationWithReqAttrs].getName
+
+ private val splitOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperation].getName
+
+ private val splitReqAttrsOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithReqAttrs]
+ .getName
+
+ private val splitWithRowIdOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithRowId].getName
+
+ test("column-update: column update schema contains a single assigned column") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: column update schema contains multiple assigned columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false),
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: column update schema is pk and dep for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ // All assignments are identity, so updatedColumns is empty; pk and dep remain because the
+ // connector declares them as base columns.
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Array(PK_FIELD, DEP_FIELD))))
+ }
+
+ test("column-update: row filter condition is orthogonal to column narrowing") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = 'engineering' WHERE pk IN (1, 3)")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: update all rows (no WHERE clause)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("salary", IntegerType, nullable = true)
+ ))))
+ }
+
+ test("column-update: column update schema excludes identity assignments in a mixed UPDATE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: CHAR and VARCHAR self-assignments are reported as updated") {
+ createTable("pk INT NOT NULL, id INT, c CHAR(3), v VARCHAR(10), dep STRING")
+ sql(s"INSERT INTO $tableNameAsString VALUES (1, 1, 'a', 'x', 'hr')")
+
+ sql(s"UPDATE $tableNameAsString SET id = id, c = c, v = v WHERE pk = 1")
+
+ checkLastUpdatedColumns("c", "v")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, 1, "a ", "x", "hr") :: Nil)
+ }
+
+ test("column-update: a LogicalWriteInfo that does not implement columnUpdateSchema throws") {
+ val info = new LogicalWriteInfo {
+ override def options(): CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty()
+ override def queryId(): String = "query"
+ override def schema(): StructType = new StructType()
+ }
+
+ checkError(
+ exception = intercept[SparkUnsupportedOperationException](info.columnUpdateSchema()),
+ condition = "DATA_SOURCE_COLUMN_UPDATE_SCHEMA_NOT_IMPLEMENTED",
+ parameters = Map("class" -> info.getClass.getName))
+ }
+
+ test("column-update: updatedColumns is empty for DELETE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"DELETE FROM $tableNameAsString WHERE dep = 'hr'")
+
+ checkLastUpdatedColumns()
+ }
+
+ test("column-update split: row-level DELETE falls back to full-width rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split", "true"))
+ checkDeleteFallsBackToFullWidthRows(Row(2, 200, "hr"))
+ }
+
+ test("column-update split: row-level DELETE by a metadata row ID falls back to full-width " +
+ "rows") {
+ // Each partition holds one row, so `_partition` identifies it.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "_partition"))
+ checkDeleteFallsBackToFullWidthRows(Row(2, 200, "software"))
+ }
+
+ test("column-update split: MERGE falls back to full-width rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split", "true"))
+ checkMergeFallsBackToFullWidthRows()
+ }
+
+ test("column-update: data correctness -- single column update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0)
+ }
+
+ test("column-update: data correctness -- update all rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 200, "hr") :: Row(2, 400, "software") :: Row(3, 600, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 3, numCopiedRows = 0)
+ }
+
+ test("column-update: data correctness -- mixed identity and real assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "engineering") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0)
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-req-attrs", reqAttrs))
+ }
+
+ test("column-update: requiredDataAttributes - data correctness") {
+ createAndInitTableWithReqAttrs("dep,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: undeclared distribution column is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "pk,salary",
+ "column-update-cluster-by", "dep"))
+
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName
+ Seq(Nil, Seq(excludeColumnPruning)).foreach { confs =>
+ withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[`dep`]"))
+ }
+ }
+ }
+
+ test("column-update split: undeclared distribution column is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-split-req-attrs", "pk,salary",
+ "column-update-cluster-by", "dep"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> splitReqAttrsOperationClass, "columns" -> "[`dep`]"))
+ }
+
+ test("column-update: undeclared write requirement columns are listed once each, sorted") {
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 1, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "pk,salary",
+ "column-update-cluster-by", "id,dep"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[`dep`, `id`]"))
+ }
+
+ test("column-update: a write requirement on a missing nested field reports why") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1 }, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "pk,s",
+ "column-update-cluster-by", "s.missing"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "columns" -> "[`s`.`missing` (FIELD_NOT_FOUND)]"))
+ }
+
+ test("column-update: an empty write requirement column name is reported as <empty>") {
+ checkError(
+ exception = QueryCompilationErrors.undeclaredWriteRequirementColumnsError(
+ reqAttrsOperationClass, Seq(Seq.empty[String] -> None)).asInstanceOf[AnalysisException],
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[<empty>]"))
+ }
+
+ test("column-update: undeclared column inside nested transforms is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "pk,salary",
+ "column-update-cluster-by", "dep",
+ "column-update-cluster-by-nested-transform", "true"))
+
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName
+ Seq(Nil, Seq(excludeColumnPruning)).foreach { confs =>
+ withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[`dep`]"))
+ }
+ }
+ }
+
+ test("column-update: distribution by an undeclared row ID column is allowed") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "salary,dep",
+ "column-update-cluster-by", "pk"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 200, "hr") :: Nil)
+ }
+
+ test("column-update: distribution by an undeclared row ID column assigned to itself is " +
+ "allowed") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "salary,dep",
+ "column-update-cluster-by", "pk"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1, pk = pk WHERE pk = 1")
+
+ checkLastUpdatedColumns("salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 200, "hr") :: Nil)
+ }
+
+ test("column-update: partition source column not in requiredDataAttributes") {
+ createAndInitTableWithReqAttrs("pk,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ executeAndCheckScanExcludes(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1", "dep")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively on the row-ID " +
+ "column without adopting the declared spelling") {
+ createAndInitTableWithReqAttrs("PK,salary,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "dep"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: undeclared, differently-cased columns in the RHS and WHERE clause " +
+ "are read with the table's own spelling") {
+ createAndInitTableWithReqAttrs("pk,dep,salary",
+ "pk INT NOT NULL, salary INT, bonus INT, extra INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "extra": 5, "dep": "hr" }
+ |""".stripMargin)
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + BONUS WHERE pk = 1 AND EXTRA > 3",
+ "pk INT, salary INT, bonus INT, extra INT, dep STRING, _partition STRING")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, 110, 10, 5, "hr") :: Nil)
+ }
+
+ test("column-update: DELETE on a scan that returns unrequested columns is unaffected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update", "true",
+ InMemoryBaseTable.SIMULATE_PARTIAL_COLUMN_PRUNING, "true"))
+
+ sql(s"DELETE FROM $tableNameAsString WHERE pk < 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(2, 200, 20, "software") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively on an assigned " +
+ "column without adopting the declared spelling") {
+ createAndInitTableWithReqAttrs("pk,SALARY,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "dep"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: case-sensitive analysis accepts columns that differ only in case") {
+ withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") {
+ createAndInitTableWithReqAttrs("pk,salary,SALARY,dep",
+ "pk INT NOT NULL, salary INT, SALARY INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "SALARY": 200, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "SALARY", "dep"))
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, -1, 200, "hr") :: Nil)
+ }
+ }
+
+ test("column-update: empty requiredDataAttributes throws AnalysisException") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }""".stripMargin,
+ java.util.Map.of("column-update-empty-req-attrs", "true"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ parameters = Map("connector" -> reqAttrsOperationClass))
+ }
+
+ test("column-update: nested requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("pk,s.c1,dep",
+ "pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map("connector" -> reqAttrsOperationClass, "nestedAttributes" -> "[`s`.`c1`]"))
+ }
+
+ test("column-update: duplicate requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("pk,pk,salary", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map("connector" -> reqAttrsOperationClass, "duplicateAttributes" -> "[`pk`]"))
+ }
+
+ test("column-update: case-only duplicate requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("PK,pk,salary", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map("connector" -> reqAttrsOperationClass, "duplicateAttributes" -> "[`PK`]"))
+ }
+
+ test("column-update: duplicate requiredDataAttributes are reported in declared order") {
+ createAndInitTableWithReqAttrs("salary,pk,SALARY,PK", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass, "duplicateAttributes" -> "[`salary`, `pk`]"))
+ }
+
+ test("column-update: metadata column in requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("pk,_partition,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET dep = 'x' WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "metadataAttributes" -> "[`_partition`]"))
+ }
+
+ test("column-update split: metadata column in requiredDataAttributes throws " +
+ "AnalysisException") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-req-attrs", "pk,_partition,dep"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET dep = 'x' WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> splitReqAttrsOperationClass,
+ "metadataAttributes" -> "[`_partition`]"))
+ }
+
+ test("column-update: requiredDataAttributes rejects a column that does not exist") {
+ createAndInitTableWithReqAttrs("pk,nonexistent_col,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "unknownAttributes" -> "[`nonexistent_col`]"))
+ }
+
+ test("column-update: column update schema excludes undeclared columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "dep", "salary"))
+ }
+
+ test("column-update: data correctness") {
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "id": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ }
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-split", "true"))
+ }
+
+ test("column-update split: column update schema has the declared columns") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update split: data correctness") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, -1, "hr") :: Nil)
+ }
+
+ test("column-update: split-updates sends a narrow UPDATE as delete and reinsert") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update", "true", "split-updates", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ deleteWriteLogEntry(id = 1, metadata = Row("hr", null)),
+ reinsertWriteLogEntry(metadata = Row("hr", null), data = Row(1, "hr", -1)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "hr") :: Nil)
+ }
+
+ test("column-update split: delete and reinsert carry the row ID and preserved metadata") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ deleteWriteLogEntry(id = 1, metadata = Row("hr", null)),
+ reinsertWriteLogEntry(metadata = Row("hr", null), data = Row(1, "hr", -1)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "hr") :: Nil)
+ }
+
+ test("column-update split: row-ID reassignment on narrow write is rejected") {
+ // `extra` is neither declared, updated, nor referenced, so requiredDataAttributes()
+ // ([pk, dep, salary]) doesn't cover every table column -- the write stays genuinely narrow
+ // and reassigning the row ID must be rejected.
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING, extra STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr", "extra": "x" }
+ |{ "pk": 2, "salary": 200, "dep": "software", "extra": "y" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT",
+ parameters = Map("connector" -> splitOperationClass, "rowIds" -> "[`pk`]"))
+ }
+
+ test("column-update split: row-ID reassignment is allowed when requiredDataAttributes " +
+ "covers every table column") {
+ createAndInitTableSplitWithReqAttrs("pk,salary,dep",
+ "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(2, 200, "software") :: Row(11, -1, "hr") :: Nil)
+ }
+
+ test("column-update split: row-ID reassignment is allowed when requiredDataAttributes " +
+ "covers every data column of a relation with metadata columns") {
+ createAndInitTableSplitWithReqAttrs("pk,salary,dep",
+ "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ // `_partition` and the coerced `salary > 0.5` make the analyzer add the metadata columns to
+ // the target relation in a later iteration. Only analyzes the plan.
+ val parsed = spark.sessionState.sqlParser.parsePlan(
+ s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 " +
+ s"WHERE dep = 'hr' AND _partition = 'hr' AND salary > 0.5")
+ val write = spark.sessionState.executePlan(parsed).analyzed.collectFirst {
+ case w: WriteDelta => w
+ }.getOrElse(fail("couldn't find WriteDelta in analyzed plan"))
+
+ assert(write.table.output.map(_.name) == Seq("pk", "salary", "dep"))
+ }
+
+ private def createAndInitTableSplitWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-split-req-attrs", reqAttrs))
+ }
+
+ test("column-update split: undeclared row-ID column on narrow write is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "pk"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED",
+ parameters = Map("connector" -> splitWithRowIdOperationClass, "rowIds" -> "[`pk`]"))
+ }
+
+ test("column-update split: metadata row ID preserved on reinsert need not be declared") {
+ // `_partition` is a required metadata attribute with PRESERVE_ON_REINSERT, so the REINSERT
+ // row carries it through the metadata projection.
+ // Each partition holds one row, so `_partition` identifies it.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "_partition"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq == Seq("salary"))
+ checkLastColumnUpdateWriteLog(
+ WriteLogEntry(operation = Delete, metadata = Some(Row("hr", null))),
+ reinsertWriteLogEntry(metadata = Row("hr", null), data = Row(-1)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update split: metadata row ID not preserved on reinsert must be declared") {
+ // `index` is a required metadata attribute, but its REINSERT value is null.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "index"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED",
+ parameters = Map("connector" -> splitWithRowIdOperationClass, "rowIds" -> "[`index`]"))
+ }
+
+ test("column-update split: metadata row ID outside requiredMetadataAttributes must be " +
+ "declared") {
+ // Without required metadata attributes, nothing projects `_partition` into the REINSERT row.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "_partition", "no-metadata", "true"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED",
+ parameters = Map("connector" -> splitWithRowIdOperationClass, "rowIds" -> "[`_partition`]"))
+ }
+
+ // ---------------------------------------------------------------------------
+ // Write relation tests: verify that only the write relation is narrowed to the
+ // declared columns, for in-place and delete+insert deltas.
+ // ---------------------------------------------------------------------------
+
+ test("column-update: write relation has the declared columns in declared order") {
Review Comment:
Done in d8a4641. All eight tests moved into `ColumnUpdateTableSharedTests`, so the no-metadata suite runs them too. The trait takes the table properties through `reqAttrsTableProps` and `fullWidthTableProps`, and the scan's metadata columns through `scanMetadataColumns`. It picks the write node name, the replacement operations and the req-attrs operation class from `deltaUpdate`.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -127,6 +128,17 @@ abstract class RowLevelOperationSuiteBase
createTable(columns, transforms)
}
+ // `extraTableProps` as the TBLPROPERTIES clause of a CREATE TABLE statement
+ protected def extraTablePropertiesClause: String = {
Review Comment:
Done in d8a4641: it sits next to `extraTableProps` and `noMetadata` now, as a pure move, and the `createTable` overloads are contiguous again.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "holdenk (via GitHub)" <gi...@apache.org>.
holdenk commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-5498457278
Hey @anuragmantri do you have the cycles to update this?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3732493322
##########
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:
I have implemented Option 1 in https://github.com/apache/spark/pull/55518/commits/cef886e16b7fb400b90841caa8975355f4c7ff42 please let me know what you think.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,47 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ *
+ * @since 4.3.0
Review Comment:
I was hoping to maybe still land this 4.3 as I requested [here](https://lists.apache.org/thread/0d1b7gg3bd7r6ml1m0fs4z0tcry23ost) since this PR has been around for a while. I don't know if we have missed the train already. I will let you and other reviewers decide. Let me know if we should move on to Spark 4.4. I will update this.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,541 @@
+/*
+ * 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, TableInfo}
+import org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity, reference}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+
+/**
+ * Tests for UPDATE statements targeting connectors that return true from
+ * [[org.apache.spark.sql.connector.write.RowLevelOperation#supportsColumnUpdates]].
Review Comment:
Updated the comments. Thanks.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +82,42 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated, copied, or reinserted record with metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must
+ * override this method when mixing in {@link SupportsColumnUpdates}; the default delegates
+ * to {@link #write(Object, Object)} so existing connectors are unaffected.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T metadata, T record) throws IOException {
+ write(metadata, record);
+ }
+
+ /**
+ * Writes one updated, copied, or reinserted record without metadata.
+ * <p>
+ * Equivalent to {@link #writeUpdate(Object, Object)} for writers that do not require metadata.
+ * The default delegates to {@link #write(Object)}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T record) throws IOException {
Review Comment:
Done. Implementation must override this or we will throw.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala:
##########
@@ -418,7 +432,29 @@ case class ReplaceData(
// validates row projection output is compatible with table attributes
private def rowAttrsResolved: Boolean = {
val inRowAttrs = DataTypeUtils.toAttributes(projections.rowProjection.schema)
- table.skipSchemaResolution || areCompatible(inRowAttrs, table.output)
+ val inUpdateAttrs = projections.updateRowProjection match {
+ case Some(projection) => DataTypeUtils.toAttributes(projection.schema)
+ case None => Nil
+ }
+ // `rowProjection` (INSERT-tagged rows) validates against `table.output` -- the full table
+ // shape. `updateRowProjection` (UPDATE/COPY-tagged rows) is narrow for column-update
+ // connectors, so it validates against `projectedDataAttrs` (the connector-declared narrow
+ // set) instead. When the connector does not mix in `SupportsColumnUpdates`,
+ // `updateRowProjection` is absent and `updateResolved` is trivially true.
+ val insertResolved = table.skipSchemaResolution || inRowAttrs.isEmpty ||
Review Comment:
Done, updated all the sites you mentioned.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedUpdateTableSuite.scala:
##########
@@ -45,11 +45,13 @@ class DeltaBasedUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase {
sql(s"SELECT * FROM $tableNameAsString"),
Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ // info.schema() reflects the row projection: `id` is non-nullable because the assignment
+ // supplies a non-null literal, matching master's projection-derived write schema semantics.
Review Comment:
Done.
##########
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)
Review Comment:
Thanks. I'm going with the smaller option (a). Rejecting the combination is contained enough for this PR and preserves narrow writes for every non-row-ID split-update case.
I added test coverage for this in `DeltaBasedColumnUpdateTableSuite`
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,47 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
Review Comment:
Added the lines on all mentioned API docs.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -267,13 +268,97 @@ abstract class RowLevelOperationSuiteBase
protected def checkLastWriteInfo(
expectedRowSchema: StructType = new StructType(),
expectedRowIdSchema: Option[StructType] = None,
- expectedMetadataSchema: Option[StructType] = None): Unit = {
+ expectedMetadataSchema: Option[StructType] = None,
+ expectedUpdateSchema: Option[StructType] = None): Unit = {
val info = table.lastWriteInfo
assert(info.schema == expectedRowSchema, "row schema must match")
val actualRowIdSchema = Option(info.rowIdSchema.orElse(null))
assert(actualRowIdSchema == expectedRowIdSchema, "row ID schema must match")
val actualMetadataSchema = Option(info.metadataSchema.orElse(null))
assert(actualMetadataSchema == expectedMetadataSchema, "metadata schema must match")
+ val actualUpdateSchema = Option(info.updateSchema.orElse(null))
+ assert(actualUpdateSchema == expectedUpdateSchema, "update schema must match")
+ }
+
+ protected def getUpdateSummary(): org.apache.spark.sql.connector.write.UpdateSummary = {
+ catalog.loadTable(ident).asInstanceOf[InMemoryTable]
+ .commits.last.writeSummary.get
+ .asInstanceOf[org.apache.spark.sql.connector.write.UpdateSummary]
+ }
+
+ /**
+ * Asserts the last UPDATE's write summary metrics. `deltaUpdate` controls the expected
+ * COPY count: MoR connectors emit only deltas (no COPY rows) so `numCopiedRows` is forced
+ * to 0; CoW connectors emit COPY rows for unchanged rows in matched groups.
+ */
+ protected def checkUpdateMetrics(
+ numUpdatedRows: Long,
+ numCopiedRows: Long,
+ deltaUpdate: Boolean = false): Unit = {
+ val summary = getUpdateSummary()
+ assert(summary.numUpdatedRows() === numUpdatedRows,
+ s"Expected numUpdatedRows=$numUpdatedRows, got ${summary.numUpdatedRows()}")
+ val expectedCopied = if (deltaUpdate) 0L else numCopiedRows
+ assert(summary.numCopiedRows() === expectedCopied,
+ s"Expected numCopiedRows=$expectedCopied, got ${summary.numCopiedRows()}")
+ }
+
+ /**
+ * Asserts that the column names in RowLevelOperationInfo.updatedColumns() received by the
+ * last operation match exactly the expected set. Order is ignored.
+ */
+ protected def checkLastUpdatedColumns(expectedNames: String*): Unit = {
+ val actual = table.lastUpdatedColumns.map(_.describe()).toSet
+ val expected = expectedNames.toSet
+ assert(actual == expected,
+ s"updatedColumns mismatch: expected ${expected.mkString("[", ", ", "]")} " +
+ s"but got ${actual.mkString("[", ", ", "]")}")
+ }
+
+ /**
+ * Asserts the set of top-level column names present in the last connector scan schema.
+ * Metadata columns are filtered out so tests can focus on data-column narrowing without
+ * having to enumerate `_partition` / `index` on every assertion.
+ */
+ protected def checkLastScanDataColumns(expectedNames: String*): Unit = {
Review Comment:
Verified there are no callers and removed this method.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3575232225
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java:
##########
@@ -105,4 +105,48 @@ default String description() {
default NamedReference[] requiredMetadataAttributes() {
return new NamedReference[0];
}
+
+ /**
+ * Controls whether to send only the required data columns to the connector rather than the
+ * full row.
+ * <p>
+ * When true, Spark narrows the data column schema ({@link LogicalWriteInfo#schema()}) to only
+ * the columns declared via {@link #requiredDataAttributes()}. Metadata columns (from
+ * {@link #requiredMetadataAttributes()}) and row ID columns (from
+ * {@link SupportsDelta#rowId()}) are unaffected and always projected separately.
+ * <p>
+ * If {@link #requiredDataAttributes()} returns a non-empty array, the write schema is exactly
+ * those columns in declared order. The connector must include all columns it wants to receive,
+ * including the columns being updated. If {@link #requiredDataAttributes()} returns an empty
+ * array, Spark sends only the non-identity assigned columns (heuristic path).
+ * <p>
+ * Currently only consulted for UPDATE operations.
+ *
+ * @since 4.3.0
+ */
+ default boolean supportsColumnUpdates() {
+ return false;
+ }
+
+ /**
+ * Returns data column references required to perform this row-level operation.
+ * <p>
+ * This method is only consulted by Spark when {@link #supportsColumnUpdates()} returns
+ * {@code true}. If {@code supportsColumnUpdates()} returns {@code false}, the returned array
+ * is ignored and the full table row is sent (the default behavior).
+ * <p>
+ * When non-empty, the returned columns become the write schema in declared order.
+ * The connector must declare all columns it wants to receive, including the columns being
+ * updated. Use {@link RowLevelOperationInfo#updatedColumns()} to learn which columns are being
+ * assigned, then add any extra columns needed for row lookup or routing (e.g., primary key).
+ * <p>
+ * When empty (the default), Spark falls back to sending only the non-identity assigned columns.
+ * <p>
+ * Currently only consulted for UPDATE operations.
+ *
+ * @since 4.3.0
+ */
+ default NamedReference[] requiredDataAttributes() {
+ return new NamedReference[0];
Review Comment:
I changed this to a mixin and added AnalysisException to prevent empty array here.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3575233851
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java:
##########
@@ -37,4 +38,20 @@ public interface RowLevelOperationInfo {
* Returns the row-level SQL command (e.g. DELETE, UPDATE, MERGE).
*/
Command command();
+
+ /**
+ * Returns the columns being updated in an UPDATE statement, as non-identity assignments.
+ *
+ * <p>For DELETE and MERGE, returns an empty array.
+ *
+ * <p>Connectors can use this to decide what {@link RowLevelOperation#requiredDataAttributes()}
+ * to declare. For instance, a connector that needs its primary key for row lookup can check
+ * whether pk is already in the updated columns list and, if not, add it to
+ * requiredDataAttributes().
+ *
+ * @since 4.3.0
+ */
+ default NamedReference[] updatedColumns() {
+ return new NamedReference[0];
Review Comment:
Added docs for nested column updates.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3785130151
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,68 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * Currently honored only for UPDATE. DELETE and MERGE ignore this interface: those operations
+ * receive full-width rows and {@link LogicalWriteInfo#updateSchema()} is absent for them.
+ *
+ * @since 4.3.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * The returned columns become the schema of updated, copied, and reinserted rows, in declared
+ * order. Implementations must include every column they want to receive (typically the columns
+ * reported by {@link RowLevelOperationInfo#updatedColumns()} plus any columns needed for row
+ * lookup or routing, e.g. a primary key).
+ * <p>
+ * If any of the columns from {@link RowLevelOperationInfo#updatedColumns()} are
+ * missing, an analysis exception is thrown.
+ * <p>
+ * For updates on nested fields such as {@code SET t.s.c1 = -1} the connector should declare the
+ * root struct column {@code s} rather than any nested field.
+ */
+ NamedReference[] requiredDataAttributes();
+
+ /**
+ * Returns additional data column references that must be present in the narrow scan but are
+ * NOT delivered to the writer.
+ * <p>
+ * Use this for columns needed only for planning -- e.g. resolving the table's partitioning
+ * expressions against the scan output, or as clustering keys returned from
+ * {@link RequiresDistributionAndOrdering#requiredDistribution()} -- where the column itself
+ * should not appear in the row passed to
+ * {@link DataWriter#writeUpdate(Object, Object)} / {@link DeltaWriter#update} /
+ * {@link DeltaWriter#reinsert}. Columns returned here are not reflected in
+ * {@link LogicalWriteInfo#updateSchema()}.
+ * <p>
+ * Must not overlap with {@link #requiredDataAttributes()}; defaults to an empty array.
+ */
+ default NamedReference[] scanOnlyDataAttributes() {
Review Comment:
Hmm.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "huaxingao (via GitHub)" <gi...@apache.org>.
huaxingao commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4395825050
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,111 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, updated, copied, and reinserted
+ * rows follow it: group-based operations receive them through
+ * {@link DataWriter#writeColumnUpdate(Object, Object)} or
+ * {@link DataWriter#writeColumnUpdate(Object)}, and operations that mix in {@link SupportsDelta}
+ * through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert}. Rows the command inserts
+ * always follow {@link LogicalWriteInfo#schema()} and arrive through
+ * {@link DataWriter#write(Object)} and {@link DeltaWriter#insert}. When
+ * {@link LogicalWriteInfo#columnUpdateSchema()} is absent, every row follows
+ * {@link LogicalWriteInfo#schema()}, as for any other operation. Connectors must decide which
+ * layout to expect by whether {@link LogicalWriteInfo#columnUpdateSchema()} is present, not by
+ * the command. A builder that wants full-width rows for some command can check
+ * {@link RowLevelOperationInfo#command()} and build an operation without this mix-in. Currently
+ * Spark narrows only UPDATE.
+ * <p>
+ * The scan builder returned by {@link #newScanBuilder} should implement
+ * {@link org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns}, so that the scan
+ * does not read columns the command neither references nor declares. Otherwise the scan reads
+ * every column. For operations that do not mix in {@link SupportsDelta}, only top-level columns
+ * are pruned: a struct column the command references through a nested field is read in full.
+ *
+ * @since 4.4.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * When Spark narrows the write, the returned columns become
+ * {@link LogicalWriteInfo#columnUpdateSchema()}, in declared order. Spark matches each declared
+ * column with the session's resolver, and {@link LogicalWriteInfo#columnUpdateSchema()} uses the
+ * table's spelling of each column, which can differ in case from the declared reference.
+ * Implementations must include every column they want to receive: every column reported by
+ * {@link RowLevelOperationInfo#updatedColumns()}, plus any columns needed for row lookup or
+ * routing. Columns that are not declared are absent from the rows the connector receives, so
+ * the connector must preserve their values itself. Spark may call this method more than once
+ * for the same operation, and every call must return the same columns.
+ * <p>
+ * If {@link RowLevelOperationInfo#updatedColumns()} is empty, the builder should build an
+ * operation without this mix-in.
Review Comment:
This sentence leaves out "for UPDATE". `updatedColumns()` is always empty for DELETE and MERGE (RowLevelOperationInfo L43-44), so read literally it tells connectors to drop the mix-in for every DELETE and MERGE, which the test fixture does not do (it drops it only for `command == UPDATE && updatedColumns().isEmpty`). It would also help to say why, since keeping the mix-in still works. Maybe replace lines 65-66 with:
```java
* For an UPDATE whose assignments are all identity assignments,
* {@link RowLevelOperationInfo#updatedColumns()} is empty and there is nothing to narrow, so the
* builder may build an operation without this mix-in to receive full-width rows.
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub) via reviews" <re...@spark.apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4429023441
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,114 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, updated, copied, and reinserted
+ * rows follow it: group-based operations receive them through
+ * {@link DataWriter#writeColumnUpdate(Object, Object)} or
+ * {@link DataWriter#writeColumnUpdate(Object)}, and operations that mix in {@link SupportsDelta}
+ * through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert}. Rows the command inserts
+ * always follow {@link LogicalWriteInfo#schema()} and arrive through
+ * {@link DataWriter#write(Object)} and {@link DeltaWriter#insert}. When
+ * {@link LogicalWriteInfo#columnUpdateSchema()} is absent, every row follows
+ * {@link LogicalWriteInfo#schema()}, as for any other operation. Connectors must decide which
+ * layout to expect by whether {@link LogicalWriteInfo#columnUpdateSchema()} is present, not by
+ * the command. A builder that wants full-width rows for some command can check
+ * {@link RowLevelOperationInfo#command()} and build an operation without this mix-in. Currently
+ * Spark narrows only UPDATE.
+ * <p>
+ * The scan builder returned by {@link #newScanBuilder} should implement
+ * {@link org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns}, so that the scan
+ * does not read columns the command neither references nor declares. Otherwise the scan reads
+ * every column. For operations that do not mix in {@link SupportsDelta}, only top-level columns
+ * are pruned: a struct column the command references through a nested field is read in full.
+ *
+ * @since 4.4.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * When Spark narrows the write, the returned columns become
+ * {@link LogicalWriteInfo#columnUpdateSchema()}, in declared order. Spark matches each declared
+ * column with the session's resolver, and {@link LogicalWriteInfo#columnUpdateSchema()} uses the
+ * table's spelling of each column, which can differ in case from the declared reference.
+ * Implementations must include every column they want to receive: every column reported by
+ * {@link RowLevelOperationInfo#updatedColumns()}, plus any columns needed for row lookup or
+ * routing. Columns that are not declared are absent from the rows the connector receives, so
+ * the connector must preserve their values itself. Spark may call this method more than once
+ * for the same operation, and every call must return the same columns.
+ * <p>
+ * If the returned array is empty, Spark does not narrow the write and the operation receives
Review Comment:
Probably reasonable after a closer look.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub) via reviews" <re...@spark.apache.org>.
aokolnychyi commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-6534378097
Taking another look soon too, sorry for the delay.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3148764907
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -50,20 +51,35 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
protected def buildOperationTable(
table: SupportsRowLevelOperations,
command: Command,
- options: CaseInsensitiveStringMap): RowLevelOperationTable = {
- val info = RowLevelOperationInfoImpl(command, options)
+ options: CaseInsensitiveStringMap,
+ updatedColumns: Seq[NamedReference] = Nil): RowLevelOperationTable = {
+ val info = RowLevelOperationInfoImpl(command, options, updatedColumns)
val operation = table.newRowLevelOperationBuilder(info).build()
RowLevelOperationTable(table, operation)
}
+ // Builds a DataSourceV2Relation for a row-level operation, optionally narrowing its output.
+ //
+ // When dataAttrs is non-empty, the relation output is narrowed to include only columns
+ // required for a column-update write. When dataAttrs is empty, the full relation.output is
+ // preserved.
protected def buildRelationWithAttrs(
relation: DataSourceV2Relation,
table: RowLevelOperationTable,
metadataAttrs: Seq[AttributeReference],
- rowIdAttrs: Seq[AttributeReference] = Nil): DataSourceV2Relation = {
-
- val attrs = dedupAttrs(relation.output ++ rowIdAttrs ++ metadataAttrs)
- relation.copy(table = table, output = attrs)
+ rowIdAttrs: Seq[AttributeReference] = Nil,
+ dataAttrs: Seq[AttributeReference] = Nil,
+ cond: Expression = TrueLiteral): DataSourceV2Relation = {
+
+ if (dataAttrs.nonEmpty) {
+ val required =
+ AttributeSet(dataAttrs) ++ AttributeSet(Seq(cond)) ++ AttributeSet(rowIdAttrs)
+ val narrowOutput = relation.output.filter(required.contains)
+ relation.copy(table = table, output = dedupAttrs(narrowOutput ++ rowIdAttrs ++ metadataAttrs))
Review Comment:
Can an attribute in `required` be missing from `relation.output`?
`rowIdAttrs` seems to be added 2 times.
If we already have a `dedupAttrs()` then probably doesn't make sense build `AttributeSet`s.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192428127
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java:
##########
@@ -105,4 +105,45 @@ default String description() {
default NamedReference[] requiredMetadataAttributes() {
return new NamedReference[0];
}
+
+
+ /**
+ * Controls whether to send only the required data columns to the connector rather than the
+ * full row.
+ * <p>
+ * When true, Spark narrows the data column schema ({@link LogicalWriteInfo#schema()}) to only
+ * the columns declared via {@link #requiredDataAttributes()}. Metadata columns (from
+ * {@link #requiredMetadataAttributes()}) and row ID columns (from
+ * {@link SupportsDelta#rowId()}) are unaffected and always projected separately.
+ * <p>
+ * If {@link #requiredDataAttributes()} returns a non-empty array, the write schema is exactly
+ * those columns in declared order. The connector must include all columns it wants to receive,
+ * including the columns being updated. If {@link #requiredDataAttributes()} returns an empty
+ * array, Spark sends only the non-identity assigned columns (heuristic path).
+ *
+ * @since 4.2.0
+ */
+ default boolean supportsColumnUpdates() {
+ return false;
+ }
+
+ /**
+ * Returns data column references required to perform this row-level operation.
+ * <p>
+ * This method is only consulted by Spark when {@link #supportsColumnUpdates()} returns
+ * {@code true}. If {@code supportsColumnUpdates()} returns {@code false}, the returned array
+ * is ignored and the full table row is sent (the default behavior).
+ * <p>
+ * When non-empty, the returned columns become the write schema in declared order.
+ * The connector must declare all columns it wants to receive, including the columns being
+ * updated. Use {@link RowLevelOperationInfo#updatedColumns()} to learn which columns are being
+ * assigned, then add any extra columns needed for row lookup or routing (e.g., primary key).
+ * <p>
+ * When empty (the default), Spark falls back to sending only the non-identity assigned columns.
+ *
+ * @since 4.2.0
+ */
+ default NamedReference[] requiredDataAttributes() {
Review Comment:
Given the scope of this PR, shall we mention that `DELETE and MERGE ignores this method`?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192381665
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java:
##########
@@ -105,4 +105,45 @@ default String description() {
default NamedReference[] requiredMetadataAttributes() {
return new NamedReference[0];
}
+
+
+ /**
+ * Controls whether to send only the required data columns to the connector rather than the
+ * full row.
+ * <p>
+ * When true, Spark narrows the data column schema ({@link LogicalWriteInfo#schema()}) to only
+ * the columns declared via {@link #requiredDataAttributes()}. Metadata columns (from
+ * {@link #requiredMetadataAttributes()}) and row ID columns (from
+ * {@link SupportsDelta#rowId()}) are unaffected and always projected separately.
+ * <p>
+ * If {@link #requiredDataAttributes()} returns a non-empty array, the write schema is exactly
+ * those columns in declared order. The connector must include all columns it wants to receive,
+ * including the columns being updated. If {@link #requiredDataAttributes()} returns an empty
+ * array, Spark sends only the non-identity assigned columns (heuristic path).
+ *
+ * @since 4.2.0
+ */
+ default boolean supportsColumnUpdates() {
+ return false;
+ }
+
+ /**
+ * Returns data column references required to perform this row-level operation.
+ * <p>
+ * This method is only consulted by Spark when {@link #supportsColumnUpdates()} returns
+ * {@code true}. If {@code supportsColumnUpdates()} returns {@code false}, the returned array
+ * is ignored and the full table row is sent (the default behavior).
+ * <p>
+ * When non-empty, the returned columns become the write schema in declared order.
+ * The connector must declare all columns it wants to receive, including the columns being
Review Comment:
This is very strong assumption, but it seems that this PR didn't have a protection. May I ask if we have some kind of assertion or a test coverage for this?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192368701
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala:
##########
@@ -470,7 +506,28 @@ case class WriteDelta(
case Some(projection) => DataTypeUtils.toAttributes(projection.schema)
case None => Nil
}
- table.skipSchemaResolution || areCompatible(inRowAttrs, outRowAttrs)
+ table.skipSchemaResolution ||
+ areCompatible(inRowAttrs, outRowAttrs) ||
+ dataAttrsResolved(inRowAttrs)
Review Comment:
nit. Please minimize the change of existing code as much as possible like the following.
```
table.skipSchemaResolution || areCompatible(inRowAttrs, outRowAttrs) ||
dataAttrsResolved(inRowAttrs)
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3779990365
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +84,51 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated, copied, or reinserted record with metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must override this
+ * method when mixing in {@link SupportsColumnUpdates}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but does not override this method.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T metadata, T record) throws IOException {
Review Comment:
I am still thinking about the best possible name for this method. What worries me a bit is that it will also pass copied or re-inserted rows. This is different from DeltaWriter where we have these different methods. I am not sure I can offer a better name, though.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3785101803
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/LogicalWriteInfo.java:
##########
@@ -65,4 +65,14 @@ default Optional<StructType> metadataSchema() {
throw new SparkUnsupportedOperationException(
"DATA_SOURCE_METADATA_SCHEMA_NOT_IMPLEMENTED", Map.of("class", getClass().getName()));
}
+
+ /**
+ * the narrow schema for updates. Present only when the connector mixes in
+ * {@link SupportsColumnUpdates}.
+ *
+ * @since 4.3.0
+ */
+ default Optional<StructType> updateSchema() {
Review Comment:
This could technically also become `columnUpdateSchema()`, if we decide to do the rename. It does indicate the schema of copied rows in CoW and reinserted rows. One thing I don't like is that it is longer.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3785330664
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,68 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * Currently honored only for UPDATE. DELETE and MERGE ignore this interface: those operations
+ * receive full-width rows and {@link LogicalWriteInfo#updateSchema()} is absent for them.
+ *
+ * @since 4.3.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * The returned columns become the schema of updated, copied, and reinserted rows, in declared
+ * order. Implementations must include every column they want to receive (typically the columns
+ * reported by {@link RowLevelOperationInfo#updatedColumns()} plus any columns needed for row
+ * lookup or routing, e.g. a primary key).
+ * <p>
+ * If any of the columns from {@link RowLevelOperationInfo#updatedColumns()} are
+ * missing, an analysis exception is thrown.
+ * <p>
+ * For updates on nested fields such as {@code SET t.s.c1 = -1} the connector should declare the
+ * root struct column {@code s} rather than any nested field.
+ */
+ NamedReference[] requiredDataAttributes();
+
+ /**
+ * Returns additional data column references that must be present in the narrow scan but are
+ * NOT delivered to the writer.
+ * <p>
+ * Use this for columns needed only for planning -- e.g. resolving the table's partitioning
+ * expressions against the scan output, or as clustering keys returned from
+ * {@link RequiresDistributionAndOrdering#requiredDistribution()} -- where the column itself
+ * should not appear in the row passed to
+ * {@link DataWriter#writeUpdate(Object, Object)} / {@link DeltaWriter#update} /
+ * {@link DeltaWriter#reinsert}. Columns returned here are not reflected in
+ * {@link LogicalWriteInfo#updateSchema()}.
+ * <p>
+ * Must not overlap with {@link #requiredDataAttributes()}; defaults to an empty array.
+ */
+ default NamedReference[] scanOnlyDataAttributes() {
Review Comment:
I am not sure adding this part is justified. If we need these columns for clustering / ordering as the doc says, this means they have to stay with the plan till the very end. If so, I would prefer to simply include all such data columns in `requiredDataColumns` and do a projection inside the connector, if needed. Iceberg has all utilities for that, for example.
At this point, I am reluctant to add the complexity to Spark for something we are not sure would be needed. It does complicate the API and also the Spark logic that is already very complicated.
We can always reconsider and add this API later.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "peter-toth (via GitHub)" <gi...@apache.org>.
peter-toth commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4053024876
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -228,4 +371,146 @@ 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, validating that every assigned column is covered. Returns `Nil` otherwise.
+ */
+ private def resolveConnectorDataAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation,
+ assignments: Seq[Assignment]): Seq[AttributeReference] = operation match {
+ case scu: SupportsColumnUpdates =>
+ val attrs = resolveRequiredDataAttrs(relation, scu)
+ validateUpdatedColumnsSubset(scu, assignments, attrs)
+ attrs
+ case _ => Nil
+ }
+
+ /**
+ * Computes the narrow set of data columns that must be present in the scan for a column-update
+ * write: connector-declared attrs (`requiredDataAttributes()`), unioned with any table columns
+ * referenced by non-identity assignment RHS expressions and the operation condition.
+ */
+ 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)
Review Comment:
**Finding 27.** A table carrying a `CHECK` constraint cannot be updated through this path at all.
`ResolveTableConstraints` wraps the rewritten query in `Filter(CheckInvariant(...), query)` and resolves the constraint's column references against `query.output` (`ResolveTableConstraints.scala:60-61`). `RowLevelOperationTable.constraints()` delegates to the real table, so the rule fires for the operation table too. Narrowing takes those columns out of both the scan and the write query, and nothing here accounts for it.
Measured on `72ec505`. Same table in every arm, `pk INT NOT NULL, id INT, dep STRING, extra INT` partitioned by `identity(dep)`, same constraint `ALTER TABLE ... ADD CONSTRAINT positive_extra CHECK (extra > 0)`, same statement `UPDATE t SET id = -1 WHERE pk = 1`:
| connector | result |
|---|---|
| wide `supports-deltas` (MoR) | works |
| wide default fixture (CoW) | works |
| `column-update` (MoR) | ``[UNRESOLVED_COLUMN.WITH_SUGGESTION] A column, variable, or function parameter with name `extra` cannot be resolved. Did you mean one of the following? [`dep`, `id`, `pk`, `index`, `_partition`]`` |
| `column-update-cow` (CoW) | same |
So it is a narrowing regression, not pre-existing. The error names a column the statement never mentions and says nothing about `requiredDataAttributes()`.
The second arm is worse. When the constraint's column *is* in the narrow scan, because the condition or an assignment RHS references it, the delta path still drops it from the write query in the `Project` below (`:320-323`), and the Resolution batch never converges:
| connector | statement | result |
|---|---|---|
| `column-update` (MoR) | `SET id = -1 WHERE extra > 3` | `RuntimeException: Max iterations (100) reached for batch Resolution` |
| `column-update` (MoR) | `SET id = extra + 1 WHERE pk = 1` | same |
| `column-update-cow` (CoW) | `SET id = -1 WHERE extra > 3` | works |
The CoW row is the control for that pair: `buildNarrowReplaceDataUpdateProjection` maps all of `plan.output`, so the column survives into `query.output` there, and only the delta builder drops it.
I applied the matching fix locally and it holds. Carrying the remaining narrow scan columns through the `Project`:
```scala
val connectorIds = connectorDataAttrs.map(_.exprId).toSet
val carryAlong = plan.output.filterNot { a =>
MetadataAttribute.isValid(a.metadata) || rowIdAttrSet.contains(a) ||
assignedKeyIds.contains(a.exprId) || connectorIds.contains(a.exprId)
}
Project(
Seq(operationType) ++ assignedValues ++ connectorPassThroughValues ++
metadataValues ++ rowIdValues ++ originalRowIdValues ++ carryAlong,
plan)
```
turns both looping arms green and keeps the payload narrow: `updateSchema` stays `struct<pk, dep, id>` and `extra` does not leak into it, because `updateRowProjection` picks `connectorDataAttrs` out of the `Project` by name. `DeltaBasedColumnUpdateTableSuite` + `GroupBasedColumnUpdateTableSuite` stay 43/43. And the constraint is genuinely enforced again afterwards: with `CHECK (id > 0)` instead, `SET id = -1 WHERE extra > 3` gives `CHECK_CONSTRAINT_VIOLATION`.
That still leaves the first arm, where the constraint column is not in the narrow scan at all. Reading it is the only way Spark can re-validate the constraint on the new row, so `computeNarrowReadAttrs` would have to include the `Check` constraints' references as well. The alternative is to refuse to narrow a table that carries `Check` constraints, with a clear error. I would widen rather than refuse, and note in the javadoc that a table with CHECK constraints narrows less.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,57 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * Currently honored only for UPDATE. DELETE and MERGE ignore this interface: those operations
+ * receive full-width rows and {@link LogicalWriteInfo#updateSchema()} is absent for them.
+ *
+ * @since 4.3.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * The returned columns become the schema of updated, copied, and reinserted rows, in declared
+ * order. Implementations must include every column they want to receive (typically the columns
+ * reported by {@link RowLevelOperationInfo#updatedColumns()} plus any columns needed for row
+ * lookup or routing, e.g. a primary key).
+ * <p>
+ * If any of the columns from {@link RowLevelOperationInfo#updatedColumns()} are
+ * missing, an analysis exception is thrown.
+ * <p>
+ * For updates on nested fields such as {@code SET t.s.c1 = -1} the connector should declare the
+ * root struct column {@code s} rather than any nested field.
+ * <p>
+ * This also covers columns needed only for planning, e.g. resolving the table's partitioning
Review Comment:
**Finding 17 (regressed).** Dropping `validatePartitionAttrsDeclared` is a fair policy call, and I agree Spark should not decide what a connector wants. But the thing that breaks when a partition source column is missing is not the connector's write requirement. It is Spark's own SPJ reporting rule, and it throws instead of degrading.
This paragraph reads as advisory, "also covers columns needed only for planning". For any partitioned table whose scan reports `KeyGroupedPartitioning` it is mandatory. Measured on `72ec505`, `requiredDataAttributes() = [pk, id]`, table `pk INT NOT NULL, id INT, dep STRING` partitioned by `identity(dep)`, `UPDATE t SET id = -1 WHERE pk = 1`:
```
org.apache.spark.sql.AnalysisException: Unable to resolve dep given [pk,id,_partition,index].
at QueryCompilationErrors$.cannotResolveAttributeError(QueryCompilationErrors.scala:1988)
at V2ExpressionUtils$.resolveRef(V2ExpressionUtils.scala:56)
at V2ExpressionUtils$.toCatalystTransformOpt(V2ExpressionUtils.scala:147)
at V2ExpressionUtils$.toCatalystOpt(V2ExpressionUtils.scala:128)
at V2ScanPartitioningAndOrdering$$anonfun$partitioning$2 ... (V2ScanPartitioningAndOrdering.scala:51)
```
The same connector against an unpartitioned table passes, so the partitioning is the trigger.
This also corrects what I wrote at round 5 ([r3735139324](https://github.com/apache/spark/pull/55518#discussion_r3735139324)), where I said the read side would "silently report no partitioning". It does not. `V2ScanPartitioningAndOrdering.partitioning` resolves the keys against `scanRelation.relation` (`ExtractV2ScanInfo`, `DataSourceV2Relation.scala:457-461`), and that relation is exactly what this PR narrows. `toCatalystTransformOpt`'s `IdentityTransform` case then calls the throwing `resolveRef`. Both the `Opt` naming and the comment two lines below the call site ("Keep the partitioning when at least one of its keys is still in the scan output") say the rule means to tolerate a pruned key. It never gets the chance. The write-side site I measured at round 4 ([r3735141175](https://github.com/apache/spark/pull/55518#discussion_r3735141175)) is still reachable too, for a connector that clusters on a data column.
Either fix closes the read side:
- resolve the keys with a non-throwing lookup and drop the partitioning when a key is absent. This asks nothing of connectors and is what the rule already intends;
- or have `computeNarrowReadAttrs` add the table's partition-transform references to the *scan*. This is not restrictive the way the removed guard was: the payload is projected down to `connectorDataAttrs` by name, so a partition column sitting in the scan never reaches the writer. I confirmed that projection behaviour while measuring finding 27, where a carried-along column stayed out of `updateSchema()`.
Whichever you take, this sentence should say the declaration is required rather than that it "covers" the case, and there should be a test for a column-update UPDATE on a partitioned table whose partition column is not declared. Every fixture in the suite declares `dep`, which is why 43 tests pass over this.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +84,51 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated, copied, or reinserted record with metadata.
+ * <p>
+ * Connectors that mix in {@link SupportsColumnUpdates} receive records here in the schema
+ * declared by {@link LogicalWriteInfo#updateSchema()}. Implementations must override this
+ * method when mixing in {@link SupportsColumnUpdates}.
+ * <p>
+ * If this method fails (by throwing an exception), {@link #abort()} will be called and this
+ * data writer is considered to have been failed.
+ *
+ * @throws IOException if failure happens during disk/network IO like writing files.
+ * @throws SparkUnsupportedOperationException if the connector mixes in
+ * {@link SupportsColumnUpdates} but does not override this method.
+ *
+ * @since 4.3.0
+ */
+ default void writeUpdate(T metadata, T record) throws IOException {
+ throw new SparkUnsupportedOperationException(
+ "DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED",
+ Map.of("class", getClass().getName()));
+ }
Review Comment:
**Finding 29.** A connector that implements only `writeUpdate(record)` still fails here.
Which overload Spark calls depends on whether the operation declares required metadata attributes (`ReplaceDataExec.writingTask`, `WriteToDataSourceV2Exec.scala:385-397`), which is not visible from the interface. The single-argument form's doc says it is "Equivalent to `writeUpdate(Object, Object)` for writers that do not require metadata", so implementing that one alone reads as sufficient, and it is not.
The sibling pair a few lines up already solves this: `write(T metadata, T record)` defaults to `write(record)` (`:83-85`). Mirroring it keeps the two channels consistent:
```suggestion
writeUpdate(record);
}
```
A connector that overrides neither form still gets `DATA_SOURCE_WRITE_UPDATE_NOT_IMPLEMENTED`, from the single-argument default. The trade-off is the one `write` already makes: a writer that needs the metadata must override the two-argument form.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,684 @@
+/*
+ * 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, TableInfo}
+import org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity, reference}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, Spark narrows the update-row projection
+ * (LogicalWriteInfo.updateSchema()) to contain only the declared columns rather than
+ * the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends RowLevelOperationSuiteBase {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ test("column-update: rowSchema contains only the single assigned column") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ // info.schema() is empty: UPDATE-only column-update writes have no INSERT-shaped rows,
+ // so nothing flows through the writer's write() path. The narrow update-row layout is
+ // carried by info.updateSchema().
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: rowSchema contains multiple assigned columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false),
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: rowSchema is empty for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ // All assignments are identity, so updatedColumns is empty -- the connector still declares
+ // pk (row lookup) and dep (write-side clustering key), so the narrow update schema is
+ // just [pk, dep].
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Array(PK_FIELD, DEP_FIELD))))
+ }
+
+ test("column-update: row filter condition is orthogonal to column narrowing") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = 'engineering' WHERE pk IN (1, 3)")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: update all rows (no WHERE clause)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("salary", IntegerType, nullable = true)
+ ))))
+ }
+
+ test("column-update: rowSchema excludes identity assignments in a mixed UPDATE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: cross-column assignment is not treated as identity") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = dep, id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: nested struct field update narrows to the root struct column") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("s"),
+ s"expected [s] in updatedColumns (root struct) but got: $updatedNames")
+
+ // info.updateSchema() carries the narrow row layout; info.schema() is the full table.
+ // `dep` is present because the connector also declares it as the write-side clustering key.
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("s"),
+ s"s must be in update schema: $updateSchema")
+ assert(updateSchema.fieldNames.contains("dep"),
+ s"dep must be in update schema: $updateSchema")
+ }
+
+ test("column-update: nested field identity update reports root struct as updated") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = s.c1 WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("s"),
+ s"nested identity is reported as an update at root-column granularity: $updatedNames")
+
+ // Data correctness: the struct is rewritten but with equal values, so rows are unchanged.
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, Row(1, 2), "hr") :: Nil)
+ }
+
+ test("column-update: updatedColumns contains non-identity assigned columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'eng' WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("id", "dep"),
+ s"expected [id, dep] in updatedColumns but got: $updatedNames")
+ }
+
+ test("column-update: updatedColumns excludes identity assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = dep WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("id"),
+ s"expected only [id] in updatedColumns but got: $updatedNames")
+ }
+
+ test("column-update: updatedColumns is empty for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ assert(table.lastUpdatedColumns.isEmpty,
+ s"expected empty updatedColumns but got: ${table.lastUpdatedColumns.mkString(", ")}")
+ }
+
+ test("column-update: updatedColumns is empty for DELETE (Javadoc contract)") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"DELETE FROM $tableNameAsString WHERE dep = 'hr'")
+
+ assert(table.lastUpdatedColumns.isEmpty,
+ s"DELETE must pass empty updatedColumns but got: ${table.lastUpdatedColumns.mkString(", ")}")
+ }
+
+ test("column-update: data correctness -- single column update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ // 1 row matched WHERE pk = 1 -> 1 update. MoR emits no COPY rows.
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0, deltaUpdate = true)
+ }
+
+ test("column-update: data correctness -- update all rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 200, "hr") :: Row(2, 400, "software") :: Row(3, 600, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 3, numCopiedRows = 0, deltaUpdate = true)
+ }
+
+ test("column-update: data correctness -- mixed identity and real assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "engineering") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0, deltaUpdate = true)
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-req-attrs", reqAttrs)
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update: requiredDataAttributes - data correctness") {
+ createAndInitTableWithReqAttrs("dep,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively on the row-ID " +
+ "column without adopting the declared spelling") {
+ // The connector declares `PK` (not `pk`); resolution must match case-insensitively but the
+ // narrow scan/write must still use the table's own column name, `pk`, not the connector's
+ // declared spelling -- otherwise column pruning can't find `PK` in the physical scan.
+ createAndInitTableWithReqAttrs("PK,salary,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("pk"),
+ s"update schema must use the table's own spelling `pk`, not `PK`: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("PK"),
+ s"update schema must not adopt the connector's declared spelling `PK`: $updateSchema")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively on an assigned " +
+ "column without adopting the declared spelling") {
+ // Same as above, but the case mismatch is on `salary` (the assigned column) rather than
+ // the row-ID column.
+ createAndInitTableWithReqAttrs("pk,SALARY,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("salary"),
+ s"update schema must use the table's own spelling `salary`, not `SALARY`: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("SALARY"),
+ s"update schema must not adopt the connector's declared spelling `SALARY`: $updateSchema")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: empty requiredDataAttributes throws AnalysisException") {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-empty-req-attrs", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(
+ StructType.fromDDL("pk INT NOT NULL, id INT, dep STRING"))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }""".stripMargin)
+
+ // A connector that mixes in SupportsColumnUpdates but returns an empty
+ // requiredDataAttributes() violates the mix-in contract -- analysis must reject the
+ // operation rather than silently falling through to the wide write path.
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ }
+ assert(ex.getCondition == "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ s"expected COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES but got: ${ex.getCondition}")
+ }
+
+ test("column-update: requiredDataAttributes throws AnalysisException for invalid column") {
+ createAndInitTableWithReqAttrs("nonexistent_col", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ }
+ assert(ex.getMessage.contains("nonexistent_col"),
+ s"Expected error about unresolvable column but got: ${ex.getMessage}")
+ }
+
+ private def createAndInitTableFromInfo(schemaString: String, jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-from-info", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update from-info: write schema is updatedColumns + pk/dep pass-through") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("salary"),
+ s"salary must be in update schema: $updateSchema")
+ assert(updateSchema.fieldNames.contains("pk"),
+ s"pk must be in update schema: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("id"),
+ s"id must not be in update schema: $updateSchema")
+ assert(updateSchema.fieldNames.contains("dep"),
+ s"dep must be in update schema: $updateSchema")
+ }
+
+ test("column-update from-info: pk already in updatedColumns is not duplicated") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ val pkCount = updateSchema.fieldNames.count(_ == "pk")
+ assert(pkCount == 1, s"pk must appear exactly once in update schema: $updateSchema")
+ }
+
+ test("column-update from-info: data correctness") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "id": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ }
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-split", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update split: write schema is narrow (assigned + pk/dep pass-through)") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ // For representUpdateAsDeleteAndInsert connectors, reinsert-tagged rows still route through
+ // writeUpdate() (info.updateSchema() = narrow), not write(). info.schema() stays empty.
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update split: data correctness") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, -1, "hr") :: Nil)
+ }
+
+ test("column-update split: row-ID reassignment on narrow write is rejected") {
+ // `extra` is neither declared, updated, nor referenced, so requiredDataAttributes()
+ // ([pk, dep, salary]) doesn't cover every table column -- the write stays genuinely narrow
+ // and reassigning the row ID must be rejected.
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING, extra STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr", "extra": "x" }
+ |{ "pk": 2, "salary": 200, "dep": "software", "extra": "y" }
+ |""".stripMargin)
+
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+ }
+ assert(ex.getCondition == "COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT",
+ s"expected COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT but got: ${ex.getCondition}")
+ }
+
+ test("column-update split: row-ID reassignment is allowed when requiredDataAttributes " +
+ "covers every table column") {
+ // Finding 22: the SPLIT_UPDATE_ROW_ID_REASSIGNMENT message advertises a second remedy --
+ // "include every table column in requiredDataAttributes()" -- but the guard never checked
+ // for it. When the declaration covers the whole relation, the REINSERT payload already is
+ // the full row with the new row-ID value, and the DELETE half still carries the original
+ // row-ID via newLazyRowIdProjection, so reassignment is safe and must be allowed.
+ createAndInitTableSplitWithReqAttrs("pk,salary,dep",
+ "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(2, 200, "software") :: Row(11, -1, "hr") :: Nil)
+ }
+
+ private def createAndInitTableSplitWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-split-req-attrs", reqAttrs)
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update split: undeclared row-ID column on narrow write is rejected") {
+ // The connector's requiredDataAttributes() excludes `pk` (the row ID). The REINSERT
+ // payload for the split-update path is projected down to requiredDataAttributes(), so
+ // without `pk` there the reinserted row would have no identity for the connector to
+ // place it by. This must be rejected at analysis time, not discovered at write time.
+ createAndInitTableSplitMissingRowId("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+ }
+ assert(ex.getCondition == "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED",
+ s"expected COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED but got: ${ex.getCondition}")
+ assert(ex.getMessage.contains("pk"),
+ s"error message must name the undeclared row ID column `pk`: ${ex.getMessage}")
+ }
+
+ private def createAndInitTableSplitMissingRowId(
+ schemaString: String, jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-split-missing-row-id", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ // ---------------------------------------------------------------------------
+ // Scan-narrowing tests: verify that when the connector opts into column updates,
+ // the physical scan reads only the columns actually needed by the operation.
+ //
+ // Assertions use `checkLastScanExcludes` (proves narrowing) and `checkLastScanIncludes`
+ // (proves transparent widening) rather than exact-set matches -- ColumnPruning is
+ // free to further tighten the scan beyond our analysis-time narrow set (e.g. when an
+ // assignment RHS is a literal, the assigned column doesn't need to be read).
+ // ---------------------------------------------------------------------------
+
+ test("column-update: scan excludes columns outside connector-declared + cond + RHS refs") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, extra STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr", "extra": "x" }
+ |{ "pk": 2, "id": 2, "dep": "software", "extra": "y" }
+ |""".stripMargin)
+
+ // requiredDataAttributes = [pk, id]; cond refs [pk] only; RHS is a literal.
+ // `dep` is scan-only declared (needed for scan-side partitioning resolution) so it must
Review Comment:
**Finding 30.** Wording left over from the removed `scanOnlyDataAttributes()`. There is no scan-only category any more: `dep` is in `requiredDataAttributes()` like every other declared column.
Here, and again at `:654-657` ("`dep` is scan-only declared (needed for scan-side partitioning resolution) so it also appears; `extra` is neither declared, scan-only, nor referenced"). Also `InMemoryRowLevelOperationTable.scala:303`, which calls `dep` "the write-side clustering key, see `clusterColumnRef`" while `clusterColumnRef` is `PARTITION_COLUMN_REF` (`_partition`), a metadata column, not `dep`.
Two test names in this file no longer match what they assert, same cause. "rowSchema contains only the single assigned column" (`:42`) asserts `[pk, dep, id]`, and "rowSchema is empty for a full identity update" (`:82`) asserts `[pk, dep]`. "rowSchema" is also the pre-redesign name for what is now `updateSchema()`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3786556681
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, any required metadata attrs and optionally connector
+ // declared attrs)
val rowAttrs = relation.output
+ val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates]
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveConnectorDataAttrs(relation, operation)
+ } else Nil
+ val scanOnlyDataAttrs = if (supportsColumnUpdate) {
+ resolveScanOnlyDataAttrs(relation, operation)
+ } else Nil
+
+ if (supportsColumnUpdate) {
+ validateUpdatedColumnsSubset(operation, assignments, connectorDataAttrs)
Review Comment:
Why not validate `validateUpdatedColumnsSubset` when we resolve the attributes?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3787781513
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, any required metadata attrs and optionally connector
+ // declared attrs)
val rowAttrs = relation.output
+ val supportsColumnUpdate = operation.isInstanceOf[SupportsColumnUpdates]
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveConnectorDataAttrs(relation, operation)
+ } else Nil
+ val scanOnlyDataAttrs = if (supportsColumnUpdate) {
+ resolveScanOnlyDataAttrs(relation, operation)
+ } else Nil
+
+ if (supportsColumnUpdate) {
+ validateUpdatedColumnsSubset(operation, assignments, connectorDataAttrs)
+ validateNoOverlap(operation, connectorDataAttrs, scanOnlyDataAttrs)
+ validatePartitionAttrsDeclared(operation, relation, connectorDataAttrs, scanOnlyDataAttrs)
+ }
+
+
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
- val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ if (supportsColumnUpdate && operation.representUpdateAsDeleteAndInsert) {
+ validateNoRowIdReassignment(operation, assignments, rowIdAttrs)
+ validateRowIdDeclared(operation, connectorDataAttrs, rowIdAttrs)
+ }
+
+ val narrowDataAttrs = if (supportsColumnUpdate) {
+ computeNarrowReadAttrs(relation, connectorDataAttrs, scanOnlyDataAttrs, assignments, cond)
+ } else {
+ relation.output
+ }
+
+ val readRelation = if (supportsColumnUpdate) {
+ buildNarrowRelationWithAttrs(relation, operationTable, narrowDataAttrs, metadataAttrs,
+ rowIdAttrs)
+ } else {
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ }
// 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 rowDeltaPlan = if (supportsColumnUpdate) {
Review Comment:
I am worried about these constant if statements. We will need to find a better structure.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4406298765
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -343,3 +353,63 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
}
}
}
+
+object RewriteRowLevelCommand extends SQLConfHelper {
+
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation, in declared order and with the relation's spelling. An empty array,
+ * nested fields, duplicate columns, columns that do not exist in the relation (including a
+ * reference with no name parts) and metadata columns are rejected with an `AnalysisException`.
+ */
+ private[sql] def resolveRequiredDataAttrs(
+ relation: LogicalPlan,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes.toImmutableArraySeq
+ if (refs.isEmpty) {
Review Comment:
**55. An operation whose declaration is exactly `updatedColumns()` fails an UPDATE whose assignments are all identity assignments, and the new fixture arms do not guard against it.** This is my answer to your question in the [latest summary](https://github.com/apache/spark/pull/55518#pullrequestreview-5672546710). `collectUpdatedAttrs` skips identity assignments (`RewriteUpdateTable.scala:252-256`), so `UPDATE t SET salary = salary WHERE pk = 1` reports no updated columns. An operation that declares exactly `info.updatedColumns()` then fails here at analysis with `COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES`, while the same statement works without the mix-in. That declaration is what `SupportsColumnUpdates.java:59-61` asks for when no lookup or routing column is needed, e.g. for a delta operation whose row ID arrives through the row ID projection, and the suite's own `DeltaBasedColumnUpdateSplitOperationWithRowId` (`InMemoryRowLevelOperationTable.scala:623-631`) is one. It keeps
the mix-in for UPDATE (`InMemoryRowLevelOperationTable.scala:204-207`).
8b1771ad6a7 answers this with a builder-side "should" (`SupportsColumnUpdates.java:65-66`) and with `updatesNothing` arms in the default fixtures (`InMemoryRowLevelOperationTable.scala:167-170` and `:190`). However, `declaredColumns` always starts with `pk` (`InMemoryRowLevelOperationTable.scala:395-400`), and every table those fixtures build has `pk`, so the arms prevent no failure. They only switch to a full-width operation, which the `extraTableProps -> None` arm of the new shared test (`ColumnUpdateTableSharedTests.scala:235`) then checks. The case `("id = id, dep = dep", Seq.empty, Seq("pk", "dep"), Row(1, 10, "hr"))`, which passed on d8a4641, was removed for the same reason.
I'd prefer that Spark not narrow when the declaration is empty, rather than asking every builder to branch:
- `resolveWriteAttrs` (`RewriteUpdateTable.scala:268-271`) keeps `validateUpdatedColumnsSubset`, so a non-identity UPDATE with an empty declaration still fails with `COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS`, and it returns `relation.output` when the declaration is empty.
- `V2Writes.hasNarrowRows` (`V2Writes.scala:144-146`) uses the same predicate. Otherwise `validateWriteRelation` (`V2Writes.scala:157-158`) re-resolves the empty declaration and fails.
- This check, its error condition, the Javadoc at `SupportsColumnUpdates.java:65-66` and the fixture arms can then go, and the removed case can come back.
Removing L369-371 alone is not enough. With an empty write relation, `buildWriteDeltaProjections` returns `rowProjection = None` (L246-251), the DELETE shape, and `DeltaWritingSparkTask` calls `project` on it for every UPDATE row (`WriteToDataSourceV2Exec.scala:880` and `:900`). Narrowing only when `updatedColumns()` is non-empty would also work, but it stops narrowing identity-only UPDATEs for operations that declare columns (the `Some(Seq("pk", "salary"))` arm at `ColumnUpdateTableSharedTests.scala:236`), and `hasNarrowRows` would need `RowLevelOperationTable.operationInfo`, which `RowLevelWrite` does not expose. If the Javadoc rule stays instead, it needs the UPDATE scope from [this comment](https://github.com/apache/spark/pull/55518#discussion_r4395825050), and it should say that an operation whose declaration can be empty must drop the mix-in, since keeping the mix-in does not work for such an operation.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala:
##########
@@ -110,8 +111,10 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla
// this rule assigns runtime filters to both scan relations (will be shared at runtime)
// and must transform the runtime filter condition to use correct expr IDs for each relation
// note this only applies to group-based row-level operations (i.e. ReplaceData)
+ // the map is built from the original table as the write table may be narrower than the
Review Comment:
**57.** Nit: this comment reads as if the original table covers every column the condition references, which does not hold for metadata columns. When the UPDATE resolves in one analyzer iteration, `AddMetadataColumns` never widens the target relation, so `write.originalTable.output` has only the data columns, while a condition such as `_partition = 'hr' AND pk IN (SELECT pk FROM s)` references `_partition` from `metadataOutput`, which this map does not cover either. Could the comment say that the map covers the table's data columns, which the write table may leave out? The metadata-column case behaves the same on master with `write.table.output`, so I'll look into that part separately.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -228,4 +244,96 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Returns the table attributes that these assignments set to a new value, skipping identity
+ * assignments.
+ */
+ private def collectUpdatedAttrs(assignments: Seq[Assignment]): Seq[AttributeReference] = {
+ assignments.collect {
+ case Assignment(key: AttributeReference, value) if !isIdentityAssignment(key, value) => key
Review Comment:
**56. `updatedColumns()` can name a metadata column, which no operation that keeps the mix-in can then declare.** This `collect` keeps metadata keys. When an UPDATE resolves only in a later Resolution iteration, for example because `salary > 0.5` needs a cast, `AddMetadataColumns` adds the metadata columns to the target relation (the shape that `DeltaBasedColumnUpdateTableSuite.scala:450-469` analyzes), and the alignment over that output (`ResolveRowLevelCommandAssignments.scala:48`) keeps an assignment to a metadata column. So `UPDATE t SET _partition = 'x' WHERE salary > 0.5` on a `column-update` or `column-update-cow` table reports `updatedColumns()` as `[_partition]`:
- The default fixtures follow `SupportsColumnUpdates.java:59-61` and declare `[pk, dep, _partition]`, which `resolveRequiredDataAttrs` rejects with `COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE` (`RewriteRowLevelCommand.scala:402-408`), as `SupportsColumnUpdates.java:77-80` says it does.
- An operation that leaves `_partition` out fails `validateUpdatedColumnsSubset` (L284-287) with `COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS` for `[_partition]`.
Both messages blame the connector, while on a full-width table the same statement analyzes and the value goes through the metadata projection. Whether such an assignment survives the alignment depends on the analyzer iteration regardless of this PR (with `salary > 0` it is dropped silently), but the dead end between the two checks is new. Could `collectUpdatedAttrs` skip `MetadataAttribute` keys, or could the rewrite reject an assignment to a metadata column with an error that names the user's assignment?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala:
##########
@@ -1500,6 +1501,21 @@ private class BufferWriter(schema: StructType) extends DataWriter[InternalRow] {
buffer.log.append(logEntry)
}
+ // Tag UPDATE/COPY rows distinctly so tests can verify that Spark dispatches
+ // through writeColumnUpdate (the SupportsColumnUpdates path) rather than write.
+ override def writeColumnUpdate(metadata: InternalRow, row: InternalRow): Unit = {
Review Comment:
**58. These overrides are on the `BufferWriter` that every in-memory writer uses, so a dispatch that sends full-width rows to `writeColumnUpdate` passes every suite.** A writer that does not override `writeColumnUpdate`, as for any operation that does not opt in, throws `DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED` there (`DataWriter.java:138-142`). With these overrides, every in-memory writer appends the row exactly like `write`, and only the log tag differs. The commits do not look at the tag: `PartitionBasedReplaceData.doCommit` commits `buffer.rows` (`InMemoryRowLevelOperationTable.scala:269-278`), and `PartitionBasedNarrowReplaceData.doCommit` overlays each tagged row on its original row with `columnUpdateSchema().orElse(schema())` (`InMemoryRowLevelOperationTable.scala:703-704` and `:756-769`), which rebuilds the same full row. So neither of these regressions would be noticed:
- `DataWithProjectionWritingSparkTask` calls `writeColumnUpdate` regardless of `useWriteColumnUpdate` (`WriteToDataSourceV2Exec.scala:836-837`). GroupBasedNoMetadataUpdateTableSuite, GroupBasedNoMetadataDeleteFromTableSuite and GroupBasedNoMetadataMergeIntoTableSuite do not check the write log, so they stay green.
- `DataSourceV2Strategy.scala:629` checks only `isInstanceOf[SupportsColumnUpdates]`, while `V2Writes.hasNarrowRows` keeps `command == UPDATE` (`V2Writes.scala:144-146`). The DELETE and MERGE rows of a `column-update-cow` table then go to `writeColumnUpdate`, and `checkDeleteFallsBackToFullWidthRows` and `checkMergeFallsBackToFullWidthRows` (`ColumnUpdateTableSharedTests.scala:278-320`) still pass, since they check only `columnUpdateSchema()`, `updatedColumns()` and the final rows.
Only the metadata-path tests that assert `Write` tags (`GroupBasedUpdateTableSuite.scala:48-50`, `GroupBasedDeleteFromTableSuite.scala:47-49` and `GroupBasedMergeIntoTableSuite.scala:75-81`) would catch the same mistake in `DataAndMetadataWritingSparkTask`. Could these overrides move to a writer that only the column-update fixtures create, e.g. a `BufferWriter` subclass used by `PartitionBasedNarrowReplaceData` and as the delegate of `RecordOnlyColumnUpdateWriter` (`InMemoryRowLevelOperationTable.scala:911`)? Alternatively, the two fallback helpers could also assert that the write log has no `WriteColumnUpdate` entry.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -273,6 +345,440 @@ class InMemoryRowLevelOperationTable private (
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
}
+ // Finds the one current row with the given row-ID values, matching a data column by value and
+ // `_partition` by the partition key.
+ private def currentRowByRowId(values: Seq[(String, Any)]): InternalRow = {
+ val matches = dataMap.values.iterator.flatten.flatMap { buf =>
+ buf.rows.filter { row =>
+ values.forall {
+ case (name, value) if name == PartitionKeyColumn.name => buf.keyString() == value.toString
+ case (name, value) =>
+ val i = schema.fieldIndex(name)
+ row.get(i, schema.fields(i).dataType) == value
+ }
+ }
+ }.toSeq
+ matches match {
+ case Seq(row) => row
+ case _ => throw new IllegalStateException(
+ s"Expected one row with row ID $values, found ${matches.size}")
+ }
+ }
+
+ // Rebuilds a full table row for a narrow column-update row: the base row's values (or the
+ // existence default for a column added after the base row was written), overlaid with the
+ // narrow row's columns.
+ private def overlayNarrowRow(
+ baseRow: Option[InternalRow],
+ narrowRow: InternalRow,
+ narrowFieldIdx: Map[String, Int]): InternalRow = {
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ val field = schema.fields(i)
+ val value = if (i < base.numFields) {
+ base.get(i, field.dataType)
Review Comment:
**59. The narrow commits read stored rows by the current schema's position and type, so they rebuild wrong rows after a column move or a type change.** `dataMap` keeps each row in the schema it was written with (`alterTableWithData` keeps the old schema unless a column is dropped, `InMemoryBaseTable.scala:491-495`), and the reader maps fields by name against that schema and casts widened types (`InMemoryBaseTable.scala:1312-1313`, `:1331` and `:1377-1378`). This overlay, `currentRowByRowId` (L356-357) and the copy-on-write `pk` lookup (L762) use the current ordinal and type instead. For example, with `(pk INT NOT NULL, salary INT, bonus INT, dep STRING)` and the rows (1, 100, 10, hr), (2, 200, 20, software) and (3, 300, 30, hr):
- After `ALTER TABLE t ALTER COLUMN salary AFTER bonus`, `UPDATE t SET salary = -1 WHERE pk = 1` would leave `bonus` = 100 for pk 1 on a `column-update` table, and `bonus` = 300 for the copied pk 3 on a `column-update-cow` table.
- After `ALTER COLUMN bonus TYPE DOUBLE`, the overlay reads the stored INT slot as a DOUBLE, so `bonus` 10 would come back as 4.9E-323 instead of 10.0.
- After moving `pk`, the first narrow UPDATE would fail at commit with "found 0".
The full-width fixtures keep the right values in all three cases. No test does this today, but since finding 25 the whole `UpdateTableSuiteBase` battery runs on these fixtures, so the next schema-evolution UPDATE test added there would fail only on the narrow suites. Could these helpers read stored rows by name through the row's own `BufferedRows.schema`, as `BufferedRowsReader` does?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala:
##########
@@ -109,22 +113,98 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper {
WriteToDataSourceV2(Some(r), microBatchWrite, newQuery, customMetrics)
case rd @ ReplaceData(r: DataSourceV2Relation, _, query, _, projections, _, None) =>
- val rowSchema = projections.rowProjection.schema
+ validateWriteRelation(rd)
+ val (rowSchema, columnUpdateSchema) =
+ writeSchemas(projections.rowProjection.schema, hasNarrowRows(rd))
val metadataSchema = projections.metadataProjection.map(_.schema)
val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap)
- val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema)
+ val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema,
+ columnUpdateSchema)
val write = writeBuilder.build()
+ validateColumnUpdateWrite(rd, write)
val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, r.funCatalog)
rd.copy(write = Some(write), query = newQuery)
case wd @ WriteDelta(r: DataSourceV2Relation, _, query, _, projections, _, None) =>
+ validateWriteRelation(wd)
val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap)
- val deltaWriteBuilder = newDeltaWriteBuilder(r.table, writeOptions, projections)
+ val deltaWriteBuilder =
+ newDeltaWriteBuilder(r.table, writeOptions, projections, hasNarrowRows(wd))
val deltaWrite = deltaWriteBuilder.build()
+ validateColumnUpdateWrite(wd, deltaWrite)
val newQuery = DistributionAndOrderingUtils.prepareQuery(deltaWrite, query, r.funCatalog)
wd.copy(write = Some(deltaWrite), query = newQuery)
}
+ /**
+ * Whether a row-level write delivers narrow rows for updated, copied, and reinserted records.
+ * This must match the writes whose write relation RewriteUpdateTable narrows, even when the
+ * declared columns cover the whole table.
+ */
+ private[v2] def hasNarrowRows(write: RowLevelWrite): Boolean = {
+ write.operation.isInstanceOf[SupportsColumnUpdates] && write.operation.command == UPDATE
+ }
+
+ /**
+ * Checks that the write relation matches the current operation: a write that delivers narrow
+ * rows must have exactly the declared columns, in declared order, and any other write the full
+ * table columns. A rule that replaces the operation after the rewrite can break this, and the
+ * connector would then decode the rows with the wrong layout.
+ */
+ private def validateWriteRelation(rowLevelWrite: RowLevelWrite): Unit = {
+ val writeNames = rowLevelWrite.table.output.map(_.name)
+ val expectedNames = rowLevelWrite.operation match {
+ case op: SupportsColumnUpdates if hasNarrowRows(rowLevelWrite) =>
+ RewriteRowLevelCommand.resolveRequiredDataAttrs(rowLevelWrite.originalTable, op)
+ .map(_.name)
+ case _ =>
+ rowLevelWrite.originalTable.output.map(_.name)
+ }
+ if (writeNames != expectedNames) {
+ throw SparkException.internalError(
+ s"The write relation of operation ${rowLevelWrite.operation.getClass.getName} has " +
+ s"columns ${writeNames.mkString("[", ", ", "]")}, but the operation expects " +
+ s"${expectedNames.mkString("[", ", ", "]")}.")
+ }
+ }
+
+ /**
+ * Checks that the distribution and ordering of a write that delivers narrow rows reference only
+ * columns passed to the write. Other columns stay in the query only until column pruning removes
+ * them, so they are rejected here whether or not pruning ran.
+ */
+ private def validateColumnUpdateWrite(rowLevelWrite: RowLevelWrite, write: Write): Unit = {
+ write match {
+ case w: RequiresDistributionAndOrdering if hasNarrowRows(rowLevelWrite) =>
+ val distribution = w.requiredDistribution match {
+ case d: ClusteredDistribution => d.clustering.toImmutableArraySeq
+ case d: OrderedDistribution => d.ordering.toImmutableArraySeq
+ case _ => Seq.empty
Review Comment:
**61.** Nit: this repeats the distribution match of `DistributionAndOrderingUtils.prepareQuery` with a different fallback. `prepareQuery` (`DistributionAndOrderingUtils.scala:41-49`), which runs right after this check (L124-125 and L134-135), matches `OrderedDistribution`, `ClusteredDistribution` and `UnspecifiedDistribution` without a wildcard, while this match ends with `case _ => Seq.empty`. If a distribution kind is added and only `prepareQuery` learns it, this check would skip that kind's columns, and an undeclared clustering column would again pass or fail depending on whether ColumnPruning ran, which item 7 removed. Could both use one helper that returns the distribution's expressions?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -72,10 +72,12 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
protected def buildOperationTable(
table: SupportsRowLevelOperations,
command: Command,
- options: CaseInsensitiveStringMap): RowLevelOperationTable = {
- val info = RowLevelOperationInfoImpl(command, options)
+ options: CaseInsensitiveStringMap,
+ updatedAttrs: Seq[AttributeReference] = Nil): RowLevelOperationTable = {
+ val updatedColumns = updatedAttrs.map(attr => FieldReference(Seq(attr.name)))
Review Comment:
**60. Nothing tests that a top-level column with a dot in its name reaches `updatedColumns()` as one part.** This line uses `FieldReference(Seq(attr.name))`, i.e. `FieldReference.column(attr.name)`, so that a column such as `` `a.b` `` is not parsed into the nested reference `` `a`.`b` `` ([earlier thread](https://github.com/apache/spark/pull/55518#discussion_r3634079860)). However, no `checkLastUpdatedColumns` call uses a dotted name, and the only `` `a.b` `` UPDATE test (`DeltaBasedColumnUpdateTableSuite.scala:364-375`) runs on the req-attrs fixture, which ignores `updatedColumns()` and builds its expected message from `a.name` (`RewriteUpdateTable.scala:284`). So a revert to `FieldReference(attr.name)` would pass every suite, while connectors would receive `` `a`.`b` `` again and the default `column-update` fixtures, which declare the reported columns, would reject ``SET `a.b` = -1`` with `COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE`. Could a test on the default fixture update
`` `a.b` `` and check `` checkLastUpdatedColumns("`a.b`") ``? `describe()` quotes a dotted part, so the check tells the two forms apart. `FieldReference.column(attr.name)` would also make the intent visible here.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -128,12 +151,61 @@ class InMemoryRowLevelOperationTable private (
copied.replacedPartitions = replacedPartitions
copied.lastWriteInfo = lastWriteInfo
copied.lastWriteLog = lastWriteLog
+ copied.lastUpdatedColumns = lastUpdatedColumns
copied
}
override def newRowLevelOperationBuilder(
info: RowLevelOperationInfo): RowLevelOperationBuilder = {
- if (properties.getOrDefault(SUPPORTS_DELTAS, "false") == "true") {
+ lastUpdatedColumns = info.updatedColumns()
+ val splitUpdates = properties.getOrDefault(SPLIT_UPDATES, "false").toBoolean
+ // the copy-on-write commit finds the original row by `pk`
+ val updatesPk = info.updatedColumns().exists(_.describe() == "pk")
+ // The default column-update fixtures drop the mix-in for an UPDATE with no updated columns,
+ // as SupportsColumnUpdates asks. The fixed-declaration fixtures, the split row-ID fixture and
+ // all fixtures for DELETE and MERGE keep it.
+ val updatesNothing = info.command == Command.UPDATE && info.updatedColumns().isEmpty
+ if (properties.getOrDefault(COLUMN_UPDATE, "false") == "true") {
+ if (updatesNothing) {
+ () => DeltaBasedOperation(info.command, info.options)
+ } else if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ }
+ } else if (properties.containsKey(COLUMN_UPDATE_REQ_ATTRS)) {
+ val reqCols = properties.get(COLUMN_UPDATE_REQ_ATTRS).split(",").map(_.trim)
+ if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperationWithReqAttrs(
+ info.command, reqCols, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperationWithReqAttrs(info.command, reqCols, info.options)
+ }
+ } else if (properties.getOrDefault(COLUMN_UPDATE_EMPTY_REQ_ATTRS, "false") == "true") {
Review Comment:
**62.** Nit: `column-update-empty-req-attrs` (L100) builds the same operation as `column-update-req-attrs` with an empty list. This branch exists only because `"".split(",")` returns `Array("")` at L179. With `.filter(_.nonEmpty)` there and at L197, "empty requiredDataAttributes throws AnalysisException" (`DeltaBasedColumnUpdateTableSuite.scala:336-350`) could use `reqAttrsTableProps("")`. With item 55's suggestion, this key would go away anyway.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RowLevelWriteColumnPruningSuite.scala:
##########
@@ -0,0 +1,278 @@
+/*
+ * 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.catalyst.optimizer
+
+import java.util
+
+import org.apache.spark.SparkException
+import org.apache.spark.sql.catalyst.ProjectingInternalRow
+import org.apache.spark.sql.catalyst.dsl.expressions._
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, EqualNullSafe, Expression, If, Literal, Not}
+import org.apache.spark.sql.catalyst.plans.{LeftOuter, PlanTest}
+import org.apache.spark.sql.catalyst.plans.logical._
+import org.apache.spark.sql.catalyst.plans.logical.MergeRows.{Copy, Keep, Update}
+import org.apache.spark.sql.catalyst.rules.RuleExecutor
+import org.apache.spark.sql.catalyst.util.{ReplaceDataProjections, WriteDeltaProjections}
+import org.apache.spark.sql.catalyst.util.RowDeltaUtils._
+import org.apache.spark.sql.connector.catalog.{Column, Table, TableCapability}
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Tests the [[ColumnPruning]] case that prunes the query columns of a [[RowLevelWrite]] that its
+ * projections don't read and moves the projections to the ordinals of the pruned query.
+ */
+class RowLevelWriteColumnPruningSuite extends PlanTest {
+
+ object Optimize extends RuleExecutor[LogicalPlan] {
+ val batches = Batch("Column Pruning", FixedPoint(10),
+ ColumnPruning, CollapseProject, RemoveNoopOperators) :: Nil
+ }
+
+ // t(pk, id, dep, salary, bonus, extra) with the metadata column _partition
+ private val read = LocalRelation(
+ $"pk".int.notNull, $"id".int, $"dep".string, $"salary".int, $"bonus".int, $"extra".int,
+ $"_partition".string.notNull)
+
+ // the write and original table relations; the pruning never inspects them
+ private val tableRelation: DataSourceV2Relation = {
Review Comment:
**64.** Nit: `TestRelation` could replace this placeholder. As L53 says, the pruning never inspects these relations, but the placeholder needs an anonymous `Table`, `DataSourceV2Relation.create` and the imports at L20, L32, L33 and L35. Catalyst tests already have a minimal `NamedRelation`, `TestRelation` (`V2WriteAnalysisSuite.scala:137`), which `InlineCTESuite.scala:39` also uses, so `TestRelation(Nil)` could serve as both `table` and `originalTable`.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -273,6 +345,440 @@ class InMemoryRowLevelOperationTable private (
override def abort(messages: Array[WriterCommitMessage]): Unit = {}
}
+ // Finds the one current row with the given row-ID values, matching a data column by value and
+ // `_partition` by the partition key.
+ private def currentRowByRowId(values: Seq[(String, Any)]): InternalRow = {
+ val matches = dataMap.values.iterator.flatten.flatMap { buf =>
+ buf.rows.filter { row =>
+ values.forall {
+ case (name, value) if name == PartitionKeyColumn.name => buf.keyString() == value.toString
+ case (name, value) =>
+ val i = schema.fieldIndex(name)
+ row.get(i, schema.fields(i).dataType) == value
+ }
+ }
+ }.toSeq
+ matches match {
+ case Seq(row) => row
+ case _ => throw new IllegalStateException(
+ s"Expected one row with row ID $values, found ${matches.size}")
+ }
+ }
+
+ // Rebuilds a full table row for a narrow column-update row: the base row's values (or the
+ // existence default for a column added after the base row was written), overlaid with the
+ // narrow row's columns.
+ private def overlayNarrowRow(
+ baseRow: Option[InternalRow],
+ narrowRow: InternalRow,
+ narrowFieldIdx: Map[String, Int]): InternalRow = {
+ val fullRow = new GenericInternalRow(schema.length)
+ baseRow.foreach { base =>
+ for (i <- schema.fields.indices) {
+ val field = schema.fields(i)
+ val value = if (i < base.numFields) {
+ base.get(i, field.dataType)
+ } else {
+ ResolveDefaultColumns.getExistenceDefaultValue(field)
+ }
+ fullRow.update(i, value)
+ }
+ }
+ schema.fields.zipWithIndex.foreach { case (field, i) =>
+ narrowFieldIdx.get(field.name).foreach { j =>
+ fullRow.update(i, narrowRow.get(j, field.dataType))
+ }
+ }
+ fullRow
+ }
+
+ // The base columns `pk` and `dep`, where the table has them, followed by the updated columns.
+ private def declaredColumns(updatedCols: Seq[NamedReference]): Array[NamedReference] = {
+ val base = Seq("pk", "dep").filter(schema.fieldNames.contains).map(FieldReference(_))
+ val baseNames = base.map(_.describe()).toSet
+ (base ++ updatedCols.filterNot(r => baseNames.contains(r.describe()))).toArray
+ }
+
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep`, where the table has them, plus whatever columns Spark reports as
+ // being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
+ command: Command,
+ updatedCols: Seq[NamedReference] = Nil,
+ options: CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty())
+ extends DeltaBasedOperation(command, options)
+ with SupportsColumnUpdates {
+ override def representUpdateAsDeleteAndInsert(): Boolean = false
+ override def requiredDataAttributes(): Array[NamedReference] = declaredColumns(updatedCols)
+
+ // The case-class equality of DeltaBasedOperation would make this operation equal to a wide
+ // one, or to another column-update operation, with the same command and options.
+ override def equals(other: Any): Boolean = other match {
+ case ref: AnyRef => this eq ref
+ case _ => false
+ }
+ override def hashCode(): Int = System.identityHashCode(this)
+ override def canEqual(other: Any): Boolean = false
+
+ private val clusterExprs: Seq[Expression] = {
+ val refs = Option(properties.get(COLUMN_UPDATE_CLUSTER_BY))
+ .map(_.split(",").map(name => FieldReference(name.trim)).toSeq)
+ .getOrElse(Seq(PARTITION_COLUMN_REF))
+ if (properties.getOrDefault(COLUMN_UPDATE_CLUSTER_BY_NESTED_TRANSFORM, "false") == "true") {
+ refs.map(ref => LogicalExpressions.apply("f", LogicalExpressions.apply("g", ref)))
+ } else {
+ refs
+ }
+ }
+
+ override def newWriteBuilder(info: LogicalWriteInfo): DeltaWriteBuilder = {
+ lastWriteInfo = info
+ val writeSchema = info.columnUpdateSchema().orElse(info.schema())
+ val batchWrite = new TestBatchWrite with DeltaBatchWrite {
+ override def createBatchWriterFactory(
+ physicalInfo: PhysicalWriteInfo): DeltaWriterFactory = {
+ newWriterFactory(info, writeSchema)
+ }
+
+ override protected def doCommit(messages: Array[WriterCommitMessage]): Unit =
+ dataMap.synchronized {
+ val newData = messages.map(_.asInstanceOf[BufferedRows])
+ commitColumnUpdate(info, writeSchema, newData)
+ }
+
+ override def abort(messages: Array[WriterCommitMessage]): Unit = {}
+ }
+ new DeltaWriteBuilder {
+ override def build(): DeltaWrite =
+ if (noMetadata && !properties.containsKey(COLUMN_UPDATE_CLUSTER_BY)) {
+ new DeltaWrite {
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ } else {
+ new DeltaWrite with RequiresDistributionAndOrdering {
+
+ override def requiredDistribution(): Distribution = {
+ Distributions.clustered(clusterExprs.toArray)
+ }
+
+ override def requiredOrdering(): Array[SortOrder] = {
+ clusterExprs.map { expr =>
+ LogicalExpressions.sort(
+ expr,
+ SortDirection.ASCENDING,
+ SortDirection.ASCENDING.defaultNullOrdering()): SortOrder
+ }.toArray
+ }
+
+ override def toBatch: DeltaBatchWrite = batchWrite
+ }
+ }
+ }
+ }
+
+ protected def newWriterFactory(
+ info: LogicalWriteInfo,
+ writeSchema: StructType): DeltaWriterFactory = {
+ new DeltaBufferedRowsWriterFactory(writeSchema)
+ }
+
+ // For column-update writes, updated rows contain only the declared columns
+ // (columnUpdateSchema from LogicalWriteInfo). We expand each row to the full
+ // table schema by overlaying them on the base row found by pk.
+ protected def commitColumnUpdate(
+ info: LogicalWriteInfo,
+ writeSchema: StructType,
+ newData: Array[BufferedRows]): Unit = {
+ val writeFieldIdx = writeSchema.fieldNames.zipWithIndex.toMap
Review Comment:
**65.** Nit: `StructType.getFieldIndex` already provides this name-to-index map. The three commits build `writeSchema.fieldNames.zipWithIndex.toMap` here, at L555 and at L751, and pass it to `overlayNarrowRow` (L371-374), which has to trust that the map belongs to `narrowRow`. Passing the narrow `StructType` and calling `getFieldIndex` (cached; `BufferedRowsReader` uses it at `InMemoryBaseTable.scala:1331`) would drop the three copies and tie the index to the row's schema. The `fieldNames.contains(name)` plus `fieldIndex(name)` pair at L567-571 would also become one `getFieldIndex` call.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/ColumnUpdateTableSharedTests.scala:
##########
@@ -0,0 +1,387 @@
+/*
+ * 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 scala.jdk.CollectionConverters._
+
+import org.apache.spark.{SparkException, SparkRuntimeException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.{InMemoryBaseTable, InMemoryRowLevelOperationTable}
+import org.apache.spark.sql.connector.write.{RowLevelOperation, RowLevelOperationTable}
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+// Column-update tests that hold for both delta-based and group-based operations; each suite
+// supplies its column-update table properties through `extraTableProps`.
+trait ColumnUpdateTableSharedTests extends UpdateTableSuiteBase {
+
+ // the properties of a column-update table whose connector declares exactly `reqAttrs`
+ protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String]
+
+ // the properties of a table of the same kind whose operation does not mix in column updates
+ protected def fullWidthTableProps: java.util.Map[String, String]
+
+ // the metadata columns the scan of a column-update table reads, in DDL form
+ protected def scanMetadataColumns: String = ", _partition STRING, index INT"
+
+ // the class name of the operation that `reqAttrsTableProps` builds
+ protected def reqAttrsOperationClass: String = if (deltaUpdate) {
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateOperationWithReqAttrs].getName
+ } else {
+ classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperationWithReqAttrs].getName
+ }
+
+ // `base` with the given table properties added; a given key replaces the same key in `base`
+ protected def tablePropsWith(
+ base: java.util.Map[String, String],
+ props: (String, String)*): java.util.Map[String, String] = {
+ (base.asScala ++ props).asJava
+ }
+
+ // t(pk, id, dep, salary, bonus, extra) with CHECK (extra > 0), partitioned by dep
+ protected def createAndInitTableWithProps(props: java.util.Map[String, String]): Unit = {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, salary INT, bonus INT, extra INT",
+ """{ "pk": 1, "id": 1, "dep": "hr", "salary": 100, "bonus": 10, "extra": 1 }
+ |{ "pk": 2, "id": 2, "dep": "hr", "salary": 200, "bonus": 20, "extra": 1 }
+ |{ "pk": 3, "id": 3, "dep": "software", "salary": 300, "bonus": 30, "extra": 1 }
+ |""".stripMargin,
+ props)
+ sql(s"ALTER TABLE $tableNameAsString ADD CONSTRAINT positive_extra CHECK (extra > 0)")
+ }
+
+ test("column-update: undeclared columns are not written and keep their values") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkLastColumnUpdateSchema("pk", "dep", "salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update: write relation has the declared columns in declared order") {
+ createAndInitTableWithProps(reqAttrsTableProps("salary,pk"))
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'")
+ }
+
+ assert(!write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "salary", "pk")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update: the plan writes the declared columns and scans the referenced ones") {
+ createAndInitTableWithProps(reqAttrsTableProps("pk,salary"))
+
+ // `id` is neither declared nor referenced; `dep`, `bonus` and `extra` are read for the
+ // condition, the assignment and the CHECK constraint
+ val write = executeAndKeepAnalyzedWrite {
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'",
+ s"pk INT, dep STRING, salary INT, bonus INT, extra INT$scanMetadataColumns")
+ }
+
+ val writeNodeName = if (deltaUpdate) "WriteDelta" else "ReplaceData"
+ val header = write.simpleString(Int.MaxValue)
+ assert(header.matches(s"""$writeNodeName RelationV2\\[pk#\\d+, salary#\\d+\\] .*"""), header)
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update: undeclared, differently-cased columns in the RHS and WHERE clause are " +
+ "read with the table's own spelling") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, extra INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "extra": 5, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "extra": 5, "dep": "hr" }
+ |""".stripMargin)
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + BONUS WHERE pk = 1 AND EXTRA > 3",
+ s"pk INT, salary INT, bonus INT, extra INT, dep STRING$scanMetadataColumns")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 110, 10, 5, "hr") :: Row(2, 200, 20, 5, "hr") :: Nil)
+ }
+
+ test("column-update: assigned column outside requiredDataAttributes is rejected") {
+ createAndInitTableWithProps(reqAttrsTableProps("pk,salary"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET bonus = 0 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "missingColumns" -> "[`bonus`]"))
+ }
+
+ test("column-update: an operation replaced after the rewrite is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ }
+ val relation = write.table.asInstanceOf[DataSourceV2Relation]
+ val operationTable = relation.table.asInstanceOf[RowLevelOperationTable]
+ val t = operationTable.table.asInstanceOf[InMemoryRowLevelOperationTable]
+ val replacements: Seq[(RowLevelOperation, String)] = if (deltaUpdate) {
+ Seq(
+ t.DeltaBasedOperation(UPDATE, CaseInsensitiveStringMap.empty()) -> "[pk, salary, dep]",
+ new t.DeltaBasedColumnUpdateOperation(UPDATE) -> "[pk, dep]")
+ } else {
+ Seq(
+ t.PartitionBasedOperation(UPDATE, CaseInsensitiveStringMap.empty()) -> "[pk, salary, dep]",
+ new t.PartitionBasedColumnUpdateOperation(UPDATE) -> "[pk, dep]")
+ }
+ replacements.foreach { case (operation, expectedNames) =>
+ assert(operation != operationTable.operation)
+ val replaced = write.withNewTable(
+ relation.copy(table = operationTable.copy(operation = operation)(
+ operationTable.operationInfo)))
+ checkError(
+ exception = intercept[SparkException](V2Writes(replaced)),
+ condition = "INTERNAL_ERROR",
+ parameters = Map("message" ->
+ (s"The write relation of operation ${operation.getClass.getName} has columns " +
+ s"[pk, dep, salary], but the operation expects $expectedNames.")))
+ }
+ }
+
+ test("column-update: nested struct field update narrows to the root struct column") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |{ "pk": 2, "s": { "c1": 3, "c2": 4 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+
+ checkLastUpdatedColumns("s")
+ checkLastColumnUpdateSchema("pk", "dep", "s")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, Row(-1, 2), "hr") :: Row(2, Row(3, 4), "hr") :: Nil)
+ }
+
+ test("column-update: excluding ColumnPruning keeps the write narrow") {
+ createAndInitTableWithProps(reqAttrsTableProps("pk,salary"))
+
+ withSQLConf(SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName) {
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'",
+ s"pk INT, id INT, dep STRING, salary INT, bonus INT, extra INT$scanMetadataColumns")
+ }
+
+ checkLastColumnUpdateSchema("pk", "salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update: reassigning a declared row-ID column keeps the row's other columns") {
+ createAndInitTableWithProps(reqAttrsTableProps("pk,salary"))
+
+ sql(s"UPDATE $tableNameAsString SET pk = 2, salary = 0 WHERE pk = 1")
+
+ // the copy-on-write fixture finds original rows by `pk`, so it falls back to full-width rows
+ assert(table.lastWriteInfo.columnUpdateSchema().isPresent == deltaUpdate)
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk, id"),
+ Row(2, 1, "hr", 0, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update: an UPDATE whose assignments are all identity assignments is narrowed " +
+ "only if the operation keeps the mix-in") {
+ Seq(
+ extraTableProps -> None,
+ reqAttrsTableProps("pk,salary") -> Some(Seq("pk", "salary"))
+ ).foreach { case (props, columnUpdateNames) =>
+ withTable(tableNameAsString) {
+ createAndInitTableWithProps(props)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary, dep = dep WHERE pk = 1")
+
+ checkLastUpdatedColumns()
+ columnUpdateNames match {
+ case Some(names) => checkLastColumnUpdateSchema(names: _*)
+ case None =>
+ assert(!table.lastWriteInfo.columnUpdateSchema().isPresent,
Review Comment:
**63.** Nit: the full-width check is written out three times with three messages. L247-248, L281-283 and L312-314 repeat `assert(!table.lastWriteInfo.columnUpdateSchema().isPresent, ...)`, and L224 checks `isPresent == deltaUpdate` without a message. Could `RowLevelOperationSuiteBase` get a `checkLastWriteIsFullWidth()` next to `checkLastColumnUpdateSchema` (item 45), with L224 picking one of the two by `deltaUpdate`?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub) via reviews" <re...@spark.apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4428977891
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,114 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, updated, copied, and reinserted
+ * rows follow it: group-based operations receive them through
+ * {@link DataWriter#writeColumnUpdate(Object, Object)} or
+ * {@link DataWriter#writeColumnUpdate(Object)}, and operations that mix in {@link SupportsDelta}
+ * through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert}. Rows the command inserts
+ * always follow {@link LogicalWriteInfo#schema()} and arrive through
+ * {@link DataWriter#write(Object)} and {@link DeltaWriter#insert}. When
+ * {@link LogicalWriteInfo#columnUpdateSchema()} is absent, every row follows
+ * {@link LogicalWriteInfo#schema()}, as for any other operation. Connectors must decide which
+ * layout to expect by whether {@link LogicalWriteInfo#columnUpdateSchema()} is present, not by
+ * the command. A builder that wants full-width rows for some command can check
+ * {@link RowLevelOperationInfo#command()} and build an operation without this mix-in. Currently
+ * Spark narrows only UPDATE.
+ * <p>
+ * The scan builder returned by {@link #newScanBuilder} should implement
+ * {@link org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns}, so that the scan
+ * does not read columns the command neither references nor declares. Otherwise the scan reads
+ * every column. For operations that do not mix in {@link SupportsDelta}, only top-level columns
+ * are pruned: a struct column the command references through a nested field is read in full.
+ *
+ * @since 4.4.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * When Spark narrows the write, the returned columns become
+ * {@link LogicalWriteInfo#columnUpdateSchema()}, in declared order. Spark matches each declared
+ * column with the session's resolver, and {@link LogicalWriteInfo#columnUpdateSchema()} uses the
+ * table's spelling of each column, which can differ in case from the declared reference.
+ * Implementations must include every column they want to receive: every column reported by
+ * {@link RowLevelOperationInfo#updatedColumns()}, plus any columns needed for row lookup or
+ * routing. Columns that are not declared are absent from the rows the connector receives, so
+ * the connector must preserve their values itself. Spark may call this method more than once
+ * for the same operation, and every call must return the same columns.
+ * <p>
+ * If the returned array is empty, Spark does not narrow the write and the operation receives
+ * full-width rows. An UPDATE whose assignments are all identity assignments reports no
+ * {@link RowLevelOperationInfo#updatedColumns()}, so an operation that declares exactly those
+ * columns takes this path for such an UPDATE. An empty array is not an opt-out: an UPDATE that
+ * sets a column is still rejected, because the array misses that column.
+ * <p>
+ * An operation that does not mix in {@link SupportsDelta} should locate original rows through
+ * metadata columns returned by {@link #requiredMetadataAttributes()} that are preserved on
+ * update ({@link org.apache.spark.sql.connector.catalog.MetadataColumn#PRESERVE_ON_UPDATE}),
+ * such as a file and a row position. If it locates rows through a data column, its builder
+ * should build an operation without this mix-in when
+ * {@link RowLevelOperationInfo#updatedColumns()} contains that column.
+ * <p>
+ * Each entry must name a top-level data column of the table. For updates on nested fields such
Review Comment:
This sounds like a temp limitation?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub) via reviews" <re...@spark.apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4428138885
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -237,4 +253,101 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Returns the table attributes that these assignments set to a new value, skipping identity
+ * assignments.
+ */
+ private def collectUpdatedAttrs(assignments: Seq[Assignment]): Seq[AttributeReference] = {
+ assignments.collect {
+ case Assignment(key: AttributeReference, value) if !isIdentityAssignment(key, value) => key
Review Comment:
**66. `updatedColumns()` reports the user's spelling when the SET clause is already aligned, so the check that `SupportsColumnUpdates.java:74-76` asks builders to make can miss the lookup column.** This line takes the reported name from the assignment key, and a key resolves to `a.withName(<typed spelling>)` (`expressions/package.scala:197`). Alignment is the only step that restores the table's attribute (`AssignmentUtils.scala:91`), and it runs only for `!u.aligned` (`ResolveRowLevelCommandAssignments.scala:51`). `AssignmentUtils.aligned` compares names with the resolver (`AssignmentUtils.scala:402`), so in a case-insensitive session a SET clause that lists every column in table order skips it, and "Column names match the case used in the table schema" (`RowLevelOperationInfo.java:53`) no longer holds. For example:
- With `t(pk INT NOT NULL, salary INT, bonus INT, dep STRING)` and the rows (1, 100, 10, hr) and (2, 200, 20, hr) on `column-update-cow-req-attrs = pk,salary`, `UPDATE t SET PK = 2, SALARY = 0, BONUS = BONUS, DEP = DEP WHERE PK = 1` reports `[PK, SALARY]`. `updatesPk` (`InMemoryRowLevelOperationTable.scala:161`) is then false, so the narrow operation is built instead of the full-width one, and the commit finds the base row by the new `pk = 2`. The updated row becomes (2, 0, 20, hr) instead of (2, 0, 10, hr), with no error. The same statement in lower case falls back and gets it right. A real connector that follows `SupportsColumnUpdates.java:74-76` with an exact comparison would look rows up the same way.
- With `t(pk INT NOT NULL, id INT, dep STRING)` on the default `column-update` fixture, `UPDATE t SET PK = PK, ID = -1, DEP = 'eng' WHERE pk = 1` declares `[pk, dep, ID, DEP]` (`InMemoryRowLevelOperationTable.scala:392-396`) and fails with `COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE`, while `SET id = -1, dep = 'eng'` works.
`("ID = -1", Seq("id"), ...)` at `DeltaBasedColumnUpdateTableSuite.scala:169` does not cover this, because assigning one of three columns goes through alignment. Since `u.aligned` guarantees one assignment per column in table order, could this use the relation's attributes, e.g. `relation.output.zip(assignments).collect { case (attr, Assignment(key: Attribute, value)) if !isIdentityAssignment(key, value) => attr }`, the same idea as the exprId map-back in `resolveRequiredDataAttrs`, and could that test get a row with upper-case keys?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -237,4 +253,101 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Returns the table attributes that these assignments set to a new value, skipping identity
+ * assignments.
+ */
+ private def collectUpdatedAttrs(assignments: Seq[Assignment]): Seq[AttributeReference] = {
+ assignments.collect {
+ case Assignment(key: AttributeReference, value) if !isIdentityAssignment(key, value) => key
Review Comment:
**67. A metadata key can still reach `updatedColumns()` on a table that accepts any schema, so item 56's dead end remains there.** a6e22b24587 removed the metadata filter on these keys, and [the reply to item 56](https://github.com/apache/spark/pull/55518#discussion_r4426725837) says SPARK-60721 handles it. That holds for other tables, but SPARK-60721's check and SPARK-60667's target exclusion both skip `skipSchemaResolution` tables (`ResolveRowLevelCommandAssignments.scala:98` and `Analyzer.scala:1198`). On a `column-update=true, accept-any-schema=true` table `t(pk INT NOT NULL, salary INT, dep STRING)`, `UPDATE t SET pk = pk, salary = -1, dep = dep, index = 5, _partition = 'x'` needs no injected rule: `AddMetadataColumns` widens the target to `[pk, salary, dep, index, _partition]`, the assignments are then aligned, and this reports `[salary, index, _partition]`. Declaring those fails with `COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE`, leaving them out fails with `COLUMN_UPDAT
E_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS`, and both blame the connector for the user's assignment. Before a6e22b24587 this reported `[salary]`. It is a rare shape, since the SET clause has to list every data and metadata column in output order, and only an error-quality problem. Could this line skip metadata keys again with `!MetadataAttribute.isValid(key.metadata)`, or could `RowLevelOperationInfo.updatedColumns()` mention the case?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/LogicalWriteInfoImpl.scala:
##########
@@ -29,7 +29,8 @@ private[sql] case class LogicalWriteInfoImpl(
schema: StructType,
options: CaseInsensitiveStringMap,
override val rowIdSchema: Optional[StructType] = Optional.empty[StructType],
- override val metadataSchema: Optional[StructType] = Optional.empty[StructType])
+ override val metadataSchema: Optional[StructType] = Optional.empty[StructType],
+ override val columnUpdateSchema: Optional[StructType] = Optional.empty[StructType])
Review Comment:
**69. The PR description says "`LogicalWriteInfoImpl.apply` takes a required `columnUpdateSchema`", but only the `Option` overload at L43 does.** The `Optional` `apply` keeps this default, so master's five-argument shape, as in `FileTable.mergedWriteInfo` before this PR, still compiles. An extension under `org.apache.spark.sql` that rebuilds the info that way hands the connector an empty `schema()` and no `columnUpdateSchema()`, while `DataSourceV2Strategy.scala:651` still sends narrow rows. Item 10's default throw does not help here, because this class overrides the method. Spark's own paths all forward it, and `LogicalWriteInfo.java:83-84` already says a wrapper must forward it, so only the description needs a change. Could the compatibility paragraph say that the compiler catches the `Option` overload but not this one?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/ColumnUpdateTableSharedTests.scala:
##########
@@ -0,0 +1,427 @@
+/*
+ * 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 scala.jdk.CollectionConverters._
+
+import org.apache.spark.{SparkException, SparkRuntimeException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.{InMemoryBaseTable, InMemoryRowLevelOperationTable}
+import org.apache.spark.sql.connector.write.{RowLevelOperation, RowLevelOperationTable}
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+// Column-update tests that hold for both delta-based and group-based operations; each suite
+// supplies its column-update table properties through `extraTableProps`.
+trait ColumnUpdateTableSharedTests extends UpdateTableSuiteBase {
+
+ // the properties of a column-update table whose connector declares exactly `reqAttrs`
+ protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String]
+
+ // the properties of a table of the same kind whose operation does not mix in column updates
+ protected def fullWidthTableProps: java.util.Map[String, String]
+
+ // the metadata columns the scan of a column-update table reads, in DDL form
+ protected def scanMetadataColumns: String = ", _partition STRING, index INT"
+
+ // the class name of the operation that `reqAttrsTableProps` builds
+ protected def reqAttrsOperationClass: String = if (deltaUpdate) {
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateOperationWithReqAttrs].getName
+ } else {
+ classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperationWithReqAttrs].getName
+ }
+
+ // `base` with the given table properties added; a given key replaces the same key in `base`
+ protected def tablePropsWith(
+ base: java.util.Map[String, String],
+ props: (String, String)*): java.util.Map[String, String] = {
+ (base.asScala ++ props).asJava
+ }
+
+ // t(pk, id, dep, salary, bonus, extra) with CHECK (extra > 0), partitioned by dep
+ protected def createAndInitTableWithProps(props: java.util.Map[String, String]): Unit = {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, salary INT, bonus INT, extra INT",
+ """{ "pk": 1, "id": 1, "dep": "hr", "salary": 100, "bonus": 10, "extra": 1 }
+ |{ "pk": 2, "id": 2, "dep": "hr", "salary": 200, "bonus": 20, "extra": 1 }
+ |{ "pk": 3, "id": 3, "dep": "software", "salary": 300, "bonus": 30, "extra": 1 }
+ |""".stripMargin,
+ props)
+ sql(s"ALTER TABLE $tableNameAsString ADD CONSTRAINT positive_extra CHECK (extra > 0)")
+ }
+
+ test("column-update: undeclared columns are not written and keep their values") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkLastColumnUpdateSchema("pk", "dep", "salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update: write relation has the declared columns in declared order") {
+ createAndInitTableWithProps(reqAttrsTableProps("salary,pk"))
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'")
+ }
+
+ assert(!write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "salary", "pk")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update: the plan writes the declared columns and scans the referenced ones") {
Review Comment:
**73.** Nit: some of the new tests overlap or only exercise the fixture.
- L89 and L105 run the same schema, statement and rows. L89 adds the declared order, the original table and the empty `schema()`, and L105 adds the exact scan and the node name, so one test with `"salary,pk"` could check all of them. The regex at L118 repeats the first assertion of `checkColumnUpdateWrite` (`RowLevelOperationSuiteBase.scala:229`).
- L141 and L292 make the same assertion from `"pk,salary"` and `""`, so `Seq(...).foreach` could keep both inputs.
- `DeltaBasedColumnUpdateTableSuite.scala:504` asserts only the text of the fixture's own exception. DELETE is never narrow (`V2Writes.scala:144`), so it covers no Spark behavior.
- `DeltaBasedColumnUpdateTableSuite.scala:572` cannot fail on its CHECK, since `extra > 0` holds for every value it writes. It shows that the column resolves rather than that the constraint is enforced.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,607 @@
+/*
+ * 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 scala.jdk.CollectionConverters._
+
+import org.apache.spark.{SparkRuntimeException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.analysis.RewriteRowLevelCommand
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+import org.apache.spark.sql.catalyst.util.RowDeltaUtils.ORIGINAL_ROW_ID_VALUE_PREFIX
+import org.apache.spark.sql.connector.catalog.{Delete, InMemoryBaseTable, InMemoryRowLevelOperationTable, Reinsert}
+import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
+import org.apache.spark.sql.connector.write.LogicalWriteInfo
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.errors.QueryCompilationErrors
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, updated rows contain only the declared columns
+ * (LogicalWriteInfo.columnUpdateSchema()) rather than the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase
+ with ColumnUpdateTableSharedTests {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ override protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String] = {
+ java.util.Map.of("column-update-req-attrs", reqAttrs)
+ }
+
+ override protected def fullWidthTableProps: java.util.Map[String, String] = {
+ java.util.Map.of("supports-deltas", "true")
+ }
+
+ override protected def scanMetadataColumns: String = ", _partition STRING"
+
+ private val splitOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperation].getName
+
+ private val splitWithRowIdOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithRowId].getName
+
+ private def splitReqAttrsTableProps(reqAttrs: String): java.util.Map[String, String] = {
+ tablePropsWith(reqAttrsTableProps(reqAttrs), "split-updates" -> "true")
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData, reqAttrsTableProps(reqAttrs))
+ }
+
+ private val splitTableProps = tablePropsWith(extraTableProps, "split-updates" -> "true")
+
+ private val splitPartitionRowIdTableProps =
+ java.util.Map.of("column-update-split-row-id", "_partition")
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData, splitTableProps)
+ }
+
+ test("column-update: column update schema contains a single assigned column") {
+ Seq(extraTableProps, splitTableProps).foreach { props =>
+ withTable(tableNameAsString) {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin,
+ props)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+ }
+ }
+
+ test("column-update: update handles metadata columns correctly") {
+ Seq(
+ "id IN (1, 100)",
+ "id IN (1, 20) AND dep IN (SELECT * FROM updated_dep)"
+ ).foreach { condition =>
+ withTable(tableNameAsString) {
+ withTempView("updated_dep") {
+ createAndInitTable("pk INT NOT NULL, id INT, salary INT, dep STRING",
+ """{ "pk": 1, "id": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "id": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+ import testImplicits._
+ Seq(Some("hr"), Some("it")).toDF().createOrReplaceTempView("updated_dep")
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE $condition")
+
+ checkLastColumnUpdateWriteLog(
+ updateWriteLogEntry(id = 1, metadata = Row("hr", null), data = Row(1, "hr", -1)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, -1, 100, "hr") :: Row(2, 2, 200, "software") :: Row(3, 3, 300, "hr") :: Nil)
+ }
+ }
+ }
+ }
+
+ test("column-update: CHAR and VARCHAR self-assignments are reported as updated") {
+ createTable("pk INT NOT NULL, id INT, c CHAR(3), v VARCHAR(10), dep STRING")
+ sql(s"INSERT INTO $tableNameAsString VALUES (1, 1, 'a', 'x', 'hr')")
+
+ sql(s"UPDATE $tableNameAsString SET id = id, c = c, v = v WHERE pk = 1")
+
+ checkLastUpdatedColumns("c", "v")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, 1, "a ", "x", "hr") :: Nil)
+ }
+
+ test("column-update: a LogicalWriteInfo that does not implement columnUpdateSchema throws") {
+ val info = new LogicalWriteInfo {
+ override def options(): CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty()
+ override def queryId(): String = "query"
+ override def schema(): StructType = new StructType()
+ }
+
+ checkError(
+ exception = intercept[SparkUnsupportedOperationException](info.columnUpdateSchema()),
+ condition = "DATA_SOURCE_COLUMN_UPDATE_SCHEMA_NOT_IMPLEMENTED",
+ parameters = Map("class" -> info.getClass.getName))
+ }
+
+ test("column-update: updatedColumns and the column update schema follow the non-identity " +
+ "assignments") {
+ Seq(
+ ("id = -1", Seq("id"), Seq("pk", "dep", "id"), Row(1, -1, "hr")),
+ ("ID = -1", Seq("id"), Seq("pk", "dep", "id"), Row(1, -1, "hr")),
+ ("id = -1, dep = 'eng'", Seq("id", "dep"), Seq("pk", "dep", "id"), Row(1, -1, "eng")),
+ ("id = -1, dep = dep", Seq("id"), Seq("pk", "dep", "id"), Row(1, -1, "hr")),
+ ("id = id, dep = 'eng'", Seq("dep"), Seq("pk", "dep"), Row(1, 10, "eng")),
+ ("id = id, dep = dep", Seq.empty, Seq("pk", "dep"), Row(1, 10, "hr")),
+ ("id = pk", Seq("id"), Seq("pk", "dep", "id"), Row(1, 1, "hr"))
+ ).foreach { case (assignments, updatedColumns, columnUpdateNames, updatedRow) =>
+ withTable(tableNameAsString) {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 10, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET $assignments WHERE pk = 1")
+
+ checkLastUpdatedColumns(updatedColumns: _*)
+ checkLastColumnUpdateSchema(columnUpdateNames: _*)
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), updatedRow :: Nil)
+ }
+ }
+ }
+
+ test("column-update split: row-level DELETE falls back to full-width rows") {
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin)
+ checkDeleteFallsBackToFullWidthRows(Row(2, 200, "hr"))
+ }
+
+ test("column-update split: row-level DELETE by a metadata row ID falls back to full-width " +
+ "rows") {
+ // Each partition holds one row, so `_partition` identifies it.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ splitPartitionRowIdTableProps)
+ checkDeleteFallsBackToFullWidthRows(Row(2, 200, "software"))
+ }
+
+ test("column-update split: MERGE falls back to full-width rows") {
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+ checkMergeFallsBackToFullWidthRows()
+ }
+
+ test("column-update: undeclared write requirement columns are rejected") {
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName
+ val splitReqAttrsOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithReqAttrs]
+ .getName
+ Seq(
+ (reqAttrsTableProps("pk,salary"), "dep", reqAttrsOperationClass, "[`dep`]"),
+ (splitReqAttrsTableProps("pk,salary"), "dep", splitReqAttrsOperationClass, "[`dep`]"),
+ (tablePropsWith(reqAttrsTableProps("pk,salary"),
+ "column-update-cluster-by-nested-transform" -> "true"), "dep", reqAttrsOperationClass,
+ "[`dep`]"),
+ (reqAttrsTableProps("pk,salary"), "id,dep", reqAttrsOperationClass, "[`dep`, `id`]"),
+ (reqAttrsTableProps("pk,salary,s"), "s.missing", reqAttrsOperationClass,
+ "[`s`.`missing` (FIELD_NOT_FOUND)]")
+ ).foreach { case (props, clusterBy, operationClass, columns) =>
+ withTable(tableNameAsString) {
+ val tableProps = tablePropsWith(props, "column-update-cluster-by" -> clusterBy)
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, s STRUCT<c1: INT>, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 1, "s": { "c1": 1 }, "dep": "hr" }
+ |""".stripMargin,
+ tableProps)
+
+ Seq(Nil, Seq(excludeColumnPruning)).foreach { confs =>
+ withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> operationClass, "columns" -> columns))
+ }
+ }
+ }
+ }
+ }
+
+ test("column-update: an empty write requirement column name is reported as <empty>") {
+ checkError(
+ exception = QueryCompilationErrors.undeclaredWriteRequirementColumnsError(
+ reqAttrsOperationClass, Seq(Seq.empty[String] -> None)).asInstanceOf[AnalysisException],
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[<empty>]"))
+ }
+
+ test("column-update: distribution by an undeclared row ID column is allowed") {
+ Seq("salary = -1", "salary = -1, pk = pk").foreach { assignments =>
+ withTable(tableNameAsString) {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin,
+ tablePropsWith(reqAttrsTableProps("salary,dep"), "column-update-cluster-by" -> "pk"))
+
+ sql(s"UPDATE $tableNameAsString SET $assignments WHERE pk = 1")
+
+ checkLastUpdatedColumns("salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 200, "hr") :: Nil)
+ }
+ }
+ }
+
+ test("column-update: an undeclared, unreferenced partition column is not read") {
+ createAndInitTableWithReqAttrs("pk,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1", "pk INT, _partition STRING")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively without adopting the " +
+ "declared spelling") {
+ createAndInitTableWithReqAttrs("PK,SALARY,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ checkLastColumnUpdateSchema("pk", "salary", "dep")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: DELETE on a scan that returns unrequested columns is unaffected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |""".stripMargin,
+ tablePropsWith(extraTableProps, InMemoryBaseTable.SIMULATE_PARTIAL_COLUMN_PRUNING -> "true"))
+
+ sql(s"DELETE FROM $tableNameAsString WHERE pk < 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(2, 200, 20, "software") :: Nil)
+ }
+
+ test("column-update: case-sensitive analysis accepts columns that differ only in case") {
+ withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") {
+ createAndInitTableWithReqAttrs("pk,salary,SALARY,dep",
+ "pk INT NOT NULL, salary INT, SALARY INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "SALARY": 200, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateSchema("pk", "salary", "SALARY", "dep")
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, -1, 200, "hr") :: Nil)
+ }
+ }
+
+ test("column-update: invalid requiredDataAttributes are rejected") {
+ Seq(
+ ("pk,s.c1,dep", "COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE",
+ "nestedAttributes" -> "[`s`.`c1`]"),
+ ("pk,pk,salary", "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ "duplicateAttributes" -> "[`pk`]"),
+ ("salary,PK,SALARY,pk", "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ "duplicateAttributes" -> "[`salary`, `PK`]"),
+ ("pk,_partition,salary", "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE",
+ "metadataAttributes" -> "[`_partition`]"),
+ ("pk,nonexistent_col,salary", "COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE",
+ "unknownAttributes" -> "[`nonexistent_col`]"),
+ ("pk,salary", "COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS",
+ "missingColumns" -> "[`a.b`]")
+ ).foreach { case (reqAttrs, condition, columns) =>
+ withTable(tableNameAsString) {
+ createAndInitTableWithReqAttrs(reqAttrs,
+ "pk INT NOT NULL, salary INT, s STRUCT<c1: INT, c2: INT>, `a.b` INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "s": { "c1": 1, "c2": 2 }, "a.b": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1, `a.b` = -1 WHERE pk = 1")
+ },
+ condition = condition,
+ parameters = Map("connector" -> reqAttrsOperationClass, columns))
+ }
+ }
+ }
+
+ test("column-update: a declared column with no name parts is reported as <empty>") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+ val t = table
+ val operation = new t.DeltaBasedColumnUpdateOperation(UPDATE) {
+ override def requiredDataAttributes(): Array[NamedReference] =
+ Array(FieldReference(Seq.empty[String]), FieldReference("pk"),
+ FieldReference(Seq.empty[String]))
+ }
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ RewriteRowLevelCommand.resolveRequiredDataAttrs(
+ spark.table(tableNameAsString).queryExecution.analyzed, operation)
+ },
+ condition = "COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> operation.getClass.getName, "unknownAttributes" -> "[<empty>, <empty>]"))
+ }
+
+ test("column-update split: delete and reinsert carry the row ID and preserved metadata") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ deleteWriteLogEntry(id = 1, metadata = Row("hr", null)),
+ reinsertWriteLogEntry(metadata = Row("hr", null), data = Row(1, "hr", -1)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "hr") :: Nil)
+ }
+
+ test("column-update split: row-ID reassignment on narrow write is rejected") {
+ // the declared [pk, dep, salary] leaves out `extra`, so the write is narrow
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING, extra STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr", "extra": "x" }
+ |{ "pk": 2, "salary": 200, "dep": "software", "extra": "y" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT",
+ parameters = Map("connector" -> splitOperationClass, "rowIds" -> "[`pk`]"))
+ }
+
+ test("column-update split: row-ID reassignment is allowed when requiredDataAttributes " +
+ "covers every table column") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ splitReqAttrsTableProps("pk,salary,dep"))
+
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(2, 200, "software") :: Row(11, -1, "hr") :: Nil)
+ }
+
+ test("column-update split: row-ID reassignment is allowed when requiredDataAttributes " +
+ "covers every data column of a relation with metadata columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ splitReqAttrsTableProps("pk,salary,dep"))
+
+ // `_partition` and the coerced `salary > 0.5` make the analyzer add the metadata columns to
Review Comment:
**71. This comment is no longer true since SPARK-60667, so the test no longer covers what its name says.** With the merge of master, `AddMetadataColumns` skips the target of a row-level command (`Analyzer.scala:1144` and `Analyzer.scala:1197-1198`), so the relation stays `[pk, salary, dep]` with no metadata columns in it. The statement should also run now: the delete and reinsert outputs of `buildDeletesAndInserts` have the same width, 6, and SPARK-60667's commit message lists this `ExpandExec` failure as fixed. `DeltaBasedUpdateAsDeleteAndInsertTableSuite.scala:98-128` already runs the same shape. As a result:
- The only assertion, at L453, checks the declaration, and the relation is now the same as in the test at L420. Removing the `filterNot` at `RewriteUpdateTable.scala:319`, the fix for [item 21](https://github.com/apache/spark/pull/55518#discussion_r4316536100), fails no test. The filter now matters only on a table that accepts any schema, where it swaps `COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT` for an `ExpandExec` failure.
- No narrow suite executes a later-iteration UPDATE that references a metadata column, because the SPARK-60667 tests are in concrete suites that the narrow suites do not extend.
- `executePlan(parsed).analyzed` at L449 leaves a transaction open, because the default mode begins one (`QueryExecution.scala:145-151`) and nothing commits or aborts it. It is harmless today only because `after` resets the catalog manager.
a6e22b24587 removed two other tests that carried this premise, and this one stayed. Could this test run the statement with `sql(...)` and check `Row(2, 200, "software") :: Row(11, -1, "hr") :: Nil`, with the name and the comment updated? `dataAttrs` at `RewriteUpdateTable.scala:319` could then be `relation.output`, which keeps the row-ID error on such a table.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -372,6 +875,66 @@ private class DeltaBufferWriter(schema: StructType) extends BufferWriter(schema)
override def commit(): WriterCommitMessage = buffer
}
+// Test-only: a writer factory for a connector that never overrides writeColumnUpdate(), to
+// exercise the DataWriter#writeColumnUpdate default that throws
+// DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED.
+// Declared top-level (no reference to the enclosing table) so it stays task-serializable.
+private class NoWriteColumnUpdateOverrideWriterFactory extends DataWriterFactory {
Review Comment:
**78.** Nit: since item 58 moved the `writeColumnUpdate` overrides off `BufferWriter`, this writer and its factory duplicate the plain writer that `TestBatchWrite` already creates (`InMemoryBaseTable.scala:1009-1011`), which ends in the same throw at `DataWriter.java:138-142`. L730 could be `super.createBatchWriterFactory(info)`, with L878-893 gone and `GroupBasedColumnUpdateTableSuite.scala:167` expecting `org.apache.spark.sql.connector.catalog.BufferWriter`. That test would then also fail if the overrides came back to `BufferWriter`, which no test catches today.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupBasedRowLevelOperationScanPlanning.scala:
##########
@@ -65,7 +65,9 @@ object GroupBasedRowLevelOperationScanPlanning extends Rule[LogicalPlan] with Pr
.mkString(", ")
}
- val (scan, output) = PushDownUtils.pruneColumns(scanBuilder, relation, relation.output, Nil)
+ val requiredAttrs =
+ if (V2Writes.hasNarrowRows(rd)) readAttrs(rd.query, relation) else relation.output
Review Comment:
**77.** Nit, and fine to leave as is or to do separately: this gate ties a scan rule to the write layout, although `readAttrs` gives the same columns for wide writes unless the condition folds to true. Copied rows of DELETE and MERGE reference every read column, and for UPDATE only `If(TrueLiteral, trueValue, _) => trueValue` drops one (`optimizer/expressions.scala:713`), in which case no COPY row needs it. So a wide `UPDATE t SET big_col = NULL` with no WHERE clause reads `big_col` for every row, while delta and narrow copy-on-write scans do not. `val requiredAttrs = readAttrs(rd.query, relation)` would remove that difference and match `ColumnPruning`, which applies to every write since item 11. It is not free of behavior change, though: a scan with unconditional `filterAttributes()` would then fail `UPDATE t SET dep = 'x'` with `DATA_SOURCE_INVALID_RUNTIME_FILTER_ATTRIBUTE.CANNOT_RESOLVE` (`DataSourceV2Strategy.scala:207`), where master runs it.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -354,3 +364,69 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
}
}
}
+
+object RewriteRowLevelCommand extends SQLConfHelper {
+
+ /**
+ * Whether the operation narrows the write: it mixes in `SupportsColumnUpdates` and declares at
+ * least one column. An operation that declares none receives full-width rows.
+ */
+ private[sql] def narrowsWrite(operation: RowLevelOperation): Boolean = operation match {
+ case op: SupportsColumnUpdates => op.requiredDataAttributes.nonEmpty
+ case _ => false
+ }
+
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation, in declared order and with the relation's spelling. Nested fields,
+ * duplicate columns, columns that do not exist in the relation (including a reference with no
+ * name parts) and metadata columns are rejected with an `AnalysisException`.
+ */
+ private[sql] def resolveRequiredDataAttrs(
+ relation: LogicalPlan,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes.toImmutableArraySeq
+ val nameless = refs.filter(_.fieldNames.isEmpty).map(_ => Seq.empty[String])
+ if (nameless.nonEmpty) {
+ throw QueryCompilationErrors.unknownRequiredDataAttributeError(
+ operation.getClass.getName, nameless)
+ }
+ val nested = refs.filter(_.fieldNames.length > 1).map(_.fieldNames.toSeq)
+ if (nested.nonEmpty) {
+ throw QueryCompilationErrors.nestedRequiredDataAttributeError(
+ operation.getClass.getName, nested)
+ }
+ val normalizedNames =
+ refs.map(ref => SchemaUtils.foldName(ref.fieldNames.head, conf.caseSensitiveAnalysis))
Review Comment:
**68. The duplicate check treats `foldName` as column identity, which its Scaladoc says it is not.** `SchemaUtils.scala:541-549` says "The fold alone is therefore not an identity" and "Callers must narrow a multi-candidate match with the resolver". In a case-insensitive session, a column named U+0130 and a column named `i` plus U+0307 fold to the same string while `equalsIgnoreCase` keeps them apart. Each reference resolves to its own column at L410, yet declaring both fails here with `COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE`, whose message says "Each column must be declared at most once". This is rare and it fails loudly, so it is low priority. Moving the check after L418 and counting `resolved` by exprId, filtering `refs.zip(resolved)` in order, would use the resolver's answer and keep `` [`salary`, `PK`] ``. Two messages would change, and no test pins either: `x,x` would become `COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE`, and case-insensitive `salary,SALARY` on a ta
ble with both columns would become `AMBIGUOUS_REFERENCE` from L410.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java:
##########
@@ -37,4 +38,22 @@ public interface RowLevelOperationInfo {
* Returns the row-level SQL command (e.g. DELETE, UPDATE, MERGE).
*/
Command command();
+
+ /**
+ * Returns the columns being updated by this operation. Currently only UPDATE populates it;
+ * other commands report an empty array.
+ * <p>
+ * A column is reported unless its assigned value is the column itself, so identity assignments
+ * such as {@code SET a = a} are excluded. With the default CHAR and VARCHAR handling, a
+ * self-assignment on a CHAR or VARCHAR column is still reported, because Spark pads or
+ * length-checks the assigned value. Nested struct field updates are reported at root-column
+ * granularity (e.g. {@code SET s.c1 = -1} returns {@code s}), even when the field is assigned to
Review Comment:
**70. The PR description's exception list for identity assignments leaves out nested self-assignments.** It says "Identity assignments such as `SET a = a` are excluded, except self-assignments on CHAR or VARCHAR columns with the default CHAR and VARCHAR handling", which reads as a closed list, but `SET s.c1 = s.c1` is reported as `s`, as this Javadoc and `ColumnUpdateTableSharedTests.scala:362-371` say. An operation that declares `[pk, dep]` therefore rejects that statement with `COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS`. Could the description mention it? Reporting `s` looks right to me: `applyFieldAssignments` (`AssignmentUtils.scala:222`) has no null guard, so a NULL `s` becomes `{null, null}`, as "update fields inside NULL structs" (`UpdateTableSuiteBase.scala:456-467`) implies. For the same reason, "although its values don't change" at `ColumnUpdateTableSharedTests.scala:367-368` holds only for a non-NULL `s`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -159,14 +164,24 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
relation: DataSourceV2Relation,
operationTable: RowLevelOperationTable,
assignments: Seq[Assignment],
- cond: Expression): WriteDelta = {
+ cond: Expression,
+ writeAttrs: Seq[AttributeReference],
+ updatedAttrs: Seq[AttributeReference]): WriteDelta = {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
// resolve all needed attrs (e.g. row ID and any required metadata attrs)
- val rowAttrs = relation.output
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
+ val splitUpdates = operation.representUpdateAsDeleteAndInsert
+
+ operation match {
+ case op: SupportsColumnUpdates if splitUpdates && RewriteRowLevelCommand.narrowsWrite(op) =>
Review Comment:
**72. No test covers `&& RewriteRowLevelCommand.narrowsWrite(op)`.** Every test of an empty declaration (`reqAttrsTableProps("")` at `ColumnUpdateTableSharedTests.scala:237` and `ColumnUpdateTableSharedTests.scala:293`) is non-split, and every split test declares at least one column, so reverting this line to `case op: SupportsColumnUpdates if splitUpdates =>` fails nothing. The difference shows on `t(pk INT NOT NULL, salary INT, dep STRING)` with `column-update-split-row-id = _partition, no-metadata = true`: `UPDATE t SET salary = salary WHERE pk = 1` analyzes to a full-width `WriteDelta` with the guard, and fails with `COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED` without it. This is the `DeltaBasedColumnUpdateSplitOperationWithRowId` shape from item 55. The fixture's commit throws `IllegalStateException` for a matching row (`InMemoryRowLevelOperationTable.scala:561-564`), so a test would have to check the analyzed plan.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/V2Writes.scala:
##########
@@ -109,22 +112,95 @@ object V2Writes extends Rule[LogicalPlan] with PredicateHelper {
WriteToDataSourceV2(Some(r), microBatchWrite, newQuery, customMetrics)
case rd @ ReplaceData(r: DataSourceV2Relation, _, query, _, projections, _, None) =>
- val rowSchema = projections.rowProjection.schema
+ validateWriteRelation(rd)
+ val (rowSchema, columnUpdateSchema) =
+ writeSchemas(projections.rowProjection.schema, hasNarrowRows(rd))
val metadataSchema = projections.metadataProjection.map(_.schema)
val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap)
- val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema)
+ val writeBuilder = newWriteBuilder(r.table, writeOptions, rowSchema, metadataSchema,
+ columnUpdateSchema)
val write = writeBuilder.build()
+ validateColumnUpdateWrite(rd, write)
val newQuery = DistributionAndOrderingUtils.prepareQuery(write, query, r.funCatalog)
rd.copy(write = Some(write), query = newQuery)
case wd @ WriteDelta(r: DataSourceV2Relation, _, query, _, projections, _, None) =>
+ validateWriteRelation(wd)
val writeOptions = mergeOptions(Map.empty, r.options.asCaseSensitiveMap.asScala.toMap)
- val deltaWriteBuilder = newDeltaWriteBuilder(r.table, writeOptions, projections)
+ val deltaWriteBuilder =
+ newDeltaWriteBuilder(r.table, writeOptions, projections, hasNarrowRows(wd))
val deltaWrite = deltaWriteBuilder.build()
+ validateColumnUpdateWrite(wd, deltaWrite)
val newQuery = DistributionAndOrderingUtils.prepareQuery(deltaWrite, query, r.funCatalog)
wd.copy(write = Some(deltaWrite), query = newQuery)
}
+ /**
+ * Whether a row-level write delivers narrow rows for updated, copied, and reinserted records.
+ * This must match the writes whose write relation RewriteUpdateTable narrows, even when the
+ * declared columns cover the whole table.
+ */
+ private[v2] def hasNarrowRows(write: RowLevelWrite): Boolean = {
+ write.operation.command == UPDATE && RewriteRowLevelCommand.narrowsWrite(write.operation)
+ }
+
+ /**
+ * Checks that the write relation matches the current operation: a write that delivers narrow
+ * rows must have exactly the declared columns, in declared order, and any other write the full
+ * table columns. A rule that replaces the operation after the rewrite can break this, and the
+ * connector would then decode the rows with the wrong layout.
+ */
+ private def validateWriteRelation(rowLevelWrite: RowLevelWrite): Unit = {
+ val writeNames = rowLevelWrite.table.output.map(_.name)
+ val expectedNames = rowLevelWrite.operation match {
+ case op: SupportsColumnUpdates if hasNarrowRows(rowLevelWrite) =>
+ RewriteRowLevelCommand.resolveRequiredDataAttrs(rowLevelWrite.originalTable, op)
+ .map(_.name)
+ case _ =>
+ rowLevelWrite.originalTable.output.map(_.name)
+ }
+ if (writeNames != expectedNames) {
+ throw SparkException.internalError(
+ s"The write relation of operation ${rowLevelWrite.operation.getClass.getName} has " +
+ s"columns ${writeNames.mkString("[", ", ", "]")}, but the operation expects " +
+ s"${expectedNames.mkString("[", ", ", "]")}.")
+ }
+ }
+
+ /**
+ * Checks that the distribution and ordering of a write that delivers narrow rows reference only
+ * columns passed to the write. Other columns stay in the query only until column pruning removes
+ * them, so they are rejected here whether or not pruning ran.
+ */
+ private def validateColumnUpdateWrite(rowLevelWrite: RowLevelWrite, write: Write): Unit = {
+ write match {
+ case w: RequiresDistributionAndOrdering if hasNarrowRows(rowLevelWrite) =>
+ val distribution =
+ DistributionAndOrderingUtils.distributionExpressions(w.requiredDistribution)
+ val readAttrs = rowLevelWrite.references.toSeq.distinctBy(_.name)
Review Comment:
**75.** Nit: `distinctBy(_.name)` and the `sortBy` at L184 do nothing this check needs. Every read ordinal comes from `indexWhere` with the resolver (`RewriteRowLevelCommand.scala:332`), so two attributes in `references` cannot share a name, and `LocalRelation(rowLevelWrite.references.toSeq)`, as in item 7, is enough. If such a path ever appeared, `distinctBy` would let the column pass here and fail with `AMBIGUOUS_REFERENCE` in `prepareQuery`. `unresolvableReferences` already walks the expressions in order and ends with `.distinct`, so the list is deterministic without the sort, and the sort makes this the only `COLUMN_UPDATE_*` list not in the connector's or the assignments' order: `column-update-cluster-by = id,dep` reports `` [`dep`, `id`] ``. Dropping it changes the expected value at `DeltaBasedColumnUpdateTableSuite.scala:228`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -237,4 +253,101 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Returns the table attributes that these assignments set to a new value, skipping identity
+ * assignments.
+ */
+ private def collectUpdatedAttrs(assignments: Seq[Assignment]): Seq[AttributeReference] = {
+ assignments.collect {
+ case Assignment(key: AttributeReference, value) if !isIdentityAssignment(key, value) => key
+ }
+ }
+
+ /**
+ * Returns the output of the write relation, which is the layout of the data rows the connector
+ * receives. If the operation opts into column updates, this is the connector's
+ * `requiredDataAttributes()`. An empty declaration yields the full relation output, as does any
+ * operation without the mix-in. A mixed-in operation's declaration must cover every updated
+ * column, even an empty one.
+ */
+ private def resolveWriteAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation,
+ updatedAttrs: Seq[AttributeReference]): Seq[AttributeReference] = operation match {
+ case op: SupportsColumnUpdates =>
+ val attrs = if (RewriteRowLevelCommand.narrowsWrite(op)) {
Review Comment:
**76.** Nit: this branch is no longer needed. Since 8d9856dd55b removed the `refs.isEmpty` throw, `resolveRequiredDataAttrs` returns an empty `Seq` for an empty declaration, so L279-283 can be `val attrs = RewriteRowLevelCommand.resolveRequiredDataAttrs(relation, op)`, with L284-285 unchanged. Similarly, the `operation match` at L178-184 can be `if (splitUpdates && RewriteRowLevelCommand.narrowsWrite(operation))`, because `narrowsWrite` takes a `RowLevelOperation` and both validators use the operation only for its class name. `writeRelation.output` at L88, L127 and L203 is `writeAttrs`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala:
##########
@@ -111,10 +112,10 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla
// and must transform the runtime filter condition to use correct expr IDs for each relation
// note this only applies to group-based row-level operations (i.e. ReplaceData)
// see RewriteUpdateTable for more details
- // the condition may also reference metadata columns resolved from the metadata output
- // of the target relation, which are not in the table output and must be mapped too
- val metadataAttrs = (cond.references -- write.table.output).toSeq
- val attrMap = buildTableToScanAttrMap(write.table.output ++ metadataAttrs, relation.output)
+ // the write table may leave out data columns, so map those of the original table
Review Comment:
**74.** Nit: since SPARK-60664 this hunk no longer changes the result, and it dropped master's comment. `attrMap` only rewrites attributes in `cond` (L119-121), and master's `cond.references -- write.table.output` already includes an undeclared column such as `id`, so `GroupBasedColumnUpdateTableSuite.scala:193` should behave the same without this change. Master's "the condition may also reference metadata columns resolved from the metadata output of the target relation, which are not in the table output and must be mapped too" is gone although L117 still does that, and [the reply to item 57](https://github.com/apache/spark/pull/55518#discussion_r4426726027) covers the code but not the comment. Could this be `val attrMap = buildTableToScanAttrMap(cond.references.toSeq, relation.output)`, which needs neither `tableAttrs`, `metadataAttrs` nor either comment? Keeping these lines and restoring the comment works too.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3634060727
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/dynamicpruning/RowLevelOperationRuntimeGroupFiltering.scala:
##########
@@ -139,17 +138,13 @@ class RowLevelOperationRuntimeGroupFiltering(optimizeSubqueries: Rule[LogicalPla
tableAttrs: Seq[Attribute],
scanAttrs: Seq[Attribute]): AttributeMap[Attribute] = {
- val attrMapping = tableAttrs.map { tableAttr =>
+ // The scan may be narrowed to exclude columns not needed by the connector.
+ // Attributes absent from the scan are skipped here; the caller must ensure
+ // that any attribute referenced in the condition is present in the scan.
+ val attrMapping = tableAttrs.flatMap { tableAttr =>
scanAttrs
.find(scanAttr => conf.resolver(scanAttr.name, tableAttr.name))
.map(scanAttr => tableAttr -> scanAttr)
- .getOrElse {
- throw new AnalysisException(
- errorClass = "_LEGACY_ERROR_TEMP_3075",
- messageParameters = Map(
- "tableAttr" -> tableAttr.toString,
- "scanAttrs" -> scanAttrs.mkString(",")))
- }
}
Review Comment:
No. Unfortunately, this PR should not remove this because the existing sanity check is used for other code path in the existing test cases. Please recover it.
I guess you may achieve your goal via the following. Please review and revise the following example for your purpose.
```scala
private def buildTableToScanAttrMap(
tableAttrs: Seq[Attribute],
scanAttrs: Seq[Attribute],
requiredAttrs: AttributeSet): AttributeMap[Attribute] = {
// Table attrs may be legitimately absent from a column-update narrowed scan, so map only
// those that have a matching scan attribute. Attrs referenced by the condition must always
// be present (computeNarrowReadAttrs keeps them in the scan); failing to map one would
// leave a dangling reference in the group filter, so keep the strict check for them.
val attrMapping = tableAttrs.flatMap { tableAttr =>
val matched = scanAttrs.find(scanAttr => conf.resolver(scanAttr.name, tableAttr.name))
if (matched.isEmpty && requiredAttrs.contains(tableAttr)) {
throw new AnalysisException(
errorClass = "_LEGACY_ERROR_TEMP_3075",
messageParameters = Map(
"tableAttr" -> tableAttr.toString,
"scanAttrs" -> scanAttrs.mkString(",")))
}
matched.map(scanAttr => tableAttr -> scanAttr)
}
AttributeMap(attrMapping)
}
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3633921145
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,540 @@
+/*
+ * 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, TableInfo}
+import org.apache.spark.sql.connector.expressions.LogicalExpressions.{identity, reference}
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+
+/**
+ * Tests for UPDATE statements targeting connectors that return true from
+ * [[org.apache.spark.sql.connector.write.RowLevelOperation#supportsColumnUpdates]].
+ *
+ * When a connector supports column updates, Spark narrows the row projection
+ * (LogicalWriteInfo.schema()) to contain only the assigned/changed columns rather than
+ * the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends RowLevelOperationSuiteBase {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ test("column-update: rowSchema contains only the single assigned column") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ // info.schema() is empty: UPDATE-only column-update writes have no INSERT-shaped rows,
+ // so nothing flows through the writer's write() path. The narrow update-row layout is
+ // carried by info.updateSchema().
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: rowSchema contains multiple assigned columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("id", IntegerType, nullable = false),
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: rowSchema is empty for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ // All assignments are identity, so updatedColumns is empty -- the connector still declares
+ // pk for row lookup, so the narrow update schema is just [pk].
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Array(PK_FIELD))))
+ }
+
+ test("column-update: row filter condition is orthogonal to column narrowing") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = 'engineering' WHERE pk IN (1, 3)")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: update all rows (no WHERE clause)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("salary", IntegerType, nullable = true)
+ ))))
+ }
+
+ test("column-update: rowSchema excludes identity assignments in a mixed UPDATE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: cross-column assignment is not treated as identity") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = dep, id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: nested struct field update narrows to the root struct column") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("s"),
+ s"expected [s] in updatedColumns (root struct) but got: $updatedNames")
+
+ // info.updateSchema() carries the narrow row layout; info.schema() is the full table.
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("s"),
+ s"s must be in update schema: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("dep"),
+ s"dep must not be in update schema: $updateSchema")
+ }
+
+ test("column-update: nested field identity update reports root struct as updated") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = s.c1 WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("s"),
+ s"nested identity is reported as an update at root-column granularity: $updatedNames")
+
+ // Data correctness: the struct is rewritten but with equal values, so rows are unchanged.
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, Row(1, 2), "hr") :: Nil)
+ }
+
+ test("column-update: updatedColumns contains non-identity assigned columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'eng' WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("id", "dep"),
+ s"expected [id, dep] in updatedColumns but got: $updatedNames")
+ }
+
+ test("column-update: updatedColumns excludes identity assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = dep WHERE pk = 1")
+
+ val updatedNames = table.lastUpdatedColumns.map(_.describe()).toSet
+ assert(updatedNames == Set("id"),
+ s"expected only [id] in updatedColumns but got: $updatedNames")
+ }
+
+ test("column-update: updatedColumns is empty for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ assert(table.lastUpdatedColumns.isEmpty,
+ s"expected empty updatedColumns but got: ${table.lastUpdatedColumns.mkString(", ")}")
+ }
+
+ test("column-update: updatedColumns is empty for DELETE (Javadoc contract)") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"DELETE FROM $tableNameAsString WHERE dep = 'hr'")
+
+ assert(table.lastUpdatedColumns.isEmpty,
+ s"DELETE must pass empty updatedColumns but got: ${table.lastUpdatedColumns.mkString(", ")}")
+ }
+
+ test("column-update: data correctness -- single column update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ // 1 row matched WHERE pk = 1 -> 1 update. MoR emits no COPY rows.
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0, deltaUpdate = true)
+ }
+
+ test("column-update: data correctness -- update all rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 200, "hr") :: Row(2, 400, "software") :: Row(3, 600, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 3, numCopiedRows = 0, deltaUpdate = true)
+ }
+
+ test("column-update: data correctness -- mixed identity and real assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "engineering") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0, deltaUpdate = true)
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-req-attrs", reqAttrs)
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update: requiredDataAttributes - data correctness") {
+ createAndInitTableWithReqAttrs("dep,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: empty requiredDataAttributes throws AnalysisException") {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-empty-req-attrs", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(
+ StructType.fromDDL("pk INT NOT NULL, id INT, dep STRING"))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }""".stripMargin)
+
+ // A connector that mixes in SupportsColumnUpdates but returns an empty
+ // requiredDataAttributes() violates the mix-in contract -- analysis must reject the
+ // operation rather than silently falling through to the wide write path.
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ }
+ assert(ex.getCondition == "EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ s"expected EMPTY_REQUIRED_DATA_ATTRIBUTES but got: ${ex.getCondition}")
+ }
+
+ test("column-update: requiredDataAttributes throws AnalysisException for invalid column") {
+ createAndInitTableWithReqAttrs("nonexistent_col", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ val ex = intercept[org.apache.spark.sql.AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ }
+ assert(ex.getMessage.contains("nonexistent_col"),
+ s"Expected error about unresolvable column but got: ${ex.getMessage}")
+ }
+
+ private def createAndInitTableFromInfo(schemaString: String, jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-from-info", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update from-info: write schema is updatedColumns + pk pass-through") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ assert(updateSchema.fieldNames.contains("salary"),
+ s"salary must be in update schema: $updateSchema")
+ assert(updateSchema.fieldNames.contains("pk"),
+ s"pk must be in update schema: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("id"),
+ s"id must not be in update schema: $updateSchema")
+ assert(!updateSchema.fieldNames.contains("dep"),
+ s"dep must not be in update schema: $updateSchema")
+ }
+
+ test("column-update from-info: pk already in updatedColumns is not duplicated") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+
+ val updateSchema = table.lastWriteInfo.updateSchema().get()
+ val pkCount = updateSchema.fieldNames.count(_ == "pk")
+ assert(pkCount == 1, s"pk must appear exactly once in update schema: $updateSchema")
+ }
+
+ test("column-update from-info: data correctness") {
+ createAndInitTableFromInfo("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "id": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ }
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String): Unit = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-split", "true")
+ val columns = CatalogV2Util.structTypeToV2Columns(StructType.fromDDL(schemaString))
+ val transforms = Array[Transform](identity(reference(Seq("dep"))))
+ val tableInfo = new TableInfo.Builder()
+ .withColumns(columns)
+ .withPartitions(transforms)
+ .withProperties(props)
+ .build()
+ catalog.createTable(ident, tableInfo)
+ append(schemaString, jsonData)
+ }
+
+ test("column-update split: write schema is narrow (assigned + pk pass-through)") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ // For representUpdateAsDeleteAndInsert connectors, reinsert-tagged rows still route through
+ // writeUpdate() (info.updateSchema() = narrow), not write(). info.schema() stays empty.
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update split: data correctness") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, -1, "hr") :: Nil)
+ }
+
+ // ---------------------------------------------------------------------------
+ // Scan-narrowing tests: verify that when the connector opts into column updates,
+ // the physical scan reads only the columns actually needed by the operation.
+ //
+ // Assertions use `checkLastScanExcludes` (proves narrowing) and `checkLastScanIncludes`
+ // (proves transparent widening) rather than exact-set matches -- ColumnPruning is
+ // free to further tighten the scan beyond our analysis-time narrow set (e.g. when an
+ // assignment RHS is a literal, the assigned column doesn't need to be read).
+ // ---------------------------------------------------------------------------
+
+ test("column-update: scan excludes columns outside connector-declared + cond + RHS refs") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ // requiredDataAttributes = [pk, id]; cond refs [pk] only; RHS is a literal.
+ // `dep` is neither declared nor referenced -- it must not appear in the scan.
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastScanExcludes("dep")
+ }
+
+ test("column-update: scan transparently widens for cond-referenced columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ // requiredDataAttributes = [pk, id]; cond refs [dep].
+ // `dep` must appear even though the connector didn't declare it.
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE dep = 'hr'")
+
+ checkLastScanIncludes("dep")
+ }
+
+ test("column-update: scan transparently widens for assignment RHS references") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin)
+
+ // requiredDataAttributes = [pk, salary]; RHS `salary + bonus` references `bonus`.
+ // `bonus` must appear in the scan even though it's not declared as required.
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE pk = 1")
+
+ checkLastScanIncludes("bonus", "salary")
+ checkLastScanExcludes("dep")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 110, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update: analysis fails when assignment key is outside requiredDataAttributes") {
+ // Connector declares only [pk] but the user assigns to `id`. We enforce
+ // updatedColumns ⊆ requiredDataAttributes at analysis time (root-column granularity).
Review Comment:
```suggestion
// updatedColumns is a subset of requiredDataAttributes at analysis time
// (root-column granularity).
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-5055078133
Thanks for the reviews @dongjoon-hyun and @peter-toth. I addressed your comments.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3575237230
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -50,20 +51,36 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
protected def buildOperationTable(
table: SupportsRowLevelOperations,
command: Command,
- options: CaseInsensitiveStringMap): RowLevelOperationTable = {
- val info = RowLevelOperationInfoImpl(command, options)
+ options: CaseInsensitiveStringMap,
+ updatedColumns: Seq[NamedReference] = Nil): RowLevelOperationTable = {
+ val info = RowLevelOperationInfoImpl(command, options, updatedColumns)
val operation = table.newRowLevelOperationBuilder(info).build()
RowLevelOperationTable(table, operation)
}
+ /**
+ * Builds a DataSourceV2Relation for a row-level operation, optionally narrowing its output.
+ *
+ * When dataAttrs is non-empty, the relation output is narrowed to include only columns
+ * required for a column-update write. When dataAttrs is empty, the full relation.output is
+ * preserved.
+ */
protected def buildRelationWithAttrs(
relation: DataSourceV2Relation,
table: RowLevelOperationTable,
metadataAttrs: Seq[AttributeReference],
- rowIdAttrs: Seq[AttributeReference] = Nil): DataSourceV2Relation = {
-
- val attrs = dedupAttrs(relation.output ++ rowIdAttrs ++ metadataAttrs)
- relation.copy(table = table, output = attrs)
+ rowIdAttrs: Seq[AttributeReference] = Nil,
+ dataAttrs: Seq[AttributeReference] = Nil,
+ cond: Expression = TrueLiteral): DataSourceV2Relation = {
Review Comment:
I added a narrow version of this method which has a different implementation.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-5225987325
> @anuragmantri, I believe you forgot to push the new commit.
Sorry, I pushed now in [774412c](https://github.com/apache/spark/pull/55518/commits/774412c5c1f3ea575faccdd068b7d2bbe417fe97) @peter-toth
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3739282713
##########
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:
I added the validation a not implicitly adding the partition column since it makes the `scanOnlyDataAttributes` API tightly scoped for this particular case.
##########
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:
Thanks for the finding, I have widened the validation and also added a test that you had described.
##########
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:
Done
##########
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:
You are right, changed the wording in both places.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3423916257
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -154,30 +193,114 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
cond: Expression): WriteDelta = {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
+ val supportsColumnUpdate = operation.supportsColumnUpdates()
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
- val rowAttrs = relation.output
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveRequiredDataAttrs(relation, operation)
+ } else Nil
+
val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ // Connector-declared attrs not being assigned are passed through so ColumnPruning
+ // keeps them in the scan and the connector receives their current values.
+ val assignedAttrs = if (supportsColumnUpdate) computeAssignedAttrs(assignments)
+ else relation.output
+ val connectorExtraAttrs: Seq[AttributeReference] = if (connectorDataAttrs.nonEmpty) {
+ val assignedAttrSet = AttributeSet(assignedAttrs)
+ connectorDataAttrs.filterNot(assignedAttrSet.contains)
+ } else Nil
+
// build a plan for updated records that match the condition
val matchedRowsPlan = Filter(cond, readRelation)
val rowDeltaPlan = if (operation.representUpdateAsDeleteAndInsert) {
buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else if (supportsColumnUpdate) {
+ buildColumnUpdateProjection(
+ matchedRowsPlan, assignments, rowIdAttrs, metadataAttrs, connectorExtraAttrs)
} else {
buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs)
}
+ val effectiveRowAttrs = if (supportsColumnUpdate && connectorDataAttrs.nonEmpty) {
+ connectorDataAttrs
+ } else if (supportsColumnUpdate) {
+ assignedAttrs
+ } else {
+ relation.output
+ }
+
// 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(
+ rowDeltaPlan, effectiveRowAttrs, rowIdAttrs, metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond)
}
+ /**
+ * Builds the WriteDelta projection for the column update path. The resulting Project
+ * references only the columns needed for the write, so ColumnPruning narrows the scan
+ * to match.
+ */
+ private def buildColumnUpdateProjection(
+ plan: LogicalPlan,
+ assignments: Seq[Assignment],
+ rowIdAttrs: Seq[Attribute],
+ metadataAttrs: Seq[Attribute],
+ connectorExtraAttrs: Seq[AttributeReference] = Nil): LogicalPlan = {
+
+ val assignedValues = assignments.collect {
+ case Assignment(key: Attribute, value) if !isIdentityAssignment(key, value) =>
+ Alias(value, key.name)()
+ }
+
+ val connectorExtraAttrSet = AttributeSet(connectorExtraAttrs)
+ val connectorPassThroughValues = plan.output.filter { a =>
+ connectorExtraAttrSet.contains(a) && !MetadataAttribute.isValid(a.metadata)
+ }
+
+ val metadataAttrSet = AttributeSet(metadataAttrs)
+ val metadataValues = plan.output.filter(metadataAttrSet.contains).map { attr =>
+ if (MetadataAttribute.isPreservedOnUpdate(attr)) {
+ attr
+ } else {
+ Alias(Literal(null, attr.dataType), attr.name)(explicitMetadata = Some(attr.metadata))
+ }
+ }
+
+ val rowIdAttrSet = AttributeSet(rowIdAttrs)
+ val rowIdValues = plan.output.filter(rowIdAttrSet.contains)
+
+ val originalRowIdValues = buildOriginalRowIdValues(rowIdAttrs, assignments)
+ val operationType = Alias(Literal(UPDATE_OPERATION), OPERATION_COLUMN)()
+
+ Project(
+ Seq(operationType) ++ assignedValues ++ connectorPassThroughValues ++
+ metadataValues ++ rowIdValues ++ originalRowIdValues,
+ plan)
+ }
+
+ // Returns the table attributes that are genuinely updated (non-identity) in this UPDATE.
+ private def computeAssignedAttrs(assignments: Seq[Assignment]): Seq[AttributeReference] = {
+ assignments.collect {
+ case Assignment(key: AttributeReference, value) if !isIdentityAssignment(key, value) => key
+ }
+ }
+
+ private def isIdentityAssignment(key: Attribute, value: Expression): Boolean = {
Review Comment:
Does this handle nested columns?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3423889877
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -50,20 +51,36 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
protected def buildOperationTable(
table: SupportsRowLevelOperations,
command: Command,
- options: CaseInsensitiveStringMap): RowLevelOperationTable = {
- val info = RowLevelOperationInfoImpl(command, options)
+ options: CaseInsensitiveStringMap,
+ updatedColumns: Seq[NamedReference] = Nil): RowLevelOperationTable = {
+ val info = RowLevelOperationInfoImpl(command, options, updatedColumns)
val operation = table.newRowLevelOperationBuilder(info).build()
RowLevelOperationTable(table, operation)
}
+ /**
+ * Builds a DataSourceV2Relation for a row-level operation, optionally narrowing its output.
+ *
+ * When dataAttrs is non-empty, the relation output is narrowed to include only columns
+ * required for a column-update write. When dataAttrs is empty, the full relation.output is
+ * preserved.
+ */
protected def buildRelationWithAttrs(
relation: DataSourceV2Relation,
table: RowLevelOperationTable,
metadataAttrs: Seq[AttributeReference],
- rowIdAttrs: Seq[AttributeReference] = Nil): DataSourceV2Relation = {
-
- val attrs = dedupAttrs(relation.output ++ rowIdAttrs ++ metadataAttrs)
- relation.copy(table = table, output = attrs)
+ rowIdAttrs: Seq[AttributeReference] = Nil,
+ dataAttrs: Seq[AttributeReference] = Nil,
+ cond: Expression = TrueLiteral): DataSourceV2Relation = {
Review Comment:
Don't think cond should be passed here, I'd compute prior to calling.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3423944549
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -106,38 +105,78 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val remainingRowsPlan = addOperationColumn(COPY_OPERATION,
Filter(remainingRowFilter, readRelation))
- // the new state is a union of updated and copied over records
- val query = Union(updatedRowsPlan, remainingRowsPlan)
+ val updatedAndRemainingRowsPlan = Union(updatedRowsPlan, remainingRowsPlan)
- // build a plan to replace read groups in the table
val writeRelation = relation.copy(table = operationTable)
- val projections = buildReplaceDataProjections(query, relation.output, metadataAttrs)
+ val projections = buildReplaceDataProjections(updatedAndRemainingRowsPlan, rowAttrs,
+ metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
- ReplaceData(writeRelation, cond, query, relation, projections, groupFilterCond)
+ ReplaceData(writeRelation, cond, updatedAndRemainingRowsPlan, relation, projections,
+ groupFilterCond)
}
- // this method assumes the assignments have been already aligned before
+ /**
+ * When the connector supports column updates and declares required data attributes,
+ * the read relation is narrowed at analysis time so that GroupBasedRowLevelOperationScanPlanning
+ * uses only the needed columns for the scan. Otherwise, the full relation output is used.
+ */
+ private def buildReplaceDataReadRelation(
+ relation: DataSourceV2Relation,
+ operationTable: RowLevelOperationTable,
+ assignments: Seq[Assignment],
+ cond: Expression): (DataSourceV2Relation, Seq[Attribute], Seq[AttributeReference]) = {
+
+ val operation = operationTable.operation
+ val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
+ val connectorDataAttrs = resolveRequiredDataAttrs(relation, operation)
+ val isNarrow = operation.supportsColumnUpdates() && connectorDataAttrs.nonEmpty
+
+ val readRelation = if (isNarrow) {
+ val allRequired = (connectorDataAttrs ++ computeAssignedAttrs(assignments)).distinct
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs, dataAttrs = allRequired,
+ cond = cond)
+ } else {
+ buildRelationWithAttrs(relation, operationTable, metadataAttrs)
+ }
+
+ val rowAttrs: Seq[Attribute] = if (isNarrow) connectorDataAttrs else relation.output
+
+ (readRelation, rowAttrs, metadataAttrs)
+ }
+
+ /**
+ * Builds the update projection for ReplaceData plans. Assumes assignments are already aligned.
+ *
+ * plan.output may be narrowed by buildReplaceDataReadRelation, so only columns present in the
+ * plan are projected.
+ */
private def buildReplaceDataUpdateProjection(
Review Comment:
This is fragile too.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3424001894
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -106,38 +105,78 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val remainingRowsPlan = addOperationColumn(COPY_OPERATION,
Filter(remainingRowFilter, readRelation))
- // the new state is a union of updated and copied over records
- val query = Union(updatedRowsPlan, remainingRowsPlan)
+ val updatedAndRemainingRowsPlan = Union(updatedRowsPlan, remainingRowsPlan)
- // build a plan to replace read groups in the table
val writeRelation = relation.copy(table = operationTable)
Review Comment:
Why don't we want to narrow down the expected schema in the write relation too in this case?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192401119
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -154,30 +206,145 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
cond: Expression): WriteDelta = {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
+ // Column-update support applies to the standard delta path and the delete+reinsert path.
+ // When representUpdateAsDeleteAndInsert is true, the REINSERT leg of the Expand already
+ // uses only assigned values, so the narrow effectiveRowAttrs applies correctly.
+ val supportsColumnUpdate = operation.supportsColumnUpdates()
// resolve all needed attrs (e.g. row ID and any required metadata attrs)
- val rowAttrs = relation.output
val rowIdAttrs = resolveRowIdAttrs(relation, operation)
val metadataAttrs = resolveRequiredMetadataAttrs(relation, operation)
- // construct a read relation and include all required metadata columns
+ // Connector-declared data attrs used to determine pass-through columns in the write plan.
+ val connectorDataAttrs = if (supportsColumnUpdate) {
+ resolveRequiredDataAttrs(relation, operation)
+ } else Nil
+
+ // MOR uses a full-schema scan; ColumnPruning narrows it via Project references.
val readRelation = buildRelationWithAttrs(relation, operationTable, metadataAttrs, rowIdAttrs)
+ // Connector-required attrs that are NOT being assigned are added as pass-throughs in the
+ // plan so that ColumnPruning keeps them in the physical scan AND the connector receives
+ // their current values via DeltaWriter.update's row argument.
+ val assignedAttrs = if (supportsColumnUpdate) computeAssignedAttrs(assignments)
+ else relation.output
+ val connectorExtraAttrs: Seq[AttributeReference] = if (connectorDataAttrs.nonEmpty) {
+ val assignedAttrSet = AttributeSet(assignedAttrs)
+ connectorDataAttrs.filterNot(assignedAttrSet.contains)
+ } else Nil
+
// build a plan for updated records that match the condition
val matchedRowsPlan = Filter(cond, readRelation)
val rowDeltaPlan = if (operation.representUpdateAsDeleteAndInsert) {
buildDeletesAndInserts(matchedRowsPlan, assignments, rowIdAttrs)
+ } else if (supportsColumnUpdate) {
+ buildColumnUpdateProjection(
+ matchedRowsPlan, assignments, rowIdAttrs, metadataAttrs, connectorExtraAttrs)
} else {
buildWriteDeltaUpdateProjection(matchedRowsPlan, assignments, rowIdAttrs)
}
+ // Effective row write schema:
+ // - Narrow path (connectorDataAttrs declared): exactly connector-declared cols in declared
+ // order. The connector must declare ALL columns it wants to receive (including updated
+ // ones). This mirrors the metadata pattern and enables strict areCompatible validation.
+ // - Heuristic path (connectorDataAttrs empty): only the assigned (changed) columns.
+ // - Full path (no column-update support): full table output.
+ val effectiveRowAttrs = if (supportsColumnUpdate && connectorDataAttrs.nonEmpty) {
+ connectorDataAttrs
+ } else if (supportsColumnUpdate) {
+ assignedAttrs
+ } else {
+ relation.output
+ }
+
// 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(
+ rowDeltaPlan, effectiveRowAttrs, rowIdAttrs, metadataAttrs)
val groupFilterCond = if (groupFilterEnabled) Some(cond) else None
WriteDelta(writeRelation, cond, rowDeltaPlan, relation, projections, groupFilterCond)
}
+ // Builds the row delta projection for the column update path.
+ //
+ // The resulting Project references only:
+ // - assigned column values (new values being written)
+ // - connector pass-through values (connector declared but not assigned)
+ // - metadata columns (nulled or preserved)
+ // - row ID columns (for delta identification)
+ // - original row ID values (only when a row ID column is being reassigned)
+ //
+ // ColumnPruning observes exactly these references and narrows the physical scan accordingly.
+ // Connectors that need additional columns in the scan (e.g., partition columns for
+ // distribution) should declare them in requiredDataAttributes().
+ //
+ // Note: AlignUpdateAssignments guarantees all assignment keys are top-level
Review Comment:
Do we have a test coverage for this, `AlignUpdateAssignments` contract?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56599][SQL] Add scan narrowing for column-level UPDATEs in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3192343349
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java:
##########
@@ -105,4 +105,45 @@ default String description() {
default NamedReference[] requiredMetadataAttributes() {
return new NamedReference[0];
}
+
+
+ /**
+ * Controls whether to send only the required data columns to the connector rather than the
+ * full row.
+ * <p>
+ * When true, Spark narrows the data column schema ({@link LogicalWriteInfo#schema()}) to only
+ * the columns declared via {@link #requiredDataAttributes()}. Metadata columns (from
+ * {@link #requiredMetadataAttributes()}) and row ID columns (from
+ * {@link SupportsDelta#rowId()}) are unaffected and always projected separately.
+ * <p>
+ * If {@link #requiredDataAttributes()} returns a non-empty array, the write schema is exactly
+ * those columns in declared order. The connector must include all columns it wants to receive,
+ * including the columns being updated. If {@link #requiredDataAttributes()} returns an empty
+ * array, Spark sends only the non-identity assigned columns (heuristic path).
+ *
+ * @since 4.2.0
Review Comment:
4.2.0 -> 4.3.0
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperation.java:
##########
@@ -105,4 +105,45 @@ default String description() {
default NamedReference[] requiredMetadataAttributes() {
return new NamedReference[0];
}
+
+
+ /**
+ * Controls whether to send only the required data columns to the connector rather than the
+ * full row.
+ * <p>
+ * When true, Spark narrows the data column schema ({@link LogicalWriteInfo#schema()}) to only
+ * the columns declared via {@link #requiredDataAttributes()}. Metadata columns (from
+ * {@link #requiredMetadataAttributes()}) and row ID columns (from
+ * {@link SupportsDelta#rowId()}) are unaffected and always projected separately.
+ * <p>
+ * If {@link #requiredDataAttributes()} returns a non-empty array, the write schema is exactly
+ * those columns in declared order. The connector must include all columns it wants to receive,
+ * including the columns being updated. If {@link #requiredDataAttributes()} returns an empty
+ * array, Spark sends only the non-identity assigned columns (heuristic path).
+ *
+ * @since 4.2.0
+ */
+ default boolean supportsColumnUpdates() {
+ return false;
+ }
+
+ /**
+ * Returns data column references required to perform this row-level operation.
+ * <p>
+ * This method is only consulted by Spark when {@link #supportsColumnUpdates()} returns
+ * {@code true}. If {@code supportsColumnUpdates()} returns {@code false}, the returned array
+ * is ignored and the full table row is sent (the default behavior).
+ * <p>
+ * When non-empty, the returned columns become the write schema in declared order.
+ * The connector must declare all columns it wants to receive, including the columns being
+ * updated. Use {@link RowLevelOperationInfo#updatedColumns()} to learn which columns are being
+ * assigned, then add any extra columns needed for row lookup or routing (e.g., primary key).
+ * <p>
+ * When empty (the default), Spark falls back to sending only the non-identity assigned columns.
+ *
+ * @since 4.2.0
Review Comment:
ditto. 4.3.0
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3786105284
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, any required metadata attrs and optionally connector
Review Comment:
It is unclear what `optionally connector declared attrs` means in this case. I'd opt to simplify.
```
// resolve all needed attrs (e.g. row ID and any required data / metadata attrs)
// resolve all needed attrs (e.g. row ID or any required metadata attrs)
```
Or similar
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub)" <gi...@apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r3786105284
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -155,29 +254,125 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val operation = operationTable.operation.asInstanceOf[SupportsDelta]
- // resolve all needed attrs (e.g. row ID and any required metadata attrs)
+ // resolve all needed attrs (e.g. row ID, any required metadata attrs and optionally connector
Review Comment:
It is unclear what `optionally connector declared attrs` means in this case. I'd opt to simplify.
```
// resolve all needed attrs (e.g. row ID or any required metadata attrs)
// resolve all needed attrs (e.g. row ID and any required data / metadata attrs)
```
Or similar
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "anuragmantri (via GitHub)" <gi...@apache.org>.
anuragmantri commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-5498692777
> Hey @anuragmantri do you have the cycles to update this?
Yes, I will update this. I will be actively working on this PR. @holdenk
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4354576399
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/RowLevelOperationTable.scala:
##########
@@ -58,4 +63,13 @@ private[sql] case class RowLevelOperationTable(
override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = {
operation.newWriteBuilder(info)
}
+
+ /**
+ * Builds the operation for another table from the same info, for a rule that redirects the
+ * write after the rewrite, such as to a branch of the table.
+ */
+ def withTable(newTable: Table with SupportsRowLevelOperations): RowLevelOperationTable = {
Review Comment:
**26. `withTable` does not say that its one result must replace the write relation and every read of the old table.** Each call builds a new operation and a new `RowLevelOperationTable`, while scan planning pairs the reads with the write by reference: `GroupBasedRowLevelOperation` collects the reads with `r.table eq table` (`patterns.scala:480-483`), and `GroupBasedRowLevelOperationScanPlanning` replaces only reads of that instance (`GroupBasedRowLevelOperationScanPlanning.scala:93`). `V2Writes.validateWriteRelation` compares only column names, so it does not notice a mismatch.
If a rule swaps only the write relation, as the new test does before running `V2Writes` alone (`GroupBasedColumnUpdateTableSuite.scala:551-555`), or calls `withTable` once per relation, the write finds no read relation, so its copy-on-write scan is never configured and runtime group filtering is skipped. The in-memory fixture then fails on `scan.data` at commit (`InMemoryRowLevelOperationTable.scala:728`). In Iceberg, `CopyOnWriteOperation.overwrittenFiles()` returns an empty set when the scan is null ([SparkWrite.java L450-452](https://github.com/apache/iceberg/blob/9303e4c89705f139290ae8c1266fb4485785e279/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkWrite.java#L450-L452)), so the commit adds the rewritten files without removing the old ones, and the rows are duplicated silently. Iceberg's `ResolveBranch` builds one instance and uses it for both the write and the query, which is the right shape. This hazard is not new with `withTable`, but after item 9 the
PR description points rule authors to this helper.
Could the Scaladoc say that a rule should call `withTable` once and use the result for the write relation and every read of the old table?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala:
##########
@@ -380,14 +380,17 @@ case class ReplaceDataExec(
write: Write,
rowLevelCommand: RowLevelOperation.Command,
tableName: String,
+ useWriteColumnUpdate: Boolean = false,
Review Comment:
**27.** Minor: this default is new in 22332b3 (the previous commit, 9b6655a, had `useWriteColumnUpdate: Boolean,`), no caller uses it, and `false` is the wrong value for every narrow write. The only construction passes the flag by name (`DataSourceV2Strategy.scala:594`), and `withTransaction` and `withNewChildInternal` copy it. A construction in the pre-PR six-argument shape still compiles and sends narrow UPDATE and COPY rows to `DataWriter.write(...)`, whose layout is the empty `schema()`, and nothing fails, because the flag cannot be derived from `write`. `DataAndMetadataWritingSparkTask` and `DataWithProjectionWritingSparkTask` already require it. Could the default be dropped?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -128,12 +153,61 @@ class InMemoryRowLevelOperationTable private (
copied.replacedPartitions = replacedPartitions
copied.lastWriteInfo = lastWriteInfo
copied.lastWriteLog = lastWriteLog
+ copied.lastUpdatedColumns = lastUpdatedColumns
copied
}
override def newRowLevelOperationBuilder(
info: RowLevelOperationInfo): RowLevelOperationBuilder = {
- if (properties.getOrDefault(SUPPORTS_DELTAS, "false") == "true") {
+ lastUpdatedColumns = info.updatedColumns()
+ val splitUpdates = properties.getOrDefault(SPLIT_UPDATES, "false").toBoolean
+ if (properties.getOrDefault(COLUMN_UPDATE, "false") == "true") {
+ if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ }
+ } else if (properties.containsKey(COLUMN_UPDATE_REQ_ATTRS)) {
+ val reqCols = properties.get(COLUMN_UPDATE_REQ_ATTRS).split(",").map(_.trim)
+ if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperationWithReqAttrs(
+ info.command, reqCols, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperationWithReqAttrs(info.command, reqCols, info.options)
+ }
+ } else if (properties.getOrDefault(COLUMN_UPDATE_EMPTY_REQ_ATTRS, "false") == "true") {
+ // Test-only: returns an empty requiredDataAttributes() so we can verify Spark rejects it.
+ () => new DeltaBasedColumnUpdateOperationWithReqAttrs(
+ info.command, Array.empty, info.options)
+ } else if (properties.getOrDefault(COLUMN_UPDATE_COW, "false") == "true") {
+ if (info.updatedColumns().exists(_.describe() == "pk")) {
+ () => PartitionBasedOperation(info.command, info.options)
+ } else {
+ () => new PartitionBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ }
+ } else if (properties.containsKey(COLUMN_UPDATE_COW_REQ_ATTRS)) {
Review Comment:
**29. Item 5's fallback covers only `column-update-cow`, so the req-attrs copy-on-write flavor still commits wrong rows when an UPDATE assigns `pk`.** L185-186 build a full-width operation when `updatedColumns()` contains `pk`, as `SupportsColumnUpdates.java:62-67` now asks, but this branch always builds the narrow operation, and its commit looks the base row up by the new `pk` value with `find` (L747). With `reqAttrsTableProps("pk,salary")` on the shared rows (1, 1, hr, 100, 10, 1), (2, 2, hr, 200, 20, 1) and (3, 3, software, 300, 30, 1), `UPDATE t SET pk = 2, salary = 0 WHERE pk = 1` is accepted because `pk` is declared. The narrow row [2, 0] is then overlaid on the existing `pk = 2` row, so the commit writes (2, 2, hr, 0, 20, 1) instead of (2, 1, hr, 0, 10, 1). With duplicate `pk` values, the overlay uses whichever row comes first, while the delta commit now fails unless exactly one row matches. No test assigns `pk` on this flavor. Could this branch use the same fallback (or fa
il clearly), and could the copy-on-write commit also require exactly one match?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -321,23 +810,61 @@ class InMemoryRowLevelOperationTable private (
}
}
-private class DeltaBufferedRowsWriterFactory(schema: StructType) extends DeltaWriterFactory {
+private class DeltaBufferedRowsWriterFactory(
+ schema: StructType,
+ intRowId: Boolean = true) extends DeltaWriterFactory {
override def createWriter(partitionId: Int, taskId: Long): DeltaWriter[InternalRow] = {
- new DeltaBufferWriter(schema)
+ new DeltaBufferWriter(schema, intRowId)
+ }
+}
+
+// Test-only: a writer factory for a connector that never overrides writeColumnUpdate(), to
+// exercise the DataWriter#writeColumnUpdate default that throws
+// DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED.
+// Declared top-level (no reference to the enclosing table) so it stays task-serializable.
+private class NoWriteColumnUpdateOverrideWriterFactory extends DataWriterFactory {
+ override def createWriter(partitionId: Int, taskId: Long): DataWriter[InternalRow] = {
+ new NoWriteColumnUpdateOverrideWriter
}
}
-private class DeltaBufferWriter(schema: StructType) extends BufferWriter(schema)
- with DeltaWriter[InternalRow] {
+private class NoWriteColumnUpdateOverrideWriter extends DataWriter[InternalRow] {
+ override def write(record: InternalRow): Unit = {}
+ override def commit(): WriterCommitMessage = new BufferedRows(Seq.empty, StructType(Nil))
+ override def abort(): Unit = {}
+ override def close(): Unit = {}
+}
+
+// Test-only: a writer that overrides writeColumnUpdate(record) but not
+// writeColumnUpdate(metadata, record), so rows with metadata reach it through the default.
+private class RecordOnlyColumnUpdateWriterFactory(schema: StructType) extends DataWriterFactory {
+ override def createWriter(partitionId: Int, taskId: Long): DataWriter[InternalRow] = {
+ new RecordOnlyColumnUpdateWriter(schema)
+ }
+}
+
+private class RecordOnlyColumnUpdateWriter(schema: StructType) extends DataWriter[InternalRow] {
+ private val delegate = new BufferWriter(schema)
+ override def write(record: InternalRow): Unit = delegate.write(record)
+ override def writeColumnUpdate(record: InternalRow): Unit = delegate.writeColumnUpdate(record)
+ override def commit(): WriterCommitMessage = delegate.commit()
+ override def abort(): Unit = delegate.abort()
+ override def close(): Unit = delegate.close()
+}
+
+// Logs the row ID of a DELETE as an integer when `intRowId` is set and as null otherwise; the
+// commit then resolves such a DELETE from its metadata.
+private class DeltaBufferWriter(schema: StructType, intRowId: Boolean = true)
+ extends BufferWriter(schema) with DeltaWriter[InternalRow] {
private final val DELETE = UTF8String.fromString(Delete.toString)
private final val UPDATE = UTF8String.fromString(Update.toString)
private final val REINSERT = UTF8String.fromString(Reinsert.toString)
private final val INSERT = UTF8String.fromString(Insert.toString)
override def delete(meta: InternalRow, id: InternalRow): Unit = {
- val pk = id.getInt(0)
- buffer.deletes += pk
+ val pk: Any = if (intRowId) id.getInt(0) else null
Review Comment:
**30.** Minor: items 14 and 23 made the split commit locate rows by `rowId()`, but this writer still drops a row ID that is not an `INT` and treats an `INT` row ID as `pk`. With `no-metadata`, a delete arrives as `delete(null, rowIdProjection)` (`WriteToDataSourceV2Exec.scala:890`). So on the suite's own `column-update-split-row-id=_partition` + `no-metadata` table (`DeltaBasedColumnUpdateTableSuite.scala:865-868`, which today runs only a rejected UPDATE), `DELETE FROM t WHERE pk = 1` logs `(DELETE, null, null, null)`, and the null-ID branch (L572-574) calls `metadataSchema.fieldIndex("_partition")` on an empty schema (L559), which fails with `FIELD_NOT_FOUND`. A MERGE with `WHEN MATCHED THEN DELETE` fails the same way. With `column-update-split-row-id=id` on `pk INT NOT NULL, id INT NOT NULL, dep STRING` and rows (1, 2, hr) and (2, 1, hr), `DELETE ... WHERE pk = 1` adds 2 to `buffer.deletes`, and `InMemoryBaseTable.scala:416` removes `pk = 2` instead. Nothing fails today, but the
fixture accepts both configurations. Could the writer keep a copy of the row ID beside the logged entry (leaving the asserted `id` slot as it is) and resolve such DELETEs through `currentRowByRowId`, or fail clearly for row-ID setups it does not support?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -141,6 +153,19 @@ abstract class RowLevelOperationSuiteBase
append(schemaString, jsonData)
}
+ protected def createAndInitTable(
Review Comment:
**39.** Nit: this overload re-inlines the column conversion (L117), the default `dep` partitioning (L127) and the `TableInfo` build of `createTable(columns, transforms)` (L142-149); only `.withProperties(tableProps)` differs, and about 30 call sites create their tables through it. `UpdateTableSuiteBase.scala:162-166` likewise rebuilds what `createTable(columns, Array.empty[Transform])` does. Round 4 ([r4074793101](https://github.com/apache/spark/pull/55518#discussion_r4074793101)) suggested a props parameter on the base helper. Could `createTable(columns, transforms)` take `tableProps: java.util.Map[String, String] = extraTableProps`, with both `createAndInitTable` overloads and `UpdateTableSuiteBase` going through it?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala:
##########
@@ -356,6 +357,32 @@ trait RowLevelWrite extends V2WriteCommand with SupportsSubquery {
def condition: Expression
def originalTable: NamedRelation
+ /** The projections that read query columns by ordinal. */
+ protected def queryProjections: Seq[ProjectingInternalRow]
+
+ /**
+ * The write reads the operation column at ordinal 0 and other query columns only through
+ * projections, by ordinal. Any other query column can be pruned with [[withPrunedQuery]].
+ */
+ override lazy val references: AttributeSet = {
+ val output = query.output
Review Comment:
**28. `references` walks the `List` behind `query.output` once per read ordinal, which is quadratic in the query width.** `query.output` is a `List` in the rewrite shapes (`Project.output` is `projectList.map`, and the relation output comes from `StructType.map`), so `output.length` (L371) and `output(ordinal)` (L376) are linear and run once per ordinal. A full-row write reads every column, so one evaluation takes about 1.5 * W^2 list steps for a W-column query, e.g. 37.5M at 5,000 columns. This now applies to every DSv2 DELETE, UPDATE and MERGE, not only to column updates: `CheckAnalysis` forces it through `missingInput`, and the new `ColumnPruning` guard forces it again on every new `ReplaceData` or `WriteDelta` copy. Master used the cached `query.outputSet`, and 9b6655a's `query.output.take(lastReadOrdinal + 1)` was linear. `Optimizer.scala:1541` repeats `output(i)` inside `output.indices.filter`.
Also, L371 checks only the upper bound, so a `-1` ordinal from a third-party builder escapes as a raw `IndexOutOfBoundsException` instead of this internal error.
Could this use `val output = query.output.toIndexedSeq` with `!output.isDefinedAt(ordinal)`, and could the `ColumnPruning` case use `w.query.output.zipWithIndex.collect { case (a, i) if w.references.contains(a) => i }`?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,1114 @@
+/*
+ * 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.{SparkException, SparkRuntimeException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+import org.apache.spark.sql.catalyst.util.RowDeltaUtils.ORIGINAL_ROW_ID_VALUE_PREFIX
+import org.apache.spark.sql.connector.catalog.{Delete, InMemoryBaseTable, InMemoryRowLevelOperationTable, Reinsert}
+import org.apache.spark.sql.connector.write.{LogicalWriteInfo, RowLevelOperationTable}
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.errors.QueryCompilationErrors
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, updated rows contain only the declared columns
+ * (LogicalWriteInfo.columnUpdateSchema()) rather than the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase
+ with ColumnUpdateTableSharedTests {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ override protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String] = {
+ java.util.Map.of("column-update-req-attrs", reqAttrs)
+ }
+
+ override protected def fullWidthTableProps: java.util.Map[String, String] = {
+ java.util.Map.of("supports-deltas", "true")
+ }
+
+ override protected def scanMetadataColumns: String = ", _partition STRING"
+
+ private val reqAttrsOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateOperationWithReqAttrs].getName
+
+ private val splitOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperation].getName
+
+ private val splitReqAttrsOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithReqAttrs]
+ .getName
+
+ private val splitWithRowIdOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithRowId].getName
+
+ test("column-update: column update schema contains a single assigned column") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: column update schema contains multiple assigned columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false),
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: column update schema is pk and dep for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ // All assignments are identity, so updatedColumns is empty; pk and dep remain because the
+ // connector declares them as base columns.
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Array(PK_FIELD, DEP_FIELD))))
+ }
+
+ test("column-update: row filter condition is orthogonal to column narrowing") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = 'engineering' WHERE pk IN (1, 3)")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: update all rows (no WHERE clause)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("salary", IntegerType, nullable = true)
+ ))))
+ }
+
+ test("column-update: column update schema excludes identity assignments in a mixed UPDATE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: CHAR and VARCHAR self-assignments are reported as updated") {
+ createTable("pk INT NOT NULL, id INT, c CHAR(3), v VARCHAR(10), dep STRING")
+ sql(s"INSERT INTO $tableNameAsString VALUES (1, 1, 'a', 'x', 'hr')")
+
+ sql(s"UPDATE $tableNameAsString SET id = id, c = c, v = v WHERE pk = 1")
+
+ checkLastUpdatedColumns("c", "v")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, 1, "a ", "x", "hr") :: Nil)
+ }
+
+ test("column-update: a LogicalWriteInfo that does not implement columnUpdateSchema throws") {
+ val info = new LogicalWriteInfo {
+ override def options(): CaseInsensitiveStringMap = CaseInsensitiveStringMap.empty()
+ override def queryId(): String = "query"
+ override def schema(): StructType = new StructType()
+ }
+
+ checkError(
+ exception = intercept[SparkUnsupportedOperationException](info.columnUpdateSchema()),
+ condition = "DATA_SOURCE_COLUMN_UPDATE_SCHEMA_NOT_IMPLEMENTED",
+ parameters = Map("class" -> info.getClass.getName))
+ }
+
+ test("column-update: updatedColumns is empty for DELETE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"DELETE FROM $tableNameAsString WHERE dep = 'hr'")
+
+ checkLastUpdatedColumns()
+ }
+
+ test("column-update split: row-level DELETE falls back to full-width rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split", "true"))
+ checkDeleteFallsBackToFullWidthRows(Row(2, 200, "hr"))
+ }
+
+ test("column-update split: row-level DELETE by a metadata row ID falls back to full-width " +
+ "rows") {
+ // Each partition holds one row, so `_partition` identifies it.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "_partition"))
+ checkDeleteFallsBackToFullWidthRows(Row(2, 200, "software"))
+ }
+
+ test("column-update split: MERGE falls back to full-width rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split", "true"))
+ checkMergeFallsBackToFullWidthRows()
+ }
+
+ test("column-update: data correctness -- single column update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0)
+ }
+
+ test("column-update: data correctness -- update all rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 200, "hr") :: Row(2, 400, "software") :: Row(3, 600, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 3, numCopiedRows = 0)
+ }
+
+ test("column-update: data correctness -- mixed identity and real assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "engineering") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0)
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-req-attrs", reqAttrs))
+ }
+
+ test("column-update: requiredDataAttributes - data correctness") {
+ createAndInitTableWithReqAttrs("dep,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: undeclared distribution column is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "pk,salary",
+ "column-update-cluster-by", "dep"))
+
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName
+ Seq(Nil, Seq(excludeColumnPruning)).foreach { confs =>
+ withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[`dep`]"))
+ }
+ }
+ }
+
+ test("column-update split: undeclared distribution column is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-split-req-attrs", "pk,salary",
+ "column-update-cluster-by", "dep"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> splitReqAttrsOperationClass, "columns" -> "[`dep`]"))
+ }
+
+ test("column-update: undeclared write requirement columns are listed once each, sorted") {
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 1, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "pk,salary",
+ "column-update-cluster-by", "id,dep"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[`dep`, `id`]"))
+ }
+
+ test("column-update: a write requirement on a missing nested field reports why") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1 }, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "pk,s",
+ "column-update-cluster-by", "s.missing"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "columns" -> "[`s`.`missing` (FIELD_NOT_FOUND)]"))
+ }
+
+ test("column-update: an empty write requirement column name is reported as <empty>") {
+ checkError(
+ exception = QueryCompilationErrors.undeclaredWriteRequirementColumnsError(
+ reqAttrsOperationClass, Seq(Seq.empty[String] -> None)).asInstanceOf[AnalysisException],
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[<empty>]"))
+ }
+
+ test("column-update: undeclared column inside nested transforms is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "pk,salary",
+ "column-update-cluster-by", "dep",
+ "column-update-cluster-by-nested-transform", "true"))
+
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName
+ Seq(Nil, Seq(excludeColumnPruning)).foreach { confs =>
+ withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[`dep`]"))
+ }
+ }
+ }
+
+ test("column-update: distribution by an undeclared row ID column is allowed") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "salary,dep",
+ "column-update-cluster-by", "pk"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 200, "hr") :: Nil)
+ }
+
+ test("column-update: distribution by an undeclared row ID column assigned to itself is " +
+ "allowed") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "salary,dep",
+ "column-update-cluster-by", "pk"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1, pk = pk WHERE pk = 1")
+
+ checkLastUpdatedColumns("salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 200, "hr") :: Nil)
+ }
+
+ test("column-update: partition source column not in requiredDataAttributes") {
+ createAndInitTableWithReqAttrs("pk,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ executeAndCheckScanExcludes(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1", "dep")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively on the row-ID " +
+ "column without adopting the declared spelling") {
+ createAndInitTableWithReqAttrs("PK,salary,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "dep"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: undeclared, differently-cased columns in the RHS and WHERE clause " +
+ "are read with the table's own spelling") {
+ createAndInitTableWithReqAttrs("pk,dep,salary",
+ "pk INT NOT NULL, salary INT, bonus INT, extra INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "extra": 5, "dep": "hr" }
+ |""".stripMargin)
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + BONUS WHERE pk = 1 AND EXTRA > 3",
+ "pk INT, salary INT, bonus INT, extra INT, dep STRING, _partition STRING")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, 110, 10, 5, "hr") :: Nil)
+ }
+
+ test("column-update: DELETE on a scan that returns unrequested columns is unaffected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update", "true",
+ InMemoryBaseTable.SIMULATE_PARTIAL_COLUMN_PRUNING, "true"))
+
+ sql(s"DELETE FROM $tableNameAsString WHERE pk < 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(2, 200, 20, "software") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively on an assigned " +
+ "column without adopting the declared spelling") {
+ createAndInitTableWithReqAttrs("pk,SALARY,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "dep"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: case-sensitive analysis accepts columns that differ only in case") {
+ withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") {
+ createAndInitTableWithReqAttrs("pk,salary,SALARY,dep",
+ "pk INT NOT NULL, salary INT, SALARY INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "SALARY": 200, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "SALARY", "dep"))
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, -1, 200, "hr") :: Nil)
+ }
+ }
+
+ test("column-update: empty requiredDataAttributes throws AnalysisException") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }""".stripMargin,
+ java.util.Map.of("column-update-empty-req-attrs", "true"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ parameters = Map("connector" -> reqAttrsOperationClass))
+ }
+
+ test("column-update: nested requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("pk,s.c1,dep",
+ "pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map("connector" -> reqAttrsOperationClass, "nestedAttributes" -> "[`s`.`c1`]"))
+ }
+
+ test("column-update: duplicate requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("pk,pk,salary", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map("connector" -> reqAttrsOperationClass, "duplicateAttributes" -> "[`pk`]"))
+ }
+
+ test("column-update: case-only duplicate requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("PK,pk,salary", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map("connector" -> reqAttrsOperationClass, "duplicateAttributes" -> "[`PK`]"))
+ }
+
+ test("column-update: duplicate requiredDataAttributes are reported in declared order") {
+ createAndInitTableWithReqAttrs("salary,pk,SALARY,PK", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass, "duplicateAttributes" -> "[`salary`, `pk`]"))
+ }
+
+ test("column-update: metadata column in requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("pk,_partition,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET dep = 'x' WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "metadataAttributes" -> "[`_partition`]"))
+ }
+
+ test("column-update split: metadata column in requiredDataAttributes throws " +
+ "AnalysisException") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-req-attrs", "pk,_partition,dep"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET dep = 'x' WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> splitReqAttrsOperationClass,
+ "metadataAttributes" -> "[`_partition`]"))
+ }
+
+ test("column-update: requiredDataAttributes rejects a column that does not exist") {
+ createAndInitTableWithReqAttrs("pk,nonexistent_col,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "unknownAttributes" -> "[`nonexistent_col`]"))
+ }
+
+ test("column-update: column update schema excludes undeclared columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "dep", "salary"))
+ }
+
+ test("column-update: data correctness") {
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "id": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ }
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-split", "true"))
+ }
+
+ test("column-update split: column update schema has the declared columns") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update split: data correctness") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, -1, "hr") :: Nil)
+ }
+
+ test("column-update: split-updates sends a narrow UPDATE as delete and reinsert") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update", "true", "split-updates", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ deleteWriteLogEntry(id = 1, metadata = Row("hr", null)),
+ reinsertWriteLogEntry(metadata = Row("hr", null), data = Row(1, "hr", -1)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "hr") :: Nil)
+ }
+
+ test("column-update split: delete and reinsert carry the row ID and preserved metadata") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ deleteWriteLogEntry(id = 1, metadata = Row("hr", null)),
+ reinsertWriteLogEntry(metadata = Row("hr", null), data = Row(1, "hr", -1)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "hr") :: Nil)
+ }
+
+ test("column-update split: row-ID reassignment on narrow write is rejected") {
+ // `extra` is neither declared, updated, nor referenced, so requiredDataAttributes()
+ // ([pk, dep, salary]) doesn't cover every table column -- the write stays genuinely narrow
+ // and reassigning the row ID must be rejected.
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING, extra STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr", "extra": "x" }
+ |{ "pk": 2, "salary": 200, "dep": "software", "extra": "y" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT",
+ parameters = Map("connector" -> splitOperationClass, "rowIds" -> "[`pk`]"))
+ }
+
+ test("column-update split: row-ID reassignment is allowed when requiredDataAttributes " +
+ "covers every table column") {
+ createAndInitTableSplitWithReqAttrs("pk,salary,dep",
+ "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(2, 200, "software") :: Row(11, -1, "hr") :: Nil)
+ }
+
+ test("column-update split: row-ID reassignment is allowed when requiredDataAttributes " +
+ "covers every data column of a relation with metadata columns") {
+ createAndInitTableSplitWithReqAttrs("pk,salary,dep",
+ "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ // `_partition` and the coerced `salary > 0.5` make the analyzer add the metadata columns to
+ // the target relation in a later iteration. Only analyzes the plan.
+ val parsed = spark.sessionState.sqlParser.parsePlan(
+ s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 " +
+ s"WHERE dep = 'hr' AND _partition = 'hr' AND salary > 0.5")
+ val write = spark.sessionState.executePlan(parsed).analyzed.collectFirst {
+ case w: WriteDelta => w
+ }.getOrElse(fail("couldn't find WriteDelta in analyzed plan"))
+
+ assert(write.table.output.map(_.name) == Seq("pk", "salary", "dep"))
+ }
+
+ private def createAndInitTableSplitWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-split-req-attrs", reqAttrs))
+ }
+
+ test("column-update split: undeclared row-ID column on narrow write is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "pk"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED",
+ parameters = Map("connector" -> splitWithRowIdOperationClass, "rowIds" -> "[`pk`]"))
+ }
+
+ test("column-update split: metadata row ID preserved on reinsert need not be declared") {
+ // `_partition` is a required metadata attribute with PRESERVE_ON_REINSERT, so the REINSERT
+ // row carries it through the metadata projection.
+ // Each partition holds one row, so `_partition` identifies it.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "_partition"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq == Seq("salary"))
+ checkLastColumnUpdateWriteLog(
+ WriteLogEntry(operation = Delete, metadata = Some(Row("hr", null))),
+ reinsertWriteLogEntry(metadata = Row("hr", null), data = Row(-1)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update split: metadata row ID not preserved on reinsert must be declared") {
+ // `index` is a required metadata attribute, but its REINSERT value is null.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "index"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED",
+ parameters = Map("connector" -> splitWithRowIdOperationClass, "rowIds" -> "[`index`]"))
+ }
+
+ test("column-update split: metadata row ID outside requiredMetadataAttributes must be " +
+ "declared") {
+ // Without required metadata attributes, nothing projects `_partition` into the REINSERT row.
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "_partition", "no-metadata", "true"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED",
+ parameters = Map("connector" -> splitWithRowIdOperationClass, "rowIds" -> "[`_partition`]"))
+ }
+
+ // ---------------------------------------------------------------------------
+ // Write relation tests: verify that only the write relation is narrowed to the
+ // declared columns, for in-place and delete+insert deltas.
+ // ---------------------------------------------------------------------------
+
+ test("column-update: write relation has the declared columns in declared order") {
Review Comment:
**36.** Minor: item 18's `ColumnUpdateTableSharedTests` took the tests listed there, but eight more UPDATE tests are still copied between this suite and `GroupBasedColumnUpdateTableSuite` (this suite's line / the group suite's line):
- "write relation has the declared columns in declared order" (L883 / L245)
- "scan reads only the referenced columns" (L915 / L307), which differ only by the trait's `scanMetadataColumns`
- "analyzed plan shows the declared columns as the write target" (L958 / L500), which differ only by `write.nodeName`
- "an operation replaced after the rewrite is rejected" (L931 / L511), both new in 22332b3
- "undeclared, differently-cased columns in the RHS and WHERE clause ..." (L476 / L598)
- the column update schema test (L653 / L47) and "data correctness" (L665 / L59)
- the assignment outside `requiredDataAttributes` (L1024 / L469)
That is about 110 lines on each side, every fix has to be made twice, and `DeltaBasedNoMetadataColumnUpdateTableSuite` runs none of them. This suite also hard-codes `java.util.Map.of("column-update-req-attrs", ...)` at L884, L916 and L959 and in `createAndInitTableWithReqAttrs` (L289-295) instead of using its own `reqAttrsTableProps`. Could these move into the trait, with hooks for the replacement operations and `reqAttrsOperationClass`?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/RowLevelOperationSuiteBase.scala:
##########
@@ -127,6 +128,17 @@ abstract class RowLevelOperationSuiteBase
createTable(columns, transforms)
}
+ // `extraTableProps` as the TBLPROPERTIES clause of a CREATE TABLE statement
+ protected def extraTablePropertiesClause: String = {
Review Comment:
**40.** Nit: this sits between `createTable(columns)` (L126-129) and the overload it delegates to (L142-149), so the `createTable` overloads, which were contiguous before this PR, are split. AGENTS.md asks to "put the new member with the code it belongs with" (as in item 19). Could it move next to `extraTableProps` (L94-96) and `noMetadata` (L101-102)?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -343,3 +365,60 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
}
}
}
+
+object RewriteRowLevelCommand extends SQLConfHelper {
+
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation, in declared order and with the relation's spelling. An empty declaration,
+ * nested fields, duplicate columns, columns that do not exist in the relation and metadata
+ * columns are rejected with an `AnalysisException`.
+ */
+ private[sql] def resolveRequiredDataAttrs(
+ relation: LogicalPlan,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes.toImmutableArraySeq
+ if (refs.isEmpty) {
+ throw QueryCompilationErrors.emptyRequiredDataAttributesError(operation.getClass.getName)
+ }
+ val nested = refs.filter(_.fieldNames.length != 1).map(_.fieldNames.toSeq)
Review Comment:
**34.** Nit: a declaration with no name parts, i.e. a `NamedReference` whose `fieldNames()` is empty, passes this `length != 1` filter, so it is reported as a nested field, and since `toSQLId(Seq())` is empty, the message reads "declared nested field(s) [] in `requiredDataAttributes()`. Declare the root struct column instead of a nested field." This came with item 17's `toSQLId` change; the same PR prints `<empty>` for this case in `undeclaredWriteRequirementColumnsError` (`QueryCompilationErrors.scala:4725`) and `V2ExpressionUtils.unresolvableColumns`. Could empty declarations be rejected separately, e.g. as an unknown column shown as `<empty>`, before L390?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/UpdateTableSuiteBase.scala:
##########
@@ -53,6 +53,68 @@ abstract class UpdateTableSuiteBase extends RowLevelOperationSuiteBase {
s"Expected numCopiedRows=$expectedCopied, got ${summary.numCopiedRows()}")
}
+ test("SPARK-58111: updatedColumns: single non-identity assignment") {
Review Comment:
**32.** Minor: 22332b3 moved these six `updatedColumns()` tests from `DeltaBasedUpdateTableSuiteBase` here, so they now run in all eight UPDATE suites, i.e. 48 table setups and UPDATE jobs instead of 18. The value they check is computed from the assignments before the operation type matters (`RewriteUpdateTable.scala:47-48`), so it does not depend on the table kind, and "no WHERE clause still reports assigned columns" (L96) checks nothing beyond L56, because the condition is not involved. Could they become one table-driven test, e.g. `Seq("id = -1" -> Seq("id"), "id = -1, dep = 'eng'" -> Seq("id", "dep"), ...)`, or stay in one suite?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -343,3 +365,60 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
}
}
}
+
+object RewriteRowLevelCommand extends SQLConfHelper {
+
+ /**
+ * Resolves the connector-declared required data attributes for column-update writes against
+ * the given relation, in declared order and with the relation's spelling. An empty declaration,
+ * nested fields, duplicate columns, columns that do not exist in the relation and metadata
+ * columns are rejected with an `AnalysisException`.
+ */
+ private[sql] def resolveRequiredDataAttrs(
+ relation: LogicalPlan,
+ operation: SupportsColumnUpdates): Seq[AttributeReference] = {
+ val refs = operation.requiredDataAttributes.toImmutableArraySeq
+ if (refs.isEmpty) {
+ throw QueryCompilationErrors.emptyRequiredDataAttributesError(operation.getClass.getName)
+ }
+ val nested = refs.filter(_.fieldNames.length != 1).map(_.fieldNames.toSeq)
+ if (nested.nonEmpty) {
+ throw QueryCompilationErrors.nestedRequiredDataAttributeError(
+ operation.getClass.getName, nested)
+ }
+ val normalizedNames = refs.map { ref =>
Review Comment:
**38.** Nit: this is `SchemaUtils.foldName` (`SchemaUtils.scala:551-553`), whose Scaladoc says that `AttributeSeq` looks attributes up by this key, which is what the duplicate check relies on. Could this be `refs.map(ref => SchemaUtils.foldName(ref.fieldNames.head, conf.caseSensitiveAnalysis))`, dropping the `java.util.Locale` import?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/v2Commands.scala:
##########
@@ -356,6 +357,32 @@ trait RowLevelWrite extends V2WriteCommand with SupportsSubquery {
def condition: Expression
def originalTable: NamedRelation
+ /** The projections that read query columns by ordinal. */
+ protected def queryProjections: Seq[ProjectingInternalRow]
Review Comment:
**35.** Nit: the PR description's compatibility paragraph says that `RowLevelWrite` has a new abstract `withPrunedQuery`, but this `queryProjections` is abstract too and drives the new `references`, so an extension that adds only `withPrunedQuery` still fails to compile. Since the description becomes the commit message, could it name both?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,637 @@
+/*
+ * 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.{SparkException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTable
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.connector.write.RowLevelOperationTable
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+class GroupBasedColumnUpdateTableSuite extends UpdateTableSuiteBase
+ with ColumnUpdateTableSharedTests {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-cow", "true")
+ props
+ }
+
+ override protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String] = {
+ java.util.Map.of("column-update-cow-req-attrs", reqAttrs)
+ }
+
+ override protected def fullWidthTableProps: java.util.Map[String, String] = java.util.Map.of()
+
+ test("column-update ReplaceData: column update schema has only the declared columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "dep", "salary"))
+ }
+
+ test("column-update ReplaceData: data correctness -- bonus preserved, salary updated") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ // CoW: pk=1 and pk=3 match (2 updates); the 'hr' partition has no other rows so no copies.
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update ReplaceData: subquery WHERE condition data correctness") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ withTempView("target_deps") {
+ import testImplicits._
+ val subqueryDF = Seq("hr").toDF()
+ subqueryDF.createOrReplaceTempView("target_deps")
+
+ sql(
+ s"""UPDATE $tableNameAsString
+ |SET salary = -1
+ |WHERE dep IN (SELECT * FROM target_deps)
+ |""".stripMargin)
+ }
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update ReplaceData: runtime group filtering data correctness") {
+ Seq(true, false).foreach { dppEnabled =>
+ Seq(true, false).foreach { aqeEnabled =>
+ withSQLConf(
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> dppEnabled.toString,
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString) {
+ withTable(tableNameAsString) {
+ withTempView("matched_pk") {
+ createAndInitTable("pk INT, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ import testImplicits._
+ val matchedPkDF = Seq(Some(1), None).toDF()
+ matchedPkDF.createOrReplaceTempView("matched_pk")
+
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = -1 " +
+ s"WHERE pk IN (SELECT * FROM matched_pk)",
+ primaryScanSchema =
+ "pk INT, salary INT, dep STRING, _partition STRING, index INT",
+ groupFilterScanSchema = Some("pk INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, 300, 30, "hr") :: Nil)
+
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("column-update ReplaceData: updated and copied rows reach writeColumnUpdate with " +
+ "metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", null), data = Row(1, "hr", -1)),
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", 1), data = Row(2, "hr", 200)))
+ }
+
+ test("column-update ReplaceData: without metadata, rows reach writeColumnUpdate(record)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "no-metadata", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Row(3, 300, 30, "software") :: Nil)
+ }
+
+ test("column-update ReplaceData: a writer overriding only writeColumnUpdate(record) " +
+ "receives rows with metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "column-update-cow-record-only-writer", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: without metadata, column update schema is the declared " +
+ "layout") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "no-metadata", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedColumnUpdateSchema = Some(StructType(Seq(PK_FIELD, DEP_FIELD,
+ StructField("salary", IntegerType, nullable = true)))))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 110, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: connector missing writeColumnUpdate override is rejected") {
+ createAndInitTableReplaceDataNoWriteColumnUpdate("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[SparkUnsupportedOperationException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED",
+ parameters = Map(
+ "class" -> "org.apache.spark.sql.connector.catalog.NoWriteColumnUpdateOverrideWriter"))
+ }
+
+ private def createAndInitTableReplaceDataNoWriteColumnUpdate(
+ schemaString: String, jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-cow-no-write-column-update", "true"))
+ }
+
+ // ---------------------------------------------------------------------------
+ // Write relation tests: verify that only the write relation is narrowed to the
+ // declared columns, on both copy-on-write plan shapes.
+ // ---------------------------------------------------------------------------
+
+ private val reqAttrsOperationClass =
+ classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperationWithReqAttrs].getName
+
+ private def createAndInitTableWithReqAttrs(reqAttrs: String): Unit = {
+ createAndInitTableWithProps(reqAttrsTableProps(reqAttrs))
+ }
+
+ test("column-update ReplaceData: write relation has the declared columns in declared order") {
+ createAndInitTableWithReqAttrs("salary,pk")
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'")
+ }
+
+ assert(!write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "salary", "pk")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update ReplaceData: subquery write relation has the declared columns") {
+ withTempView("deps") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE dep IN (SELECT dep FROM deps)")
+ }
+
+ assert(write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "pk", "salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+ }
+
+ test("column-update ReplaceData: runtime group filtering on an undeclared condition column") {
+ withTempView("matched_ids") {
+ // `dep` is declared because the scan filters groups at runtime by the partition column
+ createAndInitTableWithReqAttrs("pk,dep,salary")
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `id` is neither declared nor in the write relation; the group filter must still map it
+ // onto the second read of the table that the subquery rewrite adds
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, dep STRING, salary INT, bonus INT, extra INT, " +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = Some("id INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+
+ test("column-update ReplaceData: scan reads only the referenced columns") {
+ createAndInitTableWithReqAttrs("pk,salary")
+
+ // `id` is neither declared nor referenced; `dep`, `bonus` and `extra` are read for the
+ // condition, the assignment and the CHECK constraint
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'",
+ "pk INT, dep STRING, salary INT, bonus INT, extra INT, _partition STRING, index INT")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update ReplaceData: subquery scan reads only the referenced columns") {
+ withTempView("deps") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE dep IN (SELECT dep FROM deps)",
+ "pk INT, dep STRING, salary INT, bonus INT, extra INT, _partition STRING, index INT")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+ }
+
+ test("column-update ReplaceData: unread partition column skips runtime group filtering") {
+ withTempView("matched_ids") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `dep` is neither declared nor referenced, so the scan doesn't read it and doesn't
+ // report it from filterAttributes(); every group is read and replaced
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, salary INT, bonus INT, extra INT, " +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = None)
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr", "software"))
+ }
+ }
+
+ test("column-update ReplaceData: runtime filter attribute the scan doesn't read is rejected") {
+ withTempView("matched_ids") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, salary INT",
+ """{ "pk": 1, "id": 1, "dep": "hr", "salary": 100 }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-cow-req-attrs", "pk,salary",
+ "column-update-cow-unread-filter-attrs", "true"))
+ import testImplicits._
+ Seq(1).toDF().createOrReplaceTempView("matched_ids")
+
+ val ex = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE id IN (SELECT * FROM matched_ids)")
+ }
+ assert(ex.getCondition == "DATA_SOURCE_INVALID_RUNTIME_FILTER_ATTRIBUTE.CANNOT_RESOLVE")
Review Comment:
**31.** Minor: this checks `attribute` and `method`, but not `relationOutput`, which is the pruned scan output this test is about, nor `scanClass`, so a change in which other columns the narrow scan reads would go unnoticed. Similarly, the three CHECK tests (`ColumnUpdateTableSharedTests.scala:165-169` and `:184-188`, and `DeltaBasedColumnUpdateTableSuite.scala:1094-1098`) check only the `CHECK_CONSTRAINT_VIOLATION` condition, so they cannot tell which constraint or value failed. `DataSourceV2CatalystRuntimeFilterSuite` and `CheckConstraintSuite` use `checkError` with every parameter for these conditions. Could these use `checkError`, or `checkErrorMatchPVals` for the scan class?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -128,12 +153,61 @@ class InMemoryRowLevelOperationTable private (
copied.replacedPartitions = replacedPartitions
copied.lastWriteInfo = lastWriteInfo
copied.lastWriteLog = lastWriteLog
+ copied.lastUpdatedColumns = lastUpdatedColumns
copied
}
override def newRowLevelOperationBuilder(
info: RowLevelOperationInfo): RowLevelOperationBuilder = {
- if (properties.getOrDefault(SUPPORTS_DELTAS, "false") == "true") {
+ lastUpdatedColumns = info.updatedColumns()
+ val splitUpdates = properties.getOrDefault(SPLIT_UPDATES, "false").toBoolean
+ if (properties.getOrDefault(COLUMN_UPDATE, "false") == "true") {
+ if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ }
+ } else if (properties.containsKey(COLUMN_UPDATE_REQ_ATTRS)) {
+ val reqCols = properties.get(COLUMN_UPDATE_REQ_ATTRS).split(",").map(_.trim)
+ if (splitUpdates) {
+ () => new DeltaBasedColumnUpdateSplitOperationWithReqAttrs(
+ info.command, reqCols, info.options)
+ } else {
+ () => new DeltaBasedColumnUpdateOperationWithReqAttrs(info.command, reqCols, info.options)
+ }
+ } else if (properties.getOrDefault(COLUMN_UPDATE_EMPTY_REQ_ATTRS, "false") == "true") {
+ // Test-only: returns an empty requiredDataAttributes() so we can verify Spark rejects it.
+ () => new DeltaBasedColumnUpdateOperationWithReqAttrs(
+ info.command, Array.empty, info.options)
+ } else if (properties.getOrDefault(COLUMN_UPDATE_COW, "false") == "true") {
+ if (info.updatedColumns().exists(_.describe() == "pk")) {
+ () => PartitionBasedOperation(info.command, info.options)
+ } else {
+ () => new PartitionBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ }
+ } else if (properties.containsKey(COLUMN_UPDATE_COW_REQ_ATTRS)) {
+ val reqCols = properties.get(COLUMN_UPDATE_COW_REQ_ATTRS).split(",").map(_.trim)
+ () => new PartitionBasedColumnUpdateOperationWithReqAttrs(
+ info.command, reqCols, info.options)
+ } else if (
+ properties.getOrDefault(COLUMN_UPDATE_COW_NO_WRITE_COLUMN_UPDATE, "false") == "true") {
+ () => new PartitionBasedColumnUpdateOperation(
+ info.command, info.updatedColumns().toSeq, info.options)
+ } else if (properties.getOrDefault(COLUMN_UPDATE_SPLIT, "false") == "true") {
Review Comment:
**37.** Nit: after item 25, `column-update` and `column-update-req-attrs` with `split-updates` (L163-179) build the same operations with the same arguments as `column-update-split` and `column-update-split-req-attrs` here, so these keys are now synonyms, and the req-attrs plus `split-updates` branch is never exercised. As a result, "column-update: split-updates sends a narrow UPDATE as delete and reinsert" (`DeltaBasedColumnUpdateTableSuite.scala:718`) and "column-update split: delete and reinsert carry the row ID and preserved metadata" (`:735`) run the same DDL, data, UPDATE and write-log assertions. Likewise, `column-update-cow-no-write-column-update` is still a dispatch key (L195-198), although item 18 made it a writer flag. Could the old `column-update-split` and `column-update-split-req-attrs` keys go (keeping `column-update-split-row-id`), with the split tests spelled as `column-update` plus `split-updates`, and one of the two tests dropped?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/GroupBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,637 @@
+/*
+ * 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.{SparkException, SparkUnsupportedOperationException}
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.Union
+import org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTable
+import org.apache.spark.sql.connector.write.RowLevelOperation.Command.UPDATE
+import org.apache.spark.sql.connector.write.RowLevelOperationTable
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, V2Writes}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StructField, StructType}
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+class GroupBasedColumnUpdateTableSuite extends UpdateTableSuiteBase
+ with ColumnUpdateTableSharedTests {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update-cow", "true")
+ props
+ }
+
+ override protected def reqAttrsTableProps(reqAttrs: String): java.util.Map[String, String] = {
+ java.util.Map.of("column-update-cow-req-attrs", reqAttrs)
+ }
+
+ override protected def fullWidthTableProps: java.util.Map[String, String] = java.util.Map.of()
+
+ test("column-update ReplaceData: column update schema has only the declared columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "dep", "salary"))
+ }
+
+ test("column-update ReplaceData: data correctness -- bonus preserved, salary updated") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ // CoW: pk=1 and pk=3 match (2 updates); the 'hr' partition has no other rows so no copies.
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update ReplaceData: subquery WHERE condition data correctness") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ withTempView("target_deps") {
+ import testImplicits._
+ val subqueryDF = Seq("hr").toDF()
+ subqueryDF.createOrReplaceTempView("target_deps")
+
+ sql(
+ s"""UPDATE $tableNameAsString
+ |SET salary = -1
+ |WHERE dep IN (SELECT * FROM target_deps)
+ |""".stripMargin)
+ }
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 2, numCopiedRows = 0)
+ }
+
+ test("column-update ReplaceData: runtime group filtering data correctness") {
+ Seq(true, false).foreach { dppEnabled =>
+ Seq(true, false).foreach { aqeEnabled =>
+ withSQLConf(
+ SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> dppEnabled.toString,
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqeEnabled.toString) {
+ withTable(tableNameAsString) {
+ withTempView("matched_pk") {
+ createAndInitTable("pk INT, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ import testImplicits._
+ val matchedPkDF = Seq(Some(1), None).toDF()
+ matchedPkDF.createOrReplaceTempView("matched_pk")
+
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = -1 " +
+ s"WHERE pk IN (SELECT * FROM matched_pk)",
+ primaryScanSchema =
+ "pk INT, salary INT, dep STRING, _partition STRING, index INT",
+ groupFilterScanSchema = Some("pk INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, 300, 30, "hr") :: Nil)
+
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+ }
+ }
+ }
+ }
+
+ test("column-update ReplaceData: updated and copied rows reach writeColumnUpdate with " +
+ "metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", null), data = Row(1, "hr", -1)),
+ writeColumnUpdateWithMetadataLogEntry(metadata = Row("hr", 1), data = Row(2, "hr", 200)))
+ }
+
+ test("column-update ReplaceData: without metadata, rows reach writeColumnUpdate(record)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |{ "pk": 3, "salary": 300, "bonus": 30, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "no-metadata", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Row(3, 300, 30, "software") :: Nil)
+ }
+
+ test("column-update ReplaceData: a writer overriding only writeColumnUpdate(record) " +
+ "receives rows with metadata") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "column-update-cow-record-only-writer", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastColumnUpdateWriteLog(
+ writeColumnUpdateLogEntry(data = Row(1, "hr", -1)),
+ writeColumnUpdateLogEntry(data = Row(2, "hr", 200)))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: without metadata, column update schema is the declared " +
+ "layout") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow", "true", "no-metadata", "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedColumnUpdateSchema = Some(StructType(Seq(PK_FIELD, DEP_FIELD,
+ StructField("salary", IntegerType, nullable = true)))))
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 110, 10, "hr") :: Row(2, 200, 20, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: connector missing writeColumnUpdate override is rejected") {
+ createAndInitTableReplaceDataNoWriteColumnUpdate("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[SparkUnsupportedOperationException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED",
+ parameters = Map(
+ "class" -> "org.apache.spark.sql.connector.catalog.NoWriteColumnUpdateOverrideWriter"))
+ }
+
+ private def createAndInitTableReplaceDataNoWriteColumnUpdate(
+ schemaString: String, jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-cow-no-write-column-update", "true"))
+ }
+
+ // ---------------------------------------------------------------------------
+ // Write relation tests: verify that only the write relation is narrowed to the
+ // declared columns, on both copy-on-write plan shapes.
+ // ---------------------------------------------------------------------------
+
+ private val reqAttrsOperationClass =
+ classOf[InMemoryRowLevelOperationTable#PartitionBasedColumnUpdateOperationWithReqAttrs].getName
+
+ private def createAndInitTableWithReqAttrs(reqAttrs: String): Unit = {
+ createAndInitTableWithProps(reqAttrsTableProps(reqAttrs))
+ }
+
+ test("column-update ReplaceData: write relation has the declared columns in declared order") {
+ createAndInitTableWithReqAttrs("salary,pk")
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'")
+ }
+
+ assert(!write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "salary", "pk")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update ReplaceData: subquery write relation has the declared columns") {
+ withTempView("deps") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE dep IN (SELECT dep FROM deps)")
+ }
+
+ assert(write.query.exists(_.isInstanceOf[Union]))
+ checkColumnUpdateWrite(write, "pk", "salary")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+ }
+
+ test("column-update ReplaceData: runtime group filtering on an undeclared condition column") {
+ withTempView("matched_ids") {
+ // `dep` is declared because the scan filters groups at runtime by the partition column
+ createAndInitTableWithReqAttrs("pk,dep,salary")
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `id` is neither declared nor in the write relation; the group filter must still map it
+ // onto the second read of the table that the subquery rewrite adds
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, dep STRING, salary INT, bonus INT, extra INT, " +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = Some("id INT, dep STRING"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr"))
+ }
+ }
+
+ test("column-update ReplaceData: scan reads only the referenced columns") {
+ createAndInitTableWithReqAttrs("pk,salary")
+
+ // `id` is neither declared nor referenced; `dep`, `bonus` and `extra` are read for the
+ // condition, the assignment and the CHECK constraint
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'",
+ "pk INT, dep STRING, salary INT, bonus INT, extra INT, _partition STRING, index INT")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+
+ test("column-update ReplaceData: subquery scan reads only the referenced columns") {
+ withTempView("deps") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE dep IN (SELECT dep FROM deps)",
+ "pk INT, dep STRING, salary INT, bonus INT, extra INT, _partition STRING, index INT")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 220, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ }
+ }
+
+ test("column-update ReplaceData: unread partition column skips runtime group filtering") {
+ withTempView("matched_ids") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq(1, 100).toDF().createOrReplaceTempView("matched_ids")
+
+ // `dep` is neither declared nor referenced, so the scan doesn't read it and doesn't
+ // report it from filterAttributes(); every group is read and replaced
+ executeAndCheckScans(
+ s"UPDATE $tableNameAsString SET salary = salary + bonus " +
+ s"WHERE id IN (SELECT * FROM matched_ids)",
+ primaryScanSchema = "pk INT, id INT, salary INT, bonus INT, extra INT, " +
+ "_partition STRING, index INT",
+ groupFilterScanSchema = None)
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "hr", 110, 10, 1) ::
+ Row(2, 2, "hr", 200, 20, 1) ::
+ Row(3, 3, "software", 300, 30, 1) :: Nil)
+ checkReplacedPartitions(Seq("hr", "software"))
+ }
+ }
+
+ test("column-update ReplaceData: runtime filter attribute the scan doesn't read is rejected") {
+ withTempView("matched_ids") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING, salary INT",
+ """{ "pk": 1, "id": 1, "dep": "hr", "salary": 100 }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-cow-req-attrs", "pk,salary",
+ "column-update-cow-unread-filter-attrs", "true"))
+ import testImplicits._
+ Seq(1).toDF().createOrReplaceTempView("matched_ids")
+
+ val ex = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE id IN (SELECT * FROM matched_ids)")
+ }
+ assert(ex.getCondition == "DATA_SOURCE_INVALID_RUNTIME_FILTER_ATTRIBUTE.CANNOT_RESOLVE")
+ assert(ex.getMessageParameters.get("attribute") == "`dep`")
+ assert(ex.getMessageParameters.get("method") == "filterAttributes()")
+ }
+ }
+
+ test("column-update ReplaceData: a connector without column updates reads every column") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of[String, String]())
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = 0",
+ "pk INT, salary INT, dep STRING, _partition STRING, index INT")
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, 0, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: undeclared distribution column is rejected") {
+ // without metadata, the write clusters by `dep`
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-cow-req-attrs", "pk,salary", "no-metadata", "true"))
+
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName
+ for {
+ cond <- Seq("pk = 1", "dep = 'hr'")
+ confs <- Seq(Nil, Seq(excludeColumnPruning))
+ } withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE $cond")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[`dep`]"))
+ }
+ }
+
+ test("column-update ReplaceData: other operations read every column") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of())
+ val fullScanSchema = "pk INT, salary INT, bonus INT, dep STRING, _partition STRING, index INT"
+
+ executeAndCheckScan(
+ s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1", fullScanSchema)
+
+ executeAndCheckScan(s"DELETE FROM $tableNameAsString WHERE pk > 1", fullScanSchema)
+
+ withTempView("source") {
+ import testImplicits._
+ Seq((1, 150)).toDF("pk", "salary").createOrReplaceTempView("source")
+ executeAndCheckScan(
+ s"""MERGE INTO $tableNameAsString t
+ |USING source s
+ |ON t.pk = s.pk
+ |WHEN MATCHED THEN UPDATE SET salary = s.salary
+ |""".stripMargin,
+ fullScanSchema)
+ }
+
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, 150, 10, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: DELETE and MERGE on a column-update table read every column") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "hr" }
+ |""".stripMargin)
+ val fullScanSchema = "pk INT, salary INT, bonus INT, dep STRING, _partition STRING, index INT"
+
+ executeAndCheckScan(s"DELETE FROM $tableNameAsString WHERE pk > 1", fullScanSchema)
+
+ withTempView("source") {
+ import testImplicits._
+ Seq((1, 150)).toDF("pk", "salary").createOrReplaceTempView("source")
+ executeAndCheckScan(
+ s"""MERGE INTO $tableNameAsString t
+ |USING source s
+ |ON t.pk = s.pk
+ |WHEN MATCHED THEN UPDATE SET salary = s.salary
+ |""".stripMargin,
+ fullScanSchema)
+ }
+
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, 150, 10, "hr") :: Nil)
+ }
+
+ test("column-update ReplaceData: assigned column outside requiredDataAttributes is rejected") {
+ createAndInitTableWithReqAttrs("pk,salary")
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET bonus = 0 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "missingColumns" -> "[`bonus`]"))
+ }
+
+ test("column-update ReplaceData: assigned column outside requiredDataAttributes is rejected " +
+ "on the subquery path") {
+ withTempView("deps") {
+ createAndInitTableWithReqAttrs("pk,salary")
+ import testImplicits._
+ Seq("hr").toDF("dep").createOrReplaceTempView("deps")
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET bonus = 0 WHERE dep IN (SELECT dep FROM deps)")
+ },
+ condition = "COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "missingColumns" -> "[`bonus`]"))
+ }
+ }
+
+ test("column-update ReplaceData: analyzed plan shows the declared columns as the write target") {
+ createAndInitTableWithReqAttrs("pk,salary")
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + bonus WHERE dep = 'hr'")
+ }
+
+ val header = write.simpleString(Int.MaxValue)
+ assert(header.matches("""ReplaceData RelationV2\[pk#\d+, salary#\d+\] .*"""), header)
+ }
+
+ test("column-update ReplaceData: an operation replaced after the rewrite is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ }
+ val relation = write.table.asInstanceOf[DataSourceV2Relation]
+ val operationTable = relation.table.asInstanceOf[RowLevelOperationTable]
+ val t = table
+ Seq(
+ t.PartitionBasedOperation(UPDATE, CaseInsensitiveStringMap.empty()) -> "[pk, salary, dep]",
+ new t.PartitionBasedColumnUpdateOperation(UPDATE) -> "[pk, dep]"
+ ).foreach { case (operation, expectedNames) =>
+ val replaced = write.withNewTable(
+ relation.copy(table = operationTable.copy(operation = operation)(
+ operationTable.operationInfo)))
+ checkError(
+ exception = intercept[SparkException](V2Writes(replaced)),
+ condition = "INTERNAL_ERROR",
+ parameters = Map("message" ->
+ (s"The write relation of operation ${operation.getClass.getName} has columns " +
+ s"[pk, dep, salary], but the operation expects $expectedNames.")))
+ }
+ }
+
+ test("column-update ReplaceData: an operation rebuilt with withTable keeps the updated columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ val write = executeAndKeepAnalyzedWrite {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ }
+ val relation = write.table.asInstanceOf[DataSourceV2Relation]
+ val operationTable = relation.table.asInstanceOf[RowLevelOperationTable]
+ val t = table
+ t.lastUpdatedColumns = null
+
+ val rebuilt = operationTable.withTable(t)
+
+ assert(t.lastUpdatedColumns.map(_.describe()).toSeq == Seq("salary"))
Review Comment:
**33.** Nit: this brings back a hand-written `lastUpdatedColumns` assertion, which round 4 ([r4074793101](https://github.com/apache/spark/pull/55518#discussion_r4074793101)) asked to replace with `checkLastUpdatedColumns`. L549 resets the field to `null` precisely to catch `withTable` not calling the builder, and in that case this line fails with a `NullPointerException`, while the helper fails with "no row-level operation was built" on the same live table. Could this be `checkLastUpdatedColumns("salary")`?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4315467514
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/SupportsColumnUpdates.java:
##########
@@ -0,0 +1,94 @@
+/*
+ * 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.write;
+
+import org.apache.spark.annotation.Experimental;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+
+/**
+ * A mix-in interface for {@link RowLevelOperation}. Data sources can implement this interface to
+ * receive a narrow row containing only the columns declared via {@link #requiredDataAttributes()}
+ * for updated, copied, and reinserted records, instead of the full table row.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present, updated, copied, and reinserted
+ * rows follow it: group-based operations receive them through
+ * {@link DataWriter#writeColumnUpdate(Object, Object)} or
+ * {@link DataWriter#writeColumnUpdate(Object)}, and operations that mix in {@link SupportsDelta}
+ * through {@link DeltaWriter#update} and {@link DeltaWriter#reinsert}. Inserted rows always follow
+ * {@link LogicalWriteInfo#schema()} and arrive through {@link DataWriter#write(Object)} and
+ * {@link DeltaWriter#insert}. When {@link LogicalWriteInfo#columnUpdateSchema()} is absent, every
+ * row follows {@link LogicalWriteInfo#schema()}, as for any other operation. Connectors must
+ * decide which layout to expect by whether {@link LogicalWriteInfo#columnUpdateSchema()} is
+ * present, not by the command. A builder that wants full-width rows for some command can check
+ * {@link RowLevelOperationInfo#command()} and build an operation without this mix-in. Currently
+ * Spark narrows only UPDATE.
+ * <p>
+ * The scan builder returned by {@link #newScanBuilder} should implement
+ * {@link org.apache.spark.sql.connector.read.SupportsPushDownRequiredColumns}, so that the scan
+ * does not read columns the command neither references nor declares. Otherwise the scan reads
+ * every column.
+ *
+ * @since 4.4.0
+ */
+@Experimental
+public interface SupportsColumnUpdates extends RowLevelOperation {
+ /**
+ * Returns the data column references required to perform this row-level operation.
+ * <p>
+ * When Spark narrows the write, the returned columns become
+ * {@link LogicalWriteInfo#columnUpdateSchema()}, in declared order. Implementations must include
+ * every column they want to receive: every column reported by
+ * {@link RowLevelOperationInfo#updatedColumns()}, plus any columns needed for row lookup or
+ * routing, e.g. a primary key. Columns that are not declared are absent from the rows the
Review Comment:
**5. A copy-on-write column update that assigns the connector's lookup column leaves the connector no way to find the original row.** For a group-based operation, an updated row carries only the new values of the declared columns, and metadata columns that are not preserved on update are nulled (`RewriteUpdateTable.scala:146-150`), so the original value of an assigned column reaches the connector through no channel. This paragraph still suggests "a primary key" for row lookup, and Spark guards the analogous case only for split delta operations (`validateNoRowIdReassignment`).
The in-repo fixture shows the result. On a `column-update-cow` table `pk INT NOT NULL, id INT, dep STRING` with rows (1, 1, hr), (2, 2, software) and (3, 3, hr), `UPDATE t SET pk = 10 WHERE id = 1` delivers `[pk=10, dep=hr]` with metadata `[hr, null]`. `PartitionBasedNarrowReplaceData.doCommit` looks the base row up by the new `pk` (`InMemoryRowLevelOperationTable.scala:736-738`), finds none, and commits (10, null, hr) instead of (10, 1, hr). `SET pk = 3 WHERE id = 1` silently copies `id` from the other row and commits (3, 3, hr). The delta path is fine thanks to `__original_row_id_pk`, but `GroupBasedColumnUpdateTableSuite` does not inherit "update with assignments to row ID" from `DeltaBasedUpdateTableSuiteBase`, so no test covers this.
Could you document that a group-based connector has to locate rows through metadata preserved on update, or build an operation without this mix-in when `updatedColumns()` contains its lookup column? Could the fixture also fail loudly when the lookup misses, and could the copy-on-write suite get a port of that test?
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/LogicalWriteInfo.java:
##########
@@ -65,4 +69,18 @@ default Optional<StructType> metadataSchema() {
throw new SparkUnsupportedOperationException(
"DATA_SOURCE_METADATA_SCHEMA_NOT_IMPLEMENTED", Map.of("class", getClass().getName()));
}
+
+ /**
+ * the schema of updated, copied, and reinserted rows from Spark to data source in a
+ * column-level update. Present when the operation mixes in {@link SupportsColumnUpdates} and
+ * Spark delivers these rows with the columns of
+ * {@link SupportsColumnUpdates#requiredDataAttributes()}, which currently happens only for
+ * UPDATE. It covers every table column if every column is declared. When present,
+ * {@link #schema()} covers only newly inserted rows.
+ *
+ * @since 4.4.0
+ */
+ default Optional<StructType> columnUpdateSchema() {
Review Comment:
**10.** Minor: `columnUpdateSchema()` silently defaults to `Optional.empty()`, while its siblings `rowIdSchema()` and `metadataSchema()` (L60-71) throw `DATA_SOURCE_*_SCHEMA_NOT_IMPLEMENTED`. For a narrow UPDATE, Spark sends an empty `schema()`. A connector-side wrapper that rebuilds the info, the way `FileTable.mergedWriteInfo` does (this PR had to add the field there), and forwards the other methods but not this one makes the connector see "absent". By the documented rule it then decodes the narrow rows of `DeltaWriter.update` and `reinsert` with the empty `schema()` layout, silently. Spark reads this method only on infos it builds itself, so a throwing default would not affect any Spark path. Could the default throw like its siblings?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteRowLevelCommand.scala:
##########
@@ -253,6 +326,24 @@ trait RewriteRowLevelCommand extends Rule[LogicalPlan] {
WriteDeltaProjections(rowProjection, rowIdProjection, metadataProjection)
}
+ /**
+ * Moves the query columns that no projection reads after the ones that are read, and returns
+ * the reordered query with a map from each old ordinal to the new one. The optimizer prunes the
+ * query columns after the last one the write reads, which then never changes an ordinal.
+ */
+ protected def orderReadColumnsFirst(
Review Comment:
**11.** Minor (design, optional): keeping the ordinal-based projections valid under pruning takes three cooperating pieces: this analysis-time reorder plus `remapOrdinals` at three call sites in `RewriteUpdateTable`, the prefix `references` in `RowLevelWrite` (`v2Commands.scala:367-370`), and the new `ColumnPruning` case. All of them rely on "read columns come first", which only `RewriteUpdateTable` establishes, and only for the plan shapes `reorderOutputs` handles. `extractOutputs` handles `MergeRows` (L351), but `reorderOutputs` throws for it (L364), which matters for the MERGE follow-up in the PR description. A path that skips the reorder does not fail either; the prefix `references` silently keeps every column up to the last read ordinal. Could `references` be exactly the attributes at the read ordinals (plus ordinal 0), with the `ColumnPruning` case remapping the projections when it prunes, as it already does for `Generate.unrequiredChildIndex` (`Optimizer.scala:1543-1549`)?
The optimized plan would be the same, without the analysis-time reorder and its shape restriction.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/RowLevelOperationInfo.java:
##########
@@ -37,4 +38,16 @@ public interface RowLevelOperationInfo {
* Returns the row-level SQL command (e.g. DELETE, UPDATE, MERGE).
*/
Command command();
+
+ /**
+ * Returns the columns being updated by this operation. Currently only UPDATE populates it;
+ * other commands report an empty array.
+ * <p>
+ * A column is reported only if it is assigned a new value, so identity assignments such as
Review Comment:
**8. An explicit `SET v = v` on a CHAR/VARCHAR column is reported here.** With the default CHAR/VARCHAR handling, assignment alignment wraps the value in a write-side length check, `Alias(StaticInvoke(varcharTypeWriteSideCheck, Cast(v, string), 10))` (`TableOutputResolver.checkUpdate`), and `isIdentityAssignment` strips only one `Alias` before `semanticEquals`, so it does not recognize the identity. For `t(pk INT NOT NULL, salary INT, v VARCHAR(10))` and a connector that declares `[pk, salary]`, `UPDATE t SET salary = 1, v = v` fails with `COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS` and `[v]`, although `v` cannot change. CHAR(n) behaves the same, and on a split delta operation `SET rid = rid` on a VARCHAR row-ID column is rejected with `COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT`. Unassigned CHAR/VARCHAR columns stay bare attributes, so only explicit self-assignments are affected, and both column-update fixtures declare whatever Spark reports, so no test notices.
Could this Javadoc say that an assignment is excluded only when its aligned value is the column itself, and that `SET v = v` on CHAR/VARCHAR (like `SET s.c1 = s.c1`) may still be reported? Seeing through the wrappers would skip the padding and length checks, so documenting it looks safer.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -373,6 +904,21 @@ private class DeltaBufferWriter(schema: StructType) extends BufferWriter(schema)
}
object InMemoryRowLevelOperationTable {
+ // Global holder for the scan schema of the last row-level operation. See the class-level
+ // `lastScanSchema` def for why this is a companion-object field rather than a per-instance one.
+ // Tests that assert on scan pruning should reset this before running the operation under test
+ // (see `RowLevelOperationSuiteBase.beforeEach`).
+ private[catalog] val lastScanSchemaRef = new AtomicReference[StructType]()
Review Comment:
**15.** Minor: this JVM-global, last-writer-wins `AtomicReference`, hooked into the existing `PartitionBasedOperation` and `DeltaBasedOperation` scan builders and reset in `before {}`, duplicates `executeAndCheckScan(s)` in `RowLevelOperationSuiteBase`, which asserts the same `scan.readSchema()` from the statement's own executed plan; `GroupBasedColumnUpdateTableSuite.scala:109`, `:285` and `:344` already use it. It also keeps the previous statement's value: a statement that plans no scan, e.g. a metadata-only `DELETE ... WHERE dep = 'hr'`, leaves the last UPDATE's schema behind, and `GroupBasedColumnUpdateTableSuite.scala:430-457` checks UPDATE, DELETE and MERGE against one expected string. This is the same vacuous-pass class as round 4 ([r4074793084](https://github.com/apache/spark/pull/55518#discussion_r4074793084)), which was fixed only for `null`. Two more one-sided checks: `checkLastScanSchema` and `executeAndCheckScan` compare with `sameType`, which ignores field-name case
by default, so the "table's own spelling" tests (`DeltaBasedColumnUpdateTableSuite.scala:395-408`, `GroupBasedColumnUpdateTableSuite.scala:660-675`) cannot detect a spelling regression, and an expected-empty `checkLastUpdatedColumns()` matches the initial `Array.empty`. Could the suites use `executeAndCheckScan(s)` and drop this global?
##########
common/utils/src/main/resources/error/error-conditions.json:
##########
@@ -1253,6 +1253,60 @@
],
"sqlState" : "22003"
},
+ "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates` but declared duplicate column(s) <duplicateAttributes> in `requiredDataAttributes()`. Each column must be declared at most once."
+ ],
+ "sqlState" : "42000"
+ },
+ "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates` but returned an empty array from `requiredDataAttributes()`. Connectors that opt into column-level updates must declare at least one required data attribute."
+ ],
+ "sqlState" : "42000"
+ },
+ "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates` but declared metadata column(s) <metadataAttributes> in `requiredDataAttributes()`. Declare metadata columns in `requiredMetadataAttributes()` instead."
+ ],
+ "sqlState" : "42000"
+ },
+ "COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates` but declared nested field(s) <nestedAttributes> in `requiredDataAttributes()`. Declare the root struct column instead of a nested field."
+ ],
+ "sqlState" : "42000"
+ },
+ "COLUMN_UPDATE_REQUIRED_DATA_ATTRIBUTES_MISSING_UPDATED_COLUMNS" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates` but its `requiredDataAttributes()` does not cover every column being updated. Missing columns: <missingColumns>. Connectors must include every column reported by `RowLevelOperationInfo.updatedColumns()` in `requiredDataAttributes()`."
+ ],
+ "sqlState" : "42000"
+ },
+ "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates` and represents UPDATE as delete and insert, but row ID column(s) <rowIds> do not reach the reinserted row, which then has no identity for the connector to place it by. Declare each data row ID column in `requiredDataAttributes()`, and return each metadata row ID column from `requiredMetadataAttributes()` with `MetadataColumn.PRESERVE_ON_REINSERT` set."
+ ],
+ "sqlState" : "42000"
+ },
+ "COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates` and represents UPDATE as delete and insert, so UPDATE cannot assign row ID column(s) <rowIds>. The reinserted row carries only the columns in `requiredDataAttributes()`, and with a new row ID it cannot be matched to the row whose other columns the connector must preserve. Do not assign row ID columns, or include every table column in `requiredDataAttributes()`."
+ ],
+ "sqlState" : "42000"
+ },
+ "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS" : {
+ "message" : [
+ "Connector <connector> mixes in `SupportsColumnUpdates`, but its write requires a distribution or ordering by column(s) <columns> that the column-level UPDATE does not read. Declare each data column in `requiredDataAttributes()` and each metadata column in `requiredMetadataAttributes()`."
Review Comment:
**12.** Minor: "that the column-level UPDATE does not read" is not accurate. The check compares against the columns the write receives (`rowLevelWrite.references`), so it also fires for an undeclared column that the UPDATE reads: `GroupBasedColumnUpdateTableSuite.scala:401-409` expects this error with `[dep]` for `... WHERE dep = 'hr'`, where the scan reads `dep`. The Javadoc in `SupportsColumnUpdates.java:76-84` already describes it correctly. Could the message say something like "by column(s) <columns> that are not passed to the write"?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/RewriteUpdateTable.scala:
##########
@@ -228,4 +240,114 @@ object RewriteUpdateTable extends RewriteRowLevelCommand {
val expandOutput = generateExpandOutput(attrs, outputs)
Expand(outputs, expandOutput, matchedRowsPlan)
}
+
+ /**
+ * Returns the output of the write relation, which is the layout of the data rows the connector
+ * receives. If the operation opts into column updates, this is the connector's
+ * `requiredDataAttributes()`, validated against the assignments. Otherwise it is the full
+ * relation output.
+ */
+ private def resolveWriteAttrs(
+ relation: DataSourceV2Relation,
+ operation: RowLevelOperation,
+ assignments: Seq[Assignment]): Seq[AttributeReference] = operation match {
+ case op: SupportsColumnUpdates with SupportsDelta if op.representUpdateAsDeleteAndInsert =>
+ val attrs = resolveDeclaredAttrs(relation, op, assignments)
+ val rowIdAttrs = resolveRowIdAttrs(relation, op)
+ validateNoRowIdReassignment(op, relation, attrs, assignments, rowIdAttrs)
+ validateRowIdDeclared(op, attrs, rowIdAttrs, resolveRequiredMetadataAttrs(relation, op))
+ attrs
+ case op: SupportsColumnUpdates =>
+ resolveDeclaredAttrs(relation, op, assignments)
+ case _ =>
+ relation.output
+ }
+
+ /**
+ * Resolves the connector's `requiredDataAttributes()` and checks that it covers every assigned
+ * column.
+ */
+ private def resolveDeclaredAttrs(
+ relation: DataSourceV2Relation,
+ operation: SupportsColumnUpdates,
+ assignments: Seq[Assignment]): Seq[AttributeReference] = {
+ val attrs = resolveRequiredDataAttrs(relation, operation)
+ validateUpdatedColumnsSubset(operation, assignments, attrs)
+ attrs
+ }
+
+ /**
+ * 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 {
Review Comment:
**16.** Minor: `apply` already computes `updatedAttrs = collectUpdatedAttrs(assignments)` (L47), the list reported as `updatedColumns()` and named by this error, but passes raw `assignments` to `resolveWriteAttrs` (L50). So this method and `validateNoRowIdReassignment` (L318-322) re-derive the non-identity predicate; aokolnychyi's earlier suggestion ([r3786035399](https://github.com/apache/spark/pull/55518#discussion_r3786035399)) already noted that `collectUpdatedAttrs` could be reused in a few other places. With `updatedAttrs`, these become `updatedAttrs.filterNot(...)` and `updatedAttrs.filter(rowIdAttrSet.contains)`, and both `.distinct` calls can go, since aligned assignments have one entry per column. Similarly, the split arm resolves the row-ID and metadata attributes (L256, L258) that `buildWriteDeltaPlan` resolves again (L171-172), the same caller-side recomputation as in round 1 ([r3192416031](https://github.com/apache/spark/pull/55518#discussion_r3192416031)). In `resol
veRequiredDataAttrs`, the `getOrElse(a.exprId, a)` fallback (`RewriteRowLevelCommand.scala:164`) is unreachable now that metadata columns are rejected, the cast there is a no-op, and the duplicate list follows `groupBy`'s hash order instead of the declared order. Could you pass `updatedAttrs` and resolve those attributes once?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/DeltaBasedColumnUpdateTableSuite.scala:
##########
@@ -0,0 +1,999 @@
+/*
+ * 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.SparkRuntimeException
+import org.apache.spark.sql.{AnalysisException, Row}
+import org.apache.spark.sql.catalyst.optimizer.ColumnPruning
+import org.apache.spark.sql.catalyst.plans.logical.WriteDelta
+import org.apache.spark.sql.connector.catalog.{Delete, InMemoryBaseTable, InMemoryRowLevelOperationTable, Reinsert}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType}
+
+/**
+ * Tests for UPDATE statements targeting connectors that mix in
+ * [[org.apache.spark.sql.connector.write.SupportsColumnUpdates]].
+ *
+ * When a connector supports column updates, updated rows contain only the declared columns
+ * (LogicalWriteInfo.columnUpdateSchema()) rather than the full table row.
+ */
+class DeltaBasedColumnUpdateTableSuite extends DeltaBasedUpdateTableSuiteBase {
+
+ override protected lazy val extraTableProps: java.util.Map[String, String] = {
+ val props = new java.util.HashMap[String, String]()
+ props.put("column-update", "true")
+ props
+ }
+
+ private val reqAttrsOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateOperationWithReqAttrs].getName
+
+ private val splitOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperation].getName
+
+ private val splitReqAttrsOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithReqAttrs]
+ .getName
+
+ private val splitWithRowIdOperationClass =
+ classOf[InMemoryRowLevelOperationTable#DeltaBasedColumnUpdateSplitOperationWithRowId].getName
+
+ test("column-update: column update schema contains a single assigned column") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: column update schema contains multiple assigned columns") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false),
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update: column update schema is pk and dep for a full identity update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = dep WHERE pk = 1")
+
+ // All assignments are identity, so updatedColumns is empty; pk and dep remain because the
+ // connector unconditionally declares them as base columns.
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Array(PK_FIELD, DEP_FIELD))))
+ }
+
+ test("column-update: row filter condition is orthogonal to column narrowing") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET dep = 'engineering' WHERE pk IN (1, 3)")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+ test("column-update: update all rows (no WHERE clause)") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("salary", IntegerType, nullable = true)
+ ))))
+ }
+
+ test("column-update: column update schema excludes identity assignments in a mixed UPDATE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ StructField("dep", StringType, nullable = false)
+ ))))
+ }
+
+
+ test("column-update: nested struct field update narrows to the root struct column") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+
+ checkLastUpdatedColumns("s")
+
+ // `dep` is present because the connector unconditionally declares it as a base column.
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq == Seq("pk", "dep", "s"))
+ }
+
+ test("column-update: nested field identity update reports root struct as updated") {
+ createAndInitTable("pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET s.c1 = s.c1 WHERE pk = 1")
+
+ checkLastUpdatedColumns("s")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, Row(1, 2), "hr") :: Nil)
+ }
+
+ test("column-update: updatedColumns is empty for DELETE") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"DELETE FROM $tableNameAsString WHERE dep = 'hr'")
+
+ checkLastUpdatedColumns()
+ }
+
+ test("column-update: row-level DELETE falls back to full-width rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"DELETE FROM $tableNameAsString WHERE pk < 2")
+
+ assert(!table.lastWriteInfo.columnUpdateSchema().isPresent,
+ s"DELETE must not report a column update schema: " +
+ s"${table.lastWriteInfo.columnUpdateSchema()}")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(2, 200, "hr") :: Nil)
+ }
+
+ test("column-update: MERGE falls back to full-width rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+ checkMergeFallsBackToFullWidthRows()
+ }
+
+ test("column-update split: MERGE falls back to full-width rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split", "true"))
+ checkMergeFallsBackToFullWidthRows()
+ }
+
+ private def checkMergeFallsBackToFullWidthRows(): Unit = {
+ withTempView("source") {
+ import testImplicits._
+ Seq((1, 150, "hr"), (3, 300, "software")).toDF("pk", "salary", "dep")
+ .createOrReplaceTempView("source")
+
+ sql(
+ s"""MERGE INTO $tableNameAsString t
+ |USING source s
+ |ON t.pk = s.pk
+ |WHEN MATCHED THEN UPDATE SET salary = s.salary
+ |WHEN NOT MATCHED THEN INSERT *
+ |""".stripMargin)
+
+ assert(!table.lastWriteInfo.columnUpdateSchema().isPresent,
+ s"MERGE must not report a column update schema: " +
+ s"${table.lastWriteInfo.columnUpdateSchema()}")
+ checkLastUpdatedColumns()
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 150, "hr") :: Row(2, 200, "software") :: Row(3, 300, "software") :: Nil)
+ }
+ }
+
+ test("column-update: data correctness -- single column update") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0)
+ }
+
+ test("column-update: data correctness -- update all rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary * 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 200, "hr") :: Row(2, 400, "software") :: Row(3, 600, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 3, numCopiedRows = 0)
+ }
+
+ test("column-update: data correctness -- mixed identity and real assignments") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = id, dep = 'engineering' WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 1, "engineering") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ checkUpdateMetrics(numUpdatedRows = 1, numCopiedRows = 0)
+ }
+
+ private def createAndInitTableWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-req-attrs", reqAttrs))
+ }
+
+ test("column-update: requiredDataAttributes - data correctness") {
+ createAndInitTableWithReqAttrs("dep,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: undeclared distribution column is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "pk,salary",
+ "column-update-cluster-by", "dep"))
+
+ val excludeColumnPruning = SQLConf.OPTIMIZER_EXCLUDED_RULES.key -> ColumnPruning.ruleName
+ Seq(Nil, Seq(excludeColumnPruning)).foreach { confs =>
+ withSQLConf(confs: _*) {
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNDECLARED_WRITE_REQUIREMENT_COLUMNS",
+ parameters = Map("connector" -> reqAttrsOperationClass, "columns" -> "[dep]"))
+ }
+ }
+ }
+
+ test("column-update: distribution by an undeclared row ID column is allowed") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update-req-attrs", "salary,dep",
+ "column-update-cluster-by", "pk"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 200, "hr") :: Nil)
+ }
+
+ test("column-update: partition source column not in requiredDataAttributes") {
+ createAndInitTableWithReqAttrs("pk,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastScanExcludes("dep")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, 3, "hr") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively on the row-ID " +
+ "column without adopting the declared spelling") {
+ createAndInitTableWithReqAttrs("PK,salary,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "dep"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+
+ test("column-update: undeclared, differently-cased columns in the RHS and WHERE clause " +
+ "are read with the table's own spelling") {
+ createAndInitTableWithReqAttrs("pk,dep,salary",
+ "pk INT NOT NULL, salary INT, bonus INT, extra INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "extra": 5, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + BONUS WHERE pk = 1 AND EXTRA > 3")
+
+ checkLastScanSchema("pk INT, salary INT, bonus INT, extra INT, dep STRING, _partition STRING")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, 110, 10, 5, "hr") :: Nil)
+ }
+
+ test("column-update: scan that returns unrequested columns still writes narrow rows") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update", "true",
+ InMemoryBaseTable.SIMULATE_PARTIAL_COLUMN_PRUNING, "true"))
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ checkLastScanIncludes("bonus")
+ val columnUpdateSchema = table.lastWriteInfo.columnUpdateSchema().get()
+ assert(!columnUpdateSchema.fieldNames.contains("bonus"),
+ s"bonus must not be in column update schema: $columnUpdateSchema")
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(1, -1, 10, "hr") :: Nil)
+ }
+
+ test("column-update: DELETE on a scan that returns unrequested columns is unaffected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, bonus INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "bonus": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "bonus": 20, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of(
+ "column-update", "true",
+ InMemoryBaseTable.SIMULATE_PARTIAL_COLUMN_PRUNING, "true"))
+
+ sql(s"DELETE FROM $tableNameAsString WHERE pk < 2")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString"),
+ Row(2, 200, 20, "software") :: Nil)
+ }
+
+ test("column-update: requiredDataAttributes resolves case-insensitively on an assigned " +
+ "column without adopting the declared spelling") {
+ createAndInitTableWithReqAttrs("pk,SALARY,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "dep"))
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, 101, "hr") :: Row(2, 200, "software") :: Nil)
+ }
+
+ test("column-update: case-sensitive analysis accepts columns that differ only in case") {
+ withSQLConf(SQLConf.CASE_SENSITIVE.key -> "true") {
+ createAndInitTableWithReqAttrs("pk,salary,SALARY,dep",
+ "pk INT NOT NULL, salary INT, SALARY INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "SALARY": 200, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "salary", "SALARY", "dep"))
+ checkAnswer(sql(s"SELECT * FROM $tableNameAsString"), Row(1, -1, 200, "hr") :: Nil)
+ }
+ }
+
+ test("column-update: empty requiredDataAttributes throws AnalysisException") {
+ createAndInitTable("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }""".stripMargin,
+ java.util.Map.of("column-update-empty-req-attrs", "true"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ parameters = Map("connector" -> reqAttrsOperationClass))
+ }
+
+ test("column-update: nested requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("pk,s.c1,dep",
+ "pk INT NOT NULL, s STRUCT<c1: INT, c2: INT>, dep STRING",
+ """{ "pk": 1, "s": { "c1": 1, "c2": 2 }, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET s.c1 = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_NESTED_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map("connector" -> reqAttrsOperationClass, "nestedAttributes" -> "[s.c1]"))
+ }
+
+ test("column-update: duplicate requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("pk,pk,salary", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map("connector" -> reqAttrsOperationClass, "duplicateAttributes" -> "[pk]"))
+ }
+
+ test("column-update: case-only duplicate requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("PK,pk,salary", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = salary + 1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map("connector" -> reqAttrsOperationClass, "duplicateAttributes" -> "[PK]"))
+ }
+
+ test("column-update: metadata column in requiredDataAttributes throws AnalysisException") {
+ createAndInitTableWithReqAttrs("pk,_partition,dep", "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET dep = 'x' WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "metadataAttributes" -> "[_partition]"))
+ }
+
+ test("column-update split: metadata column in requiredDataAttributes throws " +
+ "AnalysisException") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-req-attrs", "pk,_partition,dep"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET dep = 'x' WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_METADATA_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> splitReqAttrsOperationClass,
+ "metadataAttributes" -> "[_partition]"))
+ }
+
+ test("column-update: requiredDataAttributes rejects a column that does not exist") {
+ createAndInitTableWithReqAttrs("pk,nonexistent_col,id", "pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+ },
+ condition = "COLUMN_UPDATE_UNKNOWN_REQUIRED_DATA_ATTRIBUTE",
+ parameters = Map(
+ "connector" -> reqAttrsOperationClass,
+ "unknownAttributes" -> "[nonexistent_col]"))
+ }
+
+ test("column-update: column update schema excludes undeclared columns") {
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE pk = 1")
+
+ assert(table.lastWriteInfo.columnUpdateSchema.get.fieldNames.toSeq ==
+ Seq("pk", "dep", "salary"))
+ }
+
+ test("column-update: data correctness") {
+ createAndInitTable("pk INT NOT NULL, salary INT, id INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "id": 10, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "id": 20, "dep": "software" }
+ |{ "pk": 3, "salary": 300, "id": 30, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, 10, "hr") ::
+ Row(2, 200, 20, "software") ::
+ Row(3, -1, 30, "hr") :: Nil)
+ }
+
+ private def createAndInitTableSplit(schemaString: String, jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-split", "true"))
+ }
+
+ test("column-update split: column update schema has the declared columns") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE pk = 1")
+
+ checkLastWriteInfo(
+ expectedRowIdSchema = Some(StructType(Array(PK_FIELD))),
+ expectedMetadataSchema = Some(StructType(Array(PARTITION_FIELD, INDEX_FIELD_NULLABLE))),
+ expectedColumnUpdateSchema = Some(StructType(Seq(
+ PK_FIELD,
+ DEP_FIELD,
+ StructField("id", IntegerType, nullable = false)
+ ))))
+ }
+
+ test("column-update split: data correctness") {
+ createAndInitTableSplit("pk INT NOT NULL, id INT, dep STRING",
+ """{ "pk": 1, "id": 1, "dep": "hr" }
+ |{ "pk": 2, "id": 2, "dep": "software" }
+ |{ "pk": 3, "id": 3, "dep": "hr" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET id = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(1, -1, "hr") :: Row(2, 2, "software") :: Row(3, -1, "hr") :: Nil)
+ }
+
+ test("column-update split: row-ID reassignment on narrow write is rejected") {
+ // `extra` is neither declared, updated, nor referenced, so requiredDataAttributes()
+ // ([pk, dep, salary]) doesn't cover every table column -- the write stays genuinely narrow
+ // and reassigning the row ID must be rejected.
+ createAndInitTableSplit("pk INT NOT NULL, salary INT, dep STRING, extra STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr", "extra": "x" }
+ |{ "pk": 2, "salary": 200, "dep": "software", "extra": "y" }
+ |""".stripMargin)
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_REASSIGNMENT",
+ parameters = Map("connector" -> splitOperationClass, "rowIds" -> "[pk]"))
+ }
+
+ test("column-update split: row-ID reassignment is allowed when requiredDataAttributes " +
+ "covers every table column") {
+ createAndInitTableSplitWithReqAttrs("pk,salary,dep",
+ "pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin)
+
+ sql(s"UPDATE $tableNameAsString SET pk = pk + 10, salary = -1 WHERE dep = 'hr'")
+
+ checkAnswer(
+ sql(s"SELECT * FROM $tableNameAsString ORDER BY pk"),
+ Row(2, 200, "software") :: Row(11, -1, "hr") :: Nil)
+ }
+
+ private def createAndInitTableSplitWithReqAttrs(
+ reqAttrs: String,
+ schemaString: String,
+ jsonData: String): Unit = {
+ createAndInitTable(schemaString, jsonData,
+ java.util.Map.of("column-update-split-req-attrs", reqAttrs))
+ }
+
+ test("column-update split: undeclared row-ID column on narrow write is rejected") {
+ createAndInitTable("pk INT NOT NULL, salary INT, dep STRING",
+ """{ "pk": 1, "salary": 100, "dep": "hr" }
+ |{ "pk": 2, "salary": 200, "dep": "software" }
+ |""".stripMargin,
+ java.util.Map.of("column-update-split-row-id", "pk"))
+
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql(s"UPDATE $tableNameAsString SET salary = -1 WHERE dep = 'hr'")
+ },
+ condition = "COLUMN_UPDATE_SPLIT_ROW_ID_NOT_DECLARED",
+ parameters = Map("connector" -> splitWithRowIdOperationClass, "rowIds" -> "[pk]"))
+ }
+
+ // analyzes without executing, as this fixture can't commit a write that omits `pk`
+ private def analyzeWriteDelta(sqlText: String): WriteDelta = {
+ val parsed = spark.sessionState.sqlParser.parsePlan(sqlText)
+ spark.sessionState.executePlan(parsed).analyzed.collectFirst {
+ case w: WriteDelta => w
+ }.getOrElse(fail("couldn't find WriteDelta in analyzed plan"))
+ }
+
+ test("column-update split: metadata row ID preserved on reinsert need not be declared") {
Review Comment:
**14. The split configuration that `validateRowIdDeclared` explicitly accepts is never executed.** This test (an undeclared metadata row ID preserved on reinsert) only analyzes the plan (L697), because the fixture cannot commit it: the split commit needs `pk` in the narrow row (`InMemoryRowLevelOperationTable.scala:507`), and `DeltaBufferWriter.delete` calls `id.getInt(0)` on the STRING `_partition` (`InMemoryRowLevelOperationTable.scala:870`). More broadly, no narrow delta test checks the row-ID or metadata values in the write log. So a regression where the reinserted metadata is shifted by `orderReadColumnsFirst` or pruning, or arrives as null instead of the preserved value, would go unnoticed; L716 only catches `_partition` disappearing from the schema. Could you assert the write log of a default split UPDATE with `checkLastColumnUpdateWriteLog(...)`, including the reinsert metadata, and make the fixture able to execute this test?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryRowLevelOperationTable.scala:
##########
@@ -258,6 +354,407 @@ class InMemoryRowLevelOperationTable private (
}
}
+ // A delta-based operation that supports column-level updates: Spark sends only the declared
+ // columns in the row projection instead of the full row schema. It declares `pk` (the
+ // row-lookup key) and `dep` (an unconditionally declared base column) plus whatever columns
+ // Spark reports as being assigned via `RowLevelOperationInfo#updatedColumns()`.
+ class DeltaBasedColumnUpdateOperation(
Review Comment:
**19.** Nit: this block (L357-756) sits between `DeltaBasedOperation` (L303-355) and `TestDeltaBatchWrite` (L758-771), which only `DeltaBasedOperation` uses, while the class otherwise keeps each operation next to its batch write (`PartitionBasedOperation` at L236, `PartitionBasedReplaceData` at L288). AGENTS.md asks to "put the new member with the code it belongs with". A few more nits:
- `MimaExcludes.scala:50` now reads "Write schema narrowing ...": moving the filter to `v44excludes` dropped the "scan and write" wording that matches the PR title.
- L117 refers to `InMemoryRowLevelOperationTable.lastScanSchema`, but the companion member is `lastScanSchemaRef` (L911), and L910 refers to `RowLevelOperationSuiteBase.beforeEach`, but the reset is in `before { ... }`.
- Redundant double blank lines at L549-550, `RowLevelOperationSuiteBase.scala:357-358`, `DeltaBasedColumnUpdateTableSuite.scala:161-162` and `:393-394`, and `GroupBasedColumnUpdateTableSuite.scala:560-561` and `:658-659`.
Could you move the block below `TestDeltaBatchWrite` and fix these?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala:
##########
@@ -4638,6 +4638,92 @@ private[sql] object QueryCompilationErrors extends QueryErrorsBase with Compilat
messageParameters = Map("nullableRowIdAttrs" -> nullableRowIdAttrs.mkString(", ")))
}
+ def emptyRequiredDataAttributesError(connectorClass: String): Throwable = {
+ new AnalysisException(
+ errorClass = "COLUMN_UPDATE_EMPTY_REQUIRED_DATA_ATTRIBUTES",
+ messageParameters = Map("connector" -> connectorClass))
+ }
+
+ def duplicateRequiredDataAttributeError(
+ connectorClass: String,
+ duplicateAttributes: Seq[String]): Throwable = {
+ new AnalysisException(
+ errorClass = "COLUMN_UPDATE_DUPLICATE_REQUIRED_DATA_ATTRIBUTE",
+ messageParameters = Map(
+ "connector" -> connectorClass,
+ "duplicateAttributes" -> duplicateAttributes.mkString("[", ", ", "]")))
Review Comment:
**17.** Minor: the new `COLUMN_UPDATE_*` helpers format identifier lists by hand with `mkString("[", ", ", "]")` instead of `toSQLId`, which `QueryErrorsBase` asks for ("Elements of identifiers shall be wrapped by backticks by using `toSQLId()`"), and their inputs mix `ref.describe()` (`RewriteRowLevelCommand.scala:128`, `:138`, `:146`, `:154` and `V2Writes.scala:161`) with raw names (`RewriteUpdateTable.scala:291`, `:321`, `:347`). For a top-level column named `a.b`, the UNKNOWN error shows "[`a.b`]" while MISSING_UPDATED_COLUMNS shows "[a.b]", which reads like the nested field in the NESTED error's "[s.c1]". In the same suite, `dep` appears as "`dep`" (`GroupBasedColumnUpdateTableSuite.scala:375`) and as "[dep]" (`:409`). These parameters become part of the `checkError` contract once released. Could you use `toSQLId` for all of them?
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/write/DataWriter.java:
##########
@@ -82,6 +84,60 @@ default void write(T metadata, T record) throws IOException {
write(record);
}
+ /**
+ * Writes one updated or copied record with metadata in the column update layout.
+ * <p>
+ * When {@link LogicalWriteInfo#columnUpdateSchema()} is present for a row-level operation that
+ * does not mix in {@link SupportsDelta}, Spark passes updated and copied records to this method
+ * instead of {@link #write(Object, Object)} if the operation returns a non-empty
Review Comment:
**13.** Minor: the Javadoc of the existing `write(T metadata, T record)` (L72-74) still says "This method is used by group-based row-level operations to pass back metadata for records that are updated or copied", but for a narrow UPDATE those records come here instead (`WriteToDataSourceV2Exec.scala:786-787`). A connector author who reads only `write(T, T)` gets `DATA_SOURCE_WRITE_COLUMN_UPDATE_NOT_IMPLEMENTED`. This PR adds such a pointer to `DeltaWriter.update`, `DeltaWriter.reinsert` and `LogicalWriteInfo.schema()`, but not there. Could you add a sentence to L72-74, e.g. "When {@link LogicalWriteInfo#columnUpdateSchema()} is present for an operation that does not mix in {@link SupportsDelta}, Spark passes these records to {@link #writeColumnUpdate(Object, Object)} instead."?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "dongjoon-hyun (via GitHub)" <gi...@apache.org>.
dongjoon-hyun commented on PR #55518:
URL: https://github.com/apache/spark/pull/55518#issuecomment-6523668594
Gentle ping, @anuragmantri .
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-58111][SQL] Scan and write schema narrowing for column-level UPDATE in DSv2 [spark]
Posted by "aokolnychyi (via GitHub) via reviews" <re...@spark.apache.org>.
aokolnychyi commented on code in PR #55518:
URL: https://github.com/apache/spark/pull/55518#discussion_r4429367381
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/connector/write/RowLevelOperationTable.scala:
##########
@@ -60,4 +65,14 @@ private[sql] case class RowLevelOperationTable(
override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = {
operation.newWriteBuilder(info)
}
+
+ /**
Review Comment:
Is this something that we need in Spark? If needed by Iceberg, I wouldn't add it here. That's Iceberg's problem to handle and we want to migrate to state options for branches anyway.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org