cloud-fan commented on code in PR #58045:
URL: https://github.com/apache/spark/pull/58045#discussion_r3862864751
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -162,7 +319,80 @@ 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.
+ */
+ private var definition: Expression = _
+ private var cell: CommonExpressionCell = _
+
+ private[expressions] def bindTo(exprDef: CommonExpressionDef): Unit = {
+ definition = exprDef.child
+ cell = exprDef.cell
+ }
+
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 definition's code 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 text again, which never runs
because the flag is set.
+ * The text is generated once and cached on the slots, so every copy shares
whatever mutable
+ * state the definition allocated, and a nested `With` does not grow its
code by a factor per
+ * level. The definition's locals are declared inside each guard, whose
blocks are siblings, so
+ * repeating the text declares nothing twice in one scope.
+ *
+ * Whether the isNull slot exists is decided by the definition, so it is
read off the slot rather
+ * than off this 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.
+ */
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode):
ExprCode = {
+ val slots = ctx.getCommonExpr(id.id)
+ val defGen = slots.definitionGen(ctx)
+ val assignIsNull = if (slots.value.isNull == FalseLiteral) {
+ ""
+ } else {
+ s"${slots.value.isNull} = ${defGen.isNull};"
+ }
+ ev.copy(
+ code = code"""
+ |if (!${slots.computed}) {
+ | ${defGen.code}
Review Comment:
**Blocking:**
Please emit one reusable lazy-definition block per `With` scope instead of
pasting `defGen.code` at every reference. Nested `nullif` leaves nested `With`
nodes, so each level's two references duplicate the complete inner body (`T(n)
= 2T(n-1) + O(1)`). Whole-stage codegen cannot use `Expression.reduceCodeSize`
here because it has `currentVars` and no `INPUT_ROW`; a deep nested-`nullif`
compile or source-size regression would catch this.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodegenFallback.scala:
##########
@@ -25,13 +25,32 @@ 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`.
+ *
+ * Three places reason about which subtrees run interpretively, and all of
them dispatch on the
+ * trait: `CollapseCodegenStages.supportCodegen`, which turns whole-stage
codegen off;
+ * `EquivalentExpressions.childrenToRecurse`, which keeps a subtree that
only `eval` reaches out
+ * of subexpression elimination; and `With.refUnderCodegenFallback`, which
decides whether a
+ * memoized reference is reached that way. A caller that is not a
`CodegenFallback` is invisible
+ * to all three, so it has to be named in each of them as well.
Review Comment:
**Nit:**
Please narrow this to the sites that cannot discover the caller
structurally. `With` calls `generate` without mixing in `CodegenFallback`, but
`CollapseCodegenStages` still finds the descendant fallback through its
recursive `plan.expressions.exists(_.exists(...))`; a non-trait caller
therefore does not necessarily have to be named at all three sites.
##########
sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala:
##########
@@ -430,6 +430,136 @@ class ColumnExpressionSuite extends SharedSparkSession {
checkAnswer(testData.filter($"a".between($"b", $"c")), expectAnswer)
}
+ // Runs `f` on each of the three evaluation paths, since a `With` left in a
conditional branch is
+ // evaluated by all three and the memoization is implemented separately for
interpretation and for
+ // codegen.
+ private def onEachEvalPath(f: => Unit): Unit = {
+ withSQLConf(
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
+ SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN")(f)
+ withSQLConf(
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false",
+ SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY")(f)
+ withSQLConf(
+ SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true",
+ SQLConf.CODEGEN_FACTORY_MODE.key -> "CODEGEN_ONLY")(f)
+ }
+
+ test("SPARK-58902: BETWEEN on a nondeterministic input inside a conditional
branch") {
+ onEachEvalPath {
+ // `BETWEEN` reads its input twice. Inlining the common expression into
a branch gave each
+ // read its own value, so the two comparisons saw two different ids and
6 of the 10 rows came
+ // back true where 3 is correct. A single partition makes the id
sequence 0, 1, 2, ...
+ val df = spark.range(0, 10, 1, 1)
+ checkAnswer(
+ df.selectExpr(
+ "CASE WHEN id < 0 THEN false ELSE monotonically_increasing_id()
BETWEEN 3 AND 5 END"),
+ (0 until 10).map(i => Row(i >= 3 && i <= 5)))
+ }
+ }
+
+ test("SPARK-58902: a definition in a branch is evaluated only on the rows
that reach it") {
+ onEachEvalPath {
+ // Rows 0 to 4 take the first branch, so the five rows that reach the
ELSE see ids 0 to 4 --
+ // the same values they would see if the branch were the whole
expression.
+ val df = spark.range(0, 10, 1, 1)
+ checkAnswer(
+ df.selectExpr(
+ "CASE WHEN id < 5 THEN NULL ELSE monotonically_increasing_id()
BETWEEN 1 AND 2 END"),
+ (0 until 10).map { i =>
+ if (i < 5) Row(null) else Row(i - 5 >= 1 && i - 5 <= 2)
+ })
+ }
+ }
+
+ test("SPARK-58902: a branch condition that can raise is not evaluated on
other rows") {
+ onEachEvalPath {
+ withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+ // Nothing is relocated, so a condition is evaluated only where it
always was. `a` runs
+ // -2, -1, 0, 1, 2, 3: the first two take the third branch, 0 the
first, 1 and 2 the second,
+ // and 3 falls through. `rand` is in [0, 1), so a row reaching a
BETWEEN is true whatever it
+ // draws -- the answers here do not depend on memoization. What this
rules out is the
+ // alternative that was measured and rejected: hoisting the definition
into a `Project` and
+ // guarding that column with the branch condition, which puts `6 / a`
outside conditional
+ // evaluation and raises on the `a = 0` row under ANSI.
+ val df = spark.range(0, 6, 1, 1).selectExpr("cast(id as int) - 2 as a")
+ checkAnswer(
+ df.selectExpr(
+ "CASE WHEN a = 0 THEN false " +
+ "WHEN 6 / a > 2 THEN rand(1) BETWEEN 0 AND 1 " +
+ "WHEN 6 / a < -2 THEN rand(2) BETWEEN 0 AND 1 " +
+ "ELSE false END"),
+ Seq(Row(true), Row(true), Row(false), Row(true), Row(true),
Row(false)))
+ }
+ }
+ }
+
+ test("SPARK-58902: a nondeterministic input a branch cannot pre-evaluate is
still read once") {
+ onEachEvalPath {
+ // Enough rows that the two copies an inlining implementation makes have
to fall out of step:
+ // each copy owns its own generator seeded the same way, so they only
differ once the first
+ // comparison has skipped the second copy on some row.
+ val df = spark.range(0, 20, 1, 1)
+ // `randstr(3, 0)` cannot be pre-evaluated into a project, because it
raises on a negative
Review Comment:
**Nit:**
Please replace the negative-length rationale with the nondeterministic
memoization behavior this test actually exercises. This invocation's foldable
length is literal `3`, so it cannot raise for negative length; codegen also
calls `lengthInteger()` while generating code rather than selectively by row.
##########
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
Review Comment:
**Nit:**
Please narrow this sentence to the curried `With(commonExprs: _*)(replaced)`
helper. The case class also exposes direct `apply(child, defs)` construction,
so a definition can contain a reference to its own id; because `refsToBind`
scans only `child`, that nested reference remains unbound and fails during
evaluation.
--
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]
