Added IT tests

Signed-off-by: Hendrik Saly <[email protected]>
opensearch-project · Oct 28, 2024 · 8dd9029 · 8dd9029
1 parent 1f43c23
commit 8dd9029
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 11 deletions.
diff --git a/...st/src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLBetweenITSuite.scala b/...st/src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLBetweenITSuite.scala
@@ -5,6 +5,8 @@
 
 package org.opensearch.flint.spark.ppl
 
+import java.sql.Timestamp
+
 import org.apache.spark.sql.QueryTest
 import org.apache.spark.sql.streaming.StreamTest
 
@@ -56,12 +58,12 @@ class FlintSparkPPLBetweenITSuite
                        | """.stripMargin)
 
     val results = frame.collect()
-    assert(results.length == 3)
+    assert(results.length == 1)
     assert(frame.columns.length == 6)
 
     results.foreach(row => {
       val age = row.getAs[Int]("age")
-      assert(age >= 20 && age <= 30, s"Age $age is not between 20 and 30")
+      assert(age >= 21 && age <= 29, s"Age $age is not between 21 and 29")
     })
   }
 
@@ -80,36 +82,53 @@ class FlintSparkPPLBetweenITSuite
     })
   }
 
+  test("test between should return records where NOT between two integer values") {
+    val frame = sql(s"""
+                       | source = $testTable | where NOT age between 20 and 30
+                       | """.stripMargin)
+
+    val results = frame.collect()
+    assert(results.length == 1)
+    assert(frame.columns.length == 6)
+
+    results.foreach(row => {
+      val age = row.getAs[Int]("age")
+      assert(age < 20 || age > 30, s"Age $age is not between 20 and 30")
+    })
+  }
+
   test("test between should return records between two date values") {
     val frame = sql(s"""
-                       | source = $timeSeriesTestTable | where timestamp between '2023-10-01 00:01:00' and '2023-10-01 00:10:00'
+                       | source = $timeSeriesTestTable | where time between '2023-10-01 00:01:00' and '2023-10-01 00:10:00'
                        | """.stripMargin)
 
     val results = frame.collect()
     assert(results.length == 2)
     assert(frame.columns.length == 4)
 
     results.foreach(row => {
-      val ts = row.getAs[String]("timestamp")
+      val ts = row.getAs[Timestamp]("time")
       assert(
-        ts >= "2023-10-01 00:01:00" && ts <= "2023-10-01 00:01:00",
+        !ts.before(Timestamp.valueOf("2023-10-01 00:01:00")) || !ts.after(
+          Timestamp.valueOf("2023-10-01 00:01:00")),
         s"Timestamp $ts is not between '2023-10-01 00:01:00' and '2023-10-01 00:10:00'")
     })
   }
 
   test("test between should return records NOT between two date values") {
     val frame = sql(s"""
-                       | source = $timeSeriesTestTable | where timestamp NOT between '2023-10-01 00:01:00' and '2023-10-01 00:10:00'
+                       | source = $timeSeriesTestTable | where time NOT between '2023-10-01 00:01:00' and '2023-10-01 00:10:00'
                        | """.stripMargin)
 
     val results = frame.collect()
-    assert(results.length == 4)
+    assert(results.length == 3)
     assert(frame.columns.length == 4)
 
     results.foreach(row => {
-      val ts = row.getAs[String]("timestamp")
+      val ts = row.getAs[Timestamp]("time")
       assert(
-        ts < "2023-10-01 00:01:00" || ts > "2023-10-01 00:01:00",
+        ts.before(Timestamp.valueOf("2023-10-01 00:01:00")) || ts.after(
+          Timestamp.valueOf("2023-10-01 00:01:00")),
         s"Timestamp $ts is not between '2023-10-01 00:01:00' and '2023-10-01 00:10:00'")
     })
 

diff --git a/ppl-spark-integration/src/main/antlr4/OpenSearchPPLParser.g4 b/ppl-spark-integration/src/main/antlr4/OpenSearchPPLParser.g4
@@ -386,7 +386,7 @@ logicalExpression
 comparisonExpression
    : left = valueExpression comparisonOperator right = valueExpression  # compareExpr
    | valueExpression IN valueList                                       # inExpr
-   | expr1 = functionArg BETWEEN expr2 = functionArg AND expr3 = functionArg   # between
+   | expr1 = functionArg NOT? BETWEEN expr2 = functionArg AND expr3 = functionArg   # between
    ;
 
 valueExpressionList

diff --git a/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java b/ppl-spark-integration/src/main/java/org/opensearch/sql/ppl/parser/AstExpressionBuilder.java
@@ -292,7 +292,7 @@ public UnresolvedExpression visitConvertedDataType(OpenSearchPPLParser.Converted
     @Override
     public UnresolvedExpression visitBetween(OpenSearchPPLParser.BetweenContext ctx) {
         UnresolvedExpression betweenExpr = new Between(visit(ctx.expr1),visit(ctx.expr2),visit(ctx.expr3));
-        return betweenExpr;
+        return ctx.NOT() != null ? new Not(betweenExpr) : betweenExpr;
     }
 
     private Function buildFunction(