Extract SparkPlan product and append to trace

dd-java-agent/instrumentation/spark/spark_2.12/src/main/java/datadog/trace/instrumentation/spark/Spark212Instrumentation.java

@@ -8,15 +8,20 @@
 import datadog.trace.api.Config;
 import net.bytebuddy.asm.Advice;
 import org.apache.spark.SparkContext;
+import org.apache.spark.sql.execution.SparkPlan;
+import org.apache.spark.sql.execution.SparkPlanInfo;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import scala.Predef;
+import scala.collection.JavaConverters;
 
 @AutoService(InstrumenterModule.class)
 public class Spark212Instrumentation extends AbstractSparkInstrumentation {
   @Override
   public String[] helperClassNames() {
     return new String[] {
       packageName + ".AbstractDatadogSparkListener",
+      packageName + ".AbstractSparkPlanSerializer",
       packageName + ".DatabricksParentContext",
       packageName + ".OpenlineageParentContext",
       packageName + ".DatadogSpark212Listener",
@@ -27,6 +32,7 @@ public String[] helperClassNames() {
       packageName + ".SparkSQLUtils",
       packageName + ".SparkSQLUtils$SparkPlanInfoForStage",
       packageName + ".SparkSQLUtils$AccumulatorWithStage",
+      packageName + ".Spark212PlanSerializer"
     };
   }
 
@@ -40,6 +46,13 @@ public void methodAdvice(MethodTransformer transformer) {
             .and(isDeclaredBy(named("org.apache.spark.SparkContext")))
             .and(takesNoArguments()),
         Spark212Instrumentation.class.getName() + "$InjectListener");
+
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("fromSparkPlan"))
+            .and(takesArgument(0, named("org.apache.spark.sql.execution.SparkPlan")))
+            .and(isDeclaredBy(named("org.apache.spark.sql.execution.SparkPlanInfo$"))),
+        Spark212Instrumentation.class.getName() + "$SparkPlanInfoAdvice");
   }
 
   public static class InjectListener {
@@ -78,4 +91,25 @@ public static void enter(@Advice.This SparkContext sparkContext) {
       sparkContext.listenerBus().addToSharedQueue(AbstractDatadogSparkListener.listener);
     }
   }
+
+  public static class SparkPlanInfoAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void exit(
+        @Advice.Return(readOnly = false) SparkPlanInfo planInfo,
+        @Advice.Argument(0) SparkPlan plan) {
+      if (planInfo.metadata().size() == 0
+          && (Config.get().isDataJobsParseSparkPlanEnabled()
+              || Config.get().isDataJobsExperimentalFeaturesEnabled())) {
+        Spark212PlanSerializer planUtils = new Spark212PlanSerializer();
+        planInfo =
+            new SparkPlanInfo(
+                planInfo.nodeName(),
+                planInfo.simpleString(),
+                planInfo.children(),
+                JavaConverters.mapAsScalaMap(planUtils.extractFormattedProduct(plan))
+                    .toMap(Predef.$conforms()),
+                planInfo.metrics());
+      }
+    }
+  }
 }

dd-java-agent/instrumentation/spark/spark_2.12/src/main/java/datadog/trace/instrumentation/spark/Spark212PlanSerializer.java

@@ -0,0 +1,10 @@
+package datadog.trace.instrumentation.spark;
+
+import org.apache.spark.sql.catalyst.trees.TreeNode;
+
+public class Spark212PlanSerializer extends AbstractSparkPlanSerializer {
+  @Override
+  public String getKey(int idx, TreeNode node) {
+    return "_dd.unknown_key." + idx;
+  }
+}

dd-java-agent/instrumentation/spark/spark_2.13/src/main/java/datadog/trace/instrumentation/spark/Spark213Instrumentation.java

@@ -8,15 +8,20 @@
 import datadog.trace.api.Config;
 import net.bytebuddy.asm.Advice;
 import org.apache.spark.SparkContext;
+import org.apache.spark.sql.execution.SparkPlan;
+import org.apache.spark.sql.execution.SparkPlanInfo;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import scala.collection.JavaConverters;
+import scala.collection.immutable.HashMap;
 
 @AutoService(InstrumenterModule.class)
 public class Spark213Instrumentation extends AbstractSparkInstrumentation {
   @Override
   public String[] helperClassNames() {
     return new String[] {
       packageName + ".AbstractDatadogSparkListener",
+      packageName + ".AbstractSparkPlanSerializer",
       packageName + ".DatabricksParentContext",
       packageName + ".OpenlineageParentContext",
       packageName + ".DatadogSpark213Listener",
@@ -27,6 +32,7 @@ public String[] helperClassNames() {
       packageName + ".SparkSQLUtils",
       packageName + ".SparkSQLUtils$SparkPlanInfoForStage",
       packageName + ".SparkSQLUtils$AccumulatorWithStage",
+      packageName + ".Spark213PlanSerializer"
     };
   }
 
@@ -40,6 +46,13 @@ public void methodAdvice(MethodTransformer transformer) {
             .and(isDeclaredBy(named("org.apache.spark.SparkContext")))
             .and(takesNoArguments()),
         Spark213Instrumentation.class.getName() + "$InjectListener");
+
+    transformer.applyAdvice(
+        isMethod()
+            .and(named("fromSparkPlan"))
+            .and(takesArgument(0, named("org.apache.spark.sql.execution.SparkPlan")))
+            .and(isDeclaredBy(named("org.apache.spark.sql.execution.SparkPlanInfo$"))),
+        Spark213Instrumentation.class.getName() + "$SparkPlanInfoAdvice");
   }
 
   public static class InjectListener {
@@ -79,4 +92,24 @@ public static void enter(@Advice.This SparkContext sparkContext) {
       sparkContext.listenerBus().addToSharedQueue(AbstractDatadogSparkListener.listener);
     }
   }
+
+  public static class SparkPlanInfoAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    public static void exit(
+        @Advice.Return(readOnly = false) SparkPlanInfo planInfo,
+        @Advice.Argument(0) SparkPlan plan) {
+      if (planInfo.metadata().size() == 0
+          && (Config.get().isDataJobsParseSparkPlanEnabled()
+              || Config.get().isDataJobsExperimentalFeaturesEnabled())) {
+        Spark213PlanSerializer planUtils = new Spark213PlanSerializer();
+        planInfo =
+            new SparkPlanInfo(
+                planInfo.nodeName(),
+                planInfo.simpleString(),
+                planInfo.children(),
+                HashMap.from(JavaConverters.asScala(planUtils.extractFormattedProduct(plan))),
+                planInfo.metrics());
+      }
+    }
+  }
 }

dd-java-agent/instrumentation/spark/spark_2.13/src/main/java/datadog/trace/instrumentation/spark/Spark213PlanSerializer.java

@@ -0,0 +1,10 @@
+package datadog.trace.instrumentation.spark;
+
+import org.apache.spark.sql.catalyst.trees.TreeNode;
+
+public class Spark213PlanSerializer extends AbstractSparkPlanSerializer {
+  @Override
+  public String getKey(int idx, TreeNode node) {
+    return node.productElementName(idx);
+  }
+}

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

@@ -66,7 +66,7 @@
  */
 public abstract class AbstractDatadogSparkListener extends SparkListener {
   private static final Logger log = LoggerFactory.getLogger(AbstractDatadogSparkListener.class);
-  private static final ObjectMapper objectMapper = new ObjectMapper();
+  protected static final ObjectMapper objectMapper = new ObjectMapper();
   public static volatile AbstractDatadogSparkListener listener = null;
 
   public static volatile boolean finishTraceOnApplicationEnd = true;

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

@@ -36,7 +36,8 @@ public String[] knownMatchingTypes() {
       "org.apache.spark.deploy.yarn.ApplicationMaster",
       "org.apache.spark.util.Utils",
       "org.apache.spark.util.SparkClassUtils",
-      "org.apache.spark.scheduler.LiveListenerBus"
+      "org.apache.spark.scheduler.LiveListenerBus",
+      "org.apache.spark.sql.execution.SparkPlanInfo$"
     };
   }