Merge pull request #19874 from vespa-engine/bratseth/grouping-with-types-test

Bratseth/grouping with types test

2 files changed, 142 insertions, 34 deletions

diff --git a/container-search/src/main/java/com/yahoo/search/grouping/vespa/ResultBuilder.java b/container-search/src/main/java/com/yahoo/search/grouping/vespa/ResultBuilder.java
index 81b7ae51963..c52e36986be 100644
--- a/container-search/src/main/java/com/yahoo/search/grouping/vespa/ResultBuilder.java
+++ b/container-search/src/main/java/com/yahoo/search/grouping/vespa/ResultBuilder.java
@@ -162,14 +162,14 @@ class ResultBuilder {
         }
 
         Group fill(Group group) {
-            for (AggregationResult res : this.group.getAggregationResults()) {
-                int tag = res.getTag();
-                if (res instanceof HitsAggregationResult) {
-                    group.add(newHitList(group.size(), tag, (HitsAggregationResult)res));
+            for (AggregationResult result : this.group.getAggregationResults()) {
+                int tag = result.getTag();
+                if (result instanceof HitsAggregationResult) {
+                    group.add(newHitList(group.size(), tag, (HitsAggregationResult)result));
                 } else {
-                    String label = transform.getLabel(res.getTag());
+                    String label = transform.getLabel(result.getTag());
                     if (label != null) {
-                        group.setField(label, newResult(res, tag));
+                        group.setField(label, newResult(result, tag));
                     }
                 }
             }
@@ -331,13 +331,13 @@ class ResultBuilder {
         }
 
         GroupBuilder getOrCreateGroup(com.yahoo.searchlib.aggregation.Group execGroup) {
-            ResultNode res = execGroup.getId();
-            GroupBuilder ret = childResultGroups.get(res);
+            ResultNode result = execGroup.getId();
+            GroupBuilder ret = childResultGroups.get(result);
             if (ret != null) {
                 ret.merge(execGroup);
             } else {
                 ret = new GroupBuilder(resultId.newChildId(childResultGroups.size()), execGroup, stableChildren);
-                childResultGroups.put(res, ret);
+                childResultGroups.put(result, ret);
                 childGroups.add(ret);
             }
             return ret;
diff --git a/container-search/src/test/java/com/yahoo/search/grouping/vespa/GroupingExecutorTestCase.java b/container-search/src/test/java/com/yahoo/search/grouping/vespa/GroupingExecutorTestCase.java
index 4c5f3f2bd79..2fda56c7454 100644
--- a/container-search/src/test/java/com/yahoo/search/grouping/vespa/GroupingExecutorTestCase.java
+++ b/container-search/src/test/java/com/yahoo/search/grouping/vespa/GroupingExecutorTestCase.java
@@ -35,7 +35,6 @@ import com.yahoo.searchlib.expression.StringResultNode;
 
 import org.junit.Test;
 
-import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -188,8 +187,8 @@ public class GroupingExecutorTestCase {
         );
         Execution exec = newExecution(new GroupingExecutor(),
                                       new ResultProvider(Arrays.asList(
-                                              new GroupingListHit(Arrays.asList(grpA), null),
-                                              new GroupingListHit(Arrays.asList(grpB), null))));
+                                              new GroupingListHit(List.of(grpA), null),
+                                              new GroupingListHit(List.of(grpB), null))));
         Group grp = req.getResultGroup(exec.search(query));
         assertEquals(1, grp.size());
         Hit hit = grp.get(0);
@@ -221,8 +220,8 @@ public class GroupingExecutorTestCase {
         );
         Execution exec = newExecution(new GroupingExecutor(),
                                       new ResultProvider(Arrays.asList(
-                                              new GroupingListHit(Arrays.asList(grpExpected), null),
-                                              new GroupingListHit(Arrays.asList(grpUnexpected), null))));
+                                              new GroupingListHit(List.of(grpExpected), null),
+                                              new GroupingListHit(List.of(grpUnexpected), null))));
         Group grp = req.getResultGroup(exec.search(query));
         assertEquals(1, grp.size());
         Hit hit = grp.get(0);
@@ -252,8 +251,8 @@ public class GroupingExecutorTestCase {
         ));
         Execution exec = newExecution(new GroupingExecutor(),
                                       new ResultProvider(Arrays.asList(
-                                              new GroupingListHit(Arrays.asList(grp0), null),
-                                              new GroupingListHit(Arrays.asList(grp1), null))),
+                                              new GroupingListHit(List.of(grp0), null),
+                                              new GroupingListHit(List.of(grp1), null))),
                                       new FillRequestThrower());
         Result res = exec.search(query);
         
@@ -275,7 +274,7 @@ public class GroupingExecutorTestCase {
     }
 
     @Test
-    public void requireThatUnfilledHitsRenderError() throws IOException {
+    public void requireThatUnfilledHitsRenderError() {
         Query query = newQuery();
         GroupingRequest req = GroupingRequest.newInstance(query);
         req.setRootOperation(GroupingOperation.fromString("all(group(foo) each(each(output(summary(bar)))))"));
@@ -292,8 +291,8 @@ public class GroupingExecutorTestCase {
                                                                                           .addHit(new com.yahoo.searchlib.aggregation.FS4Hit()))));
         Execution exec = newExecution(new GroupingExecutor(),
                                       new ResultProvider(Arrays.asList(
-                                              new GroupingListHit(Arrays.asList(grp0), null),
-                                              new GroupingListHit(Arrays.asList(grp1), null))),
+                                              new GroupingListHit(List.of(grp0), null),
+                                              new GroupingListHit(List.of(grp1), null))),
                                       new FillErrorProvider());
         Result res = exec.search(query);
         exec.fill(res);
@@ -318,8 +317,8 @@ public class GroupingExecutorTestCase {
                                       .addOrderBy(new AggregationRefNode(0), true)));
         Result res = newExecution(new GroupingExecutor(),
                                   new ResultProvider(Arrays.asList(
-                                          new GroupingListHit(Arrays.asList(grp), null),
-                                          new GroupingListHit(Arrays.asList(grp), null)))).search(query);
+                                          new GroupingListHit(List.of(grp), null),
+                                          new GroupingListHit(List.of(grp), null)))).search(query);
 
         GroupList groupList = (GroupList)req.getResultGroup(res).get(0);
         assertEquals(1.0, groupList.get(0).getRelevance().getScore(), 1E-6);
@@ -347,10 +346,10 @@ public class GroupingExecutorTestCase {
         Execution exec = newExecution(new GroupingExecutor(),
                                       err,
                                       new ResultProvider(Arrays.asList(
-                                              new GroupingListHit(Arrays.asList(grp0), null),
-                                              new GroupingListHit(Arrays.asList(grp1), null))));
+                                              new GroupingListHit(List.of(grp0), null),
+                                              new GroupingListHit(List.of(grp1), null))));
         Result res = exec.search(query);
-        assertTrue(res.hits().getError() != null);
+        assertNotNull(res.hits().getError());
         assertEquals(Error.TIMEOUT.code, res.hits().getError().getCode());
         assertFalse(err.continuedOnFail);
 
@@ -358,10 +357,10 @@ public class GroupingExecutorTestCase {
         exec = newExecution(new GroupingExecutor(),
                             err,
                             new ResultProvider(Arrays.asList(
-                                    new GroupingListHit(Arrays.asList(grp0), null),
-                                    new GroupingListHit(Arrays.asList(grp1), null))));
+                                    new GroupingListHit(List.of(grp0), null),
+                                    new GroupingListHit(List.of(grp1), null))));
         res = exec.search(query);
-        assertTrue(res.hits().getError() != null);
+        assertNotNull(res.hits().getError());
         assertEquals(Error.TIMEOUT.code, res.hits().getError().getCode());
         assertTrue(err.continuedOnFail);
     }
@@ -441,8 +440,8 @@ public class GroupingExecutorTestCase {
                                       .addHit(new com.yahoo.searchlib.aggregation.FS4Hit()))));
         Execution exec = newExecution(new GroupingExecutor(),
                                       new ResultProvider(Arrays.asList(
-                                              new GroupingListHit(Arrays.asList(pass0), null),
-                                              new GroupingListHit(Arrays.asList(pass1), null))));
+                                              new GroupingListHit(List.of(pass0), null),
+                                              new GroupingListHit(List.of(pass1), null))));
         Result res = exec.search(query);
         exec.fill(res);
 
@@ -462,7 +461,7 @@ public class GroupingExecutorTestCase {
                 .addChild(new com.yahoo.searchlib.aggregation.Group().setId(new StringResultNode("foo"))
                 .addAggregationResult(new HitsAggregationResult(1, "bar"))
         ));
-        GroupingListHit pass0 = new GroupingListHit(Arrays.asList(grp), null);
+        GroupingListHit pass0 = new GroupingListHit(List.of(grp), null);
 
         GlobalId gid = new GlobalId((new DocumentId("id:ns:type::1")).getGlobalId());
         grp = new Grouping(0);
@@ -470,8 +469,8 @@ public class GroupingExecutorTestCase {
                 .addChild(new com.yahoo.searchlib.aggregation.Group().setId(new StringResultNode("foo"))
                 .addAggregationResult(new HitsAggregationResult(1, "bar").addHit(new com.yahoo.searchlib.aggregation.FS4Hit(4, gid, 6)))
         ));
-        GroupingListHit pass1 = new GroupingListHit(Arrays.asList(grp), null);
-        Query queryB = newQuery(); /** required by {@link GroupingListHit#getSearchQuery()} */
+        GroupingListHit pass1 = new GroupingListHit(List.of(grp), null);
+        Query queryB = newQuery(); // required by GroupingListHit.getSearchQuery()
         pass1.setQuery(queryB);
 
         QueryMapper qm = new QueryMapper();
@@ -551,6 +550,63 @@ public class GroupingExecutorTestCase {
         assertEquals(3, message.getCode());
     }
 
+    @Test
+    public void testResultsFromMultipleDocumentTypes() {
+        Query query = newQuery();
+        GroupingRequest request = GroupingRequest.newInstance(query);
+        request.setRootOperation(GroupingOperation.fromString("all(group(foo) each(output(min(bar), max(bar))))"));
+
+        Map<String, List<GroupingListHit>> resultsByDocumentType = new HashMap<>();
+        Grouping groupA1 = new Grouping(0);
+        groupA1.setRoot(new com.yahoo.searchlib.aggregation.Group()
+                             .addChild(new com.yahoo.searchlib.aggregation.Group().setId(new StringResultNode("uniqueA")).addAggregationResult(new MaxAggregationResult().setMax(new IntegerResultNode(6)).setTag(4)))
+                             .addChild(new com.yahoo.searchlib.aggregation.Group().setId(new StringResultNode("common")).addAggregationResult(new MaxAggregationResult().setMax(new IntegerResultNode(9)).setTag(4)))
+        );
+        Grouping groupA2 = new Grouping(0);
+        groupA2.setRoot(new com.yahoo.searchlib.aggregation.Group()
+                             .addChild(new com.yahoo.searchlib.aggregation.Group().setId(new StringResultNode("uniqueB")).addAggregationResult(new MaxAggregationResult().setMax(new IntegerResultNode(9)).setTag(4)))
+                             .addChild(new com.yahoo.searchlib.aggregation.Group().setId(new StringResultNode("common")).addAggregationResult(new MinAggregationResult().setMin(new IntegerResultNode(6)).setTag(3)))
+        );
+        Grouping groupB1 = new Grouping(0);
+        groupB1.setRoot(new com.yahoo.searchlib.aggregation.Group()
+                                .addChild(new com.yahoo.searchlib.aggregation.Group().setId(new StringResultNode("uniqueA")).addAggregationResult(new MaxAggregationResult().setMax(new IntegerResultNode(2)).setTag(4)))
+                                .addChild(new com.yahoo.searchlib.aggregation.Group().setId(new StringResultNode("common")).addAggregationResult(new MaxAggregationResult().setMax(new IntegerResultNode(3)).setTag(4)))
+        );
+        Grouping groupB2 = new Grouping(0);
+        groupB2.setRoot(new com.yahoo.searchlib.aggregation.Group()
+                                .addChild(new com.yahoo.searchlib.aggregation.Group().setId(new StringResultNode("uniqueC")).addAggregationResult(new MaxAggregationResult().setMax(new IntegerResultNode(7)).setTag(4)))
+                                .addChild(new com.yahoo.searchlib.aggregation.Group().setId(new StringResultNode("common")).addAggregationResult(new MaxAggregationResult().setMax(new IntegerResultNode(11)).setTag(4)))
+        );
+        resultsByDocumentType.put("typeA", List.of(new GroupingListHit(List.of(groupA1), null),
+                                                   new GroupingListHit(List.of(groupA2), null)));
+        resultsByDocumentType.put("typeB", List.of(new GroupingListHit(List.of(groupB1), null),
+                                                   new GroupingListHit(List.of(groupB2), null)));
+        Execution execution = newExecution(new GroupingExecutor(),
+                                           new MockClusterSearcher(),
+                                           new MultiDocumentTypeResultProvider(resultsByDocumentType));
+
+        Result result = execution.search(query);
+        Group group = request.getResultGroup(result);
+        assertEquals(1, group.size());
+        Hit hit = group.get(0);
+        assertTrue(hit instanceof GroupList);
+        GroupList list = (GroupList)hit;
+
+        assertEquals(4, list.size());
+
+        assertNotNull(hit = list.get("group:string:uniqueA"));
+        assertEquals(6L, hit.getField("max(bar)"));
+
+        assertNotNull(hit = list.get("group:string:uniqueB"));
+        assertEquals(9L, hit.getField("max(bar)"));
+
+        assertNotNull(hit = list.get("group:string:common"));
+        assertEquals(11L, hit.getField("max(bar)"));
+
+        assertNotNull(hit = list.get("group:string:common"));
+        assertEquals(6L, hit.getField("min(bar)"));
+    }
+
     // --------------------------------------------------------------------------------
     //
     // Utilities
@@ -589,7 +645,6 @@ public class GroupingExecutorTestCase {
         }
     }
 
-    @SuppressWarnings("serial")
     private static class FillRequestException extends RuntimeException {
 
         final String summaryClass;
@@ -608,7 +663,6 @@ public class GroupingExecutorTestCase {
         }
     }
 
-    @SuppressWarnings("serial")
     private static class GroupingListException extends RuntimeException {
 
         final List<Grouping> lst;
@@ -755,6 +809,35 @@ public class GroupingExecutorTestCase {
         }
     }
 
+    /** Simulate multiple document types returning a grouping result */
+    @After (GroupingExecutor.COMPONENT_NAME)
+    private static class MultiDocumentTypeResultProvider extends Searcher {
+
+        final Map<String, List<GroupingListHit>> hitsByDocumentType;
+        final Map<String, Integer> passByDocumentType = new HashMap<>();
+
+        MultiDocumentTypeResultProvider(Map<String, List<GroupingListHit>> hitsByDocumentType) {
+            this.hitsByDocumentType = hitsByDocumentType;
+        }
+
+        @Override
+        public Result search(Query query, Execution execution) {
+            return result(query, execution, query.getModel().getRestrict().stream().findFirst().get());
+        }
+
+        private Result result(Query query, Execution execution, String documentType) {
+            GroupingListHit hit = hitsByDocumentType.get(documentType).get(passByDocumentType.getOrDefault(documentType, 0));
+            for (Grouping grp : hit.getGroupingList()) {
+                grp.setFirstLevel(passByDocumentType.getOrDefault(documentType, 0));
+                grp.setLastLevel(passByDocumentType.getOrDefault(documentType, 0));
+            }
+            passByDocumentType.compute(documentType, (k, v) -> v == null ? 1 : v + 1);
+            Result res = execution.search(query);
+            res.hits().add(hit);
+            return res;
+        }
+    }
+
     private static class FillErrorProvider extends Searcher {
 
         @Override
@@ -767,4 +850,29 @@ public class GroupingExecutorTestCase {
             result.hits().addError(ErrorMessage.createInternalServerError("foo"));
         }
     }
+
+    // The essence of prelude.ClusterSearcher
+    private static class MockClusterSearcher extends Searcher {
+
+        @Override
+        public Result search(Query query, Execution execution) {
+            Query queryA = query.clone();
+            queryA.getModel().setRestrict("typeA");
+            Result resultA = execution.search(queryA);
+
+            Query queryB = query.clone();
+            queryB.getModel().setRestrict("typeB");
+            Result resultB = execution.search(queryB);
+
+            Result mergedResult = new Result(query);
+            mergedResult.mergeWith(resultA);
+            mergedResult.hits().addAll(resultA.hits().asUnorderedHits());
+            mergedResult.mergeWith(resultB);
+            mergedResult.hits().addAll(resultB.hits().asUnorderedHits());
+
+            return mergedResult;
+        }
+
+    }
+
 }