cloud-fan commented on code in PR #58045:
URL: https://github.com/apache/spark/pull/58045#discussion_r3904659842
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -38,11 +73,124 @@ 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. The helper `With(commonExprs: _*)(replaced)`
builds the references
+ * outside the definitions and cannot produce 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. The case
class constructor does
+ * take `child` and `defs` directly, so a caller can hand a definition a
reference to any of these
+ * ids; that reference is then never bound, and evaluating it raises "Cannot
evaluate a common
+ * expression reference outside its With", which is the failure to want.
Scanning `children`
+ * instead would look safer and be worse -- it would bind such a reference,
and since
+ * `CommonExpressionCell.get` sets `computed` only after the nested
evaluation returns, that loud
+ * error would become 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) }
Review Comment:
**Blocking (P1):** Rebinding on entry does not preserve the outer scope when
a nested copied `With` shares reference objects. `CommonExpressionRef` has one
mutable definition/cell pair, so the inner `bindTo` overwrites it and this
method never restores the previous pair. With outer definition `1` and an inner
definition with the same id returning `10`, `outer(Add(Add(ref, inner(ref +
ref)), ref))` returns `31` instead of the lexically scoped `22`. Please save
and restore each prior binding in `finally` (or use a binding stack), and cover
this nested-copy case repeatedly.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -38,11 +73,124 @@ 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. The helper `With(commonExprs: _*)(replaced)`
builds the references
+ * outside the definitions and cannot produce 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. The case
class constructor does
+ * take `child` and `defs` directly, so a caller can hand a definition a
reference to any of these
+ * ids; that reference is then never bound, and evaluating it raises "Cannot
evaluate a common
+ * expression reference outside its With", which is the failure to want.
Scanning `children`
+ * instead would look safer and be worse -- it would bind such a reference,
and since
+ * `CommonExpressionCell.get` sets `computed` only after the nested
evaluation returns, that loud
+ * error would become 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)
Review Comment:
**Blocking (P1):** This fallback registers the same stateful `With` instance
for every generated occurrence. `GenerateOrdering` generates the key once for
row A and once for row B, so both emitted calls advance one stateful
definition. For example, a `With(MonotonicallyIncreasingID())` whose reference
is below a non-leaf `CodegenFallback` makes `compare(row, row)` observe
consecutive values and return nonzero. Please give each fallback occurrence
independent state (as the interpreted ordering does for its right side) and add
a `GenerateOrdering` regression asserting reflexivity.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -207,6 +207,126 @@ 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,
+ private var fillCode: Option[Block] = None) {
+
+ /**
+ * The code that computes the definition into the slots and sets
`computed`, which a reference
+ * emits behind that flag. Generated once and cached, so that every
reference shares whatever
+ * mutable state the definition allocated, such as an RNG, rather than
getting its own.
+ *
+ * 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 -- leaving a
reference's own code a single
+ * call: such a definition would otherwise have its body pasted once per
reference at every
+ * level. `Expression.reduceCodeSize` already keeps that from running
away, by hoisting
+ * whichever node's code first passes the same threshold, so what a method
here buys is the body
+ * once per scope rather than once per reference up to that threshold, and
a bound that no
+ * longer grows with depth wherever the threshold sits. It is only
possible where the definition
+ * reads the input row rather than local variables -- the condition
`reduceCodeSize` splits
+ * under, and for the same reason. Whole-stage codegen generates the
operators a `With` reaches
+ * with `currentVars` set, so neither applies there and the body is pasted
per reference.
+ */
+ def fill: Block = {
+ if (fillCode.isEmpty) {
+ val defGen = definition.genCode(CodegenContext.this)
Review Comment:
**Non-blocking (P2):** `fillCode` is still empty while `definition.genCode`
runs. If a directly constructed definition contains a `CommonExpressionRef` to
its own id, that reference resolves this same slot and recursively re-enters
`fill` until `StackOverflowError`; interpreted evaluation instead reaches the
controlled unbound-reference Spark error documented in `With.scala`. Please
publish an in-progress state or reject the cycle before recursing, and add a
generated-evaluation regression that explicitly rules out `StackOverflowError`.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -98,6 +104,55 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
}
}
+ /**
+ * Whether substituting this definition into its references is as good as
evaluating it once.
+ * True when it is referenced once, since substituting then evaluates it
once either way, or when
+ * it is cheap enough to evaluate repeatedly and deterministic, so the
repeated evaluations agree.
+ *
+ * `CollapseProject.isCheap` answers what one evaluation costs, not whether
a second one is
+ * allowed: it admits a `PythonUDF`, which may be nondeterministic.
Determinism is checked here
+ * rather than read into the cost test, so that a reader of either call site
can see which of the
+ * two questions is being asked.
+ */
+ private def canSubstitute(
+ child: Expression,
+ id: CommonExpressionId,
+ commonExprIdSet: Set[CommonExpressionId]): Boolean = {
+ !commonExprIdSet.contains(id) || (CollapseProject.isCheap(child) &&
child.deterministic)
+ }
+
+ /**
+ * `w` with every definition that gains nothing from being memoized inlined
into its references:
+ * one cheap enough to evaluate twice, and one that is referenced once
anyway. This is the test
+ * the main rewrite already applies before it hoists a definition into a
project.
+ *
+ * Inlining matters beyond the per-row bookkeeping it saves. A `With` is not
foldable, so it hides
+ * whatever it wraps from `ConstantFolding`, `PushFoldableIntoBranches`,
`SimplifyConditionals`
+ * and `ReplaceNullWithFalseInPredicate`, all of which run in later batches.
Dropping the `With`
+ * once nothing is left to memoize keeps `CASE WHEN c THEN nullif(1, 1) END`
folding as it did
+ * before this rule learned to leave one behind.
+ */
+ private def inlineDefsThatGainNothing(
+ w: With,
+ commonExprIdSet: Set[CommonExpressionId]): Expression = {
+ val (toInline, toKeep) = w.defs.partition { d =>
+ canSubstitute(d.child, d.id, commonExprIdSet)
+ }
+ if (toInline.isEmpty) {
+ w
+ } else {
+ val refToExpr = toInline.map(d => d.id -> d.child).toMap
+ val newChild =
w.child.transformWithPruning(_.containsPattern(COMMON_EXPR_REF)) {
+ // A ref of a definition kept here, or of an enclosing `With`, is left
for its owner.
+ case ref: CommonExpressionRef if refToExpr.contains(ref.id) =>
refToExpr(ref.id)
+ }
+ // `copy` rather than `withNewChildren`, which requires the child count
to be unchanged. The
+ // references of the kept definitions are carried over unbound; the new
`With` binds them, and
+ // `w` is discarded here, so no two `With`s hold the same reference.
+ if (toKeep.isEmpty) newChild else w.copy(child = newChild, defs = toKeep)
Review Comment:
**Non-blocking (P2):** Could we add a focused test for this partial branch?
The current branch-local cases each use one definition, so they exercise
all-inline or all-keep results but not `copy(child = newChild, defs = toKeep)`.
A two-definition `With` under `CaseWhen`/`Coalesce`, with one cheap
deterministic definition and one expensive multiply referenced definition,
should assert that only the expensive definition and its id remain, no
removed-id reference survives, and a second optimizer run is identical.
--
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]
