Return Correct Type Information for Fields

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

@@ -66,19 +66,19 @@ public class SelectResultSet extends ResultSet {
     private String indexName;
     private String typeName;
     private List<Schema.Column> columns = new ArrayList<>();
-    private ColumnTypeProvider scriptColumnType;
+    private ColumnTypeProvider outputColumnType;
 
     private List<String> head;
     private long size;
     private long totalHits;
     private List<DataRows.Row> rows;
 
-    public SelectResultSet(Client client, Query query, Object queryResult, ColumnTypeProvider scriptColumnType) {
+    public SelectResultSet(Client client, Query query, Object queryResult, ColumnTypeProvider outputColumnType) {
         this.client = client;
         this.query = query;
         this.queryResult = queryResult;
         this.selectAll = false;
-        this.scriptColumnType = scriptColumnType;
+        this.outputColumnType = outputColumnType;
 
         if (isJoinQuery()) {
             JoinSelect joinQuery = (JoinSelect) query;
@@ -328,7 +328,8 @@ private Schema.Type fetchMethodReturnType(int fieldIndex, MethodField field) {
                 if (field.getExpression() instanceof SQLCaseExpr) {
                     return Schema.Type.TEXT;
                 }
-                return SQLFunctions.getScriptFunctionReturnType(fieldIndex, field, scriptColumnType);
+                Schema.Type resolvedType = outputColumnType.get(fieldIndex);
+                return SQLFunctions.getScriptFunctionReturnType(field, resolvedType);
             }
             default:
                 throw new UnsupportedOperationException(
@@ -377,7 +378,6 @@ private List<Schema.Column> populateColumns(Query query, String[] fieldNames, Ma
              * name instead.
              */
             if (fieldMap.get(fieldName) instanceof MethodField) {
-
                 MethodField methodField = (MethodField) fieldMap.get(fieldName);
                 int fieldIndex = fieldNameList.indexOf(fieldName);
                 columns.add(

src/main/java/com/amazon/opendistroforelasticsearch/sql/plugin/RestSqlAction.java

@@ -210,7 +210,7 @@ private boolean isSQLFeatureEnabled() {
         return allowExplicitIndex && isSqlEnabled;
     }
 
-    public static ColumnTypeProvider performAnalysis(String sql) {
+    private static ColumnTypeProvider performAnalysis(String sql) {
         LocalClusterState clusterState = LocalClusterState.state();
         SqlAnalysisConfig config = new SqlAnalysisConfig(
             clusterState.getSettingValue(QUERY_ANALYSIS_ENABLED),

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

@@ -26,7 +26,6 @@
 import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
 import com.alibaba.druid.sql.ast.expr.SQLTextLiteralExpr;
 import com.alibaba.druid.sql.ast.expr.SQLVariantRefExpr;
-import com.amazon.opendistroforelasticsearch.sql.domain.ColumnTypeProvider;
 import com.amazon.opendistroforelasticsearch.sql.domain.Field;
 import com.amazon.opendistroforelasticsearch.sql.domain.KVValue;
 import com.amazon.opendistroforelasticsearch.sql.domain.MethodField;
@@ -974,25 +973,14 @@ public String getCastScriptStatement(String name, String castType, List<KVValue>
      * approach will return type of result column as DOUBLE, although there is enough information to understand that
      * it might be safely treated as INTEGER.
      */
-    public static Schema.Type getScriptFunctionReturnType(
-            int fieldIndex, MethodField field, ColumnTypeProvider scriptColumnType) {
+    public static Schema.Type getScriptFunctionReturnType(MethodField field, Schema.Type resolvedType) {
         Schema.Type returnType;
         String functionName = ((ScriptMethodField) field).getFunctionName().toLowerCase();
-        if (!numberOperators.contains(functionName) && !mathConstants.contains(functionName)
-                && !trigFunctions.contains(functionName) && !stringOperators.contains(functionName)
-                && !stringFunctions.contains(functionName) && !binaryOperators.contains(functionName)
-                && !dateFunctions.contains(functionName) && !conditionalFunctions.contains(functionName)
-                && !utilityFunctions.contains(functionName)) {
-            throw new UnsupportedOperationException(
-                    String.format(
-                            "The following method is not supported in Schema: %s",
-                            functionName));
-        }
         if (functionName.equals("cast")) {
             String castType = ((SQLCastExpr) field.getExpression()).getDataType().getName();
             return getCastFunctionReturnType(castType);
         } else {
-            returnType = scriptColumnType.get(fieldIndex);
+            returnType = resolvedType;
         }
         return returnType;
     }

src/test/java/com/amazon/opendistroforelasticsearch/sql/unittest/utils/SQLFunctionsTest.java

@@ -34,6 +34,7 @@
 import org.elasticsearch.common.collect.Tuple;
 import org.junit.Assert;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
@@ -45,18 +46,11 @@
 
 public class SQLFunctionsTest {
 
-    private FieldMaker fieldMaker;
     private SQLFunctions sqlFunctions;
 
     @Rule
     public ExpectedException exceptionRule = ExpectedException.none();
 
-    @Before
-    public void init() {
-        fieldMaker = new FieldMaker();
-    }
-    public void initSqlFunctions() { sqlFunctions = new SQLFunctions(); }
-
     @Test
     public void testAssign() throws SqlParseException {
         SQLFunctions sqlFunctions = new SQLFunctions();
@@ -83,7 +77,8 @@ public void testAbsWithIntReturnType() {
         field.setExpression(invokeExpr);
         ColumnTypeProvider columnTypeProvider = new ColumnTypeProvider(ESDataType.INTEGER);
 
-        final Schema.Type returnType = sqlFunctions.getScriptFunctionReturnType(0, field, columnTypeProvider);
+        Schema.Type resolvedType = columnTypeProvider.get(0);
+        final Schema.Type returnType = sqlFunctions.getScriptFunctionReturnType(field, resolvedType);
         Assert.assertEquals(returnType, Schema.Type.INTEGER);
     }
 
@@ -100,18 +95,8 @@ public void testCastReturnType() {
         field.setExpression(castExpr);
         ColumnTypeProvider columnTypeProvider = new ColumnTypeProvider(ESDataType.INTEGER);
 
-        final Schema.Type returnType = sqlFunctions.getScriptFunctionReturnType(0, field, columnTypeProvider);
+        Schema.Type resolvedType = columnTypeProvider.get(0);
+        final Schema.Type returnType = sqlFunctions.getScriptFunctionReturnType(field, resolvedType);
         Assert.assertEquals(returnType, Schema.Type.INTEGER);
     }
-
-    @Test
-    public void testNullScriptReturnTypeThrowsException() throws UnsupportedOperationException {
-        exceptionRule.expect(UnsupportedOperationException.class);
-        exceptionRule.expectMessage("The following method is not supported in Schema: lo");
-
-        List<KVValue> params = new ArrayList<>();
-        MethodField field = new ScriptMethodField("LO", params, null, null);
-        ColumnTypeProvider columnTypeProvider = new ColumnTypeProvider(ESDataType.INTEGER);
-        final Schema.Type type = sqlFunctions.getScriptFunctionReturnType(0, field, columnTypeProvider);
-    }
 }