Use unique group count as total-hits if available

2 files changed, 22 insertions, 7 deletions

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
index 321f86facd0..14e82e9e46d 100644
--- 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
@@ -2,6 +2,7 @@
 package com.yahoo.search.grouping.result;
 
 import com.yahoo.component.chain.dependencies.Before;
+import com.yahoo.processing.request.CompoundName;
 import com.yahoo.search.Query;
 import com.yahoo.search.Result;
 import com.yahoo.search.Searcher;
@@ -21,8 +22,11 @@ import java.util.Iterator;
 @Before(GroupingExecutor.COMPONENT_NAME)
 public class FlatteningSearcher extends Searcher {
 
+    private final CompoundName flatten = CompoundName.fromComponents("grouping", "flatten");
+
     @Override
     public Result search(Query query, Execution execution) {
+        if ( ! query.properties().getBoolean(flatten, true)) return execution.search(query);
         if ( ! query.properties().getBoolean("flatten", true)) return execution.search(query);
 
         query.trace("Flattening groups", 2);
@@ -30,16 +34,23 @@ public class FlatteningSearcher extends Searcher {
         query.setHits(0);
         Result result = execution.search(query);
         query.setHits(originalHits);
-        flatten(result.hits(), result);
+        flatten(result.hits(), 0, result);
         return result;
     }
 
-    public void flatten(HitGroup hits, Result result) {
+    public void flatten(HitGroup hits, int level, 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 we count the number of unique groups, use that as total hit count.
+            if (level == 0 && (hit instanceof RootGroup)) {
+                if (hit.fields().get("count()") != null)
+                    result.setTotalHitCount((long)hit.fields().get("count()"));
+            }
+
             if (hit instanceof HitGroup) {
-                flatten((HitGroup)hit, result);
+                flatten((HitGroup)hit, level++, result);
                 i.remove();
             } else {
                 result.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
index 9e5bfb4a9f5..c97bf150645 100644
--- 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
@@ -16,10 +16,12 @@ 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.ExpressionCountAggregationResult;
 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.aggregation.hll.SparseSketch;
 import com.yahoo.searchlib.expression.StringResultNode;
 import org.junit.Test;
 
@@ -40,10 +42,11 @@ public class FlatteningSearcherTestCase {
     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)))))"));
+        req.setRootOperation(GroupingOperation.fromString("all(group(foo) output(count()) each(each(output(summary(bar)))))"));
 
         Grouping group0 = new Grouping(0);
-        group0.setRoot(new com.yahoo.searchlib.aggregation.Group()
+        group0.setRoot(new Group()
+                               .addAggregationResult(new ExpressionCountAggregationResult(new SparseSketch(), sketch -> 69))
                                .addChild(new Group().setId(new StringResultNode("unique1"))
                                                     .addAggregationResult(new HitsAggregationResult(3, "bar")
                                                     )
@@ -53,7 +56,7 @@ public class FlatteningSearcherTestCase {
                                                     )
                                ));
         Grouping group1 = new Grouping(0);
-        group1.setRoot(new com.yahoo.searchlib.aggregation.Group()
+        group1.setRoot(new Group()
                                .addChild(new Group().setId(new StringResultNode("unique1"))
                                                     .addAggregationResult(new HitsAggregationResult(3, "bar")
                                                                                   .addHit(fs4Hit(0.7))
@@ -75,6 +78,7 @@ public class FlatteningSearcherTestCase {
         Result result = execution.search(query);
         assertEquals(5, result.hits().size());
         assertFlat(result);
+        assertEquals(2, result.getTotalHitCount());
     }
 
     private void assertFlat(Result result) {
@@ -87,7 +91,7 @@ public class FlatteningSearcherTestCase {
     }
 
     private void dump(Hit hit, String indent) {
-        System.out.println(indent + hit);
+        System.out.println(indent + hit + " (class " + hit.getClass() + ")");
         if (hit instanceof HitGroup) {
             for (var child : (HitGroup)hit)
                 dump(child, indent + "    ");