diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala
index 6c661667825b2..848966a97a70e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/SQLQueryTestSuite.scala
@@ -256,20 +256,23 @@ class SQLQueryTestSuite extends QueryTest with SharedSparkSession {
     def splitWithSemicolon(seq: Seq[String]) = {
       seq.mkString("\n").split("(?<=[^\\\\]);")
     }
-    val input = fileToString(new File(testCase.inputFile))
 
-    val (comments, code) = input.split("\n").partition { line =>
+    def splitCommentsAndCodes(input: String) = input.split("\n").partition { line =>
       val newLine = line.trim
       newLine.startsWith("--") && !newLine.startsWith("--QUERY-DELIMITER")
     }
 
+    val input = fileToString(new File(testCase.inputFile))
+
+    val (comments, code) = splitCommentsAndCodes(input)
+
     // If `--IMPORT` found, load code from another test case file, then insert them
     // into the head in this test.
     val importedTestCaseName = comments.filter(_.startsWith("--IMPORT ")).map(_.substring(9))
     val importedCode = importedTestCaseName.flatMap { testCaseName =>
       listTestCases.find(_.name == testCaseName).map { testCase =>
         val input = fileToString(new File(testCase.inputFile))
-        val (_, code) = input.split("\n").partition(_.trim.startsWith("--"))
+        val (_, code) = splitCommentsAndCodes(input)
         code
       }
     }.flatten