HyukjinKwon opened a new pull request, #198: URL: https://github.com/apache/spark-connect-rust/pull/198
### What changes were proposed in this pull request? [SPARK-60995](https://issues.apache.org/jira/browse/SPARK-60995): these functions now check their arguments before any request, with PySpark's errors. - **`when`**: `when()` and `Column.when` raise `NOT_EXPECTED_TYPE` for a condition that is not a Column. `Column.when` raises `INVALID_WHEN_USAGE` for a column that is not a `when()` without `otherwise()`. `Column.otherwise` raises `PySparkTypeError` for a column that is not a `when()`, or one that already has an `otherwise()`. - **Partition transforms**: `pyspark.sql.functions.partitioning` holds PySpark Connect's partition transforms. Its `bucket` raises `NOT_EXPECTED_TYPE` for a `numBuckets` that is neither an int nor a Column. `bucket`, `years`, `months`, `days` and `hours` of `pyspark.sql.functions` warn that they are deprecated, and call those of `partitioning`, as in PySpark. - **`pandas_udf` and `arrow_udf`** are PySpark's. `pyspark/sql/pandas/functions.py` is vendored from Apache Spark 4.3.0 and checked by the drift check. It creates the UDF with the new `pyspark.sql.connect.udf._create_udf`. So: - they raise `CANNOT_BE_NONE` without a return type; - they raise `INVALID_PANDAS_UDF_TYPE` for a function type that is not a value of `PandasUDFType` or `ArrowUDFType`; - they warn as PySpark's do. `UserDefinedFunction` checks its function, return type and eval type, and the return type for the eval type (`_check_return_type`, verbatim), as PySpark Connect's does. - **`udtf` and `arrow_udtf`** check their handler classes as PySpark does (`_validate_udtf_handler` and `_validate_arrow_udtf_handler`, verbatim): `INVALID_UDTF_HANDLER_TYPE`, `INVALID_UDTF_NO_EVAL`, `INVALID_UDTF_RETURN_TYPE`, `INVALID_UDTF_BOTH_RETURN_TYPE_AND_ANALYZE` and `INVALID_ARROW_UDTF_WITH_ANALYZE`. ### Why are the changes needed? Before, these calls built other expressions, or UDFs, without an error: - `F.when("a", 1)` was `CASE WHEN a THEN 1 END`. - `F.col("a").when(cond, 1)` started a new CASE WHEN, dropping `a`. - `F.col("a").otherwise(0)` returned `a`. - `F.bucket("n", "c")` took `n` as a column name. - `pandas_udf` took `"scalar"` or `"grouped_map"` as a function type, which PySpark does not. It ignored a `PandasUDFType`, such as `GROUPED_MAP`, and inferred the eval type instead. Without a return type, it defaulted to `StringType`. - `udtf(3)` and `arrow_udtf(3)` returned a UDTF. ### Does this PR introduce _any_ user-facing change? Yes. These functions raise the errors of PySpark 4.3.0 for these arguments. `pandas_udf` and `arrow_udf` take PySpark's function types, and no longer `"scalar"` and the other strings. `PandasUDFType` no longer has the `MAP_ITER` and `COGROUPED_MAP` that PySpark does not have. ### How was this patch tested? - A probe of 34 calls of these functions, valid and invalid, gives the same results and warnings as official PySpark 4.3.0. Before this change, 26 of them differed. One call differs only in that official PySpark needs a session to parse the DDL return type of a pandas UDF, which the probe does not have. - New offline tests check the errors of 18 of these calls, the reprs of `when` and `otherwise`, and the warnings of the partition transforms. They also pass with official PySpark 4.3.0, and fail before this change. - The offline tests of `pandas_udf` and `arrow_udf` now pass PySpark's function types and return types. - The lib tests and every integration test, e2e included, against a local Spark 4.3.0 server (652 passed, and 123 of the transport crate), the build of all tests with `wasm-udf`, the offline Python suites (818 passed, 9 skipped), and the live drop-in and Arrow suites (79 passed) and live cases (10 passed). Clippy reports no diagnostics in the changed lines. The eval type of `arrow_udtf`, with which every `arrow_udtf` fails, is [SPARK-60999](https://issues.apache.org/jira/browse/SPARK-60999). ### 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]
