Instrument QueryExecution.assertAnalyzed() to catch DataFrame analysis failures

Author: aboitreaud

dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java

@@ -644,9 +644,7 @@ public synchronized void onJobEnd(SparkListenerJobEnd jobEnd) {
       }
     } else {
       lastJobFailed = false;
-      // Note: we intentionally do NOT reset lastSqlFailed here. A successful Spark job
-      // should not erase a prior SQL failure (e.g. from cleanup operations after a failed
-      // spark.sql() call). The SQL failure is only relevant at finishApplication time.
+      lastSqlFailed = false;
     }
 
     SparkAggregatedTaskMetrics metrics = jobMetrics.remove(jobEnd.jobId());

dd-java-agent/instrumentation/spark/spark-common/src/main/java/datadog/trace/instrumentation/spark/AbstractSparkInstrumentation.java

@@ -40,7 +40,8 @@ public String[] knownMatchingTypes() {
       "org.apache.spark.util.SparkClassUtils",
       "org.apache.spark.scheduler.LiveListenerBus",
       "org.apache.spark.sql.execution.SparkPlanInfo$",
-      "org.apache.spark.sql.SparkSession"
+      "org.apache.spark.sql.SparkSession",
+      "org.apache.spark.sql.execution.QueryExecution"
     };
   }
 
@@ -77,6 +78,14 @@ public void methodAdvice(MethodTransformer transformer) {
             .and(isDeclaredBy(named("org.apache.spark.sql.SparkSession"))),
         AbstractSparkInstrumentation.class.getName() + "$SparkSqlFailureAdvice");
 
+    // QueryExecution.assertAnalyzed() — catch all Catalyst analysis failures regardless of
+    // entry point (SparkSession.sql, Dataset.select, Dataset.filter, etc.)
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("assertAnalyzed"))
+            .and(isDeclaredBy(named("org.apache.spark.sql.execution.QueryExecution"))),
+        AbstractSparkInstrumentation.class.getName() + "$QueryExecutionFailureAdvice");
+
     // LiveListenerBus class is used to manage spark listeners
     transformer.applyAdvice(
         isMethod()
@@ -89,11 +98,6 @@ public void methodAdvice(MethodTransformer transformer) {
   public static class PrepareSubmitEnvAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enter(@Advice.Argument(0) SparkSubmitArguments submitArgs) {
-      System.err.println(
-          "[DD-SPARK-DEBUG] PrepareSubmitEnvAdvice.enter: primaryResource="
-              + submitArgs.primaryResource()
-              + ", thread="
-              + Thread.currentThread().getName());
 
       // Using pyspark `python script.py`, spark JVM is launched as PythonGatewayServer, which is
       // exited using System.exit(0), leading to the exit advice not being called
@@ -104,36 +108,20 @@ public static void enter(@Advice.Argument(0) SparkSubmitArguments submitArgs) {
 
         // prepareSubmitEnvironment might be called before/after runMain depending on spark version
         AbstractDatadogSparkListener.finishTraceOnApplicationEnd = true;
-        System.err.println(
-            "[DD-SPARK-DEBUG] PrepareSubmitEnvAdvice: detected pyspark-shell, set isPysparkShell=true, finishTraceOnApplicationEnd=true");
       }
     }
   }
 
   public static class RunMainAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enter() {
-      System.err.println(
-          "[DD-SPARK-DEBUG] RunMainAdvice.enter: isPysparkShell="
-              + AbstractDatadogSparkListener.isPysparkShell
-              + ", thread="
-              + Thread.currentThread().getName());
       if (!AbstractDatadogSparkListener.isPysparkShell) {
         AbstractDatadogSparkListener.finishTraceOnApplicationEnd = false;
       }
     }
 
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
     public static void exit(@Advice.Thrown Throwable throwable) {
-      System.err.println(
-          "[DD-SPARK-DEBUG] RunMainAdvice.exit: throwable="
-              + (throwable != null
-                  ? throwable.getClass().getName() + ": " + throwable.getMessage()
-                  : "null")
-              + ", listenerNull="
-              + (AbstractDatadogSparkListener.listener == null)
-              + ", thread="
-              + Thread.currentThread().getName());
       if (AbstractDatadogSparkListener.listener != null) {
         AbstractDatadogSparkListener.listener.finishApplication(
             System.currentTimeMillis(), throwable, 0, null);
@@ -146,15 +134,6 @@ public static void exit(@Advice.Thrown Throwable throwable) {
   public static class YarnFinishAdvice {
     @Advice.OnMethodEnter(suppress = Throwable.class)
     public static void enter(@Advice.Argument(1) int exitCode, @Advice.Argument(2) String msg) {
-      System.err.println(
-          "[DD-SPARK-DEBUG] YarnFinishAdvice.enter: exitCode="
-              + exitCode
-              + ", msg="
-              + msg
-              + ", listenerNull="
-              + (AbstractDatadogSparkListener.listener == null)
-              + ", thread="
-              + Thread.currentThread().getName());
       if (AbstractDatadogSparkListener.listener != null) {
         AbstractDatadogSparkListener.listener.finishApplication(
             System.currentTimeMillis(), null, exitCode, msg);
@@ -163,19 +142,6 @@ public static void enter(@Advice.Argument(1) int exitCode, @Advice.Argument(2) S
   }
 
   public static class SparkSqlFailureAdvice {
-    @Advice.OnMethodEnter(suppress = Throwable.class)
-    public static void enter(@Advice.Argument(0) String sqlText) {
-      System.err.println(
-          "[DD-SPARK-DEBUG] SparkSqlFailureAdvice.enter: thread="
-              + Thread.currentThread().getName()
-              + ", sql="
-              + (sqlText != null && sqlText.length() > 200
-                  ? sqlText.substring(0, 200) + "..."
-                  : sqlText)
-              + ", listenerNull="
-              + (AbstractDatadogSparkListener.listener == null));
-    }
-
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
     public static void exit(@Advice.Thrown Throwable throwable) {
       if (throwable != null) {
@@ -190,15 +156,27 @@ public static void exit(@Advice.Thrown Throwable throwable) {
                 + (AbstractDatadogSparkListener.listener == null));
         if (AbstractDatadogSparkListener.listener != null) {
           AbstractDatadogSparkListener.listener.onSqlFailure(throwable);
-        } else {
-          System.err.println(
-              "[DD-SPARK-DEBUG] SparkSqlFailureAdvice.exit: listener is null, cannot record SQL failure");
         }
-      } else {
+      }
+    }
+  }
+
+  public static class QueryExecutionFailureAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void exit(@Advice.Thrown Throwable throwable) {
+      if (throwable != null) {
         System.err.println(
-            "[DD-SPARK-DEBUG] SparkSqlFailureAdvice.exit: thread="
+            "[DD-SPARK-DEBUG] QueryExecutionFailureAdvice.exit: thread="
                 + Thread.currentThread().getName()
-                + ", success (no throwable)");
+                + ", throwable="
+                + throwable.getClass().getName()
+                + ": "
+                + throwable.getMessage()
+                + ", listenerNull="
+                + (AbstractDatadogSparkListener.listener == null));
+        if (AbstractDatadogSparkListener.listener != null) {
+          AbstractDatadogSparkListener.listener.onSqlFailure(throwable);
+        }
       }
     }
   }

dd-java-agent/instrumentation/spark/spark-common/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkTest.groovy

@@ -223,6 +223,38 @@ abstract class AbstractSparkTest extends InstrumentationSpecification {
     }
   }
 
+  def "DataFrame analysis failure on unresolved column marks application span as error"() {
+    setup:
+    def sparkSession = SparkSession.builder()
+      .config("spark.master", "local[2]")
+      .getOrCreate()
+
+    try {
+      // Triggers AnalysisException via Dataset.select() -> QueryExecution.assertAnalyzed(),
+      // NOT through SparkSession.sql(). This exercises the QueryExecutionFailureAdvice.
+      sparkSession.range(1).toDF("id").select("nonexistent_column")
+    } catch (Exception ignored) {
+      // Expected: AnalysisException thrown by Catalyst analysis
+    }
+    sparkSession.stop()
+
+    expect:
+    assertTraces(1) {
+      trace(1) {
+        span {
+          operationName "spark.application"
+          resourceName "spark.application"
+          spanType "spark"
+          errored true
+          parent()
+          assert span.tags["error.type"] == "Spark SQL Failed"
+          assert span.tags["error.message"] =~ /(?i).*nonexistent_column.*/
+          assert span.tags["error.stack"] =~ /(?s).*AnalysisException.*/
+        }
+      }
+    }
+  }
+
   def "capture SparkSubmit.runMain() errors"() {
     setup:
     def sparkSession = SparkSession.builder()