HyukjinKwon opened a new pull request, #220: URL: https://github.com/apache/spark-connect-rust/pull/220
### What changes were proposed in this pull request? [SPARK-61014](https://issues.apache.org/jira/browse/SPARK-61014): `foreach` and `foreachPartition` now run as PySpark Connect's do, and `toJSON` returns a DataFrame. - `df.foreach(f)` selects a struct of the columns of each row, with a UDF of no result that calls `f` with the row. - `df.foreachPartition(f)` runs a `mapInArrow` function of no result that calls `f` with an iterator of the rows of each partition. It converts them with `spark.sql.execution.pyspark.binaryAsBytes`. - The functions are pickled and run by the server's Python worker, so they are Python, in `pyspark.sql._foreach`, as `pyspark.sql.connect.dataframe` builds them. - `df.toJSON()` returns the DataFrame of `select(to_json(struct(col("*"))).alias("value"))`. In the Rust API, `DataFrame::to_json_dataframe` builds the DataFrame of `toJSON`, and `to_json` collects its strings. ### Why are the changes needed? `foreach` and `foreachPartition` failed on the worker with `TypeError: cannot unpack non-iterable function object`, because the command of their UDF was the function rather than the pair of the function and its return type. `toJSON()` raised `NotImplementedError`. ### Does this PR introduce _any_ user-facing change? Yes. `foreach`, `foreachPartition` and `toJSON` work as in PySpark 4.3.0. ### How was this patch tested? - A new live test checks the row that `foreach` gives `f` and the rows that `foreachPartition` gives it, by errors that `f` raises on the server. It also checks the DataFrame of `toJSON`. It passes with official PySpark 4.3.0, and fails without this change. - A new offline test unpickles the functions that are sent to a mock Connect server, and calls them. The mock now serves `GetWithDefault` config requests. The test fails without this change. - A new lib test checks the plan of `to_json_dataframe`. - A probe of `foreach`, `foreachPartition` (with `binaryAsBytes` on and off) and `toJSON` matches official PySpark 4.3.0. - The lib tests and every integration test, e2e included, against a local Spark 4.3.0 server (672 passed, and 129 of the transport crate), the build of all tests with `wasm-udf`, the offline Python suites (1625 passed, 8 skipped), and the live drop-in and Arrow suites (95 passed) and live cases (10 passed). Clippy reports no diagnostics in the changed lines. ### 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]
