Introduce AllGroupsCollectorManager

lucene/CHANGES.txt

@@ -120,6 +120,8 @@ Improvements
 
 * GITHUB#15225: Improve package documentation for org.apache.lucene.util. (Syed Mohammad Saad)
 
+* GITHUB#15557: Introduce AllGroupsCollectorManager to parallelize search when using AllGroupsCollector. (Binlong Gao)
+
 Optimizations
 ---------------------
 * GITHUB#15681, GITHUB#15833: Replace pre-sized array or empty array with lambda expression to call Collection#toArray. (Zhou Hui)

lucene/grouping/src/java/org/apache/lucene/search/grouping/AllGroupsCollectorManager.java

@@ -0,0 +1,71 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.lucene.search.grouping;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.function.Supplier;
+import org.apache.lucene.search.CollectorManager;
+
+/**
+ * A CollectorManager implementation for AllGroupsCollector.
+ *
+ * @lucene.experimental
+ */
+public class AllGroupsCollectorManager<T>
+    implements CollectorManager<AllGroupsCollector<T>, Collection<T>> {
+
+  private final Supplier<GroupSelector<T>> groupSelectorFactory;
+
+  /**
+   * Creates a new AllGroupsCollectorManager.
+   *
+   * @param groupSelectorFactory factory to create group selectors for each collector
+   */
+  public AllGroupsCollectorManager(Supplier<GroupSelector<T>> groupSelectorFactory) {
+    this.groupSelectorFactory = groupSelectorFactory;
+  }
+
+  @Override
+  public AllGroupsCollector<T> newCollector() {
+    return new AllGroupsCollector<>(groupSelectorFactory.get());
+  }
+
+  @Override
+  public Collection<T> reduce(Collection<AllGroupsCollector<T>> collectors) {
+    if (collectors.isEmpty()) {
+      return Collections.emptyList();
+    }
+
+    if (collectors.size() == 1) {
+      return collectors.iterator().next().getGroups();
+    }
+
+    // Merge groups from all collectors
+    Set<T> allGroups = new HashSet<>();
+    for (AllGroupsCollector<T> collector : collectors) {
+      Collection<T> groups = collector.getGroups();
+      if (groups != null) {
+        allGroups.addAll(groups);
+      }
+    }
+
+    return allGroups;
+  }
+}

lucene/grouping/src/test/org/apache/lucene/search/grouping/TestAllGroupsCollector.java

@@ -16,6 +16,7 @@
  */
 package org.apache.lucene.search.grouping;
 
+import java.util.Collection;
 import java.util.HashMap;
 import org.apache.lucene.document.Document;
 import org.apache.lucene.document.Field;
@@ -100,17 +101,22 @@ public void testTotalGroupCount() throws Exception {
     IndexSearcher indexSearcher = newSearcher(w.getReader());
     w.close();
 
-    AllGroupsCollector<?> allGroupsCollector = createRandomCollector(groupField);
-    indexSearcher.search(new TermQuery(new Term("content", "random")), allGroupsCollector);
-    assertEquals(4, allGroupsCollector.getGroupCount());
+    AllGroupsCollectorManager<?> allGroupsCollectorManager =
+        createRandomCollectorManager(groupField);
+    Collection<?> groups =
+        indexSearcher.search(
+            new TermQuery(new Term("content", "random")), allGroupsCollectorManager);
+    assertEquals(4, groups.size());
 
-    allGroupsCollector = createRandomCollector(groupField);
-    indexSearcher.search(new TermQuery(new Term("content", "some")), allGroupsCollector);
-    assertEquals(3, allGroupsCollector.getGroupCount());
+    allGroupsCollectorManager = createRandomCollectorManager(groupField);
+    groups =
+        indexSearcher.search(new TermQuery(new Term("content", "some")), allGroupsCollectorManager);
+    assertEquals(3, groups.size());
 
-    allGroupsCollector = createRandomCollector(groupField);
-    indexSearcher.search(new TermQuery(new Term("content", "blob")), allGroupsCollector);
-    assertEquals(2, allGroupsCollector.getGroupCount());
+    allGroupsCollectorManager = createRandomCollectorManager(groupField);
+    groups =
+        indexSearcher.search(new TermQuery(new Term("content", "blob")), allGroupsCollectorManager);
+    assertEquals(2, groups.size());
 
     indexSearcher.getIndexReader().close();
     dir.close();
@@ -121,12 +127,13 @@ private void addGroupField(Document doc, String groupField, String value) {
     doc.add(new SortedDocValuesField(groupField, new BytesRef(value)));
   }
 
-  private AllGroupsCollector<?> createRandomCollector(String groupField) {
+  private AllGroupsCollectorManager<?> createRandomCollectorManager(String groupField) {
     if (random().nextBoolean()) {
-      return new AllGroupsCollector<>(new TermGroupSelector(groupField));
+      return new AllGroupsCollectorManager<>(() -> new TermGroupSelector(groupField));
     } else {
       ValueSource vs = new BytesRefFieldSource(groupField);
-      return new AllGroupsCollector<>(new ValueSourceGroupSelector(vs, new HashMap<>()));
+      return new AllGroupsCollectorManager<>(
+          () -> new ValueSourceGroupSelector(vs, new HashMap<>()));
     }
   }
 }