Add FlatteningSearcher

8 files changed, 189 insertions, 49 deletions

diff --git a/container-search/src/main/java/com/yahoo/prelude/searcher/MultipleResultsSearcher.java b/container-search/src/main/java/com/yahoo/prelude/searcher/MultipleResultsSearcher.java
index d017cce0d44..3c61a361cbb 100644
--- a/container-search/src/main/java/com/yahoo/prelude/searcher/MultipleResultsSearcher.java
+++ b/container-search/src/main/java/com/yahoo/prelude/searcher/MultipleResultsSearcher.java
@@ -13,9 +13,8 @@ import com.yahoo.search.searchchain.Execution;
 import java.util.*;
 
 /**
- * <p> Groups hits according to sddocname. </p>
- *
- * <p> For each group, the desired number of hits can be specified. </p>
+ * Groups hits according to document type.
+ * For each group, the desired number of hits can be specified.
  *
  * @author Tony Vaagenes
  */
@@ -48,7 +47,7 @@ public class MultipleResultsSearcher extends Searcher {
         }
     }
 
-    private class HitsRetriever {
+    private static class HitsRetriever {
 
         PartitionedResult partitionedResult;
 
@@ -58,12 +57,12 @@ public class MultipleResultsSearcher extends Searcher {
         private final Parameters parameters;
         private final int hits;
         private final int offset;
-        private Execution execution;
-        private Result initialResult;
+        private final Execution execution;
+        private final Result initialResult;
 
         HitsRetriever(Query query, Execution execution, Parameters parameters) throws ParameterException {
-            this.offset=query.getOffset();
-            this.hits=query.getHits();
+            this.offset = query.getOffset();
+            this.hits = query.getHits();
             this.nextOffset = query.getOffset() + query.getHits();
             this.query = query;
             this.parameters = parameters;
@@ -362,13 +361,14 @@ public class MultipleResultsSearcher extends Searcher {
         }
     }
 
-    @SuppressWarnings("serial")
     private static class ParameterException extends Exception {
+
         String msg;
 
         ParameterException(String msg) {
             this.msg = msg;
         }
+
     }
 
 }
diff --git a/container-search/src/main/java/com/yahoo/search/Query.java b/container-search/src/main/java/com/yahoo/search/Query.java
index fb7281e1f24..83fa18d847f 100644
--- a/container-search/src/main/java/com/yahoo/search/Query.java
+++ b/container-search/src/main/java/com/yahoo/search/Query.java
@@ -726,8 +726,7 @@ public class Query extends com.yahoo.processing.Request implements Cloneable {
      * if the trace level of the query is sufficiently high.
      *
      * @param message      the message to add
-     * @param includeQuery true to append the query root stringValue
-     *        at the end of the message
+     * @param includeQuery true to append the query root stringValue at the end of the message
      * @param traceLevel   the context level of the message, this method will do nothing
      *                     if the traceLevel of the query is lower than this value
      */
diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/InterleavedSearchInvoker.java b/container-search/src/main/java/com/yahoo/search/dispatch/InterleavedSearchInvoker.java
index d7c9f1dce53..5e38f6b4bdd 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/InterleavedSearchInvoker.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/InterleavedSearchInvoker.java
@@ -116,8 +116,6 @@ public class InterleavedSearchInvoker extends SearchInvoker implements ResponseM
         InvokerResult result = new InvokerResult(query, query.getHits());
         List<LeanHit> merged = Collections.emptyList();
         long nextTimeout = query.getTimeLeft();
-        boolean extraDebug = (query.getOffset() == 0) && (query.getHits() == 7) && log.isLoggable(java.util.logging.Level.FINE);
-        List<InvokerResult> processed = new ArrayList<>();
         var groupingResultAggregator = new GroupingResultAggregator();
         try {
             while (!invokers.isEmpty() && nextTimeout >= 0) {
@@ -127,9 +125,6 @@ public class InterleavedSearchInvoker extends SearchInvoker implements ResponseM
                     break;
                 } else {
                     InvokerResult toMerge = invoker.getSearchResult(execution);
-                    if (extraDebug) {
-                        processed.add(toMerge);
-                    }
                     merged = mergeResult(result.getResult(), toMerge, merged, groupingResultAggregator);
                     ejectInvoker(invoker);
                 }
@@ -143,32 +138,6 @@ public class InterleavedSearchInvoker extends SearchInvoker implements ResponseM
         insertNetworkErrors(result.getResult());
         result.getResult().setCoverage(createCoverage());
 
-        if (extraDebug && merged.size() > 0) {
-            int firstPartId = merged.get(0).getPartId();
-            for (int index = 1; index < merged.size(); index++) {
-                if (merged.get(index).getPartId() != firstPartId) {
-                    extraDebug = false;
-                    log.fine("merged["+index+"/"+merged.size()+"] from partId "+merged.get(index).getPartId()+", first "+firstPartId);
-                    break;
-                }
-            }
-        }
-        if (extraDebug) {
-            log.fine("Interleaved "+processed.size()+" results");
-            for (int pIdx = 0; pIdx < processed.size(); ++pIdx) {
-                var p = processed.get(pIdx);
-                log.fine("InvokerResult "+pIdx+" total hits "+p.getResult().getTotalHitCount());
-                var lean = p.getLeanHits();
-                for (int idx = 0; idx < lean.size(); ++idx) {
-                    var hit = lean.get(idx);
-                    log.fine("lean hit "+idx+" relevance "+hit.getRelevance()+" partid "+hit.getPartId());
-                }
-            }
-            for (int mIdx = 0; mIdx < merged.size(); ++mIdx) {
-                var hit = merged.get(mIdx);
-                log.fine("merged hit "+mIdx+" relevance "+hit.getRelevance()+" partid "+hit.getPartId());                
-            }
-        }
         int needed = query.getOffset() + query.getHits();
         for (int index = query.getOffset(); (index < merged.size()) && (index < needed); index++) {
             result.getLeanHits().add(merged.get(index));
diff --git a/container-search/src/main/java/com/yahoo/search/grouping/result/FlatteningSearcher.java b/container-search/src/main/java/com/yahoo/search/grouping/result/FlatteningSearcher.java
new file mode 100644
index 00000000000..321f86facd0
--- /dev/null
+++ b/container-search/src/main/java/com/yahoo/search/grouping/result/FlatteningSearcher.java
@@ -0,0 +1,50 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.search.grouping.result;
+
+import com.yahoo.component.chain.dependencies.Before;
+import com.yahoo.search.Query;
+import com.yahoo.search.Result;
+import com.yahoo.search.Searcher;
+import com.yahoo.search.grouping.vespa.GroupingExecutor;
+import com.yahoo.search.result.Hit;
+import com.yahoo.search.result.HitGroup;
+import com.yahoo.search.searchchain.Execution;
+
+import java.util.Iterator;
+
+/**
+ * Flattens a grouping result into a flat list of hits on the top level in the returned result.
+ * Useful when using grouping to create hits with diversity and similar.
+ *
+ * @author bratseth
+ */
+@Before(GroupingExecutor.COMPONENT_NAME)
+public class FlatteningSearcher extends Searcher {
+
+    @Override
+    public Result search(Query query, Execution execution) {
+        if ( ! query.properties().getBoolean("flatten", true)) return execution.search(query);
+
+        query.trace("Flattening groups", 2);
+        int originalHits = query.getHits();
+        query.setHits(0);
+        Result result = execution.search(query);
+        query.setHits(originalHits);
+        flatten(result.hits(), result);
+        return result;
+    }
+
+    public void flatten(HitGroup hits, Result result) {
+        int hitsLeft = hits.size(); // Iterate only through the initial size
+        for (Iterator<Hit> i = hits.iterator(); i.hasNext() && hitsLeft-- > 0;) {
+            Hit hit = i.next();
+            if (hit instanceof HitGroup) {
+                flatten((HitGroup)hit, result);
+                i.remove();
+            } else {
+                result.hits().add(hit);
+            }
+        }
+    }
+
+}
diff --git a/container-search/src/main/java/com/yahoo/search/result/HitGroup.java b/container-search/src/main/java/com/yahoo/search/result/HitGroup.java
index 6d09bf66175..efe25e04f2e 100644
--- a/container-search/src/main/java/com/yahoo/search/result/HitGroup.java
+++ b/container-search/src/main/java/com/yahoo/search/result/HitGroup.java
@@ -249,9 +249,7 @@ public class HitGroup extends Hit implements DataList<Hit>, Cloneable, Iterable<
         return hit;
     }
 
-    /**
-     * Adds a list of hits to this group, the same
-     */
+    /** Adds a list of hits to this group, the same as calling add for each item in the list. */
     public void addAll(List<Hit> hits) {
         for (Hit hit : hits)
             add(hit);
diff --git a/container-search/src/test/java/com/yahoo/search/grouping/result/FlatteningSearcherTestCase.java b/container-search/src/test/java/com/yahoo/search/grouping/result/FlatteningSearcherTestCase.java
new file mode 100644
index 00000000000..9e5bfb4a9f5
--- /dev/null
+++ b/container-search/src/test/java/com/yahoo/search/grouping/result/FlatteningSearcherTestCase.java
@@ -0,0 +1,126 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.search.grouping.result;
+
+import com.yahoo.component.ComponentId;
+import com.yahoo.component.chain.dependencies.After;
+import com.yahoo.document.GlobalId;
+import com.yahoo.prelude.fastsearch.FastHit;
+import com.yahoo.prelude.fastsearch.GroupingListHit;
+import com.yahoo.search.Query;
+import com.yahoo.search.Result;
+import com.yahoo.search.Searcher;
+import com.yahoo.search.grouping.GroupingRequest;
+import com.yahoo.search.grouping.request.GroupingOperation;
+import com.yahoo.search.grouping.vespa.GroupingExecutor;
+import com.yahoo.search.result.Hit;
+import com.yahoo.search.result.HitGroup;
+import com.yahoo.search.searchchain.Execution;
+import com.yahoo.search.searchchain.SearchChain;
+import com.yahoo.searchlib.aggregation.FS4Hit;
+import com.yahoo.searchlib.aggregation.Group;
+import com.yahoo.searchlib.aggregation.Grouping;
+import com.yahoo.searchlib.aggregation.HitsAggregationResult;
+import com.yahoo.searchlib.expression.StringResultNode;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Queue;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * @author bratseth
+ */
+public class FlatteningSearcherTestCase {
+
+    @Test
+    public void testFlatteningSearcher() {
+        Query query = new Query("?query=test");
+        GroupingRequest req = GroupingRequest.newInstance(query);
+        req.setRootOperation(GroupingOperation.fromString("all(group(foo) each(each(output(summary(bar)))))"));
+
+        Grouping group0 = new Grouping(0);
+        group0.setRoot(new com.yahoo.searchlib.aggregation.Group()
+                               .addChild(new Group().setId(new StringResultNode("unique1"))
+                                                    .addAggregationResult(new HitsAggregationResult(3, "bar")
+                                                    )
+                               )
+                               .addChild(new Group().setId(new StringResultNode("unique2"))
+                                                    .addAggregationResult(new HitsAggregationResult(3, "bar")
+                                                    )
+                               ));
+        Grouping group1 = new Grouping(0);
+        group1.setRoot(new com.yahoo.searchlib.aggregation.Group()
+                               .addChild(new Group().setId(new StringResultNode("unique1"))
+                                                    .addAggregationResult(new HitsAggregationResult(3, "bar")
+                                                                                  .addHit(fs4Hit(0.7))
+                                                                                  .addHit(fs4Hit(0.6))
+                                                                                  .addHit(fs4Hit(0.3))
+                                                    )
+                               )
+                               .addChild(new Group().setId(new StringResultNode("unique2"))
+                                                    .addAggregationResult(new HitsAggregationResult(3, "bar")
+                                                                                  .addHit(fs4Hit(0.5))
+                                                                                  .addHit(fs4Hit(0.4))
+                                                    )
+                               ));
+        Execution execution = newExecution(new FlatteningSearcher(),
+                                           new GroupingExecutor(ComponentId.fromString("grouping")),
+                                           new ResultProvider(Arrays.asList(
+                                                   new GroupingListHit(List.of(group0), null),
+                                                   new GroupingListHit(List.of(group1), null))));
+        Result result = execution.search(query);
+        assertEquals(5, result.hits().size());
+        assertFlat(result);
+    }
+
+    private void assertFlat(Result result) {
+        for (var hit : result.hits())
+            assertTrue(hit instanceof FastHit);
+    }
+
+    private FS4Hit fs4Hit(double relevance) {
+        return new FS4Hit(0, new GlobalId(new byte[GlobalId.LENGTH]), relevance);
+    }
+
+    private void dump(Hit hit, String indent) {
+        System.out.println(indent + hit);
+        if (hit instanceof HitGroup) {
+            for (var child : (HitGroup)hit)
+                dump(child, indent + "    ");
+        }
+    }
+
+    private static Execution newExecution(Searcher... searchers) {
+        return new Execution(new SearchChain(new ComponentId("foo"), Arrays.asList(searchers)),
+                             Execution.Context.createContextStub());
+    }
+
+    @After (GroupingExecutor.COMPONENT_NAME)
+    private static class ResultProvider extends Searcher {
+
+        final Queue<GroupingListHit> hits = new LinkedList<>();
+        int pass = 0;
+
+        ResultProvider(List<GroupingListHit> hits) {
+            this.hits.addAll(hits);
+        }
+
+        @Override
+        public Result search(Query query, Execution exec) {
+            GroupingListHit hit = hits.poll();
+            for (Grouping group : hit.getGroupingList()) {
+                group.setFirstLevel(pass);
+                group.setLastLevel(pass);
+            }
+            ++pass;
+            Result result = exec.search(query);
+            result.hits().add(hit);
+            return result;
+        }
+    }
+
+}
diff --git a/container-search/src/test/java/com/yahoo/search/grouping/result/HitRendererTestCase.java b/container-search/src/test/java/com/yahoo/search/grouping/result/HitRendererTestCase.java
index 7f70178bcf7..1035c9d9284 100644
--- a/container-search/src/test/java/com/yahoo/search/grouping/result/HitRendererTestCase.java
+++ b/container-search/src/test/java/com/yahoo/search/grouping/result/HitRendererTestCase.java
@@ -151,7 +151,6 @@ public class HitRendererTestCase {
         return new Group(id, new Relevance(1));
     }
 
-    @SuppressWarnings("deprecation")
     private static void assertRender(HitGroup hit, String expectedXml) {
         StringWriter str = new StringWriter();
         XMLWriter out = new XMLWriter(str, 0, -1);
diff --git a/container-search/src/test/java/com/yahoo/search/rendering/JsonRendererTestCase.java b/container-search/src/test/java/com/yahoo/search/rendering/JsonRendererTestCase.java
index 637666a6a19..1d2e3c240b9 100644
--- a/container-search/src/test/java/com/yahoo/search/rendering/JsonRendererTestCase.java
+++ b/container-search/src/test/java/com/yahoo/search/rendering/JsonRendererTestCase.java
@@ -948,7 +948,7 @@ public class JsonRendererTestCase {
             }
         });
         Group g = new Group(new StringId("Jones"), new Relevance(1.0));
-        g.setField("count()", Integer.valueOf(7));
+        g.setField("count()", 7);
         gl.add(g);
         rg.add(gl);
         r.hits().add(rg);
@@ -1443,8 +1443,7 @@ public class JsonRendererTestCase {
         ByteArrayOutputStream bs = new ByteArrayOutputStream();
         ListenableFuture<Boolean> f = renderer.render(bs, r, execution, null);
         assertTrue(f.get());
-        String summary = Utf8.toString(bs.toByteArray());
-        return summary;
+        return Utf8.toString(bs.toByteArray());
     }
 
     @SuppressWarnings("unchecked")