HyukjinKwon commented on code in PR #59283:
URL: https://github.com/apache/spark/pull/59283#discussion_r4254325346
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/package.scala:
##########
@@ -90,6 +90,9 @@ package object util extends Logging {
private def usePrettyExpression(
e: Expression,
shouldTrimTempResolvedColumn: Boolean = false): Expression = e transform
{
+ case u: UpdateFields if
u.fieldOps.exists(_.isInstanceOf[UpdateNestedField]) =>
+ PrettyAttribute(
+ u.sqlString(toPrettySQL(_, shouldTrimTempResolvedColumn), identity),
u.dataType)
Review Comment:
**Non-blocking (P2):** `u.dataType` runs `UpdateFields.fieldSources`, which
starts with an unchecked `structExpr.dataType.asInstanceOf[StructType]`.
`toSQLExpr` goes through this case when it builds the
`DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE` error for an `UpdateFields` whose
input is not a struct. That happens from `CheckAnalysis`, and also from
`fieldSources` itself via `failOnTypeCheckResult`. Either way, the cast throws
before the `AnalysisException` is constructed.
For example, `spark.range(1).select($"id".withField("a.b", lit(1)))` now
fails with a raw `java.lang.ClassCastException`. So do `$"a".withField("b.x.y",
lit(2))` and `$"a".dropFields("b.x.y")` when `a.b` is an INT. The error has no
error class, message parameters, or query context, and Spark Connect clients
get the same unclassified error. In the merge target, these calls failed
analysis with an `AnalysisException` that names the field and its type. Valid
nested updates are not affected. The existing "called on a non-StructType
column" and "intermediate field is not a struct" tests use one- and two-part
paths, which never leave an `UpdateNestedField` in the failing node, so they
still pass.
Taking the data type only from a well-typed update keeps the rendered text
and avoids the cast. `PrettyAttribute` defaults to `NullType`, and
`checkInputDataTypes()` reuses the same cached field sources that `dataType`
uses:
```scala
case u: UpdateFields if
u.fieldOps.exists(_.isInstanceOf[UpdateNestedField]) =>
val sql = u.sqlString(toPrettySQL(_, shouldTrimTempResolvedColumn),
identity)
if (u.checkInputDataTypes().isSuccess) {
PrettyAttribute(sql, u.dataType)
} else {
PrettyAttribute(sql)
}
```
**Verification:**
- **Regression:** withField and dropFields with a two-part path on a
non-struct column fail with an AnalysisException whose condition is
DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE.
- **Regression:** withField and dropFields with a path that continues at
least two parts past a non-struct intermediate field fail with an
AnalysisException, not ClassCastException.
- **Compatibility:** Generated column names for valid single nested
withField and dropFields calls still match the existing expected strings.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:
##########
@@ -806,35 +987,131 @@ case class UpdateFields(structExpr: Expression,
fieldOps: Seq[StructFieldsOperat
override def prettyName: String = "update_fields"
- private lazy val newFieldExprs: Seq[(StructField, Expression)] = {
- def getFieldExpr(i: Int): Expression = structExpr match {
- case c: CreateNamedStruct => c.valExprs(i)
- case _ => GetStructField(structExpr, i)
- }
- val fieldsWithIndex =
structExpr.dataType.asInstanceOf[StructType].fields.zipWithIndex
- val existingFieldExprs: Seq[(StructField, Expression)] =
- fieldsWithIndex.map { case (field, i) => (field, getFieldExpr(i))
}.toImmutableArraySeq
- fieldOps.foldLeft(existingFieldExprs)((exprs, op) => op(exprs))
- }
+ private lazy val newFieldSources: Seq[(StructField,
UpdateFields.FieldSource)] =
+ UpdateFields.fieldSources(structExpr, fieldOps)
- private lazy val newFields: Seq[StructField] = newFieldExprs.map(_._1)
+ private lazy val newFields: Seq[StructField] = newFieldSources.map(_._1)
- lazy val newExprs: Seq[Expression] = newFieldExprs.map(_._2)
+ /** Returns the expression for one output field without constructing the
complete struct. */
+ def newExpr(ordinal: Int): Expression = newFieldSources(ordinal)._2 match {
+ case UpdateFields.Original(sourceOrdinal) => getFieldExpr(sourceOrdinal)
+ case UpdateFields.Updated(expr) => expr
+ }
- lazy val evalExpr: Expression = {
- val createNamedStructExpr = CreateNamedStruct(newFieldExprs.flatMap {
- case (field, expr) => Seq(Literal(field.name), expr)
- })
+ private def getFieldExpr(ordinal: Int): Expression = structExpr match {
+ case c: CreateNamedStruct => c.valExprs(ordinal)
+ case _ => GetStructField(structExpr, ordinal)
+ }
- if (structExpr.nullable) {
- If(IsNull(structExpr), Literal(null, dataType), createNamedStructExpr)
+ lazy val evalExpr: Expression = {
+ def buildExpression(structRef: Expression): Expression = {
+ val source = structExpr match {
+ case namedStruct: CreateNamedStruct => namedStruct
+ case _ => structRef
+ }
+ val fieldSources = UpdateFields.fieldSources(source, fieldOps)
+ val updatedValues = ArrayBuffer.empty[Expression]
+ val outputOrdinals = fieldSources.map { case (_, source) =>
+ source match {
+ case UpdateFields.Original(ordinal) => structExpr match {
+ case namedStruct: CreateNamedStruct =>
+ updatedValues += namedStruct.valExprs(ordinal)
+ -1
+ case _ => ordinal
+ }
+ case UpdateFields.Updated(expr) =>
+ updatedValues += expr
+ -1
+ }
+ }
+ UpdateFieldsExpression(structRef, updatedValues.toSeq, outputOrdinals,
newFields)
Review Comment:
Thanks, verified. A struct(...) input now lowers to a CreateNamedStruct with
only the surviving values, and no copy of the original struct or With
definition is kept. Dropped and overwritten aggregates, Python UDFs, and
subqueries are no longer in the lowered tree. The new tests cover drops,
overwrites, nested paths, and projection inlining, and check that the discarded
value and With are gone.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4241668312","thread_id":"inline:4241668312","verdict_sha256":"376a843e0801b33f093c16ee386ce4fd6a12672aa472f4cfb03d6ca9bc61f79e"}
-->
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/ComplexTypes.scala:
##########
@@ -32,9 +32,9 @@ object SimplifyExtractValueOps extends Rule[LogicalPlan] {
// Remove redundant field extraction.
case GetStructField(createNamedStruct: CreateNamedStruct, ordinal, _) =>
createNamedStruct.valExprs(ordinal)
- case GetStructField(u: UpdateFields, ordinal, _)if
!u.structExpr.isInstanceOf[UpdateFields] =>
+ case GetStructField(u: UpdateFields, ordinal, _) =>
Review Comment:
Thanks, verified. The guard is back, so the chain merges before the
extraction is simplified. The new ANSI tests cover nested extraction,
projection inlining, and nullable inputs. Agreed that the nullable top-level
variant was already in the merge target.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4241668315","thread_id":"inline:4241668315","verdict_sha256":"376a843e0801b33f093c16ee386ce4fd6a12672aa472f4cfb03d6ca9bc61f79e"}
-->
--
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]