This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 003f608948 fix: propagate the Spark task ClassLoader to JVM UDF calls
(#5282)
003f608948 is described below
commit 003f608948ec386bbb666e32f5b14bc8cd01c0b8
Author: Andy Grove <[email protected]>
AuthorDate: Fri Aug 7 15:10:25 2026 -0600
fix: propagate the Spark task ClassLoader to JVM UDF calls (#5282)
---
.github/workflows/pr_build_linux.yml | 1 +
.github/workflows/pr_build_macos.yml | 1 +
docs/source/user-guide/latest/scala_java_udfs.md | 1 +
native/core/src/execution/jni_api.rs | 22 ++-
native/core/src/execution/planner.rs | 29 ++-
native/jni-bridge/src/comet_udf_bridge.rs | 4 +-
native/spark-expr/src/jvm_udf/mod.rs | 24 ++-
.../java/org/apache/comet/udf/CometUdfBridge.java | 38 +++-
.../scala/org/apache/comet/CometExecIterator.scala | 7 +-
spark/src/main/scala/org/apache/comet/Native.scala | 8 +-
.../comet/CometScalaUDFClassLoaderSuite.scala | 195 +++++++++++++++++++++
11 files changed, 302 insertions(+), 28 deletions(-)
diff --git a/.github/workflows/pr_build_linux.yml
b/.github/workflows/pr_build_linux.yml
index def2816607..3edaa3e569 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -402,6 +402,7 @@ jobs:
org.apache.comet.CometCodegenHOFSuite
org.apache.comet.CometFuzzMathSuite
org.apache.comet.CometCodegenFuzzSuite
+ org.apache.comet.CometScalaUDFClassLoaderSuite
org.apache.comet.codegen.CometSpecializedGettersDispatchSuite
org.apache.comet.CometStringDecodeSuite
org.apache.comet.CometWidthBucketSuite
diff --git a/.github/workflows/pr_build_macos.yml
b/.github/workflows/pr_build_macos.yml
index bac1dcd15b..42714de24c 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -218,6 +218,7 @@ jobs:
org.apache.comet.CometCodegenHOFSuite
org.apache.comet.CometFuzzMathSuite
org.apache.comet.CometCodegenFuzzSuite
+ org.apache.comet.CometScalaUDFClassLoaderSuite
org.apache.comet.codegen.CometSpecializedGettersDispatchSuite
org.apache.comet.CometStringDecodeSuite
org.apache.comet.CometWidthBucketSuite
diff --git a/docs/source/user-guide/latest/scala_java_udfs.md
b/docs/source/user-guide/latest/scala_java_udfs.md
index 1ebf79250f..3a55982bbc 100644
--- a/docs/source/user-guide/latest/scala_java_udfs.md
+++ b/docs/source/user-guide/latest/scala_java_udfs.md
@@ -54,6 +54,7 @@ When a UDF is rejected, the reason surfaces through Comet's
standard fallback di
- Non-deterministic expressions referenced from the argument tree (`rand`,
`uuid`, `monotonically_increasing_id`) produce per-partition sequences
consistent with Spark.
- `TaskContext.get()` inside the user function returns the driving Spark
task's context.
+- The Spark task thread's context ClassLoader is propagated to the thread that
runs the user function, so functions defined in jars supplied with `--jars` /
`spark.jars` resolve the same way they do under Spark's own execution.
- The user function must be closure-serializable; the same function that works
with Spark's executor execution works here.
## Known limitations
diff --git a/native/core/src/execution/jni_api.rs
b/native/core/src/execution/jni_api.rs
index f798b80335..2e98bf9aa7 100644
--- a/native/core/src/execution/jni_api.rs
+++ b/native/core/src/execution/jni_api.rs
@@ -352,6 +352,11 @@ struct ExecutionContext {
/// cheap to clone; the underlying `Global<JObject>` releases its JNI
global ref on drop
/// via `jni`'s `Drop` impl.
pub task_context: Option<Arc<Global<JObject<'static>>>>,
+ /// Context `ClassLoader` of the driving Spark task thread, captured at
`createPlan` time and
+ /// threaded into every JVM scalar UDF the planner builds; see
`CometUdfBridge.evaluate` for why
+ /// it has to travel with the plan. `None` when no driving Spark task is
present (unit tests,
+ /// direct native driver runs). Lifetime is as for `task_context` above.
+ pub class_loader: Option<Arc<Global<JObject<'static>>>>,
}
/// Accept serialized query plan and return the address of the native query
plan.
@@ -379,6 +384,7 @@ pub unsafe extern "system" fn
Java_org_apache_comet_Native_createPlan(
task_cpus: jlong,
key_unwrapper_obj: JObject,
task_context_obj: JObject,
+ class_loader_obj: JObject,
) -> jlong {
try_unwrap_or_throw(&e, |env| {
// Deserialize Spark configs
@@ -500,15 +506,21 @@ pub unsafe extern "system" fn
Java_org_apache_comet_Native_createPlan(
String::new()
};
- // Capture the driving Spark task's TaskContext as a JNI global
reference when
- // non-null. The `Arc<Global<JObject>>` releases its global ref on
drop, so cleanup
- // is automatic when the ExecutionContext drops.
+ // Capture the driving Spark task's TaskContext and context
ClassLoader as JNI global
+ // references when non-null. The `Arc<Global<JObject>>` releases
its global ref on
+ // drop, so cleanup is automatic when the ExecutionContext drops.
let task_context = if !task_context_obj.is_null() {
Some(Arc::new(jni_new_global_ref!(env, task_context_obj)?))
} else {
None
};
+ let class_loader = if !class_loader_obj.is_null() {
+ Some(Arc::new(jni_new_global_ref!(env, class_loader_obj)?))
+ } else {
+ None
+ };
+
let exec_context = Box::new(ExecutionContext {
id,
task_attempt_id,
@@ -536,6 +548,7 @@ pub unsafe extern "system" fn
Java_org_apache_comet_Native_createPlan(
),
tracing_event_name,
task_context,
+ class_loader,
});
Ok(Box::into_raw(exec_context) as i64)
@@ -782,7 +795,8 @@ pub unsafe extern "system" fn
Java_org_apache_comet_Native_executePlan(
PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx), partition)
.with_exec_id(exec_context_id)
.with_sql_text_pool(&exec_context.spark_plan)
- .with_task_context(exec_context.task_context.clone());
+ .with_task_context(exec_context.task_context.clone())
+ .with_class_loader(exec_context.class_loader.clone());
let (scans, shuffle_scans, root_op) = planner.create_plan(
&exec_context.spark_plan,
&mut exec_context.input_sources.clone(),
diff --git a/native/core/src/execution/planner.rs
b/native/core/src/execution/planner.rs
index 43f066e76c..c179c3b57c 100644
--- a/native/core/src/execution/planner.rs
+++ b/native/core/src/execution/planner.rs
@@ -250,6 +250,10 @@ pub struct PhysicalPlanner {
/// Captured at `createPlan` time on `ExecutionContext`; see that struct
for the
/// propagation rationale. `None` when no driving Spark task is available.
task_context: Option<Arc<Global<JObject<'static>>>>,
+ /// Context `ClassLoader` of the driving Spark task thread, captured at
`createPlan` time on
+ /// `ExecutionContext`; see that struct for the propagation rationale.
`None` when no driving
+ /// Spark task is available.
+ class_loader: Option<Arc<Global<JObject<'static>>>>,
}
impl Default for PhysicalPlanner {
@@ -267,6 +271,7 @@ impl PhysicalPlanner {
query_context_registry:
datafusion_comet_spark_expr::create_query_context_map(),
sql_text_pool: vec![],
task_context: None,
+ class_loader: None,
}
}
@@ -356,6 +361,17 @@ impl PhysicalPlanner {
self
}
+ /// Attach the driving Spark task thread's context `ClassLoader` as a
global reference. Mirrors
+ /// `with_task_context`: called by the JNI `executePlan` entry with
whatever was captured at
+ /// `createPlan` time, and cloned into every `JvmScalarUdfExpr` the
planner builds.
+ pub fn with_class_loader(
+ mut self,
+ class_loader: Option<Arc<Global<JObject<'static>>>>,
+ ) -> Self {
+ self.class_loader = class_loader;
+ self
+ }
+
/// Return session context of this planner.
pub fn session_ctx(&self) -> &Arc<SessionContext> {
&self.session_ctx
@@ -876,12 +892,14 @@ impl PhysicalPlanner {
to_arrow_datatype(udf.return_type.as_ref().ok_or_else(|| {
GeneralError("JvmScalarUdf missing
return_type".to_string())
})?);
- // Invariant: task_context is propagated for every
JvmScalarUdfExpr built during
- // normal execution. The TEST_EXEC_CONTEXT_ID path is the only
context in which
- // task_context may legitimately be None (unit tests, direct
native driver runs).
+ // Invariant: task_context and class_loader are propagated for
every
+ // JvmScalarUdfExpr built during normal execution. The
TEST_EXEC_CONTEXT_ID path is
+ // the only context in which they may legitimately be None
(unit tests, direct
+ // native driver runs).
debug_assert!(
- self.task_context.is_some() || self.exec_context_id ==
TEST_EXEC_CONTEXT_ID,
- "task_context must be set for non-test execution"
+ (self.task_context.is_some() &&
self.class_loader.is_some())
+ || self.exec_context_id == TEST_EXEC_CONTEXT_ID,
+ "task_context and class_loader must be set for non-test
execution"
);
Ok(Arc::new(JvmScalarUdfExpr::new(
udf.class_name.clone(),
@@ -889,6 +907,7 @@ impl PhysicalPlanner {
return_type,
udf.return_nullable,
self.task_context.clone(),
+ self.class_loader.clone(),
)))
}
expr => Err(GeneralError(format!("Not implemented: {expr:?}"))),
diff --git a/native/jni-bridge/src/comet_udf_bridge.rs
b/native/jni-bridge/src/comet_udf_bridge.rs
index e531d20cb1..bbd7465d00 100644
--- a/native/jni-bridge/src/comet_udf_bridge.rs
+++ b/native/jni-bridge/src/comet_udf_bridge.rs
@@ -41,7 +41,9 @@ impl<'a> CometUdfBridge<'a> {
method_evaluate: env.get_static_method_id(
JNIString::new(Self::JVM_CLASS),
jni::jni_str!("evaluate"),
-
jni::jni_sig!("(Ljava/lang/String;[J[JJJILorg/apache/spark/TaskContext;)V"),
+ jni::jni_sig!(
+
"(Ljava/lang/String;[J[JJJILorg/apache/spark/TaskContext;Ljava/lang/ClassLoader;)V"
+ ),
)?,
method_evaluate_ret: ReturnType::Primitive(Primitive::Void),
class,
diff --git a/native/spark-expr/src/jvm_udf/mod.rs
b/native/spark-expr/src/jvm_udf/mod.rs
index 0ca603ac9b..8e148ef40a 100644
--- a/native/spark-expr/src/jvm_udf/mod.rs
+++ b/native/spark-expr/src/jvm_udf/mod.rs
@@ -48,6 +48,11 @@ pub struct JvmScalarUdfExpr {
/// Spark task is available; the bridge then leaves whatever
`TaskContext.get()` already
/// returns in place.
task_context: Option<Arc<Global<JObject<'static>>>>,
+ /// Context `ClassLoader` of the driving Spark task thread, captured at
`createPlan` time and
+ /// threaded here by the planner. See `CometUdfBridge.evaluate`, which
installs it for the
+ /// duration of the call. `None` when no driving Spark task is available
(unit tests, direct
+ /// native driver runs); the bridge then installs nothing.
+ class_loader: Option<Arc<Global<JObject<'static>>>>,
}
impl JvmScalarUdfExpr {
@@ -57,6 +62,7 @@ impl JvmScalarUdfExpr {
return_type: DataType,
return_nullable: bool,
task_context: Option<Arc<Global<JObject<'static>>>>,
+ class_loader: Option<Arc<Global<JObject<'static>>>>,
) -> Self {
debug_assert!(
!class_name.is_empty(),
@@ -68,6 +74,7 @@ impl JvmScalarUdfExpr {
return_type,
return_nullable,
task_context,
+ class_loader,
}
}
}
@@ -197,14 +204,17 @@ impl PhysicalExpr for JvmScalarUdfExpr {
.set_region(env, 0, &in_sch_ptrs)
.map_err(|e| CometError::JNI { source: e })?;
- // Resolve the TaskContext reference once before building the arg
array so the
- // borrow lives until `call_static_method_unchecked` returns. When
no TaskContext
- // was propagated, pass a null object so the bridge's null-guard
leaves the thread-
- // local alone.
- let null_task_context = JObject::null();
+ // Resolve the TaskContext and ClassLoader references once before
building the arg
+ // array so the borrows live until `call_static_method_unchecked`
returns. Absent
+ // values are passed as a null object, which the bridge's
null-guards skip.
+ let null_obj = JObject::null();
let task_context_ref: &JObject = match &self.task_context {
Some(gref) => gref.as_obj(),
- None => &null_task_context,
+ None => &null_obj,
+ };
+ let class_loader_ref: &JObject = match &self.class_loader {
+ Some(gref) => gref.as_obj(),
+ None => &null_obj,
};
let ret = unsafe {
env.call_static_method_unchecked(
@@ -219,6 +229,7 @@ impl PhysicalExpr for JvmScalarUdfExpr {
JValue::Long(out_sch_ptr).as_jni(),
JValue::Int(batch.num_rows() as i32).as_jni(),
JValue::Object(task_context_ref).as_jni(),
+ JValue::Object(class_loader_ref).as_jni(),
],
)
};
@@ -255,6 +266,7 @@ impl PhysicalExpr for JvmScalarUdfExpr {
self.return_type.clone(),
self.return_nullable,
self.task_context.clone(),
+ self.class_loader.clone(),
)))
}
}
diff --git a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
index 9e97ef2226..d8dea73135 100644
--- a/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
+++ b/spark/src/main/java/org/apache/comet/udf/CometUdfBridge.java
@@ -31,6 +31,8 @@ import org.apache.spark.TaskContext;
import org.apache.spark.comet.CometTaskContextShim;
import org.apache.spark.util.TaskCompletionListener;
+import org.apache.comet.util.ClassLoaders;
+
/**
* JNI entry point for native execution to invoke a {@link CometUDF}. Matches
the static-method
* pattern used by CometScalarSubquery so the native side can dispatch via
@@ -91,6 +93,17 @@ public class CometUdfBridge {
* left on a worker by a previous task. Its task attempt ID also keys
the UDF-instance cache,
* so a UDF holding per-task state in fields sees a consistent instance
for every call within
* the task regardless of which Tokio worker is polling.
+ * @param classLoader context ClassLoader captured on the driving Spark task
thread, or {@code
+ * null} outside a Spark task. Installed as this thread's context
ClassLoader for the duration
+ * of the call, with the prior value restored in {@code finally}. Tokio
workers attach through
+ * JNI and an attached thread has no context ClassLoader, so without
this every lookup falls
+ * back to the ClassLoader that loaded Comet, which never holds user
jars ({@code --jars} /
+ * {@code spark.jars}). Both the {@code CometUDF} resolution below and
the closure
+ * deserialization in {@code CometScalaUDFCodegen} depend on it.
Installed per call rather
+ * than once when the worker attaches, because the Tokio runtime is
process-global: one worker
+ * interleaves work from task attempts of different jobs, and under
Spark Connect from
+ * sessions with different artifact ClassLoaders. The loader is a
property of the plan, not of
+ * the thread.
*/
public static void evaluate(
String udfClassName,
@@ -99,7 +112,8 @@ public class CometUdfBridge {
long outArrayPtr,
long outSchemaPtr,
int numRows,
- TaskContext taskContext) {
+ TaskContext taskContext,
+ ClassLoader classLoader) {
assert udfClassName != null && !udfClassName.isEmpty() : "udfClassName
must be non-empty";
assert inputArrayPtrs != null && inputSchemaPtrs != null
: "input pointer arrays must be non-null";
@@ -111,13 +125,20 @@ public class CometUdfBridge {
// Save-and-restore rather than only-install-if-null: the propagated
`taskContext` is the
// ground truth for this call. Any value already on the thread is either
(a) the same object
- // on a Spark task thread, or (b) stale from a prior task on a reused
Tokio worker.
+ // on a Spark task thread, or (b) stale from a prior task on a reused
Tokio worker. The same
+ // reasoning applies to the propagated `classLoader`.
TaskContext prior = TaskContext.get();
if (taskContext != null) {
CometTaskContextShim.set(taskContext);
assert TaskContext.get() == taskContext
: "TaskContext install did not take effect on this thread";
}
+ Thread currentThread = Thread.currentThread();
+ ClassLoader priorLoader = currentThread.getContextClassLoader();
+ if (classLoader != null) {
+ currentThread.setContextClassLoader(classLoader);
+ }
+
try {
evaluateInternal(
udfClassName,
@@ -128,6 +149,9 @@ public class CometUdfBridge {
numRows,
taskContext);
} finally {
+ // Unconditional: a no-op when nothing was installed, and it also undoes
any change the
+ // user function made to the ClassLoader of a worker that outlives this
call.
+ currentThread.setContextClassLoader(priorLoader);
if (taskContext != null) {
if (prior != null) {
CometTaskContextShim.set(prior);
@@ -175,14 +199,10 @@ public class CometUdfBridge {
udfClassName,
name -> {
try {
- // Resolve via the executor's context classloader so
user-supplied UDF jars
- // (added via spark.jars / --jars) are visible.
- ClassLoader cl =
Thread.currentThread().getContextClassLoader();
- if (cl == null) {
- cl = CometUdfBridge.class.getClassLoader();
- }
+ // Resolves through the context ClassLoader installed by
`evaluate`, so a
+ // user-supplied CometUDF shipped in a user jar is visible.
return (CometUDF)
- Class.forName(name, true,
cl).getDeclaredConstructor().newInstance();
+
ClassLoaders.loadClass(name).getDeclaredConstructor().newInstance();
} catch (ReflectiveOperationException e) {
throw new RuntimeException("Failed to instantiate CometUDF: "
+ name, e);
}
diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
index a3216404ac..57578235f7 100644
--- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
+++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala
@@ -125,8 +125,11 @@ class CometExecIterator(
taskCPUs,
keyUnwrapper,
// Propagated to Tokio workers running JVM UDFs so they see this Spark
task's
- // TaskContext. See CometUdfBridge.evaluate.
- TaskContext.get())
+ // TaskContext and context ClassLoader. Read here because this class is
only ever
+ // constructed on a Spark task thread (see `taskAttemptId` above); a
JNI-attached Tokio
+ // worker has neither. See CometUdfBridge.evaluate.
+ TaskContext.get(),
+ Thread.currentThread().getContextClassLoader)
}
private var nextBatch: Option[ColumnarBatch] = None
diff --git a/spark/src/main/scala/org/apache/comet/Native.scala
b/spark/src/main/scala/org/apache/comet/Native.scala
index 3cfa51b6e1..51b3e5e41b 100644
--- a/spark/src/main/scala/org/apache/comet/Native.scala
+++ b/spark/src/main/scala/org/apache/comet/Native.scala
@@ -48,6 +48,11 @@ class Native extends NativeBase {
* @param taskMemoryManager
* the task-level memory manager that is responsible for tracking memory
usage across JVM and
* native side.
+ * @param taskContext
+ * the driving Spark task's `TaskContext`, propagated to Tokio workers
running JVM UDFs.
+ * @param classLoader
+ * the calling Spark task thread's context ClassLoader, propagated to
Tokio workers running
+ * JVM UDFs. Must be read on the task thread; see
`CometUdfBridge.evaluate`.
* @return
* the address to native query plan.
*/
@@ -70,7 +75,8 @@ class Native extends NativeBase {
taskAttemptId: Long,
taskCPUs: Long,
keyUnwrapper: CometFileKeyUnwrapper,
- taskContext: TaskContext): Long
+ taskContext: TaskContext,
+ classLoader: ClassLoader): Long
// scalastyle:on
/**
diff --git
a/spark/src/test/scala/org/apache/comet/CometScalaUDFClassLoaderSuite.scala
b/spark/src/test/scala/org/apache/comet/CometScalaUDFClassLoaderSuite.scala
new file mode 100644
index 0000000000..17fd0fea09
--- /dev/null
+++ b/spark/src/test/scala/org/apache/comet/CometScalaUDFClassLoaderSuite.scala
@@ -0,0 +1,195 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet
+
+import java.io.File
+import java.net.URLClassLoader
+import java.nio.charset.StandardCharsets.UTF_8
+import java.nio.file.{Files, Path, Paths}
+import javax.tools.ToolProvider
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.{CometTestBase, Encoders}
+import org.apache.spark.sql.catalyst.expressions.Expression
+
+/**
+ * Regression coverage for UDF closures whose capturing class lives in a user
jar, which used to
+ * fail with:
+ *
+ * {{{
+ * 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
+ * }}}
+ *
+ * That exception 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 classloading
rationale lives on
+ * `CometUdfBridge.evaluate`.
+ *
+ * The fixture puts the capturing class on `spark.executor.extraClassPath`,
the local-mode
+ * equivalent of a `--jars` submission: `LocalSchedulerBackend` feeds it into
the executor's
+ * `MutableURLClassLoader`, which becomes the task thread's context
ClassLoader while staying
+ * invisible to the ClassLoader that loaded Comet.
+ */
+class CometScalaUDFClassLoaderSuite extends CometTestBase {
+
+ import CometScalaUDFClassLoaderSuite._
+
+ override protected def sparkConf: SparkConf =
+ super.sparkConf
+ .set(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key, "true")
+ .set("spark.executor.extraClassPath", hiddenClassesDir.toString)
+
+ private def withHiddenUdfTable(f: => Unit): Unit = {
+ spark.udf.register("hiddenUdf", hiddenFn)
+ withTable("t") {
+ sql("CREATE TABLE t (s STRING) USING parquet")
+ sql("INSERT INTO t VALUES ('a'), ('b'), (NULL)")
+ f
+ }
+ }
+
+ test("fixture: hidden class reachable from task threads, not from Comet's
ClassLoader") {
+ // Guards the tests below from passing vacuously.
`classOf[Expression].getClassLoader` is the
+ // fallback `lookupOrCompile` uses when the calling thread has no context
ClassLoader.
+ intercept[ClassNotFoundException] {
+ Class.forName(HiddenClassName, false, classOf[Expression].getClassLoader)
+ }
+ val reachable = spark
+ .range(4)
+ .repartition(2)
+ .mapPartitions(_ => Iterator(canLoadHiddenClass))(Encoders.scalaBoolean)
+ .collect()
+ assert(reachable.forall(identity), s"task threads could not load
$HiddenClassName")
+ }
+
+ // Both leaf shapes matter. With the native scan the dispatcher runs on a
Tokio worker, which has
+ // no context ClassLoader of its own and depends on the propagated one;
without it the dispatcher
+ // runs on the Spark task thread, which already has the executor's
ClassLoader installed.
+ Seq(true, false).foreach { nativeScan =>
+ test(s"ScalaUDF closure captured by a user-jar class,
nativeScan=$nativeScan") {
+ withSQLConf(CometConf.COMET_NATIVE_SCAN_ENABLED.key ->
nativeScan.toString) {
+ withHiddenUdfTable {
+ checkSparkAnswer(sql("SELECT hiddenUdf(s) FROM t"))
+ }
+ }
+ }
+ }
+
+ test("the thread running the UDF sees the task thread's context
ClassLoader") {
+ // Pins the propagation itself rather than its symptom: the UDF body
reports whether the
+ // ClassLoader installed on whatever thread invoked it can reach the user
jar.
+ spark.udf.register("loaderProbe", (_: String) => loaderReport())
+ withHiddenUdfTable {
+ val reports = sql("SELECT loaderProbe(s) FROM
t").collect().map(_.getString(0)).distinct
+ assert(
+ reports.forall(_.startsWith("loaded|")),
+ s"UDF thread could not load $HiddenClassName: ${reports.mkString(",
")}")
+ }
+ }
+}
+
+object CometScalaUDFClassLoaderSuite {
+
+ val HiddenClassName = "hidden.HiddenUdf"
+
+ /**
+ * Directory holding the compiled capturing class. Compiled once per JVM,
before the
+ * SparkSession starts, because `spark.executor.extraClassPath` is read at
session creation. A
+ * directory works as a classpath entry, so there is no need to package a
jar.
+ */
+ lazy val hiddenClassesDir: Path = compileHiddenClass()
+
+ /**
+ * The UDF, obtained through a ClassLoader over `hiddenClassesDir` alone.
Deliberately not
+ * closed: the loaded class stays live in the registered UDF for the rest of
the JVM's life.
+ */
+ lazy val hiddenFn: String => String =
+ new URLClassLoader(Array(hiddenClassesDir.toUri.toURL),
getClass.getClassLoader)
+ .loadClass(HiddenClassName)
+ .getMethod("make")
+ .invoke(null)
+ .asInstanceOf[String => String]
+
+ /** Runs on Spark threads; kept on the companion so the closures capture
nothing. */
+ def canLoadHiddenClass: Boolean =
+ try {
+ Class.forName(HiddenClassName, false,
Thread.currentThread().getContextClassLoader)
+ true
+ } catch {
+ case _: ClassNotFoundException => false
+ }
+
+ /**
+ * Reports whether the current thread's context ClassLoader reaches the user
jar, tagged with
+ * the thread name so a failure says which thread was missing it.
+ */
+ def loaderReport(): String = {
+ val status = if (canLoadHiddenClass) "loaded" else "MISSING"
+ s"$status|${Thread.currentThread().getName}"
+ }
+
+ private def compileHiddenClass(): Path = {
+ // createTempDirectory does not create its parent (java.io.tmpdir, pinned
to target/tmp by the
+ // pom), which does not exist yet on a fresh checkout, so make it up front.
+ val tmpRoot =
Files.createDirectories(Paths.get(System.getProperty("java.io.tmpdir")))
+ val workDir = Files.createTempDirectory(tmpRoot, "comet-hidden-udf")
+ val src = workDir.resolve("HiddenUdf.java")
+ // The intersection cast is what makes the lambda serializable, and
therefore what makes
+ // `hidden.HiddenUdf` the capturing class recorded in the SerializedLambda.
+ Files.write(
+ src,
+ """package hidden;
+ |
+ |import java.io.Serializable;
+ |import scala.Function1;
+ |
+ |public class HiddenUdf {
+ | public static Function1<Object, Object> make() {
+ | return (Function1<Object, Object> & Serializable)
+ | (Object o) -> (o == null ? null : "hidden:" + o);
+ | }
+ |}
+ |""".stripMargin.getBytes(UTF_8))
+
+ val classesDir = Files.createDirectories(workDir.resolve("classes"))
+ val compiler = ToolProvider.getSystemJavaCompiler
+ assert(compiler != null, "test must run on a JDK (needs the javax.tools
compiler)")
+ // Only scala-library is needed. Handing javac the whole test classpath
makes it open and index
+ // every jar on it, which costs more than the compile itself.
+ val classpath = Option(classOf[Function1[_,
_]].getProtectionDomain.getCodeSource)
+ .map(_.getLocation.getPath)
+ .getOrElse(System.getProperty("java.class.path"))
+ val rc =
+ compiler.run(null, null, null, "-cp", classpath, "-d",
classesDir.toString, src.toString)
+ assert(rc == 0, s"javac failed with exit code $rc")
+
+ deleteOnExitRecursively(workDir.toFile)
+ classesDir
+ }
+
+ /** Parents are registered before children, and deletion runs in reverse
registration order. */
+ private def deleteOnExitRecursively(file: File): Unit = {
+ file.deleteOnExit()
+ Option(file.listFiles()).foreach(_.foreach(deleteOnExitRecursively))
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
