Skip to content

[SPARK-53311][SQL][PYTHON][CORE] Make PullOutNonDeterministic use canonicalized expressions #52061

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
benrobby wants to merge 2 commits into from
Closed

[SPARK-53311][SQL][PYTHON][CORE] Make PullOutNonDeterministic use canonicalized expressions #52061

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
14a67cf
[SPARK-53311] make pulloutnondeterministic use canonicalized expressions
benhurdelhey Aug 18, 2025
08230da
fix lint
benhurdelhey Aug 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion ...n/scala/org/apache/spark/sql/catalyst/analysis/NondeterministicExpressionCollection.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ object NondeterministicExpressionCollection {
case namedExpression: NamedExpression => namedExpression
case _ => Alias(nondeterministicExpr, "_nondeterministic")()
}
nonDeterministicToAttributes.put(nondeterministicExpr, namedExpression)
nonDeterministicToAttributes.put(nondeterministicExpr.canonicalized, namedExpression)
}
}
}
Expand Down
4 changes: 2 additions & 2 deletions ...alyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/PullOutNondeterministic.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ object PullOutNondeterministic extends Rule[LogicalPlan] {
NondeterministicExpressionCollection.getNondeterministicToAttributes(a.groupingExpressions)
val newChild = Project(a.child.output ++ nondeterToAttr.values.asScala.toSeq, a.child)
val deterministicAggregate = a.transformExpressions { case e =>
Option(nondeterToAttr.get(e)).map(_.toAttribute).getOrElse(e)
Option(nondeterToAttr.get(e.canonicalized)).map(_.toAttribute).getOrElse(e)
Copy link
Contributor

@cloud-fan cloud-fan Aug 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks risky to me. We should not deduplicate nondeterministic expressions, as they might have side effects.

Copy link
Contributor

@cloud-fan cloud-fan Aug 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok we are already doing it with non-canonicalized expression. I'm fine with this change now

Copy link
Author

@benrobby benrobby Aug 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the quick review! Yes, this rule PullOutNondeterministic is already doing that. Let me fix the tests..

}.copy(child = newChild)

deterministicAggregate.groupingExpressions.foreach(expr => if (!expr.deterministic) {
Expand All @@ -69,7 +69,7 @@ object PullOutNondeterministic extends Rule[LogicalPlan] {
val nondeterToAttr =
NondeterministicExpressionCollection.getNondeterministicToAttributes(p.expressions)
val newPlan = p.transformExpressions { case e =>
Option(nondeterToAttr.get(e)).map(_.toAttribute).getOrElse(e)
Option(nondeterToAttr.get(e.canonicalized)).map(_.toAttribute).getOrElse(e)
}
val newChild = Project(p.child.output ++ nondeterToAttr.values.asScala.toSeq, p.child)
Project(p.output, newPlan.withNewChildren(newChild :: Nil))
Expand Down
5 changes: 3 additions & 2 deletions sql/core/src/test/scala/org/apache/spark/sql/IntegratedUDFTestUtils.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -474,7 +474,8 @@ object IntegratedUDFTestUtils extends SQLHelper {
* casted_col.cast(df.schema["col"].dataType)
* }}}
*/
case class TestPythonUDF(name: String, returnType: Option[DataType] = None) extends TestUDF {
case class TestPythonUDF(name: String, returnType: Option[DataType] = None,
deterministic: Boolean = true) extends TestUDF {
private[IntegratedUDFTestUtils] lazy val udf = new UserDefinedPythonFunction(
name = name,
func = SimplePythonFunction(
Expand All @@ -487,7 +488,7 @@ object IntegratedUDFTestUtils extends SQLHelper {
accumulator = null),
dataType = StringType,
pythonEvalType = PythonEvalType.SQL_BATCHED_UDF,
udfDeterministic = true) {
udfDeterministic = deterministic) {

override def builder(e: Seq[Expression]): Expression = {
assert(e.length == 1, "Defined UDF only has one column")
Expand Down
19 changes: 18 additions & 1 deletion sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
package org.apache.spark.sql.execution.python

import org.apache.spark.sql.{AnalysisException, IntegratedUDFTestUtils, QueryTest, Row}
import org.apache.spark.sql.functions.{array, col, count, transform}
import org.apache.spark.sql.functions.{array, avg, col, count, transform}
import org.apache.spark.sql.test.SharedSparkSession
import org.apache.spark.sql.types.LongType

Expand Down Expand Up @@ -139,4 +139,21 @@ class PythonUDFSuite extends QueryTest with SharedSparkSession {
checkAnswer(df, Row(0, 1, 1, 0, 1, 1))
}
}

test("SPARK-53311: Nondeterministic Python UDF pull out in aggregate with grouping") {
assume(shouldTestPythonUDFs)

// nondeterministic UDF
val pythonUDF = TestPythonUDF(name = "foo", Some(LongType), deterministic = false)

// This query should work without throwing an analysis exception
// The UDF foo(value) appears in both grouping expressions and aggregate expressions
// The fix ensures that both instances are properly mapped to the same attribute
val df = spark.range(1)
.selectExpr("id", "id % 3 as value")
.groupBy(pythonUDF(col("value")))
.agg(avg("id"), pythonUDF(col("value")))

checkAnswer(df, Row(0, 0.0, 0))
}
}