manual merge of NLPchina#152 . too many changes.

eliranmoyal · eliranmoyal · commit cdbaeb642c97 · 2016-03-18T10:59:39.000+02:00
diff --git a/src/main/java/org/nlpcn/es4sql/domain/hints/HintFactory.java b/src/main/java/org/nlpcn/es4sql/domain/hints/HintFactory.java
@@ -58,8 +58,11 @@ public static Hint getHintFromString(String hintAsString) throws SqlParseExcepti
         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);
         }
         if(hintAsString.startsWith("! ROUTINGS")) {
             String[] routings = getParamsFromHint(hintAsString,"! ROUTINGS");
diff --git a/src/main/java/org/nlpcn/es4sql/query/AggregationQueryAction.java b/src/main/java/org/nlpcn/es4sql/query/AggregationQueryAction.java
@@ -137,7 +137,8 @@ public SqlElasticSearchRequestBuilder explain() throws SqlParseException {
 			}
 		}
         
-		setLimit(getLimitFromHint());
+
+        setLimitFromHint(this.select.getHints());
 
 		request.setSearchType(SearchType.DEFAULT);
         updateRequestWithIndexAndRoutingOptions(select, request);
@@ -263,20 +264,27 @@ 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
@@ -373,6 +373,30 @@ private SqlElasticSearchRequestBuilder getSearchRequestBuilder(String query) thr
     }
 
 
+    @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 {
 		Set expectedAges = new HashSet<>(ContiguousSet.create(Range.closed(20, 40), DiscreteDomain.integers()));
