Add last_day Function To OpenSearch SQL Plugin (#1344)

Signed-off-by: GabeFernandez310 <[email protected]>
opensearch-project · Feb 16, 2023 · d9114f5 · d9114f5
1 parent 672c72f
commit d9114f5
Show file tree

Hide file tree

Showing 9 changed files with 179 additions and 0 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
@@ -361,6 +361,11 @@ public static FunctionExpression hour_of_day(Expression... expressions) {
     return compile(FunctionProperties.None, BuiltinFunctionName.HOUR_OF_DAY, expressions);
   }
 
+  public static FunctionExpression last_day(FunctionProperties functionProperties,
+                                            Expression... expressions) {
+    return compile(functionProperties, BuiltinFunctionName.LAST_DAY, expressions);
+  }
+
   public static FunctionExpression microsecond(Expression... expressions) {
     return compile(FunctionProperties.None, BuiltinFunctionName.MICROSECOND, 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
@@ -160,6 +160,7 @@ public void register(BuiltinFunctionRepository repository) {
     repository.register(get_format());
     repository.register(hour(BuiltinFunctionName.HOUR));
     repository.register(hour(BuiltinFunctionName.HOUR_OF_DAY));
+    repository.register(last_day());
     repository.register(localtime());
     repository.register(localtimestamp());
     repository.register(makedate());
@@ -566,6 +567,18 @@ private DefaultFunctionResolver hour(BuiltinFunctionName name) {
     );
   }
 
+  private  DefaultFunctionResolver last_day() {
+    return define(BuiltinFunctionName.LAST_DAY.getName(),
+        impl(nullMissingHandling(DateTimeFunction::exprLastDay), DATE, STRING),
+        implWithProperties(nullMissingHandlingWithProperties((functionProperties, arg)
+            -> DateTimeFunction.exprLastDayToday(
+            functionProperties.getQueryStartClock())), DATE, TIME),
+        impl(nullMissingHandling(DateTimeFunction::exprLastDay), DATE, DATE),
+        impl(nullMissingHandling(DateTimeFunction::exprLastDay), DATE, DATETIME),
+        impl(nullMissingHandling(DateTimeFunction::exprLastDay), DATE, TIMESTAMP)
+    );
+  }
+
   private FunctionResolver makedate() {
     return define(BuiltinFunctionName.MAKEDATE.getName(),
         impl(nullMissingHandling(DateTimeFunction::exprMakeDate), DATE, DOUBLE, DOUBLE));
@@ -1287,6 +1300,39 @@ private ExprValue exprHour(ExprValue time) {
         HOURS.between(LocalTime.MIN, time.timeValue()));
   }
 
+  /**
+   * Helper function to retrieve the last day of a month based on a LocalDate argument.
+   *
+   * @param today a LocalDate.
+   * @return a LocalDate associated with the last day of the month for the given input.
+   */
+  private LocalDate getLastDay(LocalDate today) {
+    return LocalDate.of(
+        today.getYear(),
+        today.getMonth(),
+        today.getMonth().length(today.isLeapYear()));
+  }
+
+  /**
+   * Returns a DATE for the last day of the month of a given argument.
+   *
+   * @param datetime A DATE/DATETIME/TIMESTAMP/STRING ExprValue.
+   * @return An DATE value corresponding to the last day of the month of the given argument.
+   */
+  private ExprValue exprLastDay(ExprValue datetime) {
+    return new ExprDateValue(getLastDay(datetime.dateValue()));
+  }
+
+  /**
+   * Returns a DATE for the last day of the current month.
+   *
+   * @param clock The clock for the query start time from functionProperties.
+   * @return An DATE value corresponding to the last day of the month of the given argument.
+   */
+  private ExprValue exprLastDayToday(Clock clock) {
+    return new ExprDateValue(getLastDay(formatNow(clock).toLocalDate()));
+  }
+
   /**
    * Following MySQL, function receives arguments of type double and rounds them before use.
    * Furthermore:

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
@@ -81,6 +81,7 @@ public enum BuiltinFunctionName {
   GET_FORMAT(FunctionName.of("get_format")),
   HOUR(FunctionName.of("hour")),
   HOUR_OF_DAY(FunctionName.of("hour_of_day")),
+  LAST_DAY(FunctionName.of("last_day")),
   MAKEDATE(FunctionName.of("makedate")),
   MAKETIME(FunctionName.of("maketime")),
   MICROSECOND(FunctionName.of("microsecond")),

diff --git a/core/src/test/java/org/opensearch/sql/expression/datetime/DateTimeFunctionTest.java b/core/src/test/java/org/opensearch/sql/expression/datetime/DateTimeFunctionTest.java
@@ -862,6 +862,80 @@ public void hourOfDayInvalidArguments() {
 
   }
 
+  private void checkForExpectedDay(
+      FunctionExpression functionExpression,
+      String expectedDay,
+      String testExpr) {
+    assertEquals(DATE, functionExpression.type());
+    assertEquals(new ExprDateValue(expectedDay), eval(functionExpression));
+    assertEquals(testExpr, functionExpression.toString());
+  }
+
+  private static Stream<Arguments> getTestDataForLastDay() {
+    return Stream.of(
+        Arguments.of(new ExprDateValue("2017-01-20"), "2017-01-31", "last_day(DATE '2017-01-20')"),
+        //Leap year
+        Arguments.of(new ExprDateValue("2020-02-20"), "2020-02-29", "last_day(DATE '2020-02-20')"),
+        //Non leap year
+        Arguments.of(new ExprDateValue("2017-02-20"), "2017-02-28", "last_day(DATE '2017-02-20')"),
+        Arguments.of(new ExprDateValue("2017-03-20"), "2017-03-31", "last_day(DATE '2017-03-20')"),
+        Arguments.of(new ExprDateValue("2017-04-20"), "2017-04-30", "last_day(DATE '2017-04-20')"),
+        Arguments.of(new ExprDateValue("2017-05-20"), "2017-05-31", "last_day(DATE '2017-05-20')"),
+        Arguments.of(new ExprDateValue("2017-06-20"), "2017-06-30", "last_day(DATE '2017-06-20')"),
+        Arguments.of(new ExprDateValue("2017-07-20"), "2017-07-31", "last_day(DATE '2017-07-20')"),
+        Arguments.of(new ExprDateValue("2017-08-20"), "2017-08-31", "last_day(DATE '2017-08-20')"),
+        Arguments.of(new ExprDateValue("2017-09-20"), "2017-09-30", "last_day(DATE '2017-09-20')"),
+        Arguments.of(new ExprDateValue("2017-10-20"), "2017-10-31", "last_day(DATE '2017-10-20')"),
+        Arguments.of(new ExprDateValue("2017-11-20"), "2017-11-30", "last_day(DATE '2017-11-20')"),
+        Arguments.of(new ExprDateValue("2017-12-20"), "2017-12-31", "last_day(DATE '2017-12-20')")
+    );
+  }
+
+  @ParameterizedTest(name = "{2}")
+  @MethodSource("getTestDataForLastDay")
+  public void testLastDay(ExprValue testedDateTime, String expectedResult, String expectedQuery) {
+    lenient().when(nullRef.valueOf(env)).thenReturn(nullValue());
+    lenient().when(missingRef.valueOf(env)).thenReturn(missingValue());
+
+    checkForExpectedDay(
+        DSL.last_day(functionProperties, DSL.literal(testedDateTime)),
+        expectedResult,
+        expectedQuery
+    );
+  }
+
+  @Test
+  public void testLastDayWithTimeType() {
+    lenient().when(nullRef.valueOf(env)).thenReturn(nullValue());
+    lenient().when(missingRef.valueOf(env)).thenReturn(missingValue());
+
+    FunctionExpression expression = DSL.last_day(
+        functionProperties, DSL.literal(new ExprTimeValue("12:23:34")));
+
+    LocalDate expected = LocalDate.now(functionProperties.getQueryStartClock());
+    LocalDate result = eval(expression).dateValue();
+
+
+    assertAll(
+        () -> assertEquals((expected.lengthOfMonth()), result.getDayOfMonth()),
+        () -> assertEquals("last_day(TIME '12:23:34')", expression.toString())
+    );
+  }
+
+  private void lastDay(String date) {
+    FunctionExpression expression = DSL.day_of_week(
+        functionProperties, DSL.literal(new ExprDateValue(date)));
+    eval(expression);
+  }
+
+  @Test
+  public void testLastDayInvalidArgument() {
+    lenient().when(nullRef.valueOf(env)).thenReturn(nullValue());
+    lenient().when(missingRef.valueOf(env)).thenReturn(missingValue());
+
+    assertThrows(SemanticCheckException.class, () ->  lastDay("asdfasdf"));
+  }
+
   @Test
   public void microsecond() {
     when(nullRef.type()).thenReturn(TIME);

diff --git a/docs/user/dql/functions.rst b/docs/user/dql/functions.rst
@@ -2745,6 +2745,24 @@ Example::
     | hello,world                        |
     +------------------------------------+
 
+LAST_DAY
+--------
+
+Usage: Returns the last day of the month as a DATE for a valid argument.
+
+Argument type: DATE/DATETIME/STRING/TIMESTAMP/TIME
+
+Return type: DATE
+
+Example::
+
+    os> SELECT last_day('2023-02-06');
+    fetched rows / total rows = 1/1
+    +--------------------------+
+    | last_day('2023-02-06')   |
+    |--------------------------|
+    | 2023-02-28               |
+    +--------------------------+
 
 LEFT
 ----

diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/DateTimeFunctionIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/DateTimeFunctionIT.java
@@ -513,6 +513,33 @@ public void testHourFunctionAliasesReturnTheSameResults() throws IOException {
     result1.getJSONArray("datarows").similar(result2.getJSONArray("datarows"));
   }
 
+  @Test
+  public void testLastDay() throws IOException {
+    JSONObject result = executeQuery(
+        String.format("SELECT last_day(cast(date0 as date)) FROM %s LIMIT 3",
+            TEST_INDEX_CALCS));
+    verifyDataRows(result,
+        rows("2004-04-30"),
+        rows("1972-07-31"),
+        rows("1975-11-30"));
+
+    result = executeQuery(
+        String.format("SELECT last_day(datetime(cast(date0 AS string))) FROM %s LIMIT 3",
+            TEST_INDEX_CALCS));
+    verifyDataRows(result,
+        rows("2004-04-30"),
+        rows("1972-07-31"),
+        rows("1975-11-30"));
+
+    result = executeQuery(
+        String.format("SELECT last_day(cast(date0 AS timestamp)) FROM %s LIMIT 3",
+            TEST_INDEX_CALCS));
+    verifyDataRows(result,
+        rows("2004-04-30"),
+        rows("1972-07-31"),
+        rows("1975-11-30"));
+  }
+
   @Test
   public void testMicrosecond() throws IOException {
     JSONObject result = executeQuery("select microsecond(timestamp('2020-09-16 17:30:00.123456'))");

diff --git a/sql/src/main/antlr/OpenSearchSQLLexer.g4 b/sql/src/main/antlr/OpenSearchSQLLexer.g4
@@ -217,6 +217,7 @@ GET_FORMAT:                         'GET_FORMAT';
 IF:                                 'IF';
 IFNULL:                             'IFNULL';
 ISNULL:                             'ISNULL';
+LAST_DAY:                           'LAST_DAY';
 LENGTH:                             'LENGTH';
 LN:                                 'LN';
 LOCALTIME:                          'LOCALTIME';

diff --git a/sql/src/main/antlr/OpenSearchSQLParser.g4 b/sql/src/main/antlr/OpenSearchSQLParser.g4
@@ -452,6 +452,7 @@ dateTimeFunctionName
     | FROM_UNIXTIME
     | HOUR
     | HOUR_OF_DAY
+    | LAST_DAY
     | MAKEDATE
     | MAKETIME
     | MICROSECOND

diff --git a/sql/src/test/java/org/opensearch/sql/sql/antlr/SQLSyntaxParserTest.java b/sql/src/test/java/org/opensearch/sql/sql/antlr/SQLSyntaxParserTest.java
@@ -521,6 +521,12 @@ public void can_parse_minute_of_day_function() {
     assertNotNull(parser.parse("SELECT minute_of_day('2022-12-14 12:23:34');"));;
   }
 
+  @Test
+  public void can_parse_last_day_function() {
+    assertNotNull(parser.parse("SELECT last_day(\"2017-06-20\")"));
+    assertNotNull(parser.parse("SELECT last_day('2004-01-01 01:01:01')"));
+  }
+
   @Test
   public void can_parse_wildcard_query_relevance_function() {
     assertNotNull(