You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ho...@apache.org on 2026/09/25 00:23:07 UTC
(spark) branch branch-4.3 updated: [SPARK-58626][PYTHON][SQL] Optimize evaluating inputs to transpiled UDFs
This is an automated email from the ASF dual-hosted git repository.
asf-gitbox-commits pushed a commit to branch branch-4.3
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.3 by this push:
new 9bf43b7f2cd5 [SPARK-58626][PYTHON][SQL] Optimize evaluating inputs to transpiled UDFs
9bf43b7f2cd5 is described below
commit 9bf43b7f2cd5473688c212b8fcce3e6812eca202
Author: Holden Karau <ho...@pigscanfly.ca>
AuthorDate: Thu Sep 24 17:00:07 2026 -0700
[SPARK-58626][PYTHON][SQL] Optimize evaluating inputs to transpiled UDFs
`_udf_param_N` in a transpiled option used to be replaced by a copy of the bound argument, so an
argument the body reads N times was evaluated N times. Each one is now a reference to the call's Nth
argument, the argument itself stays on the Python UDF node, and `ConvertToCatalyst` decides per call
what to do with it.
The rule: an argument the body reads more than once is computed once a row, or we don't transpile the
call at all.
However subsequent optimizations may decide to inline some deterministic expressions which will change this.
Reading an argument twice evaluated it twice. That's wasted work for anything expensive and a potentiall incorrect for non-deterministic inputs.
Yes.
- Each argument a body reads is evaluated once a row, so a body reading a `rand()` argument twice sees
one draw rather than two.
- A body that reads an expensive argument more than once runs as interpreted Python where no column
fits, instead of lowering with the argument recomputed at each read.
- A call reading both sides of a non-inner join condition, with a parameter read more than once, keeps
the Python UDF, which `ExtractPythonUDFFromJoinCondition` turns down
(`UNSUPPORTED_FEATURE.PYTHON_UDF_IN_ON_CLAUSE`) -- the same result as with transpilation off.
`ConvertToCatalystSuite` for the plan shapes: one column per distinct argument, a column per call, one
column for a call sitting in two slots of one operator, per-argument placement, and the Python UDF kept
for each position that can't hold a column. Plus two for the walkers themselves: `substitute` leaving a
nested call's options alone, and a reference to an argument the call hasn't got being a classed error
rather than an unresolved node. Its helper runs `LogicalPlanIntegrity.validateOptimizedPlan` around
every case, and walks every operator for a dangling reference itself, since
`validateNoDanglingReferences` only ever looks at the root and the root is never the Project this
inserts.
`TranspiledUDFParameterSuite` for real plans: the grouping key's column, the `Aggregate` and lambda
fallbacks, an `Aggregate` whose only call sits in a subquery, the straddling non-inner join, a draw
above a `Join` keeping the join's condition -- at the child and again with a `select` in between -- and
a MERGE instruction condition where an argument that mods by zero only raises if the column is really
there.
`test_udf_transpile_unit.py` for values rather than plans: one evaluation per parameter, a draw per
call, the fallbacks, and `orderBy(u(rand()))`, `repartition(2, u(rand()))` and `groupBy(u(rand()))`,
which used to fail with `INVALID_UDF_PARAMETER_PLACEHOLDER_INDEX`.
Yes, Claude.
Generated-by: Claude (Opus 5)
Closes #58189 from holdenk/SPARK-58626-avoid-double-eval.
Lead-authored-by: Holden Karau <ho...@pigscanfly.ca>
Co-authored-by: Holden Karau <ho...@snowflake.com>
Signed-off-by: Holden Karau <ho...@snowflake.com>
(cherry picked from commit 9907192c92f35202d21cfef1a5f1b441c5e81546)
Signed-off-by: Holden Karau <ho...@snowflake.com>
---
.../pyspark/sql/tests/test_udf_transpile_parity.py | 23 +-
.../pyspark/sql/tests/test_udf_transpile_unit.py | 369 ++++++++++++++++
python/pyspark/sql/transpile.py | 45 ++
.../ResolveTranspiledPythonUDFOptions.scala | 27 +-
.../spark/sql/catalyst/expressions/PythonUDF.scala | 135 +++++-
.../spark/sql/catalyst/optimizer/Optimizer.scala | 333 +++++++++++++-
.../spark/sql/catalyst/trees/TreePatterns.scala | 1 +
.../optimizer/ConvertToCatalystSuite.scala | 490 ++++++++++++++++++++-
.../python/UserDefinedPythonFunction.scala | 31 +-
.../python/TranspiledUDFParameterSuite.scala | 268 +++++++++++
10 files changed, 1661 insertions(+), 61 deletions(-)
diff --git a/python/pyspark/sql/tests/test_udf_transpile_parity.py b/python/pyspark/sql/tests/test_udf_transpile_parity.py
index a8276df0b162..cf3122577fff 100644
--- a/python/pyspark/sql/tests/test_udf_transpile_parity.py
+++ b/python/pyspark/sql/tests/test_udf_transpile_parity.py
@@ -20,12 +20,13 @@ transpilation enabled.
Transpilation is only attempted when both
``spark.sql.experimental.optimizer.transpilePyUDFs`` and
-``spark.sql.ansi.enabled`` are true, and it is designed to fall back to
-interpreted Python rather than risk semantic drift. These classes re-run the
-shared UDF mixins under that configuration so we can confirm that turning on the
-experimental feature does not change UDF results compared with the default
-(transpilation off) runs covered by the original concrete classes
-(``UDFTests``, ``UDFCombinationsTests``, ``UnifiedUDFTests``).
+``spark.sql.ansi.enabled`` are true. When either is off, or if transpilation
+fails, we fall back to the regular non-transpiled code path.
+
+These classes re-run the shared UDF mixins under that configuration to confirm
+turning the feature on does not change UDF results. In two cases the transpiled body
+evaluates less than interpreted Python, so if evaluating the input would cause an
+error the transpiled code may succeed where the regular path would not.
Transpilation is currently only supported in regular (non-Connect) Spark, so
these classes are guarded with ``is_remote_only()`` and are intentionally not
@@ -34,12 +35,10 @@ the transpiler directly live in ``test_udf_transpile_unit.py`` and
``test_udf_transpile_hypothesis.py``.
Note on configuration: enabling transpilation requires ANSI mode, so an "on"
-run is unavoidably also an ANSI run. All inherited tests currently pass as-is
-under this configuration, so no per-test overrides are defined here. If a future
-change makes an inherited test diverge purely due to ANSI semantics or because
-transpilation bypasses a Python-side effect (rather than a genuine result
-change), override it here with a documented ``unittest.skip`` rather than
-editing the inherited test body.
+run is unavoidably also an ANSI run.
+
+When an inherited test diverges in an acceptable way, override it here with a
+documented ``unittest.skip`` rather than editing the inherited test body.
"""
import unittest
diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py b/python/pyspark/sql/tests/test_udf_transpile_unit.py
index 298cea0b223c..33084697492b 100644
--- a/python/pyspark/sql/tests/test_udf_transpile_unit.py
+++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py
@@ -47,6 +47,12 @@ _TRANSPILE_ON = {
"spark.sql.ansi.enabled": True,
}
+# The interpreted path to compare against where ANSI is what makes the two differ.
+_TRANSPILE_OFF_ANSI_ON = {
+ "spark.sql.experimental.optimizer.transpilePyUDFs": False,
+ "spark.sql.ansi.enabled": True,
+}
+
@unittest.skipIf(
is_remote_only(),
@@ -1470,6 +1476,369 @@ class UDFTranspileUnitTests(ReusedSQLTestCase):
def _eval_python_count(df):
return df._jdf.queryExecution().executedPlan().toString().count("EvalPython")
+ @staticmethod
+ def _optimized_plan(df):
+ return df._jdf.queryExecution().optimizedPlan().toString()
+
+ # Internal plan-string details: the alias a pre-evaluated argument gets (ConvertToCatalyst
+ # names them ``_udf_param_N``), and the token a plan shows for a draw it still evaluates.
+ _SHARED_ARG_ALIAS = "_udf_param_"
+ _DRAW = "rand("
+
+ def _shares_an_argument(self, df):
+ return self._SHARED_ARG_ALIAS in self._optimized_plan(df)
+
+ def _draw_count(self, df):
+ return self._optimized_plan(df).count(self._DRAW)
+
+ def test_udf_transpile_evaluates_inputs_once(self):
+ # SPARK-58626: an argument the body uses twice is drawn once per row, so `x == x` over a
+ # random input is True on every row and `x if x > 0.5 else 0.0` never returns a value its
+ # own condition would have rejected.
+ from pyspark.sql.functions import expr, lit, rand
+
+ eq_self = lambda x: x == x # noqa: E731
+ clamp = lambda x: x if x > 0.5 else 0.0 # noqa: E731
+ square_fn = lambda x: x * x # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ eq_udf = UserDefinedFunction(eq_self, BooleanType())
+ clamp_udf = UserDefinedFunction(clamp, DoubleType())
+ self.assertTrue(eq_udf.transpiled and clamp_udf.transpiled)
+ rows = self.spark.range(100)
+
+ # `functions.rand()` bakes a literal seed in, but `expr("rand()")` arrives with the seed
+ # unresolved and ResolveRandomSeed then gives each spliced copy its OWN seed. Both have
+ # to end up sharing one evaluation, or `x == x` compares two independent draws.
+ #
+ # Sharing leaves `col = col`, which SimplifyBinaryComparison folds to true and column
+ # pruning then drops the draw entirely -- so zero draws in the plan is the proof that
+ # both sides became the same column. Two independent draws would leave both.
+ for arg in (rand(), expr("rand()")):
+ eq_df = rows.select(eq_udf(arg).alias("v"))
+ self.assertEqual(0, self._draw_count(eq_df))
+ self.assertTrue(all(r[0] for r in eq_df.collect()))
+
+ clamped = rows.select(clamp_udf(rand()).alias("v"))
+ self.assertEqual(1, self._draw_count(clamped))
+ vals = [r[0] for r in clamped.collect()]
+ self.assertTrue(all(v == 0.0 or v > 0.5 for v in vals), vals)
+ # Both branches have to run for the check above to mean anything.
+ self.assertTrue(any(v == 0.0 for v in vals) and any(v > 0.5 for v in vals), vals)
+
+ # A literal and a plain column are as cheap to read twice as to share, so the optimizer
+ # puts them back at each use site and neither grows the plan.
+ square = UserDefinedFunction(square_fn, LongType())
+ self.assertTrue(square.transpiled)
+ lit_df = rows.select(square(lit(3)).alias("v"))
+ self.assertFalse(self._shares_an_argument(lit_df))
+ self.assertEqual([9] * 100, [r[0] for r in lit_df.collect()])
+ col_df = self.spark.range(3).select(square("id").alias("v"))
+ self.assertFalse(self._shares_an_argument(col_df))
+ self.assertEqual([0, 1, 4], [r[0] for r in col_df.collect()])
+
+ def test_udf_transpile_shares_a_python_udf_argument(self):
+ # SPARK-58626: an argument that is itself a nondeterministic Python UDF gets a column like
+ # any other draw the body reads twice, even though CollapseProject.isCheap counts a
+ # PythonUDF as cheap. Cheapness is about wasted work; two calls would be two draws, and
+ # `a - a` would come out nonzero.
+ #
+ # ExtractPythonUDFs would have pointed both copies at one result attribute anyway, so this
+ # was the right answer before the column too -- by accident, from an unrelated rule, and at
+ # the cost of a second Python round trip (the dedup runs through an ExpressionSet, which
+ # holds a nondeterministic expression once per add).
+ import random
+
+ draw = lambda x: random.random() # noqa: E731
+ subtract_self = lambda a: a - a # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ rand_udf = UserDefinedFunction(draw, DoubleType()).asNondeterministic()
+ diff = UserDefinedFunction(subtract_self, DoubleType())
+ self.assertTrue(diff.transpiled)
+ df = self.spark.range(20).select(diff(rand_udf("id")))
+ self.assertTrue(self._shares_an_argument(df), self._optimized_plan(df))
+ self.assertEqual([0.0] * 20, [r[0] for r in df.collect()], "One draw, read twice")
+
+ def test_udf_transpile_repeats_a_deterministic_python_udf_argument(self):
+ # SPARK-58626: a deterministic Python UDF argument is cheap to repeat as far as
+ # CollapseProject.isCheap is concerned, so it stays at each use site instead of getting a
+ # column -- and ExtractPythonUDFs then points both copies at one eval node, so the body
+ # still reads one evaluation. What we promise is the evaluation count, not which rule keeps
+ # it, and this is the shape where another rule does. A second argument because a call whose
+ # arguments are all Python UDFs keeps the batch pipeline instead of lowering.
+ def helper(v):
+ return v + 1
+
+ opaque = lambda x: helper(x) # noqa: E731
+ first_twice = lambda a, b: a - a + b # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ inner = UserDefinedFunction(opaque, LongType())
+ self.assertEqual([], inner.transpiled, "The inner UDF has to stay Python")
+ outer = self._transpiled_udf(first_twice, LongType())
+ df = self.spark.range(5).select(outer(inner("id"), "id"))
+ plan = self._optimized_plan(df)
+ self.assertFalse(self._shares_an_argument(df), plan)
+ self.assertEqual(
+ 1,
+ self._eval_python_count(df),
+ df._jdf.queryExecution().executedPlan().toString(),
+ )
+ self.assertEqual(list(range(5)), [r[0] for r in df.collect()])
+
+ def test_udf_transpile_declines_inside_a_higher_order_function(self):
+ # SPARK-58626: a call inside a lambda is not lowered at all -- Spark applies a Python UDF
+ # over the whole array there and we leave that to it -- so the draw stays a single draw per
+ # call rather than one per read.
+ #
+ # `clamp` returns its input or 0.0, so a value in (0.0, 0.5] would be two draws: the
+ # condition saw one and the branch another.
+ from pyspark.sql.functions import array, col, lit, rand, transform
+
+ clamp = lambda x: x if x > 0.5 else 0.0 # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ c = self._transpiled_udf(clamp, DoubleType())
+ arr = self.spark.range(500).select(array(lit(1.0), lit(2.0)).alias("a"))
+ hof = arr.select(transform(col("a"), lambda e: c(rand())).alias("v"))
+ self.assertGreater(self._eval_python_count(hof), 0, self._optimized_plan(hof))
+ vals = [v for r in hof.collect() for v in r["v"]]
+ self.assertTrue(vals and all(v == 0.0 or v > 0.5 for v in vals), vals)
+
+ def test_udf_transpile_declines_under_an_aggregate(self):
+ # SPARK-58626: an Aggregate can't host the column -- a result expression no aggregate
+ # function wraps has to be built from the grouping expressions -- so a body reading a
+ # nondeterministic argument twice keeps the interpreted UDF, which draws once. `clamp`
+ # returns its input or 0.0, so a value in (0.0, 0.5] would be two draws.
+ from pyspark.sql.functions import col, rand
+
+ clamp = lambda x: x if x > 0.5 else 0.0 # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ c = self._transpiled_udf(clamp, DoubleType())
+ agg = self.spark.range(500).groupBy(col("id").alias("g")).agg(c(rand()).alias("v"))
+ self.assertGreater(self._eval_python_count(agg), 0, self._optimized_plan(agg))
+ vals = [r["v"] for r in agg.collect()]
+ self.assertTrue(all(v == 0.0 or v > 0.5 for v in vals), vals)
+ self.assertTrue(any(v == 0.0 for v in vals) and any(v > 0.5 for v in vals), vals)
+
+ def test_udf_transpile_gives_each_call_its_own_draw(self):
+ # SPARK-58626: one draw per parameter, not one per operator. Two calls each passing their
+ # own `rand()` owe two draws, so the doubled values differ on every row. Keying a shared
+ # column on the parameter index alone put both bodies on one draw and made this exactly 0.0.
+ from pyspark.sql.functions import rand
+
+ add_self = lambda x: x + x # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ u = self._transpiled_udf(add_self, DoubleType())
+ df = self.spark.range(500).select((u(rand()) - u(rand())).alias("v"))
+ self.assertEqual(0, self._eval_python_count(df))
+ self.assertEqual(2, self._optimized_plan(df).count("AS " + self._SHARED_ARG_ALIAS))
+ diffs = [r["v"] for r in df.collect()]
+ self.assertTrue(all(d != 0.0 for d in diffs), "each call draws for itself")
+
+ def test_udf_transpile_position_can_rule_out_sharing(self):
+ # SPARK-58626: inside a lambda nothing is lowered, so nothing is pre-evaluated either -- the
+ # empty array below must NOT raise even though the argument divides by zero.
+ #
+ # A conditional branch, by contrast, DOES get one evaluation, even though a bare Catalyst
+ # `when` is lazy. Measured: the interpreted Python UDF this replaces evaluates its inputs
+ # in a projection below the conditional and raises there too, so eager is the Python-parity
+ # answer as well as the cheaper one. See transpile.py.
+ from pyspark.sql.functions import array, col, lit, rand, transform, when
+
+ clamp = lambda x: x if x > 0.5 else 0.0 # noqa: E731
+ used_twice = lambda a, b: (a + a) + b # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ c = UserDefinedFunction(clamp, DoubleType())
+ self.assertTrue(c.transpiled)
+ rows = self.spark.range(200).select((col("id") + 1).cast("double").alias("x"))
+
+ nested = rows.select(when(col("x") > 0, c(rand())).otherwise(lit(-1.0)).alias("v"))
+ vals = [r[0] for r in nested.collect()]
+ # `x` is `id + 1`, so the otherwise branch never fires and every value came from the
+ # body -- which means a value in (0.0, 0.5] could only be the body seeing two draws.
+ self.assertTrue(
+ all(v == 0.0 or v > 0.5 for v in vals),
+ f"a conditional branch must share one draw too: {vals}",
+ )
+ self.assertEqual(1, self._draw_count(nested))
+ self.assertTrue(self._shares_an_argument(nested))
+
+ s = UserDefinedFunction(used_twice, DoubleType())
+ self.assertTrue(s.transpiled)
+ df = self.spark.createDataFrame([(1.0, 0.0)], "a double, b double").select(
+ col("a"), col("b"), array().cast("array<double>").alias("arr")
+ )
+ lazy = df.select(transform(col("arr"), lambda e: s(col("a") / col("b"), e)).alias("v"))
+ self.assertEqual([[]], [r[0] for r in lazy.collect()])
+
+ def test_udf_transpile_shares_a_nondeterministic_argument_in_a_predicate(self):
+ # SPARK-58626: a `where` keeps the shared column when the argument is nondeterministic.
+ # PushPredicateThroughNonJoin only pushes a Filter through a Project whose fields are all
+ # deterministic, so the draw is not substituted back and the body sees one value. A
+ # deterministic argument IS pushed back and evaluated per use -- same rows, just extra work.
+ from pyspark.sql.functions import col, rand
+
+ add_self = lambda x: x + x # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ u = UserDefinedFunction(add_self, DoubleType())
+ self.assertTrue(u.transpiled)
+
+ nd = self.spark.range(0, 10).where(u(rand()) > 0.5)
+ self.assertTrue(self._shares_an_argument(nd), self._optimized_plan(nd))
+ self.assertEqual(1, self._draw_count(nd))
+
+ det = self.spark.range(0, 10).where(u((col("id") % 3).cast("double")) > 1.0)
+ self.assertFalse(self._shares_an_argument(det), self._optimized_plan(det))
+
+ def test_udf_transpile_udf_as_a_predicate(self):
+ # A predicate is transpiled like anything else in a `where`: no Python worker, same answers.
+ # In a join condition it is not, since no column can go there -- see below.
+ from pyspark.sql.functions import col
+
+ used_twice = lambda a, b: (a + a) if b > 0 else 0.0 # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ u = UserDefinedFunction(used_twice, DoubleType())
+ self.assertTrue(u.transpiled)
+ df = self.spark.createDataFrame([(1.0, 2.0), (4.0, 2.0)], "a double, b double")
+ other = self.spark.createDataFrame([(2.0,)], "c double")
+ predicate = u(col("a") / col("b"), col("b")) > 2.0
+
+ filtered = df.where(predicate)
+ self.assertEqual(0, self._eval_python_count(filtered))
+ self.assertEqual([(4.0, 2.0)], [(r[0], r[1]) for r in filtered.collect()])
+
+ # A join has no side to hold the column, so the call goes back to Python -- the same
+ # path the predicate takes with transpilation off.
+ joined = df.join(other, predicate)
+ self.assertGreater(self._eval_python_count(joined), 0, self._optimized_plan(joined))
+ self.assertEqual([(4.0, 2.0, 2.0)], [tuple(r) for r in joined.collect()])
+
+ def test_udf_transpile_survives_an_argument_the_analyzer_moved(self):
+ # SPARK-58626: PullOutNondeterministic pulls a nondeterministic call into a projection below
+ # an operator that cannot hold one, leaving an attribute where the call was -- and so no
+ # arguments for the option's references to point at. The call keeps the Python path there
+ # rather than failing the query, which is also one evaluation: the projection computes it
+ # once. Master lowered these, with the draw spliced into the option body and read twice from
+ # the column PullOutNondeterministic left; this gives up that lowering, not correctness.
+ from pyspark.sql.functions import col, rand
+
+ add_self = lambda x: x + x # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ u = self._transpiled_udf(add_self, DoubleType())
+ rows = self.spark.range(0, 4).select(col("id").cast("double").alias("a"))
+ self.assertEqual(4, len(rows.orderBy(u(rand())).collect()))
+ self.assertEqual(4, len(rows.repartition(2, u(rand())).collect()))
+ self.assertEqual(4, len(rows.groupBy(u(rand())).count().collect()))
+
+ def test_udf_transpile_evaluates_inputs_once_in_a_group_by(self):
+ # SPARK-58626: an Aggregate is the awkward one -- a use no aggregate function wraps has to
+ # *be* a grouping expression, so it cannot read a pre-evaluated column. `x * x` reads its
+ # parameter twice and `a + 1` is not cheap, so all three shapes here keep the interpreted
+ # UDF, which evaluates the argument once itself. This pins the answers and that fallback.
+ from pyspark.sql.functions import col
+ from pyspark.sql.functions import sum as sum_
+
+ square_fn = lambda x: x * x # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ square = UserDefinedFunction(square_fn, LongType())
+ self.assertTrue(square.transpiled)
+ df = self.spark.createDataFrame([(1,), (2,), (2,)], "a long")
+
+ grouped = df.groupBy(col("a") + 1).agg(square(col("a") + 1).alias("v"))
+ self.assertEqual([(2, 4), (3, 9)], sorted((r[0], r[1]) for r in grouped.collect()))
+
+ summed = df.agg(sum_(square(col("a") + 1)).alias("s"))
+ self.assertEqual([22], [r[0] for r in summed.collect()])
+
+ mixed = df.groupBy(col("a") + 1).agg(
+ square(col("a") + 1).alias("v"), sum_(square(col("a") + 1)).alias("s")
+ )
+ self.assertEqual([(2, 4, 4), (3, 9, 18)], sorted(tuple(r) for r in mixed.collect()))
+
+ def test_udf_transpile_evaluates_inputs_once_in_a_subquery(self):
+ # SPARK-58626: end to end because decorrelation is what makes this interesting. In the last
+ # shape the shared argument reads a column of the OUTER query, so the pre-evaluated column
+ # lands inside the subquery holding an OuterReference, and decorrelation has to carry it
+ # back out. With decorrelation off we decline the column instead -- see ConvertToCatalyst.
+ square_fn = lambda x: x * x # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ square = UserDefinedFunction(square_fn, LongType())
+ self.assertTrue(square.transpiled)
+ # Both registrations go through SQLTestUtils' context managers: ReusedSQLTestCase
+ # shares one session across the class, so a failure between a bare register() and the
+ # try would leave sq_test_58626 in every later test's session.
+ with self.temp_func("sq_test_58626"), self.temp_view("t_58626"):
+ self.spark.udf.register("sq_test_58626", square)
+ self.spark.range(0, 5).selectExpr("id as a").createOrReplaceTempView("t_58626")
+ uncorrelated = self.spark.sql(
+ "SELECT a FROM t_58626 WHERE a < "
+ "(SELECT max(sq_test_58626(a + 1)) FROM t_58626)"
+ )
+ self.assertEqual([0, 1, 2, 3, 4], [r[0] for r in uncorrelated.collect()])
+
+ correlated = self.spark.sql(
+ "SELECT a FROM t_58626 o WHERE a < "
+ "(SELECT max(sq_test_58626(i.a + 1)) FROM t_58626 i WHERE i.a = o.a)"
+ )
+ self.assertEqual([0, 1, 2, 3, 4], sorted(r[0] for r in correlated.collect()))
+
+ outer_arg = self.spark.sql(
+ "SELECT a FROM t_58626 o WHERE EXISTS "
+ "(SELECT 1 FROM t_58626 i WHERE sq_test_58626(o.a + 1) = i.a)"
+ )
+ self.assertEqual([0, 1], sorted(r[0] for r in outer_arg.collect()))
+
+ def test_udf_transpile_puts_no_column_directly_above_an_aggregate(self):
+ # SPARK-58626: `splitSubquery` finds the Aggregate a correlated scalar subquery is built on
+ # by matching `Filter(_, Aggregate)` -- adjacency, not a search -- in an earlier batch than
+ # the pushdown that would collapse a Project away. A column between a HAVING and its
+ # Aggregate hid it, `mayHaveCountBug` read false, and the outer rows matching nothing were
+ # dropped instead of counted 0: measured as [] where the same plan without the column
+ # returns [3, 4]. So nothing goes directly above an Aggregate. The call keeps the Python UDF
+ # and raises here exactly as transpilation-off does, because a Python eval node lands
+ # between the Filter and the Aggregate just the same -- that half is not ours.
+ square_fn = lambda x: x * x # noqa: E731
+ legacy = {**_TRANSPILE_ON, "spark.sql.legacy.scalarSubqueryCountBugBehavior": True}
+ # `count(*) + 1` is read twice and is not cheap, so it is what would be owed a column.
+ query = (
+ "SELECT o.a FROM t_out_58626 o WHERE "
+ "(SELECT count(*) FROM t_in_58626 i WHERE i.a = o.a "
+ "HAVING sq_test_58626(count(*) + 1) > 0) = 0"
+ )
+ with self.sql_conf(legacy):
+ square = UserDefinedFunction(square_fn, LongType())
+ self.assertTrue(square.transpiled)
+ with (
+ self.temp_func("sq_test_58626"),
+ self.temp_view("t_out_58626"),
+ self.temp_view("t_in_58626"),
+ ):
+ self.spark.udf.register("sq_test_58626", square)
+ self.spark.range(0, 5).selectExpr("id as a").createOrReplaceTempView("t_out_58626")
+ self.spark.range(0, 3).selectExpr("id as a").createOrReplaceTempView("t_in_58626")
+ df = self.spark.sql(query)
+ plan = self._optimized_plan(df)
+ self.assertNotIn(self._SHARED_ARG_ALIAS, plan, plan)
+ self.assertEqual(1, self._eval_python_count(df), plan)
+
+ def test_udf_transpile_unused_argument_diverges_under_ansi(self):
+ # SPARK-58626: an argument the body never uses never reaches the option, so nothing
+ # evaluates it, where the interpreted UDF computes every argument column. Under ANSI
+ # that is an error-vs-rows difference rather than merely saved work.
+ from pyspark.sql.functions import col, lit
+
+ first = lambda a, b: a # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ u = UserDefinedFunction(first, LongType())
+ self.assertTrue(u.transpiled)
+ df = self.spark.createDataFrame([(3,)], "x long")
+ rows = df.select(u(col("x"), col("x") / lit(0))).collect()
+ self.assertEqual([3], [r[0] for r in rows])
+ with self.sql_conf(_TRANSPILE_OFF_ANSI_ON):
+ u_i = UserDefinedFunction(first, LongType())
+ df = self.spark.createDataFrame([(3,)], "x long")
+ with self.assertRaises(Exception) as ctx:
+ df.select(u_i(col("x"), col("x") / lit(0))).collect()
+ self.assertIn("DIVIDE_BY_ZERO", str(ctx.exception))
+
def test_udf_transpile_lowers_operators(self):
# Operators lower to Catalyst and match Python: modulo sign-parity,
# non-commutative -/* (parameter order), unary nesting, constant
diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py
index 11531e1b9118..83e85538564c 100644
--- a/python/pyspark/sql/transpile.py
+++ b/python/pyspark/sql/transpile.py
@@ -41,6 +41,51 @@ of its own. Passed straight to ``udf(...)``, wrapped in another call, returned
by another lambda, or sharing a line with a second lambda, nothing in the
source read back says which lambda is the UDF, so it falls back to interpreted
Python rather than risk the wrong body.
+
+A lowered UDF computes each argument it uses once per row (SPARK-58626): the
+argument becomes a column on the operator's input, and however many times the
+body reads that parameter it reads that one column. Two parameters bound to the
+same deterministic argument share the column, so ``f(a + 1, a + 1)`` computes
+``a + 1`` once.
+
+A call inside a higher-order function's lambda is never lowered: Spark already
+applies a Python UDF over the whole array there, and that path is left to it.
+Otherwise a few positions have nowhere to put the column -- under a ``groupBy``,
+in a join condition, in a command such as ``DELETE FROM``, or a draw anywhere
+above a join, where the column would cost the join its condition -- and there a
+body reading the parameter more than once stays an interpreted Python UDF, which
+computes its inputs once, rather than being lowered into an argument evaluated
+per read. One evaluation per parameter per row either way; ``f(rand(), rand())``
+is two parameters and so still two draws::
+
+ body = lambda x: x if x > 0.5 else 0.0
+ clamp = udf(body, "double")
+ df.select(clamp(rand())) # never a value the body's own condition rejects
+
+Where the column stands it is computed for every row reaching the operator, which
+makes the argument eager even where the call sits in a branch that may not run:
+under ANSI ``when(cond, f(a / b))`` can raise on a row with ``b = 0`` and ``cond``
+false. That is what the interpreted Python UDF does too, evaluating its inputs in
+a projection feeding the worker, below the conditional -- unlike a bare Catalyst
+``when``, which is lazy. Whether such an error surfaces is not a guarantee in
+either direction: an argument left inline, or inlined again by a later rule, is
+lazy once more.
+
+Two arguments no projection can hold count as those positions too: one that is
+itself an aggregate (``f(sum(a))``), and one reading an outer query's column when
+correlated-subquery decorrelation is disabled.
+
+An argument read once gets no column, since one read is one evaluation anyway,
+and neither does one as cheap to repeat as to read -- a bare column or a literal.
+Anything more, arithmetic included, is either computed once or left to
+interpreted Python; a repeated ``a + 1`` gets a column where one fits and stops
+the UDF being lowered where one does not.
+
+That column is what we emit, not what necessarily runs: later optimizer rules may
+inline a deterministic one again where that is faster, as they may for any other
+expression. A draw is never inlined, because those rules check determinism.
+
+An argument the body never reads is not computed at all.
"""
import ast
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptions.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptions.scala
index 188b3e5ae22d..d9b07093fb89 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptions.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveTranspiledPythonUDFOptions.scala
@@ -17,7 +17,7 @@
package org.apache.spark.sql.catalyst.analysis
-import org.apache.spark.sql.catalyst.expressions.TranspiledPythonUDF
+import org.apache.spark.sql.catalyst.expressions.{TranspiledPythonUDF, TranspiledUDFParameter}
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.trees.TreePattern.TRANSPILED_PYTHON_UDF
@@ -51,13 +51,26 @@ object ResolveTranspiledPythonUDFOptions extends Rule[LogicalPlan] {
// Bottom-up so a nested TranspiledPythonUDF (a transpiled UDF feeding another) is pruned
// -- and thus resolved -- before its parent's input types are inspected.
op.transformExpressionsUpWithPruning(_.containsPattern(TRANSPILED_PYTHON_UDF)) {
- case t: TranspiledPythonUDF
- if t.optionInputCategories.nonEmpty && t.pythonUDFExpr.childrenResolved =>
- val argTypes = t.pythonUDFExpr.children.map(_.dataType)
- val kept = t.transpiledOptions.zip(t.optionInputCategories).collect {
- case (option, categories) if optionMatchesTypes(categories, argTypes) => option
+ // The second half of the guard stops this firing every later iteration: categories
+ // cleared and options resolved means there's nothing to do. It would still converge
+ // without it, but it would re-walk every option each time round.
+ case t: TranspiledPythonUDF if t.arguments.forall(_.resolved) &&
+ (t.optionInputCategories.nonEmpty || !t.transpiledOptions.forall(_.resolved)) =>
+ val args = t.arguments
+ val pruned = if (t.optionInputCategories.isEmpty) {
+ t
+ } else {
+ val argTypes = args.map(_.dataType)
+ val kept = t.transpiledOptions.zip(t.optionInputCategories).collect {
+ case (option, categories) if optionMatchesTypes(categories, argTypes) => option
+ }
+ t.copy(transpiledOptions = kept, optionInputCategories = Nil)
}
- t.copy(transpiledOptions = kept, optionInputCategories = Nil)
+ // Type each `_udf_param_N` reference from the argument it stands for. The options are
+ // unresolved until this runs, so the analyzer comes back after and coerces their
+ // bodies like anything else.
+ pruned.copy(transpiledOptions =
+ pruned.transpiledOptions.map(TranspiledUDFParameter.resolveTypes(_, args)))
}
}
}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
index 48da53ccf7c2..f8dff185ac56 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
@@ -17,17 +17,20 @@
package org.apache.spark.sql.catalyst.expressions
+import scala.collection.mutable
+
import org.apache.spark.SparkException.internalError
import org.apache.spark.api.python.{PythonEvalType, PythonFunction}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.analysis.UnresolvedException
-import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateFunction
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateFunction}
import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode}
import org.apache.spark.sql.catalyst.trees.TreePattern.{PYTHON_UDF, TRANSPILED_PYTHON_UDF,
- TreePattern}
+ TRANSPILED_UDF_PARAMETER, TreePattern}
import org.apache.spark.sql.catalyst.util.toPrettySQL
import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors}
import org.apache.spark.sql.types._
+import org.apache.spark.util.Utils
/**
* Helper functions for [[PythonUDF]]
@@ -91,6 +94,109 @@ trait PythonFuncExpression extends NonSQLExpression with UserDefinedExpression {
}
+/**
+ * Stands in for a `_udf_param_N` placeholder in a transpiled option: a reference to the call's
+ * `index`th argument (SPARK-58626).
+ *
+ * A reference and not a copy, so the argument stays put in [[TranspiledPythonUDF.arguments]] and
+ * `ConvertToCatalyst` decides per call whether to compute it once in a Project below the operator.
+ * [[TranspiledUDFParameter.substitute]] copies the argument to each use site only where
+ * repeating it is as cheap as a column read -- a bare column or a literal. An argument worth
+ * more than that, a draw or a regex, is either computed once or not transpiled at all:
+ * `ConvertToCatalyst` keeps the interpreted Python UDF, which evaluates its inputs once wherever
+ * it sits. `transpile.py` says so for UDF authors.
+ */
+case class TranspiledUDFParameter(
+ index: Int,
+ paramType: Option[DataType] = None,
+ paramNullable: Boolean = true)
+ extends LeafExpression with Unevaluable {
+
+ final override val nodePatterns: Seq[TreePattern] = Seq(TRANSPILED_UDF_PARAMETER)
+
+ // Unresolved until ResolveTranspiledPythonUDFOptions reads the type off the bound argument. That
+ // keeps the whole TranspiledPythonUDF unresolved, which is what brings the analyzer back to
+ // coerce the option body once the types are in.
+ override lazy val resolved: Boolean = paramType.isDefined
+
+ override def dataType: DataType =
+ paramType.getOrElse(throw new UnresolvedException("dataType"))
+
+ override def nullable: Boolean = paramNullable
+}
+
+
+object TranspiledUDFParameter {
+
+ /**
+ * Every argument index an option reads, in tree order, one entry per read. Callers want both the
+ * set and the counts: a parameter read once is evaluated once wherever it sits, so how often it
+ * shows up is what decides whether a column buys anything.
+ */
+ def referencedIndexes(option: Expression): Seq[Int] = {
+ val indexes = mutable.ArrayBuffer.empty[Int]
+ mapRefs(option) { p => indexes += p.index; p }
+ indexes.toSeq
+ }
+
+ /**
+ * Points every reference in `option` at whatever `replacement` gives for its index.
+ *
+ * Never looks at what it splices in: an argument can hold an *enclosing* call's reference, and
+ * walking back into the replacement would find it, splice again, and run out of stack.
+ */
+ def substitute(option: Expression, replacement: Int => Expression): Expression =
+ mapRefs(option) { p =>
+ val substituted = replacement(p.index)
+ // We read a reference's type off its argument once, in analysis, and never look again, so a
+ // later rule retyping the argument would leave this stale with the body already coerced
+ // against the old type. Nothing does that today; shout if that changes. Nullability is NOT
+ // checked -- `sameType` ignores it, and `UpdateAttributeNullability` runs in a later batch
+ // and legitimately does change it. Harmless, since ConvertToCatalyst substitutes the real
+ // argument and the final plan takes its nullability from that.
+ if (Utils.isTesting && p.paramType.isDefined && substituted.resolved) {
+ assert(p.dataType.sameType(substituted.dataType),
+ s"Parameter ${p.index} was typed ${p.dataType} but its argument is " +
+ s"${substituted.dataType}")
+ }
+ substituted
+ }
+
+ /**
+ * Fills in each reference's type from the bound arguments, once those are resolved.
+ */
+ def resolveTypes(option: Expression, arguments: Seq[Expression]): Expression =
+ mapRefs(option) {
+ case p if p.paramType.isEmpty =>
+ // Only a hand-built option gets here out of range; the builder bounds-checks what it emits.
+ // Left untyped the node never resolves and CheckAnalysis blames an internal error, which
+ // tells nobody anything.
+ if (p.index < 0 || p.index >= arguments.length) {
+ throw QueryCompilationErrors.invalidUDFParameterPlaceholderIndex(
+ p.index, arguments.length)
+ }
+ val arg = arguments(p.index)
+ p.copy(paramType = Some(arg.dataType), paramNullable = arg.nullable)
+ case p => p
+ }
+
+ /**
+ * Walks this call's references. A nested call's options are skipped -- those indexes count
+ * against *its* arguments -- and the walk continues through
+ * [[TranspiledPythonUDF.pythonUDFExpr]], whose children are in our index space and may hold
+ * our references.
+ */
+ private def mapRefs(
+ option: Expression)(f: TranspiledUDFParameter => Expression): Expression = option match {
+ case _ if !option.containsPattern(TRANSPILED_UDF_PARAMETER) => option
+ case p: TranspiledUDFParameter => f(p)
+ case nested: TranspiledPythonUDF =>
+ nested.withNewChildren(mapRefs(nested.pythonUDFExpr)(f) +: nested.transpiledOptions)
+ case other => other.mapChildren(mapRefs(_)(f))
+ }
+}
+
+
case class TranspiledPythonUDF(
name: String,
pythonUDFExpr: Expression,
@@ -115,14 +221,25 @@ case class TranspiledPythonUDF(
copy(pythonUDFExpr = newChildren.head, transpiledOptions = newChildren.tail.toList)
final override val nodePatterns: Seq[TreePattern] = Seq(TRANSPILED_PYTHON_UDF)
- // True when every direct input to pythonUDFExpr is a plain PythonUDF (not a
- // TranspiledPythonUDF). Used to decide whether to preserve the UDF batch pipeline
- // rather than inserting a Catalyst node in the middle of a Python UDF chain.
+ /**
+ * The call's bound arguments, which a [[TranspiledUDFParameter]] refers to by position.
+ *
+ * Not `pythonUDFExpr.children`, which is the trap: `fromUDFExpr` wraps an aggregate UDF in an
+ * [[AggregateExpression]], whose children are the aggregate function and its filter. Read them
+ * off the wrapper and you type parameter 0 from the UDAF's return type, then substitute the whole
+ * aggregate function in for it.
+ */
+ def arguments: Seq[Expression] = pythonUDFExpr match {
+ case agg: AggregateExpression => agg.aggregateFunction.children
+ case other => other.children
+ }
+
+ // True when every argument is a plain PythonUDF (not a TranspiledPythonUDF). Used to decide
+ // whether to preserve the UDF batch pipeline rather than inserting a Catalyst node in the middle
+ // of a Python UDF chain. Reads `arguments`, not `pythonUDFExpr.children`: for an aggregate UDF
+ // the latter is the aggregate function and its filter, so this said false for every UDAF.
def hasOnlyPythonUDFInputs: Boolean =
- pythonUDFExpr.children.nonEmpty &&
- pythonUDFExpr.children.forall {
- _.isInstanceOf[PythonUDF]
- }
+ arguments.nonEmpty && arguments.forall(_.isInstanceOf[PythonUDF])
}
/**
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
index 6a607ea01aff..cf4bae09abba 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala
@@ -1035,15 +1035,276 @@ object ConvertToCatalyst extends Rule[LogicalPlan] {
// here could otherwise escape and reach execution un-stripped.
plan.transformDownWithSubqueriesAndPruning(
_.containsPattern(TRANSPILED_PYTHON_UDF), ruleId) {
- case p => p.transformExpressionsWithPruning(_.containsPattern(TRANSPILED_PYTHON_UDF)) {
- case s: TranspiledPythonUDF => applyExpr(s, parentIsUdf = false)
+ // A fast path, not a correctness guard -- `convertOperator` hands back an operator it finds
+ // nothing to do in. The pruning above only skips subtrees, so every Filter and Sort between
+ // the root and a call lands here. The bit is the cheap reject, cached per node; the walk then
+ // also rejects an operator whose only call sits in a subquery, since a SubqueryExpression
+ // reports its inner plan's bits while `Expression.exists` does not descend into the plan.
+ case p if !p.expressions.exists { e =>
+ e.containsPattern(TRANSPILED_PYTHON_UDF) && e.exists(_.isInstanceOf[TranspiledPythonUDF])
+ } => p
+
+ case p => convertOperator(p)
+ }
+ }
+
+ /**
+ * Whether the call still has its arguments. An analyzer rule may replace the Python UDF with an
+ * attribute -- `PullOutNondeterministic` does, for a nondeterministic call in an operator that
+ * cannot hold one -- and then [[TranspiledPythonUDF.arguments]] is empty while the option still
+ * refers to them by position. There is nothing to substitute, so the attribute stands.
+ */
+ private def callIsIntact(s: TranspiledPythonUDF): Boolean = s.pythonUDFExpr match {
+ case _: PythonFuncExpression => true
+ case agg: AggregateExpression => agg.aggregateFunction.isInstanceOf[PythonFuncExpression]
+ case _ => false
+ }
+
+ /**
+ * Converts the calls in one operator, putting an argument the body reads more than once in a
+ * column on the operator's child so we compute it once a row instead of once a read
+ * (SPARK-58626). Three operators can't hold that Project, and a call under one of them that owes
+ * an evaluation keeps the interpreted Python UDF instead:
+ *
+ * - More than one child: we'd have to pick a side.
+ * - A Command: the Project hides the relation DataSourceV2Strategy looks for, and
+ * `DELETE FROM t WHERE udf(a + 1) > 0` becomes an internal error.
+ * - An Aggregate: a result expression no aggregate function wraps has to be built from the
+ * grouping expressions, and a column is not one, so `groupBy(a + 1).agg(udf(a + 1))` would
+ * fail MISSING_AGGREGATION. Splitting the Aggregate to make room was tried and reverted: it
+ * reshapes correlated subqueries, which broke `splitSubquery` on a HAVING and flipped
+ * COUNT-bug handling, for a column in one uncommon shape.
+ * - Anything whose child is an Aggregate, for the second half of that same trap:
+ * `splitSubquery` finds a subquery's Aggregate by matching `Filter(_, Aggregate)`, so a
+ * Project in between hides it and the COUNT bug goes unhandled. Measured on
+ * `HAVING udf(count(*) + 1) > 0`: the outer rows that match nothing came back dropped where
+ * the same plan without the column returns them. It runs before the pushdown that would
+ * collapse the Project away, so being invisible later is no help.
+ *
+ * MergeRows is on neither list, on purpose. Its instructions are first-match-wins, so a column
+ * runs `WHEN ... AND udf(a / b) > 0` for every row of the join -- but so does the interpreted
+ * UDF, which computes its inputs below MergeRows and raises on an unmatched row under ANSI today.
+ * Turning it down there is what would diverge.
+ */
+ private def convertOperator(p: LogicalPlan): LogicalPlan =
+ // A column widens the operator's child, and we only learn whether this operator's output is a
+ // widening of its child's after building it. When it isn't, convert again without columns
+ // rather than fail the query: `preEvaluate` says yes to every single-child operator that is not
+ // a Command or an Aggregate, a longer list than anyone has thought about -- Window, Generate,
+ // Expand, CollectMetrics, ScriptTransformation, whatever an extension adds.
+ convertOperator(p, withColumns = true).getOrElse(convertOperator(p, withColumns = false).get)
+
+ private def convertOperator(p: LogicalPlan, withColumns: Boolean): Option[LogicalPlan] = {
+ val extraColumns = mutable.ArrayBuffer.empty[Alias]
+ val columnByKey = mutable.HashMap.empty[ShareKey, Alias]
+ // Per call *node*, by identity: an operator can hold one node in two expression slots -- Window
+ // copies its spec into every WindowSpecDefinition, MergeRows has five -- and each slot is
+ // visited on its own. Minting an id per visit gave one call two columns, so two draws. A plain
+ // counter, not an ExprId: nothing outside the sharing key ever reads it, and `newExprId` would
+ // bump the global counter for every call we merely looked at.
+ var nextCallId = 0L
+ val callIds = new java.util.IdentityHashMap[TranspiledPythonUDF, java.lang.Long]()
+ val preEvaluate: Option[PreEvaluate] =
+ if (withColumns && p.children.length == 1 && !p.isInstanceOf[Command] &&
+ !p.isInstanceOf[Aggregate] && !p.children.head.isInstanceOf[Aggregate]) {
+ val child = p.children.head
+ // Once per operator, not once per call: an operator can hold several, and this walks the
+ // whole child plan. See `canPlaceColumn`, which is the only reader.
+ lazy val joinBelow = child.exists(_.isInstanceOf[Join])
+ Some((call, args, option) => {
+ val callId = callIds.computeIfAbsent(call, _ => { nextCallId += 1; nextCallId })
+ def keyOf(index: Int): ShareKey = shareKey(args(index), callId, index)
+ val reads = TranspiledUDFParameter.referencedIndexes(option)
+ val readsPerKey = reads.groupBy(keyOf).map { case (k, is) => k -> is.length }
+ val owed = mustPreEvaluate(args, option)
+ // Read more than once under one key and not cheap to repeat. That covers everything we
+ // owe (`owed` counts reads of one index; a key only ever merges indexes together, so its
+ // count is at least as high) plus what is merely worth sharing: two parameters holding
+ // equal arguments, each read once. Not every repeat earns a column -- a pointless Project
+ // sticks around under anything CollapseProject can't merge through, and stops
+ // SpecialLimits matching `Project(_, Sort(...))` for TakeOrderedAndProjectExec.
+ val shared = reads.distinct.filter { index =>
+ readsPerKey(keyOf(index)) > 1 && !cheapToRepeat(args(index)) &&
+ canPlaceColumn(args(index), child, joinBelow)
+ }
+ if (!owed.forall(shared.contains)) {
+ // Something we owe one evaluation has nowhere to go, so hand the call back to Python,
+ // which evaluates each of its inputs once wherever it sits.
+ None
+ } else if (shared.isEmpty) {
+ Some(TranspiledUDFParameter.substitute(option, args))
+ } else {
+ val byIndex = shared.map { index =>
+ // The ExprId goes in the name: numbering per operator would give two operators
+ // each a `_udf_param_0`, which makes a plan string ambiguous. Same reason
+ // RewriteWithExpression has USE_COMMON_EXPR_ID_FOR_ALIAS.
+ index -> columnByKey.getOrElseUpdate(keyOf(index), {
+ val exprId = NamedExpression.newExprId
+ val alias = Alias(args(index), s"_udf_param_${exprId.id}")(exprId)
+ extraColumns += alias
+ alias
+ })
+ }.toMap
+ Some(TranspiledUDFParameter.substitute(
+ option, i => byIndex.get(i).map(_.toAttribute).getOrElse(args(i))))
+ }
+ })
+ } else {
+ None
+ }
+ // Enter each expression at its root rather than at the first TranspiledPythonUDF in it, so
+ // that `applyExpr` sees every node on the way down and can thread `parentIsUdf`. Starting
+ // at the node itself would report no enclosing UDF even when a PythonUDF wraps it, and the
+ // batch pipeline the flag exists to protect would be split anyway.
+ val converted = p.mapExpressions {
+ case e if e.containsPattern(TRANSPILED_PYTHON_UDF) =>
+ applyExpr(e, parentIsUdf = false, preEvaluate, inLambda = false)
+ case e => e
+ }
+ // A call turned down can orphan a column registered for an argument of its own. Drop those:
+ // an unread column still computes, and under ANSI can raise on a row nothing looked at.
+ val read = AttributeSet(converted.expressions.flatMap(_.references))
+ val columns = extraColumns.filter(a => read.contains(a.toAttribute)).toSeq
+ if (columns.isEmpty) {
+ Some(converted)
+ } else {
+ val child = p.children.head
+ val widened = converted.withNewChildren(Seq(Project(child.output ++ columns, child)))
+ // A Filter passes its child's output straight up, so it carries our columns now. Project them
+ // back off to keep the schema. The checks are the two RewriteWithExpression asserts, as a
+ // condition: we only ever append, so an operator whose output shrinks or loses an attribute
+ // is one we should not have offered a column to.
+ if (p.output.length > widened.output.length || !p.outputSet.subsetOf(widened.outputSet)) {
+ None
+ } else if (p.output.length < widened.output.length) {
+ Some(Project(p.output, widened))
+ } else {
+ Some(widened)
}
}
}
- def applyExpr(expression: Expression, parentIsUdf: Boolean = false): Expression = {
+ /**
+ * What decides whether two parameters share a column: an equal argument of equal nullability, or
+ * nothing -- a nondeterministic argument keys on the call it belongs to and its position in it.
+ */
+ private sealed trait ShareKey
+ // Nullability rides along because `canonicalized` drops it: an AttributeReference canonicalizes
+ // to `AttributeReference("none", dataType)(exprId)`, so `a#5` nullable and `a#5` not would
+ // otherwise share one column, and a body coerced against the nullable one would read a column
+ // declared without nulls -- enough for NullPropagation to fold an `IsNull` to false.
+ private case class DeterministicKey(canonical: Expression, nullable: Boolean) extends ShareKey
+ // A call id, not just the parameter index: `f(rand())` and `g(rand())` in one operator would
+ // otherwise collide on one column. `f(rand(), rand())` owes two draws, and so do two calls that
+ // each pass their own `rand()`.
+ private case class NondeterministicKey(callId: Long, index: Int) extends ShareKey
+
+ // One key per distinct deterministic argument in the operator: `f(a + 1, a + 1)` computes
+ // `a + 1` once, and two calls that each read it more than once share the column. The read
+ // count is still per call -- a call reading it once keeps its own copy, one evaluation.
+ private def shareKey(arg: Expression, callId: Long, index: Int): ShareKey =
+ if (arg.deterministic) DeterministicKey(arg.canonicalized, arg.nullable)
+ else NondeterministicKey(callId, index)
+
+ /**
+ * Whether `arg` can sit in a Project on `child`. Four ways it can't:
+ *
+ * - It reads something the child doesn't output. `f(g(x))` is the one you hit: `f`'s
+ * argument is `g`'s column, which this Project is still defining.
+ * - It's an aggregate, window function or generator. Only PlanHelper can say, and
+ * only about a Project, so this one is asked with a throwaway alias.
+ * - It has an OuterReference and decorrelation is off, where the fallback rewrites
+ * Filters only and would strand it. Both confs, since EXISTS and IN switch on
+ * their own even while the global one is on.
+ * - It still holds an enclosing call's reference, which only that call can
+ * substitute, and only while the reference is in the option body.
+ *
+ * A nondeterministic column anywhere above a Join takes the join's condition down with
+ * it: PushPredicateThroughNonJoin pushes nothing through a Project holding a
+ * nondeterministic field, so a Filter above it never reaches the Join and
+ * `where t1.a = t2.a and udf(rand()) < 2` plans as a cartesian product. Below anywhere,
+ * not just at the child -- one `select` in between is enough, and the Filter can sit
+ * above whatever we are rewriting.
+ *
+ * `joinBelow` is by name so a deterministic argument never pays for the walk behind it.
+ */
+ private def canPlaceColumn(arg: Expression, child: LogicalPlan, joinBelow: => Boolean): Boolean =
+ (arg.deterministic || !joinBelow) &&
+ arg.references.subsetOf(child.outputSet) &&
+ PlanHelper.specialExpressionsInUnsupportedOperator(
+ // A fixed id, not `newExprId`: this alias is thrown away with the Project and
+ // PlanHelper never reads it, so minting one would bump the global counter and
+ // shift the ExprIds in every later plan string by however many arguments we
+ // merely considered.
+ Project(Seq(Alias(arg, "_udf_param")(ExprId(0))), child)).isEmpty &&
+ ((conf.decorrelateInnerQueryEnabled &&
+ conf.decorrelateInnerQueryEnabledForExistsIn) ||
+ !arg.containsPattern(OUTER_REFERENCE)) &&
+ // Bit first, then a walk, as in `apply`: a SubqueryExpression reports its inner plan's bits
+ // but `exists` does not descend into it, so an argument holding a subquery that merely
+ // *contains* a call would fail the bit test over a reference that is not ours to substitute.
+ !(arg.containsPattern(TRANSPILED_UDF_PARAMETER) &&
+ arg.exists(_.isInstanceOf[TranspiledUDFParameter]))
+
+ /**
+ * Notes the columns to add to the operator's child, and points the option's refs at them. None
+ * means we can't compute an evaluation the body is owed here, and the caller keeps the Python
+ * UDF. The call comes along so that two visits to one node share its columns.
+ */
+ private type PreEvaluate =
+ (TranspiledPythonUDF, Seq[Expression], Expression) => Option[Expression]
+
+ /**
+ * The parameters this call owes one evaluation: read more than once, with an argument that isn't
+ * cheap to repeat. Where no column can hold one we don't transpile at all. Per parameter:
+ * `f(rand(), rand())` owes two draws, `f(rand())` owes one however often it is read.
+ */
+ private def mustPreEvaluate(args: Seq[Expression], option: Expression): Set[Int] =
+ TranspiledUDFParameter.referencedIndexes(option).groupBy(identity).collect {
+ case (index, reads) if reads.length > 1 && !cheapToRepeat(args(index)) => index
+ }.toSet
+
+ /**
+ * Whether reading an argument twice is as good as reading a column that holds it once.
+ *
+ * A nondeterministic one never is, whatever `isCheap` says of it: two copies of `rand()` are two
+ * draws, so the body sees two values for one parameter. A deterministic one is only work, but
+ * work counts, so anything `isCheap` won't vouch for -- a regex, `a + 1` -- is owed one too.
+ */
+ private def cheapToRepeat(arg: Expression): Boolean = arg match {
+ // An enclosing call's reference computes nothing of its own. What it stands for may well be
+ // expensive, but that is the enclosing call's argument, and these copies are reads of it.
+ case _: TranspiledUDFParameter => true
+ case _ => arg.deterministic && CollapseProject.isCheap(arg)
+ }
+
+ /**
+ * Converts the transpiled calls in one expression, leaving every argument at its use sites -- and
+ * so keeping the Python UDF for a call that owes one of them a single evaluation.
+ *
+ * This is what a custom transpiler's own ConvertToX gets, and it comes with no column placement
+ * at all: whether a Project fits below the operator is the operator's business, and only `apply`
+ * above, walking the plan, can know. A custom rule that wants an argument computed once has to
+ * place the column itself.
+ */
+ def applyExpr(expression: Expression, parentIsUdf: Boolean = false): Expression =
+ applyExpr(expression, parentIsUdf, None, inLambda = false)
+
+ private def applyExpr(
+ expression: Expression,
+ parentIsUdf: Boolean,
+ preEvaluate: Option[PreEvaluate],
+ inLambda: Boolean): Expression = {
+ def recurse(e: Expression, parentIsUdf: Boolean): Expression =
+ applyExpr(e, parentIsUdf, preEvaluate, inLambda)
expression match {
+ // Nothing to convert below here, so skip the subtree rather than walking it node by node.
+ case other if !other.containsPattern(TRANSPILED_PYTHON_UDF) => other
+
case s: TranspiledPythonUDF =>
+ // Every branch that turns the call down leaves the Python UDF in place and keeps walking
+ // its arguments, which may hold transpilable calls of their own.
+ def keepPython: Expression = s.pythonUDFExpr.mapChildren(recurse(_, parentIsUdf = true))
// We _shouldn't_ have these nodes if ANSI is not enabled or transpilation is disabled
// but if someone changed it while running we'll want to strip the nodes out.
if (!conf.getConf(SQLConf.ANSI_ENABLED)) {
@@ -1051,12 +1312,25 @@ object ConvertToCatalyst extends Rule[LogicalPlan] {
log"${MDC(LogKeys.CONFIG, SQLConf.ANSI_ENABLED.key)} is disabled. The transpiler " +
log"targets ANSI semantics and refuses to rewrite plans under non-ANSI mode. " +
log"Enable ANSI or disable transpilation to silence this warning.")
- s.pythonUDFExpr.mapChildren(applyExpr(_, parentIsUdf = true))
+ keepPython
} else if (!conf.getConf(SQLConf.ATTEMPT_TRANSPILATION_OF_PYTHON_UDFS)) {
logWarning(log"Skipping Python UDF transpilation: " +
log"${MDC(LogKeys.CONFIG, SQLConf.ATTEMPT_TRANSPILATION_OF_PYTHON_UDFS.key)} " +
log"is disabled but we still got TranspiledPythonUDFs in our plan.")
- s.pythonUDFExpr.mapChildren(applyExpr(_, parentIsUdf = true))
+ keepPython
+ } else if (!callIsIntact(s)) {
+ // PullOutNondeterministic moves a nondeterministic call into a projection below and
+ // leaves an attribute here, so `arguments` is empty while the option still refers to
+ // them. Keep the attribute: the call is computed once below us, which is the promise,
+ // just not in Catalyst. `orderBy(u(rand()))` and `repartition(2, u(rand()))` land here.
+ keepPython
+ } else if (inLambda ||
+ s.transpiledOptions.headOption.exists(_.containsPattern(LAMBDA_FUNCTION))) {
+ // Lambdas are out of scope for lowering: a higher-order function's lambda already has a
+ // story for a Python UDF in ExtractPythonUDFFromLambda, which applies it over the whole
+ // array. Both shapes are left to it -- a call inside a lambda, and an option body that
+ // builds one.
+ keepPython
} else if (!parentIsUdf || !s.hasOnlyPythonUDFInputs) {
// Walk the full list of transpiled options and pick the first one,
// falling back to the original Python UDF if none are available.
@@ -1070,26 +1344,59 @@ object ConvertToCatalyst extends Rule[LogicalPlan] {
// option's dataType already matches; a custom transpiler MUST do the
// same (or insert its own Cast), or it will silently change the output
// schema.
- val firstEvaluable = s.transpiledOptions.headOption
- firstEvaluable match {
+ s.transpiledOptions.headOption match {
case None =>
- s.pythonUDFExpr.mapChildren(applyExpr(_, parentIsUdf = true))
+ keepPython
case Some(catalystExpr) =>
- // Recursively apply to the children first because we may use them as inputs in parent
- catalystExpr.mapChildren(applyExpr(_, parentIsUdf = false))
+ // Arguments live on the Python UDF, not in the option, so convert them there -- one
+ // could be a transpilable call itself. Recurse on the option's root too, not just its
+ // children, or a nested call sitting at the root survives and throws at execution.
+ //
+ // Only the arguments the body actually reads. Converting an unread one would register
+ // a column for its nested call that nothing ever reads, and under ANSI that column
+ // computes -- and can raise on -- an argument `transpile.py` promises is never
+ // computed at all.
+ val reads = TranspiledUDFParameter.referencedIndexes(catalystExpr)
+ val read = reads.toSet
+ // Over `reads`, not `read`: a Set names whichever bad index it yields first, so the
+ // error moved around between JVMs.
+ reads.find(index => index < 0 || index >= s.arguments.length).foreach { index =>
+ // Only a hand-built option can be out of range: the builder bounds-checks what it
+ // emits, and a pre-typed reference slips past the analyzer's check too.
+ throw QueryCompilationErrors.invalidUDFParameterPlaceholderIndex(
+ index, s.arguments.length)
+ }
+ val args = s.arguments.zipWithIndex.map {
+ case (arg, i) if read.contains(i) => recurse(arg, parentIsUdf = false)
+ case (arg, _) => arg
+ }
+ val option = recurse(catalystExpr, parentIsUdf = false)
+ val substituted = preEvaluate match {
+ case Some(f) => f(s, args, option)
+ // No column can go here at all, so an evaluation the body is owed can only be kept
+ // by the Python UDF, which computes its inputs once in a projection of its own.
+ case None =>
+ Option.when(mustPreEvaluate(args, option).isEmpty)(
+ TranspiledUDFParameter.substitute(option, args))
+ }
+ // None from either arm means we can't keep that promise here, so don't transpile.
+ substituted.getOrElse(keepPython)
}
} else {
// We should avoid converting a UDF node where that could break pipelining.
// For example: (UDF -> UDF -> UDF) is often cheaper than UDF -> Catalyst -> UDF.
- s.pythonUDFExpr.mapChildren(applyExpr(_, parentIsUdf = true))
+ keepPython
}
case _ =>
// Not a TranspiledPythonUDF: recurse down, telling the children whether
// this node is itself a scalar Python UDF so a transpiled child can
// preserve the UDF batch pipeline (e.g. an outer UDF that could not be
// transpiled wrapping one that could).
- expression.mapChildren(
- applyExpr(_, parentIsUdf = isScalarPythonUDF(expression)))
+ //
+ // A conditional branch is not treated like a lambda, on purpose: the column stands, which
+ // makes the argument eager, and the interpreted UDF is eager there too. See `transpile.py`.
+ expression.mapChildren(applyExpr(_, parentIsUdf = isScalarPythonUDF(expression),
+ preEvaluate, inLambda || expression.isInstanceOf[LambdaFunction]))
}
}
}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala
index 39b9be7d5410..64b9e2223a6f 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/trees/TreePatterns.scala
@@ -100,6 +100,7 @@ object TreePattern extends Enumeration {
val TIME_WINDOW: Value = Value
val TIME_ZONE_AWARE_EXPRESSION: Value = Value
val TRANSPILED_PYTHON_UDF: Value = Value
+ val TRANSPILED_UDF_PARAMETER: Value = Value
val TRUE_OR_FALSE_LITERAL: Value = Value
val USER_DEFINED_AGGREGATION: Value = Value
val VARIANT_GET: Value = Value
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConvertToCatalystSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConvertToCatalystSuite.scala
index 23baf25a5897..8220cba685ca 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConvertToCatalystSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/ConvertToCatalystSuite.scala
@@ -18,20 +18,24 @@
package org.apache.spark.sql.catalyst.optimizer
import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.AnalysisException
import org.apache.spark.sql.catalyst.dsl.expressions._
import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Count}
-import org.apache.spark.sql.catalyst.plans.PlanTest
-import org.apache.spark.sql.catalyst.plans.logical.{Filter, LocalRelation, Project}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Count, Sum}
+import org.apache.spark.sql.catalyst.plans.{Inner, PlanTest}
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, DeleteFromTable, Filter, Join, JoinHint, LocalRelation, LogicalPlan, LogicalPlanIntegrity, Project, Sort}
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.types.{BooleanType, IntegerType, LongType}
+import org.apache.spark.sql.types.{ArrayType, BooleanType, IntegerType, LongType}
/**
* Unit tests for the ConvertToCatalyst optimizer rule, which rewrites
* TranspiledPythonUDF nodes to their Catalyst equivalents.
*
- * These tests exercise the rule directly via applyExpr rather than running the
- * full optimizer pipeline, which means no JVM/Python bridge is required.
+ * Nodes are built by hand, so no JVM/Python bridge is needed.
+ *
+ * Which entry point a test uses matters. `applyExpr` takes a bare expression with pre-evaluation
+ * off, so it only covers the fallback of leaving arguments at their use sites. `convert` goes
+ * through `apply`, the only caller that turns pre-evaluation on, so column tests need it.
*/
class ConvertToCatalystSuite extends PlanTest {
@@ -56,6 +60,67 @@ class ConvertToCatalystSuite extends PlanTest {
private val catalystExpr: Expression = Add(attrA, Literal(4L))
+ // A reference to the UDF's `index`th argument, typed the way the analyzer would type it.
+ private def pref(index: Int): TranspiledUDFParameter =
+ TranspiledUDFParameter(index, Some(LongType))
+
+ // A nondeterministic argument, cast so it types like every other argument here. Nondeterministic
+ // because Rand is: an expression is only deterministic if all of its children are.
+ private def draw(seed: Long = 1L): Expression = Cast(Rand(Literal(seed)), LongType)
+
+ /**
+ * The pre-evaluated argument columns ConvertToCatalyst added anywhere in `plan`, subquery plans
+ * included -- the rule converts those too (`transformDownWithSubqueriesAndPruning`), while a
+ * plain `collect` stops at the outer plan and would call a column placed inside one absent.
+ */
+ private def paramColumns(plan: LogicalPlan): Seq[Alias] =
+ plan.collectWithSubqueries { case Project(list, _) => list }.flatten.collect {
+ case a: Alias if a.name.startsWith("_udf_param_") => a
+ }
+
+ /**
+ * Runs the rule, plus the plan checks the optimizer would have run for us -- applying a rule
+ * object directly skips them, and stranding a reference or widening a schema is exactly how this
+ * rewrite would go wrong.
+ */
+ private def convert(plan: LogicalPlan): LogicalPlan = {
+ val converted = ConvertToCatalyst(plan)
+ // What the real Optimizer checks around every rule when `spark.sql.planChangeValidation` is on
+ // (it defaults to `Utils.isTesting`) and what applying a rule object directly skips: dangling
+ // references, duplicate ExprIds, special expressions in the wrong operator, aggregate shape,
+ // nullability and schema. Duplicate ExprIds is the one worth naming -- it is what catches a
+ // rewrite that mints two aliases with one id.
+ //
+ // Dangling references get their own walk, on every plan and per operator, because
+ // `validateNoDanglingReferences` is no help here: a resolved root matches its last case, so its
+ // `collectFirst` returns there and never visits a child -- and the root is never the Project
+ // this rule inserts.
+ converted.foreach {
+ case n if n.resolved && n.children.nonEmpty =>
+ assert(n.missingInput.isEmpty,
+ s"${n.missingInput.mkString(", ")} dangling in ${n.nodeName}:\n$converted")
+ case _ =>
+ }
+ // The rest of it on resolved plans only: the full check asks every output attribute for its
+ // dataType, and a hand-built plan with an unbound lambda hasn't got one.
+ if (plan.resolved) {
+ LogicalPlanIntegrity.validateOptimizedPlan(plan, converted, lightweight = false)
+ .foreach(failure => fail(s"$failure\n$converted"))
+ }
+ // A leftover reference is Unevaluable and throws at execution. Check it here so no test has to
+ // remember to -- with subqueries, since the rule converts those and a plain `exists` would let
+ // one stranded inside a subquery pass.
+ val survivors = converted.collectWithSubqueries {
+ case operator if operator.expressions.exists(_.exists {
+ case _: TranspiledPythonUDF | _: TranspiledUDFParameter => true
+ case _ => false
+ }) => operator.nodeName
+ }
+ assert(survivors.isEmpty,
+ s"A transpiled node survived conversion in ${survivors.mkString(", ")}:\n$converted")
+ converted
+ }
+
// ---- helpers ----
// Both ANSI and ATTEMPT_TRANSPILATION must be true for the transpile path to fire.
@@ -126,6 +191,10 @@ class ConvertToCatalystSuite extends PlanTest {
val zeroPyUDF = PythonUDF("udf", null, LongType, Seq.empty,
PythonEvalType.SQL_BATCHED_UDF, udfDeterministic = true)
assert(!TranspiledPythonUDF("udf", zeroPyUDF, List(Literal(42L))).hasOnlyPythonUDFInputs)
+ // An aggregate UDF's arguments sit under the AggregateExpression, not beside it: reading
+ // pythonUDFExpr.children saw (aggregateFunction, filter) and said false for every UDAF.
+ val aggOverPyUDF = makePyUDAF(innerPyUDF).toAggregateExpression()
+ assert(TranspiledPythonUDF("agg", aggOverPyUDF, List(catalystExpr)).hasOnlyPythonUDFInputs)
}
test("falls back to PythonUDF when ANSI is disabled") {
@@ -188,6 +257,29 @@ class ConvertToCatalystSuite extends PlanTest {
}
}
+ test("computes one column per distinct argument, read by every reference") {
+ transpileOn {
+ // `lambda a, b: a * a + b` over `f(a + 1, a + 1)`: three refs, two parameters, one argument.
+ // Keyed on the argument, so that's one column read three times.
+ val arg = Add(attrA, Literal(1L))
+ val option = Add(Multiply(pref(0), pref(0)), pref(1))
+ val pyUDF = PythonUDF("udf", null, LongType, Seq(arg, arg),
+ PythonEvalType.SQL_BATCHED_UDF, udfDeterministic = true)
+ val tpudf = TranspiledPythonUDF("udf", pyUDF, List(option))
+ val converted = convert(Project(Seq(Alias(tpudf, "v")()), LocalRelation(attrA)))
+ val columns = paramColumns(converted)
+ assert(columns.map(_.child) == Seq(arg),
+ s"Expected one column for the shared argument: $converted")
+ // By exprId, not name -- an alias whose id doesn't match what the body reads is the dangling
+ // attribute bug, and names wouldn't show it.
+ val reads = converted.expressions.head.collect {
+ case r: AttributeReference if r.name.startsWith("_udf_param_") => r.exprId
+ }
+ assert(reads == Seq.fill(3)(columns.head.toAttribute.exprId),
+ s"Expected all three references to read that one column: $converted")
+ }
+ }
+
test("uses pre-coerced transpiledOptions as-is (analysis is responsible for coercion)") {
// The Analyzer coerces transpiledOptions before the optimizer runs, because
// TranspiledPythonUDF.children exposes them to the resolver's generic coercion pass.
@@ -234,4 +326,390 @@ class ConvertToCatalystSuite extends PlanTest {
}
}
}
+
+ test("drops an input the option never uses") {
+ transpileOn {
+ // `lambda a, b: a` over f(a, rand()): substitution dropped b, so nothing evaluates it. The
+ // Python path computes every argument column -- an accepted difference, pinned here.
+ val unused = Rand(Literal(1L))
+ val option = Add(attrA, Literal(1L))
+ val tpudf = TranspiledPythonUDF(
+ "udf",
+ PythonUDF("udf", null, LongType, Seq(attrA, unused),
+ PythonEvalType.SQL_BATCHED_UDF, udfDeterministic = true),
+ List(option))
+ val result = ConvertToCatalyst.applyExpr(tpudf, parentIsUdf = false)
+ assert(result == option, s"Expected the option unchanged, got: $result")
+ assert(!result.exists(_.isInstanceOf[Rand]), s"Unused argument survived in: $result")
+ }
+ }
+
+ test("leaves an argument an aggregate function wraps at each use site") {
+ transpileOn {
+ // Only a custom transpiler emits an option body that is itself an aggregate. Through
+ // `applyExpr`, so this is the fallback path only -- nothing could host such a body anyway.
+ //
+ // A bare column, since anything more is owed one evaluation and would keep the Python UDF
+ // here rather than substitute.
+ val arg = attrA
+ val option = Count(Seq(Multiply(pref(0), pref(0)))).toAggregateExpression()
+ val tpudf = makeTPUDF(makePyUDF(arg), option)
+ val result = ConvertToCatalyst.applyExpr(tpudf, parentIsUdf = false)
+ // Compare the aggregate function, not the AggregateExpression: `toAggregateExpression` mints
+ // a fresh resultId per call, so the wrappers never compare equal.
+ assert(result.asInstanceOf[AggregateExpression].aggregateFunction ==
+ Count(Seq(Multiply(arg, arg))),
+ s"Expected the argument left at both use sites: $result")
+ }
+ }
+
+ test("keeps the Python UDF inside a higher-order function's lambda") {
+ transpileOn {
+ // Lambdas are out of scope for lowering, whatever the argument costs:
+ // ExtractPythonUDFFromLambda already applies a Python UDF over the whole array and this rule
+ // leaves that to it. Through `apply`, since `applyExpr` alone never sees the lambda.
+ Seq(attrA, Divide(attrA, attrA), draw()).foreach { arg =>
+ val call = makeTPUDF(makePyUDF(arg), Add(pref(0), pref(0)))
+ val lambdaVar = NamedLambdaVariable("x", LongType, nullable = false)
+ val arr = AttributeReference("arr", ArrayType(LongType))()
+ val body = ArrayTransform(arr, LambdaFunction(Add(call, lambdaVar), Seq(lambdaVar)))
+ val converted = convert(Project(Seq(Alias(body, "v")()), LocalRelation(attrA, arr)))
+ assert(paramColumns(converted).isEmpty, s"Expected no column for $arg: $converted")
+ assert(converted.expressions.head.collect { case u: PythonUDF => u }.length == 1,
+ s"Expected the Python UDF kept for $arg: $converted")
+ assert(converted.expressions.head.collect { case e if e == arg => e }.length == 1,
+ s"Expected the one evaluation the Python UDF makes for $arg: $converted")
+ }
+ }
+ }
+
+ test("keeps the Python UDF when the option body builds a lambda") {
+ transpileOn {
+ // The other half of the same rule, which only a custom transpiler can reach: an option that
+ // lowers to a higher-order function of its own. One place to reason about lambdas, not two.
+ val param = TranspiledUDFParameter(0, Some(ArrayType(LongType)))
+ val lambdaVar = NamedLambdaVariable("x", LongType, nullable = false)
+ val option = Cast(
+ Size(ArrayTransform(param, LambdaFunction(Add(lambdaVar, Literal(1L)), Seq(lambdaVar)))),
+ LongType)
+ val arr = AttributeReference("arr", ArrayType(LongType))()
+ val call = makeTPUDF(makePyUDF(arr), option)
+ val converted = convert(Project(Seq(Alias(call, "v")()), LocalRelation(arr)))
+ assert(converted.expressions.head.exists(_.isInstanceOf[PythonUDF]),
+ s"Expected the Python UDF kept: $converted")
+ }
+ }
+
+ test("gives one call in two of an operator's slots a single column") {
+ transpileOn {
+ // The same node, twice in one Sort. An operator can do this on its own -- Window copies its
+ // spec into every WindowSpecDefinition, MergeRows has five slots -- and each slot converts
+ // separately, so an id minted per visit gave this call two columns and two draws.
+ val call = makeTPUDF(makePyUDF(draw()), Add(pref(0), pref(0)))
+ val converted = convert(Sort(
+ Seq(SortOrder(call, Ascending), SortOrder(Multiply(call, Literal(2L)), Descending)),
+ global = true,
+ Project(Seq(Alias(attrA, "a")()), LocalRelation(attrA))))
+ assert(paramColumns(converted).length == 1,
+ s"Expected one column for one call: $converted")
+ }
+ }
+
+ test("gives two calls' nondeterministic arguments a column each") {
+ transpileOn {
+ // Two calls in one Project, each reading its own draw twice. Keying a column on the parameter
+ // index alone collided them and put both bodies on one draw.
+ val option = Add(pref(0), pref(0))
+ def call(seed: Long): TranspiledPythonUDF = makeTPUDF(makePyUDF(draw(seed)), option)
+ val converted = convert(
+ Project(Seq(Alias(call(1L), "x")(), Alias(call(2L), "y")()), LocalRelation(attrA)))
+ assert(paramColumns(converted).length == 2, s"Expected one column per call: $converted")
+ }
+ }
+
+ test("keeps the Python UDF under an Aggregate, where no column can go") {
+ transpileOn {
+ // An Aggregate holds no Project of its own -- a result expression no aggregate function wraps
+ // has to be built from the grouping expressions -- so a call owed an evaluation goes back to
+ // Python. Both shapes: inside an aggregate function and beside one.
+ val call = makeTPUDF(makePyUDF(Add(attrA, Literal(1L))), Add(pref(0), pref(0)))
+ val insideSum = Aggregate(Seq(attrA),
+ Seq(Alias(attrA, "a")(), Alias(Sum(call).toAggregateExpression(), "s")()),
+ LocalRelation(attrA))
+ val beside = Aggregate(Seq(attrA),
+ Seq(Alias(attrA, "a")(), Alias(call, "v")()), LocalRelation(attrA))
+ Seq(insideSum, beside).foreach { agg =>
+ val converted = convert(agg)
+ assert(paramColumns(converted).isEmpty, s"Expected no column: $converted")
+ assert(converted.expressions.exists(_.exists(_.isInstanceOf[PythonUDF])),
+ s"Expected the Python UDF kept: $converted")
+ }
+ }
+ }
+
+ test("keeps the Python UDF in a HAVING, where a column would hide the Aggregate") {
+ transpileOn {
+ // A Filter on an Aggregate could host the Project perfectly well, and that is the trap:
+ // `splitSubquery` finds a correlated subquery's Aggregate by matching `Filter(_, Aggregate)`,
+ // so a Project between them hides it and the COUNT bug goes unhandled. Measured end to end in
+ // test_udf_transpile_puts_no_column_directly_above_an_aggregate.
+ val call = makeTPUDF(makePyUDF(Add(attrA, Literal(1L))), Add(pref(0), pref(0)))
+ val agg = Aggregate(Seq(attrA), Seq(attrA), LocalRelation(attrA))
+ val converted = convert(Filter(GreaterThan(call, Literal(0L)), agg))
+ assert(paramColumns(converted).isEmpty, s"Expected no column: $converted")
+ assert(converted.expressions.exists(_.exists(_.isInstanceOf[PythonUDF])),
+ s"Expected the Python UDF kept: $converted")
+ }
+ }
+
+ test("puts one argument in a column though another can't go in one") {
+ transpileOn {
+ // `f(sum(a), a + 1)`, the aggregate read once and `a + 1` twice: no Project can hold the
+ // aggregate, which is no reason to leave `a + 1` at both use sites too.
+ val aggArg = Sum(attrA).toAggregateExpression()
+ val plainArg = Add(attrA, Literal(1L))
+ val option = Add(pref(0), Multiply(pref(1), pref(1)))
+ val pyUDF = PythonUDF("udf", null, LongType, Seq(aggArg, plainArg),
+ PythonEvalType.SQL_BATCHED_UDF, udfDeterministic = true)
+ val converted = convert(Project(
+ Seq(Alias(TranspiledPythonUDF("udf", pyUDF, List(option)), "v")()), LocalRelation(attrA)))
+ assert(paramColumns(converted).map(_.child) == Seq(plainArg),
+ s"Expected a column for the argument that can have one: $converted")
+ }
+ }
+
+ test("keeps the Python UDF in a Join condition, where no column can go") {
+ transpileOn {
+ // A join has no side to compute on, so an argument owed one evaluation goes back to Python --
+ // which is what a UDF predicate in a join condition gets without transpilation anyway.
+ // Repeating it instead would make it lazy again and swallow an ANSI error the interpreted UDF
+ // raises.
+ val arg = Add(attrA, Literal(1L))
+ val attrB = AttributeReference("b", LongType)()
+ val call = makeTPUDF(makePyUDF(arg), Multiply(pref(0), pref(0)))
+ val converted = convert(Join(LocalRelation(attrA), LocalRelation(attrB), Inner,
+ Some(GreaterThan(call, Literal(0L))), JoinHint.NONE))
+ assert(paramColumns(converted).isEmpty, s"Expected no column under a join: $converted")
+ assert(converted.expressions.exists(_.exists(_.isInstanceOf[PythonUDF])),
+ s"Expected the Python UDF kept: $converted")
+ assert(converted.expressions.head.collect { case e if e == arg => e }.length == 1,
+ s"Expected the one evaluation the Python UDF makes: $converted")
+ }
+ }
+
+ test("keeps the Python UDF when an analyzer rule has taken the arguments away") {
+ transpileOn {
+ // PullOutNondeterministic replaces a nondeterministic call with an attribute for the
+ // projection it moved the call to, leaving no arguments for the option's references. Shaped
+ // by hand here; `orderBy(u(rand()))` is the query that produces it.
+ val stripped = TranspiledPythonUDF(
+ "udf", AttributeReference("_nondeterministic", LongType)(), List(Add(pref(0), pref(0))))
+ val result = ConvertToCatalyst.applyExpr(stripped, parentIsUdf = false)
+ assert(result == stripped.pythonUDFExpr,
+ s"Expected the attribute the analyzer left, got: $result")
+ }
+ }
+
+ test("still pre-evaluates a call sitting beside an unrelated lambda") {
+ transpileOn {
+ // Pre-evaluation turns off going into a lambda, not for any expression that holds one, so a
+ // call outside it still gets a column.
+ val arg = Add(attrA, Literal(1L))
+ val call = makeTPUDF(makePyUDF(arg), Multiply(pref(0), pref(0)))
+ val lambdaVar = NamedLambdaVariable("x", LongType, nullable = false)
+ val arr = AttributeReference("arr", ArrayType(LongType))()
+ val unrelated = ArrayTransform(arr, LambdaFunction(lambdaVar, Seq(lambdaVar)))
+ val converted = convert(
+ Project(Seq(Alias(Add(call, Size(unrelated)), "v")()), LocalRelation(attrA, arr)))
+ assert(paramColumns(converted).map(_.child) == Seq(arg),
+ s"Expected the call outside the lambda to keep its column: $converted")
+ }
+ }
+
+ test("keeps the Python UDF for an argument carrying an outer reference, decorrelation off") {
+ // A column holding an OuterReference lands inside the subquery. Decorrelation carries it out;
+ // the fallback rewrites Filters only and would strand it, so with that off we turn it down.
+ val arg = Add(OuterReference(attrA), Literal(1L))
+ val tpudf = makeTPUDF(makePyUDF(arg), Multiply(pref(0), pref(0)))
+ val plan = Project(Seq(Alias(tpudf, "v")()), LocalRelation(attrA))
+ transpileOn {
+ assert(paramColumns(convert(plan)).map(_.child) == Seq(arg),
+ "Expected a column when decorrelation is on")
+ }
+ withSQLConf(
+ SQLConf.ANSI_ENABLED.key -> "true",
+ SQLConf.ATTEMPT_TRANSPILATION_OF_PYTHON_UDFS.key -> "true",
+ SQLConf.DECORRELATE_INNER_QUERY_ENABLED.key -> "false") {
+ val converted = convert(plan)
+ assert(paramColumns(converted).isEmpty, s"Expected no column: $converted")
+ // And with nowhere to put it, the call is left to Python rather than evaluating `a + 1`
+ // twice per row.
+ assert(converted.expressions.head.exists(_.isInstanceOf[PythonUDF]),
+ s"Expected the Python UDF kept: $converted")
+ }
+ }
+
+ test("keeps the Python UDF for an argument that is itself an aggregate") {
+ transpileOn {
+ // `udf(sum(a))` read twice: no Project can hold an aggregate, and running it at each use site
+ // is not the answer either, so the call keeps the Python UDF. Under a real Aggregate the
+ // operator itself is on the deny list, so the same call never gets this far.
+ val arg = Sum(attrA).toAggregateExpression()
+ val tpudf = makeTPUDF(makePyUDF(arg), Multiply(pref(0), pref(0)))
+ val converted = convert(Project(Seq(Alias(tpudf, "v")()), LocalRelation(attrA)))
+ assert(paramColumns(converted).isEmpty, s"Expected no column: $converted")
+ assert(converted.expressions.head.exists(_.isInstanceOf[PythonUDF]),
+ s"Expected the Python UDF kept: $converted")
+ assert(converted.expressions.head.collect { case a: AggregateExpression => a }.length == 1,
+ s"Expected the one aggregate the Python UDF reads: $converted")
+ }
+ }
+
+ test("leaves a Command's cheap arguments at each use site") {
+ transpileOn {
+ // No column under a Command: the Project would land between DeleteFromTable and its relation,
+ // hiding what DataSourceV2Strategy matches on. A bare column costs nothing to read twice, so
+ // the call still lowers; anything more keeps the Python UDF.
+ val arg = attrA
+ val relation = LocalRelation(attrA)
+ val option = Multiply(pref(0), pref(0))
+ val tpudf = makeTPUDF(makePyUDF(arg), option)
+ val converted = ConvertToCatalyst(DeleteFromTable(relation, GreaterThan(tpudf, Literal(0L))))
+ assert(converted == DeleteFromTable(relation, GreaterThan(Multiply(arg, arg), Literal(0L))),
+ s"Expected the argument left at both use sites: $converted")
+
+ val costly = makeTPUDF(makePyUDF(Add(attrA, Literal(1L))), option)
+ val fellBack =
+ ConvertToCatalyst(DeleteFromTable(relation, GreaterThan(costly, Literal(0L))))
+ assert(fellBack.expressions.exists(_.exists(_.isInstanceOf[PythonUDF])),
+ s"Expected the Python UDF kept for `a + 1`: $fellBack")
+ }
+ }
+
+ test("apply keeps a transpilable UDF Python when wrapped by a non-transpiled Python UDF") {
+ // `apply` threads parentIsUdf down from the top of each expression, so the mid UDF stays Python
+ // rather than splitting the batch pipeline into Python -> Catalyst -> Python. The applyExpr
+ // tests above pass parentIsUdf in directly and so bypass that threading.
+ transpileOn {
+ val plain = makePyUDF(attrA)
+ val midPy = makePyUDF(plain)
+ val midTPUDF = makeTPUDF(midPy, Add(plain, Literal(4L)))
+ assert(midTPUDF.hasOnlyPythonUDFInputs)
+ val outerPy = makePyUDF(midTPUDF)
+ val optimized = ConvertToCatalyst(Project(Seq(Alias(outerPy, "v")()), LocalRelation(attrA)))
+ val exprs = optimized.flatMap(_.expressions)
+ assert(!exprs.exists(_.exists(_.isInstanceOf[TranspiledPythonUDF])),
+ s"TranspiledPythonUDF survived: $optimized")
+ // Three Python UDFs remain (outer, mid, plain); the mid was not converted to its Add option.
+ assert(exprs.map(_.collect { case u: PythonUDF => u }.size).sum == 3,
+ s"Expected the mid UDF to stay Python (3 PythonUDFs), got: $optimized")
+ }
+ }
+
+ test("projects the pre-evaluated column away again below a Filter") {
+ transpileOn {
+ // A Filter takes its output from its child, so without the restoring Project the schema
+ // would change. A Project-hosted call never reaches that branch, hence a Filter here.
+ val arg = Add(attrA, Literal(1L))
+ val relation = LocalRelation(attrA)
+ val tpudf = makeTPUDF(makePyUDF(arg), Multiply(pref(0), pref(0)))
+ val converted = convert(Filter(GreaterThan(tpudf, Literal(0L)), relation))
+ assert(paramColumns(converted).map(_.child) == Seq(arg), s"Expected a column: $converted")
+ assert(converted.output == relation.output,
+ s"Expected the column projected away again: ${converted.output}")
+ }
+ }
+
+ test("adds no column for an argument the body never reads") {
+ transpileOn {
+ // `f(g(a + 1), a)` where f's body only ever reads parameter 1. Converting g would register a
+ // column for `a + 1` that nothing reads -- harmless enough until you remember `transpile.py`
+ // promises an unused argument is not computed at all, and that under ANSI the column runs and
+ // can raise. So we only convert the arguments the body actually asks for.
+ val unread = Add(attrA, Literal(1L))
+ val inner = makeTPUDF(makePyUDF(unread), Multiply(pref(0), pref(0)))
+ val pyUDF = PythonUDF("udf", null, LongType, Seq(inner, attrA),
+ PythonEvalType.SQL_BATCHED_UDF, udfDeterministic = true)
+ val outer = TranspiledPythonUDF("udf", pyUDF, List(Add(pref(1), pref(1))))
+ val converted = convert(Project(Seq(Alias(outer, "v")()), LocalRelation(attrA)))
+ assert(paramColumns(converted).forall(_.child != unread),
+ s"Pre-evaluated an argument the body never reads: $converted")
+ }
+ }
+
+ test("leaves an argument still holding an enclosing call's reference at each use site") {
+ transpileOn {
+ // Inverse nesting of the test below: the outer option lowers to another transpiled call whose
+ // argument is the outer call's own parameter ref. Custom transpilers only. The inner call has
+ // to give up -- aliased into a Project that ref is stranded where nobody can substitute it.
+ val arg = Add(attrA, Literal(1L))
+ val innerCall = makeTPUDF(makePyUDF(pref(0)), Multiply(pref(0), pref(0)))
+ val outerCall = makeTPUDF(makePyUDF(arg), innerCall)
+ val converted = convert(Project(Seq(Alias(outerCall, "v")()), LocalRelation(attrA)))
+ // One column, for the outer call's own argument; the inner call gave up and inlined.
+ assert(paramColumns(converted).map(_.child) == Seq(arg),
+ s"Expected only the outer argument pre-evaluated: $converted")
+ }
+ }
+
+ test("converts a nested transpiled UDF used as an argument") {
+ transpileOn {
+ // `f(g(a + 1))`, `f`'s option just `_udf_param_0` and `g`'s repeating its own parameter. We
+ // convert arguments before pre-evaluating, so `g`'s argument gets the column and `f`'s -- by
+ // then `g`'s body reading it -- can't, since a Project can't read an alias it's defining.
+ val arg = Add(attrA, Literal(1L))
+ val innerTPUDF = makeTPUDF(makePyUDF(arg), Multiply(pref(0), pref(0)))
+ val outerTPUDF = makeTPUDF(makePyUDF(innerTPUDF), pref(0))
+ val converted = convert(Project(Seq(Alias(outerTPUDF, "v")()), LocalRelation(attrA)))
+ val columns = paramColumns(converted)
+ assert(columns.map(_.child) == Seq(arg),
+ s"Expected only the inner argument pre-evaluated: $converted")
+ // Everything read has to be something the Project below actually defines.
+ val defined = columns.map(_.toAttribute.exprId).toSet
+ val read = converted.expressions.flatMap(_.collect {
+ case r: AttributeReference if r.name.startsWith("_udf_param_") => r.exprId
+ })
+ assert(read.nonEmpty && read.forall(defined.contains),
+ s"A read column was never defined: $converted")
+ }
+ }
+
+ test("substitute leaves a nested call's own options alone") {
+ // A nested call's references count against *its* arguments, so ours must not touch them.
+ // `referencedIndexes` and `resolveTypes` both stop there; `substitute` used to walk in, and
+ // reading a nested index as ours means an out-of-range argument or, worse, the wrong one.
+ // The rule converts a nested call before substituting so it never hits this, but `substitute`
+ // is what a custom transpiler's own rule calls.
+ val inner = makeTPUDF(makePyUDF(attrA), Add(pref(1), Literal(2L)))
+ val option = Add(pref(0), inner)
+ val substituted = TranspiledUDFParameter.substitute(option, Seq[Expression](attrA))
+ val nested = substituted.collectFirst { case n: TranspiledPythonUDF => n }
+ assert(nested.map(_.transpiledOptions).contains(inner.transpiledOptions),
+ s"A nested call's options were substituted out of our index space: $substituted")
+ assert(substituted.asInstanceOf[Add].left == attrA,
+ s"Our own reference was not substituted: $substituted")
+ }
+
+ test("an option reading an argument the call hasn't got is a classed error") {
+ // Only a hand-built option gets here -- the transpiler bounds-checks what it emits, and a
+ // pre-typed reference slips past the analyzer's check as well. Left to resolve on its own the
+ // node never does, and CheckAnalysis blames an internal error, which tells nobody anything.
+ val readsTooFar = makeTPUDF(makePyUDF(attrA), Add(pref(1), Literal(1L)))
+ transpileOn {
+ checkError(
+ exception = intercept[AnalysisException] {
+ convert(Project(Seq(Alias(readsTooFar, "v")()), LocalRelation(attrA)))
+ },
+ condition = "INVALID_UDF_PARAMETER_PLACEHOLDER_INDEX",
+ parameters = Map("index" -> "1", "numParams" -> "1"))
+ }
+
+ // Same error from the analyzer's side, where the type comes off the argument.
+ checkError(
+ exception = intercept[AnalysisException] {
+ TranspiledUDFParameter.resolveTypes(TranspiledUDFParameter(2, None), Seq(attrA))
+ },
+ condition = "INVALID_UDF_PARAMETER_PLACEHOLDER_INDEX",
+ parameters = Map("index" -> "2", "numParams" -> "1"))
+ }
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala
index cfd1fb4ea222..656b223192c5 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala
@@ -28,7 +28,7 @@ import net.razorvine.pickle.Pickler
import org.apache.spark.api.python.{PythonEvalType, PythonFunction, PythonWorkerUtils, SpecialLengths}
import org.apache.spark.sql.{Column, TableArg}
import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
-import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, Descending, Expression, FunctionTableSubqueryArgumentExpression, NamedArgumentExpression, NullsFirst, NullsLast, PythonUDAF, PythonUDF, PythonUDTF, PythonUDTFAnalyzeResult, PythonUDTFSelectedExpression, SortOrder, TranspiledPythonUDF, UnresolvedPolymorphicPythonUDTF, UnresolvedTableArgPlanId}
+import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, Descending, Expression, FunctionTableSubqueryArgumentExpression, NamedArgumentExpression, NullsFirst, NullsLast, PythonUDAF, PythonUDF, PythonUDTF, PythonUDTFAnalyzeResult, PythonUDTFSelectedExpression, SortOrder, TranspiledPythonUDF, TranspiledUDFParameter, UnresolvedPolymorphicPythonUDTF, UnresolvedTableArgPlanId}
import org.apache.spark.sql.catalyst.parser.ParserInterface
import org.apache.spark.sql.catalyst.plans.logical.{Generate, LogicalPlan, NamedParametersSupport, OneRowRelation}
import org.apache.spark.sql.classic.{DataFrame, Dataset, SparkSession}
@@ -126,13 +126,17 @@ case class UserDefinedPythonFunction(
if (transpiledExprsForUse.nonEmpty &&
optionInputTypesForUse.length == transpiledExprsForUse.length &&
optionInputTypesForUse.forall(_.length == e.length)) {
- val udfChildren = udfExpr.children.toArray
- // Resolve the `_udf_param_N` placeholders the transpiler emits into the bound
- // UDF arguments. Apply this ONLY to the transpiled options -- never to
- // `udfExpr` itself, whose children are the user's argument expressions. A user
- // column literally named `_udf_param_N` passed as an argument must not be
- // rewritten, so we leave `udfExpr` untouched.
- def resolveUDFParams(expression: Expression, children: Array[Expression]): Expression = {
+ // Turn the transpiler's `_udf_param_N` placeholders into references to the bound arguments.
+ // References, not copies, so the argument stays in `udfExpr`'s children and ConvertToCatalyst
+ // can compute it once in a Project below the operator (SPARK-58626). We run before the
+ // arguments are bound so the reference has no type yet; ResolveTranspiledPythonUDFOptions
+ // fills that in, which is also what gets the option body coerced.
+ //
+ // Options ONLY, never `udfExpr` itself. Somebody's column really can be called
+ // `_udf_param_0`, and if they pass it in we must not rewrite it.
+ // Only the arity is needed, so the arguments are not in scope here at all -- there is nothing
+ // to read one off the wrong node with.
+ def resolveUDFParams(expression: Expression): Expression = {
expression match {
case UnresolvedAttribute(nameParts)
if nameParts.length == 1 && nameParts.head.startsWith("_udf_param_") =>
@@ -140,17 +144,16 @@ case class UserDefinedPythonFunction(
val index = suffix.toIntOption.getOrElse {
throw QueryCompilationErrors.invalidUDFParameterPlaceholder(nameParts.head)
}
- if (index >= 0 && index < children.length) {
- children(index)
+ if (index >= 0 && index < e.length) {
+ TranspiledUDFParameter(index)
} else {
- throw QueryCompilationErrors.invalidUDFParameterPlaceholderIndex(
- index, children.length)
+ throw QueryCompilationErrors.invalidUDFParameterPlaceholderIndex(index, e.length)
}
case _ =>
- expression.mapChildren(resolveUDFParams(_, children))
+ expression.mapChildren(resolveUDFParams)
}
}
- val resolvedOptions = transpiledExprsForUse.map(resolveUDFParams(_, udfChildren))
+ val resolvedOptions = transpiledExprsForUse.map(resolveUDFParams)
TranspiledPythonUDF(name, udfExpr, resolvedOptions, optionInputTypesForUse)
} else {
udfExpr
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/TranspiledUDFParameterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/TranspiledUDFParameterSuite.scala
new file mode 100644
index 000000000000..162cb1f80247
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/TranspiledUDFParameterSuite.scala
@@ -0,0 +1,268 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.SparkArithmeticException
+import org.apache.spark.api.python.PythonEvalType
+import org.apache.spark.sql.{AnalysisException, Column, QueryTest, Row}
+import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
+import org.apache.spark.sql.catalyst.expressions.{Add, AttributeReference, Expression, Multiply, TranspiledPythonUDF, TranspiledUDFParameter}
+import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, BaseEvalPython, Join}
+import org.apache.spark.sql.classic.ClassicConversions._
+import org.apache.spark.sql.connector.catalog.InMemoryRowLevelOperationTableCatalog
+import org.apache.spark.sql.functions.{array, col, max, rand, sum, transform}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.LongType
+
+/**
+ * Tests how `UserDefinedPythonFunction.builder` turns the `_udf_param_N` placeholders into
+ * references to the UDF's arguments (SPARK-58626). The builder owns this because it is the last
+ * place that still sees the placeholders.
+ *
+ * Plus the operator shapes that decide whether an argument gets pre-evaluated: a Project and a
+ * MergeRows host the column, while an Aggregate, a join condition and a lambda cannot, so a call
+ * that needs one there keeps the Python UDF.
+ * Asserted on the plan where a deterministic argument gives the same answer either way, and on a
+ * raised error under MERGE, where an argument that mods by zero only blows up if the column is
+ * really there.
+ *
+ * A nondeterministic argument is the case where the column is not an optimization: with no column
+ * to read, a body reading the parameter twice would draw twice, so we leave the Python UDF.
+ */
+class TranspiledUDFParameterSuite extends QueryTest with SharedSparkSession {
+
+ private val argA = AttributeReference("a", LongType)()
+ private val argB = AttributeReference("b", LongType)()
+
+ private def param(index: Int): Expression = UnresolvedAttribute(s"_udf_param_$index")
+
+ /** A nondeterministic argument, LONG-typed to match what `udfWith` declares. A `def`, so two
+ * uses are two draws with their own seeds. */
+ private def draw: Column = (rand() * 3).cast(LongType)
+
+ /**
+ * A transpiled UDF of `arity` numeric parameters carrying the given option body. `func` is null
+ * because a call that transpiles never reaches Python -- so a fallback shows up as an NPE rather
+ * than as a quietly passing test.
+ */
+ private def udfWith(option: Expression, arity: Int = 2): UserDefinedPythonFunction =
+ UserDefinedPythonFunction(
+ "udf",
+ null,
+ LongType,
+ PythonEvalType.SQL_BATCHED_UDF,
+ udfDeterministic = true,
+ List(Column(option)).asJava,
+ List(List.fill(arity)("numeric").asJava).asJava)
+
+ /** The argument indexes the single option the builder produced refers to, in tree order. */
+ private def parameterIndexes(option: Expression, args: Expression*): Seq[Int] =
+ udfWith(option).builder(args) match {
+ case t: TranspiledPythonUDF =>
+ t.transpiledOptions.head.collect { case p: TranspiledUDFParameter => p.index }
+ case other => fail(s"Expected a TranspiledPythonUDF, got: $other")
+ }
+
+ private def transpileOn(f: => Unit): Unit = withSQLConf(
+ SQLConf.ATTEMPT_TRANSPILATION_OF_PYTHON_UDFS.key -> "true",
+ SQLConf.ANSI_ENABLED.key -> "true")(f)
+
+ test("refers to every parameter the body reads, once per read") {
+ // `lambda a, b: a * a + b`: two refs to `a`, one to `b`. Binding both parameters to the same
+ // argument changes nothing -- the option refers to parameters, and which argument each stands
+ // for is settled later.
+ val option = Add(Multiply(param(0), param(0)), param(1))
+ Seq(argB, argA).foreach { second =>
+ assert(parameterIndexes(option, argA, second) == Seq(0, 0, 1),
+ s"Expected a, a, b referenced when b is bound to $second")
+ }
+ }
+
+ test("drops a parameter the body never uses") {
+ // `lambda a, b: a + a`: b never reaches the option, so nothing evaluates it.
+ val option = Add(param(0), param(0))
+ assert(parameterIndexes(option, argA, argB) == Seq(0, 0))
+ }
+
+ test("pre-evaluates an argument two references read") {
+ transpileOn {
+ // `id % 3` rather than a bare column, or it's cheap and gets left inline. Count definitions,
+ // not reads: one `AS _udf_param_` read twice is the shape we want. The name carries the
+ // column's ExprId, so match the prefix rather than a number.
+ val square = udfWith(Multiply(param(0), param(0)), arity = 1)
+ val df = spark.range(0, 6).select(square(col("id") % 3).as("sq"))
+ val plan = df.queryExecution.optimizedPlan.toString
+ assert("AS _udf_param_".r.findAllIn(plan).length == 1,
+ s"Expected exactly one pre-evaluated argument column:\n$plan")
+ checkAnswer(df, Seq(Row(0L), Row(1L), Row(4L), Row(0L), Row(1L), Row(4L)))
+ }
+ }
+
+ test("keeps the Python UDF under an Aggregate") {
+ transpileOn {
+ // No Project can go under an Aggregate, so a call owed an evaluation stays Python. `func` is
+ // null here, so the plan is the only thing to assert on; the answers are pinned in
+ // test_udf_transpile_unit.py, where the UDF is real.
+ val square = udfWith(Multiply(param(0), param(0)), arity = 1)
+ Seq(spark.range(0, 6).groupBy(col("id") % 3).agg(square(col("id") % 3).as("sq")),
+ spark.range(0, 6).groupBy(col("id") % 3).agg(square(draw).as("sq"))).foreach { df =>
+ val plan = df.queryExecution.optimizedPlan
+ assert(!plan.toString.contains("_udf_param_"), s"Expected no column:\n$plan")
+ assert(plan.exists(_.isInstanceOf[BaseEvalPython]), s"Expected a Python fallback:\n$plan")
+ }
+ }
+ }
+
+ test("leaves an Aggregate alone when its only call is inside a subquery") {
+ transpileOn {
+ // A tree-pattern bit propagates out of a SubqueryExpression, so this Aggregate looks from the
+ // outside like it holds a call. Nothing here reshapes an Aggregate any more, but the shape is
+ // cheap to keep pinned: it is what a rule rebuilding one on every pass looked like, a
+ // StackOverflowError.
+ val square = udfWith(Multiply(param(0), param(0)), arity = 1)
+ val inner = spark.range(0, 3).select(square(col("id") % 3).as("m")).agg(max(col("m")))
+ val df = spark.range(0, 6).groupBy(col("id") % 2).agg(sum(col("id") + inner.scalar()))
+ assert(df.queryExecution.optimizedPlan.collect { case a: Aggregate => a }.nonEmpty)
+ assert(df.collect().length == 2)
+ }
+ }
+
+ test("a call straddling both sides of a non-inner join fails the way Python does") {
+ transpileOn {
+ // No column can go under a join, so a call owed an evaluation keeps the Python UDF, and
+ // ExtractPythonUDFFromJoinCondition turns one down on anything but an inner join. Same error
+ // the query gets with transpilation off. A body reading its parameter once lowers and runs.
+ val square = udfWith(Multiply(param(0), param(0)), arity = 2)
+ val left = spark.range(0, 3).select(col("id").as("a"))
+ val right = spark.range(0, 3).select(col("id").as("c"))
+ val straddles = square(col("a") + 1, col("c")) > 0
+ val joined = left.join(right, straddles, "left_outer")
+ checkError(
+ exception = intercept[AnalysisException](joined.collect()),
+ condition = "UNSUPPORTED_FEATURE.PYTHON_UDF_IN_ON_CLAUSE",
+ parameters = Map("joinType" -> "LEFT OUTER"),
+ sqlState = Some("0A000"))
+
+ // Read once, the same call lowers and the join runs.
+ val once = udfWith(Add(param(0), param(1)), arity = 2)
+ val lowered = left.join(right, once(col("a") + 1, col("c")) > 0, "left_outer")
+ assert(!lowered.queryExecution.optimizedPlan.exists(_.isInstanceOf[BaseEvalPython]),
+ s"Expected no Python:\n${lowered.queryExecution.optimizedPlan}")
+ assert(lowered.collect().length == 9)
+ }
+ }
+
+ test("keeps the Python UDF where a draw's column would strand a join's condition") {
+ transpileOn {
+ // `where a = c and udf(rand()) > 0` over two tables: a nondeterministic column between the
+ // Filter and the Join stops PushPredicateThroughNonJoin pushing anything through it, so
+ // PushPredicateThroughJoin never sees the Filter and the join keeps no condition at all -- a
+ // cartesian product with the right answers. The call keeps the Python UDF instead, and the
+ // equi-join survives.
+ val square = udfWith(Multiply(param(0), param(0)), arity = 1)
+ val left = spark.range(0, 4).select(col("id").as("a"))
+ val right = spark.range(0, 4).select(col("id").as("c"))
+ val df = left.join(right).where(col("a") === col("c") && (square(draw) > 0))
+ val plan = df.queryExecution.optimizedPlan
+ assert(plan.collect { case j: Join => j }.exists(_.condition.isDefined),
+ s"Expected the join to keep its condition:\n$plan")
+ assert(!plan.toString.contains("_udf_param_"), s"Expected no column:\n$plan")
+ assert(plan.exists(_.isInstanceOf[BaseEvalPython]), s"Expected a Python fallback:\n$plan")
+ }
+ }
+
+ test("keeps the Python UDF where a draw's column would strand a join under a projection") {
+ transpileOn {
+ // Same as above with a `select` in between, which is what `FROM t1, t2 WHERE ...` looks like
+ // once anything projects: the Filter's child is a Project, not the Join. The column lands
+ // below that Project and blocks the pushdown just the same, so the check has to look for a
+ // Join anywhere below, not only at the child.
+ val square = udfWith(Multiply(param(0), param(0)), arity = 1)
+ val left = spark.range(0, 4).select(col("id").as("a"))
+ val right = spark.range(0, 4).select(col("id").as("c"))
+ val df = left.join(right).select(col("c"), col("a"))
+ .where(col("a") === col("c") && (square(draw) > 0))
+ val plan = df.queryExecution.optimizedPlan
+ assert(plan.collect { case j: Join => j }.exists(_.condition.isDefined),
+ s"Expected the join to keep its condition:\n$plan")
+ assert(!plan.toString.contains("_udf_param_"), s"Expected no column:\n$plan")
+ assert(plan.exists(_.isInstanceOf[BaseEvalPython]), s"Expected a Python fallback:\n$plan")
+ }
+ }
+
+ test("keeps the Python UDF inside a higher-order function's lambda") {
+ transpileOn {
+ // Lambdas are out of scope for lowering: Spark applies a Python UDF over the whole array
+ // there and this rule leaves that to it. The Python eval node is what says we did.
+ val square = udfWith(Multiply(param(0), param(0)), arity = 1)
+ val df = spark.range(0, 6).select(
+ transform(array(col("id")), _ => square(draw)).as("sq"))
+ val plan = df.queryExecution.optimizedPlan
+ assert(plan.exists(_.isInstanceOf[BaseEvalPython]), s"Expected a Python fallback:\n$plan")
+ assert(!plan.toString.contains("_udf_param_"), s"Expected no column:\n$plan")
+ }
+ }
+
+ test("pre-evaluates a MERGE instruction condition's argument") {
+ // MergeRows is a plain unary node, so it hosts the Project -- and it has to, or a transpiled
+ // MERGE ends up lazier than an interpreted one. Measured: an interpreted UDF in `WHEN MATCHED
+ // AND` raises under ANSI on a row with no match at all, since its inputs are computed below
+ // MergeRows. First-match-wins instructions don't make it lazy.
+ //
+ // Also the only shape here whose `output` is a constructor field, so no restoring Project gets
+ // added and the `__row_id` lookup has to still find its own past the one we appended.
+ withSQLConf(
+ "spark.sql.catalog.rowcat" -> classOf[InMemoryRowLevelOperationTableCatalog].getName) {
+ transpileOn {
+ val square = udfWith(Multiply(param(0), param(0)), arity = 1)
+ val target = "rowcat.ns.t"
+ // The DataFrame API rather than SQL because `udfWith` leaves `func` null on purpose, and
+ // registering a UDF by name dereferences it. Applying the call as a Column skips that.
+ // `%` and not `/` so the argument stays LONG and matches the option's declared return
+ // type. Modulo by zero raises under ANSI too, just as REMAINDER_BY_ZERO.
+ def runMerge(): Unit =
+ sql("SELECT 1 AS id").as("s")
+ .mergeInto(target, col("s.id") === col(s"$target.id"))
+ .whenMatched(square(col(s"$target.a") % col(s"$target.b")) > 3)
+ .delete()
+ .merge()
+ def create(rows: String): Unit = {
+ sql(s"CREATE TABLE $target (id INT, a LONG, b LONG) USING foo")
+ sql(s"INSERT INTO $target VALUES $rows")
+ }
+ withTable(target) {
+ create("(1, 6, 4), (2, 6, 1)")
+ runMerge()
+ // id=1 matches, (6 % 4) squared is 4 > 3, so it goes. id=2 never matches.
+ checkAnswer(sql(s"SELECT id FROM $target"), Seq(Row(2)))
+ }
+ withTable(target) {
+ // id=3 has b=0 and no matching source row, so its condition never fires. It raises
+ // anyway, because the column is computed for every row of the join -- which is also what
+ // proves the column is there, since inline the instruction would never evaluate it.
+ create("(1, 6, 4), (3, 6, 0)")
+ val e = intercept[SparkArithmeticException](runMerge())
+ assert(e.getCondition == "REMAINDER_BY_ZERO", s"Unexpected error: ${e.getCondition}")
+ }
+ }
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org