@@ -1030,28 +1030,40 @@ class SparkContext(config: SparkConf) extends Logging {
10301030 }
10311031
10321032 /**
1033- * Support function for API backtraces.
1033+ * Set the thread-local property for overriding the call sites
1034+ * of actions and RDDs.
10341035 */
1035- def setCallSite (site : String ) {
1036- setLocalProperty(" externalCallSite " , site )
1036+ def setCallSite (shortCallSite : String ) {
1037+ setLocalProperty(CallSite . SHORT_FORM , shortCallSite )
10371038 }
10381039
10391040 /**
1040- * Support function for API backtraces.
1041+ * Set the thread-local property for overriding the call sites
1042+ * of actions and RDDs.
1043+ */
1044+ private [spark] def setCallSite (callSite : CallSite ) {
1045+ setLocalProperty(CallSite .SHORT_FORM , callSite.shortForm)
1046+ setLocalProperty(CallSite .LONG_FORM , callSite.longForm)
1047+ }
1048+
1049+ /**
1050+ * Clear the thread-local property for overriding the call sites
1051+ * of actions and RDDs.
10411052 */
10421053 def clearCallSite () {
1043- setLocalProperty(" externalCallSite" , null )
1054+ setLocalProperty(CallSite .SHORT_FORM , null )
1055+ setLocalProperty(CallSite .LONG_FORM , null )
10441056 }
10451057
10461058 /**
10471059 * Capture the current user callsite and return a formatted version for printing. If the user
1048- * has overridden the call site, this will return the user's version.
1060+ * has overridden the call site using `setCallSite()` , this will return the user's version.
10491061 */
10501062 private [spark] def getCallSite (): CallSite = {
1051- Option (getLocalProperty(" externalCallSite " )) match {
1052- case Some (callSite) => CallSite (callSite, longForm = " " )
1053- case None => Utils .getCallSite
1054- }
1063+ Option (getLocalProperty(CallSite . SHORT_FORM )).map { case shortCallSite =>
1064+ val longCallSite = Option (getLocalProperty( CallSite . LONG_FORM )).getOrElse( " " )
1065+ CallSite (shortCallSite, longCallSite)
1066+ }.getOrElse( Utils .getCallSite())
10551067 }
10561068
10571069 /**
0 commit comments