Support Cast() function (#253)

Added support for Cast() function- re-added CAST syntax into grammar file, and added functionality for JDBC format as well.

src/main/antlr/OpenDistroSqlLexer.g4

@@ -49,24 +49,30 @@ ASC:                                'ASC';
 BETWEEN:                            'BETWEEN';
 BY:                                 'BY';
 CASE:                               'CASE';
+CAST:                               'CAST';
 CROSS:                              'CROSS';
+DATETIME:                           'DATETIME';
 DELETE:                             'DELETE';
 DESC:                               'DESC';
 DESCRIBE:                           'DESCRIBE';
 DISTINCT:                           'DISTINCT';
+DOUBLE:                             'DOUBLE';
 ELSE:                               'ELSE';
 EXISTS:                             'EXISTS';
 FALSE:                              'FALSE';
+FLOAT:                              'FLOAT';
 FROM:                               'FROM';
 GROUP:                              'GROUP';
 HAVING:                             'HAVING';
 IN:                                 'IN';
 INNER:                              'INNER';
+INT:                                'INT';
 IS:                                 'IS';
 JOIN:                               'JOIN';
 LEFT:                               'LEFT';
 LIKE:                               'LIKE';
 LIMIT:                              'LIMIT';
+LONG:                               'LONG';
 MATCH:                              'MATCH';
 NATURAL:                            'NATURAL';
 NOT:                                'NOT';
@@ -79,6 +85,7 @@ REGEXP:                             'REGEXP';
 RIGHT:                              'RIGHT';
 SELECT:                             'SELECT';
 SHOW:                               'SHOW';
+STRING:                             'STRING';
 THEN:                               'THEN';
 TRUE:                               'TRUE';
 UNION:                              'UNION';

src/main/antlr/OpenDistroSqlParser.g4

@@ -241,7 +241,6 @@ dottedId
     | '.' uid
     ;
 
-
 //    Literals
 
 decimalLiteral
@@ -297,7 +296,8 @@ functionCall
     ;
 
 specificFunction
-    : CASE expression caseFuncAlternative+
+    : CAST '(' expression AS convertedDataType ')'                  #dataTypeFunctionCall
+    | CASE expression caseFuncAlternative+
       (ELSE elseArg=functionArg)? END                               #caseFunctionCall
     | CASE caseFuncAlternative+
       (ELSE elseArg=functionArg)? END                               #caseFunctionCall
@@ -308,6 +308,15 @@ caseFuncAlternative
       THEN consequent=functionArg
     ;
 
+convertedDataType
+    : typeName=DATETIME
+    | typeName=INT
+    | typeName=DOUBLE
+    | typeName=LONG
+    | typeName=FLOAT
+    | typeName=STRING
+    ;
+
 aggregateWindowedFunction
     : (AVG | MAX | MIN | SUM)
       '(' aggregator=(ALL | DISTINCT)? functionArg ')'

src/main/java/com/amazon/opendistroforelasticsearch/sql/executor/format/SelectResultSet.java

@@ -15,6 +15,7 @@
 
 package com.amazon.opendistroforelasticsearch.sql.executor.format;
 
+import com.alibaba.druid.sql.ast.expr.SQLCastExpr;
 import com.amazon.opendistroforelasticsearch.sql.domain.Field;
 import com.amazon.opendistroforelasticsearch.sql.domain.JoinSelect;
 import com.amazon.opendistroforelasticsearch.sql.domain.MethodField;
@@ -124,6 +125,9 @@ private void loadFromEsState(Query query) {
         // Assumption is all indices share the same mapping which is validated in TermFieldRewriter.
         Map<String, Map<String, FieldMappingMetaData>> indexMappings = mappings.values().iterator().next();
 
+        // if index mappings size is 0 and the expression is a cast: that means that we are casting by alias
+        // if so, add the original field that was being looked at to the mapping (how?)
+
         /*
          * There are three cases regarding type name to consider:
          * 1. If the correct type name was given, its typeMapping is retrieved
@@ -245,9 +249,13 @@ private List<Field> fetchFields(Query query) {
             List<Field> groupByFields = select.getGroupBys().isEmpty() ? new ArrayList<>() :
                     select.getGroupBys().get(0);
 
+
             for (Field selectField : select.getFields()) {
                 if (selectField instanceof MethodField && !selectField.isScriptField()) {
                     groupByFields.add(selectField);
+                } else if (selectField.isScriptField()
+                        && selectField.getAlias().equals(groupByFields.get(0).getName())) {
+                    return select.getFields();
                 }
             }
             return groupByFields;
@@ -315,7 +323,10 @@ private Schema.Type fetchMethodReturnType(Field field) {
                 // TODO: return type information is disconnected from the function definitions in SQLFunctions.
                 // Refactor SQLFunctions to have functions self-explanatory (types, scripts) and pluggable
                 // (similar to Strategy pattern)
-
+                if (field.getExpression() instanceof SQLCastExpr) {
+                    return SQLFunctions.getCastFunctionReturnType(
+                            ((SQLCastExpr) field.getExpression()).getDataType().getName());
+                }
                 return SQLFunctions.getScriptFunctionReturnType(
                         ((ScriptMethodField) field).getFunctionName());
             }

src/main/java/com/amazon/opendistroforelasticsearch/sql/query/DefaultQueryAction.java

@@ -18,6 +18,7 @@
 import com.alibaba.druid.sql.ast.SQLExpr;
 import com.alibaba.druid.sql.ast.expr.SQLBinaryOpExpr;
 import com.alibaba.druid.sql.ast.expr.SQLBinaryOperator;
+import com.alibaba.druid.sql.ast.expr.SQLCastExpr;
 import com.amazon.opendistroforelasticsearch.sql.domain.Field;
 import com.amazon.opendistroforelasticsearch.sql.domain.KVValue;
 import com.amazon.opendistroforelasticsearch.sql.domain.MethodField;
@@ -95,7 +96,6 @@ public SqlElasticSearchRequestBuilder explain() throws SqlParseException {
         setIndicesAndTypes();
 
         setFields(select.getFields());
-
         setWhere(select.getWhere());
         setSorts(select.getOrderBys());
         setLimit(select.getOffset(), select.getRowCount());
@@ -152,6 +152,9 @@ public void setFields(List<Field> fields) throws SqlParseException {
                     MethodField method = (MethodField) field;
                     if (method.getName().toLowerCase().equals("script")) {
                         handleScriptField(method);
+                        if (method.getExpression() instanceof SQLCastExpr) {
+                            includeFields.add(method.getParams().get(0).toString());
+                        }
                     } else if (method.getName().equalsIgnoreCase("include")) {
                         for (KVValue kvValue : method.getParams()) {
                             includeFields.add(kvValue.value.toString());
@@ -265,8 +268,15 @@ private String getNullOrderString(SQLBinaryOpExpr expr) {
 
     private ScriptSortType getScriptSortType(Order order) {
         ScriptSortType scriptSortType;
-        ScriptMethodField smf = (ScriptMethodField) order.getSortField();
-        Schema.Type scriptFunctionReturnType = SQLFunctions.getScriptFunctionReturnType(smf.getFunctionName());
+        Schema.Type scriptFunctionReturnType;
+        if (order.getSortField().getExpression() instanceof SQLCastExpr) {
+            scriptFunctionReturnType = SQLFunctions.getCastFunctionReturnType(
+                    ((SQLCastExpr) order.getSortField().getExpression()).getDataType().getName());
+        } else {
+            ScriptMethodField smf = (ScriptMethodField) order.getSortField();
+            scriptFunctionReturnType = SQLFunctions.getScriptFunctionReturnType(smf.getFunctionName());
+        }
+
 
         // as of now script function return type returns only text and double
         switch (scriptFunctionReturnType) {
@@ -275,9 +285,11 @@ private ScriptSortType getScriptSortType(Order order) {
                 break;
 
             case DOUBLE:
+            case FLOAT:
+            case INTEGER:
+            case LONG:
                 scriptSortType = ScriptSortType.NUMBER;
                 break;
-
             default:
                 throw new RuntimeException("unknown");
         }

src/main/java/com/amazon/opendistroforelasticsearch/sql/utils/SQLFunctions.java

@@ -739,4 +739,25 @@ public static Schema.Type getScriptFunctionReturnType(String functionName) {
                         "The following method is not supported in Schema: %s",
                         functionName));
     }
+
+    public static Schema.Type getCastFunctionReturnType(String castType) {
+        switch (StringUtils.toUpper(castType)) {
+            case "FLOAT":
+                return Schema.Type.FLOAT;
+            case "DOUBLE":
+                return Schema.Type.DOUBLE;
+            case "INT":
+                return Schema.Type.INTEGER;
+            case "STRING":
+                return Schema.Type.TEXT;
+            case "DATETIME":
+                return Schema.Type.DATE;
+            case "LONG":
+                return Schema.Type.LONG;
+            default:
+                throw new UnsupportedOperationException(
+                    StringUtils.format("The following type is not supported by cast(): %s", castType)
+                );
+        }
+    }
 }