[SPARK-42151][SQL] Align UPDATE assignments with table attributes

[aokolnychyi]

What changes were proposed in this pull request?

This PR adds a rule to align UPDATE assignments with table attributes.

Why are the changes needed?

These changes are needed so that we can rewrite UPDATE statements into executable plans for tables that support row-level operations. In particular, our row-level mutation framework assumes Spark is responsible for building an updated version of each affected row and that row is passed back to the data source.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

This PR comes with tests.

[aokolnychyi]

We need to think about a reliable way to check if default values have been resolved. Right now, it simply relies on the order of rules, which is fragile. Ideas are welcome.

[aokolnychyi]

I follow what we do for V2 inserts.

[aokolnychyi]

The logic in this method tries to follow by name resolution we have in V2 tables.

[aokolnychyi]

cc @huaxingao @cloud-fan @dongjoon-hyun @sunchao @viirya @gengliangwang

[aokolnychyi]

I am using DATATYPE_MISMATCH as it seems appropriate.

[aokolnychyi]

@cloud-fan, this doc gives a bit more details about why this PR is a prerequisite for rewriting UPDATEs. Let me know if this makes sense!

[aokolnychyi]

We can apply this rule only if table implements SupportsRowLevelOperations if that feels safer?

[viirya]

What happens if the table doesn't implement SupportsRowLevelOperations?

[aokolnychyi]

I plan to ignore such statements when rewriting UPDATEs into executable plans, like we do today for DELETE. This would allow data sources to inject their own handling.

[cloud-fan]

hmm, do we expect a data source that can directly update an inner field? For such data sources, this is a regression.

[cloud-fan]

Thinking about this more, I think this is required by the row level operation framework so we have no choice. Data sources can skip it (skipSchemaResolution return true) and use a more advanced implementation if they can.

[aokolnychyi]

Correct, the existing row-level APIs assume Spark is responsible for building an updated version of the row. That should work for Delta, Iceberg, Hudi, Hive ACID.

Once there is another use case, we should be able to extend the framework to cover it.

[aokolnychyi]

Let me also know if you think we should only apply this to implementations of SupportsRowLevelOperations.
Here is the original question.

[cloud-fan]

how about GetArrayStructField?

[cloud-fan]

For ALTER COLUMN we support a special syntax to reference any inner field, for example, array_col.element.field1, map_col.key.field2, etc. Shall we support this syntax in UPDATE as well? The related code is in StructType.findNestedField

[aokolnychyi]

We should eventually. This PR doesn't support updating arrays or maps, though. I wanted to work on it later and unblock further row-level operation development for now. For now, I throw an exception and support only nested fields in structs.

[aokolnychyi]

Actually, I can add support for those expressions here but fail temporary in the rewrite logic.

[viirya]

Maybe add a comment that this rule cannot be changed in order for now.

[aokolnychyi]

Added a comment above.

[johanl-db]

There should be a test case for UPDATE nested_struct_table SET s.n_i = 1" that ensures the struct s.n_s is preserved as a whole instead of recursing and generating assignments for each of its children.

This is important if s.n_s contain null values: the assignments must be (s.n_i = 1, s.n_s = s.n_s), not (s.n_i = 1, s.n_s.dn_i = s.n_s.dn_i, s.n_s.dn_l = s.n_s.dn_l) so that s.n_s is still null after the update.

[aokolnychyi]

Will make sure there is test case for this.

[aokolnychyi]

I have removed this resolution but the logic is same: runtime null checks, varchar/char length checks, etc.

[aokolnychyi]

@johanl-db, here is the test we talked about. If you have time to contribute any other tests or to check the alignment logic works for Delta, it would be great!

[aokolnychyi]

Failures don't seem to be related.

[dongjoon-hyun]

Thank you for updating PR, @aokolnychyi .

[cloud-fan]

A few ideas to make the code more robust:

1. I think it's better to operate on the resolved column expressions, instead of turning back the expression to Seq[String]
2. Given the parser rule for the UPDATE command, the column expression can only be AttributeReference or accessing (array of) struct's fields. We can group by expr.references.head to get a map from AttributeReference to Seq[Expression] and the corresponding update expressions.
3. We validate the map we got in step 2: for each top-level column, its expressions must be of the same tree height (to avoid updating both 'a.b' and 'a.b.c'), and must be different from each other.
4. Now it's easy to build the new update expressions: for each top-level column, if it doesn't have a match in the map, use the actual column value as the update expression, else ... (same algorithm below)

[aokolnychyi]

it's better to operate on the resolved column expressions

I agree, let's see if we can avoid the conversion to references.

We validate the map we got in step 2: for each top-level column, its expressions must be of the same tree height (to avoid updating both 'a.b' and 'a.b.c')

Could you elaborate a bit on how you see the tree height check? Like add a separate method for computing expression height? What about cases when it is OK to have different expression heights like 'a.b.n1' and 'a.c' where a, b, c are all structs?

[cloud-fan]

ah you are right, we can't simply check the tree height. I think a better way is to use a ExpressionSet to make sure these column expressions have no duplication.

[aokolnychyi]

Using ExpressionSet to detect duplicate assignments to a.b.c and a.b.c would be easy. What about cases like a.b and a.b.c where we assign a value to a struct and its field at the same time? Are you thinking of recursively adding all subparts of each column key to ExpressionSet? For instance, we would need to add a.b, a.b.c, a.b.c.d to ExpressionSet for a.b.c.d?

[aokolnychyi]

We could probably build ExpressionSet for each update key per top-level attribute and check the intersection across all ExpressionSet is empty. Let me know if that's similar to what you thought.

[cloud-fan]

yea this SGTM.

[aokolnychyi]

@cloud-fan, any ideas on how to avoid deconstructing Seq[String] when applying a set of assignments to a top-level attribute? The problem is that we recurse top to bottom in applyUpdates whereas assignment.key is a set of nested GetStructField calls with the outer expression referring the leaf column.

I can see ways to perform the validation without converting keys to Seq[String] but I don't see an easy way to avoid that in applyUpdates.

[cloud-fan]

I'm thinking about this

def alignAssignments(
    assignments: Seq[Assignment],
    attrs: Seq[Attribute]): Seq[Assignment] = {
  // use ExpressionSet to check assignments have no duplication
  ...
  attrs.map { attr =>
    Assignment(attr, applyUpdates(assignments, attr))
  }
}

def applyUpdates(
    assignments: Seq[Assignment],
    col: Expression): Expression = {
  val (exactAssigments, others) = assignments.partition { assignment =>
    assigment.key.semanticEquals(col)
  }
  val relatedAssignments = others.filter { assignment =>
    assigment.key.exists(_.semanticEquals(col))
  }
  assert(exactAssigments.length <= 1)
  if (exactAssigments.nonEmpty) {
    if (relatedAssignments.nonEmpty) fail...
    exactAssigments.head.value
  } else {
    if (relatedAssignments.isEmpty) {
      col
    } else {
      assert(col.dataType.isInstanceOf[StructType])
      CreatedStruct(col.dataType.asInstanceOf[StructType].fields.flatMap { field =>
        Literal(field.name) :: applyUpdates(relatedAssignments, GetStruct(col, field.name)) :: Nil
      })
    }
  }
}

[aokolnychyi]

Perfect, I forgot about exists. Thanks!

[aokolnychyi]

@cloud-fan, I've changed the approach to avoid deconstructing references. However, I decided to keep the validation while recursing vs doing this in a separate step as we discussed. When I tried to implement that idea, it turned out to be pretty involved with lots of edge cases. For instance, we can't have multiple assignments per top-level key but keys can reference top-level fields many times, a.b.c and a.b.d are allowed but a.b and a.b.c are not.

It felt easier to validate while recursing, similar to TableOutputResolver.

Let me know what you think.

[aokolnychyi]

Added this for better error messages.

[aokolnychyi]

@viirya, I decided not to align assignments if tables don't extend SupportsRowLevelOperations. That way, data sources using their own implementations won't be affected. They can still use AssignmentUtils.

[aokolnychyi]

Copied from ResolveOutputRelation to preserve the existing behavior for data sources that rely on custom implementations.

[aokolnychyi]

Failures in streaming tests don't seem related.

[cloud-fan]

thanks, merging to master!

[aokolnychyi]

Thanks for reviewing, @cloud-fan @huaxingao @dongjoon-hyun @viirya @johanl-db!