dongjoon-hyun opened a new pull request, #59751:
URL: https://github.com/apache/spark/pull/59751
### What changes were proposed in this pull request?
This PR aims to fix `ExtractPythonUDFs` to chain two element-wise Python
UDFs only when they have the same `elementwiseNestingDepth`.
SPARK-58728 (#57940) lets `ExtractPythonUDFFromLambda` lift a Python UDF out
of a nested higher-order function lambda into an element-wise UDF over the
whole array. `elementwiseNestingDepth` is the number of array levels that the
Python worker flattens before it calls the function. When a UDF in a nested
lambda iterates over the result of another UDF, the two lifted UDFs have
different depths.
```
transform(values, x -> transform(f(x), y -> g(y)))
f: depth 1 over `values`, because `f(x)` is the argument of the inner
`transform`
g: depth 2 over the result of `f`, because `g` is in the inner lambda
```
Previously, `ExtractPythonUDFs.shouldExtractUDFExpressionTree` chained `g`
with its child `f` because they have the same eval type. However, the Python
worker evaluates a chain as one composed function, `g(f(...))`, and
`ArrowEvalPythonExec` sends only the depth of the outermost UDF of each chain
(`udfs.map(_.elementwiseNestingDepth)`). So, the worker flattened the input of
`f` two levels instead of one.
With this PR, `g` is evaluated by a separate `ArrowEvalPython` over the
result column of `f`. This also avoids evaluating `f` again inside the chain.
### Why are the changes needed?
This is a correctness issue. Depending on the input, the query fails or
returns wrong results silently.
For example, a natural query like `transform(docs, d ->
transform(tokenize(d), t -> normalize(t)))` fails.
```python
from pyspark.sql import functions as sf
from pyspark.sql.functions import udf
df = spark.createDataFrame([(["Hello World", "Foo"],)], "docs array<string>")
tokenize = udf(lambda d: d.split(" "), "array<string>")
normalize = udf(lambda t: t.lower(), "string")
df.select(
sf.transform("docs", lambda d: sf.transform(tokenize(d), lambda t:
normalize(t))).alias("r")
).collect()
```
**BEFORE**
```
pyspark.errors.exceptions.captured.PythonException: An exception was thrown
from the Python worker:
Traceback (most recent call last):
...
File ".../pyspark/eval_handlers/_elementwise.py", line 83, in
_elementwise_input_schema
assert isinstance(data_type, ArrayType)
AssertionError
```
**AFTER**
```
[Row(r=[['hello', 'world'], ['foo']])]
```
The same query fails in the same way with `udf(..., useArrow=False)`,
`pandas_udf`, and `arrow_udf`.
When the outer array element is itself an array, `f` receives the leaf
values instead of each element, so the query can return wrong results silently.
```python
df = spark.createDataFrame([([[1, 2], [3]],)], "values array<array<int>>")
repeat = udf(lambda a: a + a, "array<int>")
plus_one = udf(lambda y: y + 1, "int")
df.select(
sf.transform("values", lambda row: sf.transform(repeat(row), lambda y:
plus_one(y))).alias("r")
).collect()
```
**BEFORE**
```
[Row(r=[[3, 5, None, None], [7, None]])]
```
**AFTER**
```
[Row(r=[[2, 3, 2, 3], [4, 4]])]
```
### Does this PR introduce _any_ user-facing change?
No. This fixes a bug in SPARK-58728, which is not released yet. It exists
only in `master` and `branch-4.x`.
### How was this patch tested?
Pass the CIs with the newly added test cases.
- `ExtractPythonUDFFromLambdaSuite` checks that `f` and `g` are evaluated by
two `ArrowEvalPython` operators, and that UDFs of the same depth are still
chained.
- `test_udf_in_higher_order_function.py` checks the result of the first
example with `udf`, `pandas_udf`, and `arrow_udf`, and the result of the second
example with `udf`. Its Spark Connect parity suite runs the same tests.
Without this fix, the new Scala test fails with `List(List(("f", 1), ("g",
2))) did not equal List(List(("g", 2)), List(("f", 1)))`, and the new Python
tests fail with the `AssertionError` and the wrong results above.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
