SOLR-18195: CombinedQueryComponent supports Collapse now (#4277)

Support for using {!collapse} with CombinedQueryComponent (RRF)

Author: ercsonusharma

changelog/unreleased/SOLR-18195-collapse-results-combined-query.yml

@@ -0,0 +1,9 @@
+# See https://github.com/apache/solr/blob/main/dev-docs/changelog.adoc
+title: Support for using {!collapse} with CombinedQueryComponent (RRF)
+type: added
+authors:
+  - name: Sonu Sharma
+    nick: ercsonusharma
+links:
+  - name: SOLR-18195
+    url: https://issues.apache.org/jira/browse/SOLR-18195

solr/core/src/java/org/apache/solr/handler/component/CombinedQueryComponent.java

@@ -31,6 +31,7 @@
 import java.util.Set;
 import java.util.stream.Collectors;
 import org.apache.lucene.search.Explanation;
+import org.apache.lucene.search.Query;
 import org.apache.lucene.search.Sort;
 import org.apache.lucene.search.SortField;
 import org.apache.solr.client.solrj.SolrServerException;
@@ -42,6 +43,7 @@
 import org.apache.solr.common.params.GroupParams;
 import org.apache.solr.common.params.ShardParams;
 import org.apache.solr.common.params.SolrParams;
+import org.apache.solr.common.util.CollectionUtil;
 import org.apache.solr.common.util.NamedList;
 import org.apache.solr.common.util.SimpleOrderedMap;
 import org.apache.solr.common.util.StrUtils;
@@ -53,8 +55,10 @@
 import org.apache.solr.response.SolrQueryResponse;
 import org.apache.solr.schema.IndexSchema;
 import org.apache.solr.schema.SchemaField;
+import org.apache.solr.search.CollapsingQParserPlugin;
 import org.apache.solr.search.DocListAndSet;
 import org.apache.solr.search.QueryResult;
+import org.apache.solr.search.SolrIndexSearcher;
 import org.apache.solr.search.SolrReturnFields;
 import org.apache.solr.search.SortSpec;
 import org.apache.solr.util.SolrResponseUtil;
@@ -236,7 +240,10 @@ private void prepareCombinedResponseBuilder(
       boolean partialResults,
       boolean segmentTerminatedEarly,
       Boolean setMaxHitsTerminatedEarly) {
-    QueryResult combinedQueryResult = QueryAndResponseCombiner.simpleCombine(queryResults);
+    SolrIndexSearcher searcher = crb.req.getSearcher();
+    List<Query> collapseFilters = getCollapseFilters(crb.getFilters());
+    QueryResult combinedQueryResult =
+        QueryAndResponseCombiner.simpleCombine(queryResults, collapseFilters, searcher);
     combinedQueryResult.setPartialResults(partialResults);
     combinedQueryResult.setSegmentTerminatedEarly(segmentTerminatedEarly);
     combinedQueryResult.setMaxHitsTerminatedEarly(setMaxHitsTerminatedEarly);
@@ -258,6 +265,21 @@ private void prepareCombinedResponseBuilder(
     }
   }
 
+  /** Extracts the list of CollapsingPostFilter query from the filter list, if present. */
+  private static List<Query> getCollapseFilters(List<Query> filters) {
+    if (CollectionUtil.isNotEmpty(filters)) {
+      return filters.stream()
+          .filter(q -> q instanceof CollapsingQParserPlugin.CollapsingPostFilter)
+          .toList();
+    }
+    return List.of();
+  }
+
+  /**
+   * Each shard response contains both "response" and "response_per_query". Only "response" is
+   * deduplicated across sub-queries, so any processing of "response_per_query" must exclude docs
+   * not present in "response" to avoid reintroducing docs that were eliminated while deduplication.
+   */
   @Override
   protected void mergeIds(ResponseBuilder rb, ShardRequest sreq) {
     SortSpec ss = rb.getSortSpec();
@@ -296,6 +318,10 @@ protected void mergeIds(ResponseBuilder rb, ShardRequest sreq) {
     long approximateTotalHits = 0;
     Map<String, List<ShardDoc>> shardDocMap = new HashMap<>();
     String[] queriesToCombineKeys = rb.req.getParams().getParams(CombinerParams.COMBINER_QUERY);
+    // Build per-shard set of doc IDs from the shard's combined (deduplicated) response.
+    // Used to filter per-query docs so that RRF doesn't reintroduce docs
+    // excluded by collapse at the shard level.
+    Map<String, Set<Object>> combinedDocIdsPerShard = HashMap.newHashMap(sreq.responses.size());
     // TODO: to be parallelized outer loop
     for (int queryIndex = 0; queryIndex < queriesToCombineKeys.length; queryIndex++) {
       int failedShardCount = 0;
@@ -377,9 +403,15 @@ protected void mergeIds(ResponseBuilder rb, ShardRequest sreq) {
                 : new NamedList<>();
         // go through every doc in this response, construct a ShardDoc, and
         // put it in the uniqueDoc to dedup
+        Set<Object> thisShardCombinedIds =
+            combinedDocIdsPerShard.computeIfAbsent(
+                srsp.getShard(), shard -> extractIdsFromCombinedResponse(rb, srsp, uniqueKeyField));
         for (int i = 0; i < docs.size(); i++) {
           SolrDocument doc = docs.get(i);
           Object id = doc.getFieldValue(uniqueKeyField.getName());
+          if (!thisShardCombinedIds.contains(id)) {
+            continue;
+          }
           ShardDoc shardDoc = new ShardDoc();
           shardDoc.id = id;
           shardDoc.orderInShard = i;
@@ -610,4 +642,19 @@ protected boolean lessThan(ShardDoc docA, ShardDoc docB) {
     for (int i = 0; i < resultSize; i++) responseDocs.add(null);
     return resultIds;
   }
+
+  /**
+   * Extracts the set of doc IDs from the shard's combined response (produced by simpleCombine).
+   * Returns an empty set if the combined response is not available.
+   */
+  private static Set<Object> extractIdsFromCombinedResponse(
+      ResponseBuilder rb, ShardResponse srsp, SchemaField uniqueKeyField) {
+    Object response = SolrResponseUtil.getSubsectionFromShardResponse(rb, srsp, "response", false);
+    if (response instanceof SolrDocumentList docList) {
+      return docList.stream()
+          .map(doc -> doc.getFieldValue(uniqueKeyField.getName()))
+          .collect(Collectors.toSet());
+    }
+    return Set.of();
+  }
 }

solr/core/src/java/org/apache/solr/handler/component/combine/QueryAndResponseCombiner.java

@@ -16,19 +16,33 @@
  */
 package org.apache.solr.handler.component.combine;
 
+import java.io.IOException;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.internal.hppc.IntDoubleHashMap;
+import org.apache.lucene.queries.function.FunctionScoreQuery;
+import org.apache.lucene.search.DoubleValues;
+import org.apache.lucene.search.DoubleValuesSource;
 import org.apache.lucene.search.Explanation;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TotalHits;
 import org.apache.solr.common.SolrException;
 import org.apache.solr.common.params.SolrParams;
+import org.apache.solr.common.util.CollectionUtil;
 import org.apache.solr.common.util.SimpleOrderedMap;
 import org.apache.solr.handler.component.ShardDoc;
 import org.apache.solr.search.DocIterator;
 import org.apache.solr.search.DocSet;
 import org.apache.solr.search.DocSlice;
+import org.apache.solr.search.QueryCommand;
 import org.apache.solr.search.QueryResult;
+import org.apache.solr.search.SolrIndexSearcher;
+import org.apache.solr.search.SortedIntDocSet;
 import org.apache.solr.util.plugin.NamedListInitializedPlugin;
 
 /**
@@ -49,12 +63,19 @@ public abstract List<ShardDoc> combine(
       Map<String, List<ShardDoc>> queriesDocMap, SolrParams solrParams);
 
   /**
-   * Simple combine query result list as a union.
+   * Combine query result list as a union, optionally deduplicating by a collapse field. When a
+   * collapse filter is provided, only one document per unique field value is kept (based on the
+   * collapse sort/score selection). This ensures that collapse semantics are preserved across
+   * combined queries.
    *
    * @param queryResults the query results to be combined
+   * @param collapseFilters the collapse post filters, or empty if no collapse dedup is needed
+   * @param searcher the searcher to read field values from, required when collapseFilters is
+   *     non-empty
    * @return the combined query result
    */
-  public static QueryResult simpleCombine(List<QueryResult> queryResults) {
+  public static QueryResult simpleCombine(
+      List<QueryResult> queryResults, List<Query> collapseFilters, SolrIndexSearcher searcher) {
     QueryResult combinedQueryResults = new QueryResult();
     DocSet combinedDocSet = null;
     Map<Integer, Float> uniqueDocIds = new HashMap<>();
@@ -71,6 +92,19 @@ public static QueryResult simpleCombine(List<QueryResult> queryResults) {
         combinedDocSet = combinedDocSet.union(queryResult.getDocSet());
       }
     }
+
+    // If collapse fields are specified, deduplicate by field value across combined queries.
+    // Each sub-query already collapsed individually, but different sub-queries may have
+    // selected different group heads for the same field value.
+    int removedByCollapse = 0;
+    if (CollectionUtil.isNotEmpty(collapseFilters) && searcher != null && queryResults.size() > 1) {
+      int preCollapseSize = uniqueDocIds.size();
+      combinedDocSet =
+          removeCollapsedDuplicatesViaSearcher(
+              collapseFilters, searcher, uniqueDocIds, combinedDocSet);
+      removedByCollapse = preCollapseSize - uniqueDocIds.size();
+    }
+
     int combinedResultsLength = uniqueDocIds.size();
     int[] combinedResultsDocIds = new int[combinedResultsLength];
     float[] combinedResultScores = new float[combinedResultsLength];
@@ -87,14 +121,64 @@ public static QueryResult simpleCombine(List<QueryResult> queryResults) {
             combinedResultsLength,
             combinedResultsDocIds,
             combinedResultScores,
-            Math.max(combinedResultsLength, totalMatches),
+            Math.max(combinedResultsLength, totalMatches - removedByCollapse),
             combinedResultScores.length > 0 ? combinedResultScores[0] : 0,
             TotalHits.Relation.GREATER_THAN_OR_EQUAL_TO);
     combinedQueryResults.setDocList(combinedResultSlice);
     combinedQueryResults.setDocSet(combinedDocSet);
     return combinedQueryResults;
   }
 
+  /**
+   * Removes collapsed duplicates across combined sub-queries. Ensures that only one document per
+   * collapse field value retained across the merged results. Entries removed by collapsing are also
+   * removed from {@code uniqueDocIds} (mutated in place).
+   *
+   * @return the collapsed combined DocSet, or null if combinedDocSet was null
+   */
+  private static DocSet removeCollapsedDuplicatesViaSearcher(
+      List<Query> collapseFilters,
+      SolrIndexSearcher searcher,
+      Map<Integer, Float> uniqueDocIds,
+      DocSet combinedDocSet) {
+    IntDoubleHashMap scoreMap = new IntDoubleHashMap(uniqueDocIds.size());
+    uniqueDocIds.forEach((doc, score) -> scoreMap.put(doc, score.doubleValue()));
+    Query baseQuery;
+    boolean needDocSet;
+    if (combinedDocSet != null) {
+      baseQuery = combinedDocSet.makeQuery();
+      needDocSet = true;
+    } else {
+      int[] queryDocIds =
+          uniqueDocIds.keySet().stream().mapToInt(Integer::intValue).sorted().toArray();
+      baseQuery = new SortedIntDocSet(queryDocIds).makeQuery();
+      needDocSet = false;
+    }
+    Query scoredQuery =
+        FunctionScoreQuery.boostByValue(baseQuery, new PrecomputedScoreValuesSource(scoreMap));
+
+    try {
+      QueryCommand cmd =
+          new QueryCommand()
+              .setQuery(scoredQuery)
+              .setFilterList(collapseFilters)
+              .setLen(uniqueDocIds.size())
+              .setNeedDocSet(needDocSet);
+      QueryResult result = searcher.search(cmd);
+
+      Set<Integer> retainedDocIds = HashSet.newHashSet(result.getDocList().size());
+      DocIterator iter = result.getDocList().iterator();
+      while (iter.hasNext()) {
+        retainedDocIds.add(iter.nextDoc());
+      }
+
+      uniqueDocIds.keySet().retainAll(retainedDocIds);
+      return needDocSet ? result.getDocSet() : null;
+    } catch (IOException e) {
+      throw new SolrException(SolrException.ErrorCode.SERVER_ERROR, e);
+    }
+  }
+
   /**
    * Retrieves a list of explanations for the given queries and results.
    *
@@ -127,4 +211,67 @@ public static QueryAndResponseCombiner getImplementation(
     throw new SolrException(
         SolrException.ErrorCode.BAD_REQUEST, "Unknown Combining algorithm: " + algorithm);
   }
+
+  /**
+   * A {@link DoubleValuesSource} backed by a global doc ID to score map. Returns pre-computed
+   * scores for specific document IDs.
+   */
+  private static class PrecomputedScoreValuesSource extends DoubleValuesSource {
+
+    private final IntDoubleHashMap scoreByDoc;
+
+    PrecomputedScoreValuesSource(IntDoubleHashMap scoreByDoc) {
+      this.scoreByDoc = scoreByDoc;
+    }
+
+    @Override
+    public DoubleValues getValues(LeafReaderContext ctx, DoubleValues existing) {
+      int base = ctx.docBase;
+      return new DoubleValues() {
+        private double currentScore;
+
+        @Override
+        public double doubleValue() {
+          return currentScore;
+        }
+
+        @Override
+        public boolean advanceExact(int doc) {
+          int globalDoc = base + doc;
+          currentScore = scoreByDoc.get(globalDoc);
+          return true;
+        }
+      };
+    }
+
+    @Override
+    public boolean needsScores() {
+      return false;
+    }
+
+    @Override
+    public DoubleValuesSource rewrite(IndexSearcher searcher) {
+      return this;
+    }
+
+    @Override
+    public boolean isCacheable(LeafReaderContext ctx) {
+      return false;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      return o instanceof PrecomputedScoreValuesSource other && scoreByDoc.equals(other.scoreByDoc);
+    }
+
+    @Override
+    public int hashCode() {
+      return scoreByDoc.hashCode();
+    }
+
+    @Override
+    public String toString() {
+      return "PrecomputedScoreValuesSource(docs=" + scoreByDoc.size() + ")";
+    }
+  }
 }

solr/core/src/test-files/solr/collection1/conf/schema-vector-catchall.xml

@@ -504,7 +504,7 @@
         </analyzer>
     </fieldType>
     <copyField source="*_commongrams" dest="*_commongrams_stop" />
-    
+
     <fieldType name="binary" class="solr.BinaryField" />
     <fieldType name="collation" class="solr.CollationField" language="en" />
     <fieldType name="icuCollation" class="solr.ICUCollationField" locale="en" />
@@ -691,6 +691,7 @@
 
 
     <dynamicField name="*_s" type="string" indexed="true" stored="true"/>
+    <dynamicField name="*_sdv" type="string" indexed="true" stored="true" docValues="true" multiValued="false"/>
     <dynamicField name="*_s1" type="string" indexed="true" stored="true" multiValued="false"/>
     <dynamicField name="*_s1_ns" type="string" indexed="true" stored="false" multiValued="false"/>
     <dynamicField name="*_l" type="long" indexed="true" stored="true"/>
@@ -857,14 +858,9 @@
     <copyField source="lowerfilt1" dest="lowerfilt1and2"/>
     <copyField source="lowerfilt" dest="lowerfilt1and2"/>
 
-    <copyField source="*" dest="text"/>
-    <copyField source="id" dest="range_facet_l"/>
-    <copyField source="id" dest="id_i1"/>
     <copyField source="range_facet_f" dest="range_facet_d"/>
     <copyField source="range_facet_f1" dest="range_facet_f1_dv"/>
 
-    <copyField source="id" dest="range_facet_l_dv"/>
-    <copyField source="id" dest="range_facet_i_dv"/>
     <copyField source="range_facet_f" dest="range_facet_f_dv"/>
     <copyField source="range_facet_f" dest="range_facet_d_dv"/>
     <copyField source="bday" dest="range_facet_dt_dv"/>