Fix after review

Author: zhenlineo

connector/connect/client/jvm/pom.xml

@@ -221,13 +221,6 @@
         <groupId>org.apache.maven.plugins</groupId>
         <artifactId>maven-jar-plugin</artifactId>
         <executions>
-          <execution>
-            <id>default-jar</id>
-            <phase>compile</phase>
-            <goals>
-              <goal>jar</goal>
-            </goals>
-          </execution>
           <execution>
             <id>prepare-test-jar</id>
             <phase>test-compile</phase>

connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/UDFClassLoadingE2ESuite.scala

@@ -46,7 +46,7 @@ class UDFClassLoadingE2ESuite extends RemoteSparkSession {
       .range(10)
       .filter(n => {
         // Try to use spark result
-        new SparkResult[Int](null, null, null)
+        new SparkResult[Int](null, null, null, "")
         n > 5
       })
       .collectAsList()
@@ -58,17 +58,20 @@ class UDFClassLoadingE2ESuite extends RemoteSparkSession {
     addClientTestArtifactInServerClasspath(session)
     val ds = session.range(10).filter(n => n % 2 == 0)
 
-    // load SparkResult as a stubbed class
+    // Load SparkResult as a stubbed class
     val rows = ds.collectAsList()
     assert(rows == Arrays.asList[Long](0, 2, 4, 6, 8))
 
     // Upload SparkResult and then SparkResult can be used in the udf
     addClientTestArtifactInServerClasspath(session, testJar = false)
-    val rows2 = session.range(10).filter(n => {
-      // Try to use spark result
-      new SparkResult[Int](null, null, null)
-      n > 5
-    }).collectAsList()
+    val rows2 = session
+      .range(10)
+      .filter(n => {
+        // Try to use spark result
+        new SparkResult[Int](null, null, null, "")
+        n > 5
+      })
+      .collectAsList()
     assert(rows2 == Arrays.asList[Long](6, 7, 8, 9))
   }
 

connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/util/IntegrationTestUtils.scala

@@ -29,8 +29,8 @@ object IntegrationTestUtils {
 
   // System properties used for testing and debugging
   private val DEBUG_SC_JVM_CLIENT = "spark.debug.sc.jvm.client"
-  // Enable this flag to print all client debug log + server logs to the console
-  private[connect] val isDebug = System.getProperty(DEBUG_SC_JVM_CLIENT, "true").toBoolean
+  // Enable this flag to print all server logs to the console
+  private[connect] val isDebug = System.getProperty(DEBUG_SC_JVM_CLIENT, "false").toBoolean
 
   private[sql] lazy val scalaVersion = {
     versionNumberString.split('.') match {

connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/util/RemoteSparkSession.scala

@@ -96,11 +96,9 @@ object SparkConnectServerUtils {
     // To find InMemoryTableCatalog for V2 writer tests
     val catalystTestJar =
       tryFindJar("sql/catalyst", "spark-catalyst", "spark-catalyst", test = true)
-        .map(clientTestJar => Seq("--jars", clientTestJar.getCanonicalPath))
+        .map(jar => Seq("--jars", jar.getCanonicalPath))
         .getOrElse(Seq.empty)
 
-    val jarsConfigs = Seq("--jars", catalystTestJar.mkString(","))
-
     // Use InMemoryTableCatalog for V2 writer tests
     val writerV2Configs = Seq(
       "--conf",
@@ -127,7 +125,7 @@ object SparkConnectServerUtils {
       Seq("--conf", s"spark.sql.catalogImplementation=$catalogImplementation")
     }
 
-    jarsConfigs ++ writerV2Configs ++ hiveTestConfigs
+    catalystTestJar ++ writerV2Configs ++ hiveTestConfigs
   }
 
   def start(): Unit = {
@@ -198,9 +196,10 @@ trait RemoteSparkSession extends ConnectFunSuite with BeforeAndAfterAll {
         debug(error)
         throw error
       }
-    }
 
-    addClientTestArtifactInServerClasspath(spark)
+      // Add client test jar into the spark session classpath
+      addClientTestArtifactInServerClasspath(spark)
+    }
   }
 
   // For UDF maven E2E tests, the server needs the client test code to find the UDFs defined in
@@ -214,8 +213,8 @@ trait RemoteSparkSession extends ConnectFunSuite with BeforeAndAfterAll {
       // So we skip building or finding this jar for SBT.
       "sbt-tests-do-not-need-this-jar",
       "spark-connect-client-jvm",
-      test = testJar
-    ).foreach(clientTestJar => session.addArtifact(clientTestJar.getCanonicalPath))
+      test = testJar).foreach(clientTestJar =>
+      session.addArtifact(clientTestJar.getCanonicalPath))
   }
 
   override def afterAll(): Unit = {

connector/connect/server/src/main/scala/org/apache/spark/sql/connect/artifact/SparkConnectArtifactManager.scala

@@ -164,7 +164,10 @@ class SparkConnectArtifactManager(sessionHolder: SessionHolder) extends Logging
     val urls = getSparkConnectAddedJars :+ classDir.toUri.toURL
     val stubClassLoader =
       StubClassLoader(null, SparkEnv.get.conf.get(CONNECT_SCALA_UDF_STUB_CLASSES))
-    new ChildFirstURLClassLoader(urls.toArray, stubClassLoader, Utils.getContextOrSparkClassLoader)
+    new ChildFirstURLClassLoader(
+      urls.toArray,
+      stubClassLoader,
+      Utils.getContextOrSparkClassLoader)
   }
 
   /**

connector/connect/server/src/test/scala/org/apache/spark/sql/connect/artifact/StubClassLoaderSuite.scala

@@ -61,11 +61,8 @@ class StubClassLoaderSuite extends SparkFunSuite {
     val stubClassLoader = new RecordedStubClassLoader(null, _ => true)
 
     // Install artifact without class A.
-    val sessionClassLoader = new ChildFirstURLClassLoader(
-      Array(udfNoAJar),
-      stubClassLoader,
-      sysClassLoader
-    )
+    val sessionClassLoader =
+      new ChildFirstURLClassLoader(Array(udfNoAJar), stubClassLoader, sysClassLoader)
     // Load udf with A used in the same class.
     deserializeUdf(sessionClassLoader)
     // Class A should be stubbed.
@@ -77,25 +74,22 @@ class StubClassLoaderSuite extends SparkFunSuite {
     val sysClassLoader = getClass.getClassLoader()
     val stubClassLoader = new RecordedStubClassLoader(null, _ => true)
 
-    val cl1 = new ChildFirstURLClassLoader(
-      Array.empty,
-      stubClassLoader,
-      sysClassLoader)
+    val cl1 = new ChildFirstURLClassLoader(Array.empty, stubClassLoader, sysClassLoader)
 
     // Failed to load dummy udf
-    intercept[Exception]{
+    intercept[Exception] {
       deserializeUdf(cl1)
     }
     // Successfully stubbed the missing class.
-    assert(stubClassLoader.lastStubbed ===
-      "org.apache.spark.sql.connect.artifact.StubClassDummyUdf")
+    assert(
+      stubClassLoader.lastStubbed ===
+        "org.apache.spark.sql.connect.artifact.StubClassDummyUdf")
 
     // Creating a new class loader will unpack the udf correctly.
     val cl2 = new ChildFirstURLClassLoader(
       Array(udfNoAJar),
       stubClassLoader, // even with the same stub class loader.
-      sysClassLoader
-    )
+      sysClassLoader)
     // Should be able to load after the artifact is added
     deserializeUdf(cl2)
   }
@@ -105,37 +99,32 @@ class StubClassLoaderSuite extends SparkFunSuite {
     val sysClassLoader = getClass.getClassLoader()
     val stubClassLoader = new RecordedStubClassLoader(null, _ => true)
 
-    val sessionClassLoader = new ChildFirstURLClassLoader(
-      Array.empty,
-      stubClassLoader,
-      sysClassLoader)
+    val sessionClassLoader =
+      new ChildFirstURLClassLoader(Array.empty, stubClassLoader, sysClassLoader)
 
     // Failed to load dummy udf
-    val exception = intercept[Exception]{
+    val exception = intercept[Exception] {
       deserializeUdf(sessionClassLoader)
     }
     // Successfully stubbed the missing class.
-    assert(stubClassLoader.lastStubbed ===
-      "org.apache.spark.sql.connect.artifact.StubClassDummyUdf")
+    assert(
+      stubClassLoader.lastStubbed ===
+        "org.apache.spark.sql.connect.artifact.StubClassDummyUdf")
     // But failed to find the method on the stub class.
     val cause = exception.getCause
     assert(cause.isInstanceOf[NoSuchMethodException])
     assert(
       cause.getMessage.contains("org.apache.spark.sql.connect.artifact.StubClassDummyUdf"),
-      cause.getMessage
-    )
+      cause.getMessage)
   }
 
   private def deserializeUdf(sessionClassLoader: ClassLoader): UdfPacket = {
-    Utils.deserialize[UdfPacket](
-      udfByteArray,
-      sessionClassLoader
-    )
+    Utils.deserialize[UdfPacket](udfByteArray, sessionClassLoader)
   }
 }
 
 class RecordedStubClassLoader(parent: ClassLoader, shouldStub: String => Boolean)
-  extends StubClassLoader(parent, shouldStub) {
+    extends StubClassLoader(parent, shouldStub) {
   var lastStubbed: String = _
 
   override def findClass(name: String): Class[_] = {

core/src/main/java/org/apache/spark/util/ChildFirstURLClassLoader.java

@@ -42,7 +42,7 @@ public ChildFirstURLClassLoader(URL[] urls, ClassLoader parent) {
 
   /**
    * Specify the grandparent if there is a need to load in the order of
-   * `grandparent -> urls (child) -> parent`.
+   * `grandparent -> urls (child) -> parent`.
    */
   public ChildFirstURLClassLoader(URL[] urls, ClassLoader parent, ClassLoader grandparent) {
     super(urls, grandparent);

core/src/main/scala/org/apache/spark/executor/Executor.scala

@@ -554,7 +554,8 @@ private[spark] class Executor(
           taskDescription.artifacts.jars,
           taskDescription.artifacts.archives,
           isolatedSession)
-        // Always reset the thread class loader
+        // Always reset the thread class loader to ensure if any updates, all threads (not only
+        // the thread that updated the dependencies) can update to the new class loader.
         Thread.currentThread.setContextClassLoader(isolatedSession.replClassLoader)
         task = ser.deserialize[Task[Any]](
           taskDescription.serializedTask, Thread.currentThread.getContextClassLoader)
@@ -1076,7 +1077,7 @@ private[spark] class Executor(
     } else {
       parent
     }
-    logInfo(s"Created or updated repl class loader for $sessionUUID.")
+    logInfo(s"Created or updated repl class loader $classLoader for $sessionUUID.")
     classLoader
   }
 
@@ -1155,7 +1156,8 @@ private[spark] class Executor(
         }
       }
       if (updated) {
-        // TODO: only update the class loader if the stub class should be unloaded.
+        // When a new url is added for non-default class loader, recreate the class loader
+        // to ensure all classes are updated.
         state.urlClassLoader = createClassLoader(state.urlClassLoader.getURLs, useStub = true)
         state.replClassLoader =
           addReplClassLoaderIfNeeded(state.urlClassLoader, state.replClassDirUri, state.sessionUUID)

core/src/main/scala/org/apache/spark/internal/config/package.scala

@@ -2552,10 +2552,10 @@ package object config {
     ConfigBuilder("spark.connect.scalaUdf.stubClasses")
       .internal()
       .doc("""
-          |Comma separated list of binary names of classes/packages that should be stub during the
-          |Scala UDF serdeser and execution if not found on the server classpath.
+          |Comma-separated list of binary names of classes/packages that should be stubbed during
+          |the Scala UDF serde and execution if not found on the server classpath.
           |An empty list effectively disables stubbing for all missing classes.
-          |By default the server stubs classes from the Scala client package.
+          |By default, the server stubs classes from the Scala client package.
           |""".stripMargin)
       .version("3.5.0")
       .stringConf

core/src/main/scala/org/apache/spark/util/StubClassLoader.scala

@@ -33,6 +33,8 @@ class StubClassLoader(parent: ClassLoader, shouldStub: String => Boolean)
       throw new ClassNotFoundException(name)
     }
     val bytes = StubClassLoader.generateStub(name)
+    // scalastyle:off println
+    println("###stub class: " + name)
     defineClass(name, bytes, 0, bytes.length)
   }
 }