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