[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 apache#46524 from zhengruifeng/missing_check_map_in_xxx. Authored-by: Ruifeng Zheng <[email protected]> Signed-off-by: Ruifeng Zheng <[email protected]>
panbingkun · May 10, 2024 · 259760a · 259760a
1 parent 256a238
commit 259760a
Showing 1 changed file with 2 additions and 0 deletions.
diff --git 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,