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