You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2023/07/13 05:26:23 UTC
[spark] branch master updated: [SPARK-43964][SQL][PYTHON] Support arrow-optimized Python UDTFs
This is an automated email from the ASF dual-hosted git repository.
gurwls223 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 3dbcd833983 [SPARK-43964][SQL][PYTHON] Support arrow-optimized Python UDTFs
3dbcd833983 is described below
commit 3dbcd8339831e1d211e23d83729460636f74a91e
Author: allisonwang-db <al...@databricks.com>
AuthorDate: Thu Jul 13 14:26:00 2023 +0900
[SPARK-43964][SQL][PYTHON] Support arrow-optimized Python UDTFs
### What changes were proposed in this pull request?
This PR supports arrow-optimized Python user-defined table functions (UDTFs). Similar to Python UDFs, arrow optimization for Python UDTFs can be enabled/disabled using:
1. The `spark.sql.execution.pythonUDTF.arrow.enabled` config (default: True)
2. The `useArrow` parameter of the `udtf` function (default: None)
Note that **arrow optimization is enabled by default** for Python UDTFs. If the user does not have the necessary pandas or pyarrow dependencies, it will fall back to using regular Python UDTF.
### Why are the changes needed?
To make Python UDTFs more performant.
### Does this PR introduce _any_ user-facing change?
Yes. This PR allows users to create arrow-optimized Python UDTFs.
### How was this patch tested?
New unit tests.
Closes #41867 from allisonwang-db/spark-43964-arrow-udtf.
Authored-by: allisonwang-db <al...@databricks.com>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
.../org/apache/spark/api/python/PythonRunner.scala | 2 +
python/pyspark/errors/error_classes.py | 25 ++
python/pyspark/rdd.py | 2 +
python/pyspark/sql/_typing.pyi | 1 +
python/pyspark/sql/functions.py | 47 ++--
python/pyspark/sql/pandas/serializers.py | 155 ++++++++---
python/pyspark/sql/tests/test_udtf.py | 291 +++++++++++++++++++--
python/pyspark/sql/udtf.py | 119 ++++++++-
python/pyspark/worker.py | 145 +++++++---
.../spark/sql/catalyst/expressions/PythonUDF.scala | 3 +-
.../plans/logical/pythonLogicalOperators.scala | 19 ++
.../org/apache/spark/sql/internal/SQLConf.scala | 7 +
.../spark/sql/execution/SparkStrategies.scala | 3 +
.../execution/python/ArrowEvalPythonUDTFExec.scala | 94 +++++++
.../execution/python/ArrowPythonUDTFRunner.scala | 68 +++++
.../execution/python/BatchEvalPythonUDTFExec.scala | 116 ++------
...thonUDTFExec.scala => EvalPythonUDTFExec.scala} | 110 ++------
.../sql/execution/python/ExtractPythonUDFs.scala | 8 +-
.../sql/execution/python/PythonArrowInput.scala | 10 +-
.../python/UserDefinedPythonFunction.scala | 2 +
.../apache/spark/sql/IntegratedUDFTestUtils.scala | 2 +
.../sql/execution/python/PythonUDTFSuite.scala | 38 +++
22 files changed, 968 insertions(+), 299 deletions(-)
diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
index 90225208499..ffb10985768 100644
--- a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
+++ b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
@@ -57,6 +57,7 @@ private[spark] object PythonEvalType {
val SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE = 208
val SQL_TABLE_UDF = 300
+ val SQL_ARROW_TABLE_UDF = 301
def toString(pythonEvalType: Int): String = pythonEvalType match {
case NON_UDF => "NON_UDF"
@@ -72,6 +73,7 @@ private[spark] object PythonEvalType {
case SQL_MAP_ARROW_ITER_UDF => "SQL_MAP_ARROW_ITER_UDF"
case SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE => "SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE"
case SQL_TABLE_UDF => "SQL_TABLE_UDF"
+ case SQL_ARROW_TABLE_UDF => "SQL_ARROW_TABLE_UDF"
}
}
diff --git a/python/pyspark/errors/error_classes.py b/python/pyspark/errors/error_classes.py
index e87d37c63e7..57447d56892 100644
--- a/python/pyspark/errors/error_classes.py
+++ b/python/pyspark/errors/error_classes.py
@@ -202,6 +202,11 @@ ERROR_CLASSES_JSON = """
" must be set to `true` to enable Python profile."
]
},
+ "INVALID_ARROW_UDTF_RETURN_TYPE" : {
+ "message" : [
+ "The return type of the arrow-optimized Python UDTF should be of type 'pandas.DataFrame', but the function returned a value of type <type_name> with value: <value>."
+ ]
+ },
"INVALID_BROADCAST_OPERATION": {
"message": [
"Broadcast can only be <operation> in driver."
@@ -262,6 +267,16 @@ ERROR_CLASSES_JSON = """
"Eval type for UDF must be <eval_type>."
]
},
+ "INVALID_UDTF_EVAL_TYPE" : {
+ "message" : [
+ "The eval type for the UDTF '<name>' is invalid. It must be one of <eval_type>."
+ ]
+ },
+ "INVALID_UDTF_NO_EVAL" : {
+ "message" : [
+ "The UDTF '<name>' is invalid. It does not implement the required 'eval' method. Please implement the 'eval' method in '<name>' and try again."
+ ]
+ },
"INVALID_WHEN_USAGE": {
"message": [
"when() can only be applied on a Column previously generated by when() function, and cannot be applied once otherwise() is applied."
@@ -643,6 +658,16 @@ ERROR_CLASSES_JSON = """
"Expected <expected> values for `<item>`, got <actual>."
]
},
+ "UDTF_EXEC_ERROR" : {
+ "message" : [
+ "User defined table function encountered an error in the '<method_name>' method: <error>"
+ ]
+ },
+ "UDTF_RETURN_SCHEMA_MISMATCH" : {
+ "message" : [
+ "The number of columns in the result does not match the specified schema. Expected column count: <expected>, Actual column count: <actual>. Please make sure the values returned by the function have the same number of columns as specified in the output schema."
+ ]
+ },
"UNEXPECTED_RESPONSE_FROM_SERVER" : {
"message" : [
"Unexpected response from iterator server."
diff --git a/python/pyspark/rdd.py b/python/pyspark/rdd.py
index 30d6c1556e7..8ea9a310222 100644
--- a/python/pyspark/rdd.py
+++ b/python/pyspark/rdd.py
@@ -114,6 +114,7 @@ if TYPE_CHECKING:
AtomicValue,
RowLike,
SQLArrowBatchedUDFType,
+ SQLArrowTableUDFType,
SQLBatchedUDFType,
SQLTableUDFType,
)
@@ -159,6 +160,7 @@ class PythonEvalType:
SQL_GROUPED_MAP_PANDAS_UDF_WITH_STATE: "PandasGroupedMapUDFWithStateType" = 208
SQL_TABLE_UDF: "SQLTableUDFType" = 300
+ SQL_ARROW_TABLE_UDF: "SQLArrowTableUDFType" = 301
def portable_hash(x: Hashable) -> int:
diff --git a/python/pyspark/sql/_typing.pyi b/python/pyspark/sql/_typing.pyi
index 7b5c206c98f..3d095f55709 100644
--- a/python/pyspark/sql/_typing.pyi
+++ b/python/pyspark/sql/_typing.pyi
@@ -59,6 +59,7 @@ RowLike = TypeVar("RowLike", List[Any], Tuple[Any, ...], pyspark.sql.types.Row)
SQLBatchedUDFType = Literal[100]
SQLArrowBatchedUDFType = Literal[101]
SQLTableUDFType = Literal[300]
+SQLArrowTableUDFType = Literal[301]
class SupportsOpen(Protocol):
def open(self, partition_id: int, epoch_id: int) -> bool: ...
diff --git a/python/pyspark/sql/functions.py b/python/pyspark/sql/functions.py
index b7d1204deef..d7a2f529fa6 100644
--- a/python/pyspark/sql/functions.py
+++ b/python/pyspark/sql/functions.py
@@ -49,7 +49,7 @@ from pyspark.sql.types import ArrayType, DataType, StringType, StructType, _from
# Keep UserDefinedFunction import for backwards compatible import; moved in SPARK-22409
from pyspark.sql.udf import UserDefinedFunction, _create_py_udf # noqa: F401
-from pyspark.sql.udtf import UserDefinedTableFunction, _create_udtf
+from pyspark.sql.udtf import UserDefinedTableFunction, _create_py_udtf
# Keep pandas_udf and PandasUDFType import for backwards compatible import; moved in SPARK-28264
from pyspark.sql.pandas.functions import pandas_udf, PandasUDFType # noqa: F401
@@ -15514,6 +15514,7 @@ def udtf(
cls: Optional[Type] = None,
*,
returnType: Union[StructType, str],
+ useArrow: Optional[bool] = None,
) -> Union[UserDefinedTableFunction, functools.partial]:
"""Creates a user defined table function (UDTF).
@@ -15526,29 +15527,20 @@ def udtf(
returnType : :class:`pyspark.sql.types.StructType` or str
the return type of the user-defined table function. The value can be either a
:class:`pyspark.sql.types.StructType` object or a DDL-formatted struct type string.
+ useArrow : bool or None, optional
+ whether to use Arrow to optimize the (de)serializations. When it's set to None, the
+ Spark config "spark.sql.execution.pythonUDTF.arrow.enabled" is used.
Examples
--------
- Implement the UDTF class
+ Implement the UDTF class and create a UDTF:
>>> class TestUDTF:
... def eval(self, *args: Any):
... yield "hello", "world"
-
- Create the UDTF
-
+ ...
>>> from pyspark.sql.functions import udtf
>>> test_udtf = udtf(TestUDTF, returnType="c1: string, c2: string")
-
- Create the UDTF using the decorator
-
- >>> @udtf(returnType="c1: int, c2: int")
- ... class PlusOne:
- ... def eval(self, x: int):
- ... yield x, x + 1
-
- Invoke the UDTF
-
>>> test_udtf().show()
+-----+-----+
| c1| c2|
@@ -15556,8 +15548,13 @@ def udtf(
|hello|world|
+-----+-----+
- Invoke the UDTF with parameters
+ UDTF can also be created using the decorator syntax:
+ >>> @udtf(returnType="c1: int, c2: int")
+ ... class PlusOne:
+ ... def eval(self, x: int):
+ ... yield x, x + 1
+ ...
>>> from pyspark.sql.functions import lit
>>> PlusOne(lit(1)).show()
+---+---+
@@ -15566,6 +15563,20 @@ def udtf(
| 1| 2|
+---+---+
+ Arrow optimization can be explicitly enabled when creating UDTFs:
+
+ >>> @udtf(returnType="c1: int, c2: int", useArrow=True)
+ ... class ArrowPlusOne:
+ ... def eval(self, x: int):
+ ... yield x, x + 1
+ ...
+ >>> ArrowPlusOne(lit(1)).show()
+ +---+---+
+ | c1| c2|
+ +---+---+
+ | 1| 2|
+ +---+---+
+
Notes
-----
User-defined table functions (UDTFs) are considered deterministic by default.
@@ -15599,9 +15610,9 @@ def udtf(
User-defined table functions do not accept keyword arguments on the calling side.
"""
if cls is None:
- return functools.partial(_create_udtf, returnType=returnType)
+ return functools.partial(_create_py_udtf, returnType=returnType, useArrow=useArrow)
else:
- return _create_udtf(cls=cls, returnType=returnType)
+ return _create_py_udtf(cls=cls, returnType=returnType, useArrow=useArrow)
def _test() -> None:
diff --git a/python/pyspark/sql/pandas/serializers.py b/python/pyspark/sql/pandas/serializers.py
index 4c095249957..f835ea57b77 100644
--- a/python/pyspark/sql/pandas/serializers.py
+++ b/python/pyspark/sql/pandas/serializers.py
@@ -360,6 +360,56 @@ class ArrowStreamPandasUDFSerializer(ArrowStreamPandasSerializer):
)
return s
+ def _create_struct_array(self, df, arrow_struct_type):
+ """
+ Create an Arrow StructArray from the given pandas.DataFrame and arrow struct type.
+
+ Parameters
+ ----------
+ df : pandas.DataFrame
+ A pandas DataFrame
+ arrow_struct_type : pyarrow.DataType
+ pyarrow struct type
+
+ Returns
+ -------
+ pyarrow.Array
+ """
+ import pyarrow as pa
+
+ # Input partition and result pandas.DataFrame empty, make empty Arrays with struct
+ if len(df) == 0 and len(df.columns) == 0:
+ arrs_names = [
+ (pa.array([], type=field.type), field.name) for field in arrow_struct_type
+ ]
+ # Assign result columns by schema name if user labeled with strings
+ elif self._assign_cols_by_name and any(isinstance(name, str) for name in df.columns):
+ arrs_names = [
+ (
+ self._create_array(df[field.name], field.type, arrow_cast=self._arrow_cast),
+ field.name,
+ )
+ for field in arrow_struct_type
+ ]
+ # Assign result columns by position
+ else:
+ arrs_names = [
+ # the selected series has name '1', so we rename it to field.name
+ # as the name is used by _create_array to provide a meaningful error message
+ (
+ self._create_array(
+ df[df.columns[i]].rename(field.name),
+ field.type,
+ arrow_cast=self._arrow_cast,
+ ),
+ field.name,
+ )
+ for i, field in enumerate(arrow_struct_type)
+ ]
+
+ struct_arrs, struct_names = zip(*arrs_names)
+ return pa.StructArray.from_arrays(struct_arrs, struct_names)
+
def _create_batch(self, series):
"""
Create an Arrow record batch from the given pandas.Series pandas.DataFrame
@@ -389,44 +439,14 @@ class ArrowStreamPandasUDFSerializer(ArrowStreamPandasSerializer):
arrs = []
for s, t in series:
if self._struct_in_pandas == "dict" and t is not None and pa.types.is_struct(t):
+ # A pandas UDF should return pd.DataFrame when the return type is a struct type.
+ # If it returns a pd.Series, it should throw an error.
if not isinstance(s, pd.DataFrame):
raise PySparkValueError(
"A field of type StructType expects a pandas.DataFrame, "
"but got: %s" % str(type(s))
)
-
- # Input partition and result pandas.DataFrame empty, make empty Arrays with struct
- if len(s) == 0 and len(s.columns) == 0:
- arrs_names = [(pa.array([], type=field.type), field.name) for field in t]
- # Assign result columns by schema name if user labeled with strings
- elif self._assign_cols_by_name and any(isinstance(name, str) for name in s.columns):
- arrs_names = [
- (
- self._create_array(
- s[field.name], field.type, arrow_cast=self._arrow_cast
- ),
- field.name,
- )
- for field in t
- ]
- # Assign result columns by position
- else:
- arrs_names = [
- # the selected series has name '1', so we rename it to field.name
- # as the name is used by _create_array to provide a meaningful error message
- (
- self._create_array(
- s[s.columns[i]].rename(field.name),
- field.type,
- arrow_cast=self._arrow_cast,
- ),
- field.name,
- )
- for i, field in enumerate(t)
- ]
-
- struct_arrs, struct_names = zip(*arrs_names)
- arrs.append(pa.StructArray.from_arrays(struct_arrs, struct_names))
+ arrs.append(self._create_struct_array(s, t))
else:
arrs.append(self._create_array(s, t, arrow_cast=self._arrow_cast))
@@ -454,6 +474,75 @@ class ArrowStreamPandasUDFSerializer(ArrowStreamPandasSerializer):
return "ArrowStreamPandasUDFSerializer"
+class ArrowStreamPandasUDTFSerializer(ArrowStreamPandasUDFSerializer):
+ """
+ Serializer used by Python worker to evaluate Arrow-optimized Python UDTFs.
+ """
+
+ def __init__(self, timezone, safecheck):
+ super(ArrowStreamPandasUDTFSerializer, self).__init__(
+ timezone=timezone,
+ safecheck=safecheck,
+ # The output pandas DataFrame's columns are unnamed.
+ assign_cols_by_name=False,
+ # Set to 'False' to avoid converting struct type inputs into a pandas DataFrame.
+ df_for_struct=False,
+ # Defines how struct type inputs are converted. If set to "row", struct type inputs
+ # are converted into Rows. Without this setting, a struct type input would be treated
+ # as a dictionary. For example, for named_struct('name', 'Alice', 'age', 1),
+ # if struct_in_pandas="dict", it becomes {"name": "Alice", "age": 1}
+ # if struct_in_pandas="row", it becomes Row(name="Alice", age=1)
+ struct_in_pandas="row",
+ # When dealing with array type inputs, Arrow converts them into numpy.ndarrays.
+ # To ensure consistency across regular and arrow-optimized UDTFs, we further
+ # convert these numpy.ndarrays into Python lists.
+ ndarray_as_list=True,
+ # Enables explicit casting for mismatched return types of Arrow Python UDTFs.
+ arrow_cast=True,
+ )
+
+ def _create_batch(self, series):
+ """
+ Create an Arrow record batch from the given pandas.Series pandas.DataFrame
+ or list of Series or DataFrame, with optional type.
+
+ Parameters
+ ----------
+ series : pandas.Series or pandas.DataFrame or list
+ A single series or dataframe, list of series or dataframe,
+ or list of (series or dataframe, arrow_type)
+
+ Returns
+ -------
+ pyarrow.RecordBatch
+ Arrow RecordBatch
+ """
+ import pandas as pd
+ import pyarrow as pa
+
+ # Make input conform to [(series1, type1), (series2, type2), ...]
+ if not isinstance(series, (list, tuple)) or (
+ len(series) == 2 and isinstance(series[1], pa.DataType)
+ ):
+ series = [series]
+ series = ((s, None) if not isinstance(s, (list, tuple)) else s for s in series)
+
+ arrs = []
+ for s, t in series:
+ if not isinstance(s, pd.DataFrame):
+ raise PySparkValueError(
+ "Output of an arrow-optimized Python UDTFs expects "
+ f"a pandas.DataFrame but got: {type(s)}"
+ )
+
+ arrs.append(self._create_struct_array(s, t))
+
+ return pa.RecordBatch.from_arrays(arrs, ["_%d" % i for i in range(len(arrs))])
+
+ def __repr__(self):
+ return "ArrowStreamPandasUDTFSerializer"
+
+
class CogroupUDFSerializer(ArrowStreamPandasUDFSerializer):
def load_stream(self, stream):
"""
diff --git a/python/pyspark/sql/tests/test_udtf.py b/python/pyspark/sql/tests/test_udtf.py
index 43ab0795042..e4db542cbb7 100644
--- a/python/pyspark/sql/tests/test_udtf.py
+++ b/python/pyspark/sql/tests/test_udtf.py
@@ -21,13 +21,25 @@ from typing import Iterator
from py4j.protocol import Py4JJavaError
-from pyspark.errors import PythonException, AnalysisException
-from pyspark.sql.functions import lit, udtf
+from pyspark.errors import (
+ PySparkAttributeError,
+ PythonException,
+ PySparkTypeError,
+ AnalysisException,
+)
+from pyspark.rdd import PythonEvalType
+from pyspark.sql.functions import lit, udf, udtf
from pyspark.sql.types import Row
-from pyspark.testing.sqlutils import ReusedSQLTestCase
+from pyspark.testing.sqlutils import (
+ have_pandas,
+ have_pyarrow,
+ pandas_requirement_message,
+ pyarrow_requirement_message,
+ ReusedSQLTestCase,
+)
-class UDTFTestsMixin:
+class BaseUDTFTestsMixin:
def test_simple_udtf(self):
class TestUDTF:
def eval(self):
@@ -369,19 +381,115 @@ class UDTFTestsMixin:
):
TestUDTF(rand(0) * 100).collect()
- def test_udtf_no_eval(self):
- @udtf(returnType="a: int, b: int")
+ def test_udtf_with_struct_input_type(self):
+ @udtf(returnType="x: string")
class TestUDTF:
- def run(self, a: int):
- yield a, a + 1
+ def eval(self, person):
+ yield f"{person.name}: {person.age}",
+
+ self.spark.udtf.register("test_udtf", TestUDTF)
+ self.assertEqual(
+ self.spark.sql(
+ "select * from test_udtf(named_struct('name', 'Alice', 'age', 1))"
+ ).collect(),
+ [Row(x="Alice: 1")],
+ )
+ def test_udtf_with_array_input_type(self):
+ @udtf(returnType="x: string")
+ class TestUDTF:
+ def eval(self, args):
+ yield str(args),
+
+ self.spark.udtf.register("test_udtf", TestUDTF)
+ self.assertEqual(
+ self.spark.sql("select * from test_udtf(array(1, 2, 3))").collect(),
+ [Row(x="[1, 2, 3]")],
+ )
+
+ def test_udtf_with_map_input_type(self):
+ @udtf(returnType="x: string")
+ class TestUDTF:
+ def eval(self, m):
+ yield str(m),
+
+ self.spark.udtf.register("test_udtf", TestUDTF)
+ self.assertEqual(
+ self.spark.sql("select * from test_udtf(map('key', 'value'))").collect(),
+ [Row(x="{'key': 'value'}")],
+ )
+
+ def test_udtf_with_struct_output_types(self):
+ @udtf(returnType="x: struct<a:int,b:int>")
+ class TestUDTF:
+ def eval(self, x: int):
+ yield {"a": x, "b": x + 1},
+
+ self.assertEqual(TestUDTF(lit(1)).collect(), [Row(x=Row(a=1, b=2))])
+
+ def test_udtf_with_array_output_types(self):
+ @udtf(returnType="x: array<int>")
+ class TestUDTF:
+ def eval(self, x: int):
+ yield [x, x + 1, x + 2],
+
+ self.assertEqual(TestUDTF(lit(1)).collect(), [Row(x=[1, 2, 3])])
+
+ def test_udtf_with_map_output_types(self):
+ @udtf(returnType="x: map<int,string>")
+ class TestUDTF:
+ def eval(self, x: int):
+ yield {x: str(x)},
+
+ self.assertEqual(TestUDTF(lit(1)).collect(), [Row(x={1: "1"})])
+
+ def test_udtf_with_pandas_input_type(self):
+ import pandas as pd
+
+ @udtf(returnType="corr: double")
+ class TestUDTF:
+ def eval(self, s1: pd.Series, s2: pd.Series):
+ yield s1.corr(s2)
+
+ self.spark.udtf.register("test_udtf", TestUDTF)
+ # TODO(SPARK-43968): check during compile time instead of runtime
with self.assertRaisesRegex(
- PythonException,
- "Failed to execute the user defined table function because it has not "
- "implemented the 'eval' method. Please add the 'eval' method and try the "
- "query again.",
+ PythonException, "AttributeError: 'int' object has no attribute 'corr'"
):
- TestUDTF(lit(1)).collect()
+ self.spark.sql(
+ "select * from values (1, 2), (2, 3) t(a, b), lateral test_udtf(a, b)"
+ ).collect()
+
+ def test_udtf_register_error(self):
+ @udf
+ def upper(s: str):
+ return s.upper()
+
+ with self.assertRaises(PySparkTypeError) as e:
+ self.spark.udtf.register("test_udf", upper)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_EVAL_TYPE",
+ message_parameters={
+ "name": "test_udf",
+ "eval_type": "SQL_TABLE_UDF, SQL_ARROW_TABLE_UDF",
+ },
+ )
+
+ def test_udtf_no_eval(self):
+ with self.assertRaises(PySparkAttributeError) as e:
+
+ @udtf(returnType="a: int, b: int")
+ class TestUDTF:
+ def run(self, a: int):
+ yield a, a + 1
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_NO_EVAL",
+ message_parameters={"name": "TestUDTF"},
+ )
def test_udtf_with_no_handler_class(self):
err_msg = "the function handler must be a class"
@@ -589,12 +697,165 @@ class UDTFTestsMixin:
],
)
+ def test_docstring(self):
+ class TestUDTF:
+ """A UDTF for test."""
+
+ def __init__(self):
+ """Initialize the UDTF"""
+ ...
+
+ def eval(self, x: int):
+ """Evaluate the input row."""
+ yield x + 1,
+
+ def terminate(self):
+ """Terminate the UDTF."""
+ ...
+
+ cls = udtf(TestUDTF, returnType="y: int").func
+ self.assertIn("A UDTF for test", cls.__doc__)
+ self.assertIn("Initialize the UDTF", cls.__init__.__doc__)
+ self.assertIn("Evaluate the input row", cls.eval.__doc__)
+ self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
-class UDTFTests(UDTFTestsMixin, ReusedSQLTestCase):
+
+class UDTFTests(BaseUDTFTestsMixin, ReusedSQLTestCase):
@classmethod
def setUpClass(cls):
super(UDTFTests, cls).setUpClass()
- cls.spark.conf.set("spark.sql.execution.pythonUDF.arrow.enabled", "false")
+ cls.spark.conf.set("spark.sql.execution.pythonUDTF.arrow.enabled", "false")
+
+ @classmethod
+ def tearDownClass(cls):
+ try:
+ cls.spark.conf.unset("spark.sql.execution.pythonUDTF.arrow.enabled")
+ finally:
+ super(UDTFTests, cls).tearDownClass()
+
+
+@unittest.skipIf(
+ not have_pandas or not have_pyarrow, pandas_requirement_message or pyarrow_requirement_message
+)
+class UDTFArrowTestsMixin(BaseUDTFTestsMixin):
+ def test_eval_type(self):
+ def upper(x: str):
+ return upper(x)
+
+ class TestUDTF:
+ def eval(self, x: str):
+ return upper(x)
+
+ self.assertEqual(
+ udtf(TestUDTF, returnType="x: string", useArrow=False).evalType,
+ PythonEvalType.SQL_TABLE_UDF,
+ )
+
+ self.assertEqual(
+ udtf(TestUDTF, returnType="x: string", useArrow=True).evalType,
+ PythonEvalType.SQL_ARROW_TABLE_UDF,
+ )
+
+ def test_udtf_eval_returning_non_tuple(self):
+ class TestUDTF:
+ def eval(self, a: int):
+ yield a
+
+ func = udtf(TestUDTF, returnType="a: int")
+ # When arrow is enabled, it can handle non-tuple return value.
+ self.assertEqual(func(lit(1)).collect(), [Row(a=1)])
+
+ def test_udtf_eval_returning_non_generator(self):
+ class TestUDTF:
+ def eval(self, a: int):
+ return (a,)
+
+ func = udtf(TestUDTF, returnType="a: int")
+ self.assertEqual(func(lit(1)).collect(), [Row(a=1)])
+
+ def test_udtf_eval_with_no_return(self):
+ @udtf(returnType="a: int")
+ class TestUDTF:
+ def eval(self, a: int):
+ ...
+
+ self.assertEqual(TestUDTF(lit(1)).collect(), [])
+
+ @udtf(returnType="a: int")
+ class TestUDTF:
+ def eval(self, a: int):
+ return
+
+ self.assertEqual(TestUDTF(lit(1)).collect(), [])
+
+ def test_udtf_terminate_with_wrong_num_output(self):
+ # The error message for arrow-optimized UDTF is different from regular UDTF.
+ err_msg = "The number of columns in the result does not match the specified schema."
+
+ @udtf(returnType="a: int, b: int")
+ class TestUDTF:
+ def eval(self, a: int):
+ yield a, a + 1
+
+ def terminate(self):
+ yield 1, 2, 3
+
+ with self.assertRaisesRegex(PythonException, err_msg):
+ TestUDTF(lit(1)).show()
+
+ @udtf(returnType="a: int, b: int")
+ class TestUDTF:
+ def eval(self, a: int):
+ yield a, a + 1
+
+ def terminate(self):
+ yield 1,
+
+ with self.assertRaisesRegex(PythonException, err_msg):
+ TestUDTF(lit(1)).show()
+
+ def test_udtf_with_empty_yield(self):
+ @udtf(returnType="a: int")
+ class TestUDTF:
+ def eval(self, a: int):
+ yield
+
+ # Arrow-optimized UDTF can support this.
+ self.assertEqual(TestUDTF(lit(1)).collect(), [Row(a=None)])
+
+ def test_udtf_with_wrong_num_output(self):
+ # The error message for arrow-optimized UDTF is different.
+ err_msg = "The number of columns in the result does not match the specified schema."
+
+ @udtf(returnType="a: int, b: int")
+ class TestUDTF:
+ def eval(self, a: int):
+ yield a,
+
+ with self.assertRaisesRegex(PythonException, err_msg):
+ TestUDTF(lit(1)).collect()
+
+ @udtf(returnType="a: int")
+ class TestUDTF:
+ def eval(self, a: int):
+ yield a, a + 1
+
+ with self.assertRaisesRegex(PythonException, err_msg):
+ TestUDTF(lit(1)).collect()
+
+
+class UDTFArrowTests(UDTFArrowTestsMixin, ReusedSQLTestCase):
+ @classmethod
+ def setUpClass(cls):
+ super(UDTFArrowTests, cls).setUpClass()
+ cls.spark.conf.set("spark.sql.execution.pythonUDTF.arrow.enabled", "true")
+
+ @classmethod
+ def tearDownClass(cls):
+ try:
+ cls.spark.conf.unset("spark.sql.execution.pythonUDTF.arrow.enabled")
+ finally:
+ super(UDTFArrowTests, cls).tearDownClass()
if __name__ == "__main__":
diff --git a/python/pyspark/sql/udtf.py b/python/pyspark/sql/udtf.py
index 3bf7bc977c3..b0f373430be 100644
--- a/python/pyspark/sql/udtf.py
+++ b/python/pyspark/sql/udtf.py
@@ -18,12 +18,15 @@
User-defined table function related classes and functions
"""
import sys
-from typing import Type, TYPE_CHECKING, Optional, Union
+import warnings
+from typing import Iterator, Type, TYPE_CHECKING, Optional, Union
from py4j.java_gateway import JavaObject
-from pyspark.errors import PySparkTypeError
+from pyspark.errors import PySparkAttributeError, PySparkTypeError
+from pyspark.rdd import PythonEvalType
from pyspark.sql.column import _to_java_column, _to_seq
+from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version
from pyspark.sql.types import StructType, _parse_datatype_string
from pyspark.sql.udf import _wrap_function
@@ -39,15 +42,99 @@ def _create_udtf(
cls: Type,
returnType: Union[StructType, str],
name: Optional[str] = None,
+ evalType: int = PythonEvalType.SQL_TABLE_UDF,
deterministic: bool = True,
) -> "UserDefinedTableFunction":
- """Create a Python UDTF."""
+ """Create a Python UDTF with the given eval type."""
udtf_obj = UserDefinedTableFunction(
- cls, returnType=returnType, name=name, deterministic=deterministic
+ cls, returnType=returnType, name=name, evalType=evalType, deterministic=deterministic
)
+
return udtf_obj
+def _create_py_udtf(
+ cls: Type,
+ returnType: Union[StructType, str],
+ name: Optional[str] = None,
+ deterministic: bool = True,
+ useArrow: Optional[bool] = None,
+) -> "UserDefinedTableFunction":
+ """Create a regular or an Arrow-optimized Python UDTF."""
+ # Determine whether to create Arrow-optimized UDTFs.
+ if useArrow is not None:
+ arrow_enabled = useArrow
+ else:
+ from pyspark.sql import SparkSession
+
+ session = SparkSession._instantiatedSession
+ arrow_enabled = (
+ session.conf.get("spark.sql.execution.pythonUDTF.arrow.enabled") == "true"
+ if session is not None
+ else True
+ )
+
+ # Create a regular Python UDTF and check for invalid handler class.
+ regular_udtf = _create_udtf(cls, returnType, name, PythonEvalType.SQL_TABLE_UDF, deterministic)
+
+ if arrow_enabled:
+ try:
+ require_minimum_pandas_version()
+ require_minimum_pyarrow_version()
+ except ImportError as e:
+ warnings.warn(
+ f"Arrow optimization for Python UDTFs cannot be enabled: {str(e)}. "
+ f"Falling back to using regular Python UDTFs.",
+ UserWarning,
+ )
+ return regular_udtf
+ return _create_arrow_udtf(regular_udtf)
+ else:
+ return regular_udtf
+
+
+def _create_arrow_udtf(regular_udtf: "UserDefinedTableFunction") -> "UserDefinedTableFunction":
+ """Create an Arrow-optimized Python UDTF."""
+ import pandas as pd
+
+ cls = regular_udtf.func
+
+ class VectorizedUDTF:
+ def __init__(self) -> None:
+ self.func = cls()
+
+ def eval(self, *args: pd.Series) -> Iterator[pd.DataFrame]:
+ if len(args) == 0:
+ yield pd.DataFrame(self.func.eval())
+ else:
+ # Create tuples from the input pandas Series, each tuple
+ # represents a row across all Series.
+ row_tuples = zip(*args)
+ for row in row_tuples:
+ yield pd.DataFrame(self.func.eval(*row))
+
+ def terminate(self) -> Iterator[pd.DataFrame]:
+ if hasattr(self.func, "terminate"):
+ yield pd.DataFrame(self.func.terminate())
+
+ vectorized_udtf = VectorizedUDTF
+ vectorized_udtf.__name__ = cls.__name__
+ vectorized_udtf.__module__ = cls.__module__
+ vectorized_udtf.__doc__ = cls.__doc__
+ vectorized_udtf.__init__.__doc__ = cls.__init__.__doc__
+ vectorized_udtf.eval.__doc__ = getattr(cls, "eval").__doc__
+ if hasattr(cls, "terminate"):
+ getattr(vectorized_udtf, "terminate").__doc__ = getattr(cls, "terminate").__doc__
+
+ return _create_udtf(
+ cls=vectorized_udtf,
+ returnType=regular_udtf.returnType,
+ name=regular_udtf._name,
+ evalType=PythonEvalType.SQL_ARROW_TABLE_UDF,
+ deterministic=regular_udtf.deterministic,
+ )
+
+
class UserDefinedTableFunction:
"""
User-defined table function in Python
@@ -67,6 +154,7 @@ class UserDefinedTableFunction:
func: Type,
returnType: Union[StructType, str],
name: Optional[str] = None,
+ evalType: int = PythonEvalType.SQL_TABLE_UDF,
deterministic: bool = True,
):
@@ -84,8 +172,14 @@ class UserDefinedTableFunction:
self._inputTypes_placeholder = None
self._judtf_placeholder = None
self._name = name or func.__name__
+ self.evalType = evalType
self.deterministic = deterministic
+ if not hasattr(func, "eval"):
+ raise PySparkAttributeError(
+ error_class="INVALID_UDTF_NO_EVAL", message_parameters={"name": self._name}
+ )
+
@property
def returnType(self) -> StructType:
# `_parse_datatype_string` accesses to JVM for parsing a DDL formatted string.
@@ -122,7 +216,7 @@ class UserDefinedTableFunction:
jdt = spark._jsparkSession.parseDataType(self.returnType.json())
assert sc._jvm is not None
judtf = sc._jvm.org.apache.spark.sql.execution.python.UserDefinedPythonTableFunction(
- self._name, wrapped_func, jdt, self.deterministic
+ self._name, wrapped_func, jdt, self.evalType, self.deterministic
)
return judtf
@@ -203,15 +297,24 @@ class UDTFRegistration:
>>> spark.sql("SELECT * FROM VALUES (0, 1), (1, 2) t(x, y), LATERAL plus_one(x)").collect()
[Row(x=0, y=1, c1=0, c2=1), Row(x=1, y=2, c1=1, c2=2)]
"""
+ if f.evalType not in [PythonEvalType.SQL_TABLE_UDF, PythonEvalType.SQL_ARROW_TABLE_UDF]:
+ raise PySparkTypeError(
+ error_class="INVALID_UDTF_EVAL_TYPE",
+ message_parameters={
+ "name": name,
+ "eval_type": "SQL_TABLE_UDF, SQL_ARROW_TABLE_UDF",
+ },
+ )
+
register_udtf = _create_udtf(
- f.func,
+ cls=f.func,
returnType=f.returnType,
name=name,
+ evalType=f.evalType,
deterministic=f.deterministic,
)
- return_udtf = f
self.sparkSession._jsparkSession.udtf().registerPython(name, register_udtf._judtf)
- return return_udtf
+ return register_udtf
def _test() -> None:
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 1d28e6add2e..8c12312da27 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -56,6 +56,7 @@ from pyspark.serializers import (
)
from pyspark.sql.pandas.serializers import (
ArrowStreamPandasUDFSerializer,
+ ArrowStreamPandasUDTFSerializer,
CogroupUDFSerializer,
ArrowStreamUDFSerializer,
ApplyInPandasWithStateSerializer,
@@ -64,7 +65,7 @@ from pyspark.sql.pandas.types import to_arrow_type
from pyspark.sql.types import StructType, _parse_datatype_json_string
from pyspark.util import fail_on_stopiteration, try_simplify_traceback
from pyspark import shuffle
-from pyspark.errors import PySparkRuntimeError
+from pyspark.errors import PySparkRuntimeError, PySparkTypeError
pickleSer = CPickleSerializer()
utf8_deserializer = UTF8Deserializer()
@@ -461,24 +462,49 @@ def assign_cols_by_name(runner_conf):
# ensure the UDTF is valid. This function also prepares a mapper function for applying
# the UDTF logic to input rows.
def read_udtf(pickleSer, infile, eval_type):
+ if eval_type == PythonEvalType.SQL_ARROW_TABLE_UDF:
+ runner_conf = {}
+ # Load conf used for arrow evaluation.
+ num_conf = read_int(infile)
+ for i in range(num_conf):
+ k = utf8_deserializer.loads(infile)
+ v = utf8_deserializer.loads(infile)
+ runner_conf[k] = v
+
+ # NOTE: if timezone is set here, that implies respectSessionTimeZone is True
+ timezone = runner_conf.get("spark.sql.session.timeZone", None)
+ safecheck = (
+ runner_conf.get("spark.sql.execution.pandas.convertToArrowArraySafely", "false").lower()
+ == "true"
+ )
+ ser = ArrowStreamPandasUDTFSerializer(timezone, safecheck)
+ else:
+ # Each row is a group so do not batch but send one by one.
+ ser = BatchedSerializer(CPickleSerializer(), 1)
+
# See `PythonUDTFRunner.PythonUDFWriterThread.writeCommand'
num_arg = read_int(infile)
arg_offsets = [read_int(infile) for _ in range(num_arg)]
handler = read_command(pickleSer, infile)
- return_type = _parse_datatype_json_string(utf8_deserializer.loads(infile))
if not isinstance(handler, type):
raise PySparkRuntimeError(
f"Invalid UDTF handler type. Expected a class (type 'type'), but "
f"got an instance of {type(handler).__name__}."
)
+ return_type = _parse_datatype_json_string(utf8_deserializer.loads(infile))
+ if not type(return_type) == StructType:
+ raise PySparkRuntimeError(
+ f"The return type of a UDTF must be a struct type, but got {type(return_type)}."
+ )
+
# Instantiate the UDTF class.
try:
udtf = handler()
except Exception as e:
raise PySparkRuntimeError(
- f"User defined table function encountered an error in "
- f"the '__init__' method: {str(e)}"
+ error_class="UDTF_EXEC_ERROR",
+ message_parameters={"method_name": "__init__", "error": str(e)},
)
# Validate the UDTF
@@ -489,40 +515,95 @@ def read_udtf(pickleSer, infile, eval_type):
"the query again."
)
- # Wrap the UDTF and convert the results.
- def wrap_udtf(f, return_type):
- if return_type.needConversion():
- toInternal = return_type.toInternal
- return lambda *a: map(toInternal, f(*a))
+ if eval_type == PythonEvalType.SQL_ARROW_TABLE_UDF:
+
+ def wrap_arrow_udtf(f, return_type):
+ arrow_return_type = to_arrow_type(return_type)
+
+ def verify_result(result):
+ import pandas as pd
+
+ if not isinstance(result, pd.DataFrame):
+ raise PySparkTypeError(
+ error_class="INVALID_ARROW_UDTF_RETURN_TYPE",
+ message_parameters={
+ "type_name": type(result).__name_,
+ "value": str(result),
+ },
+ )
+
+ # Check when the dataframe has both rows and columns.
+ if not result.empty or len(result.columns) != 0:
+ if len(result.columns) != len(return_type):
+ raise PySparkRuntimeError(
+ error_class="UDTF_RETURN_SCHEMA_MISMATCH",
+ message_parameters={
+ "expected": str(len(return_type)),
+ "actual": str(len(result.columns)),
+ },
+ )
+
+ # Verify the type and the schema of the result.
+ verify_pandas_result(result, return_type, assign_cols_by_name=False)
+ return result
+
+ return lambda *a: map(lambda res: (res, arrow_return_type), map(verify_result, f(*a)))
+
+ eval = wrap_arrow_udtf(getattr(udtf, "eval"), return_type)
+
+ if hasattr(udtf, "terminate"):
+ terminate = wrap_arrow_udtf(getattr(udtf, "terminate"), return_type)
else:
- return lambda *a: f(*a)
+ terminate = None
+
+ def mapper(_, it):
+ try:
+ for a in it:
+ # The eval function yields an iterator. Each element produced by this
+ # iterator is a tuple in the form of (pandas.DataFrame, arrow_return_type).
+ yield from eval(*[a[o] for o in arg_offsets])
+ finally:
+ if terminate is not None:
+ try:
+ yield from terminate()
+ except BaseException as e:
+ raise PySparkRuntimeError(
+ error_class="UDTF_EXEC_ERROR",
+ message_parameters={"method_name": "terminate", "error": str(e)},
+ )
- eval = wrap_udtf(getattr(udtf, "eval"), return_type)
+ return mapper, None, ser, ser
- if hasattr(udtf, "terminate"):
- terminate = wrap_udtf(getattr(udtf, "terminate"), return_type)
else:
- terminate = None
- # Return an iterator of iterators.
- def mapper(_, it):
- try:
- for a in it:
- yield tuple(eval(*[a[o] for o in arg_offsets]))
- finally:
- if terminate is not None:
- try:
- yield tuple(terminate())
- except BaseException as e:
- raise PySparkRuntimeError(
- f"User defined table function encountered an error in "
- f"the 'terminate' method: {str(e)}"
- )
+ def wrap_udtf(f, return_type):
+ assert return_type.needConversion()
+ toInternal = return_type.toInternal
+ return lambda *a: map(toInternal, f(*a))
- # Each row is a group so do not batch but send one by one.
- ser = BatchedSerializer(CPickleSerializer(), 1)
+ eval = wrap_udtf(getattr(udtf, "eval"), return_type)
+
+ if hasattr(udtf, "terminate"):
+ terminate = wrap_udtf(getattr(udtf, "terminate"), return_type)
+ else:
+ terminate = None
+
+ # Return an iterator of iterators.
+ def mapper(_, it):
+ try:
+ for a in it:
+ yield tuple(eval(*[a[o] for o in arg_offsets]))
+ finally:
+ if terminate is not None:
+ try:
+ yield tuple(terminate())
+ except BaseException as e:
+ raise PySparkRuntimeError(
+ error_class="UDTF_EXEC_ERROR",
+ message_parameters={"method_name": "terminate", "error": str(e)},
+ )
- return mapper, None, ser, ser
+ return mapper, None, ser, ser
def read_udfs(pickleSer, infile, eval_type):
@@ -933,7 +1014,7 @@ def main(infile, outfile):
eval_type = read_int(infile)
if eval_type == PythonEvalType.NON_UDF:
func, profiler, deserializer, serializer = read_command(pickleSer, infile)
- elif eval_type == PythonEvalType.SQL_TABLE_UDF:
+ elif eval_type in (PythonEvalType.SQL_TABLE_UDF, PythonEvalType.SQL_ARROW_TABLE_UDF):
func, profiler, deserializer, serializer = read_udtf(pickleSer, infile, eval_type)
else:
func, profiler, deserializer, serializer = read_udfs(pickleSer, infile, eval_type)
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 829438ccec9..bb4be6d4d11 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
@@ -161,12 +161,11 @@ case class PythonUDTF(
func: PythonFunction,
elementSchema: StructType,
children: Seq[Expression],
+ evalType: Int,
udfDeterministic: Boolean,
resultId: ExprId = NamedExpression.newExprId)
extends UnevaluableGenerator with PythonFuncExpression {
- override def evalType: Int = PythonEvalType.SQL_TABLE_UDF
-
override lazy val canonicalized: Expression = {
val canonicalizedChildren = children.map(_.canonicalized)
// `resultId` can be seen as cosmetic variation in PythonUDTF, as it doesn't affect the result.
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/pythonLogicalOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/pythonLogicalOperators.scala
index f79c360a313..593ed619cb3 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/pythonLogicalOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/pythonLogicalOperators.scala
@@ -204,6 +204,25 @@ case class BatchEvalPythonUDTF(
copy(child = newChild)
}
+/**
+ * A logical plan that evaluates a [[PythonUDTF]] using Apache Arrow.
+ *
+ * @param udtf the user-defined Python function
+ * @param requiredChildOutput the required output of the child plan. It's used for omitting data
+ * generation that will be discarded next by a projection.
+ * @param resultAttrs the output schema of the Python UDTF.
+ * @param child the child plan
+ */
+case class ArrowEvalPythonUDTF(
+ udtf: PythonUDTF,
+ requiredChildOutput: Seq[Attribute],
+ resultAttrs: Seq[Attribute],
+ child: LogicalPlan,
+ evalType: Int) extends BaseEvalPythonUDTF {
+ override protected def withNewChildInternal(newChild: LogicalPlan): ArrowEvalPythonUDTF =
+ copy(child = newChild)
+}
+
/**
* A logical plan that adds a new long column with the name `name` that
* increases one by one. This is for 'distributed-sequence' default index
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 9093b0745a6..4d1d4b45fa3 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -2923,6 +2923,13 @@ object SQLConf {
.booleanConf
.createWithDefault(false)
+ val PYTHON_TABLE_UDF_ARROW_ENABLED =
+ buildConf("spark.sql.execution.pythonUDTF.arrow.enabled")
+ .doc("Enable Arrow optimization for Python UDTFs.")
+ .version("3.5.0")
+ .booleanConf
+ .createWithDefault(true)
+
val PANDAS_GROUPED_MAP_ASSIGN_COLUMNS_BY_NAME =
buildConf("spark.sql.legacy.execution.pandas.groupedMap.assignColumnsByName")
.internal()
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
index a4cca1a248b..0eca4a1b04a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
@@ -749,6 +749,9 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
BatchEvalPythonExec(udfs, output, planLater(child)) :: Nil
case BatchEvalPythonUDTF(udtf, requiredChildOutput, resultAttrs, child) =>
BatchEvalPythonUDTFExec(udtf, requiredChildOutput, resultAttrs, planLater(child)) :: Nil
+ case ArrowEvalPythonUDTF(udtf, requiredChildOutput, resultAttrs, child, evalType) =>
+ ArrowEvalPythonUDTFExec(
+ udtf, requiredChildOutput, resultAttrs, planLater(child), evalType) :: Nil
case _ =>
Nil
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonUDTFExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonUDTFExec.scala
new file mode 100644
index 00000000000..9c9ae0fca2d
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonUDTFExec.scala
@@ -0,0 +1,94 @@
+/*
+ * 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.collection.JavaConverters._
+
+import org.apache.spark.{JobArtifactSet, TaskContext}
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.util.ArrowUtils
+import org.apache.spark.sql.vectorized.{ArrowColumnVector, ColumnarBatch}
+
+/**
+ * A physical plan that evaluates a [[PythonUDTF]] using Apache Arrow.
+ * This is similar to [[ArrowEvalPythonExec]].
+ *
+ * @param udtf the user-defined Python function.
+ * @param requiredChildOutput the required output of the child plan. It's used for omitting data
+ * generation that will be discarded next by a projection.
+ * @param resultAttrs the output schema of the Python UDTF.
+ * @param child the child plan.
+ * @param evalType the Python eval type.
+ */
+case class ArrowEvalPythonUDTFExec(
+ udtf: PythonUDTF,
+ requiredChildOutput: Seq[Attribute],
+ resultAttrs: Seq[Attribute],
+ child: SparkPlan,
+ evalType: Int)
+ extends EvalPythonUDTFExec with PythonSQLMetrics {
+
+ private val batchSize = conf.arrowMaxRecordsPerBatch
+ private val sessionLocalTimeZone = conf.sessionLocalTimeZone
+ private val largeVarTypes = conf.arrowUseLargeVarTypes
+ private val pythonRunnerConf = ArrowUtils.getPythonRunnerConfMap(conf)
+ private[this] val jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid)
+
+ override protected def evaluate(
+ argOffsets: Array[Int],
+ iter: Iterator[InternalRow],
+ schema: StructType,
+ context: TaskContext): Iterator[Iterator[InternalRow]] = {
+
+ val batchIter = if (batchSize > 0) new BatchIterator(iter, batchSize) else Iterator(iter)
+
+ val outputTypes = resultAttrs.map(_.dataType)
+
+ val columnarBatchIter = new ArrowPythonUDTFRunner(
+ udtf,
+ evalType,
+ argOffsets,
+ schema,
+ sessionLocalTimeZone,
+ largeVarTypes,
+ pythonRunnerConf,
+ pythonMetrics,
+ jobArtifactUUID).compute(batchIter, context.partitionId(), context)
+
+ columnarBatchIter.map { batch =>
+ // UDTF returns a StructType column in ColumnarBatch. Flatten the columnar batch here.
+ val columnVector = batch.column(0).asInstanceOf[ArrowColumnVector]
+ val outputVectors = resultAttrs.indices.map(columnVector.getChild)
+ val flattenedBatch = new ColumnarBatch(outputVectors.toArray)
+
+ val actualDataTypes = (0 until flattenedBatch.numCols()).map(
+ i => flattenedBatch.column(i).dataType())
+ assert(outputTypes == actualDataTypes, "Invalid schema from arrow-enabled Python UDTF: " +
+ s"expected ${outputTypes.mkString(", ")}, got ${actualDataTypes.mkString(", ")}")
+
+ flattenedBatch.setNumRows(batch.numRows())
+ flattenedBatch.rowIterator().asScala
+ }
+ }
+
+ override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan =
+ copy(child = newChild)
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala
new file mode 100644
index 00000000000..1dd06c2dc73
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonUDTFRunner.scala
@@ -0,0 +1,68 @@
+/*
+ * 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 java.io.DataOutputStream
+
+import org.apache.spark.api.python._
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.PythonUDTF
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types._
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+/**
+ * Similar to [[ArrowPythonRunner]], but for [[PythonUDTF]]s.
+ */
+class ArrowPythonUDTFRunner(
+ udtf: PythonUDTF,
+ evalType: Int,
+ offsets: Array[Int],
+ protected override val schema: StructType,
+ protected override val timeZoneId: String,
+ protected override val largeVarTypes: Boolean,
+ protected override val workerConf: Map[String, String],
+ val pythonMetrics: Map[String, SQLMetric],
+ jobArtifactUUID: Option[String])
+ extends BasePythonRunner[Iterator[InternalRow], ColumnarBatch](
+ Seq(ChainedPythonFunctions(Seq(udtf.func))), evalType, Array(offsets), jobArtifactUUID)
+ with BasicPythonArrowInput
+ with BasicPythonArrowOutput {
+
+ override protected def writeUDF(
+ dataOut: DataOutputStream,
+ funcs: Seq[ChainedPythonFunctions],
+ argOffsets: Array[Array[Int]]): Unit = {
+ PythonUDTFRunner.writeUDTF(dataOut, udtf, offsets)
+ }
+
+ override val pythonExec: String =
+ SQLConf.get.pysparkWorkerPythonExecutable.getOrElse(
+ funcs.head.funcs.head.pythonExec)
+
+ override val errorOnDuplicatedFieldNames: Boolean = true
+
+ override val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ override val bufferSize: Int = SQLConf.get.pandasUDFBufferSize
+ require(
+ bufferSize >= 4,
+ "Pandas execution requires more than 4 bytes. Please set higher buffer. " +
+ s"Please change '${SQLConf.PANDAS_UDF_BUFFER_SIZE.key}'.")
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonUDTFExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonUDTFExec.scala
index feb3dfd0ba3..7d9fdb3bf48 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonUDTFExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonUDTFExec.scala
@@ -17,24 +17,22 @@
package org.apache.spark.sql.execution.python
-import java.io.{DataOutputStream, File}
+import java.io.DataOutputStream
import java.net.Socket
+import java.nio.charset.StandardCharsets.UTF_8
import scala.collection.JavaConverters._
-import scala.collection.mutable.ArrayBuffer
import net.razorvine.pickle.Unpickler
-import org.apache.spark.{ContextAwareIterator, JobArtifactSet, SparkEnv, TaskContext}
+import org.apache.spark.{JobArtifactSet, SparkEnv, TaskContext}
import org.apache.spark.api.python.{ChainedPythonFunctions, PythonEvalType}
-import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.util.GenericArrayData
-import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode}
+import org.apache.spark.sql.execution.SparkPlan
import org.apache.spark.sql.execution.metric.SQLMetric
-import org.apache.spark.sql.types.{DataType, StructField, StructType}
-import org.apache.spark.util.Utils
+import org.apache.spark.sql.types.StructType
/**
* A physical plan that evaluates a [[PythonUDTF]]. This is similar to [[BatchEvalPythonExec]].
@@ -50,92 +48,15 @@ case class BatchEvalPythonUDTFExec(
requiredChildOutput: Seq[Attribute],
resultAttrs: Seq[Attribute],
child: SparkPlan)
- extends UnaryExecNode with PythonSQLMetrics {
+ extends EvalPythonUDTFExec with PythonSQLMetrics {
private[this] val jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid)
- override def output: Seq[Attribute] = requiredChildOutput ++ resultAttrs
-
- override def producedAttributes: AttributeSet = AttributeSet(resultAttrs)
-
- protected override def doExecute(): RDD[InternalRow] = {
- val inputRDD = child.execute().map(_.copy())
-
- inputRDD.mapPartitions { iter =>
- val context = TaskContext.get()
- val contextAwareIterator = new ContextAwareIterator(context, iter)
-
- // The queue used to buffer input rows so we can drain it to
- // combine input with output from Python.
- val queue = HybridRowQueue(context.taskMemoryManager(),
- new File(Utils.getLocalDir(SparkEnv.get.conf)), child.output.length)
- context.addTaskCompletionListener[Unit] { ctx =>
- queue.close()
- }
-
- // flatten all the arguments
- val allInputs = new ArrayBuffer[Expression]
- val dataTypes = new ArrayBuffer[DataType]
- val argOffsets = udtf.children.map { e =>
- if (allInputs.exists(_.semanticEquals(e))) {
- allInputs.indexWhere(_.semanticEquals(e))
- } else {
- allInputs += e
- dataTypes += e.dataType
- allInputs.length - 1
- }
- }.toArray
- val projection = MutableProjection.create(allInputs.toSeq, child.output)
- projection.initialize(context.partitionId())
- val schema = StructType(dataTypes.zipWithIndex.map { case (dt, i) =>
- StructField(s"_$i", dt)
- }.toArray)
-
- // Add rows to the queue to join later with the result.
- // Also keep track of the number rows added to the queue.
- // This is needed to process extra output rows from the `terminate()` call of the UDTF.
- var count = 0L
- val projectedRowIter = contextAwareIterator.map { inputRow =>
- queue.add(inputRow.asInstanceOf[UnsafeRow])
- count += 1
- projection(inputRow)
- }
-
- val outputRowIterator = evaluate(argOffsets, projectedRowIter, schema, context)
-
- val pruneChildForResult: InternalRow => InternalRow =
- if (child.outputSet == AttributeSet(requiredChildOutput)) {
- identity
- } else {
- UnsafeProjection.create(requiredChildOutput, child.output)
- }
-
- val joined = new JoinedRow
- val resultProj = UnsafeProjection.create(output, output)
-
- outputRowIterator.flatMap { outputRows =>
- // If `count` is greater than zero, it means there are remaining input rows in the queue.
- // In this case, the output rows of the UDTF are joined with the corresponding input row
- // in the queue.
- if (count > 0) {
- val left = queue.remove()
- count -= 1
- joined.withLeft(pruneChildForResult(left))
- }
- // If `count` is zero, it means all input rows have been consumed. Any additional rows
- // from the UDTF are from the `terminate()` call. We leave the left side as the last
- // element of its child output to keep it consistent with the Generate implementation
- // and Hive UDTFs.
- outputRows.map(r => resultProj(joined.withRight(r)))
- }
- }
- }
-
/**
* Evaluates a Python UDTF. It computes the results using the PythonUDFRunner, and returns
* an iterator of internal rows for every input row.
*/
- private def evaluate(
+ override protected def evaluate(
argOffsets: Array[Int],
iter: Iterator[InternalRow],
schema: StructType,
@@ -190,14 +111,23 @@ class PythonUDTFRunner(
new PythonUDFWriterThread(env, worker, inputIterator, partitionIndex, context) {
protected override def writeCommand(dataOut: DataOutputStream): Unit = {
- dataOut.writeInt(argOffsets.length)
- argOffsets.foreach { offset =>
- dataOut.writeInt(offset)
- }
- dataOut.writeInt(udtf.func.command.length)
- dataOut.write(udtf.func.command.toArray)
- writeUTF(udtf.elementSchema.json, dataOut)
+ PythonUDTFRunner.writeUDTF(dataOut, udtf, argOffsets)
}
}
}
}
+
+object PythonUDTFRunner {
+
+ def writeUDTF(dataOut: DataOutputStream, udtf: PythonUDTF, argOffsets: Array[Int]): Unit = {
+ dataOut.writeInt(argOffsets.length)
+ argOffsets.foreach { offset =>
+ dataOut.writeInt(offset)
+ }
+ dataOut.writeInt(udtf.func.command.length)
+ dataOut.write(udtf.func.command.toArray)
+ val schemaBytes = udtf.elementSchema.json.getBytes(UTF_8)
+ dataOut.writeInt(schemaBytes.length)
+ dataOut.write(schemaBytes)
+ }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonUDTFExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvalPythonUDTFExec.scala
similarity index 54%
copy from sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonUDTFExec.scala
copy to sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvalPythonUDTFExec.scala
index feb3dfd0ba3..827b2fc2bb3 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonUDTFExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvalPythonUDTFExec.scala
@@ -17,47 +17,39 @@
package org.apache.spark.sql.execution.python
-import java.io.{DataOutputStream, File}
-import java.net.Socket
+import java.io.File
-import scala.collection.JavaConverters._
import scala.collection.mutable.ArrayBuffer
-import net.razorvine.pickle.Unpickler
-
-import org.apache.spark.{ContextAwareIterator, JobArtifactSet, SparkEnv, TaskContext}
-import org.apache.spark.api.python.{ChainedPythonFunctions, PythonEvalType}
+import org.apache.spark.{ContextAwareIterator, SparkEnv, TaskContext}
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.util.GenericArrayData
-import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode}
-import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.execution.UnaryExecNode
import org.apache.spark.sql.types.{DataType, StructField, StructType}
import org.apache.spark.util.Utils
/**
- * A physical plan that evaluates a [[PythonUDTF]]. This is similar to [[BatchEvalPythonExec]].
- *
- * @param udtf the user-defined Python function
- * @param requiredChildOutput the required output of the child plan. It's used for omitting data
- * generation that will be discarded next by a projection.
- * @param resultAttrs the output schema of the Python UDTF.
- * @param child the child plan
+ * A physical plan that evaluates a [[PythonUDTF]], one partition of tuples at a time.
+ * This is similar to [[EvalPythonExec]].
*/
-case class BatchEvalPythonUDTFExec(
- udtf: PythonUDTF,
- requiredChildOutput: Seq[Attribute],
- resultAttrs: Seq[Attribute],
- child: SparkPlan)
- extends UnaryExecNode with PythonSQLMetrics {
+trait EvalPythonUDTFExec extends UnaryExecNode {
+ def udtf: PythonUDTF
+
+ def requiredChildOutput: Seq[Attribute]
- private[this] val jobArtifactUUID = JobArtifactSet.getCurrentJobArtifactState.map(_.uuid)
+ def resultAttrs: Seq[Attribute]
override def output: Seq[Attribute] = requiredChildOutput ++ resultAttrs
override def producedAttributes: AttributeSet = AttributeSet(resultAttrs)
+ protected def evaluate(
+ argOffsets: Array[Int],
+ iter: Iterator[InternalRow],
+ schema: StructType,
+ context: TaskContext): Iterator[Iterator[InternalRow]]
+
protected override def doExecute(): RDD[InternalRow] = {
val inputRDD = child.execute().map(_.copy())
@@ -130,74 +122,4 @@ case class BatchEvalPythonUDTFExec(
}
}
}
-
- /**
- * Evaluates a Python UDTF. It computes the results using the PythonUDFRunner, and returns
- * an iterator of internal rows for every input row.
- */
- private def evaluate(
- argOffsets: Array[Int],
- iter: Iterator[InternalRow],
- schema: StructType,
- context: TaskContext): Iterator[Iterator[InternalRow]] = {
- EvaluatePython.registerPicklers() // register pickler for Row
-
- // Input iterator to Python.
- val inputIterator = BatchEvalPythonExec.getInputIterator(iter, schema)
-
- // Output iterator for results from Python.
- val outputIterator =
- new PythonUDTFRunner(udtf, argOffsets, pythonMetrics, jobArtifactUUID)
- .compute(inputIterator, context.partitionId(), context)
-
- val unpickle = new Unpickler
-
- // The return type of a UDTF is an array of struct.
- val resultType = udtf.dataType
- val fromJava = EvaluatePython.makeFromJava(resultType)
-
- outputIterator.flatMap { pickedResult =>
- val unpickledBatch = unpickle.loads(pickedResult)
- unpickledBatch.asInstanceOf[java.util.ArrayList[Any]].asScala
- }.map { results =>
- assert(results.getClass.isArray)
- val res = results.asInstanceOf[Array[_]]
- pythonMetrics("pythonNumRowsReceived") += res.length
- fromJava(results).asInstanceOf[GenericArrayData]
- .array.map(_.asInstanceOf[InternalRow]).toIterator
- }
- }
-
- override protected def withNewChildInternal(newChild: SparkPlan): BatchEvalPythonUDTFExec =
- copy(child = newChild)
-}
-
-class PythonUDTFRunner(
- udtf: PythonUDTF,
- argOffsets: Array[Int],
- pythonMetrics: Map[String, SQLMetric],
- jobArtifactUUID: Option[String])
- extends BasePythonUDFRunner(
- Seq(ChainedPythonFunctions(Seq(udtf.func))),
- PythonEvalType.SQL_TABLE_UDF, Array(argOffsets), pythonMetrics, jobArtifactUUID) {
-
- protected override def newWriterThread(
- env: SparkEnv,
- worker: Socket,
- inputIterator: Iterator[Array[Byte]],
- partitionIndex: Int,
- context: TaskContext): WriterThread = {
- new PythonUDFWriterThread(env, worker, inputIterator, partitionIndex, context) {
-
- protected override def writeCommand(dataOut: DataOutputStream): Unit = {
- dataOut.writeInt(argOffsets.length)
- argOffsets.foreach { offset =>
- dataOut.writeInt(offset)
- }
- dataOut.writeInt(udtf.func.command.length)
- dataOut.write(udtf.func.command.toArray)
- writeUTF(udtf.elementSchema.json, dataOut)
- }
- }
- }
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala
index f6993eaf6bd..711da3ff3f1 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ExtractPythonUDFs.scala
@@ -316,8 +316,12 @@ object ExtractPythonUDTFs extends Rule[LogicalPlan] {
case s: Subquery if s.correlated => plan
case _ => plan.transformUpWithPruning(_.containsPattern(GENERATE)) {
- case g @ Generate(func: PythonUDTF, _, _, _, _, child) =>
- BatchEvalPythonUDTF(func, g.requiredChildOutput, g.generatorOutput, child)
+ case g @ Generate(func: PythonUDTF, _, _, _, _, child) => func.evalType match {
+ case PythonEvalType.SQL_TABLE_UDF =>
+ BatchEvalPythonUDTF(func, g.requiredChildOutput, g.generatorOutput, child)
+ case PythonEvalType.SQL_ARROW_TABLE_UDF =>
+ ArrowEvalPythonUDTF(func, g.requiredChildOutput, g.generatorOutput, child, func.evalType)
+ }
}
}
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowInput.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowInput.scala
index c78ea564f18..5c99a3f9808 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowInput.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowInput.scala
@@ -23,7 +23,7 @@ import org.apache.arrow.vector.VectorSchemaRoot
import org.apache.arrow.vector.ipc.ArrowStreamWriter
import org.apache.spark.{SparkEnv, TaskContext}
-import org.apache.spark.api.python.{BasePythonRunner, PythonRDD}
+import org.apache.spark.api.python.{BasePythonRunner, ChainedPythonFunctions, PythonRDD}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.execution.arrow.ArrowWriter
import org.apache.spark.sql.execution.metric.SQLMetric
@@ -54,6 +54,12 @@ private[python] trait PythonArrowInput[IN] { self: BasePythonRunner[IN, _] =>
dataOut: DataOutputStream,
inputIterator: Iterator[IN]): Unit
+ protected def writeUDF(
+ dataOut: DataOutputStream,
+ funcs: Seq[ChainedPythonFunctions],
+ argOffsets: Array[Array[Int]]): Unit =
+ PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets)
+
protected def handleMetadataBeforeExec(stream: DataOutputStream): Unit = {
// Write config for the worker as a number of key -> value pairs of strings
stream.writeInt(workerConf.size)
@@ -73,7 +79,7 @@ private[python] trait PythonArrowInput[IN] { self: BasePythonRunner[IN, _] =>
protected override def writeCommand(dataOut: DataOutputStream): Unit = {
handleMetadataBeforeExec(dataOut)
- PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets)
+ writeUDF(dataOut, funcs, argOffsets)
}
protected override def writeIteratorToStream(dataOut: DataOutputStream): Unit = {
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 08c7c8730d3..4d21349d025 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
@@ -64,6 +64,7 @@ case class UserDefinedPythonTableFunction(
name: String,
func: PythonFunction,
returnType: StructType,
+ pythonEvalType: Int,
udfDeterministic: Boolean) {
def builder(e: Seq[Expression]): LogicalPlan = {
@@ -72,6 +73,7 @@ case class UserDefinedPythonTableFunction(
func = func,
elementSchema = returnType,
children = e,
+ evalType = pythonEvalType,
udfDeterministic = udfDeterministic)
Generate(
udtf,
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/IntegratedUDFTestUtils.scala b/sql/core/src/test/scala/org/apache/spark/sql/IntegratedUDFTestUtils.scala
index 00962e77185..ad52b6bab5b 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/IntegratedUDFTestUtils.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/IntegratedUDFTestUtils.scala
@@ -392,6 +392,7 @@ object IntegratedUDFTestUtils extends SQLHelper {
name: String,
pythonScript: String,
returnType: StructType,
+ evalType: Int = PythonEvalType.SQL_TABLE_UDF,
deterministic: Boolean = true): UserDefinedPythonTableFunction = {
UserDefinedPythonTableFunction(
name = name,
@@ -404,6 +405,7 @@ object IntegratedUDFTestUtils extends SQLHelper {
broadcastVars = List.empty[Broadcast[PythonBroadcast]].asJava,
accumulator = null),
returnType = returnType,
+ pythonEvalType = evalType,
udfDeterministic = deterministic)
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDTFSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDTFSuite.scala
index 9bd0a13a3a1..f38c32862f1 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDTFSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDTFSuite.scala
@@ -17,6 +17,7 @@
package org.apache.spark.sql.execution.python
+import org.apache.spark.api.python.PythonEvalType
import org.apache.spark.sql.{IntegratedUDFTestUtils, QueryTest, Row}
import org.apache.spark.sql.functions.lit
import org.apache.spark.sql.test.SharedSparkSession
@@ -37,11 +38,31 @@ class PythonUDTFSuite extends QueryTest with SharedSparkSession {
| yield a, b, b - a
|""".stripMargin
+ private val arrowPythonScript: String =
+ """
+ |import pandas as pd
+ |class VectorizedUDTF:
+ | def eval(self, a: pd.Series, b: pd.Series):
+ | data = [
+ | [a, b, a + b],
+ | [a, b, a - b],
+ | [a, b, b - a],
+ | ]
+ | yield pd.DataFrame(data)
+ |""".stripMargin
+
private val returnType: StructType = StructType.fromDDL("a int, b int, c int")
private val pythonUDTF: UserDefinedPythonTableFunction =
createUserDefinedPythonTableFunction("SimpleUDTF", pythonScript, returnType)
+ private val arrowPythonUDTF: UserDefinedPythonTableFunction =
+ createUserDefinedPythonTableFunction(
+ "VectorizedUDTF",
+ arrowPythonScript,
+ returnType,
+ evalType = PythonEvalType.SQL_ARROW_TABLE_UDF)
+
test("Simple PythonUDTF") {
assume(shouldTestPythonUDFs)
val df = pythonUDTF(spark, lit(1), lit(2))
@@ -69,4 +90,21 @@ class PythonUDTFSuite extends QueryTest with SharedSparkSession {
Seq(Row(6), Row(null)))
}
}
+
+ test("Arrow optimized UDTF") {
+ assume(shouldTestPandasUDFs)
+ val df = arrowPythonUDTF(spark, lit(1), lit(2))
+ checkAnswer(df, Seq(Row(1, 2, -1), Row(1, 2, 1), Row(1, 2, 3)))
+ }
+
+ test("arrow optimized UDTF with lateral join") {
+ assume(shouldTestPandasUDFs)
+ withTempView("t") {
+ spark.udtf.registerPython("testUDTF", arrowPythonUDTF)
+ Seq((0, 1), (1, 2)).toDF("a", "b").createOrReplaceTempView("t")
+ checkAnswer(
+ sql("SELECT t.*, f.c FROM t, LATERAL testUDTF(a, b) f"),
+ sql("SELECT * FROM t, LATERAL explode(array(a + b, a - b, b - a)) t(c)"))
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org