You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@lucene.apache.org by ja...@apache.org on 2026/06/03 08:57:41 UTC
(lucene) branch branch_10x updated: Introduce FirstPassGroupingCollectorManager (#15574)
This is an automated email from the ASF dual-hosted git repository.
javanna pushed a commit to branch branch_10x
in repository https://gitbox.apache.org/repos/asf/lucene.git
The following commit(s) were added to refs/heads/branch_10x by this push:
new 5c701b3f553 Introduce FirstPassGroupingCollectorManager (#15574)
5c701b3f553 is described below
commit 5c701b3f55319a2f576e4f6f93dc75f7d50235c4
Author: gaobinlong <gb...@amazon.com>
AuthorDate: Wed Jun 3 00:34:46 2026 +0800
Introduce FirstPassGroupingCollectorManager (#15574)
This pr introduces `FirstPassGroupingCollectorManager` and switches `TestGrouping` and `BaseGroupSelectorTestCase` to use search concurrency and move away from the deprecated `search(Query, Collector)` method.
Relates to #12892.
Signed-off-by: Binlong Gao <gb...@amazon.com>
---
lucene/CHANGES.txt | 4 +-
.../FirstPassGroupingCollectorManager.java | 127 ++++++++
.../search/grouping/BaseGroupSelectorTestCase.java | 16 +-
.../lucene/search/grouping/TestGrouping.java | 350 +++++++++++++++------
4 files changed, 386 insertions(+), 111 deletions(-)
diff --git a/lucene/CHANGES.txt b/lucene/CHANGES.txt
index d46fe8a62a5..4795fa0e585 100644
--- a/lucene/CHANGES.txt
+++ b/lucene/CHANGES.txt
@@ -100,7 +100,7 @@ Improvements
* GITHUB#15976: Store DocValuesSkipper data in their own file. (Alan Woodward)
-* GITHUB#15574, GITHUB#15995: Introduce TopGroupsCollectorManager to parallelize search when using TopGroupsCollector. (Binlong Gao)
+* GITHUB#15603, GITHUB#15995: Introduce TopGroupsCollectorManager to parallelize search when using TopGroupsCollector. (Binlong Gao)
* GITHUB#15565: Introduce AllGroupHeadsCollectorManager to parallelize search when using AllGroupHeadsCollector. (Binlong Gao)
@@ -109,6 +109,8 @@ Improvements
* GITHUB#16000: Clarify that Accountable#ramBytesUsed() reports JVM heap memory only. (Luca Cavanna)
+* GITHUB#15574: Introduce FirstPassGroupingCollectorManager to parallelize search when using FirstPassGroupingCollector. (Binlong Gao)
+
* GITHUB#15660: Introduce LargeNumHitsTopDocsCollectorManager to parallelize search when using LargeNumHitsTopDocsCollector. (Binlong Gao)
* GITHUB#16088: SearchGroup#merge now returns an empty collection instead of null when topGroups is empty. (Luca Cavanna)
diff --git a/lucene/grouping/src/java/org/apache/lucene/search/grouping/FirstPassGroupingCollectorManager.java b/lucene/grouping/src/java/org/apache/lucene/search/grouping/FirstPassGroupingCollectorManager.java
new file mode 100644
index 00000000000..3c62ad12194
--- /dev/null
+++ b/lucene/grouping/src/java/org/apache/lucene/search/grouping/FirstPassGroupingCollectorManager.java
@@ -0,0 +1,127 @@
+/*
+ * 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.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.function.Supplier;
+import org.apache.lucene.search.CollectorManager;
+import org.apache.lucene.search.Sort;
+
+/**
+ * A CollectorManager implementation for {@link FirstPassGroupingCollector} that supports parallel
+ * collection and merges results across segments.
+ *
+ * <p>Example usage:
+ *
+ * <pre class="prettyprint">
+ * IndexSearcher searcher = new IndexSearcher(reader);
+ * Sort groupSort = Sort.RELEVANCE;
+ * int topNGroups = 10;
+ *
+ * FirstPassGroupingCollectorManager<BytesRef> manager =
+ * new FirstPassGroupingCollectorManager<>(
+ * () -> new TermGroupSelector("category"),
+ * groupSort,
+ * 0,
+ * topNGroups);
+ *
+ * Collection<SearchGroup<BytesRef>> searchGroups = searcher.search(query, manager);
+ *
+ * // searchGroups can then be passed to a second pass collector manager like TopGroupsCollectorManager for full group results
+ * </pre>
+ *
+ * @lucene.experimental
+ */
+public class FirstPassGroupingCollectorManager<T>
+ implements CollectorManager<FirstPassGroupingCollector<T>, Collection<SearchGroup<T>>> {
+
+ private final Supplier<GroupSelector<T>> groupSelectorFactory;
+ private final Sort groupSort;
+ private final int groupOffset;
+ private final int topNGroups;
+ private final boolean ignoreDocsWithoutGroupField;
+
+ /**
+ * Creates a new FirstPassGroupingCollectorManager.
+ *
+ * @param groupSelectorFactory factory to create group selectors for each collector
+ * @param groupSort the sort to use for groups
+ * @param groupOffset The offset in the collected groups
+ * @param topNGroups the number of top groups to collect
+ */
+ public FirstPassGroupingCollectorManager(
+ Supplier<GroupSelector<T>> groupSelectorFactory,
+ Sort groupSort,
+ int groupOffset,
+ int topNGroups) {
+ this(groupSelectorFactory, groupSort, groupOffset, topNGroups, false);
+ }
+
+ /**
+ * Creates a new FirstPassGroupingCollectorManager.
+ *
+ * @param groupSelectorFactory factory to create group selectors for each collector
+ * @param groupSort the sort to use for groups
+ * @param groupOffset The offset in the collected groups
+ * @param topNGroups the number of top groups to collect
+ * @param ignoreDocsWithoutGroupField whether to ignore documents without a group field
+ */
+ public FirstPassGroupingCollectorManager(
+ Supplier<GroupSelector<T>> groupSelectorFactory,
+ Sort groupSort,
+ int groupOffset,
+ int topNGroups,
+ boolean ignoreDocsWithoutGroupField) {
+ if (groupOffset < 0) {
+ throw new IllegalArgumentException("groupOffset must be >= 0 (got " + groupOffset + ")");
+ }
+ if (topNGroups < 1) {
+ throw new IllegalArgumentException("topNGroups must be >= 1 (got " + topNGroups + ")");
+ }
+ this.groupSelectorFactory = groupSelectorFactory;
+ this.groupSort = groupSort;
+ this.groupOffset = groupOffset;
+ this.topNGroups = topNGroups;
+ this.ignoreDocsWithoutGroupField = ignoreDocsWithoutGroupField;
+ }
+
+ @Override
+ public FirstPassGroupingCollector<T> newCollector() throws IOException {
+ return new FirstPassGroupingCollector<>(
+ groupSelectorFactory.get(),
+ groupSort,
+ groupOffset + topNGroups,
+ ignoreDocsWithoutGroupField);
+ }
+
+ @Override
+ public Collection<SearchGroup<T>> reduce(Collection<FirstPassGroupingCollector<T>> collectors)
+ throws IOException {
+ final List<Collection<SearchGroup<T>>> allGroups = new ArrayList<>();
+ for (FirstPassGroupingCollector<T> collector : collectors) {
+ Collection<SearchGroup<T>> groups = collector.getTopGroups(0);
+ if (groups != null) {
+ allGroups.add(groups);
+ }
+ }
+
+ return SearchGroup.merge(allGroups, groupOffset, topNGroups, groupSort);
+ }
+}
diff --git a/lucene/grouping/src/test/org/apache/lucene/search/grouping/BaseGroupSelectorTestCase.java b/lucene/grouping/src/test/org/apache/lucene/search/grouping/BaseGroupSelectorTestCase.java
index c5d16009c6b..1743b96f2c9 100644
--- a/lucene/grouping/src/test/org/apache/lucene/search/grouping/BaseGroupSelectorTestCase.java
+++ b/lucene/grouping/src/test/org/apache/lucene/search/grouping/BaseGroupSelectorTestCase.java
@@ -309,17 +309,17 @@ public abstract class BaseGroupSelectorTestCase<T> extends AbstractGroupingTestC
// A grouped query run in two phases against the control should give us the same
// result as the query run against shards and merged back together after each phase.
- FirstPassGroupingCollector<T> singletonFirstPass =
- new FirstPassGroupingCollector<>(getGroupSelector(), sort, 5);
- control.getIndexSearcher().search(topLevel, singletonFirstPass);
- Collection<SearchGroup<T>> singletonGroups = singletonFirstPass.getTopGroups(0);
+ FirstPassGroupingCollectorManager<T> firstPassGroupingCollectorManager =
+ new FirstPassGroupingCollectorManager<>(this::getGroupSelector, sort, 0, 5);
+ Collection<SearchGroup<T>> singletonGroups =
+ control.getIndexSearcher().search(topLevel, firstPassGroupingCollectorManager);
List<Collection<SearchGroup<T>>> shardGroups = new ArrayList<>();
for (Shard shard : shards) {
- FirstPassGroupingCollector<T> fc =
- new FirstPassGroupingCollector<>(getGroupSelector(), sort, 5);
- shard.getIndexSearcher().search(topLevel, fc);
- shardGroups.add(fc.getTopGroups(0));
+ FirstPassGroupingCollectorManager<T> fcm =
+ new FirstPassGroupingCollectorManager<>(this::getGroupSelector, sort, 0, 5);
+ Collection<SearchGroup<T>> topGroups = shard.getIndexSearcher().search(topLevel, fcm);
+ shardGroups.add(topGroups);
}
Collection<SearchGroup<T>> mergedGroups = SearchGroup.merge(shardGroups, 0, 5, sort);
assertEquals(singletonGroups, mergedGroups);
diff --git a/lucene/grouping/src/test/org/apache/lucene/search/grouping/TestGrouping.java b/lucene/grouping/src/test/org/apache/lucene/search/grouping/TestGrouping.java
index c9ed7bf5b02..341a96de006 100644
--- a/lucene/grouping/src/test/org/apache/lucene/search/grouping/TestGrouping.java
+++ b/lucene/grouping/src/test/org/apache/lucene/search/grouping/TestGrouping.java
@@ -46,6 +46,7 @@ import org.apache.lucene.queries.function.ValueSource;
import org.apache.lucene.queries.function.valuesource.BytesRefFieldSource;
import org.apache.lucene.search.CachingCollector;
import org.apache.lucene.search.Collector;
+import org.apache.lucene.search.CollectorManager;
import org.apache.lucene.search.DocIdSetIterator;
import org.apache.lucene.search.FieldDoc;
import org.apache.lucene.search.IndexSearcher;
@@ -57,6 +58,8 @@ import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.search.Sort;
import org.apache.lucene.search.SortField;
import org.apache.lucene.search.TermQuery;
+import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.search.TopFieldDocs;
import org.apache.lucene.search.TotalHits;
import org.apache.lucene.search.Weight;
import org.apache.lucene.search.similarities.BM25Similarity;
@@ -145,15 +148,28 @@ public class TestGrouping extends LuceneTestCase {
final Sort groupSort = Sort.RELEVANCE;
- final FirstPassGroupingCollector<?> c1 =
- createRandomFirstPassCollector(groupField, groupSort, 10);
- indexSearcher.search(new TermQuery(new Term("content", "random")), c1);
+ // TermGroupSelector or ValueSourceGroupSelector
+ final boolean isTermGroupSelector = random().nextBoolean();
+ final FirstPassGroupingCollectorManager<?> firstPassGroupingCollectorManager =
+ createFirstPassCollectorManager(isTermGroupSelector, groupField, groupSort, 0, 10);
+ Collection<?> topGroups =
+ indexSearcher.search(
+ new TermQuery(new Term("content", "random")), firstPassGroupingCollectorManager);
+
+ final TopGroupsCollectorManager<?> topGroupsCollectorManager =
+ createSecondPassCollectorManager(
+ isTermGroupSelector,
+ groupField,
+ toByteRefSearchGroups(topGroups),
+ groupSort,
+ Sort.RELEVANCE,
+ 0,
+ 5,
+ true);
+ final TopGroups<?> groups =
+ indexSearcher.search(
+ new TermQuery(new Term("content", "random")), topGroupsCollectorManager);
- final TopGroupsCollector<?> c2 =
- createSecondPassCollector(c1, groupSort, Sort.RELEVANCE, 0, 5, true);
- indexSearcher.search(new TermQuery(new Term("content", "random")), c2);
-
- final TopGroups<?> groups = c2.getTopGroups(0);
assertFalse(Float.isNaN(groups.maxScore));
assertEquals(7, groups.totalHitCount);
@@ -223,19 +239,20 @@ public class TestGrouping extends LuceneTestCase {
IndexSearcher searcher = newSearcher(reader);
// Test default behavior (include null group)
- FirstPassGroupingCollector<BytesRef> collector1 =
- new FirstPassGroupingCollector<>(new TermGroupSelector(groupField), Sort.RELEVANCE, 10);
- searcher.search(MatchAllDocsQuery.INSTANCE, collector1);
- Collection<SearchGroup<BytesRef>> groups1 = collector1.getTopGroups(0);
+ FirstPassGroupingCollectorManager<BytesRef> firstPassGroupingCollectorManager1 =
+ new FirstPassGroupingCollectorManager<>(
+ () -> new TermGroupSelector(groupField), Sort.RELEVANCE, 0, 10);
+ Collection<SearchGroup<BytesRef>> groups1 =
+ searcher.search(MatchAllDocsQuery.INSTANCE, firstPassGroupingCollectorManager1);
assertEquals(3, groups1.size()); // Should include null group
// Test ignoring docs without group field
- FirstPassGroupingCollector<BytesRef> collector2 =
- new FirstPassGroupingCollector<>(
- new TermGroupSelector(groupField), Sort.RELEVANCE, 10, true);
- searcher.search(MatchAllDocsQuery.INSTANCE, collector2);
- Collection<SearchGroup<BytesRef>> groups2 = collector2.getTopGroups(0);
+ FirstPassGroupingCollectorManager<BytesRef> firstPassGroupingCollectorManager2 =
+ new FirstPassGroupingCollectorManager<>(
+ () -> new TermGroupSelector(groupField), Sort.RELEVANCE, 0, 10, true);
+ Collection<SearchGroup<BytesRef>> groups2 =
+ searcher.search(MatchAllDocsQuery.INSTANCE, firstPassGroupingCollectorManager2);
assertEquals(2, groups2.size()); // Should exclude null group
@@ -261,12 +278,57 @@ public class TestGrouping extends LuceneTestCase {
IndexSearcher searcher = newSearcher(reader);
// Test ignoring docs without group field when all docs lack the field
- FirstPassGroupingCollector<BytesRef> collector =
- new FirstPassGroupingCollector<>(new TermGroupSelector("group"), Sort.RELEVANCE, 10, true);
- searcher.search(MatchAllDocsQuery.INSTANCE, collector);
- Collection<SearchGroup<BytesRef>> groups = collector.getTopGroups(0);
+ FirstPassGroupingCollectorManager<BytesRef> firstPassGroupingCollectorManager =
+ new FirstPassGroupingCollectorManager<>(
+ () -> new TermGroupSelector("group"), Sort.RELEVANCE, 0, 10, true);
+ Collection<SearchGroup<BytesRef>> groups =
+ searcher.search(MatchAllDocsQuery.INSTANCE, firstPassGroupingCollectorManager);
- assertNull(groups); // Should return null when no groups found
+ assertNotNull(groups);
+ assertTrue(groups.isEmpty());
+
+ reader.close();
+ dir.close();
+ }
+
+ public void testFirstPassGroupingCollectorManagerConstructor() {
+ IllegalArgumentException e1 =
+ expectThrows(
+ IllegalArgumentException.class,
+ () ->
+ new FirstPassGroupingCollectorManager<>(
+ () -> new TermGroupSelector("group"), Sort.RELEVANCE, -1, 10));
+ assertTrue(e1.getMessage().contains("groupOffset must be >= 0"));
+
+ IllegalArgumentException e2 =
+ expectThrows(
+ IllegalArgumentException.class,
+ () ->
+ new FirstPassGroupingCollectorManager<>(
+ () -> new TermGroupSelector("group"), Sort.RELEVANCE, 0, 0));
+ assertTrue(e2.getMessage().contains("topNGroups must be >= 1"));
+ }
+
+ public void testFirstPassGroupingCollectorManagerReduceEmptyResult() throws IOException {
+ Directory dir = newDirectory();
+ RandomIndexWriter w =
+ new RandomIndexWriter(random(), dir, newIndexWriterConfig(new MockAnalyzer(random())));
+ Document doc = new Document();
+ addGroupField(doc, "author", "author1");
+ doc.add(new TextField("content", "random text", Field.Store.NO));
+ w.addDocument(doc);
+
+ DirectoryReader reader = w.getReader();
+ w.close();
+ IndexSearcher searcher = newSearcher(reader);
+
+ // query matches nothing — reduce should return empty, not null
+ FirstPassGroupingCollectorManager<?> manager =
+ createFirstPassCollectorManager(random().nextBoolean(), "author", Sort.RELEVANCE, 0, 10);
+ Collection<?> result =
+ searcher.search(new TermQuery(new Term("content", "nonexistent")), manager);
+ assertNotNull(result);
+ assertTrue(result.isEmpty());
reader.close();
dir.close();
@@ -298,28 +360,32 @@ public class TestGrouping extends LuceneTestCase {
String groupField, Sort groupSort, int topDocs) throws IOException {
if (random().nextBoolean()) {
ValueSource vs = new BytesRefFieldSource(groupField);
- return new FirstPassGroupingCollector<>(
- new ValueSourceGroupSelector(vs, new HashMap<>()), groupSort, topDocs);
+ return new FirstPassGroupingCollectorManager<>(
+ () -> new ValueSourceGroupSelector(vs, new HashMap<>()), groupSort, 0, topDocs)
+ .newCollector();
} else {
- return new FirstPassGroupingCollector<>(
- new TermGroupSelector(groupField), groupSort, topDocs);
+ return new FirstPassGroupingCollectorManager<>(
+ () -> new TermGroupSelector(groupField), groupSort, 0, topDocs)
+ .newCollector();
}
}
- private FirstPassGroupingCollector<?> createFirstPassCollector(
+ private FirstPassGroupingCollectorManager<?> createFirstPassCollectorManager(
+ boolean isTermGroupSelector,
String groupField,
Sort groupSort,
- int topDocs,
- FirstPassGroupingCollector<?> firstPassGroupingCollector)
- throws IOException {
- GroupSelector<?> selector = firstPassGroupingCollector.getGroupSelector();
- if (TermGroupSelector.class.isAssignableFrom(selector.getClass())) {
- ValueSource vs = new BytesRefFieldSource(groupField);
- return new FirstPassGroupingCollector<>(
- new ValueSourceGroupSelector(vs, new HashMap<>()), groupSort, topDocs);
+ int groupOffset,
+ int topNGroups) {
+ if (isTermGroupSelector) {
+ return new FirstPassGroupingCollectorManager<>(
+ () -> new TermGroupSelector(groupField), groupSort, groupOffset, topNGroups);
} else {
- return new FirstPassGroupingCollector<>(
- new TermGroupSelector(groupField), groupSort, topDocs);
+ ValueSource vs = new BytesRefFieldSource(groupField);
+ return new FirstPassGroupingCollectorManager<>(
+ () -> new ValueSourceGroupSelector(vs, new HashMap<>()),
+ groupSort,
+ groupOffset,
+ topNGroups);
}
}
@@ -344,24 +410,25 @@ public class TestGrouping extends LuceneTestCase {
}
// Basically converts searchGroups from MutableValue to BytesRef if grouping by ValueSource
- @SuppressWarnings("unchecked")
- private TopGroupsCollector<?> createSecondPassCollector(
- FirstPassGroupingCollector<?> firstPassGroupingCollector,
+ private TopGroupsCollectorManager<?> createSecondPassCollectorManager(
+ boolean isTermGroupSelector,
String groupField,
Collection<SearchGroup<BytesRef>> searchGroups,
Sort groupSort,
Sort sortWithinGroup,
+ int withinGroupOffset,
int maxDocsPerGroup,
boolean getMaxScores)
throws IOException {
- if (firstPassGroupingCollector
- .getGroupSelector()
- .getClass()
- .isAssignableFrom(TermGroupSelector.class)) {
- GroupSelector<BytesRef> selector =
- (GroupSelector<BytesRef>) firstPassGroupingCollector.getGroupSelector();
- return new TopGroupsCollector<>(
- selector, searchGroups, groupSort, sortWithinGroup, maxDocsPerGroup, getMaxScores);
+ if (isTermGroupSelector) {
+ return new TopGroupsCollectorManager<>(
+ () -> new TermGroupSelector(groupField),
+ searchGroups,
+ groupSort,
+ sortWithinGroup,
+ withinGroupOffset,
+ maxDocsPerGroup,
+ getMaxScores);
} else {
ValueSource vs = new BytesRefFieldSource(groupField);
List<SearchGroup<MutableValue>> mvalSearchGroups = new ArrayList<>(searchGroups.size());
@@ -377,9 +444,14 @@ public class TestGrouping extends LuceneTestCase {
sg.sortValues = mergedTopGroup.sortValues;
mvalSearchGroups.add(sg);
}
- ValueSourceGroupSelector selector = new ValueSourceGroupSelector(vs, new HashMap<>());
- return new TopGroupsCollector<>(
- selector, mvalSearchGroups, groupSort, sortWithinGroup, maxDocsPerGroup, getMaxScores);
+ return new TopGroupsCollectorManager<>(
+ () -> new ValueSourceGroupSelector(vs, new HashMap<>()),
+ mvalSearchGroups,
+ groupSort,
+ sortWithinGroup,
+ withinGroupOffset,
+ maxDocsPerGroup,
+ getMaxScores);
}
}
@@ -443,6 +515,55 @@ public class TestGrouping extends LuceneTestCase {
return null;
}
+ @SuppressWarnings("unchecked")
+ private Collection<SearchGroup<BytesRef>> toByteRefSearchGroups(Collection<?> groups) {
+ if (groups == null
+ || groups.isEmpty()
+ || !(((SearchGroup<?>) groups.iterator().next()).groupValue instanceof MutableValue)) {
+ return (Collection<SearchGroup<BytesRef>>) groups;
+ }
+ List<SearchGroup<BytesRef>> result = new ArrayList<>(groups.size());
+ for (Object g : groups) {
+ SearchGroup<?> sg = (SearchGroup<?>) g;
+ SearchGroup<BytesRef> out = new SearchGroup<>();
+ MutableValue mv = (MutableValue) sg.groupValue;
+ out.groupValue = mv.exists() ? ((MutableValueStr) mv).value.get() : null;
+ out.sortValues = sg.sortValues;
+ result.add(out);
+ }
+ return result;
+ }
+
+ @SuppressWarnings("unchecked")
+ private TopGroups<BytesRef> toByteTopGroups(TopGroups<?> topGroups) {
+ if (topGroups == null
+ || topGroups.groups.length == 0
+ || !(topGroups.groups[0].groupValue() instanceof MutableValue)) {
+ return (TopGroups<BytesRef>) topGroups;
+ }
+ List<GroupDocs<BytesRef>> groups = new ArrayList<>(topGroups.groups.length);
+ final GroupDocs<MutableValue>[] mvalTopGroups = (GroupDocs<MutableValue>[]) topGroups.groups;
+ for (GroupDocs<MutableValue> mvalGd : mvalTopGroups) {
+ BytesRef groupValue =
+ mvalGd.groupValue().exists() ? ((MutableValueStr) mvalGd.groupValue()).value.get() : null;
+ groups.add(
+ new GroupDocs<>(
+ Float.NaN,
+ mvalGd.maxScore(),
+ mvalGd.totalHits(),
+ mvalGd.scoreDocs(),
+ groupValue,
+ mvalGd.groupSortValues()));
+ }
+ return new TopGroups<>(
+ topGroups.groupSort,
+ topGroups.withinGroupSort,
+ topGroups.totalHitCount,
+ topGroups.totalGroupedHitCount,
+ groups.toArray(GroupDocs[]::new),
+ Float.NaN);
+ }
+
@SuppressWarnings({"unchecked", "rawtypes"})
private TopGroups<BytesRef> getTopGroups(TopGroupsCollector c, int withinGroupOffset) {
if (c.getGroupSelector().getClass().isAssignableFrom(TermGroupSelector.class)) {
@@ -1152,7 +1273,6 @@ public class TestGrouping extends LuceneTestCase {
final TopGroups<BytesRef> topGroupsShards =
searchShards(
- s,
shards.subSearchers,
query,
groupSort,
@@ -1361,7 +1481,6 @@ public class TestGrouping extends LuceneTestCase {
// Get shard'd block grouping result:
final TopGroups<BytesRef> topGroupsBlockShards =
searchShards(
- sBlocks,
shardsBlocks.subSearchers,
query,
groupSort,
@@ -1443,7 +1562,6 @@ public class TestGrouping extends LuceneTestCase {
}
private TopGroups<BytesRef> searchShards(
- IndexSearcher topSearcher,
ShardSearcher[] subSearchers,
Query query,
Sort groupSort,
@@ -1470,59 +1588,43 @@ public class TestGrouping extends LuceneTestCase {
+ canUseIDV);
}
// Run 1st pass collector to get top groups per shard
- final Weight w =
- topSearcher.createWeight(
- topSearcher.rewrite(query),
- groupSort.needsScores() || docSort.needsScores() || getMaxScores
- ? ScoreMode.COMPLETE
- : ScoreMode.COMPLETE_NO_SCORES,
- 1);
final List<Collection<SearchGroup<BytesRef>>> shardGroups = new ArrayList<>();
- List<FirstPassGroupingCollector<?>> firstPassGroupingCollectors = new ArrayList<>();
- FirstPassGroupingCollector<?> firstPassCollector = null;
-
String groupField = "group";
+ // TermGroupSelector or ValueSourceGroupSelector
+ final boolean isTermGroupSelector = random().nextBoolean();
+
for (int shardIDX = 0; shardIDX < subSearchers.length; shardIDX++) {
+ final FirstPassGroupingCollectorManager<?> firstPassGroupingCollectorManager =
+ createFirstPassCollectorManager(
+ isTermGroupSelector, groupField, groupSort, 0, groupOffset + topNGroups);
- // First shard determines whether we use IDV or not;
- // all other shards match that:
- if (firstPassCollector == null) {
- firstPassCollector =
- createRandomFirstPassCollector(groupField, groupSort, groupOffset + topNGroups);
- } else {
- firstPassCollector =
- createFirstPassCollector(
- groupField, groupSort, groupOffset + topNGroups, firstPassCollector);
- }
if (VERBOSE) {
System.out.println(" shard=" + shardIDX + " groupField=" + groupField);
- System.out.println(" 1st pass collector=" + firstPassCollector);
+ System.out.println(" 1st pass collector manager=" + firstPassGroupingCollectorManager);
}
- firstPassGroupingCollectors.add(firstPassCollector);
- subSearchers[shardIDX].search(w, firstPassCollector);
- final Collection<SearchGroup<BytesRef>> topGroups = getSearchGroups(firstPassCollector, 0);
- if (topGroups != null) {
- if (VERBOSE) {
+ final Collection<SearchGroup<BytesRef>> topGroups =
+ toByteRefSearchGroups(
+ subSearchers[shardIDX].search(query, firstPassGroupingCollectorManager));
+ if (VERBOSE) {
+ System.out.println(
+ " shard "
+ + shardIDX
+ + " s="
+ + subSearchers[shardIDX]
+ + " totalGroupedHitCount=?"
+ + " "
+ + topGroups.size()
+ + " groups:");
+ for (SearchGroup<BytesRef> group : topGroups) {
System.out.println(
- " shard "
- + shardIDX
- + " s="
- + subSearchers[shardIDX]
- + " totalGroupedHitCount=?"
- + " "
- + topGroups.size()
- + " groups:");
- for (SearchGroup<BytesRef> group : topGroups) {
- System.out.println(
- " "
- + groupToString(group.groupValue)
- + " groupSort="
- + Arrays.toString(group.sortValues));
- }
+ " "
+ + groupToString(group.groupValue)
+ + " groupSort="
+ + Arrays.toString(group.sortValues));
}
- shardGroups.add(topGroups);
}
+ shardGroups.add(topGroups);
}
final Collection<SearchGroup<BytesRef>> mergedTopGroups =
@@ -1544,17 +1646,20 @@ public class TestGrouping extends LuceneTestCase {
@SuppressWarnings({"unchecked", "rawtypes"})
final TopGroups<BytesRef>[] shardTopGroups = new TopGroups[subSearchers.length];
for (int shardIDX = 0; shardIDX < subSearchers.length; shardIDX++) {
- final TopGroupsCollector<?> secondPassCollector =
- createSecondPassCollector(
- firstPassGroupingCollectors.get(shardIDX),
+ final TopGroupsCollectorManager<?> secondPassCollectorManager =
+ createSecondPassCollectorManager(
+ isTermGroupSelector,
groupField,
mergedTopGroups,
groupSort,
docSort,
+ 0,
docOffset + topNDocs,
getMaxScores);
- subSearchers[shardIDX].search(w, secondPassCollector);
- shardTopGroups[shardIDX] = getTopGroups(secondPassCollector, 0);
+
+ final TopGroups<?> topGroupsRaw =
+ subSearchers[shardIDX].search(query, secondPassCollectorManager);
+ shardTopGroups[shardIDX] = toByteTopGroups(topGroupsRaw);
if (VERBOSE) {
System.out.println(
" " + shardTopGroups[shardIDX].groups.length + " shard[" + shardIDX + "] groups:");
@@ -1674,8 +1779,49 @@ public class TestGrouping extends LuceneTestCase {
this.ctx = ctx;
}
- public void search(Weight weight, Collector collector) throws IOException {
+ @Override
+ public <C extends Collector, T> T search(Query query, CollectorManager<C, T> collectorManager)
+ throws IOException {
+ final C collector = collectorManager.newCollector();
+ final Weight weight = createWeight(rewrite(query), collector.scoreMode(), 1);
searchLeaf(ctx, 0, DocIdSetIterator.NO_MORE_DOCS, weight, collector);
+ return collectorManager.reduce(Collections.singletonList(collector));
+ }
+
+ @Override
+ public void search(Query query, Collector collector) {
+ throw new UnsupportedOperationException("Use search(Weight, CollectorManager) instead");
+ }
+
+ @Override
+ public TopDocs search(Query query, int n) {
+ throw new UnsupportedOperationException("Use search(Weight, CollectorManager) instead");
+ }
+
+ @Override
+ public TopFieldDocs search(Query query, int n, Sort sort) {
+ throw new UnsupportedOperationException("Use search(Weight, CollectorManager) instead");
+ }
+
+ @Override
+ public TopFieldDocs search(Query query, int n, Sort sort, boolean doDocScores) {
+ throw new UnsupportedOperationException("Use search(Weight, CollectorManager) instead");
+ }
+
+ @Override
+ public TopDocs searchAfter(ScoreDoc after, Query query, int n) {
+ throw new UnsupportedOperationException("Not implemented yet");
+ }
+
+ @Override
+ public TopDocs searchAfter(ScoreDoc after, Query query, int n, Sort sort) {
+ throw new UnsupportedOperationException("Not implemented yet");
+ }
+
+ @Override
+ public TopFieldDocs searchAfter(
+ ScoreDoc after, Query query, int n, Sort sort, boolean doDocScores) {
+ throw new UnsupportedOperationException("Not implemented yet");
}
@Override