Fix for propagating filters from compound to inner retrievers (#117914)

Author: pmpailis

docs/changelog/117914.yaml

@@ -0,0 +1,5 @@
+pr: 117914
+summary: Fix for propagating filters from compound to inner retrievers
+area: Ranking
+type: bug
+issues: []

server/src/main/java/org/elasticsearch/search/retriever/CompoundRetrieverBuilder.java

@@ -20,6 +20,7 @@
 import org.elasticsearch.action.search.SearchRequest;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.action.search.TransportMultiSearchAction;
+import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.rest.RestStatus;
@@ -46,6 +47,8 @@
  */
 public abstract class CompoundRetrieverBuilder<T extends CompoundRetrieverBuilder<T>> extends RetrieverBuilder {
 
+    public static final NodeFeature INNER_RETRIEVERS_FILTER_SUPPORT = new NodeFeature("inner_retrievers_filter_support");
+
     public record RetrieverSource(RetrieverBuilder retriever, SearchSourceBuilder source) {}
 
     protected final int rankWindowSize;
@@ -64,9 +67,9 @@ public T addChild(RetrieverBuilder retrieverBuilder) {
 
     /**
      * Returns a clone of the original retriever, replacing the sub-retrievers with
-     * the provided {@code newChildRetrievers}.
+     * the provided {@code newChildRetrievers} and the filters with the {@code newPreFilterQueryBuilders}.
      */
-    protected abstract T clone(List<RetrieverSource> newChildRetrievers);
+    protected abstract T clone(List<RetrieverSource> newChildRetrievers, List<QueryBuilder> newPreFilterQueryBuilders);
 
     /**
      * Combines the provided {@code rankResults} to return the final top documents.
@@ -85,13 +88,25 @@ public final RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOExceptio
         }
 
         // Rewrite prefilters
-        boolean hasChanged = false;
+        // We eagerly rewrite prefilters, because some of the innerRetrievers
+        // could be compound too, so we want to propagate all the necessary filter information to them
+        // and have it available as part of their own rewrite step
         var newPreFilters = rewritePreFilters(ctx);
-        hasChanged |= newPreFilters != preFilterQueryBuilders;
+        if (newPreFilters != preFilterQueryBuilders) {
+            return clone(innerRetrievers, newPreFilters);
+        }
 
+        boolean hasChanged = false;
         // Rewrite retriever sources
         List<RetrieverSource> newRetrievers = new ArrayList<>();
         for (var entry : innerRetrievers) {
+            // we propagate the filters only for compound retrievers as they won't be attached through
+            // the createSearchSourceBuilder.
+            // We could remove this check, but we would end up adding the same filters
+            // multiple times in case an inner retriever rewrites itself, when we re-enter to rewrite
+            if (entry.retriever.isCompound() && false == preFilterQueryBuilders.isEmpty()) {
+                entry.retriever.getPreFilterQueryBuilders().addAll(preFilterQueryBuilders);
+            }
             RetrieverBuilder newRetriever = entry.retriever.rewrite(ctx);
             if (newRetriever != entry.retriever) {
                 newRetrievers.add(new RetrieverSource(newRetriever, null));
@@ -106,7 +121,7 @@ public final RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOExceptio
             }
         }
         if (hasChanged) {
-            return clone(newRetrievers);
+            return clone(newRetrievers, newPreFilters);
         }
 
         // execute searches
@@ -166,12 +181,7 @@ public void onFailure(Exception e) {
             });
         });
 
-        return new RankDocsRetrieverBuilder(
-            rankWindowSize,
-            newRetrievers.stream().map(s -> s.retriever).toList(),
-            results::get,
-            newPreFilters
-        );
+        return new RankDocsRetrieverBuilder(rankWindowSize, newRetrievers.stream().map(s -> s.retriever).toList(), results::get);
     }
 
     @Override

server/src/main/java/org/elasticsearch/search/retriever/KnnRetrieverBuilder.java

@@ -184,8 +184,7 @@ public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException {
                     ll.onResponse(null);
                 }));
             });
-            var rewritten = new KnnRetrieverBuilder(this, () -> toSet.get(), null);
-            return rewritten;
+            return new KnnRetrieverBuilder(this, () -> toSet.get(), null);
         }
         return super.rewrite(ctx);
     }

server/src/main/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilder.java

@@ -33,19 +33,13 @@ public class RankDocsRetrieverBuilder extends RetrieverBuilder {
     final List<RetrieverBuilder> sources;
     final Supplier<RankDoc[]> rankDocs;
 
-    public RankDocsRetrieverBuilder(
-        int rankWindowSize,
-        List<RetrieverBuilder> sources,
-        Supplier<RankDoc[]> rankDocs,
-        List<QueryBuilder> preFilterQueryBuilders
-    ) {
+    public RankDocsRetrieverBuilder(int rankWindowSize, List<RetrieverBuilder> sources, Supplier<RankDoc[]> rankDocs) {
         this.rankWindowSize = rankWindowSize;
         this.rankDocs = rankDocs;
         if (sources == null || sources.isEmpty()) {
             throw new IllegalArgumentException("sources must not be null or empty");
         }
         this.sources = sources;
-        this.preFilterQueryBuilders = preFilterQueryBuilders;
     }
 
     @Override
@@ -73,10 +67,6 @@ private boolean sourceShouldRewrite(QueryRewriteContext ctx) throws IOException
     @Override
     public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException {
         assert false == sourceShouldRewrite(ctx) : "retriever sources should be rewritten first";
-        var rewrittenFilters = rewritePreFilters(ctx);
-        if (rewrittenFilters != preFilterQueryBuilders) {
-            return new RankDocsRetrieverBuilder(rankWindowSize, sources, rankDocs, rewrittenFilters);
-        }
         return this;
     }
 
@@ -94,7 +84,7 @@ public QueryBuilder topDocsQuery() {
                 boolQuery.should(query);
             }
         }
-        // ignore prefilters of this level, they are already propagated to children
+        // ignore prefilters of this level, they were already propagated to children
         return boolQuery;
     }
 
@@ -133,7 +123,7 @@ public void extractToSearchSourceBuilder(SearchSourceBuilder searchSourceBuilder
         } else {
             rankQuery = new RankDocsQueryBuilder(rankDocResults, null, false);
         }
-        // ignore prefilters of this level, they are already propagated to children
+        // ignore prefilters of this level, they were already propagated to children
         searchSourceBuilder.query(rankQuery);
         if (sourceHasMinScore()) {
             searchSourceBuilder.minScore(this.minScore() == null ? Float.MIN_VALUE : this.minScore());

server/src/test/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilderTests.java

@@ -95,12 +95,7 @@ private List<QueryBuilder> preFilters(QueryRewriteContext queryRewriteContext) t
     }
 
     private RankDocsRetrieverBuilder createRandomRankDocsRetrieverBuilder(QueryRewriteContext queryRewriteContext) throws IOException {
-        return new RankDocsRetrieverBuilder(
-            randomIntBetween(1, 100),
-            innerRetrievers(queryRewriteContext),
-            rankDocsSupplier(),
-            preFilters(queryRewriteContext)
-        );
+        return new RankDocsRetrieverBuilder(randomIntBetween(1, 100), innerRetrievers(queryRewriteContext), rankDocsSupplier());
     }
 
     public void testExtractToSearchSourceBuilder() throws IOException {

server/src/test/java/org/elasticsearch/search/vectors/TestQueryVectorBuilderPlugin.java

@@ -27,9 +27,9 @@
 /**
  * A SearchPlugin to exercise query vector builder
  */
-class TestQueryVectorBuilderPlugin implements SearchPlugin {
+public class TestQueryVectorBuilderPlugin implements SearchPlugin {
 
-    static class TestQueryVectorBuilder implements QueryVectorBuilder {
+    public static class TestQueryVectorBuilder implements QueryVectorBuilder {
         private static final String NAME = "test_query_vector_builder";
 
         private static final ParseField QUERY_VECTOR = new ParseField("query_vector");
@@ -47,11 +47,11 @@ static class TestQueryVectorBuilder implements QueryVectorBuilder {
 
         private List<Float> vectorToBuild;
 
-        TestQueryVectorBuilder(List<Float> vectorToBuild) {
+        public TestQueryVectorBuilder(List<Float> vectorToBuild) {
             this.vectorToBuild = vectorToBuild;
         }
 
-        TestQueryVectorBuilder(float[] expected) {
+        public TestQueryVectorBuilder(float[] expected) {
             this.vectorToBuild = new ArrayList<>(expected.length);
             for (float f : expected) {
                 vectorToBuild.add(f);

test/framework/src/main/java/org/elasticsearch/search/retriever/TestCompoundRetrieverBuilder.java

@@ -10,6 +10,7 @@
 package org.elasticsearch.search.retriever;
 
 import org.apache.lucene.search.ScoreDoc;
+import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.search.rank.RankDoc;
 import org.elasticsearch.xcontent.XContentBuilder;
 
@@ -23,16 +24,17 @@ public class TestCompoundRetrieverBuilder extends CompoundRetrieverBuilder<TestC
     public static final String NAME = "test_compound_retriever_builder";
 
     public TestCompoundRetrieverBuilder(int rankWindowSize) {
-        this(new ArrayList<>(), rankWindowSize);
+        this(new ArrayList<>(), rankWindowSize, new ArrayList<>());
     }
 
-    TestCompoundRetrieverBuilder(List<RetrieverSource> childRetrievers, int rankWindowSize) {
+    TestCompoundRetrieverBuilder(List<RetrieverSource> childRetrievers, int rankWindowSize, List<QueryBuilder> preFilterQueryBuilders) {
         super(childRetrievers, rankWindowSize);
+        this.preFilterQueryBuilders = preFilterQueryBuilders;
     }
 
     @Override
-    protected TestCompoundRetrieverBuilder clone(List<RetrieverSource> newChildRetrievers) {
-        return new TestCompoundRetrieverBuilder(newChildRetrievers, rankWindowSize);
+    protected TestCompoundRetrieverBuilder clone(List<RetrieverSource> newChildRetrievers, List<QueryBuilder> newPreFilterQueryBuilders) {
+        return new TestCompoundRetrieverBuilder(newChildRetrievers, rankWindowSize, newPreFilterQueryBuilders);
     }
 
     @Override

x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/rules/retriever/QueryRuleRetrieverBuilder.java

@@ -110,12 +110,14 @@ public QueryRuleRetrieverBuilder(
         Map<String, Object> matchCriteria,
         List<RetrieverSource> retrieverSource,
         int rankWindowSize,
-        String retrieverName
+        String retrieverName,
+        List<QueryBuilder> preFilterQueryBuilders
     ) {
         super(retrieverSource, rankWindowSize);
         this.rulesetIds = rulesetIds;
         this.matchCriteria = matchCriteria;
         this.retrieverName = retrieverName;
+        this.preFilterQueryBuilders = preFilterQueryBuilders;
     }
 
     @Override
@@ -156,8 +158,15 @@ public void doToXContent(XContentBuilder builder, Params params) throws IOExcept
     }
 
     @Override
-    protected QueryRuleRetrieverBuilder clone(List<RetrieverSource> newChildRetrievers) {
-        return new QueryRuleRetrieverBuilder(rulesetIds, matchCriteria, newChildRetrievers, rankWindowSize, retrieverName);
+    protected QueryRuleRetrieverBuilder clone(List<RetrieverSource> newChildRetrievers, List<QueryBuilder> newPreFilterQueryBuilders) {
+        return new QueryRuleRetrieverBuilder(
+            rulesetIds,
+            matchCriteria,
+            newChildRetrievers,
+            rankWindowSize,
+            retrieverName,
+            newPreFilterQueryBuilders
+        );
     }
 
     @Override

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankRetrieverBuilder.java

@@ -129,7 +129,10 @@ public TextSimilarityRankRetrieverBuilder(
     }
 
     @Override
-    protected TextSimilarityRankRetrieverBuilder clone(List<RetrieverSource> newChildRetrievers) {
+    protected TextSimilarityRankRetrieverBuilder clone(
+        List<RetrieverSource> newChildRetrievers,
+        List<QueryBuilder> newPreFilterQueryBuilders
+    ) {
         return new TextSimilarityRankRetrieverBuilder(
             newChildRetrievers,
             inferenceId,
@@ -138,7 +141,7 @@ protected TextSimilarityRankRetrieverBuilder clone(List<RetrieverSource> newChil
             rankWindowSize,
             minScore,
             retrieverName,
-            preFilterQueryBuilders
+            newPreFilterQueryBuilders
         );
     }
 

x-pack/plugin/rank-rrf/src/internalClusterTest/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilderIT.java

@@ -33,6 +33,7 @@
 import org.elasticsearch.search.sort.SortOrder;
 import org.elasticsearch.search.vectors.KnnVectorQueryBuilder;
 import org.elasticsearch.search.vectors.QueryVectorBuilder;
+import org.elasticsearch.search.vectors.TestQueryVectorBuilderPlugin;
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.test.hamcrest.ElasticsearchAssertions;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -57,7 +58,6 @@
 public class RRFRetrieverBuilderIT extends ESIntegTestCase {
 
     protected static String INDEX = "test_index";
-    protected static final String ID_FIELD = "_id";
     protected static final String DOC_FIELD = "doc";
     protected static final String TEXT_FIELD = "text";
     protected static final String VECTOR_FIELD = "vector";
@@ -743,6 +743,42 @@ public void extractToSearchSourceBuilder(SearchSourceBuilder searchSourceBuilder
         expectThrows(UnsupportedOperationException.class, () -> client().prepareSearch(INDEX).setSource(source).get());
     }
 
+    public void testRRFFiltersPropagatedToKnnQueryVectorBuilder() {
+        final int rankWindowSize = 100;
+        final int rankConstant = 10;
+        SearchSourceBuilder source = new SearchSourceBuilder();
+        // this will retriever all but 7 only due to top-level filter
+        StandardRetrieverBuilder standardRetriever = new StandardRetrieverBuilder(QueryBuilders.matchAllQuery());
+        // this will too retrieve just doc 7
+        KnnRetrieverBuilder knnRetriever = new KnnRetrieverBuilder(
+            "vector",
+            null,
+            new TestQueryVectorBuilderPlugin.TestQueryVectorBuilder(new float[] { 3 }),
+            10,
+            10,
+            null
+        );
+        source.retriever(
+            new RRFRetrieverBuilder(
+                Arrays.asList(
+                    new CompoundRetrieverBuilder.RetrieverSource(standardRetriever, null),
+                    new CompoundRetrieverBuilder.RetrieverSource(knnRetriever, null)
+                ),
+                rankWindowSize,
+                rankConstant
+            )
+        );
+        source.retriever().getPreFilterQueryBuilders().add(QueryBuilders.boolQuery().must(QueryBuilders.termQuery(DOC_FIELD, "doc_7")));
+        source.size(10);
+        SearchRequestBuilder req = client().prepareSearch(INDEX).setSource(source);
+        ElasticsearchAssertions.assertResponse(req, resp -> {
+            assertNull(resp.pointInTimeId());
+            assertNotNull(resp.getHits().getTotalHits());
+            assertThat(resp.getHits().getTotalHits().value(), equalTo(1L));
+            assertThat(resp.getHits().getHits()[0].getId(), equalTo("doc_7"));
+        });
+    }
+
     public void testRewriteOnce() {
         final float[] vector = new float[] { 1 };
         AtomicInteger numAsyncCalls = new AtomicInteger();