Merge pull request #21748 from vespa-engine/bjorncs/default-max-grouping

Introduce query parameters for default max hits/groups in grouping qu…

7 files changed, 93 insertions, 9 deletions

diff --git a/annotations/src/main/java/com/yahoo/api/annotations/Beta.java b/annotations/src/main/java/com/yahoo/api/annotations/Beta.java
index c555fbbcb57..4d15bb377d3 100644
--- a/annotations/src/main/java/com/yahoo/api/annotations/Beta.java
+++ b/annotations/src/main/java/com/yahoo/api/annotations/Beta.java
@@ -18,6 +18,7 @@ import java.lang.annotation.Target;
 @Target({
         ElementType.CONSTRUCTOR,
         ElementType.METHOD,
-        ElementType.TYPE
+        ElementType.TYPE,
+        ElementType.FIELD
 })
 public @interface Beta {}
diff --git a/container-search/abi-spec.json b/container-search/abi-spec.json
index 9b560e86d86..8ac5aaa127d 100644
--- a/container-search/abi-spec.json
+++ b/container-search/abi-spec.json
@@ -2318,7 +2318,9 @@
       "public static final java.lang.String SELECT_PARAMETER_PARSING",
       "public static final com.yahoo.processing.request.CompoundName PARAM_CONTINUE",
       "public static final com.yahoo.processing.request.CompoundName PARAM_REQUEST",
-      "public static final com.yahoo.processing.request.CompoundName PARAM_TIMEZONE"
+      "public static final com.yahoo.processing.request.CompoundName PARAM_TIMEZONE",
+      "public static final com.yahoo.processing.request.CompoundName PARAM_DEFAULT_MAX_HITS",
+      "public static final com.yahoo.processing.request.CompoundName PARAM_DEFAULT_MAX_GROUPS"
     ]
   },
   "com.yahoo.search.grouping.GroupingRequest": {
@@ -2336,6 +2338,10 @@
       "public com.yahoo.search.grouping.GroupingRequest setTimeZone(java.util.TimeZone)",
       "public com.yahoo.search.grouping.result.RootGroup getResultGroup(com.yahoo.search.Result)",
       "public java.util.List continuations()",
+      "public java.util.OptionalInt defaultMaxHits()",
+      "public void setDefaultMaxHits(int)",
+      "public java.util.OptionalInt defaultMaxGroups()",
+      "public void setDefaultMaxGroups(int)",
       "public static com.yahoo.search.grouping.GroupingRequest newInstance(com.yahoo.search.Query)",
       "public java.lang.String toString()"
     ],
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 f4603e5ea28..0de4e36eae5 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
@@ -1,6 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.search.grouping;
 
+import com.yahoo.api.annotations.Beta;
 import com.yahoo.component.chain.dependencies.After;
 import com.yahoo.component.chain.dependencies.Before;
 import com.yahoo.component.chain.dependencies.Provides;
@@ -37,6 +38,8 @@ public class GroupingQueryParser extends Searcher {
     public static final CompoundName PARAM_CONTINUE = new CompoundName("continue");
     public static final CompoundName PARAM_REQUEST = new CompoundName(Select.SELECT);
     public static final CompoundName PARAM_TIMEZONE = new CompoundName("timezone");
+    @Beta public static final CompoundName PARAM_DEFAULT_MAX_HITS = new CompoundName("grouping.defaultMaxHits");
+    @Beta public static final CompoundName PARAM_DEFAULT_MAX_GROUPS = new CompoundName("grouping.defaultMaxGroups");
     private static final ThreadLocal<ZoneCache> zoneCache = new ThreadLocal<>();
 
     @Override
@@ -52,6 +55,8 @@ public class GroupingQueryParser extends Searcher {
                 grpRequest.setRootOperation(op);
                 grpRequest.setTimeZone(zone);
                 grpRequest.continuations().addAll(continuations);
+                grpRequest.setDefaultMaxGroups(query.properties().getInteger(PARAM_DEFAULT_MAX_GROUPS, -1));
+                grpRequest.setDefaultMaxHits(query.properties().getInteger(PARAM_DEFAULT_MAX_HITS, -1));
             }
             return execution.search(query);
         }
diff --git a/container-search/src/main/java/com/yahoo/search/grouping/GroupingRequest.java b/container-search/src/main/java/com/yahoo/search/grouping/GroupingRequest.java
index 6b21b2c6b5c..0c163aaacae 100644
--- a/container-search/src/main/java/com/yahoo/search/grouping/GroupingRequest.java
+++ b/container-search/src/main/java/com/yahoo/search/grouping/GroupingRequest.java
@@ -1,6 +1,7 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.search.grouping;
 
+import com.yahoo.api.annotations.Beta;
 import com.yahoo.net.URI;
 import com.yahoo.search.Query;
 import com.yahoo.search.Result;
@@ -13,6 +14,7 @@ import com.yahoo.search.result.Hit;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.OptionalInt;
 import java.util.TimeZone;
 
 /**
@@ -30,6 +32,8 @@ public class GroupingRequest {
     private final List<Continuation> continuations = new ArrayList<>();
     private GroupingOperation root;
     private TimeZone timeZone;
+    private int defaultMaxHits = -1;
+    private int defaultMaxGroups = -1;
 
     private GroupingRequest(Select parent) {
         this.parent = parent;
@@ -38,16 +42,20 @@ public class GroupingRequest {
     private GroupingRequest(Select parent,
                             List<Continuation> continuations,
                             GroupingOperation root,
-                            TimeZone timeZone) {
+                            TimeZone timeZone,
+                            int defaultMaxHits,
+                            int defaultMaxGroups) {
         this.parent = parent;
         continuations.forEach(item -> this.continuations.add(item.copy()));
         this.root = root != null ? root.copy(null) : null;
         this.timeZone = timeZone;
+        this.defaultMaxHits = defaultMaxHits;
+        this.defaultMaxGroups = defaultMaxGroups;
     }
 
     /** Returns a deep copy of this */
     public GroupingRequest copy(Select parentOfCopy) {
-        return new GroupingRequest(parentOfCopy, continuations, root, timeZone);
+        return new GroupingRequest(parentOfCopy, continuations, root, timeZone, defaultMaxHits, defaultMaxGroups);
     }
 
     /**
@@ -131,6 +139,20 @@ public class GroupingRequest {
         return continuations;
     }
 
+    @Beta
+    public OptionalInt defaultMaxHits() {
+        return defaultMaxHits >= 0 ? OptionalInt.of(defaultMaxHits) : OptionalInt.empty();
+    }
+
+    @Beta public void setDefaultMaxHits(int v) { this.defaultMaxHits = v; }
+
+    @Beta
+    public OptionalInt defaultMaxGroups() {
+        return defaultMaxGroups >= 0 ? OptionalInt.of(defaultMaxGroups) : OptionalInt.empty();
+    }
+
+    @Beta public void setDefaultMaxGroups(int v) { this.defaultMaxGroups = v; }
+
     /**
      * Creates a new grouping request and adds it to the query.getSelect().getGrouping() list
      *
diff --git a/container-search/src/main/java/com/yahoo/search/grouping/vespa/GroupingExecutor.java b/container-search/src/main/java/com/yahoo/search/grouping/vespa/GroupingExecutor.java
index 3e707b0cd38..c09502110b1 100644
--- a/container-search/src/main/java/com/yahoo/search/grouping/vespa/GroupingExecutor.java
+++ b/container-search/src/main/java/com/yahoo/search/grouping/vespa/GroupingExecutor.java
@@ -150,6 +150,8 @@ public class GroupingExecutor extends Searcher {
         builder.setDefaultSummaryName(query.getPresentation().getSummary());
         builder.setTimeZone(req.getTimeZone());
         builder.addContinuations(req.continuations());
+        req.defaultMaxGroups().ifPresent(builder::setDefaultMaxGroups);
+        req.defaultMaxHits().ifPresent(builder::setDefaultMaxHits);
         builder.build();
 
         RequestContext ctx = new RequestContext(req, builder.getTransform());
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 e636c0d63a6..5ee6f7bc604 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
@@ -27,6 +27,8 @@ class RequestBuilder {
     private final GroupingTransform transform;
     private GroupingOperation root;
     private int tag = 0;
+    private int defaultMaxHits = -1;
+    private int defaultMaxGroups = -1;
 
     /**
      * Constructs a new instance of this class.
@@ -138,6 +140,10 @@ class RequestBuilder {
         return this;
     }
 
+    public RequestBuilder setDefaultMaxGroups(int v) { this.defaultMaxGroups = v; return this; }
+
+    public RequestBuilder setDefaultMaxHits(int v) { this.defaultMaxHits = v; return this; }
+
     private void processRequestNode(BuildFrame frame) {
         int level = frame.astNode.getLevel();
         if (level > 2) {
@@ -173,10 +179,16 @@ class RequestBuilder {
                 grpLevel.setPrecision(frame.state.precision + offset);
                 frame.state.precision = null;
             }
+            int max = -1;
             if (frame.state.max != null) {
-                transform.putMax(tag, frame.state.max, "group list");
-                grpLevel.setMaxGroups(LOOKAHEAD + frame.state.max + offset);
+                max = frame.state.max;
                 frame.state.max = null;
+            } else if (defaultMaxGroups >= 0) {
+                max = defaultMaxGroups;
+            }
+            if (max >= 0) {
+                transform.putMax(tag, max, "group list");
+                grpLevel.setMaxGroups(LOOKAHEAD + max + offset);
             }
             frame.grouping.getLevels().add(grpLevel);
         }
@@ -285,11 +297,17 @@ class RequestBuilder {
                 throw new UnsupportedOperationException("Can not label expression '" + exp + "'.");
             }
             HitsAggregationResult hits = (HitsAggregationResult)result;
+            int max = -1;
             if (frame.state.max != null) {
-                transform.putMax(tag, frame.state.max, "hit list");
-                int offset = transform.getOffset(tag);
-                hits.setMaxHits(LOOKAHEAD + frame.state.max + offset);
+                max = frame.state.max;
                 frame.state.max = null;
+            } else if (defaultMaxHits >= 0) {
+                max = defaultMaxHits;
+            }
+            if (max >= 0) {
+                transform.putMax(tag, max, "hit list");
+                int offset = transform.getOffset(tag);
+                hits.setMaxHits(LOOKAHEAD + max + offset);
             }
             transform.putLabel(group.getTag(), tag, frame.state.label, "hit list");
         } 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 4397831af6e..29d9f7802c1 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
@@ -698,6 +698,36 @@ public class RequestBuilderTestCase {
         assertOutput(test);
     }
 
+    @Test
+    public void require_that_default_max_values_from_request_builder_restricts_max_groups_and_hits() {
+        int defaultMaxHits = 19;
+        int defaultMaxGroups = 7;
+        RequestBuilder builder = new RequestBuilder(0)
+                .setDefaultMaxGroups(defaultMaxGroups)
+                .setDefaultMaxHits(defaultMaxHits)
+                .setRootOperation(GroupingOperation.fromString("all(group(foo)each(each(output(summary()))))"));
+        builder.build();
+        List<Grouping> requests = builder.getRequestList();
+        assertEquals(defaultMaxGroups + 1, requests.get(0).getLevels().get(0).getMaxGroups());
+        HitsAggregationResult hitsAggregation =
+                (HitsAggregationResult)requests.get(0).getLevels().get(0).getGroupPrototype().getAggregationResults().get(0);
+        assertEquals(defaultMaxHits + 1, hitsAggregation.getMaxHits());
+    }
+
+    @Test
+    public void require_that_default_max_values_from_request_builder_restricts_respects_explicit_max() {
+        RequestBuilder builder = new RequestBuilder(0)
+                .setDefaultMaxGroups(7)
+                .setDefaultMaxHits(19)
+                .setRootOperation(GroupingOperation.fromString("all(group(foo)max(11)each(max(21)each(output(summary()))))"));
+        builder.build();
+        List<Grouping> requests = builder.getRequestList();
+        assertEquals(12, requests.get(0).getLevels().get(0).getMaxGroups());
+        HitsAggregationResult hitsAggregation =
+                (HitsAggregationResult)requests.get(0).getLevels().get(0).getGroupPrototype().getAggregationResults().get(0);
+        assertEquals(22, hitsAggregation.getMaxHits());
+    }
+
     private static CompositeContinuation newComposite(EncodableContinuation... conts) {
         CompositeContinuation ret = new CompositeContinuation();
         for (EncodableContinuation cont : conts) {