viirya opened a new pull request, #7182:
URL: https://github.com/apache/datafusion-comet/pull/7182
## Which issue does this PR close?
Closes #7177.
## Rationale for this change
When a scalar `@arrow_udf` running on the native path (#6130) raises, the
error reaches the JVM as a `CometNativeException` carrying the Python
traceback. Spark's own path raises
`org.apache.spark.api.python.PythonException`, which PySpark converts to
`pyspark.errors.PythonException`. Code that catches
`pyspark.errors.PythonException` around a query therefore behaved differently
once the native path was enabled.
Changing the exception class alone is not enough, and what PySpark checks
differs by version:
- **Spark 4.1:** the job fails with `SparkException("Job aborted ...")`
caused by the task's exception. PySpark's `convert_exception` maps a
`PythonException` cause only when one of the cause's stack frames contains
`org.apache.spark.sql.execution.python`.
- **Spark 4.2:** `PythonException` carries the `PYTHON_EXCEPTION` error
condition (SPARK-54599). `DAGScheduler.abortStage` fails the job with a
user-facing `SparkThrowable` directly instead of wrapping it, and
`convert_exception` matches the class of that top-level exception.
## What changes are included in this PR?
- A new `SparkError::PythonException { traceback }` (error class
`PYTHON_EXCEPTION`). `python_traceback_error` in `python_udf.rs` now returns it
inside `ArrowError::ExternalError`, and `ArrowPythonUdfExec` unwraps it into
`DataFusionError::External`. That is how other `SparkError`s reach JNI, so the
error takes the existing `CometQueryExecutionException` JSON route to
`SparkErrorConverter`. The traceback is no longer prefixed with `Arrow UDF
Python error:`, so the message matches what Spark's runners report.
- `unwrap_arrow_external_error` moves from `spark-expr` to the common crate,
next to `SparkError`, so the operator can reuse it.
- The 4.x `ShimSparkErrorConverter` handles `PythonException` through a
per-version `CometPythonException` factory in
`org.apache.spark.sql.execution.python`:
- Spark 4.0/4.1 build `new PythonException(traceback, null)`, the same as
Spark's runner. Because the factory lives in that package, its frame is on the
exception's stack trace, which is what PySpark 4.1 checks for. That is also
where Spark's own runners create the exception. The 4.0 copy exists only so the
shared 4.x shim compiles, since the native path requires Spark 4.1.
- Spark 4.2 builds it with `errorClass = "PYTHON_EXCEPTION"` and the same
`msg`/`traceback` parameters Spark's runner uses.
- The user guide no longer lists this as a limitation.
Errors that are not raised by the UDF's own code are unchanged and still
reported as Comet native errors. This covers a result of the wrong length or
type, a failed PyArrow cast of the result, and an unpickling failure. In Spark
these also come back from the Python worker as `PythonException`; the user
guide says so.
## How are these changes tested?
- Rust: `test_python_exception_json` checks the JSON payload.
`reports_python_traceback_for_udf_errors` and
`stops_calling_python_after_the_first_error` now assert the typed
`SparkError::PythonException`, at the UDF level and at the operator boundary
respectively.
- `CometArrowPythonUdfSuite`:
- A unit test converts the JSON payload and checks that the result is a
`PythonException` with a frame in `org.apache.spark.sql.execution.python`.
- An end-to-end test runs a failing UDF on both the native path and
Spark's own path. On each path it checks the exception PySpark would convert:
the job failure itself on 4.2, or its cause together with the frame check on
4.1.
- `test_pyarrow_udf.py::test_scalar_arrow_udf_error_raises_python_exception`
catches `pyspark.errors.PythonException` from a failing native UDF. It also
runs Spark's own path as the reference.
Ran locally:
- Spark 4.1.3 with Python 3.11: Rust unit tests, `CometArrowPythonUdfSuite`,
and the scalar Arrow UDF pytests.
- Spark 4.2.0 with Python 3.13: the same tests.
- `-Pspark-4.0` compiles.
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]
