dongjoon-hyun commented on code in PR #59209:
URL: https://github.com/apache/spark/pull/59209#discussion_r4199845624
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -67,6 +67,16 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
// changed anything: `mapExpressions` and `withNewChildren` preserve
reference equality when
// they rewrite nothing, which is what makes this detectable.
if ((rewrittenAgg eq agg) && (rewrittenProj eq proj)) p else
rewrittenProj
+ // The children of a `CreateVariable` are the `ResolvedIdentifier`s
naming the variables, not
+ // rows to evaluate the default over, and `V2CommandStrategy` plans the
command only while
+ // they stay so. `CreateVariableExec` evaluates the default once, so a
`With` in it is left in
+ // place to memoize its definition, as one in a conditional branch is.
Review Comment:
The `With` kept here is still in the tree when the operator-optimization
batches run, and some of their rules copy their input without a determinism or
cheapness check. Each copy then evaluates the definition again in
`CreateVariableExec`: `prepareExpressions` gives every position its own copy,
and `With.eval` clears its cells on entry. So the default is not always
evaluated once.
For example, `LikeSimplification` rewrites a single `'p%s'` pattern into
`OctetLength(input) >= n AND StartsWith(input, ...) AND EndsWith(input, ...)`
(`expressions.scala` L926 and L968; only the `LikeAll`/`LikeAny` cases check
`isCheap`). With the counting UDF from the new test,
```sql
DECLARE OR REPLACE VARIABLE v5 BOOLEAN DEFAULT
nullif(CAST(plus_one_counted(21) AS STRING), '') LIKE '2%2'
```
calls the UDF 3 times, while the same expression in a `SELECT` calls it
once, because there the definition is already a column of the `Project` below
when `LikeSimplification` runs. With a nondeterministic UDF `pick()` that
returns `'ay'` or `'bz'`, `nullif(pick(), '') LIKE 'a%z'` can be true here
(`StartsWith` sees `'ay'` and `EndsWith` sees `'bz'`), which the `SELECT` form
never returns. `SimplifyExtractValueOps` (a non-foldable `map(...)[key]`) and
`UnwrapCastInBinaryComparison` (`CAST(ts AS DATE) = DATE '...'`) copy their
input the same way.
The copying itself isn't new: since SPARK-58818 it already applies to a
`With` kept in a conditional branch, and it applies to a plain nondeterministic
`LIKE` input. But this case makes the top-level form reachable, so I'd soften
"evaluates the default once ... to memoize its definition", or note that a
later rule that copies the `With` evaluates the definition once per copy.
Guarding the single-pattern case of `LikeSimplification` could be a separate
JIRA.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -67,6 +67,16 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
// changed anything: `mapExpressions` and `withNewChildren` preserve
reference equality when
// they rewrite nothing, which is what makes this detectable.
if ((rewrittenAgg eq agg) && (rewrittenProj eq proj)) p else
rewrittenProj
+ // The children of a `CreateVariable` are the `ResolvedIdentifier`s
naming the variables, not
+ // rows to evaluate the default over, and `V2CommandStrategy` plans the
command only while
+ // they stay so. `CreateVariableExec` evaluates the default once, so a
`With` in it is left in
+ // place to memoize its definition, as one in a conditional branch is.
+ case c: CreateVariable if c.defaultExpr.containsPattern(WITH_EXPRESSION)
=>
+ val default = c.defaultExpr.child
+ val rewritten =
default.transformUpWithPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case w: With => inlineDefsThatGainNothing(w)
+ }
+ if (rewritten eq default) c else c.copy(defaultExpr =
c.defaultExpr.copy(child = rewritten))
case p if p.expressions.exists(_.containsPattern(WITH_EXPRESSION)) =>
Review Comment:
(Pre-existing, not introduced by this PR) This arm also handles a
`Command`'s own expressions, and for some commands the hoisted column ends up
somewhere it is never computed. The worst case I found is `ADD CONSTRAINT ...
CHECK`:
```sql
CREATE TABLE testcat.ns.t (id INT, name STRING) USING foo;
ALTER TABLE testcat.ns.t ADD CONSTRAINT chk CHECK (length(name) BETWEEN 1
AND 10);
INSERT INTO testcat.ns.t VALUES (1, 'abc');
```
`applyInternal` moves `length(name)` into a `Project` over the child and
rewrites the `CheckConstraint` to `_common_expr_N >= 1 AND _common_expr_N <=
10`. `DataSourceV2Strategy` builds `check.toV2Constraint` from the optimized
expression (L810-L817), so the catalog stores a predicate over a column named
`_common_expr_N`. The ALTER succeeds, but every later write resolves
`predicate()` before `predicateSql` (`ResolveTableConstraints` L77) and fails
with `UNRESOLVED_COLUMN` until the constraint is dropped, while `DESCRIBE` and
`SHOW CREATE TABLE` still show the original SQL. The same CHECK in `CREATE
TABLE` works because it is converted during analysis.
This and my other comments on L70, L80 and L81 have the root cause this PR
fixes for `CreateVariable`: a `Project` under a `Command` (or `Call`) whose
expressions are never evaluated over the child's rows. `ConvertToCatalyst`
already refuses to put a `Project` under any `Command` for this reason
(`Optimizer.scala` L1079-L1080 and L1117). Doing the same here, i.e. inlining
instead of hoisting for a `Command` or `Call` and keeping the `With` only for
`CreateVariable` as this PR does, would cover all of them. Inlining is safe
there because those expressions must be deterministic (CHECK, DELETE
conditions, generated columns, defaults) or foldable (CALL arguments). Fine to
leave for a separate JIRA if you'd rather keep this PR small.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -67,6 +67,16 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
// changed anything: `mapExpressions` and `withNewChildren` preserve
reference equality when
// they rewrite nothing, which is what makes this detectable.
if ((rewrittenAgg eq agg) && (rewrittenProj eq proj)) p else
rewrittenProj
+ // The children of a `CreateVariable` are the `ResolvedIdentifier`s
naming the variables, not
+ // rows to evaluate the default over, and `V2CommandStrategy` plans the
command only while
+ // they stay so. `CreateVariableExec` evaluates the default once, so a
`With` in it is left in
+ // place to memoize its definition, as one in a conditional branch is.
+ case c: CreateVariable if c.defaultExpr.containsPattern(WITH_EXPRESSION)
=>
+ val default = c.defaultExpr.child
+ val rewritten =
default.transformUpWithPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case w: With => inlineDefsThatGainNothing(w)
+ }
+ if (rewritten eq default) c else c.copy(defaultExpr =
c.defaultExpr.copy(child = rewritten))
case p if p.expressions.exists(_.containsPattern(WITH_EXPRESSION)) =>
applyInternal(p)
Review Comment:
(Pre-existing) The same happens to a v2 `ALTER TABLE ... ADD COLUMN` or
`ALTER COLUMN ... SET DEFAULT` whose default has a non-cheap definition, with a
worse symptom.
```sql
ALTER TABLE testcat.ns.t ADD COLUMN c INT DEFAULT nullif(nullif(1, 2), 3)
```
should fold to 1. The analyzer doesn't check foldability
(`ResolveDefaultColumnsUtil` L588), and this batch runs before
`ConstantFolding`, so the outer definition is the inner `With`, which isn't
foldable and so isn't cheap. It has no references, so it goes into a `Project`
over the `ResolvedTable`, and `a.table.asInstanceOf[ResolvedTable]`
(`DataSourceV2Strategy` L823) throws a raw `java.lang.ClassCastException`;
`QueryExecution.isInternalError` only wraps NPE, `AssertionError` and
`MatchError`. `DEFAULT nullif(my_udf(1), 0)` gets the same exception instead of
the not-constant error.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -67,6 +67,16 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
// changed anything: `mapExpressions` and `withNewChildren` preserve
reference equality when
// they rewrite nothing, which is what makes this detectable.
if ((rewrittenAgg eq agg) && (rewrittenProj eq proj)) p else
rewrittenProj
+ // The children of a `CreateVariable` are the `ResolvedIdentifier`s
naming the variables, not
+ // rows to evaluate the default over, and `V2CommandStrategy` plans the
command only while
+ // they stay so. `CreateVariableExec` evaluates the default once, so a
`With` in it is left in
+ // place to memoize its definition, as one in a conditional branch is.
+ case c: CreateVariable if c.defaultExpr.containsPattern(WITH_EXPRESSION)
=>
+ val default = c.defaultExpr.child
+ val rewritten =
default.transformUpWithPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case w: With => inlineDefsThatGainNothing(w)
+ }
+ if (rewritten eq default) c else c.copy(defaultExpr =
c.defaultExpr.copy(child = rewritten))
case p if p.expressions.exists(_.containsPattern(WITH_EXPRESSION)) =>
applyInternal(p)
Review Comment:
(Pre-existing) For a group-based DELETE, `RewriteDeleteFromTable` uses the
same condition in the query, in `ReplaceData.condition` and in
`groupFilterCondition`. This arm picks `query` as the input, since it outputs
the table columns, and rewrites both `ReplaceData` conditions to read
`_common_expr_N` from a `Project` over `query`. But
`RowLevelOperationRuntimeGroupFiltering` applies `groupFilterCondition` to the
original relation (`Filter(cond, relation)`, L105), and the relation doesn't
produce that column. So with the runtime group filter on (the default), on a
table like the one in `GroupBasedDeleteFromTableSuite`,
```sql
DELETE FROM cat.ns1.test_table WHERE salary + 1 BETWEEN 300 AND 301
```
should fail while binding `_common_expr_N` in the runtime filter subquery.
Plan change validation skips plans rooted at a `Command` (`LogicalPlan.scala`
L438), so tests wouldn't catch it. The delta-based path (`WriteDelta`) has no
group filter and isn't affected.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -67,6 +67,16 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
// changed anything: `mapExpressions` and `withNewChildren` preserve
reference equality when
// they rewrite nothing, which is what makes this detectable.
if ((rewrittenAgg eq agg) && (rewrittenProj eq proj)) p else
rewrittenProj
+ // The children of a `CreateVariable` are the `ResolvedIdentifier`s
naming the variables, not
+ // rows to evaluate the default over, and `V2CommandStrategy` plans the
command only while
+ // they stay so. `CreateVariableExec` evaluates the default once, so a
`With` in it is left in
+ // place to memoize its definition, as one in a conditional branch is.
+ case c: CreateVariable if c.defaultExpr.containsPattern(WITH_EXPRESSION)
=>
+ val default = c.defaultExpr.child
+ val rewritten =
default.transformUpWithPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case w: With => inlineDefsThatGainNothing(w)
+ }
+ if (rewritten eq default) c else c.copy(defaultExpr =
c.defaultExpr.copy(child = rewritten))
Review Comment:
Since `CreateVariable.expressions` is just `Seq(defaultExpr)`, this case
could be
```scala
case c: CreateVariable =>
c.transformExpressionsUpWithPruning(_.containsPattern(WITH_EXPRESSION)) {
case w: With => inlineDefsThatGainNothing(w)
}
```
`mapExpressions` returns the same instance when nothing changes, which is
what the fixed-point batch needs (the comment on L64-L68 relies on the same
thing), and `DefaultValueExpression.withNewChildInternal` is `copy(child =
newChild)`, so `originalSQL` and `analyzedChild` are kept as they are now. The
guard is always true here anyway: the pruning on L43 only reaches a
`CreateVariable` whose default has a `With`, because its children are leaf
`ResolvedIdentifier`s. This also avoids depending on `DefaultValueExpression`'s
fields, and matches how `ReplaceExpressions` rewrites this node
(`p.mapExpressions(replace)`).
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -67,6 +67,16 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
// changed anything: `mapExpressions` and `withNewChildren` preserve
reference equality when
// they rewrite nothing, which is what makes this detectable.
if ((rewrittenAgg eq agg) && (rewrittenProj eq proj)) p else
rewrittenProj
+ // The children of a `CreateVariable` are the `ResolvedIdentifier`s
naming the variables, not
Review Comment:
(Pre-existing) `CreateTable` and `ReplaceTable` also have a
`ResolvedIdentifier` child, but that one carries the table columns as output
(SPARK-51834, `ResolveCatalogs` L116-L171), and generation expressions are
resolved against those attributes (SPARK-55347). So the generic arm on L80
picks that `ResolvedIdentifier` for a definition that reads a column:
```sql
CREATE TABLE testcat.ns.t (a INT, b INT GENERATED ALWAYS AS (nullif(a + 1,
0))) USING foo
```
gets a `Project` over the `ResolvedIdentifier`, `case c @
CreateTable(ResolvedIdentifier(catalog, ident), ...)` (`DataSourceV2Strategy`
L255) no longer matches, and it fails with the same "No plan for CreateTable"
internal error as this JIRA. `BETWEEN`, `REPLACE TABLE` and realistic forms
like `nullif(trim(name), '')` behave the same; `nullif(a, 0)` works because a
bare column is cheap. A column default with a nested NULLIF or BETWEEN, such as
`DEFAULT nullif(nullif(1, 2), 3)`, hits it as well, since a definition without
references matches any child. So the reasoning in this comment applies to these
commands too.
##########
sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala:
##########
@@ -645,6 +645,38 @@ class ColumnExpressionSuite extends SharedSparkSession {
}
}
+ test("SPARK-59954: a With in a DECLARE VARIABLE default is evaluated once") {
+ // `CreateVariableExec` evaluates the default on the driver, so the
accumulator needs no job to
+ // read. `plus_one_counted` is deterministic but not cheap, so its
definition is worth
+ // memoizing, and the first two defaults below each read it twice.
+ val calls = sparkContext.longAccumulator("calls")
+ spark.udf.register("plus_one_counted", (i: Int) => { calls.add(1); i + 1 })
+ spark.udf.register("plus_one_nondet",
+ udf((i: Int) => { calls.add(1); i + 1 }).asNondeterministic())
+ withUserDefinedFunction("plus_one_counted" -> true, "plus_one_nondet" ->
true) {
+ withSessionVariable("v1", "v2", "v3", "v4") {
+ sql("DECLARE OR REPLACE VARIABLE v1 INT DEFAULT
nullif(plus_one_counted(1), 3)")
+ assert(calls.value == 1, s"the NULLIF default called the UDF
${calls.value} times")
+ checkAnswer(sql("SELECT v1"), Row(2))
+
+ calls.reset()
+ sql("DECLARE OR REPLACE VARIABLE v2 BOOLEAN DEFAULT
plus_one_counted(1) BETWEEN 1 AND 3")
+ assert(calls.value == 1, s"the BETWEEN default called the UDF
${calls.value} times")
+ checkAnswer(sql("SELECT v2"), Row(true))
+
+ // A nondeterministic definition is evaluated once as well.
+ calls.reset()
+ sql("DECLARE OR REPLACE VARIABLE v3 INT DEFAULT
nullif(plus_one_nondet(1), 3)")
+ assert(calls.value == 1, s"the nondeterministic default called it
${calls.value} times")
+ checkAnswer(sql("SELECT v3"), Row(2))
+
+ // A seeded form of the SPARK-59954 repro; its `Rand` is initialized
through the `With`.
+ sql("DECLARE OR REPLACE VARIABLE v4 DOUBLE DEFAULT nullif(rand(5),
0.5)")
+ checkAnswer(sql("SELECT v4 IS NULL OR (v4 >= 0 AND v4 < 1)"),
Row(true))
Review Comment:
`v4 IS NULL OR (v4 >= 0 AND v4 < 1)` holds for every value `nullif(rand(5),
0.5)` can produce, so this checks that the DECLARE is planned and that the
`Rand` inside the definition is initialized, but not the value. An exact value
wouldn't catch a double evaluation either, since each fresh copy of `Rand(5)`
is initialized with seed 5 and partition 0 and draws the same first number.
Maybe say in the comment that this is a planning and initialization check,
since the counters above cover the once-only part.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -67,6 +67,16 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
// changed anything: `mapExpressions` and `withNewChildren` preserve
reference equality when
// they rewrite nothing, which is what makes this detectable.
if ((rewrittenAgg eq agg) && (rewrittenProj eq proj)) p else
rewrittenProj
+ // The children of a `CreateVariable` are the `ResolvedIdentifier`s
naming the variables, not
+ // rows to evaluate the default over, and `V2CommandStrategy` plans the
command only while
+ // they stay so. `CreateVariableExec` evaluates the default once, so a
`With` in it is left in
+ // place to memoize its definition, as one in a conditional branch is.
+ case c: CreateVariable if c.defaultExpr.containsPattern(WITH_EXPRESSION)
=>
+ val default = c.defaultExpr.child
+ val rewritten =
default.transformUpWithPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case w: With => inlineDefsThatGainNothing(w)
+ }
+ if (rewritten eq default) c else c.copy(defaultExpr =
c.defaultExpr.copy(child = rewritten))
case p if p.expressions.exists(_.containsPattern(WITH_EXPRESSION)) =>
Review Comment:
(Pre-existing) `Call` isn't a `Command`, but it has the same shape: its
child is a `ResolvedProcedure` with no output, and `InvokeProcedures` runs the
optimizer during analysis and then expects `Call(ResolvedProcedure(...), ...)`.
With the `sum` procedure from `ProcedureSuite`,
```sql
CALL cat.ns.sum(nullif((SELECT 5), 3), 1)
```
should return 6, like `CALL cat.ns.sum((SELECT 5), 1)` does. But this batch
runs before `OptimizeOneRowRelationSubquery` can fold the subquery, the
subquery's empty references match the `ResolvedProcedure`, and the resulting
`Project` makes `InvokeProcedures` throw "Unexpected plan for optimized CALL
statement" (L42). A non-foldable argument such as `nullif(CAST(rand() * 10 AS
INT), 3)` also gets that internal error instead of "args must be foldable".
Keeping the `With` as this PR does wouldn't work for `Call`, since a `With`
isn't foldable; inlining would.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -67,6 +67,16 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
// changed anything: `mapExpressions` and `withNewChildren` preserve
reference equality when
// they rewrite nothing, which is what makes this detectable.
if ((rewrittenAgg eq agg) && (rewrittenProj eq proj)) p else
rewrittenProj
+ // The children of a `CreateVariable` are the `ResolvedIdentifier`s
naming the variables, not
+ // rows to evaluate the default over, and `V2CommandStrategy` plans the
command only while
+ // they stay so. `CreateVariableExec` evaluates the default once, so a
`With` in it is left in
+ // place to memoize its definition, as one in a conditional branch is.
+ case c: CreateVariable if c.defaultExpr.containsPattern(WITH_EXPRESSION)
=>
+ val default = c.defaultExpr.child
+ val rewritten =
default.transformUpWithPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case w: With => inlineDefsThatGainNothing(w)
+ }
+ if (rewritten eq default) c else c.copy(defaultExpr =
c.defaultExpr.copy(child = rewritten))
case p if p.expressions.exists(_.containsPattern(WITH_EXPRESSION)) =>
applyInternal(p)
Review Comment:
(Pre-existing) A `DELETE FROM` that stays a `DeleteFromTable`, on a v1 table
or on a v2 table that only implements `SupportsDelete` like the test catalog's
`InMemoryTable`, gets a `Project` between the command and its relation:
```sql
DELETE FROM testcat.ns.t WHERE (id + 1) BETWEEN 1 AND 3
```
Nothing removes that `Project`, so `DataSourceV2Strategy` falls through to
`throw SparkException.internalError("Unexpected table relation: " + other)`
(L582) instead of raising `_LEGACY_ERROR_TEMP_1110` here or
`UNSUPPORTED_FEATURE.TABLE_OPERATION` for a v1 table, or deleting on a
connector that accepts the predicate. This is the case the `ConvertToCatalyst`
doc describes: "the Project hides the relation DataSourceV2Strategy looks for"
(`Optimizer.scala` L1079-L1080).
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -67,6 +67,16 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
// changed anything: `mapExpressions` and `withNewChildren` preserve
reference equality when
// they rewrite nothing, which is what makes this detectable.
if ((rewrittenAgg eq agg) && (rewrittenProj eq proj)) p else
rewrittenProj
+ // The children of a `CreateVariable` are the `ResolvedIdentifier`s
naming the variables, not
+ // rows to evaluate the default over, and `V2CommandStrategy` plans the
command only while
Review Comment:
(Pre-existing, independent of `With`) Only this rule keeps this invariant.
`ExtractPythonUDFs`, in a later batch, also treats the `ResolvedIdentifier`s as
inputs: a Python UDF without references is a subset of every child's output
(L327), so it wraps each child in a `BatchEvalPython` (L353). With a regular
Python UDF,
```sql
DECLARE OR REPLACE VARIABLE v INT DEFAULT nullif(py_udf(1), 3)
```
still fails with "No plan for CreateVariable", and so does a plain `DEFAULT
py_udf(1)`. Even without the extraction, `CreateVariableExec` couldn't evaluate
a `PythonUDF`, so a clear analysis error next to the subquery check
(`CheckAnalysis` L1039) looks like the right fix. Not for this PR, but since
the title says NULLIF defaults are fixed, it may be worth a separate JIRA.
##########
sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala:
##########
@@ -645,6 +645,38 @@ class ColumnExpressionSuite extends SharedSparkSession {
}
}
+ test("SPARK-59954: a With in a DECLARE VARIABLE default is evaluated once") {
+ // `CreateVariableExec` evaluates the default on the driver, so the
accumulator needs no job to
+ // read. `plus_one_counted` is deterministic but not cheap, so its
definition is worth
+ // memoizing, and the first two defaults below each read it twice.
+ val calls = sparkContext.longAccumulator("calls")
+ spark.udf.register("plus_one_counted", (i: Int) => { calls.add(1); i + 1 })
+ spark.udf.register("plus_one_nondet",
+ udf((i: Int) => { calls.add(1); i + 1 }).asNondeterministic())
Review Comment:
Nit: these registrations run before `withUserDefinedFunction` sets up the
cleanup, so if the second one threw, `plus_one_counted` would stay registered
in the shared session. Registering inside the block, as other sql/core tests
do, costs nothing.
##########
sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala:
##########
@@ -645,6 +645,38 @@ class ColumnExpressionSuite extends SharedSparkSession {
}
}
+ test("SPARK-59954: a With in a DECLARE VARIABLE default is evaluated once") {
+ // `CreateVariableExec` evaluates the default on the driver, so the
accumulator needs no job to
+ // read. `plus_one_counted` is deterministic but not cheap, so its
definition is worth
+ // memoizing, and the first two defaults below each read it twice.
+ val calls = sparkContext.longAccumulator("calls")
+ spark.udf.register("plus_one_counted", (i: Int) => { calls.add(1); i + 1 })
+ spark.udf.register("plus_one_nondet",
+ udf((i: Int) => { calls.add(1); i + 1 }).asNondeterministic())
+ withUserDefinedFunction("plus_one_counted" -> true, "plus_one_nondet" ->
true) {
+ withSessionVariable("v1", "v2", "v3", "v4") {
+ sql("DECLARE OR REPLACE VARIABLE v1 INT DEFAULT
nullif(plus_one_counted(1), 3)")
+ assert(calls.value == 1, s"the NULLIF default called the UDF
${calls.value} times")
+ checkAnswer(sql("SELECT v1"), Row(2))
+
+ calls.reset()
+ sql("DECLARE OR REPLACE VARIABLE v2 BOOLEAN DEFAULT
plus_one_counted(1) BETWEEN 1 AND 3")
+ assert(calls.value == 1, s"the BETWEEN default called the UDF
${calls.value} times")
+ checkAnswer(sql("SELECT v2"), Row(true))
+
+ // A nondeterministic definition is evaluated once as well.
+ calls.reset()
+ sql("DECLARE OR REPLACE VARIABLE v3 INT DEFAULT
nullif(plus_one_nondet(1), 3)")
+ assert(calls.value == 1, s"the nondeterministic default called it
${calls.value} times")
Review Comment:
This takes the same path as v1. `canSubstitute` is
`!multiplyReferenced.contains(id) || (CollapseProject.isCheap(child) &&
child.deterministic)`, and a `ScalaUDF` isn't foldable, so `isCheap` is already
false and `deterministic` is never checked. The block would still pass without
`&& child.deterministic`. The determinism check is covered at the plan level by
"SPARK-58902: a cheap definition is only inlined if it is also deterministic"
in `RewriteWithExpressionSuite`, so this block could be dropped, or the comment
could say it's the same path as v1.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpression.scala:
##########
@@ -67,6 +67,16 @@ object RewriteWithExpression extends Rule[LogicalPlan] {
// changed anything: `mapExpressions` and `withNewChildren` preserve
reference equality when
// they rewrite nothing, which is what makes this detectable.
if ((rewrittenAgg eq agg) && (rewrittenProj eq proj)) p else
rewrittenProj
+ // The children of a `CreateVariable` are the `ResolvedIdentifier`s
naming the variables, not
+ // rows to evaluate the default over, and `V2CommandStrategy` plans the
command only while
+ // they stay so. `CreateVariableExec` evaluates the default once, so a
`With` in it is left in
+ // place to memoize its definition, as one in a conditional branch is.
+ case c: CreateVariable if c.defaultExpr.containsPattern(WITH_EXPRESSION)
=>
+ val default = c.defaultExpr.child
+ val rewritten =
default.transformUpWithPruning(_.containsPattern(WITH_EXPRESSION)) {
+ case w: With => inlineDefsThatGainNothing(w)
+ }
Review Comment:
Nit: this is the same block as L388-L390, and those are the only two callers
of `inlineDefsThatGainNothing`. A small helper, such as `private def
inlineWithsThatGainNothing(e: Expression): Expression`, would keep the policy
(inline only what gains nothing, bottom-up for nested `With`) and the "Use
transformUp to handle nested With" note in one place. Combined with my previous
comment, this case would become `c.mapExpressions(inlineWithsThatGainNothing)`.
##########
sql/core/src/test/scala/org/apache/spark/sql/ColumnExpressionSuite.scala:
##########
@@ -645,6 +645,38 @@ class ColumnExpressionSuite extends SharedSparkSession {
}
}
+ test("SPARK-59954: a With in a DECLARE VARIABLE default is evaluated once") {
+ // `CreateVariableExec` evaluates the default on the driver, so the
accumulator needs no job to
+ // read. `plus_one_counted` is deterministic but not cheap, so its
definition is worth
+ // memoizing, and the first two defaults below each read it twice.
+ val calls = sparkContext.longAccumulator("calls")
+ spark.udf.register("plus_one_counted", (i: Int) => { calls.add(1); i + 1 })
+ spark.udf.register("plus_one_nondet",
+ udf((i: Int) => { calls.add(1); i + 1 }).asNondeterministic())
+ withUserDefinedFunction("plus_one_counted" -> true, "plus_one_nondet" ->
true) {
Review Comment:
The SPARK-58902 tests above wrap their bodies in
`withSQLConf(SQLConf.ALWAYS_INLINE_COMMON_EXPR.key -> "false")`, "so a
suite-level override cannot turn this into a test of the inlining path".
`NullIf` and `Between` read that conf when they are built, so with it set to
true they build no `With`, and the three `calls.value == 1` assertions see 2.
Nothing sets it today, but pinning it the same way would keep this test on the
`With` path.
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteWithExpressionSuite.scala:
##########
@@ -486,6 +486,21 @@ class RewriteWithExpressionSuite extends PlanTest {
)
}
+ test("SPARK-59954: a With in a DECLARE VARIABLE default leaves the variable
names alone") {
+ val ident = ResolvedIdentifier(FakeSystemCatalog,
Identifier.of(Array("session"), "v"))
+ def declare(default: Expression): CreateVariable =
+ CreateVariable(Seq(ident), DefaultValueExpression(default, "v"), replace
= true)
+
+ // The children are the variable names, not rows, so there is no `Project`
to pre-evaluate a
+ // definition in. One worth memoizing stays in the default for
`CreateVariableExec` to evaluate.
+ val kept = declare(With(Rand(Literal(0L))) { case Seq(ref) => ref * ref })
+ assert(Optimizer.execute(kept) == kept)
Review Comment:
Nit: the rest of this suite compares plans with its `comparePlans` override,
e.g. `comparePlans(Optimizer.execute(keptPlan), keptPlan)` on L537, which also
runs the analyzer's sanity check and prints a tree diff on failure. Same for
L500-L501.
--
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]
