Support more PPL builtin functions by adding a name mapping (opensear…

…ch-project#504) * Support more builtin functions by adding a name mapping Signed-off-by: Lantao Jin <[email protected]> * shorten the map declaration Signed-off-by: Lantao Jin <[email protected]> --------- Signed-off-by: Lantao Jin <[email protected]>
seankao-az · Aug 2, 2024 · d3e54d4 · d3e54d4
1 parent 06ff8be
commit d3e54d4
Show file tree

Hide file tree

Showing 9 changed files with 540 additions and 187 deletions.
diff --git a/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkSuite.scala b/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkSuite.scala
@@ -188,6 +188,28 @@ trait FlintSparkSuite extends QueryTest with FlintSuite with OpenSearchSuite wit
            | """.stripMargin)
   }
 
+  protected def createNullableStateCountryTable(testTable: String): Unit = {
+    sql(s"""
+           | CREATE TABLE $testTable
+           | (
+           |   name STRING,
+           |   age INT,
+           |   state STRING,
+           |   country STRING
+           | )
+           | USING $tableType $tableOptions
+           |""".stripMargin)
+
+    sql(s"""
+           | INSERT INTO $testTable
+           | VALUES ('Jake', 70, 'California', 'USA'),
+           |        ('Hello', 30, 'New York', 'USA'),
+           |        ('John', 25, 'Ontario', 'Canada'),
+           |        ('Jane', 20, 'Quebec', 'Canada'),
+           |        (null, 10, null, 'Canada')
+           | """.stripMargin)
+  }
+
   protected def createOccupationTable(testTable: String): Unit = {
     sql(s"""
       | CREATE TABLE $testTable

diff --git a/...t/src/test/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLBuiltinFunctionITSuite.scala b/...t/src/test/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLBuiltinFunctionITSuite.scala
diff --git a/...ark-integration/src/main/java/org/opensearch/sql/ppl/utils/BuiltinFunctionTranslator.java b/...ark-integration/src/main/java/org/opensearch/sql/ppl/utils/BuiltinFunctionTranslator.java
@@ -5,25 +5,78 @@
 
 package org.opensearch.sql.ppl.utils;
 
+import com.google.common.collect.ImmutableMap;
 import org.apache.spark.sql.catalyst.analysis.UnresolvedFunction;
 import org.apache.spark.sql.catalyst.expressions.Expression;
 import org.opensearch.sql.expression.function.BuiltinFunctionName;
 
 import java.util.List;
-import java.util.Locale;
+import java.util.Map;
 
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.ADD;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.SUBTRACT;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.MULTIPLY;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.DIVIDE;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.MODULUS;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.DAY_OF_WEEK;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.DAY_OF_MONTH;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.DAY_OF_YEAR;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.WEEK_OF_YEAR;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.WEEK;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.MONTH_OF_YEAR;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.HOUR_OF_DAY;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.MINUTE_OF_HOUR;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.SECOND_OF_MINUTE;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.SUBDATE;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.ADDDATE;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.DATEDIFF;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.LOCALTIME;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.IS_NULL;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.IS_NOT_NULL;
 import static org.opensearch.sql.ppl.utils.DataTypeTransformer.seq;
 import static scala.Option.empty;
 
 public interface BuiltinFunctionTranslator {
 
+    /**
+     * The name mapping between PPL builtin functions to Spark builtin functions.
+     */
+    static final Map<BuiltinFunctionName, String> SPARK_BUILTIN_FUNCTION_NAME_MAPPING
+        = new ImmutableMap.Builder<BuiltinFunctionName, String>()
+            // arithmetic operators
+            .put(ADD, "+")
+            .put(SUBTRACT, "-")
+            .put(MULTIPLY, "*")
+            .put(DIVIDE, "/")
+            .put(MODULUS, "%")
+            // time functions
+            .put(DAY_OF_WEEK, "dayofweek")
+            .put(DAY_OF_MONTH, "dayofmonth")
+            .put(DAY_OF_YEAR, "dayofyear")
+            .put(WEEK_OF_YEAR, "weekofyear")
+            .put(WEEK, "weekofyear")
+            .put(MONTH_OF_YEAR, "month")
+            .put(HOUR_OF_DAY, "hour")
+            .put(MINUTE_OF_HOUR, "minute")
+            .put(SECOND_OF_MINUTE, "second")
+            .put(SUBDATE, "date_sub") // only maps subdate(date, days)
+            .put(ADDDATE, "date_add") // only maps adddate(date, days)
+            .put(DATEDIFF, "datediff")
+            .put(LOCALTIME, "localtimestamp")
+            //condition functions
+            .put(IS_NULL, "isnull")
+            .put(IS_NOT_NULL, "isnotnull")
+            .build();
+
     static Expression builtinFunction(org.opensearch.sql.ast.expression.Function function, List<Expression> args) {
         if (BuiltinFunctionName.of(function.getFuncName()).isEmpty()) {
             // TODO change it when UDF is supported
             // TODO should we support more functions which are not PPL builtin functions. E.g Spark builtin functions
             throw new UnsupportedOperationException(function.getFuncName() + " is not a builtin function of PPL");
         } else {
-            String name = BuiltinFunctionName.of(function.getFuncName()).get().name().toLowerCase(Locale.ROOT);
+            BuiltinFunctionName builtin = BuiltinFunctionName.of(function.getFuncName()).get();
+            String name = SPARK_BUILTIN_FUNCTION_NAME_MAPPING
+                .getOrDefault(builtin, builtin.getName().getFunctionName());
             return new UnresolvedFunction(seq(name), seq(args), false, empty(),false);
         }
     }

diff --git a/.../org/opensearch/flint/spark/ppl/PPLLogicalPlanAggregationQueriesTranslatorTestSuite.scala b/.../org/opensearch/flint/spark/ppl/PPLLogicalPlanAggregationQueriesTranslatorTestSuite.scala
@@ -5,25 +5,26 @@
 
 package org.opensearch.flint.spark.ppl
 
-import org.junit.Assert.assertEquals
 import org.opensearch.flint.spark.ppl.PlaneUtils.plan
 import org.opensearch.sql.ppl.{CatalystPlanContext, CatalystQueryPlanVisitor}
 import org.scalatest.matchers.should.Matchers
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedFunction, UnresolvedRelation, UnresolvedStar}
 import org.apache.spark.sql.catalyst.expressions.{Alias, Ascending, Divide, EqualTo, Floor, GreaterThanOrEqual, Literal, Multiply, SortOrder, TimeWindow}
+import org.apache.spark.sql.catalyst.plans.PlanTest
 import org.apache.spark.sql.catalyst.plans.logical._
 
 class PPLLogicalPlanAggregationQueriesTranslatorTestSuite
     extends SparkFunSuite
+    with PlanTest
     with LogicalPlanTestUtils
     with Matchers {
 
   private val planTransformer = new CatalystQueryPlanVisitor()
   private val pplParser = new PPLSyntaxParser()
 
-  test("test average price  ") {
+  test("test average price") {
     // if successful build ppl logical plan and translate to catalyst logical plan
     val context = new CatalystPlanContext
     val logPlan =
@@ -38,10 +39,10 @@ class PPLLogicalPlanAggregationQueriesTranslatorTestSuite
     val aggregatePlan = Aggregate(Seq(), aggregateExpressions, tableRelation)
     val expectedPlan = Project(star, aggregatePlan)
 
-    assertEquals(compareByString(expectedPlan), compareByString(logPlan))
+    comparePlans(expectedPlan, logPlan, false)
   }
 
-  ignore("test average price with Alias") {
+  test("test average price with Alias") {
     // if successful build ppl logical plan and translate to catalyst logical plan
     val context = new CatalystPlanContext
     val logPlan = planTransformer.visit(
@@ -57,7 +58,7 @@ class PPLLogicalPlanAggregationQueriesTranslatorTestSuite
     val aggregatePlan = Aggregate(Seq(), aggregateExpressions, tableRelation)
     val expectedPlan = Project(star, aggregatePlan)
 
-    assertEquals(compareByString(expectedPlan), compareByString(logPlan))
+    comparePlans(expectedPlan, logPlan, false)
   }
 
   test("test average price group by product ") {
@@ -81,7 +82,7 @@ class PPLLogicalPlanAggregationQueriesTranslatorTestSuite
       Aggregate(groupByAttributes, Seq(aggregateExpressions, productAlias), tableRelation)
     val expectedPlan = Project(star, aggregatePlan)
 
-    assertEquals(compareByString(expectedPlan), compareByString(logPlan))
+    comparePlans(expectedPlan, logPlan, false)
   }
 
   test("test average price group by product and filter") {
@@ -109,7 +110,7 @@ class PPLLogicalPlanAggregationQueriesTranslatorTestSuite
       Aggregate(groupByAttributes, Seq(aggregateExpressions, productAlias), filterPlan)
     val expectedPlan = Project(star, aggregatePlan)
 
-    assertEquals(compareByString(expectedPlan), compareByString(logPlan))
+    comparePlans(expectedPlan, logPlan, false)
   }
 
   test("test average price group by product and filter sorted") {
@@ -144,7 +145,7 @@ class PPLLogicalPlanAggregationQueriesTranslatorTestSuite
         global = true,
         aggregatePlan)
     val expectedPlan = Project(star, sortedPlan)
-    assertEquals(compareByString(expectedPlan), compareByString(logPlan))
+    comparePlans(expectedPlan, logPlan, false)
   }
   test("create ppl simple avg age by span of interval of 10 years query test ") {
     val context = new CatalystPlanContext
@@ -164,7 +165,7 @@ class PPLLogicalPlanAggregationQueriesTranslatorTestSuite
     val aggregatePlan = Aggregate(Seq(span), Seq(aggregateExpressions, span), tableRelation)
     val expectedPlan = Project(star, aggregatePlan)
 
-    assert(compareByString(expectedPlan) === compareByString(logPlan))
+    comparePlans(expectedPlan, logPlan, false)
   }
 
   test("create ppl simple avg age by span of interval of 10 years query with sort test ") {
@@ -190,7 +191,7 @@ class PPLLogicalPlanAggregationQueriesTranslatorTestSuite
       Sort(Seq(SortOrder(UnresolvedAttribute("age"), Ascending)), global = true, aggregatePlan)
     val expectedPlan = Project(star, sortedPlan)
 
-    assert(compareByString(expectedPlan) === compareByString(logPlan))
+    comparePlans(expectedPlan, logPlan, false)
   }
 
   test("create ppl simple avg age by span of interval of 10 years by country query test ") {
@@ -219,7 +220,7 @@ class PPLLogicalPlanAggregationQueriesTranslatorTestSuite
       tableRelation)
     val expectedPlan = Project(star, aggregatePlan)
 
-    assert(compareByString(expectedPlan) === compareByString(logPlan))
+    comparePlans(expectedPlan, logPlan, false)
   }
   test("create ppl query count sales by weeks window and productId with sorting test") {
     val context = new CatalystPlanContext
@@ -257,7 +258,7 @@ class PPLLogicalPlanAggregationQueriesTranslatorTestSuite
 
     val expectedPlan = Project(star, sortedPlan)
     // Compare the two plans
-    assert(compareByString(expectedPlan) === compareByString(logPlan))
+    comparePlans(expectedPlan, logPlan, false)
   }
 
   test("create ppl query count sales by days window and productId with sorting test") {
@@ -296,7 +297,7 @@ class PPLLogicalPlanAggregationQueriesTranslatorTestSuite
       aggregatePlan)
     val expectedPlan = Project(star, sortedPlan)
     // Compare the two plans
-    assert(compareByString(expectedPlan) === compareByString(logPlan))
+    comparePlans(expectedPlan, logPlan, false)
   }
   test("create ppl query count status amount by day window and group by status test") {
     val context = new CatalystPlanContext
@@ -331,7 +332,7 @@ class PPLLogicalPlanAggregationQueriesTranslatorTestSuite
     val planWithLimit = GlobalLimit(Literal(100), LocalLimit(Literal(100), aggregatePlan))
     val expectedPlan = Project(star, planWithLimit)
     // Compare the two plans
-    assert(compareByString(expectedPlan) === compareByString(logPlan))
+    comparePlans(expectedPlan, logPlan, false)
   }
   test(
     "create ppl query count only error (status >= 400) status amount by day window and group by status test") {
@@ -368,7 +369,7 @@ class PPLLogicalPlanAggregationQueriesTranslatorTestSuite
     val planWithLimit = GlobalLimit(Literal(100), LocalLimit(Literal(100), aggregatePlan))
     val expectedPlan = Project(star, planWithLimit)
     // Compare the two plans
-    assert(compareByString(expectedPlan) === compareByString(logPlan))
+    comparePlans(expectedPlan, logPlan, false)
   }
 
 }
diff --git a/.../scala/org/opensearch/flint/spark/ppl/PPLLogicalPlanBasicQueriesTranslatorTestSuite.scala b/.../scala/org/opensearch/flint/spark/ppl/PPLLogicalPlanBasicQueriesTranslatorTestSuite.scala
@@ -5,19 +5,19 @@
 
 package org.opensearch.flint.spark.ppl
 
-import org.junit.Assert.assertEquals
 import org.opensearch.flint.spark.ppl.PlaneUtils.plan
 import org.opensearch.sql.ppl.{CatalystPlanContext, CatalystQueryPlanVisitor}
 import org.scalatest.matchers.should.Matchers
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.catalyst.analysis.{UnresolvedAttribute, UnresolvedRelation, UnresolvedStar}
 import org.apache.spark.sql.catalyst.expressions.{Ascending, AttributeReference, Descending, Literal, NamedExpression, SortOrder}
+import org.apache.spark.sql.catalyst.plans.PlanTest
 import org.apache.spark.sql.catalyst.plans.logical._
-import org.apache.spark.sql.types.IntegerType
 
 class PPLLogicalPlanBasicQueriesTranslatorTestSuite
     extends SparkFunSuite
+    with PlanTest
     with LogicalPlanTestUtils
     with Matchers {
 
@@ -31,7 +31,7 @@ class PPLLogicalPlanBasicQueriesTranslatorTestSuite
 
     val projectList: Seq[NamedExpression] = Seq(UnresolvedStar(None))
     val expectedPlan = Project(projectList, UnresolvedRelation(Seq("table")))
-    assertEquals(expectedPlan, logPlan)
+    comparePlans(expectedPlan, logPlan, false)
   }
 
   test("test simple search with escaped table name") {
@@ -41,7 +41,7 @@ class PPLLogicalPlanBasicQueriesTranslatorTestSuite
 
     val projectList: Seq[NamedExpression] = Seq(UnresolvedStar(None))
     val expectedPlan = Project(projectList, UnresolvedRelation(Seq("table")))
-    assertEquals(expectedPlan, logPlan)
+    comparePlans(expectedPlan, logPlan, false)
   }
 
   test("test simple search with schema.table and no explicit fields (defaults to all fields)") {
@@ -51,7 +51,7 @@ class PPLLogicalPlanBasicQueriesTranslatorTestSuite
 
     val projectList: Seq[NamedExpression] = Seq(UnresolvedStar(None))
     val expectedPlan = Project(projectList, UnresolvedRelation(Seq("schema", "table")))
-    assertEquals(expectedPlan, logPlan)
+    comparePlans(expectedPlan, logPlan, false)
 
   }
 
@@ -62,7 +62,7 @@ class PPLLogicalPlanBasicQueriesTranslatorTestSuite
 
     val projectList: Seq[NamedExpression] = Seq(UnresolvedAttribute("A"))
     val expectedPlan = Project(projectList, UnresolvedRelation(Seq("schema", "table")))
-    assertEquals(expectedPlan, logPlan)
+    comparePlans(expectedPlan, logPlan, false)
   }
 
   test("test simple search with only one table with one field projected") {
@@ -72,7 +72,7 @@ class PPLLogicalPlanBasicQueriesTranslatorTestSuite
 
     val projectList: Seq[NamedExpression] = Seq(UnresolvedAttribute("A"))
     val expectedPlan = Project(projectList, UnresolvedRelation(Seq("table")))
-    assertEquals(expectedPlan, logPlan)
+    comparePlans(expectedPlan, logPlan, false)
   }
 
   test("test simple search with only one table with two fields projected") {
@@ -82,7 +82,7 @@ class PPLLogicalPlanBasicQueriesTranslatorTestSuite
     val table = UnresolvedRelation(Seq("t"))
     val projectList = Seq(UnresolvedAttribute("A"), UnresolvedAttribute("B"))
     val expectedPlan = Project(projectList, table)
-    assertEquals(expectedPlan, logPlan)
+    comparePlans(expectedPlan, logPlan, false)
   }
 
   test("test simple search with one table with two fields projected sorted by one field") {
@@ -97,7 +97,7 @@ class PPLLogicalPlanBasicQueriesTranslatorTestSuite
     val sorted = Sort(sortOrder, true, table)
     val expectedPlan = Project(projectList, sorted)
 
-    assert(compareByString(expectedPlan) === compareByString(logPlan))
+    comparePlans(expectedPlan, logPlan, false)
   }
 
   test(
@@ -111,7 +111,7 @@ class PPLLogicalPlanBasicQueriesTranslatorTestSuite
     val planWithLimit =
       GlobalLimit(Literal(5), LocalLimit(Literal(5), Project(projectList, table)))
     val expectedPlan = Project(Seq(UnresolvedStar(None)), planWithLimit)
-    assertEquals(expectedPlan, logPlan)
+    comparePlans(expectedPlan, logPlan, false)
   }
 
   test(
@@ -129,8 +129,7 @@ class PPLLogicalPlanBasicQueriesTranslatorTestSuite
 
     val planWithLimit = GlobalLimit(Literal(5), LocalLimit(Literal(5), projectAB))
     val expectedPlan = Project(Seq(UnresolvedStar(None)), planWithLimit)
-
-    assertEquals(compareByString(expectedPlan), compareByString(logPlan))
+    comparePlans(expectedPlan, logPlan, false)
   }
 
   test(
@@ -152,7 +151,7 @@ class PPLLogicalPlanBasicQueriesTranslatorTestSuite
     val expectedPlan =
       Union(Seq(projectedTable1, projectedTable2), byName = true, allowMissingCol = true)
 
-    assertEquals(expectedPlan, logPlan)
+    comparePlans(expectedPlan, logPlan, false)
   }
 
   test("Search multiple tables - translated into union call with fields") {
@@ -172,6 +171,6 @@ class PPLLogicalPlanBasicQueriesTranslatorTestSuite
     val expectedPlan =
       Union(Seq(projectedTable1, projectedTable2), byName = true, allowMissingCol = true)
 
-    assertEquals(expectedPlan, logPlan)
+    comparePlans(expectedPlan, logPlan, false)
   }
 }