cloud-fan commented on code in PR #58045:
URL: https://github.com/apache/spark/pull/58045#discussion_r3854922288
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodegenFallback.scala:
##########
@@ -25,13 +25,31 @@ import
org.apache.spark.sql.catalyst.expressions.codegen.Block._
*/
trait CodegenFallback extends Expression {
- protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode =
+ CodegenFallback.generate(this, ctx, ev)
+}
+
+object CodegenFallback {
+
+ /**
+ * Generates code that evaluates `e` by calling its `eval`, rather than by
generating code for it.
+ * This is what [[CodegenFallback]] gives the expressions that mix it in,
and is also used by an
+ * expression that can generate code in general but has to fall back for a
particular shape of
+ * its own tree -- see `With.doGenCode`.
+ *
+ * The two places that reason about which subtrees run interpretively both
dispatch on the trait:
Review Comment:
**Nit:**
Please add `EquivalentExpressions.childrenToRecurse` to this maintenance
list. It also dispatches on `CodegenFallback` to keep interpretive subtrees out
of CSE, so the current `two places` claim is incomplete.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala:
##########
@@ -226,15 +228,85 @@ class RewriteWithExpressionSuite extends PlanTest {
)
}
+ test("SPARK-58902: a With left in a conditional branch of an aggregate still
converges") {
+ val Seq(a, b) = testRelation.output
+ // Not cheap and referenced twice, so it stays a memoizing `With` rather
than being inlined.
+ val inBranch = With(a + b) { case Seq(ref) => ref * ref }
+ val plan = testRelation.groupBy(a)(max(Coalesce(Seq(a,
inBranch))).as("col"))
+ // The `PhysicalAggregation` arm restructures the aggregate into a
`Project` above it, and its
+ // guard is "the expressions contain a `With`", which a surviving one
keeps true on every
+ // iteration of this fixed-point batch. Without the eq check in the rule
this raises
+ // `Max iterations (5) reached for batch Rewrite With expression`, one
`Project` per iteration.
+ val rewritten = Optimizer.execute(plan)
+ // Idempotent: running the batch again changes nothing.
+ comparePlans(Optimizer.execute(rewritten), rewritten)
+ assert(rewritten.collect { case p: Project => p }.size <= 1,
+ s"the rule stacked a Project per iteration:\n$rewritten")
+ }
+
+ test("SPARK-58902: a cheap or single-reference definition in a branch is
still inlined") {
+ val Seq(a, b) = testRelation.output
+ // A bare attribute is cheap, so inlining it costs nothing and keeps the
branch foldable.
Review Comment:
**Nit:**
`Attribute.foldable` is false, so inlining `a` does not keep this branch
foldable.
```suggestion
// A bare attribute is cheap, so inlining it costs nothing.
```
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/WithExpressionEvalSuite.scala:
##########
@@ -0,0 +1,194 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.catalyst.expressions
+
+import org.apache.spark.SparkFunSuite
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext,
CodegenFallback, ExprCode, GenerateMutableProjection}
+import org.apache.spark.sql.types.{DataType, IntegerType}
+
+/**
+ * Evaluation of [[With]] and the memoization it gives a
[[CommonExpressionRef]]. The rewrite that
+ * decides which `With`s reach evaluation at all is covered by
`RewriteWithExpressionSuite`.
+ */
+class WithExpressionEvalSuite extends SparkFunSuite {
+
+ /**
+ * A stand-in for a stateful generator: every evaluation returns the next
integer, so a second
+ * evaluation within one row is directly observable. Only used on the
interpreted path.
+ */
+ private case class Counter() extends LeafExpression with Nondeterministic {
+ @transient private var n = 0
+ override def stateful: Boolean = true
+ override def dataType: DataType = IntegerType
+ override def nullable: Boolean = false
+ override protected def initializeInternal(partitionIndex: Int): Unit = {}
+ override protected def evalInternal(input: InternalRow): Any = { n += 1; n
}
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode):
ExprCode =
+ throw new UnsupportedOperationException
+ }
+
+ private def counter(): Counter = {
+ val c = Counter()
+ c.initialize(0)
+ c
+ }
+
+ /**
+ * A node that has to be evaluated interpretively even when its parent is
generated, so that a
+ * reference below it takes the interpreted path out of generated code.
+ */
+ private case class Fallback(child: Expression) extends UnaryExpression with
CodegenFallback {
+ override def dataType: DataType = child.dataType
+ override def eval(input: InternalRow): Any = child.eval(input)
+ override protected def withNewChildInternal(newChild: Expression):
Fallback =
+ copy(child = newChild)
+ }
+
+ test("a definition is evaluated once per row however many references read
it") {
+ // `ref + ref` is the shape `BETWEEN` produces. Memoized, both references
read one value, so the
+ // sum is 2n on the nth row; inlined it would be n + (n + 1).
+ val w = With(counter()) { case Seq(ref) => Add(ref, ref) }
+ assert((1 to 4).map(_ => w.eval(InternalRow.empty)) == Seq(2, 4, 6, 8))
+ }
+
+ test("a definition is not evaluated on a row that reaches no reference") {
+ val c = counter()
+ // The branch is inside the `With`, so the `With` is entered on every row
and does clear its
+ // cells. What it must not do is evaluate the definition before a
reference is reached; putting
+ // the branch outside would only test that `If` does not evaluate the arm
it did not take.
+ val w = With(c) { case Seq(ref) => If(Literal.TrueLiteral, Literal(-1),
Add(ref, ref)) }
+ assert((1 to 3).map(_ => w.eval(InternalRow.empty)) == Seq(-1, -1, -1))
+ // The counter is still at 0, so the first row that does reach a reference
sees 1.
+ assert(With(c) { case Seq(ref) => Add(ref, ref) }.eval(InternalRow.empty)
== 2)
+ }
+
+ test("a reference behind a short-circuiting operator is not read when the
left side is false") {
+ // Neither pre-evaluating the definition into a project nor guarding that
column by the branch
+ // can express this: both evaluate on every row the branch is reached on,
while `And` short
+ // circuits before the reference. The `And` is inside the `With` so that
the `With` is entered
Review Comment:
**Nit:**
`short-circuit` is a compound verb here.
```suggestion
// can express this: both evaluate on every row the branch is reached
on, while `And`
// short-circuits before the reference. The `And` stays inside the
`With` so it is entered
```
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -38,11 +73,121 @@ case class With(child: Expression, defs:
Seq[CommonExpressionDef])
override def dataType: DataType = child.dataType
override def nullable: Boolean = child.nullable
override def children: Seq[Expression] = child +: defs
+
+ /**
+ * The references in `child` that name one of these definitions, paired with
the definition each
+ * names. The list is found once, since the tree does not change between
evaluations.
+ *
+ * Only `child` is scanned, which relies on a reference to one of these
definitions never living
+ * inside another one of them: `With.apply` cannot build that, and
`RewriteWithExpression`, which
+ * does rewrite inside a `With`, only ever replaces a reference with its
definition's child or
+ * with an attribute -- it never puts a reference inside a definition.
Scanning `children` instead
+ * would look safer and be worse -- it would also bind a reference sitting
inside the definition
+ * it names, and since `CommonExpressionCell.get` sets `computed` only after
the nested evaluation
+ * returns, that turns today's loud "outside its With" error into a
StackOverflowError. A nested
+ * `With` is not affected either way: `children` is `child +: defs`, so this
scan already descends
+ * into an inner `With`'s own definitions.
+ */
+ @transient private lazy val refsToBind: Seq[(CommonExpressionRef,
CommonExpressionDef)] = {
+ val idToDef = defs.map(d => d.id -> d).toMap
+ child.collect { case r: CommonExpressionRef if idToDef.contains(r.id) =>
(r, idToDef(r.id)) }
+ }
+
+ /**
+ * Binds this `With`'s references to its own cells, clears them, and
evaluates the child. A
+ * reference reached by that evaluation computes its definition once and
every later reference
+ * reads the value back, so a definition is evaluated where the child would
have evaluated it,
+ * once, rather than once per reference. See [[CommonExpressionCell]].
+ *
+ * The binding is redone on every evaluation rather than once, because a
reference can be reached
+ * from two `With`s. `withNewChildrenInternal` cannot hand the new `With`
its own references: a
+ * rebuilt reference compares equal to the one it replaces, since the
binding it carries is not
+ * part of its equality, so `transform` keeps the original. Binding once
would then leave the
+ * `With` that bound last deciding what both of them read. Rebinding costs
one pass over the
+ * references, two for a `BETWEEN`, and makes the `With` currently
evaluating always the owner.
+ */
+ override def eval(input: InternalRow): Any = {
+ refsToBind.foreach { case (ref, exprDef) => ref.bindTo(exprDef) }
+ defs.foreach(_.cell.clear())
+ child.eval(input)
+ }
+
+ // The cells are cleared on entry, so this holds state for the duration of
one evaluation.
+ override def stateful: Boolean = true
+
+ /**
+ * Whether one of this `With`'s references sits somewhere that will be
evaluated interpretively
+ * even though this `With` is generated. Two shapes do that: a
[[CodegenFallback]], which is
+ * evaluated by calling `eval` on it from the generated code, and a nested
`With` that itself
+ * takes the fallback below -- `With` does not mix in `CodegenFallback`, so
it has to be named
+ * here rather than matched as one. A reference reached that way needs its
cell bound and
+ * cleared, which the generated code does not do: it clears the codegen
flags.
+ *
+ * This is the same shape `EquivalentExpressions.childrenToRecurse` already
refuses to look past,
+ * for the same reason.
+ *
+ * Each level memoizes, but `holdsMyRef` runs again at every nested `With`
the scan passes, so a
+ * chain of them nested in each other's `child` costs on the order of the
square of the depth.
+ * `nullif(a, nullif(b, c))` does produce such a chain -- only the memoized
input becomes a
+ * definition, the rest stays in `child` -- but these chains are shallow in
practice. Reading a
+ * nested `With`'s own `lazy val` from here also takes its monitor while
holding this one; the
+ * edges only ever run from an ancestor to a proper descendant of an
immutable tree, so the order
+ * is a strict partial one and cannot deadlock. `canonicalizationIdMap`
below relies on the same.
+ */
+ @transient private lazy val refUnderCodegenFallback: Boolean = {
+ val ids = defs.map(_.id).toSet
+ def holdsMyRef(e: Expression): Boolean = e.exists {
+ case r: CommonExpressionRef => ids.contains(r.id)
+ case _ => false
+ }
+ child.exists {
+ case f: CodegenFallback => holdsMyRef(f)
+ case w: With if w.refUnderCodegenFallback => holdsMyRef(w)
+ case _ => false
+ }
+ }
+
+ /**
+ * Clears each definition's flag, then generates the child. The flags are
cleared in the same
+ * block the child is generated into, so a reference cannot run against a
flag left set by an
+ * earlier row: on a row that does not reach the branch holding this `With`,
neither the clearing
+ * nor any reference runs.
+ *
+ * When a reference sits under a [[CodegenFallback]], or inside a nested
`With` that itself falls
+ * back, the whole `With` is evaluated interpretively instead. Generating
the child would leave
+ * that reference reading a cell nobody bound and nobody clears, and
generating part of it is
+ * worse still: a definition reached from both sides would be computed once
through the flags and
+ * once through the cell, holding two values for one row. [[eval]] binds and
clears both, so
+ * handing it the whole subtree keeps one mechanism in play. `ctx.INPUT_ROW`
is available on that
+ * path because `CollapseCodegenStages.supportCodegen` turns whole-stage
codegen off for a plan
+ * whose expressions hold the offending `CodegenFallback` -- it is visible
there, since a `With`
+ * in a conditional branch reaches execution inside `plan.expressions` like
any other expression.
+ */
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode):
ExprCode = {
+ if (refUnderCodegenFallback) {
+ return CodegenFallback.generate(this, ctx, ev)
+ }
+ ctx.withCommonExprs(defs) { slots =>
+ val clearFlags = slots.map(s => s"${s.computed} = false;").mkString("\n")
+ val childGen = child.genCode(ctx)
+ ev.copy(
+ code = code"""
+ |$clearFlags
+ |${childGen.code}
+ """.stripMargin,
+ isNull = childGen.isNull,
+ value = childGen.value)
+ }
+ }
+
override protected def withNewChildrenInternal(
newChildren: IndexedSeq[Expression]): Expression = {
val newDefs = newChildren.tail.map(_.asInstanceOf[CommonExpressionDef])
// If any `CommonExpressionDef` has been updated (data type or
nullability), also update its
- // `CommonExpressionRef` in the `child`.
+ // `CommonExpressionRef` in the `child`. This cannot be used to hand the
new `With` its own
+ // reference objects: a rebuilt reference is `==` the one it replaces,
since the binding it
+ // carries is not part of its equality, so `transform` keeps the original.
`eval` rebinds
+ // instead of relying on the references being unshared -- see
[[refsToBind]].
Review Comment:
**Nit:**
Scaladoc links are not parsed in `//` comments, so this remains literal
bracket markup.
```suggestion
// instead of relying on the references being unshared -- see
`refsToBind`.
```
--
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]
