[AURON #2386] Resolve deserialized expression classes with an explicit class loader

.github/workflows/tpcds-reusable.yml

@@ -64,6 +64,32 @@ on:
         required: false
         type: string
         default: ''
+      assert-log-matches:
+        description: >-
+          Optional grep -E pattern that the TPC-DS run log must contain. Use it to assert that
+          the plan shape a job exists to cover was actually produced, so the job cannot pass
+          because the optimizer stopped generating it. Implies --print-plan.
+        required: false
+        type: string
+        default: ''
+      assert-log-not-matches:
+        description: >-
+          Optional grep -E pattern that must not appear in the TPC-DS run log. Spark retries
+          absorb task-level failures, so queries can report PASS while throwing hundreds of
+          exceptions; this asserts on the log instead. Keep the pattern specific to the defect
+          under test, otherwise unrelated failures are attributed to it.
+        required: false
+        type: string
+        default: ''
+      jar-on-system-classpath:
+        description: >-
+          Whether to copy the Auron jar into $SPARK_HOME/jars, where the application class
+          loader defines it. When false, Auron is defined by Spark's MutableURLClassLoader
+          instead, which is what spark-submit --jars and the application jar both feed. Some
+          class-loading defects only reproduce in the latter configuration.
+        required: false
+        type: string
+        default: 'true'
       queries:
         description: 'Optional list of queries to run'
         required: false
@@ -257,12 +283,18 @@ jobs:
           path: dev/tpcds_1g
 
       - name: Install Auron JAR
+        env:
+          JAR_ON_SYSTEM_CLASSPATH: ${{ inputs.jar-on-system-classpath }}
         run: |
           ls -la
           jar=$(ls -1 auron-${{ inputs.sparkver }}_${{ inputs.scalaver }}*.jar | head -n1)
           [ -n "$jar" ] || { echo "No jar matched: auron-${{ inputs.sparkver }}_${{ inputs.scalaver }}*.jar"; exit 1; }
           echo "AURON_SPARK_JAR=$jar" >> "$GITHUB_ENV"
-          cp "$jar" spark-bin-${{ inputs.sparkver }}_${{ inputs.scalaver }}/jars/
+          if [ "$JAR_ON_SYSTEM_CLASSPATH" = "false" ]; then
+            echo "Auron stays off the system classpath (MutableURLClassLoader defines it)"
+          else
+            cp "$jar" spark-bin-${{ inputs.sparkver }}_${{ inputs.scalaver }}/jars/
+          fi
 
       - name: Setup Java and Maven cache
         uses: actions/setup-java@v5
@@ -378,15 +410,65 @@ jobs:
           SPARK_VERSION: ${{ inputs.sparkver }}
           SCALA_VERSION: ${{ inputs.scalaver }}
           SPARK_HOME: spark-bin-${{ inputs.sparkver }}_${{ inputs.scalaver }}
+          PRINT_PLAN: ${{ inputs.assert-log-matches != '' && '--print-plan' || '' }}
         run: |
           ls -la
+          set -o pipefail
           dev/auron-it/run-it.sh  \
             ${{ inputs.extrasparkconf }} \
             --type tpcds \
             --data-location dev/tpcds_1g \
             --query-filter ${{ matrix.query }} \
             --result-check \
-            --plan-check
+            $PRINT_PLAN \
+            --plan-check 2>&1 | tee tpcds-run-${{ matrix.query }}.log
+
+      # The auron-it jar depends on the Auron uber jar and so bundles Auron classes itself, which
+      # means skipping the copy into $SPARK_HOME/jars is not on its own proof of how Auron was
+      # loaded. Assert the property the job actually depends on: that some loader other than the
+      # application class loader defined Auron. auron-it prints this at startup.
+      - name: Assert Auron is not on the system classpath
+        if: ${{ inputs.jar-on-system-classpath == 'false' }}
+        env:
+          QUERY_LOG: tpcds-run-${{ matrix.query }}.log
+        run: |
+          grep -E '^Auron (Class Loader|Code Source|On System Classpath):' "$QUERY_LOG" || true
+          if ! grep -qx 'Auron On System Classpath: false' "$QUERY_LOG"; then
+            echo "::error::Auron was defined by the application class loader, or auron-it did not report it"
+            exit 1
+          fi
+
+      # Guards against the job silently testing nothing: --plan-check only compares golden plans
+      # on Spark 3.5 (see PlanStabilityChecker), so on other versions a job whose purpose is a
+      # particular plan shape would still pass if the optimizer stopped producing it.
+      - name: Assert the run log matches ${{ inputs.assert-log-matches }}
+        if: ${{ inputs.assert-log-matches != '' }}
+        env:
+          QUERY_LOG: tpcds-run-${{ matrix.query }}.log
+          PATTERN: ${{ inputs.assert-log-matches }}
+        run: |
+          count=$(grep -c -E "$PATTERN" "$QUERY_LOG" || true)
+          if [ "$count" -eq 0 ]; then
+            echo "::error::expected the run log to match '$PATTERN', found no occurrence"
+            exit 1
+          fi
+          echo "Matched '$PATTERN' $count time(s)."
+
+      # Task-level failures are absorbed by Spark's task retries, so the queries can still
+      # report PASS while throwing hundreds of exceptions. Assert on the log directly.
+      - name: Assert the run log does not match ${{ inputs.assert-log-not-matches }}
+        if: ${{ inputs.assert-log-not-matches != '' }}
+        env:
+          QUERY_LOG: tpcds-run-${{ matrix.query }}.log
+          PATTERN: ${{ inputs.assert-log-not-matches }}
+        run: |
+          count=$(grep -c -E "$PATTERN" "$QUERY_LOG" || true)
+          if [ "$count" -gt 0 ]; then
+            echo "::error::run log matched '$PATTERN' $count time(s)"
+            grep -m5 -E "$PATTERN" "$QUERY_LOG" || true
+            exit 1
+          fi
+          echo "No occurrence of '$PATTERN'."
 
       - name: Upload RSS log
         if: ${{ failure() && (inputs.celebornver != '' || inputs.unifflever != '') }}

.github/workflows/tpcds.yml

@@ -113,6 +113,33 @@ jobs:
       hadoop-profile: 'hadoop3'
       sparktests: 'true'
 
+  test-spark-41-jdk17-scala-2-13-with-bloomfilter-optimizer-enabled:
+    name: Test spark-4.1 JDK17 Scala-2.13 with bloomFilter optimizer enabled
+    uses: ./.github/workflows/tpcds-reusable.yml
+    with:
+      sparkver: spark-4.1
+      javaver: '17'
+      scalaver: '2.13'
+      hadoop-profile: 'hadoop3'
+      sparktests: 'true'
+      # Keep Auron off the system classpath so Spark's MutableURLClassLoader defines it rather
+      # than the application class loader. The runtime bloom-filter ScalarSubquery
+      # deserialization defect (AURON #2386) cannot reproduce when the jar is in
+      # $SPARK_HOME/jars. The job asserts this arrangement rather than assuming it.
+      jar-on-system-classpath: 'false'
+      # q1, q2 and q3 are the queries the runtime bloom filter rewrites into the scalar
+      # subquery that triggered AURON #2386; running the other 96 adds no coverage here.
+      queries: '["q1,q2,q3"]'
+      # might_contain is the runtime bloom-filter probe, so its presence proves the optimizer
+      # still produced the plan this job exists to cover.
+      assert-log-matches: 'might_contain'
+      # The AURON #2386 signature specifically, not any ClassCastException.
+      assert-log-not-matches: 'ClassCastException.*DefaultSerializationProxy'
+      extrasparkconf: >-
+        --conf spark.sql.optimizer.runtime.bloomFilter.enabled=true
+        --conf spark.sql.optimizer.runtime.bloomFilter.applicationSideScanSizeThreshold=1B
+        --conf spark.sql.autoBroadcastJoinThreshold=-1
+
   test-spark-42-jdk21-scala-2-13:
     name: Test spark-4.2 JDK21 Scala-2.13
     uses: ./.github/workflows/tpcds-reusable.yml

dev/auron-it/src/main/scala/org/apache/auron/integration/Main.scala

@@ -64,6 +64,9 @@ object Main {
       opt[Unit]("regen-golden")
         .action((_, c) => c.copy(regenGoldenFiles = true))
         .text("regenerate golden files"),
+      opt[Unit]("print-plan")
+        .action((_, c) => c.copy(printPlan = true))
+        .text("print the Auron physical plan of every query (default: false)"),
       help('h', "help"))
   }
 
@@ -121,8 +124,27 @@ object Main {
     if (args.auronOnly) println("Mode: Auron-only (skip baseline)")
     if (args.enablePlanCheck) println("Plan Check: Enabled")
     if (args.regenGoldenFiles) println("Regenerate golden files: Enabled")
+    printClassLoaderSummary()
     println("-" * 60)
     println("")
   }
+
+  /**
+   * Reports which loader actually defined Auron. Auron classes reach the JVM either through the
+   * application class loader (a jar in $SPARK_HOME/jars) or through Spark's MutableURLClassLoader
+   * (spark-submit --jars, and the application jar itself). The two behave differently during
+   * deserialization, so tests covering class-loading defects such as AURON #2386 need to state
+   * which arrangement they ran under rather than infer it from how the job was set up.
+   */
+  private def printClassLoaderSummary(): Unit = {
+    val auronClass = Shims.get.getClass
+    val loader = auronClass.getClassLoader
+    val codeSource = Option(auronClass.getProtectionDomain.getCodeSource)
+      .map(_.getLocation.toString)
+      .getOrElse("unknown")
+    println(s"Auron Class Loader: ${loader.getClass.getName}")
+    println(s"Auron Code Source: $codeSource")
+    println(s"Auron On System Classpath: ${loader eq ClassLoader.getSystemClassLoader}")
+  }
 }
 // scalastyle:on

dev/auron-it/src/main/scala/org/apache/auron/integration/Suite.scala

@@ -23,7 +23,8 @@ case class SuiteArgs(
     extraSparkConf: Map[String, String] = Map.empty,
     auronOnly: Boolean = false,
     enablePlanCheck: Boolean = false,
-    regenGoldenFiles: Boolean = false)
+    regenGoldenFiles: Boolean = false,
+    printPlan: Boolean = false)
 
 abstract class Suite(val args: SuiteArgs) {
   protected lazy val sessions: SessionManager = new SessionManager(args.extraSparkConf)

dev/auron-it/src/main/scala/org/apache/auron/integration/tpcds/TPCDSSuite.scala

@@ -53,6 +53,14 @@ class TPCDSSuite(args: SuiteArgs) extends Suite(args) with TPCDSFeatures {
     setupTables(args.dataLocation, sessions.auronSession)
     val auronResults = queryRunner.runQueries(sessions.auronSession, queries)
 
+    if (args.printPlan) {
+      queries.foreach { queryId =>
+        println(s"=== Auron physical plan for $queryId ===")
+        println(auronResults(queryId).plan)
+        println(s"=== End of Auron physical plan for $queryId ===")
+      }
+    }
+
     val baseComparisons: Seq[ComparisonResult] =
       if (args.auronOnly) {
         queries.map { qid =>

pom.xml

@@ -424,6 +424,9 @@
 
             <arg>-Wconf:msg=method newInstance in class Class is deprecated:s</arg>
             <arg>-Wconf:msg=class ThreadDeath in package lang is deprecated:s</arg>
+            <!-- Proxy.getProxyClass is the only way to obtain a proxy Class for a pinned loader,
+                 which NativeConverters needs when deserializing expression graphs. -->
+            <arg>-Wconf:msg=method getProxyClass in class Proxy is deprecated:s</arg>
             <!-- Auron implements Spark plugin SPIs whose contracts use `private[spark]` types
                  (e.g. ElementTrackingStore, SparkUI, IndexShuffleBlockResolver,
                  ShuffleWriteMetricsReporter). Exposing them is required by Spark, not a defect. -->
@@ -1326,6 +1329,8 @@
                 <arg>-Wconf:msg=object JavaConverters in package collection is deprecated:s</arg>
                 <arg>-Wconf:msg=method newInstance in class Class is deprecated:s</arg>
                 <arg>-Wconf:msg=class ThreadDeath in package lang is deprecated:s</arg>
+                <!-- See note in base profile: Proxy.getProxyClass is required by NativeConverters. -->
+                <arg>-Wconf:msg=method getProxyClass in class Proxy is deprecated:s</arg>
                 <arg>-Wconf:cat=unchecked&amp;msg=outer reference:s</arg>
                 <arg>-Wconf:cat=unchecked&amp;msg=eliminated by erasure:s</arg>
                 <arg>-Wconf:cat=unused-nowarn:s</arg>

spark-extension-shims-spark/src/test/scala/org/apache/auron/NativeConvertersSuite.scala

@@ -16,6 +16,9 @@
  */
 package org.apache.auron
 
+import java.io.{ByteArrayOutputStream, InputStream, ObjectOutputStream}
+import java.lang.reflect.{InvocationHandler, Method, Proxy}
+
 import org.apache.spark.sql.AuronQueryTest
 import org.apache.spark.sql.auron.NativeConverters
 import org.apache.spark.sql.catalyst.expressions.{Cast, Literal}
@@ -83,4 +86,114 @@ class NativeConvertersSuite
     assert(!childExpr.hasScalarFunction)
     assert(childExpr.hasLiteral)
   }
+
+  /** Writes the expression/payload pair that deserializeExpression expects to read back. */
+  private def serializeObjects(expr: AnyRef, payload: AnyRef): Array[Byte] = {
+    val bos = new ByteArrayOutputStream()
+    val oos = new ObjectOutputStream(bos)
+    try {
+      oos.writeObject(expr)
+      oos.writeObject(payload)
+    } finally {
+      oos.close()
+    }
+    bos.toByteArray
+  }
+
+  test("deserializeExpression resolves proxy interfaces with the pinned class loader") {
+    val appLoader = getClass.getClassLoader
+    val interfaceName = classOf[ProxiedPayload].getName
+    // redefines only ProxiedPayload, so a proxy resolved through this loader carries an
+    // interface Class distinct from the one the application class loader hands out
+    val pinnedLoader = new SingleClassRedefiningLoader(interfaceName, appLoader)
+
+    val proxy: AnyRef = Proxy.newProxyInstance(
+      appLoader,
+      Array[Class[_]](classOf[ProxiedPayload]),
+      new ConstantPayloadHandler("auron"))
+    val serialized = serializeObjects(Literal(1), proxy)
+
+    val previousLoader = Thread.currentThread().getContextClassLoader
+    Thread.currentThread().setContextClassLoader(pinnedLoader)
+    try {
+      val (_, payload) =
+        NativeConverters.deserializeExpression[Literal, java.io.Serializable](serialized)
+      val resolvedInterface = payload.getClass.getInterfaces
+        .find(_.getName == interfaceName)
+        .getOrElse(fail(s"deserialized proxy does not implement $interfaceName"))
+      assert(
+        resolvedInterface.getClassLoader eq pinnedLoader,
+        "proxy interface was resolved by a class loader taken from the call stack " +
+          s"(${resolvedInterface.getClassLoader}) instead of the pinned loader")
+    } finally {
+      Thread.currentThread().setContextClassLoader(previousLoader)
+    }
+  }
+
+  test("deserializeExpression resolves primitive type descriptors") {
+    // no class loader resolves "int" by name, so this only works via the default fallback
+    val serialized = serializeObjects(Literal(1), java.lang.Integer.TYPE)
+    val (_, payload) = NativeConverters.deserializeExpression[Literal, Class[_]](serialized)
+    assert(payload eq java.lang.Integer.TYPE)
+  }
+}
+
+/** Serializable interface backing the dynamic proxy used by [[NativeConvertersSuite]]. */
+trait ProxiedPayload extends java.io.Serializable {
+  def payload(): String
+}
+
+/** Serializable [[InvocationHandler]] so the proxy itself can be written to a stream. */
+class ConstantPayloadHandler(value: String) extends InvocationHandler with java.io.Serializable {
+  override def invoke(proxy: AnyRef, method: Method, args: Array[AnyRef]): AnyRef =
+    method.getName match {
+      case "payload" => value
+      case "toString" => s"ProxiedPayload($value)"
+      case "hashCode" => Integer.valueOf(value.hashCode)
+      case "equals" => java.lang.Boolean.valueOf(proxy eq args(0))
+      case other => throw new UnsupportedOperationException(other)
+    }
+}
+
+/**
+ * Defines one named class from its own bytes and delegates every other name to the parent, so
+ * that class resolves to a Class object distinct from the parent's copy.
+ */
+class SingleClassRedefiningLoader(targetName: String, parent: ClassLoader)
+    extends ClassLoader(parent) {
+
+  override def loadClass(name: String, resolve: Boolean): Class[_] = synchronized {
+    if (name != targetName) {
+      return super.loadClass(name, resolve)
+    }
+    val loaded = findLoadedClass(name)
+    val cls = if (loaded != null) loaded else defineFromParent(name)
+    if (resolve) {
+      resolveClass(cls)
+    }
+    cls
+  }
+
+  private def defineFromParent(name: String): Class[_] = {
+    val resource = name.replace('.', '/') + ".class"
+    val in = parent.getResourceAsStream(resource)
+    require(in != null, s"cannot read class bytes of $name")
+    try {
+      val bytes = readFully(in)
+      defineClass(name, bytes, 0, bytes.length)
+    } finally {
+      in.close()
+    }
+  }
+
+  private def readFully(in: InputStream): Array[Byte] = {
+    val out = new ByteArrayOutputStream()
+    val buf = new Array[Byte](8192)
+    var read = in.read(buf)
+    while (read >= 0) {
+      out.write(buf, 0, read)
+      read = in.read(buf)
+    }
+    out.toByteArray
+  }
 }