update some tests

Signed-off-by: YANGDB <[email protected]>
opensearch-project · Sep 11, 2023 · 69df8ad · 69df8ad
1 parent 17e93fb
commit 69df8ad
Show file tree

Hide file tree

Showing 7 changed files with 19 additions and 95 deletions.
diff --git a/...k-integration/src/main/scala/org/opensearch/flint/spark/FlintGenericSparkExtensions.scala b/...k-integration/src/main/scala/org/opensearch/flint/spark/FlintGenericSparkExtensions.scala
diff --git a/...t-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSparkParserChain.scala b/...t-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSparkParserChain.scala
diff --git a/...spark-integration/src/main/scala/org/opensearch/flint/spark/sql/FlintSparkSqlParser.scala b/...spark-integration/src/main/scala/org/opensearch/flint/spark/sql/FlintSparkSqlParser.scala
@@ -53,7 +53,12 @@ class FlintSparkSqlParser(sparkParser: ParserInterface) extends ParserInterface
   private val flintAstBuilder = new FlintSparkSqlAstBuilder()
 
   override def parsePlan(sqlText: String): LogicalPlan = parse(sqlText) { flintParser =>
-    flintAstBuilder.visit(flintParser.singleStatement())
+    try {
+      flintAstBuilder.visit(flintParser.singleStatement())
+    } catch {
+      // Fall back to Spark parse plan logic if flint cannot parse
+      case _: ParseException => sparkParser.parsePlan(sqlText)
+    }
   }
 
   override def parseExpression(sqlText: String): Expression = sparkParser.parseExpression(sqlText)

diff --git a/flint-spark-integration/src/test/scala/org/apache/spark/FlintSuite.scala b/flint-spark-integration/src/test/scala/org/apache/spark/FlintSuite.scala
@@ -11,7 +11,7 @@ import org.apache.spark.sql.flint.config.FlintConfigEntry
 import org.apache.spark.sql.flint.config.FlintSparkConf.HYBRID_SCAN_ENABLED
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
-import org.opensearch.flint.spark.FlintGenericSparkExtensions
+import org.opensearch.flint.spark.FlintSparkExtensions
 
 trait FlintSuite extends SharedSparkSession {
   override protected def sparkConf = {
@@ -24,7 +24,7 @@ trait FlintSuite extends SharedSparkSession {
       // this rule may potentially block testing of other optimization rules such as
       // ConstantPropagation etc.
       .set(SQLConf.OPTIMIZER_EXCLUDED_RULES.key, ConvertToLocalRelation.ruleName)
-      .set("spark.sql.extensions", classOf[FlintGenericSparkExtensions].getName)
+      .set("spark.sql.extensions", classOf[FlintSparkExtensions].getName)
     conf
   }
 

diff --git a/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkPPLITSuite.scala b/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkPPLITSuite.scala
@@ -15,7 +15,7 @@ import org.apache.spark.sql.{QueryTest, Row}
 class FlintSparkPPLITSuite
   extends QueryTest
     with LogicalPlanTestUtils
-    with FlintSuite
+    with FlintPPLSuite
     with StreamTest {
 
   /** Test table and index name */
@@ -309,16 +309,11 @@ class FlintSparkPPLITSuite
   }
 
   test("create ppl simple age avg group by country query test ") {
-    val checkData = sql(s"SELECT country, AVG(age) AS avg_age FROM $testTable group by country");
-    checkData.show()
-    checkData.queryExecution.logical.show()
-
     val frame = sql(
       s"""
          | source = $testTable| stats avg(age) by country
          | """.stripMargin)
 
-
     // Retrieve the results
     val results: Array[Row] = frame.collect()
     // Define the expected results

diff --git a/...spark-integration/src/main/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLParser.scala b/...spark-integration/src/main/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLParser.scala
@@ -50,10 +50,15 @@ class FlintSparkPPLParser(sparkParser: ParserInterface) extends ParserInterface
   private val pplParser = new PPLSyntaxParser()
 
   override def parsePlan(sqlText: String): LogicalPlan = {
-    // if successful build ppl logical plan and translate to catalyst logical plan
-    val context = new CatalystPlanContext
-    planTrnasormer.visit(plan(pplParser, sqlText, false), context)
-    context.getPlan
+    try {
+      // if successful build ppl logical plan and translate to catalyst logical plan
+      val context = new CatalystPlanContext
+      planTrnasormer.visit(plan(pplParser, sqlText, false), context)
+      context.getPlan
+    } catch {
+      // Fall back to Spark parse plan logic if flint cannot parse
+      case _: ParseException | _: SyntaxCheckException => sparkParser.parsePlan(sqlText)
+    }
   }
 
   override def parseExpression(sqlText: String): Expression = sparkParser.parseExpression(sqlText)

diff --git a/spark-sql-application/src/main/scala/org/opensearch/sql/SQLJob.scala b/spark-sql-application/src/main/scala/org/opensearch/sql/SQLJob.scala
@@ -35,7 +35,7 @@ object SQLJob {
 
     val conf: SparkConf = new SparkConf()
       .setAppName("SQLJob")
-      .set("spark.sql.extensions", "org.opensearch.flint.spark.FlintGenericSparkExtensions")
+      .set("spark.sql.extensions", "org.opensearch.flint.spark.FlintSparkExtensions")
       .set("spark.datasource.flint.host", host)
       .set("spark.datasource.flint.port", port)
       .set("spark.datasource.flint.scheme", scheme)