error handling

Author: andygrove

spark/src/test/scala/org/apache/spark/sql/comet/CometPlanStabilitySuite.scala

@@ -199,29 +199,40 @@ trait CometPlanStabilitySuite extends DisableAdaptiveExecutionSuite with TPCDSBa
      * "sum(sr_return_amt#14)", so we remove all of these using regex
      */
     def cleanUpReferences(references: AttributeSet): String = {
-      referenceRegex.replaceAllIn(references.map(_.name).mkString(","), "")
+      try {
+        referenceRegex.replaceAllIn(references.map(_.name).mkString(","), "")
+      } catch {
+        case e: Exception =>
+          s"[Error processing references: ${e.getMessage}]"
+      }
     }
 
     /**
      * Generate a simplified plan as a string Example output: TakeOrderedAndProject
      * [c_customer_id] WholeStageCodegen Project [c_customer_id]
      */
     def simplifyNode(node: SparkPlan, depth: Int): String = {
-      val padding = "  " * depth
-      var thisNode = node.nodeName
-      if (node.references.nonEmpty) {
-        thisNode += s" [${cleanUpReferences(node.references)}]"
-      }
-      if (node.producedAttributes.nonEmpty) {
-        thisNode += s" [${cleanUpReferences(node.producedAttributes)}]"
-      }
-      val id = getId(node)
-      if (id > 0) {
-        thisNode += s" #$id"
+      try {
+        val padding = "  " * depth
+        var thisNode = node.nodeName
+        if (node.references.nonEmpty) {
+          thisNode += s" [${cleanUpReferences(node.references)}]"
+        }
+        if (node.producedAttributes.nonEmpty) {
+          thisNode += s" [${cleanUpReferences(node.producedAttributes)}]"
+        }
+        val id = getId(node)
+        if (id > 0) {
+          thisNode += s" #$id"
+        }
+        val childrenSimplified = node.children.map(simplifyNode(_, depth + 1))
+        val subqueriesSimplified = node.subqueries.map(simplifyNode(_, depth + 1))
+        s"$padding$thisNode\n${subqueriesSimplified.mkString("")}${childrenSimplified.mkString("")}"
+      } catch {
+        case e: Exception =>
+          val padding = "  " * depth
+          s"$padding[Error processing node ${node.getClass.getSimpleName}: ${e.getMessage}]"
       }
-      val childrenSimplified = node.children.map(simplifyNode(_, depth + 1))
-      val subqueriesSimplified = node.subqueries.map(simplifyNode(_, depth + 1))
-      s"$padding$thisNode\n${subqueriesSimplified.mkString("")}${childrenSimplified.mkString("")}"
     }
 
     simplifyNode(plan, 0)