[SPARK-53301][PYTHON] Differentiate type hints of Pandas UDF and Arrow UDF

### What changes were proposed in this pull request?
Differentiate type hints of Pandas UDF and Arrow UDF

### Why are the changes needed?
The `arrow_udf` can works with a pandas udf, and the `pandas_udf` can works with a arrow udf, because the eval type inference didn't differentiate the pandas udf and arrow udf.
But this is supposed to fail.

before:
```
In [1]: import pyarrow as pa
   ...:
   ...: from pyspark.sql import functions as sf
   ...: from pyspark.sql.functions import arrow_udf, pandas_udf
   ...:
   ...: df = spark.range(10).withColumn("v", sf.col("id") + 1)
   ...:
   ...:
   ...: pandas_udf("long")
   ...: def multiply_arrow_func(a: pa.Array, b: pa.Array) -> pa.Array:
   ...:     assert isinstance(a, pa.Array)
   ...:     assert isinstance(b, pa.Array)
   ...:     return pa.compute.multiply(a, b)
   ...:
   ...:

In [2]: df.select("id", "v", multiply_arrow_func("id", "v").alias("m")).show()
   ...:
+---+---+---+
| id|  v|  m|
+---+---+---+
|  0|  1|  0|
|  1|  2|  2|
|  2|  3|  6|
|  3|  4| 12|
|  4|  5| 20|
|  5|  6| 30|
|  6|  7| 42|
|  7|  8| 56|
|  8|  9| 72|
|  9| 10| 90|
+---+---+---+
```

after
```
In [2]:    ...: pandas_udf("long")
   ...:    ...: def multiply_arrow_func(a: pa.Array, b: pa.Array) -> pa.Array:
   ...:    ...:     assert isinstance(a, pa.Array)
   ...:    ...:     assert isinstance(b, pa.Array)
   ...:    ...:     return pa.compute.multiply(a, b)
   ...:
---------------------------------------------------------------------------
PySparkNotImplementedError                Traceback (most recent call last)

...

PySparkNotImplementedError: [UNSUPPORTED_SIGNATURE] Unsupported signature: (a: pyarrow.lib.Array, b: pyarrow.lib.Array) -> pyarrow.lib.Array.

```

### Does this PR introduce _any_ user-facing change?
no, arrow_udf is not yet released

### How was this patch tested?
new tests

### Was this patch authored or co-authored using generative AI tooling?
no

Closes apache#52054 from zhengruifeng/arrow_pandas_type_hint.

Authored-by: Ruifeng Zheng <ruifengz@apache.org>
Signed-off-by: Hyukjin Kwon <gurwls223@apache.org>

Author: zhengruifeng

python/pyspark/sql/pandas/functions.py

@@ -807,7 +807,7 @@ def _validate_vectorized_udf(f, evalType, kind: str = "pandas") -> int:
             type_hints = get_type_hints(f)
         except NameError:
             type_hints = {}
-        evalType = infer_eval_type(signature(f), type_hints)
+        evalType = infer_eval_type(signature(f), type_hints, kind)
         assert evalType is not None
 
     if evalType is None:

python/pyspark/sql/pandas/typehints.py

@@ -31,27 +31,19 @@
     )
 
 
-def infer_eval_type(
-    sig: Signature, type_hints: Dict[str, Any]
-) -> Union[
-    "PandasScalarUDFType",
-    "PandasScalarIterUDFType",
-    "PandasGroupedAggUDFType",
-    "ArrowScalarUDFType",
-    "ArrowScalarIterUDFType",
-    "ArrowGroupedAggUDFType",
-]:
+def infer_pandas_eval_type(
+    sig: Signature,
+    type_hints: Dict[str, Any],
+) -> Optional[Union["PandasScalarUDFType", "PandasScalarIterUDFType", "PandasGroupedAggUDFType"]]:
     """
     Infers the evaluation type in :class:`pyspark.util.PythonEvalType` from
     :class:`inspect.Signature` instance and type hints.
     """
-    from pyspark.sql.pandas.functions import PandasUDFType, ArrowUDFType
+    from pyspark.sql.pandas.functions import PandasUDFType
 
     require_minimum_pandas_version()
-    require_minimum_pyarrow_version()
 
     import pandas as pd
-    import pyarrow as pa
 
     annotations = {}
     for param in sig.parameters.values():
@@ -85,9 +77,8 @@ def infer_eval_type(
         )
         for a in parameters_sig
     ) and (return_annotation == pd.Series or return_annotation == pd.DataFrame)
-
-    # pa.Array, ... -> pa.Array
-    is_arrow_array = all(a == pa.Array for a in parameters_sig) and (return_annotation == pa.Array)
+    if is_series_or_frame:
+        return PandasUDFType.SCALAR
 
     # Iterator[Tuple[Series, Frame or Union[DataFrame, Series], ...] -> Iterator[Series or Frame]
     is_iterator_tuple_series_or_frame = (
@@ -110,21 +101,8 @@ def infer_eval_type(
             return_annotation, parameter_check_func=lambda a: a == pd.DataFrame or a == pd.Series
         )
     )
-
-    # Iterator[Tuple[pa.Array, ...] -> Iterator[pa.Array]
-    is_iterator_tuple_array = (
-        len(parameters_sig) == 1
-        and check_iterator_annotation(  # Iterator
-            parameters_sig[0],
-            parameter_check_func=lambda a: check_tuple_annotation(  # Tuple
-                a,
-                parameter_check_func=lambda ta: (ta == Ellipsis or ta == pa.Array),
-            ),
-        )
-        and check_iterator_annotation(
-            return_annotation, parameter_check_func=lambda a: a == pa.Array
-        )
-    )
+    if is_iterator_tuple_series_or_frame:
+        return PandasUDFType.SCALAR_ITER
 
     # Iterator[Series, Frame or Union[DataFrame, Series]] -> Iterator[Series or Frame]
     is_iterator_series_or_frame = (
@@ -143,18 +121,8 @@ def infer_eval_type(
             return_annotation, parameter_check_func=lambda a: a == pd.DataFrame or a == pd.Series
         )
     )
-
-    # Iterator[pa.Array] -> Iterator[pa.Array]
-    is_iterator_array = (
-        len(parameters_sig) == 1
-        and check_iterator_annotation(
-            parameters_sig[0],
-            parameter_check_func=lambda a: a == pa.Array,
-        )
-        and check_iterator_annotation(
-            return_annotation, parameter_check_func=lambda a: a == pa.Array
-        )
-    )
+    if is_iterator_series_or_frame:
+        return PandasUDFType.SCALAR_ITER
 
     # Series, Frame or Union[DataFrame, Series], ... -> Any
     is_series_or_frame_agg = all(
@@ -173,32 +141,141 @@ def infer_eval_type(
         and not check_iterator_annotation(return_annotation)
         and not check_tuple_annotation(return_annotation)
     )
+    if is_series_or_frame_agg:
+        return PandasUDFType.GROUPED_AGG
+
+    return None
+
+
+def infer_arrow_eval_type(
+    sig: Signature, type_hints: Dict[str, Any]
+) -> Optional[Union["ArrowScalarUDFType", "ArrowScalarIterUDFType", "ArrowGroupedAggUDFType"]]:
+    """
+    Infers the evaluation type in :class:`pyspark.util.PythonEvalType` from
+    :class:`inspect.Signature` instance and type hints.
+    """
+    from pyspark.sql.pandas.functions import ArrowUDFType
+
+    require_minimum_pyarrow_version()
+
+    import pyarrow as pa
+
+    annotations = {}
+    for param in sig.parameters.values():
+        if param.annotation is not param.empty:
+            annotations[param.name] = type_hints.get(param.name, param.annotation)
+
+    # Check if all arguments have type hints
+    parameters_sig = [
+        annotations[parameter] for parameter in sig.parameters if parameter in annotations
+    ]
+    if len(parameters_sig) != len(sig.parameters):
+        raise PySparkValueError(
+            errorClass="TYPE_HINT_SHOULD_BE_SPECIFIED",
+            messageParameters={"target": "all parameters", "sig": str(sig)},
+        )
+
+    # Check if the return has a type hint
+    return_annotation = type_hints.get("return", sig.return_annotation)
+    if sig.empty is return_annotation:
+        raise PySparkValueError(
+            errorClass="TYPE_HINT_SHOULD_BE_SPECIFIED",
+            messageParameters={"target": "the return type", "sig": str(sig)},
+        )
+
+    # pa.Array, ... -> pa.Array
+    is_arrow_array = all(a == pa.Array for a in parameters_sig) and (return_annotation == pa.Array)
+    if is_arrow_array:
+        return ArrowUDFType.SCALAR
+
+    # Iterator[Tuple[pa.Array, ...] -> Iterator[pa.Array]
+    is_iterator_tuple_array = (
+        len(parameters_sig) == 1
+        and check_iterator_annotation(  # Iterator
+            parameters_sig[0],
+            parameter_check_func=lambda a: check_tuple_annotation(  # Tuple
+                a,
+                parameter_check_func=lambda ta: (ta == Ellipsis or ta == pa.Array),
+            ),
+        )
+        and check_iterator_annotation(
+            return_annotation, parameter_check_func=lambda a: a == pa.Array
+        )
+    )
+    if is_iterator_tuple_array:
+        return ArrowUDFType.SCALAR_ITER
+
+    # Iterator[pa.Array] -> Iterator[pa.Array]
+    is_iterator_array = (
+        len(parameters_sig) == 1
+        and check_iterator_annotation(
+            parameters_sig[0],
+            parameter_check_func=lambda a: a == pa.Array,
+        )
+        and check_iterator_annotation(
+            return_annotation, parameter_check_func=lambda a: a == pa.Array
+        )
+    )
+    if is_iterator_array:
+        return ArrowUDFType.SCALAR_ITER
 
     # pa.Array, ... -> Any
     is_array_agg = all(a == pa.Array for a in parameters_sig) and (
         return_annotation != pa.Array
         and not check_iterator_annotation(return_annotation)
         and not check_tuple_annotation(return_annotation)
     )
-
-    if is_series_or_frame:
-        return PandasUDFType.SCALAR
-    elif is_arrow_array:
-        return ArrowUDFType.SCALAR
-    elif is_iterator_tuple_series_or_frame or is_iterator_series_or_frame:
-        return PandasUDFType.SCALAR_ITER
-    elif is_iterator_tuple_array or is_iterator_array:
-        return ArrowUDFType.SCALAR_ITER
-    elif is_series_or_frame_agg:
-        return PandasUDFType.GROUPED_AGG
-    elif is_array_agg:
+    if is_array_agg:
         return ArrowUDFType.GROUPED_AGG
+
+    return None
+
+
+def infer_eval_type(
+    sig: Signature,
+    type_hints: Dict[str, Any],
+    kind: str = "all",
+) -> Union[
+    "PandasScalarUDFType",
+    "PandasScalarIterUDFType",
+    "PandasGroupedAggUDFType",
+    "ArrowScalarUDFType",
+    "ArrowScalarIterUDFType",
+    "ArrowGroupedAggUDFType",
+]:
+    """
+    Infers the evaluation type in :class:`pyspark.util.PythonEvalType` from
+    :class:`inspect.Signature` instance and type hints.
+    """
+    assert kind in ["pandas", "arrow", "all"], "kind should be either 'pandas', 'arrow' or 'all'"
+
+    eval_type: Optional[
+        Union[
+            "PandasScalarUDFType",
+            "PandasScalarIterUDFType",
+            "PandasGroupedAggUDFType",
+            "ArrowScalarUDFType",
+            "ArrowScalarIterUDFType",
+            "ArrowGroupedAggUDFType",
+        ]
+    ] = None
+    if kind == "pandas":
+        eval_type = infer_pandas_eval_type(sig, type_hints)
+    elif kind == "arrow":
+        eval_type = infer_arrow_eval_type(sig, type_hints)
     else:
+        eval_type = infer_pandas_eval_type(sig, type_hints) or infer_arrow_eval_type(
+            sig, type_hints
+        )
+
+    if eval_type is None:
         raise PySparkNotImplementedError(
             errorClass="UNSUPPORTED_SIGNATURE",
             messageParameters={"signature": str(sig)},
         )
 
+    return eval_type
+
 
 def check_tuple_annotation(
     annotation: Any, parameter_check_func: Optional[Callable[[Any], bool]] = None

python/pyspark/sql/tests/arrow/test_arrow_udf_typehints.py

@@ -20,6 +20,8 @@
 
 from pyspark.sql import functions as sf
 from pyspark.testing.utils import (
+    have_pandas,
+    pandas_requirement_message,
     have_pyarrow,
     pyarrow_requirement_message,
     have_numpy,
@@ -323,6 +325,19 @@ def func(col: Union["pa.Array", "pa.Array"], *, col2: "pa.Array") -> "pa.Array":
             infer_eval_type(signature(func), get_type_hints(func)), ArrowUDFType.SCALAR
         )
 
+    @unittest.skipIf(not have_pandas, pandas_requirement_message)
+    def test_negative_with_pandas_udf(self):
+        import pandas as pd
+
+        with self.assertRaisesRegex(
+            Exception,
+            "Unsupported signature:.*pandas.core.series.Series.",
+        ):
+
+            @arrow_udf("long")
+            def multiply_pandas(a: pd.Series, b: pd.Series) -> pd.Series:
+                return a * b
+
 
 if __name__ == "__main__":
     from pyspark.sql.tests.arrow.test_arrow_udf_typehints import *  # noqa: #401

python/pyspark/sql/tests/pandas/test_pandas_udf_typehints.py

@@ -377,6 +377,19 @@ def func(col: "Union[pd.Series, pd.DataFrame]", *, col2: "pd.DataFrame") -> "pd.
             infer_eval_type(signature(func), get_type_hints(func)), PandasUDFType.SCALAR
         )
 
+    @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message)
+    def test_negative_with_arrow_udf(self):
+        import pyarrow as pa
+
+        with self.assertRaisesRegex(
+            Exception,
+            "Unsupported signature:.*pyarrow.lib.Array.",
+        ):
+
+            @pandas_udf("long")
+            def multiply_arrow(a: pa.Array, b: pa.Array) -> pa.Array:
+                return pa.compute.multiply(a, b)
+
 
 if __name__ == "__main__":
     from pyspark.sql.tests.pandas.test_pandas_udf_typehints import *  # noqa: #401