HyukjinKwon commented on code in PR #59283:
URL: https://github.com/apache/spark/pull/59283#discussion_r4263829487
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/UpdateFields.scala:
##########
@@ -39,37 +39,54 @@ object OptimizeUpdateFields extends Rule[LogicalPlan] {
}
}
+ private def optimizeFieldOps(
+ fieldOps: Seq[StructFieldsOperation]): Seq[StructFieldsOperation] = {
+ val operations = fieldOps.map {
+ case nested: UpdateNestedField =>
+ nested.copy(fieldOps = optimizeFieldOps(nested.fieldOps))
+ case other => other
+ }
+ if (!operations.forall {
+ case _: WithField | _: UpdateNestedField => true
+ case _ => false
+ }) {
+ operations
+ } else {
+ val combined = mutable.LinkedHashMap.empty[String, StructFieldsOperation]
+ operations.foreach { operation =>
+ val name = operation match {
+ case WithField(name, _) => name
+ case UpdateNestedField(name, _) => name
+ }
+ val key = if (conf.caseSensitiveAnalysis) name else
name.toLowerCase(Locale.ROOT)
+ val updated = operation match {
+ case nested: UpdateNestedField => combined.get(key) match {
+ case Some(previous: UpdateNestedField) =>
+ nested.copy(fieldOps = optimizeFieldOps(previous.fieldOps ++
nested.fieldOps))
Review Comment:
**Nit (P3):** When a key already holds an `UpdateNestedField`, this runs
`optimizeFieldOps` again on `previous.fieldOps ++ nested.fieldOps`. That
re-maps every child already optimized and rebuilds the `LinkedHashMap` over the
whole accumulated list. The flatten case first turns
`names.foldLeft(col("s"))((c, n) => c.withField(s"a.$n", v))` into a single
`UpdateFields(s, [UpdateNestedField(a, [f1]), ..., UpdateNestedField(a,
[fm])])`. The Substitution-batch application then does about m^2/2 operation
visits, each with a key normalization, where the merge-target dedupe was a
single linear pass. Programmatic updates of many fields under one parent are
the wide-schema workload this PR targets, and the new linear-growth test bounds
node counts rather than merge work. Only planning time is affected; results are
unchanged.
Accumulating each key's nested operations in order and optimizing each key's
list once, or merging into a per-key builder as operations arrive, would keep
this linear. The output should stay exactly as today. That includes the point
where a nested list containing `DropField` stops further deduplication, the
handling of a preceding or following `WithField` on the same key, and the
last-seen casing.
**Verification:**
- **Compatibility:** Optimized operation lists for mixed
WithField/UpdateNestedField chains, including overwrites, nested lists
containing DropField, and case-insensitive keys, are identical to the current
output.
- **Regression:** A chain of thousands of same-parent nested withField calls
analyzes with merge work that grows linearly with chain length.
##########
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:
**Non-blocking (P2):** `fieldSources` finds the parent field with the
`UpdateNestedField`'s own resolver (`ExtractValue(..., nested.resolver)` a few
lines up). It then applies the update through this new `WithField`, whose
`resolver` is `SQLConf.get.resolver` evaluated when it is constructed (line
710). `fieldSources` runs once at analysis (`newFieldSources`, which backs
`dataType` and `newFields`) and again when `ReplaceUpdateFieldsExpression`
lowers the node (`evalExpr` -> `buildExpression`). `buildExpression` pairs the
recomputed ordinals with the analysis-time `newFields`.
So the two passes disagree when `spark.sql.caseSensitive` changes between
analysis and execution. Take a struct column `s` with field `a` under
case-insensitive resolution: `val df2 = df.select(col("s").withField("A.x",
lit(9)))` analyzes to `struct<A: struct<..., x>>`. After
`spark.conf.set("spark.sql.caseSensitive", "true")`, `df2.collect()` optimizes
under the new setting. The nested lookup still finds ordinal 0, but
`WithField("A")` no longer matches `a`, so it appends. That gives
`outputOrdinals = [0, -1]` against a single output field.
`UpdateFieldsExpression.dataType` zips them down to the original `a`, so `x` is
silently dropped and the optimized type no longer matches the analyzed
attribute. Going the other way on a struct that has both `a` and `A` overwrites
both. A temp view that stores the analyzed plan reaches the same path. Before
this PR, the nested `WithField` was built once at `Column` conversion, so this
also needed an optimizer rewrite of that operation. This af
fects classic execution only; Spark Connect re-analyzes on every execution.
The parent's ordinal is already known at this point. Replacing
`current(field.ordinal)` in place with `StructField(nested.name,
updated.dataType, updated.nullable) -> Updated(updated)` would keep the
analysis-time layout, and the rename to the call's casing, without a second
name match.
**Verification:**
- **Regression:** Apply withField("A.x", ...) to a struct with field a under
case-insensitive analysis, then execute after switching to case-sensitive: the
result has the analyzed schema with x populated.
- **Compatibility:** Nested withField and dropFields results, schemas, and
generated names under case-sensitive and case-insensitive resolution are
unchanged when the setting does not change.
--
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]
