Github user kiszk commented on a diff in the pull request:
https://github.com/apache/spark/pull/19971#discussion_r156830943
--- Diff:
sql/core/src/main/scala/org/apache/spark/sql/execution/debug/package.scala ---
@@ -84,10 +79,31 @@ package object debug {
}
codegenSubtrees.toSeq.map { subtree =>
val (_, source) = subtree.doCodeGen()
- (subtree.toString, CodeFormatter.format(source))
+ (subtree, source)
}
}
+ /**
+ * Get WholeStageCodegenExec subtrees and the codegen in a query plan
+ *
+ * @param plan the query plan for codegen
+ * @return Sequence of WholeStageCodegen subtrees and corresponding
codegen
+ */
+ def codegenStringSeq(plan: SparkPlan): Seq[(String, String)] = {
+ codegenSubtreeSourceSeq(plan).map(s => (s._1.toString,
CodeFormatter.format(s._2)))
+ }
+
+
+ /**
+ * Get WholeStageCodegenExec subtrees' CodeAndComment in a query plan
+ *
+ * @param plan the query plan for CodeAndComment
+ * @return Sequence of WholeStageCodegen subtrees' `CodeAndComment`
+ */
+ def codegenCodeAndCommentSeq(plan: SparkPlan): Seq[CodeAndComment] = {
+ codegenSubtreeSourceSeq(plan).map(_._2)
--- End diff --
Thanks, we can do. I wanted to reduce duplicated code.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
