andygrove opened a new pull request, #5282:
URL: https://github.com/apache/datafusion-comet/pull/5282
## Which issue does this PR close?
Closes #5281.
## Rationale for this change
A `ScalaUDF` (or Java UDF) whose closure was captured by a class from a user
jar (`--jars` / `spark.jars`) fails at execution time:
```
java.lang.ClassCastException: cannot assign instance of
java.lang.invoke.SerializedLambda to field
org.apache.spark.sql.catalyst.expressions.ScalaUDF.f of type
scala.Function1
...
at
org.apache.comet.udf.codegen.CometScalaUDFCodegen.lookupOrCompile(CometScalaUDFCodegen.scala:160)
at
org.apache.comet.udf.CometUdfBridge.evaluateInternal(CometUdfBridge.java:203)
```
That `ClassCastException` is a **masked `ClassNotFoundException`**. When the
deserializing classloader cannot resolve a lambda's *capturing class*,
`ObjectInputStream` records the CNFE against the object handle, therefore skips
`SerializedLambda.readResolve()`, and the raw `SerializedLambda` then fails the
field-type check in `defaultCheckFieldValues`. The CNFE never surfaces.
The classloader is wrong because of which thread runs the deserialization.
`CometUdfBridge.evaluate` is invoked from a Tokio worker, which attaches to the
JVM through JNI, and **an attached thread has no context classloader**.
`lookupOrCompile` therefore fell back to `classOf[Expression].getClassLoader` —
the loader that loaded spark-catalyst and Comet, which never contains user
jars. Spark installs the executor's `MutableURLClassLoader` on task threads
only.
Measured by returning thread identity from inside a UDF body:
| stage leaf | thread running the UDF | context classloader |
| --- | --- | --- |
| `CometNativeScan` | `Thread-28..32` (Tokio workers) | **null** |
| shuffle read / join build side | `Executor task launch worker …` |
`MutableURLClassLoader` |
| native scan disabled | `Executor task launch worker …` |
`MutableURLClassLoader` |
So any plan whose stage leaf is a native scan (`CometNativeScan`,
`CometIcebergNativeScan`) hits this on every UDF call. `CometUdfBridge` already
propagates `TaskContext` across the JNI boundary for exactly this reason; the
classloader was not propagated.
This is invisible to Comet's existing test suites because they register UDFs
from test classes already on the application classpath.
## What changes are included in this PR?
The task thread's context classloader is now propagated to the Tokio worker,
following the same path as `TaskContext`:
- `CometExecIterator` captures
`Thread.currentThread().getContextClassLoader` (it is constructed on the Spark
task thread) and passes it to `createPlan`.
- `Native.createPlan` gains a `classLoader` parameter; `jni_api.rs` holds it
as a JNI global ref on `ExecutionContext` and hands it to the planner via
`with_class_loader`.
- `PhysicalPlanner` clones it into every `JvmScalarUdfExpr`, under the same
`debug_assert` invariant already used for `task_context`.
- `JvmScalarUdfExpr` passes it through the bridge call; the JNI signature
gains `Ljava/lang/ClassLoader;`.
- `CometUdfBridge.evaluate` installs it as the calling thread's context
classloader for the duration of the call, restoring the prior value in
`finally` — the same save-and-restore shape used for `TaskContext`.
Installing it on the thread rather than passing it down to the deserializer
fixes both affected lookups at once: the closure deserialization in
`CometScalaUDFCodegen`, and the `Class.forName(udfClassName)` in
`CometUdfBridge`, which had the same latent bug for a user-supplied `CometUDF`
implementation shipped in a user jar. It also means user code inside a UDF that
relies on the context classloader behaves as it does under Spark's own
execution.
One line was added to the "Behavior" section of
`docs/source/user-guide/latest/scala_java_udfs.md`, alongside the existing
`TaskContext.get()` guarantee.
## How are these changes tested?
New suite `CometScalaUDFClassLoaderSuite`, which compiles a class holding a
serializable `scala.Function1` lambda into a jar at test time and wires it in
via `spark.executor.extraClassPath` — the local-mode equivalent of a `--jars`
submission, since `LocalSchedulerBackend` feeds it into the executor's
`MutableURLClassLoader`. Four tests:
1. sanity — the class is loadable from task threads and raises
`ClassNotFoundException` from `classOf[Expression].getClassLoader` (the old
fallback), so the fixture reproduces the production classloader topology.
1. the failing query: `SELECT hiddenUdf(s) FROM t` with a `CometNativeScan`
leaf.
1. the propagation itself: the UDF body reports whether the classloader
installed on whatever thread invoked it can reach the user jar.
1. control: the same UDF with the native scan disabled, which runs on the
task thread and passed even before this change — isolating the failure to
classloader propagation rather than class availability.
The two new behavioral tests were confirmed to fail without the fix
(reverting only the loader install), reporting `MISSING|Thread-53,
MISSING|Thread-28, …` and naming the Tokio workers.
Existing coverage run locally on Spark 4.1 / Scala 2.13:
`CometCodegenSuite`, `CometCodegenHOFSuite`, `CometCodegenSourceSuite`,
`CometCodegenFuzzSuite` plus the new suite (180 tests), and `CometExecSuite`
(143 tests), which exercises the changed `createPlan` signature broadly. `cargo
fmt --check` and `cargo clippy --all-targets -D warnings` are clean.
--
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]
