This is an automated email from the ASF dual-hosted git repository.
asf-gitbox-commits pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 9d64ebdf12db [SPARK-58650][PYTHON][SQL] Improve source extraction
during transpilation
9d64ebdf12db is described below
commit 9d64ebdf12db8b75f78e6e388b94236b0250724c
Author: Holden Karau <[email protected]>
AuthorDate: Mon Aug 24 15:20:20 2026 -0700
[SPARK-58650][PYTHON][SQL] Improve source extraction during transpilation
### What changes were proposed in this pull request?
Change how we find the source code for lambdas in transpilation
### Why are the changes needed?
Multiple lambdas on one line could cause the transpilation logic to fail /
fall back.
### Does this PR introduce _any_ user-facing change?
More transpilation targets
### How was this patch tested?
Added tests with more transpilation of multiple lambdas on the same line.
### Was this patch authored or co-authored using generative AI tooling?
Yes Claude 5 did much of the authoring
Closes #58188 from holdenk/SPARK-58650-there-can-only-be-one-lambda.
Lead-authored-by: Holden Karau <[email protected]>
Co-authored-by: Holden Karau <[email protected]>
Signed-off-by: Holden Karau <[email protected]>
---
.../pyspark/sql/tests/test_udf_transpile_unit.py | 399 +++++++++++++++++++--
python/pyspark/sql/transpile.py | 319 ++++++++++++----
2 files changed, 609 insertions(+), 109 deletions(-)
diff --git a/python/pyspark/sql/tests/test_udf_transpile_unit.py
b/python/pyspark/sql/tests/test_udf_transpile_unit.py
index de753b0e7bc0..5593345837f6 100644
--- a/python/pyspark/sql/tests/test_udf_transpile_unit.py
+++ b/python/pyspark/sql/tests/test_udf_transpile_unit.py
@@ -870,24 +870,19 @@ class UDFTranspileUnitTests(ReusedSQLTestCase):
[result] = df.select(pudf("a")).collect()
self.assertEqual(result[0], 1)
- def test_udf_transpile_falls_back_for_plain_self_param(self):
- # A plain function whose first parameter is literally named `self`
- # must not be confused with a bound __call__ method: stripping it
- # previously emitted `_udf_param_-1` and threw an AnalysisException
- # at call construction.
- import warnings as _warnings
-
+ def test_udf_transpile_plain_self_param_is_an_ordinary_param(self):
+ # A plain function whose first parameter is literally named `self` is
not a
+ # bound receiver -- the call site supplies it. Stripping it emitted
+ # `_udf_param_-1` and threw at call construction, so it used to be
refused
+ # outright; the receiver is decided by dispatch now, so this lowers.
def weird(self, other):
return self + other
- with self.sql_conf(_TRANSPILE_ON):
- with _warnings.catch_warnings(record=True):
- _warnings.simplefilter("always")
- pudf = UserDefinedFunction(weird, LongType())
- self.assertEqual([], pudf.transpiled, "plain 'self' param must not
transpile")
- df = self.spark.createDataFrame([Row(a=2, b=3)])
- [result] = df.select(pudf("a", "b")).collect()
- self.assertEqual(result[0], 5)
+ self.assertEqual(
+ self._vals(weird, LongType(), "a long, b long", [(2, 3)]),
+ [5],
+ "both parameters are supplied at the call site, so both get
placeholders",
+ )
def test_udf_transpile_falls_back_for_wraps_decorated_function(self):
# inspect.getsource follows __wrapped__, so a functools.wraps-decorated
@@ -1362,17 +1357,55 @@ class UDFTranspileUnitTests(ReusedSQLTestCase):
# contains "UDF" (so the "UDF" substring is unreliable).
# ------------------------------------------------------------------
- def _vals(self, func, return_type, schema, rows):
- with self.sql_conf(_TRANSPILE_ON):
+ @staticmethod
+ def _udf_and_warnings(func, return_type):
+ """Build a UDF, returning it with the text of any warnings it emitted.
+
+ ``udf.py`` reports WHY a UDF fell back only as a warning, so every
test that
+ cares about a fallback needs them captured. The caller must already be
inside
+ ``sql_conf(_TRANSPILE_ON)`` -- this does not set the conf.
+ """
+ import warnings
+
+ with warnings.catch_warnings(record=True) as caught:
+ warnings.simplefilter("always")
u = UserDefinedFunction(func, return_type)
- self.assertTrue(u.transpiled, str(func))
+ return u, " ".join(str(w.message) for w in caught)
+
+ def _fallback_reason(self, func, return_type=LongType()):
+ """Assert ``func`` produced no options, and hand back the reason it
reported."""
+ u, reasons = self._udf_and_warnings(func, return_type)
+ self.assertEqual([], u.transpiled, f"{func} must fall back")
+ self.assertTrue(reasons, f"{func} fell back without saying why")
+ return u, reasons
+
+ def _transpiled_udf(self, func, return_type):
+ """A UDF asserted to have produced options; naming the fallback reason
if not.
+
+ Without the captured warning an empty ``transpiled`` asserts as bare
"[] is
+ not true", and from CI the text is only in a credentialed log
artifact. The
+ caller must already be inside ``sql_conf(_TRANSPILE_ON)``.
+ """
+ u, reasons = self._udf_and_warnings(func, return_type)
+ self.assertTrue(u.transpiled, f"{func} produced no transpiled options:
{reasons}")
+ return u
+
+ def _vals(self, func, return_type, schema, rows, require_lowered=True):
+ with self.sql_conf(_TRANSPILE_ON):
+ u = self._transpiled_udf(func, return_type)
df = self.spark.createDataFrame(rows, schema)
- return [r[0] for r in df.select(u(*df.columns)).collect()]
+ projected = df.select(u(*df.columns))
+ # ``u.transpiled`` only says options were PRODUCED; the JVM may
still
+ # discard them and run interpreted Python, returning the right
value and
+ # hiding a wrong lowering. Pass ``require_lowered=False`` only
where the
+ # JVM is EXPECTED to discard them.
+ if require_lowered:
+ self.assertEqual(0, self._eval_python_count(projected),
str(func))
+ return [r[0] for r in projected.collect()]
def _raises(self, func, schema, rows, needle="numeric"):
with self.sql_conf(_TRANSPILE_ON):
- u = UserDefinedFunction(func, LongType())
- self.assertTrue(u.transpiled, str(func))
+ u = self._transpiled_udf(func, LongType())
df = self.spark.createDataFrame(rows, schema)
with self.assertRaises(Exception) as ctx:
df.select(u(*df.columns)).collect()
@@ -1441,9 +1474,10 @@ class UDFTranspileUnitTests(ReusedSQLTestCase):
with self.subTest(case=i):
self.assertEqual(self._vals(func, rt, schema, rows), expected,
f"case {i}: {rows}")
- def test_udf_transpile_callable_object_self_offset(self):
- # A callable instance carries `self`; the extractor offsets it so a/b
- # map to _udf_param_0/_udf_param_1 (non-commutative body proves order).
+ def test_udf_transpile_callable_object_drops_its_receiver(self):
+ # A callable instance's `self` is dropped before anything indexes the
param
+ # list, so a/b are _udf_param_0/_udf_param_1 with no offsetting
anywhere
+ # downstream (the non-commutative body proves the order).
class SubAB:
def __call__(self, a, b):
return a - b
@@ -1529,7 +1563,9 @@ class UDFTranspileUnitTests(ReusedSQLTestCase):
base = lambda v, w: v + w # noqa: E731
percent_fmt = lambda x: "n=%d" % x # noqa: E731
with self.sql_conf(_TRANSPILE_ON):
- # inline / wrapped / partial lambdas -> source can't be extracted
+ # An inline or wrapped lambda is a call ARGUMENT, so the source is
+ # read fine but parses as ``Call`` rather than a definition we can
+ # unwrap; ``functools.partial`` has no reachable source at all.
self.assertEqual([], UserDefinedFunction(lambda v: v + 1,
LongType()).transpiled)
self.assertEqual(
[], UserDefinedFunction(wrapper(lambda v: v + 1),
LongType()).transpiled
@@ -1553,6 +1589,307 @@ class UDFTranspileUnitTests(ReusedSQLTestCase):
[num.select(wd("a")).first()[0], num.select(wd("a",
"a")).first()[0]], [5, 55]
)
+ def
test_udf_transpile_falls_back_when_a_sibling_lambda_shares_the_line(self):
+ # ``inspect.getsource`` works in whole lines, so every lambda on a
line hands
+ # back the same source and nothing in it says which one we hold. Taking
+ # whichever came first was right by position rather than by identity:
+ # ``minus_one`` lowered ``x + 1``, and the lambda in the decorator
lowered
+ # the decorated ``def``'s body -- both silently wrong. Refusing costs
+ # ``plus_one``, which the first-match rule happened to get right;
being right
+ # for 1-of-N by position is not something a caller can rely on
(SPARK-58650).
+ #
+ # ``fmt: off`` keeps the two on one line; asserted below, since the
formatter
+ # would otherwise split them and quietly make this test vacuous.
+ # fmt: off
+ plus_one = lambda x: x + 1; minus_one = lambda x: x - 1 # noqa:
E702,E731
+ # fmt: on
+ self.assertEqual(
+ plus_one.__code__.co_firstlineno,
+ minus_one.__code__.co_firstlineno,
+ "fixture must keep both lambdas on ONE line or it proves nothing",
+ )
+
+ captured = []
+
+ def capture(g):
+ captured.append(g)
+ return lambda fn: fn
+
+ @capture(lambda x: x + 1)
+ def unrelated(a):
+ return a * 12345
+
+ (decorator_lambda,) = captured
+
+ # A class body is the same hazard: master read the whole line and
lowered
+ # ``helper``, computing 5*100 for a UDF that really returns 5*2.
+ class TwoOnALine:
+ helper, __call__ = lambda self, x: x * 100, lambda self, x: x * 2
+
+ self.assertEqual(10, TwoOnALine()(5))
+
+ num = self.spark.createDataFrame([(5,)], "a long")
+ sibling = "put each lambda on its own line"
+ not_a_statement = "does not define it as a statement of its own"
+ with self.sql_conf(_TRANSPILE_ON):
+ for label, func, expected, needle in [
+ # Lowered the WRONG body before -- the bug.
+ ("second lambda on the line", minus_one, 4, sibling),
+ ("lambda inside a decorator", decorator_lambda, 6,
not_a_statement),
+ ("__call__ sharing a class-body line", TwoOnALine(), 10,
not_a_statement),
+ # Lowered correctly before; refused now as acknowledged
collateral.
+ ("first lambda on the line", plus_one, 6, sibling),
+ ]:
+ with self.subTest(case=label):
+ u, reasons = self._fallback_reason(func)
+ # Assert the REASON, so relaxing the guard fails here
loudly
+ # rather than passing for a new and unrelated cause.
+ self.assertIn(needle, reasons)
+ self.assertEqual(expected, num.select(u("a")).first()[0])
+
+ # The ``def`` itself is NOT collateral: we know we are holding it, so a
+ # lambda in its decorator cannot be the body and does not block
lowering.
+ self.assertEqual(self._vals(unrelated, LongType(), "a long", [(5,)]),
[61725])
+
+ def
test_udf_transpile_refuses_a_lambda_whose_source_now_reads_as_a_def(self):
+ # ``inspect.getsource`` reads the file as it is NOW, so an edit after
import
+ # can hand back source holding no lambda at all. Refusing only when a
SIBLING
+ # lambda is in view failed open here, lowering an unrelated ``def`` as
the
+ # body: 5 * 12345 for a lambda Python evaluates to 6.
+ import importlib.util
+ import os
+ import tempfile
+
+ with tempfile.TemporaryDirectory() as tmp:
+ path = os.path.join(tmp, "drifted_lambda.py")
+ with open(path, "w") as handle:
+ handle.write("f = lambda x: x + 1\n")
+ spec = importlib.util.spec_from_file_location("drifted_lambda",
path)
+ module = importlib.util.module_from_spec(spec)
+ spec.loader.exec_module(module)
+ with open(path, "w") as handle:
+ handle.write("def unrelated(a):\n return a * 12345\n")
+
+ with self.sql_conf(_TRANSPILE_ON):
+ u, reasons = self._fallback_reason(module.f)
+ # Pin the REASON, or a source read that merely failed would
pass too.
+ self.assertIn("which lambda to lower cannot be determined",
reasons)
+ num = self.spark.createDataFrame([(5,)], "a long")
+ self.assertEqual(6, num.select(u("a")).first()[0])
+
+ def test_udf_transpile_ambiguity_check_sees_only_rival_lambdas(self):
+ # Two ways the check used to misfire or not fire at all.
+ from pyspark.sql.transpile import _held_code
+
+ # A lambda nested in the held lambda's own body is not a rival -- it
can
+ # never be the UDF, and the user cannot split it onto another line. It
fell
+ # back with the sibling message, advice that could not be acted on.
+ nested = lambda x: (lambda y: y + 1)(x) # noqa: E731
+ with self.sql_conf(_TRANSPILE_ON):
+ _, reasons = self._fallback_reason(nested)
+ self.assertIn("Call", reasons, "must fall back for the body, not
for ambiguity")
+ self.assertNotIn("put each lambda on its own line", reasons)
+
+ # The mirror: a lambda RETURNED by a one-line lambda. Here the outer
one is
+ # what the source read locates and the inner one is what we hold, so
treating
+ # nested lambdas as never-rivals let this through -- and it proceeded
with the
+ # outer signature, reporting `n` as the UDF's parameter for a UDF
whose only
+ # parameter is `x`. Matching the located lambda's parameters against
the held
+ # code object's is what separates this from the case above.
+ # fmt: off
+ make_adder = lambda n: lambda x: x + n # noqa: E731
+ # fmt: on
+ add_three = make_adder(3)
+ self.assertEqual(8, add_three(5))
+ with self.sql_conf(_TRANSPILE_ON):
+ u, reasons = self._fallback_reason(add_three)
+ self.assertIn("takes different parameters", reasons)
+ self.assertEqual([], u._transpiled_param_names or [])
+ self.assertEqual(
+ 8, self.spark.createDataFrame([(5,)], "a
long").select(u("a")).first()[0]
+ )
+
+ # ``staticmethod``/``classmethod`` hide ``__code__`` behind the
descriptor, so
+ # reading it off them left the guard inactive -- skipping the check
for a shape
+ # it exists to catch. Now that these lower at all, the skip would be a
wrong
+ # answer: ``helper`` shares the line and takes the same parameter, so
it is a
+ # true rival and would be lowered instead (5 * 9, not 5 * 2).
+ class Wrapped:
+ # fmt: off
+ helper = lambda x: x * 9; __call__ = staticmethod(lambda x: x * 2)
# noqa: E702,E731
+ # fmt: on
+
+ self.assertEqual(
+ Wrapped.helper.__code__.co_firstlineno,
+ Wrapped.__call__.__code__.co_firstlineno,
+ "fixture must keep both lambdas on ONE line or it proves nothing",
+ )
+ self.assertEqual(
+ "<lambda>",
+ getattr(_held_code(Wrapped()), "co_name", None),
+ "the held code must be found inside the descriptor",
+ )
+ self.assertEqual(10, Wrapped()(5))
+ with self.sql_conf(_TRANSPILE_ON):
+ u, reasons = self._fallback_reason(Wrapped())
+ self.assertIn("put each lambda on its own line", reasons)
+ num = self.spark.createDataFrame([(5,)], "a long")
+ self.assertEqual(10, num.select(u("a")).first()[0])
+
+ def test_udf_transpile_lowers_an_annotated_lambda_binding(self):
+ # An annotated binding is the same shape as a plain one, and the form
a typed
+ # codebase writes. Only ``ast.Assign`` was unwrapped, so this was
refused as
+ # "not a statement of its own" -- while the module docstring told
users to
+ # bind the lambda to a name and give it a line, which is exactly this.
+ from typing import Callable
+
+ annotated: Callable[[int], int] = lambda x: x + 1 # noqa: E731
+ self.assertEqual(self._vals(annotated, LongType(), "a long", [(5,)]),
[6])
+
+ def test_udf_transpile_recovers_shapes_with_an_unheld_lambda_in_view(self):
+ # The ambiguity check applies only when the callable we hold IS a
lambda. The
+ # lambdas below belong to a ``def`` we are not lowering, so refusing
on their
+ # account would cost lowering for nothing.
+ from typing import Annotated
+
+ def annotated(x: Annotated[int, lambda v: v > 0]) -> int:
+ return x + 1
+
+ def returns_annotated(x) -> Annotated[int, lambda v: v > 0]:
+ return x + 2
+
+ for label, func, expected in [
+ ("lambda in a parameter annotation", annotated, 6),
+ ("lambda in the return annotation", returns_annotated, 7),
+ ]:
+ with self.subTest(case=label):
+ self.assertEqual(self._vals(func, LongType(), "a long",
[(5,)]), [expected])
+
+ def test_udf_transpile_resolves_call_on_the_type_not_the_instance(self):
+ # Python's call protocol looks ``__call__`` up on the TYPE, so an
instance
+ # attribute of that name is never what runs. ``getattr(obj,
"__call__")``
+ # finds it anyway, so the transpiler used to lower the shadowing body
and
+ # return 5*99 where Python returns 5*4 -- silently wrong. The type's
+ # ``__call__`` must win, and it must still lower.
+ class Shadowed:
+ def __call__(self, x):
+ return x * 4
+
+ shadowed = Shadowed()
+ # Alone on its line: a leading statement on the same line would make
the
+ # shadowing body unreachable for an unrelated reason and prove nothing.
+ shadowed.__call__ = lambda x: x * 99
+
+ self.assertEqual(20, shadowed(5), "the type's __call__ is what Python
runs")
+ self.assertEqual(self._vals(shadowed, LongType(), "a long", [(5,)]),
[20])
+
+ def test_udf_transpile_refuses_a_class_object(self):
+ # Calling a CLASS whose metaclass is ``type`` runs ``__init__`` and
yields an
+ # instance, so its own ``__call__`` is never the body -- but that is
the body
+ # the old ``getattr(func, "__call__")`` found. Pinned for the
+ # dynamically-created case too, where ``getsource`` cannot fall back
to a
+ # ``ClassDef``. (A class with a custom metaclass IS callable through
+ # ``Meta.__call__``, and is resolved through it rather than refused.)
+ class Lexical:
+ def __init__(self, x):
+ self.v = x * 7
+
+ def __call__(self, x):
+ return x * 1000
+
+ def impl(self, x):
+ return x * 1000
+
+ Dynamic = type("Dynamic", (), {"__call__": impl})
+
+ with self.sql_conf(_TRANSPILE_ON):
+ for label, cls in [("lexical class", Lexical), ("type() class",
Dynamic)]:
+ with self.subTest(case=label):
+ self.assertEqual([], UserDefinedFunction(cls,
LongType()).transpiled)
+
+ def
test_udf_transpile_strips_a_bound_receiver_by_dispatch_not_by_name(self):
+ # The receiver used to be dropped only when literally named ``self``,
so a
+ # bound ``__call__(this, x)`` or ``@classmethod f(cls, x)`` kept it in
the
+ # public parameter list. That declares one parameter too many and
shifts
+ # every ``_udf_param_N``: a two-column call returned column b's value
where
+ # Python raises TypeError.
+ class Recv:
+ def __call__(this, x):
+ return x + 1
+
+ class Meth:
+ def act(this, x):
+ return x + 2
+
+ class Cls:
+ @classmethod
+ def act(kls, x):
+ return x + 3
+
+ # A ``__call__`` that is a classmethod, or one that is ALREADY a bound
method,
+ # also has its receiver spoken for -- the class and the method's own
+ # ``__self__`` respectively. Both used to keep it in the public list,
which
+ # shifts every placeholder. Each expectation below is what Python
returns.
+ class ClsCall:
+ @classmethod
+ def __call__(kls, x):
+ return x + 4
+
+ class Helper:
+ def impl(self, x):
+ return x + 5
+
+ class BoundCall:
+ __call__ = Helper().impl
+
+ # And the two compose: a ``staticmethod`` prepends nothing, but the
method it
+ # wraps is already bound, so one parameter is still spoken for.
Counting only
+ # what the descriptor prepends declared a parameter too many here.
+ class StaticBound:
+ __call__ = staticmethod(Helper().impl)
+
+ for label, func, expected in [
+ ("__call__ receiver not named self", Recv(), 6),
+ ("bound method receiver not named self", Meth().act, 7),
+ ("classmethod receiver not named self", Cls.act, 8),
+ ("classmethod as __call__", ClsCall(), 9),
+ ("already-bound method as __call__", BoundCall(), 10),
+ ("staticmethod over a bound method", StaticBound(), 10),
+ ]:
+ with self.subTest(case=label):
+ self.assertEqual(func(5), expected, "fixture must match
Python's own answer")
+ self.assertEqual(self._vals(func, LongType(), "a long",
[(5,)]), [expected])
+
+ # Both at once is a callable Python itself rejects: the classmethod
prepends
+ # the class ON TOP of the method's own receiver, leaving no parameter
for the
+ # call site. Counting one receiver returned a value where Python
raises.
+ class ClassBound:
+ __call__ = classmethod(Helper().impl)
+
+ with self.assertRaises(TypeError):
+ ClassBound()(5)
+ with self.sql_conf(_TRANSPILE_ON):
+ _, reasons = self._fallback_reason(ClassBound())
+ self.assertIn("leaves no parameter for the call site", reasons)
+
+ # The mirror: a ``staticmethod`` ``__call__`` prepends nothing, so its
leading
+ # ``self`` IS supplied at the call site and both parameters are public.
+ # Resolving ``__call__`` on the type (rather than via getattr on the
instance,
+ # which fires the descriptor) hands back the raw ``staticmethod``,
which
+ # carries a ``__wrapped__`` of its own -- so for a while on this
branch the
+ # wraps guard refused every one of them for a decorator that is not
there.
+ class Static:
+ @staticmethod
+ def __call__(self, x):
+ return self + x
+
+ self.assertEqual(14, Static()(5, 9), "staticmethod __call__ binds no
receiver")
+ self.assertEqual(
+ self._vals(Static(), LongType(), "a long, b long", [(5, 9)]),
+ [14],
+ "both parameters come from the call site",
+ )
+
def test_udf_transpile_known_value_divergences(self):
# Transpile but DIVERGE from Python (documented in transpile.py;
pinned so
# a future fix is noticed): unguarded arithmetic on NULL yields NULL
@@ -1566,12 +1903,14 @@ class UDFTranspileUnitTests(ReusedSQLTestCase):
self.assertEqual(
self._vals(nan_gt, BooleanType(), "a double", [(float("nan"),),
(1.0,)]), [True, True]
)
- # `x == "5"` used to be pinned as a coercion divergence (int == "5" ->
- # True). The eq category gate now drops the numeric variant, so on a
- # long column the string option is pruned and the UDF falls back to
- # interpreted Python -- matching Python's cross-type == (always False).
+ # `x == "5"` used to be pinned as a coercion divergence (int == "5" ->
True).
+ # The eq category gate now drops the numeric variant, so on a long
column the
+ # string option is pruned, nothing is left to lower (hence
require_lowered
+ # =False), and the UDF falls back to interpreted Python -- matching
Python's
+ # cross-type == (always False).
self.assertEqual(
- self._vals(eq_strlit, BooleanType(), "a long", [(5,), (3,)]),
[False, False]
+ self._vals(eq_strlit, BooleanType(), "a long", [(5,), (3,)],
require_lowered=False),
+ [False, False],
)
def test_udf_transpile_overflow_and_modulo_zero_raise(self):
diff --git a/python/pyspark/sql/transpile.py b/python/pyspark/sql/transpile.py
index ef4ca9a90894..52d92aedfa99 100644
--- a/python/pyspark/sql/transpile.py
+++ b/python/pyspark/sql/transpile.py
@@ -34,13 +34,24 @@ UDF's parameters (e.g. ``def f(a: int, b: str)``) pins each
category and
keeps the option matrix small; prefer doing so. To bound plan growth,
functions with more than three untyped parameters only emit the
all-numeric and all-string variants.
+
+A lambda is lowered only when its source names it directly and alone: bind it
+to a name (``f = lambda x: x + 1``, annotated if you like) and give it a line
+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.
"""
import ast
-from typing import Any, Callable, List, Optional, Tuple, TYPE_CHECKING
+from typing import Any, Callable, Iterator, List, Optional, Tuple,
TYPE_CHECKING
+import contextlib
import inspect
import itertools
+import sys
import textwrap
+import threading
+import warnings
from pyspark.errors import UnsupportedOperationException
from pyspark.sql.column import Column
from pyspark.sql.types import (
@@ -90,6 +101,15 @@ class AbstractTranspiler(object):
returnType: "DataTypeOrString",
param_categories: Optional[dict] = None,
) -> Optional[Column]:
+ """Lower ``function_ast`` to a :class:`Column`, or return ``None`` to
decline.
+
+ The override point for
``spark.sql.experimental.optimizer.pyTranspilers``.
+
+ ``params`` is the CALLER-FACING parameter list: a receiver already
bound, as
+ on a method or callable instance, has been removed, so ``params[i]``
is the
+ name bound to placeholder ``_udf_param_i`` with no offsetting needed.
It is
+ also the list ``param_categories`` is keyed by.
+ """
pass
@@ -381,10 +401,9 @@ class CatalystTranspiler(AbstractTranspiler):
"category; falling back to interpreted Python"
)
case ast.Name(id=name) if name in params:
- index = params.index(name)
- if params and params[0] == "self":
- index -= 1
- return self._param_categories.get(index, "numeric")
+ # ``params`` is the caller-facing list, so its indexes are
already
+ # the ``_udf_param_N`` / category indexes -- see
``_transpile_func``.
+ return self._param_categories.get(params.index(name),
"numeric")
case ast.BinOp(left=left, op=op, right=right):
lc = self._category(params, left)
rc = self._category(params, right)
@@ -686,24 +705,11 @@ class CatalystTranspiler(AbstractTranspiler):
case ast.Name(id=name, ctx=ast.Load()):
# Insert columns referencing the param indexes for children
if name in params:
- param_index = params.index(name)
- # Special hack for self on callables
- if params[0] == "self":
- # A body that references the receiver itself (e.g.
- # ``return self``) has no column equivalent: ``self``
- # is not an argument at the call site, and offsetting
- # would emit ``_udf_param_-1``, which the JVM builder
- # rejects with an AnalysisException at call
- # construction instead of falling back. Refuse so the
- # UDF stays interpreted.
- if name == "self":
- raise UnsupportedOperationException(
- "references to `self` in a callable's body "
- "are not supported by the transpiler; falling "
- "back to interpreted Python"
- )
- param_index -= 1
- return col(f"_udf_param_{param_index}")
+ # ``params`` excludes any bound receiver (see
``_transpile_func``),
+ # so its indexes ARE the placeholder indexes. A body
referencing
+ # the receiver (``return self``) is not in this list and
so takes
+ # the branch below, which refuses -- there is no column
for it.
+ return col(f"_udf_param_{params.index(name)}")
else:
# TODO (SPARK-55207): Handle assignments, class vars, and
closures
# via scope evaluation.
@@ -863,8 +869,94 @@ def _param_category_combos(function_ast: ast.FunctionDef,
public_params: List[st
return [{i: choice[i] for i in range(n)} for choice in
itertools.product(*candidates)] or [{}]
+def _call_dunder(func: Callable) -> Any:
+ """The ``__call__`` entry from ``func``'s type.
+
+ Not ``getattr(func, "__call__")``, which is wrong in two ways that both
end with
+ lowering a body that never runs: an instance attribute ``obj.__call__ = f``
+ shadows the type's for ``getattr`` but is ignored when ``obj`` is called,
and on
+ a CLASS object it finds the ``__call__`` its instances use while calling
the
+ class runs ``__init__``.
+
+ ``getattr_static`` looks the name up without firing the descriptor
protocol, so
+ deciding what to transpile never runs user code -- a custom descriptor
used as
+ ``__call__`` would otherwise have its ``__get__`` called here.
+
+ Everything comes back undisturbed, so a ``staticmethod`` or ``classmethod``
+ arrives as the descriptor rather than the function inside it -- see
+ ``_call_impl``. There is always something to return: ``getattr_static`` on
a type
+ falls through to the metatype, so the floor is ``type.__call__``.
+ """
+ return inspect.getattr_static(type(func), "__call__")
+
+
+def _call_impl(entry: Any) -> Any:
+ """The function inside a ``staticmethod`` / ``classmethod``, else
``entry`` itself.
+
+ Both get in the way, in opposite directions: they do not forward the
wrapped
+ function's ``__code__``, and they synthesize a ``__wrapped__`` pointing at
it even
+ when no decorator is involved. So asking the descriptor directly finds no
code
+ object and a wraps decorator that is not there -- unwrap before either
question.
+
+ Narrow on purpose: unwrapping any ``__func__`` would follow the attribute
on
+ unrelated callables that expose one, and read the wrong code object.
+ """
+ return entry.__func__ if isinstance(entry, (staticmethod, classmethod))
else entry
+
+
+def _held_code(func: Callable) -> Any:
+ """The code object that runs when ``func`` is called, or ``None``.
+
+ A function or method runs its own ``__code__``; anything else runs its
type's
+ ``__call__``. Used only to ask whether we are holding a lambda.
+ """
+ target = func if (inspect.isfunction(func) or inspect.ismethod(func)) else
_call_dunder(func)
+ return getattr(_call_impl(target), "__code__", None)
+
+
+_WARNINGS_LOCK = threading.Lock()
+
+
[email protected]
+def _syntax_warnings_suppressed() -> Iterator[None]:
+ """Parse without re-emitting, or tripping over, the source's own
SyntaxWarnings.
+
+ The import already reported them. Without this, ``udf()`` repeats the
warning,
+ and under warnings-as-errors the parse raises and lowering silently turns
off.
+ Before 3.12 an invalid escape sequence was a DeprecationWarning, so ignore
that
+ too rather than lose lowering on the oldest Python we support.
+
+ The lock serializes our own use of ``warnings``, whose state is
process-global.
+ It cannot serialize anyone else's: a thread entering ``catch_warnings``
while
+ this is open has its filters restored from our older snapshot on exit, and
+ entering at all bumps the filter version, so a "once"-filtered warning
+ elsewhere in the process can fire again. Both are inherent to the stdlib
API,
+ and are why the parse is the only thing inside here.
+ """
+ with _WARNINGS_LOCK:
+ with warnings.catch_warnings():
+ warnings.filterwarnings("ignore", category=SyntaxWarning)
+ if sys.version_info < (3, 12):
+ warnings.filterwarnings(
+ "ignore", message="invalid escape sequence",
category=DeprecationWarning
+ )
+ yield
+
+
def _get_src_ast_from_func(func: Callable) -> Tuple[Optional[str],
Optional[ast.AST]]:
- """Try and get the AST from a given callable"""
+ """Try and get the AST from a given callable
+
+ KNOWN LIMITATION: this is the source on disk NOW, not necessarily the
source
+ ``func`` was compiled from. ``inspect.getsource`` reads through
``linecache``,
+ which re-reads an edited file while the code object stays as it was at
import,
+ so editing a module in a long-lived driver and then building a UDF from a
+ function imported earlier lowers the NEW body while Python runs the old
one --
+ verified: rewriting ``lambda x: x + 1`` to ``x * 9`` gives Python 6, Spark
45.
+ Closing it needs the parsed node checked against the held code object; it
is
+ not tracked separately, being part of the experimental transpiler
+ (SPARK-54783). Until then, transpilation assumes source files are not
edited
+ underneath a running session.
+ """
# Note: consider maybe dill? (see the JYTHON PR)
# inspect getsource does not work for functions defined in vanilla
# repl, but does for those in files or in ipython.
@@ -872,14 +964,14 @@ def _get_src_ast_from_func(func: Callable) ->
Tuple[Optional[str], Optional[ast.
try:
src = inspect.getsource(func)
src = textwrap.dedent(src).strip()
- ast_info = ast.parse(src)
+ with _syntax_warnings_suppressed():
+ ast_info = ast.parse(src)
except Exception:
try:
- # getattr keeps mypy happy: `__call__` on a bare Callable is
- # not attribute-accessible in the type system.
- src = inspect.getsource(getattr(func, "__call__"))
+ src = inspect.getsource(_call_dunder(func))
src = textwrap.dedent(src).strip()
- ast_info = ast.parse(src)
+ with _syntax_warnings_suppressed():
+ ast_info = ast.parse(src)
except Exception:
# No usable source (REPL/stdin definition, builtin, ...) --
# return cleanly so the caller reports "cannot transpile"
@@ -893,59 +985,97 @@ def _get_parameter_list(node: ast.FunctionDef) ->
list[str]:
return [arg.arg for arg in node.args.args]
-def _get_function_from_ast(body: ast.AST) -> ast.FunctionDef | None:
+def _get_function_from_ast(body: ast.AST, held_code: Any) ->
Tuple[Optional[ast.FunctionDef], str]:
"""
Extract a :class:`ast.FunctionDef` node from an AST produced by
``ast.parse(inspect.getsource(udf_func))``.
Handles the following source patterns (in order):
- * ``f = lambda x: x + 1`` -- lambda bound directly to a name
+ * ``f = lambda x: x + 1`` -- lambda bound to a name, annotated or not
* ``lambda x: x + 1`` -- bare expression (getsource on a raw lambda)
* ``def f(x): ... return x + 1``
* a class with a ``__call__`` method
- Returns ``None`` when no single unambiguous function can be identified --
- notably, a lambda wrapped in a call such as
- ``f = some_wrapper(lambda x: x + 1)`` parses as
``Assign(value=Call(...))``,
- which is not unwrapped here and so falls back to interpreted Python. Local
- class variables are likewise unsupported.
+ ``held_code`` is the code object that runs when the callable is called; a
+ ``co_name`` of ``<lambda>`` is what makes the ambiguity checks below
apply, and
+ its parameter names are what tell a located lambda apart from a rival.
+
+ Returns the node and an empty reason, or ``None`` and why -- paired so no
refusal
+ reaches the caller unexplained.
"""
if not hasattr(body, "body") or not body.body:
- return None
+ return None, "no statement was found in the source read for this
callable"
stmt = body.body[0]
- # Grab the value side of a top level assign (e.g. x = lambda ...)
+ # Grab the value side of a top level assign (e.g. x = lambda ...). An
annotated
+ # binding is the same shape, and the form a typed codebase writes.
if isinstance(stmt, ast.Assign):
stmt = stmt.value
+ elif isinstance(stmt, ast.AnnAssign) and stmt.value is not None:
+ stmt = stmt.value
# Bare ``lambda x: ...`` (when ``inspect.getsource`` returns a raw
# lambda expression at module top level) parses as ``Expr(Lambda)``.
if isinstance(stmt, ast.Expr) and isinstance(stmt.value, ast.Lambda):
stmt = stmt.value
+ # ``inspect.getsource`` works in whole lines, so refuse unless the lambda
located
+ # here IS the one we hold: anything else lowers a body that never runs
+ # (SPARK-58650).
+ if getattr(held_code, "co_name", None) == "<lambda>":
+ if not isinstance(stmt, ast.Lambda):
+ return None, (
+ "the source read for this lambda does not define it as a
statement of "
+ "its own -- it is wrapped in a call or a tuple assignment, or
a "
+ "surrounding definition, or the file has changed since import
and no "
+ "longer holds it -- so which lambda to lower cannot be
determined"
+ )
+ # The located lambda must take the parameters the held one does, or it
is a
+ # different lambda that merely sits where ours was read from. This is
what
+ # separates a lambda nested in the body of the one we hold (fine: it
can never
+ # be the UDF) from one that RETURNED the lambda we hold, as in the
one-line
+ # ``make_adder = lambda n: lambda x: x + n`` -- there the outer lambda
is
+ # located and the inner is held, and lowering the outer would be wrong.
+ located_args = [arg.arg for arg in stmt.args.args]
+ if located_args != list(held_code.co_varnames[:
held_code.co_argcount]):
+ return None, (
+ "the lambda defined in the source read for this one takes
different "
+ f"parameters ({', '.join(located_args) or 'none'}), so it is
not the "
+ "lambda being transpiled -- a lambda returning another lambda
on one "
+ "line, or a file changed since import; put each lambda on its
own line"
+ )
+ # Only lambdas OUTSIDE ``stmt`` are rivals; one in its body cannot be
the UDF,
+ # and the user could not split it onto another line.
+ own = set(map(id, ast.walk(stmt)))
+ if any(id(node) not in own for node in ast.walk(body) if
isinstance(node, ast.Lambda)):
+ return None, (
+ "more than one lambda is visible in the source line(s) this
one was "
+ "read from, and nothing there says which is the UDF, so it is
not "
+ "safe to lower; put each lambda on its own line to transpile
it"
+ )
+
if isinstance(stmt, ast.Lambda):
# Synthesize a one-statement FunctionDef wrapping the lambda body so
# the rest of the transpiler can treat lambdas and ``def`` uniformly.
- # ``ast.FunctionDef``'s overloads in mypy's typeshed require
- # keyword-only ``type_params`` on 3.12+, which doesn't exist at
- # runtime on every Python we support (the field was added in
- # 3.12 -- before that, passing it raises). Drop to ``Any`` so we
- # avoid the overload resolution entirely; constructing the node
- # via keyword args is well-defined at runtime even when the typed
- # overloads disagree.
fn_ctor: Any = ast.FunctionDef
- return fn_ctor(
+ synthesized = fn_ctor(
name="<lambda>",
args=stmt.args,
body=[ast.Return(value=stmt.body)],
decorator_list=[],
)
+ # A node without ``lineno`` cannot be unparsed or compiled; seed from
the
+ # lambda so positions point at real source rather than line 1.
+ return ast.fix_missing_locations(ast.copy_location(synthesized,
stmt)), ""
if isinstance(stmt, ast.FunctionDef):
- return stmt
- return None
+ return stmt, ""
+ return None, (
+ f"the source read for this callable is a {type(stmt).__name__}, which
the "
+ "transpiler cannot reduce to a single function definition"
+ )
def _transpile_func(
@@ -960,8 +1090,8 @@ def _transpile_func(
-------
list of transpiled options (one per backend x input-type variant)
list of errors as strings
- list of positional parameter names (excluding ``self`` for callable
- instances) -- needed so the caller can resolve named-argument
+ list of positional parameter names (excluding a receiver already bound, as
on a
+ method or callable instance) -- needed so the caller can resolve
named-argument
invocations to positional order at call time, since the ``_udf_param_N``
substitution in :class:`UserDefinedPythonFunction` is positional.
list of per-option input-type categories (``"numeric"`` / ``"string"`` per
@@ -999,9 +1129,12 @@ def _transpile_func(
# the WRAPPED function's source (getsource follows ``__wrapped__``),
# while the UDF actually executes the wrapper. Transpiling would
# silently reproduce the wrong behavior, so refuse and fall back.
+ # ``_call_impl`` first: a ``staticmethod`` / ``classmethod`` exposes a
+ # ``__wrapped__`` of its own, so asking the descriptor refuses every
one of
+ # them for a wraps decorator that is not there.
if (
getattr(func, "__wrapped__", None) is not None
- or getattr(getattr(func, "__call__", None), "__wrapped__", None)
is not None
+ or getattr(_call_impl(_call_dunder(func)), "__wrapped__", None) is
not None
):
return (
[],
@@ -1013,13 +1146,14 @@ def _transpile_func(
[],
[],
)
- src, ast = _get_src_ast_from_func(func)
- if ast is None:
+ # Not ``ast``: that name would shadow the module for this whole
function.
+ src, ast_info = _get_src_ast_from_func(func)
+ if ast_info is None:
return ([], ["Error getting ast for function, cannot transpile"],
[], [])
# Get the lambda body and parameters
- function_ast = _get_function_from_ast(ast)
+ function_ast, extraction_error = _get_function_from_ast(ast_info,
_held_code(func))
if function_ast is None:
- return ([], ["Error extracting function body from ast, cannot
transpile"], [], [])
+ return ([], [extraction_error], [], [])
# Default, variadic (``*args`` / ``**kwargs``), keyword-only, and
# positional-only parameters can't be represented by the positional
# ``_udf_param_N`` placeholder scheme: a call site may omit a
@@ -1034,7 +1168,7 @@ def _transpile_func(
or fn_args.kwonlyargs
or fn_args.vararg is not None
or fn_args.kwarg is not None
- or getattr(fn_args, "posonlyargs", [])
+ or fn_args.posonlyargs
):
return (
[],
@@ -1046,28 +1180,55 @@ def _transpile_func(
[],
)
params = _get_parameter_list(function_ast)
- # The transpiler strips a leading ``self`` on the assumption that the
- # source came from a bound ``__call__`` / method whose receiver is not
- # supplied at the call site. A PLAIN function whose first parameter
- # happens to be named ``self`` breaks that assumption: every arg IS
- # supplied at the call site, and stripping would misnumber the
- # ``_udf_param_N`` placeholders (emitting ``_udf_param_-1``). Refuse
- # and fall back rather than guess.
- if params and params[0] == "self" and inspect.isfunction(func):
- return (
- [],
- [
- "plain function with first parameter named 'self' is "
- "ambiguous to the transpiler's self-stripping; falling "
- "back to interpreted Python"
- ],
- [],
- [],
- )
- # Strip ``self`` for the caller-facing param list -- callers will
- # match user-supplied kwargs against this, and the user doesn't
- # name ``self`` at the call site.
- public_params = params[1:] if params and params[0] == "self" else
list(params)
+ # Drop a receiver that is already bound, so what is left is what the
call
+ # site supplies. Decided by HOW ``func`` dispatches, not by the
parameter's
+ # name: a bound ``__call__(this, x)`` or ``@classmethod f(cls, x)``
has a
+ # receiver not named ``self``, while a plain ``def f(self, x)``
supplies its
+ # ``self`` at the call site. Going by the name misnumbered every
+ # ``_udf_param_N`` -- a two-column call on ``__call__(this, x)`` read
column b
+ # for ``x`` and returned a value where Python raises TypeError. Asking
+ # ``inspect.signature`` is both weaker (a ``__signature__`` off by
exactly one
+ # is undetectable) and worse behaved (it runs user code).
+ #
+ # For a callable instance, two things can consume a leading parameter,
and
+ # they compose: what the descriptor prepends when Python looks
``__call__``
+ # up -- the instance for a plain function, the class for a
``classmethod``,
+ # nothing for a ``staticmethod`` or for an already-bound method, whose
+ # ``__get__`` returns itself -- and whatever the callable already has
bound.
+ # Each count below is checked against what Python returns for that
shape.
+ if inspect.isfunction(func):
+ spoken_for = 0
+ elif inspect.ismethod(func):
+ spoken_for = 1
+ else:
+ call_entry = _call_dunder(func)
+ call_target = _call_impl(call_entry)
+ if not (inspect.isfunction(call_target) or
inspect.ismethod(call_target)):
+ # A slot wrapper, property, partial, or other descriptor: what
it
+ # prepends is not knowable from here.
+ return (
+ [],
+ [
+ f"a {type(call_entry).__name__} as __call__ does not
say which "
+ "parameters the call site supplies, so the placeholder
"
+ "positions cannot be assigned"
+ ],
+ [],
+ [],
+ )
+ spoken_for = int(
+ inspect.isfunction(call_entry) or isinstance(call_entry,
classmethod)
+ ) + int(inspect.ismethod(call_target))
+ if spoken_for > 1 or spoken_for > len(params):
+ # Two receivers at once -- a ``classmethod`` over an already-bound
method
+ # prepends the class ON TOP of the method's own ``__self__`` -- or
one with
+ # no parameter to hold it. Python raises for whatever the call
site passes,
+ # so there is nothing correct to lower.
+ return ([], ["callable leaves no parameter for the call site to
bind"], [], [])
+ # Caller-facing params: callers match user-supplied kwargs against
this,
+ # and the receiver is not named at the call site. Everything downstream
+ # indexes off THIS list, so the placeholder numbering needs no offset.
+ public_params = params[spoken_for:]
transpiled: list[Column] = []
input_categories: list[list[str]] = []
errors = []
@@ -1081,7 +1242,7 @@ def _transpile_func(
for combo in combos:
try:
transpiled_column = transpiler._transpile_from_ast(
- src, ast, function_ast, params, returnType, combo
+ src, ast_info, function_ast, public_params,
returnType, combo
)
if transpiled_column is not None:
transpiled.append(transpiled_column)
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
