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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
