[SPARK-55108][PYTHON] Use the latest pandas-stubs for type check

### What changes were proposed in this pull request?

Upgrade `pandas-stubs` version (basically do not pin it anymore) and fix all the mypy errors.

### Why are the changes needed?

`pandas-stubs` sync with `pandas` - in theory we should use the same version for both. The stubs for numpy is smarter that it actually depends on a specific version. `pandas-stubs` does not do that. However, we should definitely not pin it on a 1.x version.

Some of the ignore comments might because we have a low mypy version (we have to do one at a time). I did not really "fix" all typing issues - I explicitly ignored a lot of them. Some of the errors are because our type hint is wrong or inaccurate.

One step at a time. I think we are moving to the right direction.

### Does this PR introduce _any_ user-facing change?

No

### How was this patch tested?

Let's check CI.

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

No

Closes #53877 from gaogaotiantian/upgrade-pandas-stubs.

Authored-by: Tian Gao <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>

Author: gaogaotiantian

dev/requirements.txt

@@ -26,7 +26,7 @@ ruff==0.14.8
 mypy==1.8.0
 pytest-mypy-plugins==1.9.3
 # See SPARK-38680.
-pandas-stubs<1.2.0.54
+pandas-stubs>=2.2.0
 scipy-stubs; python_version>='3.10'
 types-PyYAML
 

dev/spark-test-image/lint/Dockerfile

@@ -94,7 +94,7 @@ RUN python3.11 -m pip install \
     'numpy==2.0.2' \
     'numpydoc' \
     'pandas' \
-    'pandas-stubs==1.2.0.53' \
+    'pandas-stubs' \
     'plotly>=4.8' \
     'pyarrow>=22.0.0' \
     'pytest-mypy-plugins==1.9.3' \

python/pyspark/instrumentation_utils.py

@@ -124,7 +124,7 @@ def _attach(
     logger_module: Union[str, ModuleType],
     modules: List[ModuleType],
     classes: List[Type[Any]],
-    missings: List[Tuple[Type[Any], Type[Any]]],
+    missings: List[Tuple[Union[ModuleType, Type[Any]], Type[Any]]],
 ) -> None:
     if isinstance(logger_module, str):
         logger_module = importlib.import_module(logger_module)

python/pyspark/ml/functions.py

@@ -241,7 +241,7 @@ def _validate_and_transform_single_input(
         # tensor columns
         if len(batch.columns) == 1:
             # one tensor column and one expected input, vstack rows
-            single_input = np.vstack(batch.iloc[:, 0])
+            single_input = np.vstack(batch.iloc[:, 0])  # type: ignore[call-overload]
         else:
             raise ValueError(
                 "Multiple input columns found, but model expected a single "

python/pyspark/pandas/accessors.py

@@ -579,7 +579,7 @@ def new_func(o: Any) -> Union[pd.DataFrame, pd.Series]:
             return original_func(o, *args, **kwargs)
 
         def apply_func(pdf: pd.DataFrame) -> pd.DataFrame:
-            return new_func(pdf).to_frame()
+            return new_func(pdf).to_frame()  # type: ignore[operator]
 
         def pandas_series_func(
             f: Callable[[pd.DataFrame], pd.DataFrame], return_type: DataType

python/pyspark/pandas/base.py

@@ -26,7 +26,7 @@
 
 import numpy as np
 import pandas as pd
-from pandas.api.types import is_list_like, CategoricalDtype  # type: ignore[attr-defined]
+from pandas.api.types import is_list_like, CategoricalDtype
 
 from pyspark.sql import functions as F, Column, Window
 from pyspark.sql.types import LongType, BooleanType, NumericType

python/pyspark/pandas/categorical.py

@@ -17,7 +17,7 @@
 from typing import Any, Callable, List, Optional, Union, TYPE_CHECKING, cast
 
 import pandas as pd
-from pandas.api.types import (  # type: ignore[attr-defined]
+from pandas.api.types import (
     CategoricalDtype,
     is_dict_like,
     is_list_like,

python/pyspark/pandas/data_type_ops/base.py

@@ -116,7 +116,7 @@ def _should_return_all_false(left: IndexOpsLike, right: Any) -> bool:
     based on incompatible dtypes: non-numeric vs. numeric (including bools).
     """
     from pyspark.pandas.base import IndexOpsMixin
-    from pandas.api.types import is_list_like  # type: ignore[attr-defined]
+    from pandas.api.types import is_list_like
 
     def are_both_numeric(left_dtype: Dtype, right_dtype: Dtype) -> bool:
         return is_numeric_dtype(left_dtype) and is_numeric_dtype(right_dtype)

python/pyspark/pandas/data_type_ops/boolean_ops.py

@@ -19,7 +19,7 @@
 from typing import Any, Union
 
 import pandas as pd
-from pandas.api.types import CategoricalDtype, is_integer_dtype  # type: ignore[attr-defined]
+from pandas.api.types import CategoricalDtype, is_integer_dtype
 from pandas.core.dtypes.common import is_numeric_dtype
 
 from pyspark.pandas.base import column_op, IndexOpsMixin

python/pyspark/pandas/data_type_ops/categorical_ops.py

@@ -16,11 +16,11 @@
 #
 
 from itertools import chain
-from typing import cast, Any, Union
+from typing import cast, Any, Sequence, Union
 
 import pandas as pd
 import numpy as np
-from pandas.api.types import is_list_like, CategoricalDtype  # type: ignore[attr-defined]
+from pandas.api.types import is_list_like, CategoricalDtype
 
 from pyspark.pandas._typing import Dtype, IndexOpsLike, SeriesOrIndex
 from pyspark.pandas.base import IndexOpsMixin
@@ -43,7 +43,7 @@ def restore(self, col: pd.Series) -> pd.Series:
         """Restore column when to_pandas."""
         return pd.Series(
             pd.Categorical.from_codes(
-                col.replace(np.nan, -1).astype(int),
+                cast(Sequence[int], col.replace(np.nan, -1).astype(int)),
                 categories=cast(CategoricalDtype, self.dtype).categories,
                 ordered=cast(CategoricalDtype, self.dtype).ordered,
             )