Merge pull request NLPchina#152 from ottboy4/elastic2.1.1

eliranmoyal · eliranmoyal · commit 773164e143f6 · 2016-02-05T19:47:00.000+02:00
When doing an aggregation with docs, added the ability to set the 'from'
diff --git a/src/main/java/org/nlpcn/es4sql/domain/hints/HintFactory.java b/src/main/java/org/nlpcn/es4sql/domain/hints/HintFactory.java
@@ -51,8 +51,11 @@ public static Hint getHintFromString(String hintAsString){
         if(hintAsString.startsWith("! DOCS_WITH_AGGREGATION")) {
             String[] number = getParamsFromHint(hintAsString,"! DOCS_WITH_AGGREGATION");
             //todo: check if numbers etc..
-            int docsWithAggregation = Integer.parseInt(number[0]);
-            return new Hint(HintType.DOCS_WITH_AGGREGATION,new Object[]{docsWithAggregation});
+            Integer[] params = new Integer[number.length];
+            for (int i = 0; i < params.length; i++) {
+                params[i] = Integer.parseInt(number[i]);
+            }
+            return new Hint(HintType.DOCS_WITH_AGGREGATION, params);
         }
 
         return null;
diff --git a/src/main/java/org/nlpcn/es4sql/query/AggregationQueryAction.java b/src/main/java/org/nlpcn/es4sql/query/AggregationQueryAction.java
@@ -139,8 +139,8 @@ public SqlElasticSearchRequestBuilder explain() throws SqlParseException {
 				}
 			}
 		}
-        
-		setLimit(getLimitFromHint());
+
+		setLimitFromHint(this.select.getHints());
 
 		request.setSearchType(SearchType.DEFAULT);
         updateWithIndicesOptionsIfNeeded(select,request);
@@ -266,20 +266,26 @@ private void setIndicesAndTypes() {
 		}
 	}
 
-	private void setLimit( int size) {
-		request.setFrom(0);
-
-		if (size > -1) {
-			request.setSize(size);
-		}
-	}
-
-    public int getLimitFromHint() {
-        for(Hint hint : this.select.getHints()){
-            if(hint.getType() == HintType.DOCS_WITH_AGGREGATION){
-                return (int) hint.getParams()[0];
+    private void setLimitFromHint(List<Hint> hints) {
+        int from = 0;
+        int size = 0;
+        for (Hint hint : hints) {
+            if (hint.getType() == HintType.DOCS_WITH_AGGREGATION) {
+                Integer[] params = (Integer[]) hint.getParams();
+                if (params.length > 1) {
+                    // if 2 or more are given, use the first as the from and the second as the size
+                    // so it is the same as LIMIT from,size
+                    // except written as /*! DOCS_WITH_AGGREGATION(from,size) */
+                    from = params[0];
+                    size = params[1];
+                } else if (params.length == 1) {
+                    // if only 1 parameter is given, use it as the size with a from of 0
+                    size = params[0];
+                }
+                break;
             }
         }
-        return 0;
+        request.setFrom(from);
+        request.setSize(size);
     }
 }
diff --git a/src/test/java/org/nlpcn/es4sql/AggregationTest.java b/src/test/java/org/nlpcn/es4sql/AggregationTest.java
@@ -363,6 +363,29 @@ private SqlElasticSearchRequestBuilder getSearchRequestBuilder(String query) thr
         return (SqlElasticSearchRequestBuilder) searchDao.explain(query).explain();
     }
 
+    @Test
+    public void testFromSizeWithAggregations() throws Exception {
+        final String query1 = String.format("SELECT /*! DOCS_WITH_AGGREGATION(0,1) */" +
+                " account_number FROM %s/account GROUP BY gender", TEST_INDEX);
+        SearchResponse response1 = (SearchResponse) getSearchRequestBuilder(query1).get();
+
+        Assert.assertEquals(1, response1.getHits().getHits().length);
+        Terms gender1 = response1.getAggregations().get("gender");
+        Assert.assertEquals(2, gender1.getBuckets().size());
+        Object account1 = response1.getHits().getHits()[0].getSource().get("account_number");
+
+        final String query2 = String.format("SELECT /*! DOCS_WITH_AGGREGATION(1,1) */" +
+                " account_number FROM %s/account GROUP BY gender", TEST_INDEX);
+        SearchResponse response2 = (SearchResponse) getSearchRequestBuilder(query2).get();
+
+        Assert.assertEquals(1, response2.getHits().getHits().length);
+        Terms gender2 = response2.getAggregations().get("gender");
+        Assert.assertEquals(2, gender2.getBuckets().size());
+        Object account2 = response2.getHits().getHits()[0].getSource().get("account_number");
+
+        Assert.assertEquals(response1.getHits().getTotalHits(), response2.getHits().getTotalHits());
+        Assert.assertNotEquals(account1, account2);
+    }
 
     @Test
 	public void testSubAggregations() throws  Exception {

Original file line number	Diff line number	Diff line change
`@@ -139,8 +139,8 @@ public SqlElasticSearchRequestBuilder explain() throws SqlParseException {`
`139`	`139`	`}`
`140`	`140`	`}`
`141`	`141`	`}`
`142`		`-`
`143`		`- setLimit(getLimitFromHint());`
	`142`	`+`
	`143`	`+ setLimitFromHint(this.select.getHints());`
`144`	`144`
`145`	`145`	`request.setSearchType(SearchType.DEFAULT);`
`146`	`146`	`updateWithIndicesOptionsIfNeeded(select,request);`
`@@ -266,20 +266,26 @@ private void setIndicesAndTypes() {`
`266`	`266`	`}`
`267`	`267`	`}`
`268`	`268`
`269`		`- private void setLimit( int size) {`
`270`		`- request.setFrom(0);`
`271`		`-`
`272`		`- if (size > -1) {`
`273`		`- request.setSize(size);`
`274`		`- }`
`275`		`- }`
`276`		`-`
`277`		`- public int getLimitFromHint() {`
`278`		`- for(Hint hint : this.select.getHints()){`
`279`		`- if(hint.getType() == HintType.DOCS_WITH_AGGREGATION){`
`280`		`- return (int) hint.getParams()[0];`
	`269`	`+ private void setLimitFromHint(List<Hint> hints) {`
	`270`	`+ int from = 0;`
	`271`	`+ int size = 0;`
	`272`	`+ for (Hint hint : hints) {`
	`273`	`+ if (hint.getType() == HintType.DOCS_WITH_AGGREGATION) {`
	`274`	`+ Integer[] params = (Integer[]) hint.getParams();`
	`275`	`+ if (params.length > 1) {`
	`276`	`+ // if 2 or more are given, use the first as the from and the second as the size`
	`277`	`+ // so it is the same as LIMIT from,size`
	`278`	`+ // except written as /! DOCS_WITH_AGGREGATION(from,size) /`
	`279`	`+ from = params[0];`
	`280`	`+ size = params[1];`
	`281`	`+ } else if (params.length == 1) {`
	`282`	`+ // if only 1 parameter is given, use it as the size with a from of 0`
	`283`	`+ size = params[0];`
	`284`	`+ }`
	`285`	`+ break;`
`281`	`286`	`}`
`282`	`287`	`}`
`283`		`- return 0;`
	`288`	`+ request.setFrom(from);`
	`289`	`+ request.setSize(size);`
`284`	`290`	`}`
`285`	`291`	`}`