HyukjinKwon opened a new pull request, #184: URL: https://github.com/apache/spark-connect-rust/pull/184
### What changes were proposed in this pull request? [SPARK-60971](https://issues.apache.org/jira/browse/SPARK-60971). **This builds on #179 (SPARK-60957) and includes its commit; the last commit is this PR's.** **Eval types from type hints.** `GroupedData.applyInPandas` and `applyInArrow` now take the eval type of the function from its type hints, as PySpark 4.3.0 does. The inference uses `infer_group_pandas_eval_type_from_func` and `infer_group_arrow_eval_type_from_func` of the vendored `pyspark.sql.pandas.typehints`: | Function | Eval type | | --- | --- | | `(Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]`, optionally after a key | `SQL_GROUPED_MAP_PANDAS_ITER_UDF` (216) | | `(Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]`, optionally after a key | `SQL_GROUPED_MAP_ARROW_ITER_UDF` (215) | | other hints, or none | `SQL_GROUPED_MAP_PANDAS_UDF` (201) or `SQL_GROUPED_MAP_ARROW_UDF` (209), as before | When the hints cannot be read, these functions warn as PySpark does (`Cannot infer the eval type from type hints.`) and use the default form. PySpark 4.3.0's `applyInArrow` then raises `UnboundLocalError`, as its `eval_type` is unbound there; here it runs with the default form, as `applyInPandas` does. **Validation.** The map functions now validate their function as PySpark does, before they build the plan or analyze the input columns. They use `_validate_vectorized_udf`, which `pyspark.sql.pandas.functions` now has verbatim from PySpark. This covers: - the grouped and cogrouped `applyInPandas` and `applyInArrow`; - `applyInPandasWithState`; - `mapInPandas` and `mapInArrow`. For example, a grouped function of three arguments raises `INVALID_PANDAS_UDF`, and a value that is not callable raises `TypeError: unsupported callable`. ### Why are the changes needed? - A function of an iterator, which PySpark has supported since 4.1.0, got the eval type of the other form, so it failed on the Python worker with `PythonException`. For example: `def f(it: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]` with `df.groupBy("k").applyInPandas(f, schema)`. - A function that PySpark rejects up front was sent to the server and failed there. ### Does this PR introduce _any_ user-facing change? Yes, these functions infer the eval type and validate the function as in PySpark 4.3.0. ### How was this patch tested? - Against a local Spark 4.3.0 server, a probe of 18 cases gives the same results and warnings as official PySpark 4.3.0. The cases cover the iterator and plain forms, with and without keys and hints, partial hints, wrong arities, and values that are not callable, across the six functions. The one difference is the `applyInArrow` case above. Before this change, 15 of the 18 differed. - New offline tests check: - the eval type sent for each form; - the warning for partial hints; - the validation errors, raised before any request to the server. Without this change, 14 of their 18 cases fail. - A new live test runs the iterator forms, with and without keys. It also passes with official PySpark 4.3.0 and on a server started with `SPARK_SKIP_CONNECT_COMPAT_TESTS`, as CI's is. - With each module under `pyspark.sql` imported first, `_validate_vectorized_udf`, the UDF types and `pyspark.sql.conversion` import. - The lib tests and every integration test, e2e included, against the local server (643 passed), the build of all tests with `wasm-udf`, the offline Python suites (634 passed, 8 live cases skipped), and the live drop-in and Arrow suites (58 passed) and live cases (10 passed). - Clippy reports no new diagnostics. It flags the `input_columns` closures of two changed lines, which the diff shows as moved; it reports the same lint at every such closure in `group.rs`. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Opus 5.5 This pull request and its description were written by Isaac. -- 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]
