cloud-fan commented on code in PR #58045:
URL: https://github.com/apache/spark/pull/58045#discussion_r3910503167
##########
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)
Review Comment:
**Blocking (P1):** `commonExprIdSet` describes the expression before this
bottom-up rewrite, but rewriting an inner `With` can duplicate a reference
owned by this outer `With`. With
`spark.sql.optimizer.avoidCollapseUDFWithExpensiveExpr=false`, an outer
nondeterministic Python UDF used once inside a cheap deterministic inner Python
UDF becomes two outer references when that inner definition is substituted at
two sites; this stale single-use decision then substitutes the nondeterministic
definition twice, so the consumers can observe different values. Please base
this classification on reference multiplicity in the rewritten `With` and add
the nested regression.
**Recommended change:** Recompute reference multiplicity for each rewritten
With immediately before classifying its definitions, and add a focused
nested-UDF regression.
**Why this works:** Count CommonExpressionRef occurrences from the
post-transform child and definitions so references introduced by inner
substitution prevent an outer nondeterministic definition from being inlined.
**Scope:** RewriteWithExpression and its focused Catalyst optimizer suite.
**Compatibility:** No public API or data-format change; this only corrects
optimizer behavior for nested With expressions.
**Risks:** The recount must respect With binding boundaries while including
references introduced into the rewritten node. More definitions may remain
memoized in nested shapes where the original count was stale.
**Constraints:** Preserve the existing cheap deterministic and true
single-use substitutions when their post-rewrite counts permit them. Preserve
branch-local laziness and optimizer idempotence.
**Success:** The nested configuration invokes the outer nondeterministic
definition once per reached With entry, and the optimizer's existing
substitution and convergence tests continue to pass.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -207,6 +207,150 @@ 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: every one of those takes part in a case
class's `equals`,
+ // `hashCode`, `copy` and `toString`, so a mutable one would make a slot's
hash change as it is
+ // filled, and would hand a `copy` the code generated for another scope.
+ 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. Generated once and cached, so that every
reference shares whatever
+ * mutable state the definition allocated, such as an RNG, rather than
getting its own.
+ *
+ * A definition that references its own id would re-enter this while the
first call is still
+ * generating it, since `fillCode` is set only after `definition.genCode`
returns. `filling`
+ * turns that into the error `With.refsToBind` documents for the
interpreted path, rather than a
+ * StackOverflowError from generating the same definition inside itself
forever.
+ *
+ * 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) {
+ 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
+ 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]) ||
+ body.length > SQLConf.get.methodSplitThreshold
+ if (canPutInMethod && worthAMethod) {
+ val funcName = freshName("computeCommonExpr")
+ val funcFullName = addNewFunction(funcName,
+ s"""
+ |private void $funcName(InternalRow $INPUT_ROW) {
+ | $body
+ |}
+ """.stripMargin)
+ code"$funcFullName($INPUT_ROW);"
+ } else {
+ body
+ }
+ }
+ }
+
+ /**
+ * Holding a map of the common expressions of the `With` expressions
currently being generated,
+ * the same way [[currentLambdaVars]] holds the variables of the enclosing
lambdas.
+ */
+ var currentCommonExprs: mutable.Map[Long, CommonExprSlots] =
mutable.HashMap.empty
+
+ /**
+ * Allocates a value slot and a `computed` flag per definition, generates
`f` with them in scope,
+ * then takes them out of scope again. A reference generated inside `f`
reads the slots back by
+ * id and fills them the first time it is reached on a row -- the enclosing
`With` only clears the
+ * flags.
+ */
+ def withCommonExprs(defs: Seq[CommonExpressionDef])(f: Seq[CommonExprSlots]
=> ExprCode)
+ : ExprCode = {
+ val slots = defs.map { d =>
Review Comment:
**Non-blocking (P2):** Please add interpreted and generated execution
coverage for one `With` containing two surviving, distinguishable definitions,
each referenced twice. Every current runtime case has one definition, and the
optimizer's two-definition case inlines one before evaluation, so a slot-index
or cleanup bug that aliases these per-definition states would not be observed.
Exact values and once-per-definition advancement across rows would cover both
allocation and clearing.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/With.scala:
##########
@@ -162,7 +373,79 @@ 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
+ }
+
+ 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
+ * through a sibling, which recurses in `fill` before any Java exists -- so
repeating it declares
+ * nothing twice in one scope. See `CommonExprSlots.fill`, which also says
what bounds the code
+ * when a method is not possible.
+ */
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode):
ExprCode = {
+ val slots = ctx.getCommonExpr(id.id)
+ ev.copy(
+ code = code"""
+ |if (!${slots.computed}) {
Review Comment:
**Non-blocking (P2):** Could the generated path get a case that enters this
`With` but takes a child branch containing no reference? The current
no-reference checks use interpreted `Counter`, while the SQL skip case avoids
entering the `With` altogether. A `GenerateMutableProjection` over a stateful
definition with inputs `false, true, false, true` should return `-1, 0, -1, 1`;
that would catch an eager fill here that consumes values on the false rows.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -181,15 +236,12 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
rewriteWithExprAndInputPlans(
_, inputPlans, commonExprsPerChild, commonExprIdSet, isNestedWith))
val newExpr = c.withNewAlwaysEvaluatedInputs(newAlwaysEvaluatedInputs)
- // Use transformUp to handle nested With.
+ // A `With` in a conditional branch cannot go into a project, which is
always evaluated
+ // while the branch may not be. It stays where it is and memoizes its
definition per row
Review Comment:
**Nit (P3):** This memoization is per `With` entry/evaluation, not per row:
`With` clears the cell or generated flag whenever it is entered, and the same
expression object can occur at two positions for one row. Please use
`per-entry` here (and in the corresponding test/PR description wording) so the
documented lifetime matches `CommonExpressionCell` and `With.eval`.
##########
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.
Review Comment:
**Nit (P3):** Discarding this parent does not guarantee unique reference
objects. If the same `With` instance is reused at two parent positions, each
transform can build a distinct copy while retaining the same unchanged
references for `toKeep`. The new bind/restore logic makes that DAG-shaped reuse
safe, so could this comment describe that mechanism instead of claiming two
`With` expressions cannot hold the same reference?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -207,6 +207,150 @@ 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: every one of those takes part in a case
class's `equals`,
+ // `hashCode`, `copy` and `toString`, so a mutable one would make a slot's
hash change as it is
+ // filled, and would hand a `copy` the code generated for another scope.
+ 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. Generated once and cached, so that every
reference shares whatever
+ * mutable state the definition allocated, such as an RNG, rather than
getting its own.
+ *
+ * A definition that references its own id would re-enter this while the
first call is still
+ * generating it, since `fillCode` is set only after `definition.genCode`
returns. `filling`
+ * turns that into the error `With.refsToBind` documents for the
interpreted path, rather than a
+ * StackOverflowError from generating the same definition inside itself
forever.
+ *
+ * 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
Review Comment:
**Nit (P3):** Whole-stage mode does not always imply `currentVars` is
populated. For example, `SortMergeJoinExec.createJoinKey` generates keys with
`INPUT_ROW` set and `currentVars` cleared, which satisfies the helper-method
predicate below. Please describe the exact `INPUT_ROW != null && currentVars ==
null` boundary and identify local-variable whole-stage `Project` paths as the
case where extraction is unavailable.
##########
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` draws a new string per row, so inlining gives the two
comparisons of the BETWEEN
+ // two draws, while memoizing gives them one and makes the answer inside
a branch the same as
+ // outside one. Its value is a string rather than a primitive, so this
also exercises a slot
+ // that cannot be an inlined field. `id` is selected alongside so that
`checkAnswer`, which
+ // compares rows as a bag, is comparing per row rather than just
counting trues.
+ val inBranch = df.selectExpr(
+ "id", "CASE WHEN id < 0 THEN false ELSE randstr(3, 0) BETWEEN 'a' AND
'zzzz' END")
+ checkAnswer(
+ inBranch, df.selectExpr("id", "randstr(3, 0) BETWEEN 'a' AND
'zzzz'").collect().toSeq)
+
+ // The same for a nondeterministic input wrapped in arithmetic.
Comparing against the
+ // branch-free form under a fixed seed is what makes this bite: inlined,
the second comparison
+ // draws again, so the row a given draw lands on shifts. The window also
has to be one the
+ // first comparison can fail: `BETWEEN 0 AND 1` over `rand` is satisfied
by every draw, which
+ // keeps the two copies in lockstep and makes the comparison hold either
way. The assertion
+ // below is what keeps that true for these rows rather than merely
intended.
+ val branchFree = df.selectExpr("id", "(rand(7) / 1.0) BETWEEN 0.3 AND
0.6").collect()
+ assert(branchFree.exists(!_.getBoolean(1)) &&
branchFree.exists(_.getBoolean(1)),
+ s"the window is not selective over these rows:
${branchFree.mkString(", ")}")
+ val inBranchRand = df.selectExpr(
+ "id", "CASE WHEN id < 0 THEN false ELSE (rand(7) / 1.0) BETWEEN 0.3
AND 0.6 END")
+ checkAnswer(inBranchRand, branchFree.toSeq)
+ }
+ }
+
+ test("SPARK-58902: a nested With agrees across the evaluation paths") {
+ // The outer `nullif`'s definition is the inner `nullif`, which is itself
a `With`, and neither
+ // one is cheap or read once, so both survive inside the branch. That
nesting is where each
+ // reference's copy of the definition code matters: the code is generated
once per definition
+ // and the text reused, so a nested `With` does not grow its code by a
factor per level.
Review Comment:
**Nit (P3):** This is true only when the generated body can be extracted
into a helper method. In a whole-stage `Project`, `currentVars` prevents
extraction and each reference pastes the cached inner Block, so the focused
source-size test still measures `2, 4, 8, ...` growth. Could this comment
distinguish generating the definition once for shared state from emitting it
once, and scope the bounded-source claim to the row-based standalone path?
--
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]
