bhollis-dbx commented on code in PR #59283:
URL: https://github.com/apache/spark/pull/59283#discussion_r4250537112
##########
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:
Fixed in 049f52ef8. Named-struct inputs now lower to a `CreateNamedStruct`
containing only surviving values; the original struct and unused `With`
definitions are removed. Added failing-aggregate regressions covering
dropped/overwritten fields, nested paths, and projection inlining, plus
inspection assertions that discarded values and `With` nodes are absent.
##########
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:
Restored the guard in 049f52ef8 so chains merge before extraction
simplification. Added ANSI-failure regressions for nested extraction and
projection-inlined inputs, plus direct extracted-nested chains with nullable
inputs. This fixes the newly introduced regression; it does not claim to fix
the preexisting nullable-top-level variant.
--
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]
