Carry over parameters when grouping from YQL

5 files changed, 80 insertions, 25 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/grouping/GroupingQueryParser.java b/container-search/src/main/java/com/yahoo/search/grouping/GroupingQueryParser.java
index 36a8b7812ba..54194221958 100644
--- a/container-search/src/main/java/com/yahoo/search/grouping/GroupingQueryParser.java
+++ b/container-search/src/main/java/com/yahoo/search/grouping/GroupingQueryParser.java
@@ -50,25 +50,14 @@ public class GroupingQueryParser extends Searcher {
     @Override
     public Result search(Query query, Execution execution) {
         try {
-            if (query.getHttpRequest().getProperty(GROUPING_GLOBAL_MAX_GROUPS.toString()) != null) {
-                throw new IllegalInputException(GROUPING_GLOBAL_MAX_GROUPS + " must be specified in a query profile.");
-            }
+            validate(query);
 
             String reqParam = query.properties().getString(PARAM_REQUEST);
             if (reqParam == null) return execution.search(query);
 
             List<Continuation> continuations = getContinuations(query.properties().getString(PARAM_CONTINUE));
-            TimeZone zone = getTimeZone(query.properties().getString(PARAM_TIMEZONE, "utc"));
-            for (GroupingOperation op : GroupingOperation.fromStringAsList(reqParam)) {
-                GroupingRequest grpRequest = GroupingRequest.newInstance(query);
-                grpRequest.setRootOperation(op);
-                grpRequest.setTimeZone(zone);
-                grpRequest.continuations().addAll(continuations);
-                intProperty(query, PARAM_DEFAULT_MAX_GROUPS).ifPresent(grpRequest::setDefaultMaxGroups);
-                intProperty(query, PARAM_DEFAULT_MAX_HITS).ifPresent(grpRequest::setDefaultMaxHits);
-                longProperty(query, GROUPING_GLOBAL_MAX_GROUPS).ifPresent(grpRequest::setGlobalMaxGroups);
-                doubleProperty(query, PARAM_DEFAULT_PRECISION_FACTOR).ifPresent(grpRequest::setDefaultPrecisionFactor);
-            }
+            for (GroupingOperation operation : GroupingOperation.fromStringAsList(reqParam))
+                createGroupingRequestIn(query, operation, continuations);
             return execution.search(query);
         }
         catch (IllegalArgumentException e) {
@@ -76,6 +65,22 @@ public class GroupingQueryParser extends Searcher {
         }
     }
 
+    public static void validate(Query query) {
+        if (query.getHttpRequest().getProperty(GROUPING_GLOBAL_MAX_GROUPS.toString()) != null)
+            throw new IllegalInputException(GROUPING_GLOBAL_MAX_GROUPS + " must be specified in a query profile.");
+    }
+
+    public static void createGroupingRequestIn(Query query, GroupingOperation operation, List<Continuation> continuations) {
+        GroupingRequest request = GroupingRequest.newInstance(query);
+        request.setRootOperation(operation);
+        request.setTimeZone(getTimeZone(query.properties().getString(PARAM_TIMEZONE, "utc")));
+        request.continuations().addAll(continuations);
+        intProperty(query, PARAM_DEFAULT_MAX_GROUPS).ifPresent(request::setDefaultMaxGroups);
+        intProperty(query, PARAM_DEFAULT_MAX_HITS).ifPresent(request::setDefaultMaxHits);
+        longProperty(query, GROUPING_GLOBAL_MAX_GROUPS).ifPresent(request::setGlobalMaxGroups);
+        doubleProperty(query, PARAM_DEFAULT_PRECISION_FACTOR).ifPresent(request::setDefaultPrecisionFactor);
+    }
+
     private List<Continuation> getContinuations(String param) {
         if (param == null) {
             return Collections.emptyList();
@@ -87,7 +92,7 @@ public class GroupingQueryParser extends Searcher {
         return ret;
     }
 
-    private TimeZone getTimeZone(String name) {
+    private static TimeZone getTimeZone(String name) {
         ZoneCache cache = zoneCache.get();
         if (cache == null) {
             cache = new ZoneCache();
diff --git a/container-search/src/main/java/com/yahoo/search/grouping/vespa/RequestBuilder.java b/container-search/src/main/java/com/yahoo/search/grouping/vespa/RequestBuilder.java
index e43067ade62..2280ea01263 100644
--- a/container-search/src/main/java/com/yahoo/search/grouping/vespa/RequestBuilder.java
+++ b/container-search/src/main/java/com/yahoo/search/grouping/vespa/RequestBuilder.java
@@ -406,8 +406,9 @@ class RequestBuilder {
     }
 
     private void validateGlobalMax() {
-        this.totalGroupsAndSummaries = -1;
         if (globalMaxGroups < 0) return;
+
+        this.totalGroupsAndSummaries = -1;
         int totalGroupsAndSummaries = 0;
         for (Grouping grp : requestList) {
             int levelMultiplier = 1;
diff --git a/container-search/src/main/java/com/yahoo/search/yql/MinimalQueryInserter.java b/container-search/src/main/java/com/yahoo/search/yql/MinimalQueryInserter.java
index 48c48748563..e844bac21e8 100644
--- a/container-search/src/main/java/com/yahoo/search/yql/MinimalQueryInserter.java
+++ b/container-search/src/main/java/com/yahoo/search/yql/MinimalQueryInserter.java
@@ -10,7 +10,7 @@ import com.yahoo.search.Query;
 import com.yahoo.search.Result;
 import com.yahoo.search.Searcher;
 import com.yahoo.processing.request.CompoundName;
-import com.yahoo.search.grouping.GroupingRequest;
+import com.yahoo.search.grouping.GroupingQueryParser;
 import com.yahoo.search.query.QueryTree;
 import com.yahoo.search.query.parser.Parsable;
 import com.yahoo.search.query.parser.ParserEnvironment;
@@ -116,11 +116,11 @@ public class MinimalQueryInserter extends Searcher {
         }
         query.getModel().getQueryTree().setRoot(newTree.getRoot());
         query.getPresentation().getSummaryFields().addAll(parser.getYqlSummaryFields());
-        for (VespaGroupingStep step : parser.getGroupingSteps()) {
-            GroupingRequest.newInstance(query)
-                    .setRootOperation(step.getOperation())
-                    .continuations().addAll(step.continuations());
-        }
+
+        GroupingQueryParser.validate(query);
+        for (VespaGroupingStep step : parser.getGroupingSteps())
+            GroupingQueryParser.createGroupingRequestIn(query, step.getOperation(), step.continuations());
+
         if (parser.getYqlSources().size() == 0) {
             query.getModel().getSources().clear();
         } else {
diff --git a/container-search/src/test/java/com/yahoo/search/grouping/vespa/RequestBuilderTestCase.java b/container-search/src/test/java/com/yahoo/search/grouping/vespa/RequestBuilderTestCase.java
index cfcbc435ba5..5882c56c688 100644
--- a/container-search/src/test/java/com/yahoo/search/grouping/vespa/RequestBuilderTestCase.java
+++ b/container-search/src/test/java/com/yahoo/search/grouping/vespa/RequestBuilderTestCase.java
@@ -784,6 +784,16 @@ public class RequestBuilderTestCase {
     }
 
     @Test
+    public void require_that_total_groups_restriction_can_be_disabled() {
+        assertTotalGroupsAndSummaries(3+3*5+3*5*7+3*5*7*100,
+                                      -1, // disable
+                                      "all( group(a) max(3) each(output(count())" +
+                                      "     all(group(b) max(5) each(output(count())" +
+                                      "         all(group(c) max(7) each(max(100) output(count())" +
+                                      "             each(output(summary()))))))))");
+    }
+
+    @Test
     public void require_that_unbounded_queries_fails_when_global_max_is_enabled() {
         assertQueryFailsOnGlobalMax(4, "all(group(a) max(5) each(output(count())))", "5 > 4");
         assertQueryFailsOnGlobalMax(Long.MAX_VALUE, "all(group(a) each(output(count())))", "unbounded number of groups");
@@ -810,10 +820,15 @@ public class RequestBuilderTestCase {
     }
 
     private static void assertTotalGroupsAndSummaries(long expected, String query) {
+        assertTotalGroupsAndSummaries(expected, Long.MAX_VALUE, query);
+    }
+
+    private static void assertTotalGroupsAndSummaries(long expected, long globalMaxGroups, String query) {
         RequestBuilder builder = new RequestBuilder(0)
-                .setRootOperation(GroupingOperation.fromString(query)).setGlobalMaxGroups(Long.MAX_VALUE);
+                .setRootOperation(GroupingOperation.fromString(query)).setGlobalMaxGroups(globalMaxGroups);
         builder.build();
-        assertEquals(expected, builder.totalGroupsAndSummaries().orElseThrow());
+        if (globalMaxGroups >= 0) // otherwise, totalGroupsAndSummaries is not computed
+            assertEquals(expected, builder.totalGroupsAndSummaries().orElseThrow());
     }
 
     private static void assertQueryFailsOnGlobalMax(long globalMax, String query, String errorSubstring) {
@@ -938,11 +953,13 @@ public class RequestBuilderTestCase {
         String expectedOutput;
         OutputWriter outputWriter;
         Continuation continuation;
+
     }
 
-    private static interface OutputWriter {
+    private interface OutputWriter {
 
         String write(List<Grouping> groupingList, GroupingTransform transform);
+
     }
 
     private static class OffsetWriter implements OutputWriter {
diff --git a/container-search/src/test/java/com/yahoo/search/yql/MinimalQueryInserterTestCase.java b/container-search/src/test/java/com/yahoo/search/yql/MinimalQueryInserterTestCase.java
index 890f847ebfe..1ef9fe5832f 100644
--- a/container-search/src/test/java/com/yahoo/search/yql/MinimalQueryInserterTestCase.java
+++ b/container-search/src/test/java/com/yahoo/search/yql/MinimalQueryInserterTestCase.java
@@ -5,6 +5,7 @@ import com.google.common.base.Charsets;
 import com.yahoo.component.chain.Chain;
 import com.yahoo.language.Language;
 import com.yahoo.language.simple.SimpleLinguistics;
+import com.yahoo.processing.IllegalInputException;
 import com.yahoo.search.Query;
 import com.yahoo.search.Result;
 import com.yahoo.search.Searcher;
@@ -29,6 +30,7 @@ import java.util.List;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 /**
  * Smoke test for first generation YQL+ integration.
@@ -350,6 +352,36 @@ public class MinimalQueryInserterTestCase {
     }
 
     @Test
+    public void globalMaxGroupsIsCarriedOver() {
+        URIBuilder builder = new URIBuilder();
+        builder.setPath("search/");
+        builder.setParameter("yql", "select foo from bar where baz contains 'cox' " +
+                                    "| all(group(a) each(output(count())))");
+        Query query = new Query(builder.toString());
+        query.properties().set("grouping.globalMaxGroups", -1);
+        execution.search(query);
+        assertEquals(1, query.getSelect().getGrouping().size());
+        assertEquals(-1L, query.getSelect().getGrouping().get(0).globalMaxGroups().getAsLong());
+    }
+
+    @Test
+    public void globalMaxGroupsCannotBeSetInRequest() {
+        try {
+            URIBuilder builder = new URIBuilder();
+            builder.setPath("search/");
+            builder.setParameter("yql", "select foo from bar where baz contains 'cox' " +
+                                        "| all(group(a) each(output(count())))");
+            builder.setParameter("grouping.globalMaxGroups", "-1");
+            Query query = new Query(builder.toString());
+            execution.search(query);
+            fail();
+        }
+        catch (IllegalInputException e) {
+            assertEquals("grouping.globalMaxGroups must be specified in a query profile.", e.getCause().getMessage());
+        }
+    }
+
+    @Test
     public void verifyThatWarmupIsSane() {
         assertTrue(MinimalQueryInserter.warmup());
     }