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