Merge pull request NLPchina#310 from jasonliaoxiaoge/master

allwefantasy · web-flow · commit d02f32c40960 · 2016-10-20T21:22:08.000+08:00
fix bug in aggregation query when select has script
diff --git a/src/main/java/org/nlpcn/es4sql/query/AggregationQueryAction.java b/src/main/java/org/nlpcn/es4sql/query/AggregationQueryAction.java
@@ -57,22 +57,7 @@ public SqlElasticSearchRequestBuilder explain() throws SqlParseException {
 
 
                 //make groupby can reference to field alias
-                boolean refrence = false;
-                for (Field temp : select.getFields()) {
-                    if (temp instanceof MethodField && temp.getName().equals("script")) {
-                        MethodField scriptField = (MethodField) temp;
-                        for (KVValue kv : scriptField.getParams()) {
-                            if (kv.value.equals(field.getName())) {
-                                lastAgg = aggMaker.makeGroupAgg(scriptField);
-                                refrence = true;
-                                break;
-                            }
-                        }
-                    }
-                }
-
-                if (!refrence) lastAgg = aggMaker.makeGroupAgg(field);
-
+                lastAgg = getGroupAgg(field, select);
 
                 if (lastAgg != null && lastAgg instanceof TermsBuilder && !(field instanceof MethodField)) {
                     //if limit size is too small, increasing shard  size is required
@@ -115,7 +100,7 @@ public SqlElasticSearchRequestBuilder explain() throws SqlParseException {
 
                 for (int i = 1; i < groupBy.size(); i++) {
                     field = groupBy.get(i);
-                    AggregationBuilder<?> subAgg = aggMaker.makeGroupAgg(field);
+                    AggregationBuilder<?> subAgg = getGroupAgg(field, select);
                     if (subAgg instanceof TermsBuilder && !(field instanceof MethodField)) {
                         ((TermsBuilder) subAgg).size(0);
                     }
@@ -193,6 +178,26 @@ public SqlElasticSearchRequestBuilder explain() throws SqlParseException {
         SqlElasticSearchRequestBuilder sqlElasticRequestBuilder = new SqlElasticSearchRequestBuilder(request);
         return sqlElasticRequestBuilder;
     }
+	private AggregationBuilder<?> getGroupAgg(Field field, Select select2) throws SqlParseException {
+    	boolean refrence = false;
+    	AggregationBuilder<?> lastAgg = null;
+        for (Field temp : select.getFields()) {
+            if (temp instanceof MethodField && temp.getName().equals("script")) {
+                MethodField scriptField = (MethodField) temp;
+                for (KVValue kv : scriptField.getParams()) {
+                    if (kv.value.equals(field.getName())) {
+                        lastAgg = aggMaker.makeGroupAgg(scriptField);
+                        refrence = true;
+                        break;
+                    }
+                }
+            }
+        }
+
+        if (!refrence) lastAgg = aggMaker.makeGroupAgg(field);
+        
+        return lastAgg;
+	}
 
     private AbstractAggregationBuilder wrapNestedIfNeeded(AggregationBuilder nestedBuilder, boolean reverseNested) {
         if (!reverseNested) return nestedBuilder;
diff --git a/src/test/java/org/nlpcn/es4sql/ExplainTest.java b/src/test/java/org/nlpcn/es4sql/ExplainTest.java
@@ -28,6 +28,14 @@ public void searchSanity() throws IOException, SqlParseException, NoSuchMethodEx
 
 		assertThat(result, equalTo(expectedOutput));
 	}
+	
+    @Test
+    public void aggregationQuery() throws IOException, SqlParseException, NoSuchMethodException, IllegalAccessException, SQLFeatureNotSupportedException, InvocationTargetException {
+        String expectedOutput = Files.toString(new File("src/test/resources/expectedOutput/aggregation_query_explain.json"), StandardCharsets.UTF_8).replaceAll("\r","");
+        String result = explain(String.format("SELECT a, CASE WHEN gender='0' then 'aaa' else 'bbb'end a2345,count(c) FROM %s GROUP BY terms('field'='a'),a2345", TEST_INDEX));
+
+        assertThat(result, equalTo(expectedOutput));
+    }
 
 	@Test
 	public void searchSanityFilter() throws IOException, SqlParseException, NoSuchMethodException, IllegalAccessException, SQLFeatureNotSupportedException, InvocationTargetException {
diff --git a/src/test/resources/expectedOutput/aggregation_query_explain.json b/src/test/resources/expectedOutput/aggregation_query_explain.json
@@ -0,0 +1,40 @@
+{
+  "from" : 0,
+  "size" : 0,
+  "_source" : {
+    "includes" : [ ],
+    "excludes" : [ ]
+  },
+  "fields" : [ "a", "a2345" ],
+  "script_fields" : {
+    "a2345" : {
+      "script" : {
+        "inline" : "if( (doc['gender'].value=='0')){'aaa'} else {'bbb'}"
+      }
+    }
+  },
+  "aggregations" : {
+    "terms(field=a)" : {
+      "terms" : {
+        "field" : "a"
+      },
+      "aggregations" : {
+        "a2345" : {
+          "terms" : {
+            "script" : {
+              "inline" : "if( (doc['gender'].value=='0')){'aaa'} else {'bbb'}"
+            },
+            "size" : 0
+          },
+          "aggregations" : {
+            "COUNT(c)" : {
+              "value_count" : {
+                "field" : "c"
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}
