LuciferYang commented on code in PR #59209:
URL: https://github.com/apache/spark/pull/59209#discussion_r4204441401
##########
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:
Dropped the block.
##########
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:
Done, here and on the inlined case.
##########
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:
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]
