From 259760a5c5e26e33b2ee46282aeb63e4ea701020 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Fri, 10 May 2024 18:44:53 +0800 Subject: [PATCH] [SPARK-48228][PYTHON][CONNECT][FOLLOWUP] Also apply `_validate_pandas_udf` in MapInXXX ### What changes were proposed in this pull request? Also apply `_validate_pandas_udf` in MapInXXX ### Why are the changes needed? to make sure validation in `pandas_udf` is also applied in MapInXXX ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? ci ### Was this patch authored or co-authored using generative AI tooling? no Closes #46524 from zhengruifeng/missing_check_map_in_xxx. Authored-by: Ruifeng Zheng Signed-off-by: Ruifeng Zheng --- python/pyspark/sql/connect/dataframe.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py index 3c9415adec2dd..ccaaa15f3190c 100644 --- a/python/pyspark/sql/connect/dataframe.py +++ b/python/pyspark/sql/connect/dataframe.py @@ -83,6 +83,7 @@ ) from pyspark.sql.connect.functions import builtin as F from pyspark.sql.pandas.types import from_arrow_schema +from pyspark.sql.pandas.functions import _validate_pandas_udf # type: ignore[attr-defined] if TYPE_CHECKING: @@ -1997,6 +1998,7 @@ def _map_partitions( ) -> ParentDataFrame: from pyspark.sql.connect.udf import UserDefinedFunction + _validate_pandas_udf(func, evalType) udf_obj = UserDefinedFunction( func, returnType=schema,