Add Timestampdiff Function To OpenSearch SQL (#1467)

* Add `Timestampdiff` Function To OpenSearch SQL Added Tests And Implementation For Timestampdiff --------- Signed-off-by: GabeFernandez310 <[email protected]> * Removed Import Signed-off-by: GabeFernandez310 <[email protected]> --------- Signed-off-by: GabeFernandez310 <[email protected]> (cherry picked from commit 31148da)
opensearch-project · Mar 24, 2023 · b9a568b · b9a568b
1 parent 49a82ae
commit b9a568b
Show file tree

Hide file tree

Showing 11 changed files with 534 additions and 13 deletions.
diff --git a/core/src/main/java/org/opensearch/sql/expression/DSL.java b/core/src/main/java/org/opensearch/sql/expression/DSL.java
@@ -924,6 +924,12 @@ public static FunctionExpression timestampadd(FunctionProperties functionPropert
     return compile(functionProperties, BuiltinFunctionName.TIMESTAMPADD, expressions);
   }
 
+  public static FunctionExpression timestampdiff(FunctionProperties functionProperties,
+                                                Expression... expressions) {
+    return compile(functionProperties, BuiltinFunctionName.TIMESTAMPDIFF, expressions);
+  }
+
+
   public static FunctionExpression utc_date(FunctionProperties functionProperties,
                                                 Expression... args) {
     return compile(functionProperties, BuiltinFunctionName.UTC_DATE, args);

diff --git a/core/src/main/java/org/opensearch/sql/expression/datetime/DateTimeFunction.java b/core/src/main/java/org/opensearch/sql/expression/datetime/DateTimeFunction.java
@@ -235,6 +235,7 @@ public void register(BuiltinFunctionRepository repository) {
     repository.register(timediff());
     repository.register(timestamp());
     repository.register(timestampadd());
+    repository.register(timestampdiff());
     repository.register(utc_date());
     repository.register(utc_time());
     repository.register(utc_timestamp());
@@ -899,6 +900,15 @@ private DefaultFunctionResolver timestamp() {
             TIMESTAMP, TIMESTAMP, TIMESTAMP));
   }
 
+  /**
+   * Adds an interval of time to the provided DATE/DATETIME/TIME/TIMESTAMP/STRING argument.
+   * The interval of time added is determined by the given first and second arguments.
+   * The first argument is an interval type, and must be one of the tokens below...
+   * [MICROSECOND, SECOND, MINUTE, HOUR, DAY, WEEK, MONTH, QUARTER, YEAR]
+   * The second argument is the amount of the interval type to be added.
+   * The third argument is the DATE/DATETIME/TIME/TIMESTAMP/STRING to add to.
+   * @return The DATETIME representing the summed DATE/DATETIME/TIME/TIMESTAMP and interval.
+   */
   private DefaultFunctionResolver timestampadd() {
     return define(BuiltinFunctionName.TIMESTAMPADD.getName(),
         impl(nullMissingHandling(DateTimeFunction::exprTimestampAdd),
@@ -915,6 +925,35 @@ private DefaultFunctionResolver timestampadd() {
             DATETIME, STRING, INTEGER, TIME));
   }
 
+  /**
+   * Finds the difference between provided DATE/DATETIME/TIME/TIMESTAMP/STRING arguments.
+   * The first argument is an interval type, and must be one of the tokens below...
+   * [MICROSECOND, SECOND, MINUTE, HOUR, DAY, WEEK, MONTH, QUARTER, YEAR]
+   * The second argument the DATE/DATETIME/TIME/TIMESTAMP/STRING representing the start time.
+   * The third argument is the DATE/DATETIME/TIME/TIMESTAMP/STRING representing the end time.
+   * @return A LONG representing the difference between arguments, using the given interval type.
+   */
+  private DefaultFunctionResolver timestampdiff() {
+    return define(BuiltinFunctionName.TIMESTAMPDIFF.getName(),
+        impl(nullMissingHandling(DateTimeFunction::exprTimestampDiff),
+            DATETIME, STRING, DATETIME, DATETIME),
+        impl(nullMissingHandling(DateTimeFunction::exprTimestampDiff),
+            DATETIME, STRING, DATETIME, TIMESTAMP),
+        impl(nullMissingHandling(DateTimeFunction::exprTimestampDiff),
+            DATETIME, STRING, TIMESTAMP, DATETIME),
+        impl(nullMissingHandling(DateTimeFunction::exprTimestampDiff),
+            DATETIME, STRING, TIMESTAMP, TIMESTAMP),
+        implWithProperties(
+            nullMissingHandlingWithProperties(
+                (functionProperties, part, startTime, endTime) -> exprTimestampDiffForTimeType(
+                    functionProperties,
+                    part,
+                    startTime,
+                    endTime)),
+            DATETIME, STRING, TIME, TIME)
+    );
+  }
+
   /**
    * TO_DAYS(STRING/DATE/DATETIME/TIMESTAMP). return the day number of the given date.
    */
@@ -1871,6 +1910,62 @@ private ExprValue exprTimestampAddForTimeType(Clock clock,
     return exprTimestampAdd(partExpr, amountExpr, new ExprDatetimeValue(datetime));
   }
 
+  private ExprValue getTimeDifference(String part, LocalDateTime startTime, LocalDateTime endTime) {
+    long returnVal;
+    switch (part) {
+      case "MICROSECOND":
+        returnVal = MICROS.between(startTime, endTime);
+        break;
+      case "SECOND":
+        returnVal = SECONDS.between(startTime, endTime);
+        break;
+      case "MINUTE":
+        returnVal = MINUTES.between(startTime, endTime);
+        break;
+      case "HOUR":
+        returnVal = HOURS.between(startTime, endTime);
+        break;
+      case "DAY":
+        returnVal = DAYS.between(startTime, endTime);
+        break;
+      case "WEEK":
+        returnVal = WEEKS.between(startTime, endTime);
+        break;
+      case "MONTH":
+        returnVal = MONTHS.between(startTime, endTime);
+        break;
+      case "QUARTER":
+        returnVal = MONTHS.between(startTime, endTime) / 3;
+        break;
+      case "YEAR":
+        returnVal = YEARS.between(startTime, endTime);
+        break;
+      default:
+        return ExprNullValue.of();
+    }
+    return new ExprLongValue(returnVal);
+  }
+
+  private ExprValue exprTimestampDiff(
+                                      ExprValue partExpr,
+                                      ExprValue startTimeExpr,
+                                      ExprValue endTimeExpr) {
+    return getTimeDifference(
+        partExpr.stringValue(),
+        startTimeExpr.datetimeValue(),
+        endTimeExpr.datetimeValue());
+  }
+
+  private ExprValue exprTimestampDiffForTimeType(FunctionProperties fp,
+                                                 ExprValue partExpr,
+                                                 ExprValue startTimeExpr,
+                                                 ExprValue endTimeExpr) {
+    return getTimeDifference(
+        partExpr.stringValue(),
+        extractDateTime(startTimeExpr, fp),
+        extractDateTime(endTimeExpr, fp));
+  }
+
   /**
    * UTC_DATE implementation for ExprValue.
    *

diff --git a/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java b/core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java
@@ -108,6 +108,7 @@ public enum BuiltinFunctionName {
   TIME_TO_SEC(FunctionName.of("time_to_sec")),
   TIMESTAMP(FunctionName.of("timestamp")),
   TIMESTAMPADD(FunctionName.of("timestampadd")),
+  TIMESTAMPDIFF(FunctionName.of("timestampdiff")),
   TIME_FORMAT(FunctionName.of("time_format")),
   TO_DAYS(FunctionName.of("to_days")),
   TO_SECONDS(FunctionName.of("to_seconds")),