bhollis-dbx commented on code in PR #59283:
URL: https://github.com/apache/spark/pull/59283#discussion_r4273020527
##########
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`. */
+ private case class Updated(expr: Expression) extends FieldSource
+
+ /**
+ * Applies the requested field operations while retaining whether each
output field is copied or
+ * computed. This provenance lets evaluation copy unchanged fields by
ordinal and lets extraction
+ * return one field without constructing the complete struct.
+ */
+ private def fieldSources(
+ structExpr: Expression,
+ fieldOps: Seq[StructFieldsOperation]): Seq[(StructField, FieldSource)] =
{
+ val structType = structExpr.dataType.asInstanceOf[StructType]
+ def expression(source: FieldSource): Expression = source match {
+ case Original(ordinal) => structExpr match {
+ case namedStruct: CreateNamedStruct => namedStruct.valExprs(ordinal)
+ case _ => GetStructField(structExpr, ordinal)
+ }
+ case Updated(expr) => expr
+ }
+ val fields: Seq[(StructField, FieldSource)] =
structType.fields.zipWithIndex.map {
+ case (field, ordinal) =>
+ field -> Original(ordinal)
+ }.toImmutableArraySeq
+ def applyOperation(
+ current: Seq[(StructField, FieldSource)],
+ operation: StructFieldsOperation): Seq[(StructField, FieldSource)] =
operation match {
+ case op: WithField =>
+ val updated = StructField(op.name, op.valExpr.dataType,
op.valExpr.nullable) ->
+ Updated(op.valExpr)
+ val result = ArrayBuffer.empty[(StructField, FieldSource)]
+ var hasMatch = false
+ current.foreach { fieldSource =>
+ if (op.resolver(fieldSource._1.name, op.name)) {
+ hasMatch = true
+ result += updated
+ } else {
+ result += fieldSource
+ }
+ }
+ if (!hasMatch) result += updated
+ result.toSeq
+ case op: DropField =>
+ current.filterNot { case (field, _) => op.resolver(field.name,
op.name) }
+ case op =>
+ throw SparkException.internalError(
+ s"Unsupported StructFieldsOperation: ${op.getClass.getName}")
+ }
+ fieldOps.foldLeft(fields) { (current, operation) =>
+ val expanded = operation match {
+ case nested: UpdateNestedField =>
+ val currentType = StructType(current.map(_._1))
+ val field = ExtractValue(
+ Literal.create(null, currentType), Literal(nested.name),
nested.resolver)
+ .asInstanceOf[GetStructField]
+ val updated = CurrentOrigin.withOrigin(nested.origin) {
+ expression(current(field.ordinal)._2) match {
+ case UpdateFields(input, previousOps) =>
+ UpdateFields(input, previousOps ++ nested.fieldOps)
+ case input => UpdateFields(input, nested.fieldOps)
+ }
+ }
+ if (updated.checkInputDataTypes().isFailure) {
+ TypeCoercionValidation.failOnTypeCheckResult(updated)
+ }
+ WithField(nested.name, updated)
Review Comment:
Fixed in 3fe6dd456. Nested updates replace the resolved parent by ordinal
instead of matching its name again. Tests cover both case-sensitivity
transitions, withField/dropFields, analyzed DataFrames, and temp views,
checking results and schemas.
--
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]
