ueshin opened a new pull request #35184:
URL: https://github.com/apache/spark/pull/35184
### What changes were proposed in this pull request?
Makes `pandas_udf` to take type annotations with future annotations enabled.
### Why are the changes needed?
When using `from __future__ import annotations`, the type hints will be all
strings, then pandas UDF type inference won't work as follows:
```py
>>> from __future__ import annotations
>>> from typing import Union
>>> import pandas as pd
>>> from pyspark.sql.functions import pandas_udf
>>> @pandas_udf("long")
... def plus_one(v: Union[pd.Series, pd.DataFrame]) -> pd.Series:
... return v + 1
Traceback (most recent call last):
...
NotImplementedError: Unsupported signature: (v: 'Union[pd.Series,
pd.DataFrame]') -> 'pd.Series'.
```
### Does this PR introduce _any_ user-facing change?
Yes. Users can use type annotations for pandas UDFs with the future
annotations flag.
```py
>>> from __future__ import annotations
>>> from typing import Union
>>> import pandas as pd
>>> from pyspark.sql.functions import pandas_udf
>>> @pandas_udf("long")
... def plus_one(v: Union[pd.Series, pd.DataFrame]) -> pd.Series:
... return v + 1
...
>>> df = spark.range(10).selectExpr("id", "id as v")
>>> df.select(plus_one(df.v).alias("plus_one")).show()
+--------+
|plus_one|
+--------+
| 1|
| 2|
| 3|
| 4|
| 5|
| 6|
| 7|
| 8|
| 9|
| 10|
+--------+
```
### How was this patch tested?
Added tests with the future annotations enabled.
--
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]
