Extract SparkPlan product and append to trace #9783
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,9 @@ | ||
| package datadog.trace.instrumentation.spark; | ||
|
|
||
| import org.apache.spark.sql.catalyst.trees.TreeNode; | ||
|
|
||
| public class Spark212PlanUtils extends AbstractSparkPlanUtils { | ||
| public String getKey(int idx, TreeNode node) { | ||
| return String.format("_dd.unknown_key.%d", idx); | ||
| } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,9 @@ | ||
| package datadog.trace.instrumentation.spark; | ||
|
|
||
| import org.apache.spark.sql.catalyst.trees.TreeNode; | ||
|
|
||
| public class Spark213PlanUtils extends AbstractSparkPlanUtils { | ||
| public String getKey(int idx, TreeNode node) { | ||
| return node.productElementName(idx); | ||
| } | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,171 @@ | ||
| package datadog.trace.instrumentation.spark; | ||
|
|
||
| import com.fasterxml.jackson.databind.DeserializationFeature; | ||
| import com.fasterxml.jackson.databind.ObjectMapper; | ||
| import java.io.IOException; | ||
| import java.lang.reflect.InvocationTargetException; | ||
| import java.util.ArrayList; | ||
| import java.util.HashMap; | ||
| import java.util.Iterator; | ||
| import java.util.Map; | ||
| import org.apache.spark.Partitioner; | ||
| import org.apache.spark.sql.catalyst.expressions.Attribute; | ||
| import org.apache.spark.sql.catalyst.plans.JoinType; | ||
| import org.apache.spark.sql.catalyst.plans.QueryPlan; | ||
| import org.apache.spark.sql.catalyst.plans.physical.BroadcastMode; | ||
| import org.apache.spark.sql.catalyst.plans.physical.Partitioning; | ||
| import org.apache.spark.sql.catalyst.trees.TreeNode; | ||
| import scala.Option; | ||
| import scala.collection.Iterable; | ||
| import scala.collection.JavaConverters; | ||
|
|
||
| // An extension of how Spark translates `SparkPlan`s to `SparkPlanInfo`, see here: | ||
| // https://github.com/apache/spark/blob/v3.5.0/sql/core/src/main/scala/org/apache/spark/sql/execution/SparkPlanInfo.scala#L54 | ||
| public abstract class AbstractSparkPlanUtils { | ||
| private final int MAX_DEPTH = 4; | ||
| private final int MAX_LENGTH = 50; | ||
| private final ObjectMapper mapper = | ||
pawel-big-lebowski marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| new ObjectMapper().configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false); | ||
|
|
||
| private final Class[] SAFE_CLASSES = { | ||
| Attribute.class, // simpleString appends data type, avoid by using toString | ||
| JoinType.class, // enum | ||
| Partitioner.class, // not a product or TreeNode | ||
| BroadcastMode.class, // not a product or TreeNode | ||
| maybeGetClass("org.apache.spark.sql.execution.exchange.ShuffleOrigin"), // enum (v3+) | ||
| maybeGetClass("org.apache.spark.sql.catalyst.optimizer.BuildSide"), // enum (v3+) | ||
| maybeGetClass( | ||
| "org.apache.spark.sql.execution.ShufflePartitionSpec"), // not a product or TreeNode (v3+) | ||
| }; | ||
|
|
||
| public abstract String getKey(int idx, TreeNode node); | ||
|
|
||
| public Map<String, String> extractFormattedProduct(TreeNode plan) { | ||
| HashMap<String, String> result = new HashMap<>(); | ||
| extractPlanProduct(plan, 0) | ||
| .forEach( | ||
| (key, value) -> { | ||
| result.put(key, writeObjectToString(value)); | ||
| }); | ||
| return result; | ||
| } | ||
|
|
||
| protected Map<String, Object> extractPlanProduct(TreeNode node, int depth) { | ||
| HashMap<String, Object> args = new HashMap<>(); | ||
| HashMap<String, String> unparsed = new HashMap<>(); | ||
|
|
||
| int i = 0; | ||
| for (Iterator<Object> it = JavaConverters.asJavaIterator(node.productIterator()); | ||
| it.hasNext(); ) { | ||
| Object obj = it.next(); | ||
|
|
||
| Object val = parsePlanProduct(obj, depth); | ||
| if (val != null) { | ||
| args.put(getKey(i, node), val); | ||
| } else { | ||
| unparsed.put(getKey(i, node), obj.getClass().getName()); | ||
| } | ||
|
|
||
| i++; | ||
| } | ||
|
|
||
| if (unparsed.size() > 0) { | ||
| // For now, place what we can't parse here with the types so we're aware of them | ||
| args.put("_dd.unparsed", unparsed); | ||
| } | ||
| return args; | ||
| } | ||
|
|
||
| // Should only call on final values being written to `meta` | ||
| protected String writeObjectToString(Object value) { | ||
| try { | ||
| return mapper.writeValueAsString(value); | ||
| } catch (IOException e) { | ||
| return null; | ||
| } | ||
| } | ||
|
|
||
| protected Object parsePlanProduct(Object value, int depth) { | ||
| // This function MUST not arbitrarily serialize the object as we can't be sure what it is. | ||
| // A null return indicates object is unserializable, otherwise it should really only return | ||
| // valid | ||
| // JSON types (Array, Map, String, Boolean, Number, null) | ||
|
|
||
| if (value == null) { | ||
| return "null"; | ||
| } else if (value instanceof String | ||
| || value instanceof Boolean | ||
| || Number.class.isInstance(value)) { | ||
| return value; | ||
| } else if (value instanceof Option) { | ||
| return parsePlanProduct(((Option) value).getOrElse(() -> "none"), depth); | ||
| } else if (value instanceof QueryPlan) { | ||
| // don't duplicate child nodes | ||
| return null; | ||
| } else if (value instanceof Iterable && depth < MAX_DEPTH) { | ||
| ArrayList<Object> list = new ArrayList<>(); | ||
| for (Object item : JavaConverters.asJavaIterable((Iterable) value)) { | ||
| Object res = parsePlanProduct(item, depth + 1); | ||
| if (list.size() < MAX_LENGTH && res != null) { | ||
| list.add(res); | ||
| } | ||
| } | ||
| return list; | ||
| } else if (value instanceof Partitioning) { | ||
| if (value instanceof TreeNode && depth + 1 < MAX_DEPTH) { | ||
| HashMap<String, Object> inner = new HashMap<>(); | ||
| inner.put( | ||
| value.getClass().getSimpleName(), extractPlanProduct(((TreeNode) value), depth + 2)); | ||
| return inner; | ||
| } else { | ||
| return value.toString(); | ||
| } | ||
| } else if (instanceOf(value, SAFE_CLASSES)) { | ||
| return value.toString(); | ||
| } else if (value instanceof TreeNode) { | ||
| // fallback case, leave at bottom | ||
| return getSimpleString((TreeNode) value); | ||
| } | ||
|
|
||
| return null; | ||
| } | ||
|
|
||
| private String getSimpleString(TreeNode value) { | ||
| try { | ||
| // in Spark v3+, the signature of `simpleString` includes an int parameter for `maxFields` | ||
| return TreeNode.class | ||
| .getDeclaredMethod("simpleString", new Class[] {int.class}) | ||
| .invoke(value, MAX_LENGTH) | ||
| .toString(); | ||
| } catch (NoSuchMethodException | IllegalAccessException | InvocationTargetException exception) { | ||
| try { | ||
| // Attempt the Spark v2 `simpleString` signature | ||
| return TreeNode.class.getDeclaredMethod("simpleString").invoke(value).toString(); | ||
| } catch (NoSuchMethodException | ||
| | IllegalAccessException | ||
| | InvocationTargetException innerException) { | ||
| } | ||
|
|
||
| return null; | ||
| } | ||
| } | ||
|
|
||
| // Use reflection rather than native `instanceof` for classes added in later Spark versions | ||
| private boolean instanceOf(Object value, Class[] classes) { | ||
| for (Class cls : classes) { | ||
| if (cls != null && cls.isInstance(value)) { | ||
| return true; | ||
| } | ||
| } | ||
|
|
||
| return false; | ||
| } | ||
|
|
||
| private Class maybeGetClass(String cls) { | ||
| try { | ||
| return Class.forName(cls); | ||
| } catch (ClassNotFoundException e) { | ||
| return null; | ||
| } | ||
| } | ||
| } | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.