Skip to content

Add SparkPlanInfo constructor compatible with Databricks' Spark fork #9888

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
charlesmyu merged 2 commits into from
Nov 6, 2025
Merged

Add SparkPlanInfo constructor compatible with Databricks' Spark fork #9888

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e82dd20
Use SparkPlanInfo constructor that is compatible with Databricks' Spa…
charlesmyu Oct 30, 2025
0bfaa06
Centralize logic in a common AbstractSparkPlanUtils class
charlesmyu Nov 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Centralize logic in a common AbstractSparkPlanUtils class
  • Loading branch information
@charlesmyu
charlesmyu committed Nov 6, 2025
commit 0bfaa064cdc78d1d0b372f6c94246d83efd8fc26
4 changes: 3 additions & 1 deletion ...spark_2.12/src/main/java/datadog/trace/instrumentation/spark/Spark212Instrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ public String[] helperClassNames() {
return new String[] {
packageName + ".AbstractDatadogSparkListener",
packageName + ".AbstractSparkPlanSerializer",
packageName + ".AbstractSparkPlanUtils",
packageName + ".DatabricksParentContext",
packageName + ".OpenlineageParentContext",
packageName + ".DatadogSpark212Listener",
Expand Down Expand Up @@ -109,7 +110,8 @@ public static void exit(
JavaConverters.mapAsScalaMap(planSerializer.extractFormattedProduct(plan))
.toMap(Predef.$conforms());

SparkPlanInfo newPlanInfo = Spark212PlanUtils.upsertSparkPlanInfoMetadata(planInfo, meta);
SparkPlanInfo newPlanInfo =
new Spark212PlanUtils().upsertSparkPlanInfoMetadata(planInfo, meta);
if (newPlanInfo != null) {
planInfo = newPlanInfo;
}
Expand Down
62 changes: 15 additions & 47 deletions ...spark/spark_2.12/src/main/java/datadog/trace/instrumentation/spark/Spark212PlanUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,20 +1,11 @@
package datadog.trace.instrumentation.spark;

import datadog.trace.util.MethodHandles;
import java.lang.invoke.MethodHandle;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import org.apache.spark.sql.execution.SparkPlanInfo;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Option;
import scala.collection.immutable.Map;

public class Spark212PlanUtils {
private static final Logger log = LoggerFactory.getLogger(Spark212PlanUtils.class);

private static final MethodHandles methodLoader =
new MethodHandles(ClassLoader.getSystemClassLoader());
public class Spark212PlanUtils extends AbstractSparkPlanUtils {
private static final MethodHandle constructor =
methodLoader.constructor(
SparkPlanInfo.class,
Expand All @@ -35,43 +26,20 @@ public class Spark212PlanUtils {
String.class,
Option.class);

public static SparkPlanInfo upsertSparkPlanInfoMetadata(
SparkPlanInfo planInfo, scala.collection.immutable.Map<String, String> meta) {
// Attempt to create a new SparkPlanInfo with additional metadata replaced
// Since the fields are immutable we must instantiate a new SparkPlanInfo to do this

Object[] standardArgs =
new Object[] {
planInfo.nodeName(),
planInfo.simpleString(),
planInfo.children(),
meta,
planInfo.metrics()
};

if (databricksConstructor != null) {
List<Object> databricksArgs = new ArrayList<>(Arrays.asList(standardArgs));
try {
databricksArgs.add(SparkPlanInfo.class.getMethod("estRowCount").invoke(planInfo));
databricksArgs.add(SparkPlanInfo.class.getMethod("rddScopeId").invoke(planInfo));
databricksArgs.add(SparkPlanInfo.class.getMethod("explainId").invoke(planInfo));
} catch (Throwable t) {
log.warn("Error obtaining Databricks-specific SparkPlanInfo args", t);
}

SparkPlanInfo newPlan = methodLoader.invoke(databricksConstructor, databricksArgs.toArray());
if (newPlan != null) {
return newPlan;
}
}
@Override
protected MethodHandle getConstructor() {
return constructor;
}

if (constructor != null) {
SparkPlanInfo newPlan = methodLoader.invoke(constructor, standardArgs);
if (newPlan != null) {
return newPlan;
}
}
@Override
protected MethodHandle getDatabricksConstructor() {
return databricksConstructor;
}

return null;
@Override
protected Object[] getStandardArgs(SparkPlanInfo planInfo, Map meta) {
return new Object[] {
planInfo.nodeName(), planInfo.simpleString(), planInfo.children(), meta, planInfo.metrics()
};
}
}
4 changes: 3 additions & 1 deletion ...spark_2.13/src/main/java/datadog/trace/instrumentation/spark/Spark213Instrumentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ public String[] helperClassNames() {
return new String[] {
packageName + ".AbstractDatadogSparkListener",
packageName + ".AbstractSparkPlanSerializer",
packageName + ".AbstractSparkPlanUtils",
packageName + ".DatabricksParentContext",
packageName + ".OpenlineageParentContext",
packageName + ".DatadogSpark213Listener",
Expand Down Expand Up @@ -109,7 +110,8 @@ public static void exit(
Map<String, String> meta =
HashMap.from(JavaConverters.asScala(planSerializer.extractFormattedProduct(plan)));

SparkPlanInfo newPlanInfo = Spark213PlanUtils.upsertSparkPlanInfoMetadata(planInfo, meta);
SparkPlanInfo newPlanInfo =
new Spark213PlanUtils().upsertSparkPlanInfoMetadata(planInfo, meta);
if (newPlanInfo != null) {
planInfo = newPlanInfo;
}
Expand Down
55 changes: 15 additions & 40 deletions ...spark/spark_2.13/src/main/java/datadog/trace/instrumentation/spark/Spark213PlanUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,20 +1,11 @@
package datadog.trace.instrumentation.spark;

import datadog.trace.util.MethodHandles;
import java.lang.invoke.MethodHandle;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import org.apache.spark.sql.execution.SparkPlanInfo;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Option;
import scala.collection.immutable.Map;

public class Spark213PlanUtils {
private static final Logger log = LoggerFactory.getLogger(Spark213PlanUtils.class);

private static final MethodHandles methodLoader =
new MethodHandles(ClassLoader.getSystemClassLoader());
public class Spark213PlanUtils extends AbstractSparkPlanUtils {
private static final MethodHandle constructor =
methodLoader.constructor(
SparkPlanInfo.class,
Expand All @@ -35,36 +26,20 @@ public class Spark213PlanUtils {
String.class,
Option.class);

public static SparkPlanInfo upsertSparkPlanInfoMetadata(
SparkPlanInfo planInfo, scala.collection.immutable.Map<String, String> meta) {
// Attempt to create a new SparkPlanInfo with additional metadata replaced
// Since the fields are immutable we must instantiate a new SparkPlanInfo to do this

Object[] standardArgs =
new Object[] {
planInfo.nodeName(),
planInfo.simpleString(),
planInfo.children(),
meta,
planInfo.metrics()
};

if (databricksConstructor != null) {
List<Object> databricksArgs = new ArrayList<>(Arrays.asList(standardArgs));
try {
databricksArgs.add(SparkPlanInfo.class.getMethod("estRowCount").invoke(planInfo));
databricksArgs.add(SparkPlanInfo.class.getMethod("rddScopeId").invoke(planInfo));
databricksArgs.add(SparkPlanInfo.class.getMethod("explainId").invoke(planInfo));
} catch (Throwable t) {
log.warn("Error obtaining Databricks-specific SparkPlanInfo args", t);
}
@Override
protected MethodHandle getConstructor() {
return constructor;
}

SparkPlanInfo newPlan = methodLoader.invoke(databricksConstructor, databricksArgs.toArray());
if (newPlan != null) {
return newPlan;
}
}
@Override
protected MethodHandle getDatabricksConstructor() {
return databricksConstructor;
}

return null;
@Override
protected Object[] getStandardArgs(SparkPlanInfo planInfo, Map meta) {
return new Object[] {
planInfo.nodeName(), planInfo.simpleString(), planInfo.children(), meta, planInfo.metrics()
};
}
}
57 changes: 57 additions & 0 deletions ...ation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractSparkPlanUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
package datadog.trace.instrumentation.spark;

import datadog.trace.util.MethodHandles;
import java.lang.invoke.MethodHandle;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import org.apache.spark.sql.execution.SparkPlanInfo;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.collection.immutable.Map;

abstract class AbstractSparkPlanUtils {
private static final Logger log = LoggerFactory.getLogger(AbstractSparkPlanUtils.class);

protected static final MethodHandles methodLoader =
new MethodHandles(ClassLoader.getSystemClassLoader());

protected abstract MethodHandle getConstructor();

protected abstract MethodHandle getDatabricksConstructor();

// Deals with Seq which changed from Scala 2.12 to 2.13, so delegate to version-specific classes
protected abstract Object[] getStandardArgs(SparkPlanInfo planInfo, Map meta);

// Attempt to create a new SparkPlanInfo with additional metadata replaced
// Since the fields are immutable we must instantiate a new SparkPlanInfo to do this
public SparkPlanInfo upsertSparkPlanInfoMetadata(
SparkPlanInfo planInfo, scala.collection.immutable.Map<String, String> meta) {
if (getDatabricksConstructor() != null) {
List<Object> databricksArgs = new ArrayList<>(Arrays.asList(getStandardArgs(planInfo, meta)));
try {
databricksArgs.add(SparkPlanInfo.class.getMethod("estRowCount").invoke(planInfo));
databricksArgs.add(SparkPlanInfo.class.getMethod("rddScopeId").invoke(planInfo));
databricksArgs.add(SparkPlanInfo.class.getMethod("explainId").invoke(planInfo));
} catch (Throwable t) {
log.warn("Error obtaining Databricks-specific SparkPlanInfo args", t);
}

SparkPlanInfo newPlan =
methodLoader.invoke(getDatabricksConstructor(), databricksArgs.toArray());
if (newPlan != null) {
return newPlan;
}
}

if (getConstructor() != null) {
SparkPlanInfo newPlan =
methodLoader.invoke(getConstructor(), getStandardArgs(planInfo, meta));
if (newPlan != null) {
return newPlan;
}
}

return null;
}
}
Loading