fix: type infer

Author: Yicong-Huang

python/pyspark/sql/pandas/functions.py

@@ -322,6 +322,7 @@ def calculate(iterator: Iterator[pa.Array]) -> Iterator[pa.Array]:
     pyspark.sql.GroupedData.applyInArrow
     pyspark.sql.PandasCogroupedOps.applyInArrow
     pyspark.sql.UDFRegistration.register
+    pyspark.sql.GroupedData.applyInPandas
     """
     require_minimum_pyarrow_version()
 
@@ -346,6 +347,9 @@ def pandas_udf(f=None, returnType=None, functionType=None):
     .. versionchanged:: 4.0.0
         Supports keyword-arguments in SCALAR and GROUPED_AGG type.
 
+    .. versionchanged:: 4.1.0
+        Supports iterator API in GROUPED_MAP type.
+
     Parameters
     ----------
     f : function, optional

python/pyspark/sql/pandas/typehints.py

@@ -456,8 +456,22 @@ def infer_group_pandas_eval_type(
     if is_iterator_dataframe or is_iterator_dataframe_with_keys:
         return PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF
 
-    # Default to non-iterator (standard grouped map)
-    return PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF
+    # pd.DataFrame -> pd.DataFrame
+    is_dataframe = (
+        len(parameters_sig) == 1
+        and parameters_sig[0] == pd.DataFrame
+        and return_annotation == pd.DataFrame
+    )
+    # Tuple[Any, ...], pd.DataFrame -> pd.DataFrame
+    is_dataframe_with_keys = (
+        len(parameters_sig) == 2
+        and parameters_sig[1] == pd.DataFrame
+        and return_annotation == pd.DataFrame
+    )
+    if is_dataframe or is_dataframe_with_keys:
+        return PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF
+
+    return None
 
 
 def infer_group_pandas_eval_type_from_func(

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

@@ -26,9 +26,10 @@
     pandas_requirement_message,
     pyarrow_requirement_message,
 )
-from pyspark.sql.pandas.typehints import infer_eval_type
+from pyspark.sql.pandas.typehints import infer_eval_type, infer_group_pandas_eval_type
 from pyspark.sql.pandas.functions import pandas_udf, PandasUDFType
 from pyspark.sql import Row
+from pyspark.util import PythonEvalType
 
 if have_pandas:
     import pandas as pd
@@ -186,6 +187,49 @@ def func() -> float:
             PandasUDFType.GROUPED_AGG,
         )
 
+    def test_type_annotation_group_map(self):
+        # pd.DataFrame -> pd.DataFrame
+        def func(col: pd.DataFrame) -> pd.DataFrame:
+            pass
+
+        self.assertEqual(
+            infer_group_pandas_eval_type(signature(func), get_type_hints(func)),
+            PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF,
+        )
+
+        # Tuple[Any, ...], pd.DataFrame -> pd.DataFrame
+        def func(key: Tuple, col: pd.DataFrame) -> pd.DataFrame:
+            pass
+
+        self.assertEqual(
+            infer_group_pandas_eval_type(signature(func), get_type_hints(func)),
+            PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF,
+        )
+
+        # Iterator[pd.DataFrame] -> Iterator[pd.DataFrame]
+        def func(col: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]:
+            pass
+
+        self.assertEqual(
+            infer_group_pandas_eval_type(signature(func), get_type_hints(func)),
+            PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF,
+        )
+
+        # Tuple[Any, ...], Iterator[pd.DataFrame] -> Iterator[pd.DataFrame]
+        def func(key: Tuple, col: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]:
+            pass
+
+        self.assertEqual(
+            infer_group_pandas_eval_type(signature(func), get_type_hints(func)),
+            PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF,
+        )
+
+        # Should return None for unsupported signatures
+        def func(col: Iterator[pd.Series]) -> Iterator[pd.Series]:
+            pass
+
+        self.assertEqual(infer_group_pandas_eval_type(signature(func), get_type_hints(func)), None)
+
     def test_type_annotation_negative(self):
         def func(col: str) -> pd.Series:
             pass