diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometPlanStabilitySuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometPlanStabilitySuite.scala index 5117859448..0b4ec12809 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometPlanStabilitySuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometPlanStabilitySuite.scala @@ -89,79 +89,64 @@ trait CometPlanStabilitySuite extends DisableAdaptiveExecutionSuite with TPCDSBa new File(goldenFilePath, goldenFileName) } - private def isApproved( - dir: File, - actualSimplifiedPlan: String, - actualExplain: String): Boolean = { - val simplifiedFile = new File(dir, "simplified.txt") - val expectedSimplified = FileUtils.readFileToString(simplifiedFile, StandardCharsets.UTF_8) - val explainFile = new File(dir, "explain.txt") - val expectedExplain = FileUtils.readFileToString(explainFile, StandardCharsets.UTF_8) - expectedSimplified == actualSimplifiedPlan && expectedExplain == actualExplain - } - /** * Serialize and save this SparkPlan. The resulting file is used by [[checkWithApproved]] to * check stability. * - * @param plan - * the SparkPlan - * @param name - * the name of the query + * @param dir + * the directory to write to + * @param simplified + * the simplified plan * @param explain * the full explain output; this is saved to help debug later as the simplified plan is not * too useful for debugging */ - private def generateGoldenFile(plan: SparkPlan, name: String, explain: String): Unit = { - val dir = getDirForTest(name) - val simplified = getSimplifiedPlan(plan) - val foundMatch = dir.exists() && isApproved(dir, simplified, explain) - - if (!foundMatch) { - FileUtils.deleteDirectory(dir) - if (!dir.mkdirs()) { - fail(s"Could not create dir: $dir") - } - - val file = new File(dir, "simplified.txt") - FileUtils.writeStringToFile(file, simplified, StandardCharsets.UTF_8) - val fileOriginalPlan = new File(dir, "explain.txt") - FileUtils.writeStringToFile(fileOriginalPlan, explain, StandardCharsets.UTF_8) - logDebug(s"APPROVED: $file $fileOriginalPlan") + private def generateGoldenFile(dir: File, simplified: String, explain: String): Unit = { + FileUtils.deleteDirectory(dir) + if (!dir.mkdirs()) { + fail(s"Could not create dir: $dir") } + + val file = new File(dir, "simplified.txt") + FileUtils.writeStringToFile(file, simplified, StandardCharsets.UTF_8) + val fileOriginalPlan = new File(dir, "explain.txt") + FileUtils.writeStringToFile(fileOriginalPlan, explain, StandardCharsets.UTF_8) + logDebug(s"APPROVED: $file $fileOriginalPlan") } - private def checkWithApproved(plan: SparkPlan, name: String, explain: String): Unit = { - val dir = getDirForTest(name) - val tempDir = FileUtils.getTempDirectory - val actualSimplified = getSimplifiedPlan(plan) - val foundMatch = isApproved(dir, actualSimplified, explain) + private def checkWithApproved( + dir: File, + name: String, + simplified: String, + explain: String): Unit = { - if (!foundMatch) { - // show diff with last approved - val approvedSimplifiedFile = new File(dir, "simplified.txt") - val approvedExplainFile = new File(dir, "explain.txt") + val approvedSimplifiedFile = new File(dir, "simplified.txt") + val approvedExplainFile = new File(dir, "explain.txt") - val actualSimplifiedFile = new File(tempDir, s"$name.actual.simplified.txt") - val actualExplainFile = new File(tempDir, s"$name.actual.explain.txt") + // write actual files out for debugging + val tempDir = FileUtils.getTempDirectory + val actualSimplifiedFile = new File(tempDir, s"$name.actual.simplified.txt") + val actualExplainFile = new File(tempDir, s"$name.actual.explain.txt") + FileUtils.writeStringToFile(actualSimplifiedFile, simplified, StandardCharsets.UTF_8) + FileUtils.writeStringToFile(actualExplainFile, explain, StandardCharsets.UTF_8) - val approvedSimplified = - FileUtils.readFileToString(approvedSimplifiedFile, StandardCharsets.UTF_8) - // write out for debugging - FileUtils.writeStringToFile(actualSimplifiedFile, actualSimplified, StandardCharsets.UTF_8) - FileUtils.writeStringToFile(actualExplainFile, explain, StandardCharsets.UTF_8) + comparePlans("simplified", approvedSimplifiedFile, actualSimplifiedFile) + comparePlans("explain", approvedExplainFile, actualExplainFile) + } + private def comparePlans(planType: String, expectedFile: File, actualFile: File): Unit = { + val expected = FileUtils.readFileToString(expectedFile, StandardCharsets.UTF_8) + val actual = FileUtils.readFileToString(actualFile, StandardCharsets.UTF_8) + if (expected != actual) { fail(s""" - |Plans did not match: - |last approved simplified plan: ${approvedSimplifiedFile.getAbsolutePath} - |last approved explain plan: ${approvedExplainFile.getAbsolutePath} - | - |$approvedSimplified - | - |actual simplified plan: ${actualSimplifiedFile.getAbsolutePath} - |actual explain plan: ${actualExplainFile.getAbsolutePath} - | - |$actualSimplified + |Plans did not match: + |last approved $planType plan: ${expectedFile.getAbsolutePath} + | + |$expected + | + |actual $planType plan: ${actualFile.getAbsolutePath} + | + |$actual """.stripMargin) } } @@ -274,13 +259,16 @@ trait CometPlanStabilitySuite extends DisableAdaptiveExecutionSuite with TPCDSBa val qe = sql(queryString).queryExecution val plan = qe.executedPlan val explain = normalizeLocation(normalizeIds(qe.explainString(FormattedMode))) - + val simplified = getSimplifiedPlan(plan) assert(ValidateRequirements.validate(plan)) + val name = query + suffix + val dir = getDirForTest(name) + if (regenerateGoldenFiles) { - generateGoldenFile(plan, query + suffix, explain) + generateGoldenFile(dir, simplified, explain) } else { - checkWithApproved(plan, query + suffix, explain) + checkWithApproved(dir, name, simplified, explain) } } }