HyukjinKwon commented on code in PR #59283:
URL: https://github.com/apache/spark/pull/59283#discussion_r4277519373
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/EquivalentExpressions.scala:
##########
@@ -173,6 +172,29 @@ class EquivalentExpressions(
case other => other.children
}
+ private def alwaysEvaluatedDefinitions(withExpr: With): Seq[Expression] = {
+ val definitions = withExpr.defs.map(definition => definition.id ->
definition).toMap
+ val required = mutable.HashSet.empty[CommonExpressionId]
+ def collect(expression: Expression, shadowed: Set[CommonExpressionId]):
Unit = {
+ val alwaysEvaluated = skipForShortcut(expression)
+ if (alwaysEvaluated.containsPattern(COMMON_EXPR_REF)) {
+ alwaysEvaluated match {
+ case ref: CommonExpressionRef if !shadowed.contains(ref.id) =>
+ definitions.get(ref.id).foreach { definition =>
+ if (required.add(ref.id)) collect(definition.child, Set.empty)
+ }
+ case nested: With =>
+ val innerIds = shadowed ++ nested.defs.map(_.id)
+ childrenToRecurse(nested).foreach(collect(_, innerIds))
Review Comment:
**Non-blocking (P2):** `childrenToRecurse(nested)` returns `nested.child +:
alwaysEvaluatedDefinitions(nested)`, and computing those definitions already
walks `nested.child` once. The `foreach(collect(_, innerIds))` here then walks
`nested.child` again. Nothing is memoized, so each nested `With` doubles the
work below it. `updateExprTree` also calls `childrenToRecurse` at every `With`
it visits, which repeats the walk for each inner level.
The new lowering produces exactly these chains. For an input that is not
`struct(...)`, each level of a dotted path becomes
`With(GetStructField(outerRef, i)) { UpdateFieldsExpression(innerRef, ...) }`
inside the parent level's updated values. When the path is non-nullable,
`UpdateFieldsExpression` lists its updated values as always evaluated, so
`collect` reaches every level. Take a schema whose nested struct fields are
declared non-nullable many levels deep, from an explicit schema or NOT NULL
source columns, and call `col("s").withField("l1.l2.....lN.x", lit(1))` or
`dropFields` with that path. Each CSE pass over the projection then does work
that roughly doubles with every extra level. At about 20 levels that is on the
order of a million walks of the remaining body, on the driver for whole-stage
codegen and wherever projections are built. Results stay correct, but planning
and projection setup become effectively unbounded. The doubling reaches only as
deep as the path stays no
n-nullable, because a nullable level's updated values are conditional. Flat
updates and `struct(...)` inputs build no nested `With`. The merge target
stopped at every `With`, and its lowering had none, so its CSE walk was linear.
The PR's linearity tests count plan nodes and merge visits, not CSE work, so
none of them catches this.
Handling the nested `With` in one pass would fix it. Stop if
`refUnderCodegenFallback` is set. Otherwise walk `nested.child` once with its
definition ids in scope, and the first time an always-evaluated reference
reaches one of those definitions, walk that definition's child in the enclosing
scope so its outer references still mark outer definitions as required.
Memoizing `alwaysEvaluatedDefinitions` per `With` instance within one traversal
would also work.
**Verification:**
- **Regression:** Run CSE over a nested With chain modeled on nested-path
withField lowering over a non-nullable struct, deep enough (for example 30
levels) that per-level doubling would be infeasible. It must complete with
bounded work that grows with depth at most polynomially, and still report the
expected common subexpressions.
- **Compatibility:** The existing With subexpression-elimination cases
(nested scopes, guards, short-circuit peel, fallback, higher-order lambdas,
shadowing, lexical binding, and the SPARK-58902 unbound-reference check) yield
the same common subexpressions and the same rejected candidates.
- **Behavior:** Nested-path withField/dropFields queries still evaluate a
shared struct input and a shared always-evaluated value once per row on the
interpreted, codegen-only and whole-stage paths, with correct results.
##########
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:
Thanks, verified on the current head. When CSE reaches a With, it now also
walks the definitions the body always evaluates, so the struct input of a
nested-path update can be shared with other uses in the same projection.
supportedExpression still rejects candidates that hold unbound references, and
guarded, unused, CodegenFallback, lambda and shadowed definitions are not
hoisted. The new withField/dropFields test checks that a shared struct input is
evaluated once per row on the interpreted, codegen-only and whole-stage paths.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4274818970","thread_id":"inline:4274818970","verdict_sha256":"da4e342b83744e2a784cb95fd720aa39782cab886b1272d5195846b31f427d38"}
-->
##########
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:
Thanks, verified. The new test uses a non-nullable nested parent, so the
inner update evaluates both values unconditionally. It shares one value across
two nested updates and a separate column, and checks the results and
once-per-row evaluation on all three evaluation paths. Without the With-body
traversal, the shared value would be computed three times per row, so the test
now catches that regression.
<!-- SPARK_DEV_REVIEW_REPLY
{"feedback_id":"inline:4274818989","thread_id":"inline:4274818989","verdict_sha256":"da4e342b83744e2a784cb95fd720aa39782cab886b1272d5195846b31f427d38"}
-->
--
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]
