Merge pull request #152 from ottboy4/elastic2.1.1

When doing an aggregation with docs, added the ability to set the 'from'

Author: eliranmoyal

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;

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);
     }
 }

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 {