Remove of null column in JDBC output when script function is present (#…

…116) Fix of JDBC output when there's script function in group by. It is done by additional analysis in schema preparation step to differentiate script and method fields. * Fixed issue with unnecessary column in the JDBC result * Added unit test * Fixed test, simplified code a bit
opendistro-for-elasticsearch · Jul 16, 2019 · 09134c9 · 09134c9
1 parent ddfbcef
commit 09134c9
Show file tree

Hide file tree

Showing 5 changed files with 54 additions and 25 deletions.
diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/sql/domain/Field.java b/src/main/java/com/amazon/opendistroforelasticsearch/sql/domain/Field.java
@@ -15,6 +15,7 @@
 
 package com.amazon.opendistroforelasticsearch.sql.domain;
 
+import com.alibaba.druid.sql.ast.SQLExpr;
 import com.amazon.opendistroforelasticsearch.sql.parser.ChildrenType;
 import com.amazon.opendistroforelasticsearch.sql.parser.NestedType;
 
@@ -35,8 +36,9 @@ public class Field implements Cloneable{
 	private String alias;
     private NestedType nested;
     private ChildrenType children;
+    private SQLExpr expression;
 
-	public Field(String name, String alias) {
+    public Field(String name, String alias) {
 		this.name = name;
 		this.alias = alias;
         this.nested = null;
@@ -127,4 +129,12 @@ protected Object clone() throws CloneNotSupportedException {
     public boolean isScriptField() {
         return false;
     }
+
+    public void setExpression(SQLExpr expression) {
+        this.expression = expression;
+    }
+
+    public SQLExpr getExpression() {
+        return expression;
+    }
 }
diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/sql/executor/format/SelectResultSet.java b/src/main/java/com/amazon/opendistroforelasticsearch/sql/executor/format/SelectResultSet.java
@@ -226,23 +226,23 @@ private boolean containsWildcard(Query query) {
      */
     private List<Field> fetchFields(Query query) {
         Select select = (Select) query;
-        List<Field> fields;
+
         if (queryResult instanceof Aggregations) {
-            fields = select.getGroupBys().isEmpty() ? new ArrayList<>() : select.getGroupBys().get(0);
-            for (Field field : select.getFields()) {
-                if (field instanceof MethodField) {
-                    fields.add(field);
+            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 (query instanceof TableOnJoinSelect) {
-                fields = ((TableOnJoinSelect) query).getSelectedFields();
-            } else {
-                fields = select.getFields();
-            }
+            return groupByFields;
         }
 
-        return fields;
+        if (query instanceof TableOnJoinSelect) {
+            return ((TableOnJoinSelect) query).getSelectedFields();
+        }
+
+        return select.getFields();
     }
 
     private String[] fetchFieldsAsArray(Query query) {

diff --git a/src/main/java/com/amazon/opendistroforelasticsearch/sql/parser/FieldMaker.java b/src/main/java/com/amazon/opendistroforelasticsearch/sql/parser/FieldMaker.java
@@ -38,13 +38,24 @@
  */
 public class FieldMaker {
     public static Field makeField(SQLExpr expr, String alias, String tableAlias) throws SqlParseException {
+        Field field = makeFieldImpl(expr, alias, tableAlias);
+
+        // why we may get null as a field???
+        if (field != null) {
+            field.setExpression(expr);
+        }
+
+        return field;
+    }
+
+    private static Field makeFieldImpl(SQLExpr expr, String alias, String tableAlias) throws SqlParseException {
         if (expr instanceof SQLIdentifierExpr || expr instanceof SQLPropertyExpr || expr instanceof SQLVariantRefExpr) {
             return handleIdentifier(expr, alias, tableAlias);
         } else if (expr instanceof SQLQueryExpr) {
             throw new SqlParseException("unknow field name : " + expr);
         } else if (expr instanceof SQLBinaryOpExpr) {
             //make a SCRIPT method field;
-            return makeField(makeBinaryMethodField((SQLBinaryOpExpr) expr, alias, true), alias, tableAlias);
+            return makeFieldImpl(makeBinaryMethodField((SQLBinaryOpExpr) expr, alias, true), alias, tableAlias);
 
         } else if (expr instanceof SQLAllColumnExpr) {
             return Field.STAR;

diff --git a/src/test/java/com/amazon/opendistroforelasticsearch/sql/esintgtest/JdbcTestIT.java b/src/test/java/com/amazon/opendistroforelasticsearch/sql/esintgtest/JdbcTestIT.java
@@ -27,31 +27,39 @@ protected void init() throws Exception {
     }
 
     public void testDateTimeInQuery() {
-        String query = "SELECT date_format(insert_time, 'dd-MM-YYYY') FROM elasticsearch-sql_test_index_online limit 1";
-
-        JSONObject parsed = new JSONObject(executeJdbcRequest(query));
+        JSONObject response = executeJdbcRequest(
+                "SELECT date_format(insert_time, 'dd-MM-YYYY') FROM elasticsearch-sql_test_index_online limit 1");
 
         assertThat(
-                parsed.getJSONArray("datarows")
+                response.getJSONArray("datarows")
                         .getJSONArray(0)
                         .getString(0),
                 equalTo("17-08-2014"));
     }
 
     public void testDivisionInQuery() {
-        String query = "SELECT all_client/10 from elasticsearch-sql_test_index_online ORDER BY all_client/10 desc limit 1";
-
-        JSONObject parsed = new JSONObject(executeJdbcRequest(query));
+        JSONObject response = executeJdbcRequest(
+                "SELECT all_client/10 from elasticsearch-sql_test_index_online ORDER BY all_client/10 desc limit 1");
 
         assertThat(
-                parsed.getJSONArray("datarows")
+                response.getJSONArray("datarows")
                         .getJSONArray(0)
                         .getDouble(0),
                 equalTo(16827.0));
     }
 
+    public void testGroupByInQuery() {
+        JSONObject response = executeJdbcRequest(
+                "SELECT date_format(insert_time, 'YYYY-MM-dd'), COUNT(*) " +
+                        "FROM elasticsearch-sql_test_index_online " +
+                        "GROUP BY date_format(insert_time, 'YYYY-MM-dd')"
+        );
+
+        assertThat(response.getJSONArray("schema").length(), equalTo(2));
+        assertThat(response.getJSONArray("datarows").length(), equalTo(8));
+    }
 
-    private String executeJdbcRequest(String query){
-        return executeQuery(query, "jdbc");
+    private JSONObject executeJdbcRequest(String query){
+        return new JSONObject(executeQuery(query, "jdbc"));
     }
 }
diff --git a/src/test/java/com/amazon/opendistroforelasticsearch/sql/unittest/DateFormatTest.java b/src/test/java/com/amazon/opendistroforelasticsearch/sql/unittest/DateFormatTest.java
@@ -91,7 +91,7 @@ public void orderByTest() {
     @Ignore("06/27/2019: During implementing of ORDER BY date_format found that this fails as well. " +
             "Will investigate after order by fix submitted, to scope amount of fixes.")
     public void orderByWithGroupByTest() {
-        String query = "SELECT ip, count(ip) " +
+        String query = "SELECT date_format(utc_time, 'dd-MM-YYYY'), count(*) " +
                 "FROM kibana_sample_data_logs " +
                 "GROUP BY date_format(utc_time, 'dd-MM-YYYY') " +
                 "ORDER BY date_format(utc_time, 'dd-MM-YYYY') DESC";