HyukjinKwon commented on code in PR #59283:
URL: https://github.com/apache/spark/pull/59283#discussion_r4266184308
##########
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:
Thanks, verified on the current head. usePrettyExpression now reads dataType
only after checkInputDataTypes succeeds, so invalid nested paths through a
non-struct value fail with DATATYPE_MISMATCH.UNEXPECTED_INPUT_TYPE instead of
ClassCastException. The new test covers two-part and deeper paths for both
withField and dropFields, including the query context.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4254325346","thread_id":"inline:4254325346","verdict_sha256":"ab7c8fca5d849377fb4ce9962d2913f156f779311c0512e9b3e488f367c05a58"}
-->
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala:
##########
@@ -806,35 +989,151 @@ 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)
+ private[sql] def sqlString(
+ render: Expression => String,
+ quoteField: String => String): String = {
+ def updateSql(inputSql: String, operations: Seq[StructFieldsOperation]):
String = {
+ val operationsSql = operations.map {
+ case nested: UpdateNestedField =>
+ val fieldSql = if (structExpr.resolved) {
+ s"$inputSql.${quoteField(nested.name)}"
+ } else {
+ s"$inputSql[${render(Literal(nested.name))}]"
+ }
+ s"WithField(${updateSql(fieldSql, nested.fieldOps)})"
+ case operation => render(operation)
+ }
+ s"$prettyName(${(inputSql +: operationsSql).mkString(", ")})"
}
- 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))
+ updateSql(render(structExpr), fieldOps)
}
- private lazy val newFields: Seq[StructField] = newFieldExprs.map(_._1)
+ override def sql: String = sqlString(_.sql, quoteIdentifier)
- lazy val newExprs: Seq[Expression] = newFieldExprs.map(_._2)
+ private lazy val newFieldSources: Seq[(StructField,
UpdateFields.FieldSource)] =
+ UpdateFields.fieldSources(structExpr, fieldOps)
- lazy val evalExpr: Expression = {
- val createNamedStructExpr = CreateNamedStruct(newFieldExprs.flatMap {
- case (field, expr) => Seq(Literal(field.name), expr)
- })
+ private lazy val newFields: Seq[StructField] = newFieldSources.map(_._1)
- if (structExpr.nullable) {
- If(IsNull(structExpr), Literal(null, dataType), createNamedStructExpr)
- } else {
- createNamedStructExpr
+ /** 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
+ }
+
+ private def getFieldExpr(ordinal: Int): Expression = structExpr match {
+ case c: CreateNamedStruct => c.valExprs(ordinal)
+ case _ => GetStructField(structExpr, ordinal)
+ }
+
+ lazy val evalExpr: Expression = {
+ def buildExpression(structRef: Expression): Expression = {
+ val fieldSources = UpdateFields.fieldSources(structRef, fieldOps)
+ val updatedValues = ArrayBuffer.empty[Expression]
+ val outputOrdinals = fieldSources.map {
+ case (_, UpdateFields.Original(ordinal)) => ordinal
+ case (_, UpdateFields.Updated(expr)) =>
+ updatedValues += expr
+ -1
+ }
+ UpdateFieldsExpression(structRef, updatedValues.toSeq, outputOrdinals,
newFields)
+ }
+ structExpr match {
+ case namedStruct: CreateNamedStruct =>
+ val values = namedStruct.valExprs.toArray
+ CreateNamedStruct(newFieldSources.flatMap { case (field, source) =>
+ val value = source match {
+ case UpdateFields.Original(ordinal) => values(ordinal)
+ case UpdateFields.Updated(expr) => expr
+ }
+ Seq(Literal(field.name), value)
+ })
+ case _ if fieldOps.exists(_.isInstanceOf[UpdateNestedField]) =>
+ With(structExpr) { case Seq(structRef) => buildExpression(structRef) }
+ case _ =>
+ buildExpression(structExpr)
}
}
}
object UpdateFields {
+ /** Describes where an output field gets its value without creating an
expression for a copy. */
+ private sealed trait FieldSource
+
+ /** Copies the field at `ordinal` from the input struct. */
+ private case class Original(ordinal: Int) extends FieldSource
+
+ /** Computes the field from an expression supplied to `withField`. */
Review Comment:
**Nit (P3):** This says every `Updated` value is a `withField` expression,
but `fieldSources` also stores `Updated(updated)` for each `UpdateNestedField`
(line 1131), where `updated` is an `UpdateFields` built over the current parent
field. That includes nested `dropFields` paths such as
`col("s").dropFields("a.b")`, which have no `withField` value at all. On the
`With` lowering path these values hold `GetStructField(structRef, ...)`, so
they depend on the bound input. A reader who trusts this comment could treat
them as independent user expressions and move them out of that scope. Only the
comment is affected; runtime behavior is unchanged. Maybe something like:
```suggestion
/** Computes the field from a `withField` value or a nested update of the
current field. */
```
**Verification:**
- **Inspection:** The Updated doc matches every place fieldSources
constructs an Updated source, including the UpdateNestedField branch.
--
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]
