Add Support For TIME Type in "*_OF_YEAR" Functions (#199) (#1223)

Added Support And Tests For Time Type in day_of_year, week_of_year, month_of_year Functions Signed-off-by: GabeFernandez310 <[email protected]> (cherry picked from commit 6e72f18)
opensearch-project · Jan 9, 2023 · 5894ff1 · 5894ff1
1 parent fa1083b
commit 5894ff1
Show file tree

Hide file tree

Showing 7 changed files with 461 additions and 208 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
@@ -335,8 +335,9 @@ public static FunctionExpression dayofyear(Expression... expressions) {
     return compile(FunctionProperties.None, BuiltinFunctionName.DAYOFYEAR, expressions);
   }
 
-  public static FunctionExpression day_of_year(Expression... expressions) {
-    return compile(FunctionProperties.None, BuiltinFunctionName.DAY_OF_YEAR, expressions);
+  public static FunctionExpression day_of_year(
+      FunctionProperties functionProperties, Expression... expressions) {
+    return compile(functionProperties, BuiltinFunctionName.DAY_OF_YEAR, expressions);
   }
 
   public static FunctionExpression day_of_week(
@@ -372,8 +373,9 @@ public static FunctionExpression month(Expression... expressions) {
     return compile(FunctionProperties.None, BuiltinFunctionName.MONTH, expressions);
   }
 
-  public static FunctionExpression month_of_year(Expression... expressions) {
-    return compile(FunctionProperties.None, BuiltinFunctionName.MONTH_OF_YEAR, expressions);
+  public static FunctionExpression month_of_year(
+      FunctionProperties functionProperties, Expression... expressions) {
+    return compile(functionProperties, BuiltinFunctionName.MONTH_OF_YEAR, expressions);
   }
 
   public static FunctionExpression monthname(Expression... expressions) {
@@ -416,12 +418,14 @@ public static FunctionExpression to_days(Expression... expressions) {
     return compile(FunctionProperties.None, BuiltinFunctionName.TO_DAYS, expressions);
   }
 
-  public static FunctionExpression week(Expression... expressions) {
-    return compile(FunctionProperties.None, BuiltinFunctionName.WEEK, expressions);
+  public static FunctionExpression week(
+      FunctionProperties functionProperties, Expression... expressions) {
+    return compile(functionProperties, BuiltinFunctionName.WEEK, expressions);
   }
 
-  public static FunctionExpression week_of_year(Expression... expressions) {
-    return compile(FunctionProperties.None, BuiltinFunctionName.WEEK_OF_YEAR, expressions);
+  public static FunctionExpression week_of_year(
+      FunctionProperties functionProperties, Expression... expressions) {
+    return compile(functionProperties, BuiltinFunctionName.WEEK_OF_YEAR, expressions);
   }
 
   public static FunctionExpression year(Expression... expressions) {

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
@@ -89,6 +89,9 @@ public class DateTimeFunction {
   // 32536771199.999999, or equivalent '3001-01-18 23:59:59.999999' UTC
   private static final Double MYSQL_MAX_TIMESTAMP = 32536771200d;
 
+  // Mode used for week/week_of_year function by default when no argument is provided
+  private static final ExprIntegerValue DEFAULT_WEEK_OF_YEAR_MODE = new ExprIntegerValue(0);
+
   /**
    * Register Date and Time Functions.
    *
@@ -472,6 +475,9 @@ private DefaultFunctionResolver dayOfWeek(FunctionName name) {
    */
   private DefaultFunctionResolver dayOfYear(BuiltinFunctionName dayOfYear) {
     return define(dayOfYear.getName(),
+        implWithProperties(nullMissingHandlingWithProperties((functionProperties, arg)
+            -> DateTimeFunction.dayOfYearToday(
+            functionProperties.getQueryStartClock())), INTEGER, TIME),
         impl(nullMissingHandling(DateTimeFunction::exprDayOfYear), INTEGER, DATE),
         impl(nullMissingHandling(DateTimeFunction::exprDayOfYear), INTEGER, DATETIME),
         impl(nullMissingHandling(DateTimeFunction::exprDayOfYear), INTEGER, TIMESTAMP),
@@ -559,6 +565,9 @@ private DefaultFunctionResolver minute_of_day() {
    */
   private DefaultFunctionResolver month(BuiltinFunctionName month) {
     return define(month.getName(),
+        implWithProperties(nullMissingHandlingWithProperties((functionProperties, arg)
+            -> DateTimeFunction.monthOfYearToday(
+            functionProperties.getQueryStartClock())), INTEGER, TIME),
         impl(nullMissingHandling(DateTimeFunction::exprMonth), INTEGER, DATE),
         impl(nullMissingHandling(DateTimeFunction::exprMonth), INTEGER, DATETIME),
         impl(nullMissingHandling(DateTimeFunction::exprMonth), INTEGER, TIMESTAMP),
@@ -783,10 +792,18 @@ private DefaultFunctionResolver utc_timestamp() {
    */
   private DefaultFunctionResolver week(BuiltinFunctionName week) {
     return define(week.getName(),
+        implWithProperties(nullMissingHandlingWithProperties((functionProperties, arg)
+            -> DateTimeFunction.weekOfYearToday(
+            DEFAULT_WEEK_OF_YEAR_MODE,
+            functionProperties.getQueryStartClock())), INTEGER, TIME),
         impl(nullMissingHandling(DateTimeFunction::exprWeekWithoutMode), INTEGER, DATE),
         impl(nullMissingHandling(DateTimeFunction::exprWeekWithoutMode), INTEGER, DATETIME),
         impl(nullMissingHandling(DateTimeFunction::exprWeekWithoutMode), INTEGER, TIMESTAMP),
         impl(nullMissingHandling(DateTimeFunction::exprWeekWithoutMode), INTEGER, STRING),
+        implWithProperties(nullMissingHandlingWithProperties((functionProperties, time, modeArg)
+            -> DateTimeFunction.weekOfYearToday(
+                modeArg,
+            functionProperties.getQueryStartClock())), INTEGER, TIME, INTEGER),
         impl(nullMissingHandling(DateTimeFunction::exprWeek), INTEGER, DATE, INTEGER),
         impl(nullMissingHandling(DateTimeFunction::exprWeek), INTEGER, DATETIME, INTEGER),
         impl(nullMissingHandling(DateTimeFunction::exprWeek), INTEGER, TIMESTAMP, INTEGER),
@@ -827,6 +844,15 @@ private DefaultFunctionResolver date_format() {
     );
   }
 
+  private ExprValue dayOfYearToday(Clock clock) {
+    return new ExprIntegerValue(LocalDateTime.now(clock).getDayOfYear());
+  }
+
+  private ExprValue weekOfYearToday(ExprValue mode, Clock clock) {
+    return new ExprIntegerValue(
+        CalendarLookup.getWeekNumber(mode.integerValue(), LocalDateTime.now(clock).toLocalDate()));
+  }
+
   /**
    * Day of Week implementation for ExprValue when passing in an arguemt of type TIME.
    *
@@ -1519,6 +1545,10 @@ private ExprValue exprYear(ExprValue date) {
     return new ExprIntegerValue(date.dateValue().getYear());
   }
 
+  private ExprValue monthOfYearToday(Clock clock) {
+    return new ExprIntegerValue(LocalDateTime.now(clock).getMonthValue());
+  }
+
   private LocalDateTime formatNow(Clock clock) {
     return formatNow(clock, 0);
   }

diff --git a/core/src/main/java/org/opensearch/sql/expression/function/FunctionDSL.java b/core/src/main/java/org/opensearch/sql/expression/function/FunctionDSL.java
@@ -141,11 +141,11 @@ public String toString() {
    * Implementation of a function that takes two arguments, returns a value, and
    * requires FunctionProperties to complete.
    *
-   * @param function   {@link ExprValue} based unary function.
+   * @param function   {@link ExprValue} based Binary function.
    * @param returnType return type.
    * @param args1Type first argument type.
    * @param args2Type second argument type.
-   * @return Unary Function Implementation.
+   * @return Binary Function Implementation.
    */
   public static SerializableFunction<FunctionName, Pair<FunctionSignature, FunctionBuilder>>
       implWithProperties(