Add GET_FORMAT Function To OpenSearch SQL Plugin

Adds the get_format function to the OpenSearch SQL Plugin. It's implementation is aligned with the MySQL Documentation. Note: This currently only supports a USA format for DATE, TIME, TIMESTAMP and DATETIME. Signed-off-by: GabeFernandez310 <[email protected]>
ykmr1224 · Feb 15, 2023 · 72547f4 · 72547f4
1 parent 2323592
commit 72547f4
Show file tree

Hide file tree

Showing 11 changed files with 177 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
@@ -349,6 +349,10 @@ public static FunctionExpression from_days(Expression... expressions) {
     return compile(FunctionProperties.None, BuiltinFunctionName.FROM_DAYS, expressions);
   }
 
+  public static FunctionExpression get_format(Expression... expressions) {
+    return compile(FunctionProperties.None, BuiltinFunctionName.GET_FORMAT, expressions);
+  }
+
   public static FunctionExpression hour(Expression... expressions) {
     return compile(FunctionProperties.None, BuiltinFunctionName.HOUR, 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
@@ -34,6 +34,8 @@
 import static org.opensearch.sql.utils.DateTimeUtils.extractDate;
 import static org.opensearch.sql.utils.DateTimeUtils.extractDateTime;
 
+import com.google.common.collect.ImmutableTable;
+import com.google.common.collect.Table;
 import java.math.BigDecimal;
 import java.math.RoundingMode;
 import java.text.DecimalFormat;
@@ -101,6 +103,16 @@ public class DateTimeFunction {
   // 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);
 
+  // Map used to determine format output for the get_format function
+  private static final Table<String, String, String> formats =
+      ImmutableTable.<String, String, String>builder()
+      //TODO: Add support for other formats
+      .put("date", "usa", "%m.%d.%Y")
+      .put("time", "usa", "%h:%i:%s %p")
+      .put("datetime", "usa", "%Y-%m-%d %H.%i.%s")
+      .put("timestamp", "usa", "%Y-%m-%d %H.%i.%s")
+      .build();
+
   /**
    * Register Date and Time Functions.
    *
@@ -130,6 +142,7 @@ public void register(BuiltinFunctionRepository repository) {
     repository.register(dayOfYear(BuiltinFunctionName.DAY_OF_YEAR));
     repository.register(from_days());
     repository.register(from_unixtime());
+    repository.register(get_format());
     repository.register(hour(BuiltinFunctionName.HOUR));
     repository.register(hour(BuiltinFunctionName.HOUR_OF_DAY));
     repository.register(localtime());
@@ -518,6 +531,12 @@ private FunctionResolver from_unixtime() {
             STRING, DOUBLE, STRING));
   }
 
+  private DefaultFunctionResolver get_format() {
+    return define(BuiltinFunctionName.GET_FORMAT.getName(),
+        impl(nullMissingHandling(DateTimeFunction::exprGetFormat), STRING, STRING, STRING)
+    );
+  }
+
   /**
    * HOUR(STRING/TIME/DATETIME/DATE/TIMESTAMP). return the hour value for time.
    */
@@ -1193,6 +1212,23 @@ private ExprValue exprFromUnixTimeFormat(ExprValue time, ExprValue format) {
     return DateTimeFormatterUtil.getFormattedDate(value, format);
   }
 
+  /**
+   * get_format implementation for ExprValue.
+   *
+   * @param type ExprValue of the type.
+   * @param format ExprValue of Time/String type
+   * @return ExprValue..
+   */
+  private ExprValue exprGetFormat(ExprValue type, ExprValue format) {
+    if (formats.contains(type.stringValue().toLowerCase(), format.stringValue().toLowerCase())) {
+      return new ExprStringValue(formats.get(
+          type.stringValue().toLowerCase(),
+          format.stringValue().toLowerCase()));
+    }
+
+    return ExprNullValue.of();
+  }
+
   /**
    * Hour 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
@@ -77,6 +77,7 @@ public enum BuiltinFunctionName {
   DAY_OF_YEAR(FunctionName.of("day_of_year")),
   FROM_DAYS(FunctionName.of("from_days")),
   FROM_UNIXTIME(FunctionName.of("from_unixtime")),
+  GET_FORMAT(FunctionName.of("get_format")),
   HOUR(FunctionName.of("hour")),
   HOUR_OF_DAY(FunctionName.of("hour_of_day")),
   MAKEDATE(FunctionName.of("makedate")),

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
@@ -41,6 +41,7 @@
 import org.opensearch.sql.data.model.ExprDateValue;
 import org.opensearch.sql.data.model.ExprDatetimeValue;
 import org.opensearch.sql.data.model.ExprLongValue;
+import org.opensearch.sql.data.model.ExprStringValue;
 import org.opensearch.sql.data.model.ExprTimeValue;
 import org.opensearch.sql.data.model.ExprTimestampValue;
 import org.opensearch.sql.data.model.ExprValue;
@@ -674,6 +675,47 @@ public void from_days() {
     assertEquals(new ExprDateValue("2000-07-03"), expression.valueOf(env));
   }
 
+  private static Stream<Arguments> getTestDataForGetFormat() {
+    return Stream.of(
+        Arguments.of("DATE", "USA", "%m.%d.%Y"),
+        Arguments.of("DATETIME", "USA", "%Y-%m-%d %H.%i.%s"),
+        Arguments.of("TIMESTAMP", "USA", "%Y-%m-%d %H.%i.%s"),
+        Arguments.of("TIME", "USA", "%h:%i:%s %p")
+    );
+  }
+
+  private void getFormatQuery(LiteralExpression argType,
+                               LiteralExpression namedFormat,
+                               String expectedResult) {
+    FunctionExpression expr = DSL.get_format(argType, namedFormat);
+    assertEquals(STRING, expr.type());
+    assertEquals(expectedResult, eval(expr).stringValue());
+  }
+
+  @ParameterizedTest(name = "{0}{1}")
+  @MethodSource("getTestDataForGetFormat")
+  public void testGetFormat(String arg,
+                             String format,
+                             String expectedResult) {
+    lenient().when(nullRef.valueOf(env)).thenReturn(nullValue());
+    lenient().when(missingRef.valueOf(env)).thenReturn(missingValue());
+
+    getFormatQuery(
+        DSL.literal(arg),
+        DSL.literal(new ExprStringValue(format)),
+        expectedResult);
+  }
+
+  @Test
+  public void testGetFormatInvalidFormat() {
+    lenient().when(nullRef.valueOf(env)).thenReturn(nullValue());
+    lenient().when(missingRef.valueOf(env)).thenReturn(missingValue());
+    FunctionExpression expr = DSL.get_format(
+        DSL.literal("DATE"),
+        DSL.literal("1SA"));
+    assertEquals(nullValue(), eval(expr));
+  }
+
   @Test
   public void hour() {
     when(nullRef.type()).thenReturn(TIME);

diff --git a/docs/user/dql/functions.rst b/docs/user/dql/functions.rst
@@ -1762,6 +1762,29 @@ Examples::
     +-----------------------------------+
 
 
+GET_FORMAT
+----------
+
+Description
+>>>>>>>>>>>
+
+Usage: Returns a string value containing string format specifiers based on the input arguments.
+
+Argument type: TYPE, STRING
+TYPE must be one of the following tokens: [DATE, TIME, DATETIME, TIMESTAMP].
+STRING must be one of the following tokens: ["USA"] (" can be replaced by ').
+
+Examples::
+
+    os> select GET_FORMAT(DATE, 'USA');
+    fetched rows / total rows = 1/1
+    +---------------------------+
+    | GET_FORMAT(DATE, 'USA')   |
+    |---------------------------|
+    | %m.%d.%Y                  |
+    +---------------------------+
+
+
 HOUR
 ----
 

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
@@ -1195,6 +1195,12 @@ public void testFromUnixTime() throws IOException {
         rows("1976-05-07 07:00:00", "1970-01-01 03:23:44.12", "01:41:56"));
   }
 
+  @Test
+  public void testGetFormatAsArgument() throws IOException{
+    var result = executeQuery("SELECT DATE_FORMAT('2003-10-03',GET_FORMAT(DATE,'USA'))");
+    verifyDataRows(result, rows("10.03.2003"));
+  }
+
   @Test
   public void testUnixTimeStamp() throws IOException {
     var result = executeQuery(

diff --git a/sql/src/main/antlr/OpenSearchSQLLexer.g4 b/sql/src/main/antlr/OpenSearchSQLLexer.g4
@@ -213,6 +213,7 @@ EXPM1:                              'EXPM1';
 FLOOR:                              'FLOOR';
 FROM_DAYS:                          'FROM_DAYS';
 FROM_UNIXTIME:                      'FROM_UNIXTIME';
+GET_FORMAT:                         'GET_FORMAT';
 IF:                                 'IF';
 IFNULL:                             'IFNULL';
 ISNULL:                             'ISNULL';

diff --git a/sql/src/main/antlr/OpenSearchSQLParser.g4 b/sql/src/main/antlr/OpenSearchSQLParser.g4
@@ -309,6 +309,18 @@ functionCall
     | relevanceFunction                                             #relevanceFunctionCall
     | highlightFunction                                             #highlightFunctionCall
     | positionFunction                                              #positionFunctionCall
+    | getFormatFunction                                             #getFormatFunctionCall
+    ;
+
+getFormatFunction
+    : GET_FORMAT LR_BRACKET getFormatType COMMA functionArg RR_BRACKET
+    ;
+
+getFormatType
+    : DATE
+    | DATETIME
+    | TIME
+    | TIMESTAMP
     ;
 
 

diff --git a/sql/src/main/java/org/opensearch/sql/sql/parser/AstExpressionBuilder.java b/sql/src/main/java/org/opensearch/sql/sql/parser/AstExpressionBuilder.java
@@ -33,6 +33,7 @@
 import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.FilterClauseContext;
 import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.FilteredAggregationFunctionCallContext;
 import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.FunctionArgContext;
+import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.GetFormatFunctionCallContext;
 import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.HighlightFunctionCallContext;
 import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.InPredicateContext;
 import static org.opensearch.sql.sql.antlr.parser.OpenSearchSQLParser.IsNullPredicateContext;
@@ -150,6 +151,13 @@ public UnresolvedExpression visitScalarFunctionCall(ScalarFunctionCallContext ct
     return buildFunction(ctx.scalarFunctionName().getText(), ctx.functionArgs().functionArg());
   }
 
+  @Override
+  public UnresolvedExpression visitGetFormatFunctionCall(GetFormatFunctionCallContext ctx) {
+    return new Function(
+        ctx.getFormatFunction().GET_FORMAT().toString(),
+        getFormatFunctionArguments(ctx));
+  }
+
   @Override
   public UnresolvedExpression visitHighlightFunctionCall(
       HighlightFunctionCallContext ctx) {
@@ -555,6 +563,15 @@ private List<UnresolvedExpression> multiFieldRelevanceArguments(
     return builder.build();
   }
 
+  private List<UnresolvedExpression> getFormatFunctionArguments(
+      GetFormatFunctionCallContext ctx) {
+    List<UnresolvedExpression> args = Arrays.asList(
+        new Literal(ctx.getFormatFunction().getFormatType().getText(), DataType.STRING),
+        visitFunctionArg(ctx.getFormatFunction().functionArg())
+    );
+    return args;
+  }
+
   /**
    * Adds support for multi_match alternate syntax like
    * MULTI_MATCH('query'='Dale', 'fields'='*name').

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
@@ -194,6 +194,33 @@ public void can_parse_now_like_functions(String name, Boolean hasFsp, Boolean ha
     assertNotNull(parser.parse("SELECT id FROM test WHERE " + String.join(" AND ", calls)));
   }
 
+  private static Stream<Arguments> get_format_arguments() {
+    Stream.Builder<Arguments> args = Stream.builder();
+    String[] types = {"DATE", "DATETIME", "TIME", "TIMESTAMP"};
+    String[] formats = {"'USA'", "'JIS'", "'ISO'", "'EUR'", "'INTERNAL'"};
+
+    for (String type : types) {
+      for (String format : formats) {
+        args.add(Arguments.of(type, format));
+      }
+    }
+
+    return args.build();
+  }
+
+  @ParameterizedTest(name = "{0}")
+  @MethodSource("get_format_arguments")
+  public void can_parse_get_format_function(String type, String format) {
+    assertNotNull(parser.parse(String.format("SELECT GET_FORMAT(%s, %s)", type, format)));
+  }
+
+  @Test
+  public void cannot_parse_get_format_function_with_bad_arg() {
+    assertThrows(
+        SyntaxCheckException.class,
+        () -> parser.parse("GET_FORMAT(NONSENSE_ARG,'INTERNAL')"));
+  }
+
   @Test
   public void can_parse_hour_functions() {
     assertNotNull(parser.parse("SELECT hour('2022-11-18 12:23:34')"));

diff --git a/sql/src/test/java/org/opensearch/sql/sql/parser/AstExpressionBuilderTest.java b/sql/src/test/java/org/opensearch/sql/sql/parser/AstExpressionBuilderTest.java
@@ -192,6 +192,14 @@ public void canBuildFunctionCall() {
     );
   }
 
+  @Test
+  public void canBuildGetFormatFunctionCall() {
+    assertEquals(
+        function("get_format", stringLiteral("DATE"), stringLiteral("USA")),
+        buildExprAst("get_format(DATE,\"USA\")")
+    );
+  }
+
   @Test
   public void canBuildNestedFunctionCall() {
     assertEquals(