cloud-fan commented on code in PR #58045:
URL: https://github.com/apache/spark/pull/58045#discussion_r3966357431
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -162,7 +354,85 @@ case class CommonExpressionDef(child: Expression, id:
CommonExpressionId = new C
* referenced, so that we can determine the data type and nullable of the
reference node.
*/
case class CommonExpressionRef(id: CommonExpressionId, dataType: DataType,
nullable: Boolean)
- extends LeafExpression with Unevaluable {
+ extends LeafExpression {
def this(exprDef: CommonExpressionDef) = this(exprDef.id, exprDef.dataType,
exprDef.nullable)
+
+ /**
+ * The definition this reference names, and the cell holding its value for
the current row. Both
+ * are wired by the enclosing [[With]] before it evaluates its child, and
are left out of the case
+ * class parameters so that equality and canonicalization are unchanged --
and so that a rule
+ * comparing two references does not compare their cells.
+ *
+ * `@transient` for the reason [[CommonExpressionCell]]'s own fields are: a
binding lives only for
+ * the duration of one `With.eval`, which restores it on the way out, so
serializing a task sees
+ * null. A serialization that happens while an evaluation is still on the
stack would otherwise
+ * capture a live binding, and the deserialized reference would evaluate
that definition rather
+ * than raise, dragging the definition's whole subtree along with it.
+ */
+ @transient private var definition: Expression = _
+ @transient private var cell: CommonExpressionCell = _
+
+ private[expressions] def bindTo(exprDef: CommonExpressionDef): Unit = {
+ definition = exprDef.child
+ cell = exprDef.cell
+ }
+
+ private[expressions] def boundDefinition: Expression = definition
+ private[expressions] def boundCell: CommonExpressionCell = cell
+
+ private[expressions] def bindTo(
+ newDefinition: Expression,
+ newCell: CommonExpressionCell): Unit = {
+ definition = newDefinition
+ cell = newCell
+ }
+
override val nodePatterns: Seq[TreePattern] = Seq(COMMON_EXPR_REF)
+
+ // The cell is cleared by the enclosing `With` on every entry, so this reads
mutable state.
+ override def stateful: Boolean = true
+
+ /**
+ * A copy must not carry this reference's binding: the copy belongs to a
different `With`, which
+ * wires it to its own cell. `LeafLike` returns `this` here, which would
hand two `With`s one
+ * reference object and let whichever wires last decide what both of them
read --
+ * `NamedLambdaVariable` overrides this for the same reason.
+ */
+ override def withNewChildrenInternal(
+ newChildren: IndexedSeq[Expression]): CommonExpressionRef = copy()
+
+ override def eval(input: InternalRow): Any = {
+ if (cell == null) {
+ throw SparkException.internalError(
+ s"Cannot evaluate a common expression reference outside its With:
$this")
+ }
+ cell.get(definition, input)
+ }
+
+ /**
+ * Computes the definition into the shared slots if this row has not done so
yet, then reads them.
+ * The code that computes it is emitted here rather than by the enclosing
`With`, so it runs where
+ * the first reference is reached -- behind a short-circuiting operator or a
nested conditional,
+ * if that is where the reference sits.
+ *
+ * A second reference emits the same code again, which never runs because
the flag is set. What
+ * that code is depends on the definition: a call, where the definition can
be put in a method, so
+ * that a definition that is or holds another `With` is not pasted once per
reference at every
+ * level; the body itself otherwise, whose locals are declared inside each
guard. No copy of one
+ * body encloses another -- for that, a definition would have to reference
its own id, directly or
Review Comment:
**Nit (P3):** An acyclic sibling dependency can enclose another fill body
without recursing: `withCommonExprs` registers every sibling slot before
generating the child, so `d2 = Add(ref(d1), 1)` resolves `d1` and completes
normally. Only self-reference or a dependency cycle re-enters a slot already
being filled. Could we state that distinction here and narrow the earlier
`refsToBind` claim that codegen refuses every definition-reference shape?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -207,6 +207,163 @@ class CodegenContext extends Logging {
throw QueryExecutionErrors.lambdaVariableNotDefinedError(id))
}
+ /**
+ * The slots a `CommonExpressionRef` reads: the value and its nullness, plus
the flag saying
+ * whether this row has computed them yet, and the definition to compute
them from.
+ */
+ case class CommonExprSlots(
+ value: ExprCode,
+ computed: String,
+ definition: Expression) {
+
+ // Not constructor parameters: a mutable one takes part in
`equals`/`hashCode`/`copy`, so a
+ // slot's hash would change as it is filled and a `copy` would carry
another scope's code.
+ private var fillCode: Option[Block] = None
+ private var filling: Boolean = false
+
+ /**
+ * The code that computes the definition into the slots and sets
`computed`, which a reference
+ * emits behind that flag. Cached, so every reference shares whatever
mutable state the
+ * definition allocated, such as an RNG.
+ *
+ * The cache lives on the slot, so it lasts exactly as long as the scope:
the code names the
+ * `INPUT_ROW` and `currentVars` in effect when it was generated.
`GenerateOrdering` generates
+ * its key once per comparison side under a different row variable, so a
slot shared between the
+ * sides would read the wrong row, or not compile where
`Expression.reduceCodeSize` has hoisted
+ * the reference into a method taking one row.
+ *
+ * `filling` catches a definition that references its own id, which would
otherwise re-enter and
+ * recurse, since `fillCode` is set only after `definition.genCode`
returns.
+ *
+ * The body goes into a method where it can and is worth it -- a
definition that is or holds
+ * another `With`, or a body past the split threshold -- so it is emitted
once per scope rather
+ * than once per reference, which for nested `With`s would double per
level. A method is only
+ * possible where the definition reads the input row rather than local
variables, the condition
+ * `reduceCodeSize` splits under. That is not the same as whole-stage
codegen being off: a
+ * whole-stage `Project` or `Filter` passes local variables, while
+ * `SortMergeJoinExec.createJoinKey` and the aggregate output paths
generate against a row.
+ */
+ def fill: Block = {
+ if (fillCode.isEmpty) {
+ if (filling) {
+ throw SparkException.internalError(
+ "Cannot generate a common expression whose definition references
it: " +
+ definition.toString)
+ }
+ filling = true
+ try {
+ fillCode = Some(build)
+ } finally {
+ filling = false
+ }
+ }
+ fillCode.get
+ }
+
+ private def build: Block = {
+ val defGen = definition.genCode(CodegenContext.this)
+ // Whether the isNull slot exists is decided by the definition, so it is
read off the slot
+ // rather than off a reference's own `nullable`: taking it from both
would let the two
+ // disagree, and either emit `false = <isNull>;`, which does not
compile, or leave the slot
+ // holding the previous row's nullness.
+ val assignIsNull = if (value.isNull == FalseLiteral) {
+ ""
+ } else {
+ s"${value.isNull} = ${defGen.isNull};"
+ }
+ val body = code"""
+ |${defGen.code}
+ |$assignIsNull
+ |${value.value} = ${defGen.value};
+ |$computed = true;
+ """.stripMargin
+ // TODO(SPARK-59295): cover the local-variable case too, by passing the
`currentVars` values a
+ // definition reads into the method as parameters, the way
+ // `subexpressionEliminationForWholeStageCodegen` does. It needs a
decision first:
+ // `getLocalInputVariableValues` hoists an input variable that is not
evaluated yet to
+ // before the call, which for a reference behind a branch means
evaluating it on rows that
+ // never reach the reference.
+ val canPutInMethod = INPUT_ROW != null && currentVars == null
+ // A definition that is or holds another `With` is the shape whose code
doubles per level,
+ // and what this is aimed at. It is not the only one -- a definition
referencing a sibling
+ // definition of the same `With` doubles the same way, and codegen
accepts that, since the
+ // sibling's slots are in scope while this definition is generated
(`With.refsToBind` says
+ // why nothing builds that tree, and that evaluating one raises). What
bounds those is not
+ // the length arm below: `body` is assembled after `definition.genCode`
already ran
+ // `reduceCodeSize`, so the arm fires only in the band just under the
threshold. It is
+ // `reduceCodeSize` itself, which hoists whichever node's code first
passes the threshold as
+ // generation walks up, capping what one level contributes, so the code
stays linear in the
+ // depth either way. The length arm just keeps the same body from being
split once per
+ // reference, which leaves the methods small and the code as large.
+ val worthAMethod = definition.exists(_.isInstanceOf[With]) ||
Review Comment:
**Nit (P3):** `With` already sets `WITH_EXPRESSION` in `nodePatterns`, and
those bits include descendants, so this recursive `exists` walk repeats work
for every definition. Could we import the pattern and use
`definition.containsPattern(WITH_EXPRESSION)` here? It preserves the same
root-or-descendant test without revisiting the expression subtree.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -38,11 +70,154 @@ 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, found once since the tree does not change between evaluations.
Only `child` is scanned,
+ * which is why the checks below refuse a definition holding a reference of
this scope.
+ */
+ @transient private lazy val refsToBind: IndexedSeq[(CommonExpressionRef,
CommonExpressionDef)] = {
+ // Three shapes whose bindings cannot be kept straight, each of which this
path used to answer
Review Comment:
**Nit (P3):** `this path used to answer for with` is not grammatical, so the
outcome of these three shapes is hard to follow. Could this say that the path
previously evaluated them incorrectly or overflowed the stack?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/WithExpressionEvalSuite.scala:
##########
@@ -0,0 +1,372 @@
+/*
+ * 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.{SparkException, 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.catalyst.plans.SQLHelper
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{BooleanType, 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 with SQLHelper {
+
+ /**
+ * 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)) }
+ // -1 whatever the definition does, since the taken arm is a literal: the
assertion below is the
+ // load-bearing one, and this only says the evaluation ran.
+ 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. Eager filling
+ // would have consumed three values and this would read 8.
+ 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 and only the reference is skipped.
+ val c = counter()
+ val w = With(c) { case Seq(ref) =>
+ And(Literal.FalseLiteral, GreaterThanOrEqual(ref, Literal(1)))
+ }
+ // False whatever the definition does, as in the case above: the counter
read is what bites.
+ // This is not that case repeated -- there the arm is skipped by `If`,
here by `And`'s short
+ // circuit, and the two are separate paths in both evaluation and codegen.
+ assert((1 to 3).map(_ => w.eval(InternalRow.empty)) == Seq(false, false,
false))
+ assert(With(c) { case Seq(ref) => Add(ref, ref) }.eval(InternalRow.empty)
== 2)
+ }
+
+ test("a nested With memoizes each scope separately") {
+ val outer = counter()
+ val inner = counter()
+ // The outer reference is read twice, and one of its uses wraps an inner
`With` whose own
+ // reference is also read twice, so each row is `o + (2i + o)` with `o`
and `i` both n.
+ val w = With(outer) { case Seq(o) =>
+ Add(o, With(inner) { case Seq(i) => Add(Add(i, i), o) })
+ }
+ assert((1 to 3).map(_ => w.eval(InternalRow.empty)) == Seq(4, 8, 12))
+ }
+
+ test("a With rebuilt by a rule still memoizes") {
+ val w = With(counter()) { case Seq(ref) => Add(ref, ref) }
+ // Evaluate the original first. A rule rebuilds the definition, and
therefore its cell, while
+ // handing the new `With` the original's reference objects, so an
implementation that bound the
+ // references once would leave the rebuilt `With` reading this
evaluation's cell.
+ assert(w.eval(InternalRow.empty) == 2)
+ // The rule has to hand back a node that is not `==` the one it replaces,
or `transformUp` keeps
+ // the original -- `Counter()` is a case class with no parameters, so a
fresh one compares equal
+ // to the old one. Wrapping it changes the tree while leaving the
definition's value sequence,
+ // its data type and its nullability alone, so the references are carried
over rather than
+ // rebuilt, which is the shape that needs the per-evaluation rebinding.
+ val rewritten = w.transformUp { case _: Counter => Add(counter(),
Literal(0)) }
+ .asInstanceOf[With]
+ assert(rewritten ne w, "the rule has to have rebuilt something")
+ assert((1 to 3).map(_ => rewritten.eval(InternalRow.empty)) == Seq(2, 4,
6))
+ }
+
+ test("SPARK-58902: a copy of a With owns its references and its cells") {
+ // A deliberately non-stateful definition, which is what makes
`CommonExpressionDef.stateful`
+ // load-bearing: `mapChildren` finds nothing changed below it, so without
the override the copy
+ // gets the original definition object, cell included. (A `Rand`
definition would hide that.)
+ // Likewise `CommonExpressionRef.stateful`, since
`LeafLike.withNewChildrenInternal` returns
+ // `this` and the two `With`s would otherwise share references while
owning two cells.
+ val w1 = With(Literal(1)) { case Seq(ref) => Add(ref, ref) }
+ val w2 = w1.freshCopyIfContainsStatefulExpression().asInstanceOf[With]
+ def refOf(e: Expression): CommonExpressionRef =
+ e.collect { case r: CommonExpressionRef => r }.head
+ assert(w1 ne w2, "a stateful With has to be copied")
+ assert(refOf(w1.child) ne refOf(w2.child), "the copy has to own its
references")
+ assert(w1.defs.head.cell ne w2.defs.head.cell, "the copy has to own its
cell")
+ }
+
+ test("SPARK-58902: two Withs over one set of references each read their own
definition") {
+ // The shape a rule produces when it rewrites only the definition: the new
`With` is built over
+ // the original's references, since a rebuilt reference compares equal to
the one it replaces.
+ // Visibly different definitions show which cell the references actually
read.
+ val w1 = With(counter()) { case Seq(ref) => Add(ref, ref) }
+ val w2 = w1.withNewChildren(
+ IndexedSeq(w1.child, CommonExpressionDef(Literal(100),
w1.defs.head.id))).asInstanceOf[With]
+
+ assert(w1.eval(InternalRow.empty) == 2, "w1 reads its own counter")
+ assert(w2.eval(InternalRow.empty) == 200, "w2 reads its own literal")
+ assert(w1.eval(InternalRow.empty) == 4, "w1 read the value w2 memoized")
+ }
+
+ test("SPARK-58902: a definition's code is generated once however many
references read it") {
+ // Short and holding no `With`, so the body is emitted inline at each
reference. It has to be
+ // the same text, generated once: otherwise each copy calls
`addMutableState` again and the
+ // definition owns one counter per reference, which for references in
mutually exclusive
+ // positions hands out one value on two rows.
+ val ctx = new CodegenContext
+ With(MonotonicallyIncreasingID()) { case Seq(ref) => Add(ref, ref)
}.genCode(ctx)
+ val counters = ctx.inlinedMutableStates.count { case (_, name) =>
name.startsWith("count") }
+ assert(counters == 1, s"the definition was generated $counters times")
+ }
+
+ test("SPARK-58902: nesting does not multiply a definition's body when it can
go in a method") {
+ // Each level reads its definition twice, so pasting the body doubles it
per level;
+ // `CommonExprSlots.fill` putting it in a method makes it 2 at any depth.
Not covered here: the
+ // shape where the input arrives as local variables, as a whole-stage
`Project` or `Filter`
+ // passes it, where no method is possible and the count is 2, 4, 8, ...
256 at depths 1 to 8 --
+ // see SPARK-59295. A bare `CodegenContext` has `INPUT_ROW` set and
`currentVars` null.
+ val marker = 1234567
+ def nested(depth: Int): Expression = {
+ val leaf: Expression = Add(BoundReference(0, IntegerType, nullable =
false), Literal(marker))
+ (1 to depth).foldLeft(leaf) { (inner, _) =>
+ With(inner) { case Seq(ref) => Add(ref, ref) }
+ }
+ }
+ def markerCount(depth: Int): Int = {
+ val ctx = new CodegenContext
+ val source = nested(depth).genCode(ctx).code.toString +
ctx.declareAddedFunctions()
+ marker.toString.r.findAllMatchIn(source).size
+ }
+ // Pin the threshold: below the innermost body's length the fill would go
into a method of its
+ // own and the count would be 1, for a reason unrelated to nesting.
+ withSQLConf(SQLConf.CODEGEN_METHOD_SPLIT_THRESHOLD.key -> "1024") {
+ (1 to 6).foreach { depth =>
+ val count = markerCount(depth)
+ assert(count == 2, s"the innermost body was emitted $count times at
depth $depth")
+ }
+ // The 2 is the innermost body pasted at its two references, at every
depth; what stops the
+ // doubling from depth 2 on is that each enclosing definition is itself
a `With` and so goes
+ // into a method, whose call sites carry no marker.
+ //
+ // The values double per level, which a deterministic leaf does whether
a reference reads a
+ // slot or recomputes -- an arithmetic and compile check, not a second
reading of the count.
+ // Four levels keeps the product inside Int, which ANSI `Add` would
raise on past depth 10.
+ val proj = GenerateMutableProjection.generate(Seq(nested(4)))
+ assert(proj(InternalRow(1)).getInt(0) == (1 + marker) * 16)
+ }
+ }
+
+ test("SPARK-58902: a nested With sharing a reference object restores the
outer binding") {
+ // Two `With`s over one reference object, the inner one nested in the
outer one's child and
+ // redefining the id the outer one defines. Only a caller building the
case class directly
+ // produces this -- `withNewChildrenInternal` shares references between a
`With` and its
+ // replacement, which are siblings, and no rule nests a redefinition of a
live id -- so this is
+ // hardening rather than a reachable wrong answer. Binding on entry alone
is not enough: the
+ // inner `With` rebinds the shared reference to its own definition, so a
read after the inner
+ // one returned would answer with the inner definition unless the outer
binding is put back.
+ val id = new CommonExpressionId()
+ val outerDef = CommonExpressionDef(counter(), id)
+ val innerDef = CommonExpressionDef(Literal(10), id)
+ val ref = new CommonExpressionRef(outerDef)
+ val inner = With(Add(ref, ref), Seq(innerDef))
+ val outer = With(Add(Add(ref, inner), ref), Seq(outerDef))
+ // The outer definition counts, the inner one does not, so row n is n +
(10 + 10) + n. A counter
+ // is what makes the repetition worth something: leaving the inner binding
in place gives
+ // n + 20 + 10, and losing memoization of the outer definition gives n +
20 + (n + 1).
+ assert((1 to 3).map(_ => outer.eval(InternalRow.empty)) == Seq(22, 24, 26))
+
+ // The generated path never has to answer this: one id in two nested
scopes is refused while
+ // generating, so it cannot quietly disagree with the values above.
+ val literalOuter = CommonExpressionDef(Literal(1), new
CommonExpressionId())
+ val literalRef = new CommonExpressionRef(literalOuter)
+ val nested = With(
+ Add(literalRef, With(Add(literalRef, literalRef),
+ Seq(CommonExpressionDef(Literal(10), literalOuter.id)))),
+ Seq(literalOuter))
+ val generated =
intercept[SparkException](GenerateMutableProjection.generate(Seq(nested)))
+ assert(generated.getMessage.contains("is already being generated"))
+ }
+
+ test("SPARK-58902: a definition that references its own id fails without
recursing") {
+ // Only a caller building the case class directly can produce this, and it
is caught rather than
+ // left to recurse: generating the definition re-enters the same slot, and
evaluating it would
Review Comment:
**Nit (P3):** The generated half does re-enter the same slot, but the
interpreted half of this first construction uses two distinct reference
objects. Without the new validation, `refsToBind` would bind only the child
reference and the reference inside `selfDef` would fail as unbound; only the
shared-reference variant below re-enters one cell. Could we split those two
failure mechanisms in this comment?
--
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]
