hzmsh opened a new pull request, #59558: URL: https://github.com/apache/spark/pull/59558
### What changes were proposed in this pull request? Fixes [SPARK-39142](https://issues.apache.org/jira/browse/SPARK-39142) and [SPARK-42235](https://issues.apache.org/jira/browse/SPARK-42235). It also fixes [SPARK-43189](https://issues.apache.org/jira/browse/SPARK-43189), which reports the same error for the docs example. This PR takes over [#39974](https://github.com/apache/spark/pull/39974) by wayneguow (HyukjinKwon approved an earlier version; it was closed as stale after the follow-up changes got no further review). He [agreed to hand it over](https://github.com/apache/spark/pull/39974#issuecomment-6216815975). His five commits come first, with him as their author, so the merge script lists him as co-author; two of them needed conflict fixes to apply on current master. Fix: - Add three `pandas_udf` overloads without `functionType` to `python/pyspark/sql/pandas/functions.pyi`, shaped like the [`udf` overloads](https://github.com/apache/spark/blob/ca6524e34ce487bb61ab1fda4b440a86e2c2177b/python/pyspark/sql/functions/builtin.py#L34219-L34241): `pandas_udf(f, returnType)`, `@pandas_udf(returnType)` and `@pandas_udf(returnType=...)`. The return type can be any `DataTypeOrString`, which covers HyukjinKwon's review comments on the original PR ([all data types](https://github.com/apache/spark/pull/39974#discussion_r1103910747), [struct and map types](https://github.com/apache/spark/pull/39974#discussion_r1107041878)). Unlike `udf` it has no default, because `pandas_udf` raises `CANNOT_BE_NONE` without one. The result is a `UserDefinedFunctionLike`, as for `udf`. - The function parameter is `Callable[..., Any]`, as in `udf`. Reusing aliases such as `PandasScalarToScalarFunction` would add little: mypy reads their unparameterized TypeVars as `Any`, so any function with one to three positional parameters passes and they would catch little beyond functions with no parameters, and every overload error would get much longer. - These three replace the ten overloads from the original PR (a call form and a decorator form for each of five function shapes). A decorator form only sees the return type, so it cannot pick the UDF type from the function: with mypy 1.19.1 two of the ten can never match, and two return a grouped map UDF, which `pandas_udf` never infers from type hints. - From the original PR, `MapType` is kept as a return type of the legacy `SCALAR` and `SCALAR_ITER` overloads (the ones with `functionType`). Its `StructType` for `SCALAR_ITER` and `MapType` for `GROUPED_AGG` are taken out again in a later commit: the function aliases of those overloads only allow `Iterator[Series]` output and scalar results, so they accepted UDFs that fail at runtime and rejected the ones that work (see Tests). - `UserDefinedFunctionLike.__call__` (`sql/_typing.pyi` for classic, `sql/connect/_typing.py` for Connect) also takes `**kwargs: ColumnOrName`, like the runtime `UserDefinedFunction.__call__` in both. Keyword calls are supported since 4.0.0 and used in the `pandas_udf` docstring (`calc(b=..., a=...)`); with a typed result they would otherwise become `call-arg` errors. This is a separate commit. - Remove the 24 `# type: ignore[call-overload]` on `pandas_udf` calls that are now unused (`python/mypy.ini` sets `warn_unused_ignores`): 5 in `examples/src/main/python/sql/arrow.py`, which the Arrow/pandas user guide includes, 17 in `pyspark.pandas` and 2 in `pyspark.ml`. `ruff format` joined some of the shortened calls onto one line. - Annotate `unary_np_spark_mappings` in `pyspark/pandas/numpy_compat.py` as `Dict[str, Callable[[Column], Column]]`. Once `pandas_udf` results are typed, its values would otherwise be inferred as `object`, and `column_op` reports `arg-type`. - Add the typing test case `pandasUDFTypeHints` to `python/pyspark/sql/tests/typing/test_udf.yml`. ### Why are the changes needed? Type checkers reject pandas UDFs written the way the docs recommend, with Python type hints and no `functionType`, and the UDF is typed as `Any`: ```python import pandas as pd from pyspark.sql.functions import pandas_udf @pandas_udf("long") def plus_one(s: pd.Series) -> pd.Series: return s + 1 reveal_type(plus_one) plus_one(42) # master, mypy 1.19.1: # error: No overload variant of "pandas_udf" matches argument type "str" [call-overload] # (followed by notes listing all 20 overloads) # note: Revealed type is "Any" # no error for plus_one(42) # with this PR: # note: Revealed type is "pyspark.sql._typing.UserDefinedFunctionLike" # error: Argument 1 to "__call__" of "UserDefinedFunctionLike" has incompatible type "int"; expected "Column | str" [arg-type] ``` The same happens for `@pandas_udf(returnType=IntegerType())`, `pandas_udf(func, returnType=LongType())` and `@pandas_udf("col1 string, col2 long")` (the SPARK-43189 example). Only the legacy form with `functionType` type-checks, which the docstring says "will be deprecated in the future releases". When it started: - Type hints for pandas UDFs came in 3.0.0 ([SPARK-28264](https://issues.apache.org/jira/browse/SPARK-28264)), and the docstring prefers them to `functionType`. The type stubs came in 3.1.0 with the pyspark-stubs port ([SPARK-32714](https://issues.apache.org/jira/browse/SPARK-32714)). Every `pandas_udf` overload there, then and now, requires `functionType`. `functions.pyi` is the same in 4.2.0 and master. Reports from users: - [SPARK-39142](https://issues.apache.org/jira/browse/SPARK-39142) (Philip Kahn, 3.2.1): static type checkers flag `pandas_udf`-decorated functions as incorrect, which he traced to the `returnType` types in this stub's overloads; his example is a Pylance issue he filed ([microsoft/pylance-release#2789](https://github.com/microsoft/pylance-release/issues/2789)). - [SPARK-42235](https://issues.apache.org/jira/browse/SPARK-42235) (Donggu Kang, 3.3.0): `@pandas_udf(returnType=LongType())` has no matching overload. - [SPARK-43189](https://issues.apache.org/jira/browse/SPARK-43189) (Andrew Grigorev, 3.2.4 to 3.4.0): mypy prints a long `No overload variant of "pandas_udf" matches argument type "str"` error for the docs example, and the docs examples work around it with `# type: ignore[call-overload]`. - After the original PR was closed, four people asked for it there: dennisobrien ([2024-01](https://github.com/apache/spark/pull/39974#issuecomment-1874760361)) so that mypy and pyright stop reporting errors, pontusvision ([2024-08](https://github.com/apache/spark/pull/39974#issuecomment-2271554177)) with Pylance, bthorsted ([2024-11](https://github.com/apache/spark/pull/39974#issuecomment-2497469887)) asking to reopen and merge it, and AkiraVoid ([2025-07](https://github.com/apache/spark/pull/39974#issuecomment-3069084201)) saying the problem remains. - Spark's own code has 24 `# type: ignore[call-overload]` on `pandas_udf` calls for the same reason; this PR removes them. What goes wrong: - In [`functions.pyi`](https://github.com/apache/spark/blob/ca6524e34ce487bb61ab1fda4b440a86e2c2177b/python/pyspark/sql/pandas/functions.pyi#L101-L189), every `pandas_udf` overload requires `functionType` (a `PandasUDFType` value), positionally or by keyword. A call with only a return type matches none of them, so mypy reports `call-overload` and types the result as `Any`. At runtime, [`vectorized_udf`](https://github.com/apache/spark/blob/ca6524e34ce487bb61ab1fda4b440a86e2c2177b/python/pyspark/sql/pandas/functions.py#L791-L815) accepts these calls and infers the UDF type from the type hints. ### Does this PR introduce _any_ user-facing change? Yes, relative to 4.2.0 as well as master. Only static type checking changes: the `.py` edits only remove comments, rewrap lines, add one variable annotation and widen a Protocol signature in the typing-only `sql/connect/_typing.py`. Behavior change: - The documented forms without `functionType` now type-check: `@pandas_udf("long")`, `@pandas_udf(IntegerType())`, `@pandas_udf(returnType=...)`, `pandas_udf(func, "long")` and `pandas_udf(func, returnType=...)`, with any `DataType` or DDL string. Before, mypy reported `call-overload` (output above). - Their result is a `UserDefinedFunctionLike` instead of `Any`, so errors that `Any` hid are now reported, such as `plus_one(42)` above. Attributes the protocol does not declare, such as `__name__`, now give `attr-defined`, as they already do for `udf` results. - `# type: ignore[call-overload]` comments that users added on these calls become unused; with `warn_unused_ignores`, mypy reports `Unused "type: ignore" comment [unused-ignore]`. - Keyword calls on UDF results now type-check, for `udf` results too. Before, for example, `add(a=col("id"), b="id")` on a `@udf("long")` function gave `Unexpected keyword argument "a" for "__call__" of "UserDefinedFunctionLike" [call-arg]`. - Legacy forms with `functionType` type-check as before, except that `MapType` is now accepted as the return type of `SCALAR` and `SCALAR_ITER` UDFs (both run fine). Not covered: - `arrow_udf` has the same gap: its docstring uses `@arrow_udf("long")`, and its stub has no overload without `functionType`. Left for a follow-up. - The new overloads accept any callable, as `udf` does, so some UDFs that fail at runtime still type-check: a function annotated `int -> str` or on `pyarrow.Array` (`UNSUPPORTED_SIGNATURE`), one with no parameters (`INVALID_PANDAS_UDF`; the existing aliases would catch only this one), and a grouped aggregate returning a dict with a `StructType` return type (`NOT_IMPLEMENTED`). Master rejects them only because it rejects every call without `functionType`. - The legacy forms still reject `SCALAR_ITER` with `StructType` and `GROUPED_AGG` with `MapType`, as on master, though both run with the right function shape (`Iterator[DataFrame]` output, a dict result). The function aliases of these overloads cannot express those shapes; the type-hint form accepts them. - Left as on master: under `mypy --strict` (`--no-implicit-reexport`), `from pyspark.sql.functions import pandas_udf` fails with `attr-defined` (importing from `pyspark.sql.pandas.functions` works); `functionType=None` and `PandasUDFType.GROUPED_AGG_ITER` have no legacy overload (with this PR their type-hint forms type-check instead). - 8 `# type: ignore` stay in `sql/pandas/group_ops.py` and `sql/pandas/map_ops.py`, where internal code passes `PythonEvalType` values as `functionType`. - The `pandas_udf` docstring Notes still say the functions "do not take keyword arguments on the calling side", which contradicts its 4.0.0 `versionchanged` note. Not changed here. - Type checkers: CI checks with mypy only. pyright was run locally on scratch files (see Tests); Pylance itself was not tried. ### How was this patch tested? Tests: - New case `pandasUDFTypeHints` in `test_udf.yml`: `@pandas_udf(MapType(...))`, `@pandas_udf(returnType="string")`, `pandas_udf(lambda x: x, "string")`, `pandas_udf(lambda x: x, returnType=IntegerType())` and the keyword call `f(x="id")` must type-check, and calling the results with an int (`f(42)`, `g(x=42)`, `h(42)`) must fail with `arg-type`, so the return type of each new overload is checked. Each of the three new overloads is needed: removing any one of them makes the case fail with `call-overload`. - With master's stub, the case fails: `No overload variant of "pandas_udf" matches ... [call-overload]` on both decorators and both calls, and none of `f(42)`, `g(x=42)`, `h(42)` is reported because the results are `Any`. Without the `**kwargs` commit, it fails with `Unexpected keyword argument "x" for "__call__" of "UserDefinedFunctionLike" [call-arg]`. - Putting back any one of the 24 removed ignores gives exactly one `Unused "type: ignore" comment` error. - Legacy `functionType` forms with map and struct return types were run in classic PySpark (local[1], a local build of master, Python 3.12): `SCALAR` and `SCALAR_ITER` with `MapType` work. With the original PR's additions, `SCALAR_ITER` with `StructType` and `Iterator[Series]` output type-checked but fails with `[UDF_RETURN_TYPE]`, and `GROUPED_AGG` with `MapType` returning a float type-checked but fails with `AttributeError: 'float' object has no attribute 'items'`. Both are rejected again here, as on master. - A scratch file with 15 `pandas_udf` forms from the docstring, `arrow.py` and a docs notebook: master rejects 14 of them (all but the legacy `PandasUDFType.SCALAR` one), this PR accepts all 15. 9 wrong uses are rejected, for example an int as the return type, a missing return type, an unknown function type (999), `groupBy().apply()` with a type-hint UDF and calling the UDF with an int. - pyright 1.1.414 (standard mode) on a separate set of 108 forms accepts and rejects the same ones as mypy on this branch. - No `python/run-tests` modules were run: the `.py` changes only remove comments, rewrap lines, add one annotation and widen a typing-only Protocol, and `pyspark.pandas.numpy_compat` imports fine. - With mypy 1.19.1, pytest-mypy-plugins 3.2.0 and ruff 0.14.8 (the CI lint versions), `dev/lint-python` passes with every step run, none skipped: mypy annotations `Success: no issues found in 1335 source files`, examples `Success: no issues found in 104 source files`, data tests 40 passed (39 before, plus the new case). - Commands: ``` dev/lint-python MYPYPATH=python pytest -o "python_files=test_*.yml" --rootdir python --mypy-only-local-stub --mypy-ini-file python/mypy.ini python/pyspark/sql/tests/typing/test_udf.yml ``` ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Claude Opus 5.5) Co-developed with Claude Code (Claude Opus 5.5): the commits after wayneguow's, the tests and this description. The checks above were run locally. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
