LuciferYang commented on code in PR #59209:
URL: https://github.com/apache/spark/pull/59209#discussion_r4204439260
##########
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:
Fixed by #59217; `DataSourceV2SQLSuiteV1Filter` there now gets
`_LEGACY_ERROR_TEMP_1110`.
##########
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:
#59217 now covers `Call` too: the case matches `Command | Call`, and this
CALL returns 6 (new test in `ProcedureSuite`). The nondeterministic form gets
`args must be foldable`.
##########
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:
Reproduced, with and without this PR: `DEFAULT py_udf(1)` and `DEFAULT
nullif(py_udf(1), '3')` fail with `No plan for CreateVariable` over a
`BatchEvalPython` wrapping the `ResolvedIdentifier`, while `SELECT py_udf(1)`
and `SET VAR v = py_udf(1)` work. Filed as SPARK-60030.
##########
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:
Done.
##########
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:
Done, as `inlineWithsThatGainNothing`, shared with the conditional-branch
path.
##########
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:
Done; the comment now says it checks planning and `Rand` initialization, not
the value.
##########
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:
Done.
--
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]
