HyukjinKwon commented on code in PR #59283:
URL: https://github.com/apache/spark/pull/59283#discussion_r4274818970
##########
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) }
Review Comment:
**Nit (P3):** This `With` keeps the struct input only in a definition, and
`EquivalentExpressions.childrenToRecurse` visits `With.child` but never the
definitions. The input therefore stops being a subexpression-elimination
candidate. Before, the lowered update read it through `GetStructField` copies
(and `IsNull` for a nullable input), which CSE could share with other uses in
the same projection.
For example, take `val s = from_json($"j", schema)` or a deterministic
struct UDF. `df.select(s.withField("a.b", lit(1)), s.getField("c"))` used to
evaluate `s` once per row. It now evaluates it twice: once for the `With`
definition and once for the extraction. Results are the same, but an expensive
input costs twice as much on a path meant to make updates cheaper. Flat updates
and `struct(...)` inputs are not affected. The new 'field updates should
preserve shared value subexpression elimination' test shares only `withField`
values, not the input, so it does not catch this.
**Recommended change:** When subexpression elimination recurses into a With,
also recurse into each definition whose reference the With body reaches through
always-evaluated positions, so the nested-path lowering's struct input can be
shared with other uses again. Add subexpression-elimination unit coverage and
Column withField/dropFields execution tests for a shared struct input.
**Why this works:** Walk the With child with the existing always-evaluated
rules (ConditionalExpression.alwaysEvaluatedInputs, short-circuit peeling,
CodegenFallback and HigherOrderFunction stops, nested With scopes) and collect
the CommonExpressionRef ids reached. Recurse into the definitions with those
ids as always-evaluated children. supportedExpression still rejects candidates
that hold unbound references, so a hoisted definition child is evaluated before
the With exactly once, and codegen and interpreted With evaluation pick up the
shared value.
**Scope:** Let CSE see With definitions that the body always evaluates, and
cover a shared struct input end to end.
**Compatibility:** Definitions that the With body reaches only lazily stay
unevaluated when their reference is not reached.
**Risks:** Treating a conditionally evaluated definition as always evaluated
would hoist it eagerly and could raise ANSI errors or run side effects; the
walk must use exactly the existing always-evaluated rules.
**Constraints:** Keep the existing With scope-safety tests passing: lazy
definitions, CodegenFallback children and guarded values must still produce no
common subexpression.
**Success:** With subexpression elimination enabled, a deterministic struct
input read by a nested-path withField/dropFields and by another expression in
the same projection is evaluated once per row on interpreted and code-generated
paths. A With definition that the body references only under a conditional
branch or a CodegenFallback, or does not reference, is still not hoisted.
**Verification:**
- **Behavior:** A counting deterministic struct-producing UDF shared between
a nested-path withField and a separate extraction from the same input is
evaluated once per row on each evaluation path, with correct results.
- **Regression:** Lazy, CodegenFallback and guarded With definitions still
yield no common subexpression, and no candidate holds an unbound reference.
- **Compatibility:** Existing nested withField/dropFields results,
nullability and nondeterministic once-per-row evaluation are unchanged.
##########
sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala:
##########
@@ -445,6 +446,428 @@ class ColumnExpressionSuite extends SharedSparkSession {
SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY")(f)
}
+ test("withField should evaluate a nondeterministic struct expression once") {
+ onEachEvalPath {
+ val counter = new java.util.concurrent.atomic.AtomicLong()
+ val nextStruct = udf(() => {
+ val value = counter.getAndIncrement()
+ (value, value)
+ }).asNondeterministic()
+ val structExpr = nextStruct()
+ val df = spark.range(0, 10, 1, 1)
+ df.select(structExpr.withField("copy", structExpr)).collect().foreach {
row =>
+ val result = row.getStruct(0)
+ val copy = result.getStruct(2)
+ assert(result.getLong(0) == result.getLong(1))
+ assert(copy.getLong(0) == copy.getLong(1))
+ assert(result.getLong(0) != copy.getLong(0))
+ }
+ }
+ }
+
+ test("withField should evaluate a reused nondeterministic value separately")
{
+ onEachEvalPath {
+ val counter = new java.util.concurrent.atomic.AtomicLong()
+ val nextStruct = udf(() => {
+ val value = counter.getAndIncrement()
+ (value, value)
+ }).asNondeterministic()
+ val structExpr = nextStruct().expr
+ val updated = Column(UpdateFields(structExpr,
+ Seq(WithField("copy", GetStructField(structExpr, 0)))))
+ val df = spark.range(0, 10, 1, 1)
+ df.select(updated).collect().foreach { row =>
+ val result = row.getStruct(0)
+ assert(result.getLong(0) == result.getLong(1))
+ assert(result.getLong(0) != result.getLong(2))
+ }
+ }
+ }
+
+ test("withField should evaluate a reused nested update separately") {
+ onEachEvalPath {
+ val counter = new java.util.concurrent.atomic.AtomicLong()
+ val nextStruct = udf(() => {
+ val value = counter.getAndIncrement()
+ (value, value)
+ }).asNondeterministic()
+ val structExpr = nextStruct().expr
+ val nestedUpdate = UpdateFields(structExpr, Seq(WithField("extra",
Literal(1))))
+ val updated = Column(UpdateFields(structExpr,
+ Seq(WithField("copy", GetStructField(nestedUpdate, 0)))))
+ val df = spark.range(0, 10, 1, 1)
+ df.select(updated).collect().foreach { row =>
+ val result = row.getStruct(0)
+ assert(result.getLong(0) == result.getLong(1))
+ assert(result.getLong(0) != result.getLong(2))
+ }
+ }
+ }
+
+ test("field updates should preserve shared value subexpression elimination")
{
+ withSQLConf(SQLConf.SUBEXPRESSION_ELIMINATION_ENABLED.key -> "true",
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") {
+ for {
+ nested <- Seq(false, true)
+ external <- Seq(false, true)
+ opaque <- Seq(false, true)
+ } {
+ val calls = spark.sparkContext.longAccumulator
+ val sourceCalls = spark.sparkContext.longAccumulator
+ val value = udf((id: Long) => {
+ calls.add(1)
+ id + 1
+ }).apply($"id")
+ val source = if (opaque) {
+ udf((id: Long) => {
+ sourceCalls.add(1)
+ ((id, id), id)
+ }).asNonNullable().asNondeterministic()($"id")
+ } else {
+ struct(struct($"id", $"id").as("_1"), $"id".as("_2"))
+ }
+ val prefix = if (nested) "_1." else ""
+ val updated = source.withField(s"${prefix}first", value)
+ .withField(s"${prefix}second", value)
+ val columns = if (external) Seq(updated, value) else Seq(updated)
+ val rows = spark.range(0, 10, 1, 1).select(columns: _*).collect()
+ rows.zipWithIndex.foreach { case (row, index) =>
+ val result = if (nested) row.getStruct(0).getStruct(0) else
row.getStruct(0)
+ assert(result.getLong(2) == index + 1L)
+ assert(result.getLong(3) == index + 1L)
+ if (external) assert(row.getLong(1) == index + 1L)
+ }
+ val expectedCalls = if (nested && opaque) 10 * (if (external) 3 else
2) else 10
Review Comment:
**Nit (P3):** This expectation is the same whether or not CSE descends into
`With` bodies. `struct(...)` sources lower to a `CreateNamedStruct` with no
`With`, and opaque flat updates build no `With`. In the nested opaque case, the
tuple encoder makes `_1` nullable, so the inner `UpdateFieldsExpression` keeps
both values conditional, and they are evaluated 20/30 times either way.
Reverting `childrenToRecurse` to `case _: With => Nil` would still pass every
case here. Only the map-level `SubexpressionEliminationSuite` checks would
notice.
As a result, no test executes generated or interpreted code that hoists a
value out of a `With` body, which is the new behavior. A case with a
non-`struct(...)` input whose nested parent field is non-nullable, sharing one
value between two `withField` calls and another column, would give that path a
failure signal.
**Verification:**
- **Regression:** The new test fails when With bodies are excluded from
subexpression-elimination traversal.
- **Behavior:** The shared value is evaluated once per row and every updated
field holds the expected value on interpreted, codegen-only and whole-stage
paths.
##########
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:
Thanks, verified on the current head. fieldSources now replaces the resolved
parent at its ordinal instead of matching its name again, and the new test
checks results and schemas after both case-sensitivity transitions for
withField and dropFields, on the DataFrame and on a temp view.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4263829473","thread_id":"inline:4263829473","verdict_sha256":"3f2093f210bfe4bb52b320d8780eb19308936a0ca46433dce90c2bb41bccd99f"}
-->
##########
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:
Thanks, verified. optimizeFieldOps now collects each parent's nested
operations and optimizes the mergeable prefix once, keeping overwrite order,
the drop barrier and the last-seen casing, and the 2,000-update test bounds the
merge to linear operation visits.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4263829487","thread_id":"inline:4263829487","verdict_sha256":"3f2093f210bfe4bb52b320d8780eb19308936a0ca46433dce90c2bb41bccd99f"}
-->
--
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]
