[Hybrid Query] Gate collapse distinct-groups collection behind an opt-in index setting

[Ystk-hsn]

Description

Implements the decision from #1947: both collapse behaviors make sense but are mutually exclusive, so the existing behavior stays the default and the distinct-groups collection is opt-in. This adds

index.neural_search.hybrid_collapse_distinct_groups_enabled   (boolean, default false, IndexScope + Dynamic)

By default nothing changes — the collector keeps the top-size documents per sub-query, preserving score parity with the same hybrid query without collapse. Over the issue's dataset (6 groups on one shard, groupA owning the three top-scoring documents), size: 5 returns 4 hits:

groupA, groupB, groupC, groupD        (groupE crowded out after deduplication)

Turning the setting on (dynamic, takes effect on the next request):

PUT /my-index/_settings
{"index.neural_search.hybrid_collapse_distinct_groups_enabled": true}

makes the same search return 5 hits:

groupA, groupB, groupC, groupD, groupE

The two modes cannot be combined: they hand different document sets to per-sub-query normalization, so the same document can get different normalized scores under each mode — which is why this is a switch rather than a fix of the default.

Implementation notes

• HybridCollapsingTopGroupsCollector (new) — collects the top-size distinct groups per sub-query, mirroring the bookkeeping of Lucene's FirstPassGroupingCollector: a group map plus an ordered set once size groups exist, evicting the weakest group when a new competitive one arrives. When sorting by score, comparators are wrapped with HybridLeafFieldComparator so they read the sub-query's individual score rather than the sum. Emits one FieldDoc per surviving group, in the same CollapseTopFieldDocs encoding the existing collector uses — the coordinator-side normalization/deduplication pipeline is unchanged and no version gating is needed.
• HybridCollectorFactory — reads the setting and picks the collector; this is the only decision point, next to where the deprecated hybrid_collapse_docs_per_group_per_subquery setting is already read.
• HybridCollectorManager — registers the new collector type.
• HybridCollapsingTopDocsCollector, HybridLeafFieldComparator and the existing collector tests are byte-identical to main — the default path is untouched.

Testing

• Unit: HybridCollapsingTopDocsCollectorTests (default behavior, identical to main), HybridCollapsingTopGroupsCollectorTests (opt-in behavior), factory tests for the setting switch
• Integ: HybridCollapseIT covers both modes over the same skewed-groups dataset, including a test pinning the by-design default from the issue discussion

Related Issues

Resolves #1947

Check List

• New functionality includes testing.
• New functionality has been documented.
• API changes companion pull request created.
• Commits are signed per the DCO using --signoff.
• Public documentation issue/PR created.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit bc1a0ac)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ Recommended focus areas for review Possible Issue In topDocs(), per-sub-query representatives are filtered out when subQueryScore <= 0. If a matching sub-query legitimately produces a score of exactly 0 (e.g., certain function_score or rank_feature configurations), that representative will be silently dropped from that sub-query's list, and collectedHitsPerSubQuery will be inconsistent with the emitted docs. Consider tracking sub-query matches explicitly (e.g., a boolean per sub-query recorded at collect time) rather than inferring a match from score>0. for (CollectedGroup<T> group : orderedGroups) { float subQueryScore = group.scoresPerSubQuery[subQuery]; if (subQueryScore <= 0) { continue; } Object[] fields = new Object[numComparators]; for (int k = 0; k < numComparators; k++) { fields[k] = comparators[k].value(group.comparatorSlot); } fieldDocs.add(new FieldDoc(group.topDoc, subQueryScore, fields)); if (group.groupValue instanceof BytesRef) { collapseValues.add(BytesRef.deepCopyOf((BytesRef) group.groupValue)); } else { collapseValues.add(group.groupValue); } } Possible Issue scoresPerSubQuery = subScoresByQuery.clone() is executed on every collectExistingGroup call to stash sub-query scores for the doc being compared. However, this happens BEFORE the tie-break/comparison decides whether the new doc actually wins. On a losing comparison (c < 0 at the last comparator, or tie), the method returns without reverting group.scoresPerSubQuery, so a losing doc's per-sub-query scores can overwrite the winning representative's scores. Move the scoresPerSubQuery assignment to after the win is confirmed (alongside the topDoc and slot swap). private void collectExistingGroup(int doc, float[] subScoresByQuery, CollectedGroup<T> group) throws IOException { for (int compIDX = 0;; compIDX++) { leafComparators[compIDX].copy(spareSlot, doc); final int c = reversed[compIDX] * comparators[compIDX].compare(group.comparatorSlot, spareSlot); if (c < 0) { return; } else if (c > 0) { for (int compIDX2 = compIDX + 1; compIDX2 < comparators.length; compIDX2++) { leafComparators[compIDX2].copy(spareSlot, doc); } break; } else if (compIDX == compIDXEnd) { // Ties lose: docs are visited in doc id order return; } } // Remove before mutating — the sorted set locates elements by comparing slots if (Objects.nonNull(orderedGroups)) { orderedGroups.remove(group); } group.topDoc = docBase + doc; group.scoresPerSubQuery = subScoresByQuery.clone(); // The staged spare slot becomes the group's slot, the old slot becomes spare final int tmpSlot = spareSlot; spareSlot = group.comparatorSlot; group.comparatorSlot = tmpSlot; if (Objects.nonNull(orderedGroups)) { orderedGroups.add(group); setBottomToWeakestGroup(); } }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to bc1a0ac

Explore these optional code suggestions:

Category	Suggestion	Impact
General	Distinguish unmatched from zero-score hits Filtering with subQueryScore <= 0 will incorrectly drop legitimate hits whose relevance score is exactly 0.0f (e.g., function_score queries or constant_score with zero boost). Since the elected representative was matched by this sub-query only if its score is strictly greater than 0 during collection (tracked via scoresPerSubQuery initialized to 0 and updated by clone from subScoresByQuery), consider using a distinct sentinel or an explicit "matched" bitset to distinguish "did not match" from "matched with score 0". src/main/java/org/opensearch/neuralsearch/search/collector/HybridCollapsingTopGroupsCollector.java [169-172] for (CollectedGroup<T> group : orderedGroups) { float subQueryScore = group.scoresPerSubQuery[subQuery]; - if (subQueryScore <= 0) { + // Skip only when this sub-query did not match the representative at all + if (subQueryScore <= 0f) { continue; } Suggestion importance[1-10]: 5 __ Why: Valid concern: using <= 0 to detect "did not match" can drop legitimate zero-score hits from function_score queries. However, the suggestion's improved_code is identical to existing_code and only describes the issue rather than fixing it, reducing its actionable value.	Low
	Verify totalHits semantics per sub-query collectedHitsPerSubQuery is incremented for every collected document, but this count is later reported as TotalHits per sub-query in topDocs(). Since the collector collapses documents into groups, this per-sub-query totalHits value reflects raw matches, not distinct groups, which may confuse downstream consumers expecting totalHits to align with scoreDocs.length semantics used by the collapse contract. Verify this matches the semantics expected by the normalization pipeline. src/main/java/org/opensearch/neuralsearch/search/collector/HybridCollapsingTopGroupsCollector.java [238-244] +for (int subQuery = 0; subQuery < subScoresByQuery.length; subQuery++) { + float score = subScoresByQuery[subQuery]; + if (score > 0) { + collectedHitsPerSubQuery[subQuery]++; + maxScore = Math.max(score, maxScore); + } +} - Suggestion importance[1-10]: 2 __ Why: The suggestion only asks the author to verify existing behavior without proposing a concrete change, and improved_code is identical to existing_code. Low impact.	Low

---

Previous suggestions

Suggestions up to commit 44d5bfb

Category	Suggestion	Impact
Possible issue	Fix inverted score threshold skip condition The threshold skip condition uses score <= 0 which is inverted — a doc should be skipped when its score is BELOW the minimum threshold regardless of sign, not only when it is non-positive. The score <= 0 clause makes the threshold effectively unreachable for typical positive scores, defeating the pruning optimization. Remove the score <= 0 condition so any score below the threshold is skipped. src/main/java/org/opensearch/neuralsearch/search/collector/HybridCollapsingTopGroupsCollector.java [208-210] -if (isSortByScore && score <= 0 && score < minScoreThresholds[subQuery]) { +if (isSortByScore && score < minScoreThresholds[subQuery]) { continue; } Suggestion importance[1-10]: 6 __ Why: The score <= 0 clause combined with score < minScoreThresholds[subQuery] is suspicious since thresholds are updated with positive scores during eviction, making the pruning path only fire for non-positive scores. However, the impact is limited to a missed optimization rather than correctness, and the exact intent may need verification.	Low
	Use negative infinity for initial threshold Float.MIN_VALUE is the smallest positive normal float (~1.4E-45), not the most negative float. Using it as an initial score threshold means any doc with a non-positive score would be considered below the threshold before any eviction occurs. Use Float.NEGATIVE_INFINITY (or -Float.MAX_VALUE) so the initial threshold admits all scores. src/main/java/org/opensearch/neuralsearch/search/collector/HybridCollapsingTopGroupsCollector.java [179-181] minScoreThresholds = new float[getCompoundQueryScorer().getNumOfSubQueries()]; -Arrays.fill(minScoreThresholds, Float.MIN_VALUE); +Arrays.fill(minScoreThresholds, Float.NEGATIVE_INFINITY); Suggestion importance[1-10]: 6 __ Why: Float.MIN_VALUE is indeed the smallest positive value, not the most negative, so using it as an initial "admit all" threshold is technically incorrect. However, since the pruning skip is also guarded by score <= 0, the practical impact is limited, but the fix improves correctness of the threshold semantics.	Low
General	Validate topNGroups is positive topNGroups is used as the size of a bounded queue and as slot indices (numHits + 1), but there is no validation that it is positive. A zero or negative value would produce invalid comparator slot allocations and undefined behavior. Add a precondition check to fail fast with a clear error. src/main/java/org/opensearch/neuralsearch/search/collector/HybridCollapsingTopGroupsCollector.java [72-78] HybridCollapsingTopGroupsCollector( GroupSelector<T> groupSelector, String collapseField, @NonNull Sort groupSort, int topNGroups, HitsThresholdChecker hitsThresholdChecker ) { + if (topNGroups < 1) { + throw new IllegalArgumentException("topNGroups must be >= 1 (got " + topNGroups + ")"); + } Suggestion importance[1-10]: 3 __ Why: Adding input validation is a minor defensive improvement; the caller HybridCollectorFactory typically ensures a valid numHits, so the practical impact is low.	Low

Suggestions up to commit 0cbec23

Category	Suggestion	Impact
Possible issue	Re-establish bottom slot on new segments When setNextReader is called for a new segment, orderedGroups may already reference bottom slots from the previous segment's comparators. The new leaf comparators are created without calling setBottom on the weakest slot, so isCompetitive/compareBottom may operate on an uninitialized bottom for the new leaf. After creating the new leaf comparators, if orderedGroups is non-null, call setBottomToWeakestGroup() to re-establish the bottom on the new leaf. src/main/java/org/opensearch/neuralsearch/search/collector/HybridCollapsingTopDocsCollector.java [278-293] void setNextReader(LeafReaderContext ctx, HybridSubQueryScorer compoundQueryScorer) throws IOException { leafComparators = new LeafFieldComparator[comparators.length]; scoreLeafComparator = null; SortField[] sortFields = sort.getSort(); for (int i = 0; i < comparators.length; i++) { LeafFieldComparator leafComparator = comparators[i].getLeafComparator(ctx); + if (SortField.Type.SCORE.equals(sortFields[i].getType())) { + HybridLeafFieldComparator wrappedComparator = new HybridLeafFieldComparator(leafComparator); + scoreLeafComparator = wrappedComparator; + leafComparator = wrappedComparator; + } + leafComparator.setScorer(compoundQueryScorer); + leafComparators[i] = leafComparator; + } + if (Objects.nonNull(orderedGroups)) { + setBottomToWeakestGroup(); + } +} Suggestion importance[1-10]: 7 __ Why: Valid concern: when transitioning to a new segment via setNextReader, the new leaf comparators need setBottom called if orderedGroups is already established, otherwise compareBottom may operate on an uninitialized bottom slot for the new leaf, potentially causing incorrect competitiveness checks.	Medium
	Fix unreachable min-score skip condition The condition score <= 0 && score < minScoreThresholds[subQuery] is likely incorrect and prevents the intended non-competitive doc skipping. Since a score == 0 case is already filtered above, the intent is probably to skip when score < minScoreThresholds[subQuery] (regardless of sign), similar to how HybridBulkScorer uses these thresholds. The <= 0 guard makes this branch unreachable for typical positive scores. src/main/java/org/opensearch/neuralsearch/search/collector/HybridCollapsingTopDocsCollector.java [200-202] -if (isSortByScore && score <= 0 && score < minScoreThresholds[subQuery]) { +if (isSortByScore && score < minScoreThresholds[subQuery]) { continue; } Suggestion importance[1-10]: 6 __ Why: The condition score <= 0 && score < minScoreThresholds[subQuery] does appear suspicious since scores are typically positive, making this branch effectively unreachable for competitive skipping. However, this condition existed in the pre-PR code as well, so it may be intentional or out of scope. The observation is valid but the fix requires domain verification.	Low
General	Handle null collapse group values groupSelector.currentValue() may return null when the document has no value for the collapse field. Using null as a HashMap key will conflate all such documents into a single pseudo-group, and groupSelector.copyValue() semantics for null are not well-defined. Consider skipping documents whose group value is null to avoid corrupt group bookkeeping. src/main/java/org/opensearch/neuralsearch/search/collector/HybridCollapsingTopDocsCollector.java [307-312] -CollectedGroup<T> group = groupMap.get(groupSelector.currentValue()); +T currentGroupValue = groupSelector.currentValue(); +if (Objects.isNull(currentGroupValue)) { + return; +} +CollectedGroup<T> group = groupMap.get(currentGroupValue); if (Objects.isNull(group)) { collectNewGroup(doc, score, compoundQueryScorer); } else { collectExistingGroup(doc, score, group); } Suggestion importance[1-10]: 4 __ Why: Reasonable defensive check for null group values, though GroupSelector typically handles missing values via a MISSING sentinel. Impact is uncertain without knowledge of the upstream groupSelector semantics.	Low

Suggestions up to commit 3a3e174

Category	Suggestion	Impact
Possible issue	Fix incorrect score threshold check The condition score <= 0 && score < minScoreThresholds[subQuery] is likely a bug — it only skips non-positive scores, but the intent (matching the removed logic) is to skip any score below the threshold. Since score == 0 is already handled above, this should be score < minScoreThresholds[subQuery] to properly prune non-competitive positive scores. src/main/java/org/opensearch/neuralsearch/search/collector/HybridCollapsingTopDocsCollector.java [200-202] -if (isSortByScore && score <= 0 && score < minScoreThresholds[subQuery]) { +if (isSortByScore && score < minScoreThresholds[subQuery]) { continue; } Suggestion importance[1-10]: 7 __ Why: The condition score <= 0 && score < minScoreThresholds[subQuery] appears preserved from prior code but is logically odd — since minScoreThresholds accumulates evicted scores (which can be positive), the score <= 0 guard likely prevents legitimate pruning of positive non-competitive scores. However, this was the pre-existing behavior and not introduced by the PR, so impact is moderate.	Medium
General	Skip documents without collapse value groupSelector.currentValue() may return null for documents that have no value for the collapse field (e.g., missing field). Using null as a HashMap key will collapse all such docs into a single group, which may be intended, but copyValue() later stores a null group value which then causes downstream issues (e.g., BytesRef cast, null in collapseValues array). Consider explicitly skipping docs without a group value. src/main/java/org/opensearch/neuralsearch/search/collector/HybridCollapsingTopDocsCollector.java [307-313] -CollectedGroup<T> group = groupMap.get(groupSelector.currentValue()); +T currentGroup = groupSelector.currentValue(); +if (currentGroup == null) { + return; +} +CollectedGroup<T> group = groupMap.get(currentGroup); if (Objects.isNull(group)) { collectNewGroup(doc, score, compoundQueryScorer); } else { collectExistingGroup(doc, score, group); } Suggestion importance[1-10]: 4 __ Why: Handling null collapse values is a reasonable defensive check, but it depends on how groupSelector behaves for missing fields. This is an edge case not clearly demonstrated as a bug in the PR.	Low
	Bound comparator iteration loop The loop uses compIDX without bounds checking against comparators.length. If compIDXEnd computation is ever off (e.g., empty sort array making it -1), this would loop forever or ArrayIndexOutOfBoundsException. Add a guard or ensure the sort array is non-empty at construction time. src/main/java/org/opensearch/neuralsearch/search/collector/HybridCollapsingTopDocsCollector.java [315-327] private boolean isCompetitive(int doc) throws IOException { - for (int compIDX = 0;; compIDX++) { + for (int compIDX = 0; compIDX < leafComparators.length; compIDX++) { final int c = reversed[compIDX] * leafComparators[compIDX].compareBottom(doc); Suggestion importance[1-10]: 2 __ Why: The unbounded loop mirrors Lucene's FirstPassGroupingCollector pattern and is safe because compIDXEnd is set from the sort fields array (which is guaranteed non-empty in this context). The concern is theoretical.	Low

Suggestions up to commit 06aa2f9

Category	Suggestion	Impact
Possible issue	Fix min-score threshold skip condition The condition score <= 0 && score < minScoreThresholds[subQuery] is likely incorrect and effectively dead — since scores of 0 are already filtered out just above, this branch only activates for negative scores. It should probably be score < minScoreThresholds[subQuery] to actually skip non-competitive docs based on the minimum score threshold. src/main/java/org/opensearch/neuralsearch/search/collector/HybridCollapsingTopDocsCollector.java [199-201] -if (isSortByScore && score <= 0 && score < minScoreThresholds[subQuery]) { +if (isSortByScore && score < minScoreThresholds[subQuery]) { continue; } Suggestion importance[1-10]: 6 __ Why: The observation is valid: with score == 0 already filtered above, the score <= 0 condition makes the threshold check effectively dead for non-negative scores. However, this appears to be pre-existing behavior, and the fix's correctness depends on whether scores can legitimately be negative in this context.	Low
General	Defensively deep-copy BytesRef group values The previous implementation explicitly deep-copied BytesRef group values before returning them. The comment claims GroupSelector#copyValue already deep-copies, but this depends on the specific GroupSelector implementation; for safety and to preserve prior behavior, consider deep-copying BytesRef values here to guard against mutation by downstream consumers or reuse of buffers. src/main/java/org/opensearch/neuralsearch/search/collector/HybridCollapsingTopDocsCollector.java [447-457] int index = 0; for (CollectedGroup<T> group : orderedGroups) { Object[] fields = new Object[numComparators]; for (int k = 0; k < numComparators; k++) { fields[k] = comparators[k].value(group.comparatorSlot); } fieldDocs[index] = new FieldDoc(group.topDoc, group.score, fields); - // Group values were deep-copied by GroupSelector#copyValue when stored - collapseValues[index] = group.groupValue; + Object groupValue = group.groupValue; + collapseValues[index] = (groupValue instanceof BytesRef) + ? BytesRef.deepCopyOf((BytesRef) groupValue) + : groupValue; index++; } Suggestion importance[1-10]: 5 __ Why: Reasonable defensive suggestion to preserve prior deep-copy behavior for BytesRef values, guarding against buffer reuse. Impact is moderate since GroupSelector#copyValue is generally expected to deep-copy already.	Low
	Verify cross-segment comparator slot handling setNextReader is invoked once via ensureLeafComparatorsInitialized guarded by leafComparatorsInitialized, but that flag lives per-LeafCollector. If the collector processes multiple segments (a new getLeafCollector per segment), each segment's collector creates its own flag — that's fine — but any groups collected in earlier segments still reference the previous segment's LeafFieldComparators until setNextReader is called on the same SubQueryGroupCollector. Verify that setBottom/copy invocations on new-segment leafComparators correctly interact with slots populated in prior segments (comparators must persist slot values across segment transitions). src/main/java/org/opensearch/neuralsearch/search/collector/HybridCollapsingTopDocsCollector.java [277-282] +void setNextReader(LeafReaderContext ctx, HybridSubQueryScorer compoundQueryScorer) throws IOException { + leafComparators = new LeafFieldComparator[comparators.length]; + scoreLeafComparator = null; + SortField[] sortFields = sort.getSort(); + for (int i = 0; i < comparators.length; i++) { + LeafFieldComparator leafComparator = comparators[i].getLeafComparator(ctx); - Suggestion importance[1-10]: 4 __ Why: Raises a valid concern about cross-segment state, but is a verification request without a concrete fix; also the improved_code is identical to existing_code.	Low

[github-actions]

Persistent review updated to latest commit 3a3e174

[codecov]

Codecov Report

❌ Patch coverage is 90.82569% with 20 lines in your changes missing coverage. Please review.
✅ Project coverage is 83.66%. Comparing base (3e52b21) to head (44d5bfb).

Files with missing lines	Patch %	Lines
.../collector/HybridCollapsingTopGroupsCollector.java	90.68%	9 Missing and 10 partials ⚠️
...earch/search/collector/HybridCollectorFactory.java	92.30%	0 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #1956      +/-   ##
============================================
+ Coverage     83.56%   83.66%   +0.10%     
- Complexity     3956     3970      +14     
============================================
  Files           298      299       +1     
  Lines         14039    14251     +212     
  Branches       2323     2367      +44     
============================================
+ Hits          11731    11923     +192     
- Misses         1475     1486      +11     
- Partials        833      842       +9     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Persistent review updated to latest commit 0cbec23

[github-actions]

Persistent review updated to latest commit 44d5bfb

[martin-gaievski]

Seems per-leg representative election still splits a group's score. CollectedGroup holds int topDoc, each SubQueryGroupCollector is per sub-query and its score comparator is wrapped to read that leg's individual score (HybridLeafFieldComparator fed by setCurrentSubQueryScore(score)). Election happens per leg in collectExistingGroup , same for collectNewGroup.
This scenario will be relatively common, it needs only 2+ legs and one group containing two documents that legs rank differently. Legs ranking documents differently is the point of hybrid search, so any multi-doc group will potential have this problem.

One way to fix this, and it's easier with the new opt-in mode - leg-independent election can go straight in new collector collect(int doc) method:

// in the outer collect(), before the per-sub-query loop
float fusionProxy = 0f;
for (float s : subScoresByQuery) fusionProxy += s;   // leg-independent ranking key

Pass that down so every SubQueryGroupCollector elects the same doc id for a group while group.score stays the leg's own score. Later join adds legs for one doc id instead of splitting them.

[martin-gaievski]

please increase coverage, currently there is no test that can catch the missing functionality from my other code level comment, existing tests are only for per-leg representative election.

• shared helper collectWithGroupDerivedScores writes only getSubQueryScores()[0], that's one leg, so there is nothing to disagree on
• testCollapse_whenMultipleSubQueries_thenEachSubQueryHasResults does build disagreeing legs ("high scores for even docs" / "high scores for odd docs"), so actual condition is present in its data. But its assertions are for existence-only, so it cannot fail on a wrong representative

Ideally would be great to have a two-leg IT asserting group order, not just group count

[Ystk-hsn]

Added in 5ff02f0:

• testCollapse_whenLegsDisagreeWithinGroup_thenSameRepresentativeElectedAcrossLegs, two legs ranking a group's documents differently, asserting the same representative in every leg with each leg's own score, and leg independent group order. This one fails on the previous code.
• testCollapse_whenGroupStrongInOneLegOnly_thenEvictionUsesSummedScore, eviction is also leg independent.
• testCollapse_whenSortByFieldAndLegsMatchDifferentDocs_thenSameRepresentativeElected, the same guarantee under field sort.
• The two leg IT you suggested, testCollapse_whenLegsDisagreeAndDistinctGroupsEnabled_thenGroupsOrderedByFusedScore, legs scoring different fields and the full group order asserted.
• testCollapse_whenMultipleSubQueries_thenEachSubQueryHasResults is updated to the new contract, a sub-query that matched none of the elected representatives emits an empty list while its total hits still count its own matches.

[github-actions]

Persistent review updated to latest commit bc1a0ac