Skip to content

Commit 7965c91

Browse files
mstewart141HyukjinKwon
authored andcommitted
[SPARK-23569][PYTHON] Allow pandas_udf to work with python3 style type-annotated functions
## What changes were proposed in this pull request? Check python version to determine whether to use `inspect.getargspec` or `inspect.getfullargspec` before applying `pandas_udf` core logic to a function. The former is python2.7 (deprecated in python3) and the latter is python3.x. The latter correctly accounts for type annotations, which are syntax errors in python2.x. ## How was this patch tested? Locally, on python 2.7 and 3.6. Author: Michael (Stu) Stewart <mstewart141@gmail.com> Closes #20728 from mstewart141/pandas_udf_fix.
1 parent a89cdf5 commit 7965c91

2 files changed

Lines changed: 26 additions & 1 deletion

File tree

python/pyspark/sql/tests.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4381,6 +4381,24 @@ def test_timestamp_dst(self):
43814381
result = df.withColumn('time', foo_udf(df.time))
43824382
self.assertEquals(df.collect(), result.collect())
43834383

4384+
@unittest.skipIf(sys.version_info[:2] < (3, 5), "Type hints are supported from Python 3.5.")
4385+
def test_type_annotation(self):
4386+
from pyspark.sql.functions import pandas_udf
4387+
# Regression test to check if type hints can be used. See SPARK-23569.
4388+
# Note that it throws an error during compilation in lower Python versions if 'exec'
4389+
# is not used. Also, note that we explicitly use another dictionary to avoid modifications
4390+
# in the current 'locals()'.
4391+
#
4392+
# Hyukjin: I think it's an ugly way to test issues about syntax specific in
4393+
# higher versions of Python, which we shouldn't encourage. This was the last resort
4394+
# I could come up with at that time.
4395+
_locals = {}
4396+
exec(
4397+
"import pandas as pd\ndef noop(col: pd.Series) -> pd.Series: return col",
4398+
_locals)
4399+
df = self.spark.range(1).select(pandas_udf(f=_locals['noop'], returnType='bigint')('id'))
4400+
self.assertEqual(df.first()[0], 0)
4401+
43844402

43854403
@unittest.skipIf(
43864404
not _have_pandas or not _have_pyarrow,

python/pyspark/sql/udf.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,10 +42,17 @@ def _create_udf(f, returnType, evalType):
4242
PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF):
4343

4444
import inspect
45+
import sys
4546
from pyspark.sql.utils import require_minimum_pyarrow_version
4647

4748
require_minimum_pyarrow_version()
48-
argspec = inspect.getargspec(f)
49+
50+
if sys.version_info[0] < 3:
51+
# `getargspec` is deprecated since python3.0 (incompatible with function annotations).
52+
# See SPARK-23569.
53+
argspec = inspect.getargspec(f)
54+
else:
55+
argspec = inspect.getfullargspec(f)
4956

5057
if evalType == PythonEvalType.SQL_SCALAR_PANDAS_UDF and len(argspec.args) == 0 and \
5158
argspec.varargs is None:

0 commit comments

Comments
 (0)