Query rules retriever (#114855)

Author: kderusso

docs/changelog/114855.yaml

@@ -0,0 +1,5 @@
+pr: 114855
+summary: Add query rules retriever
+area: Relevance
+type: enhancement
+issues: [ ]

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -184,6 +184,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion INDEX_MODE_LOOKUP = def(8_779_00_0);
     public static final TransportVersion INDEX_REQUEST_REMOVE_METERING = def(8_780_00_0);
     public static final TransportVersion CPU_STAT_STRING_PARSING = def(8_781_00_0);
+    public static final TransportVersion QUERY_RULES_RETRIEVER = def(8_782_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -218,6 +218,10 @@ public void setRankDocs(RankDoc[] rankDocs) {
         this.rankDocs = rankDocs;
     }
 
+    public RankDoc[] getRankDocs() {
+        return rankDocs;
+    }
+
     /**
      * Gets the filters for this retriever.
      */

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

@@ -0,0 +1,142 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.search.retriever;
+
+import org.elasticsearch.action.ActionRequestValidationException;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.QueryRewriteContext;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.search.rank.RankDoc;
+import org.elasticsearch.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.List;
+
+/**
+ * A wrapper that can be used to modify the behaviour of an existing {@link RetrieverBuilder}.
+ */
+public abstract class RetrieverBuilderWrapper<T extends RetrieverBuilder> extends RetrieverBuilder {
+    protected final RetrieverBuilder in;
+
+    protected RetrieverBuilderWrapper(RetrieverBuilder in) {
+        this.in = in;
+    }
+
+    protected abstract T clone(RetrieverBuilder sub);
+
+    @Override
+    public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException {
+        var inRewrite = in.rewrite(ctx);
+        if (inRewrite != in) {
+            return clone(inRewrite);
+        }
+        return this;
+    }
+
+    @Override
+    public QueryBuilder topDocsQuery() {
+        return in.topDocsQuery();
+    }
+
+    @Override
+    public RetrieverBuilder minScore(Float minScore) {
+        return in.minScore(minScore);
+    }
+
+    @Override
+    public List<QueryBuilder> getPreFilterQueryBuilders() {
+        return in.preFilterQueryBuilders;
+    }
+
+    @Override
+    public ActionRequestValidationException validate(
+        SearchSourceBuilder source,
+        ActionRequestValidationException validationException,
+        boolean isScroll,
+        boolean allowPartialSearchResults
+    ) {
+        return in.validate(source, validationException, isScroll, allowPartialSearchResults);
+    }
+
+    @Override
+    public RetrieverBuilder retrieverName(String retrieverName) {
+        return in.retrieverName(retrieverName);
+    }
+
+    @Override
+    public void setRankDocs(RankDoc[] rankDocs) {
+        in.setRankDocs(rankDocs);
+    }
+
+    @Override
+    public RankDoc[] getRankDocs() {
+        return in.getRankDocs();
+    }
+
+    @Override
+    public boolean isCompound() {
+        return in.isCompound();
+    }
+
+    @Override
+    public QueryBuilder explainQuery() {
+        return in.explainQuery();
+    }
+
+    @Override
+    public Float minScore() {
+        return in.minScore();
+    }
+
+    @Override
+    public boolean isFragment() {
+        return in.isFragment();
+    }
+
+    @Override
+    public String toString() {
+        return in.toString();
+    }
+
+    @Override
+    public String retrieverName() {
+        return in.retrieverName();
+    }
+
+    @Override
+    public void extractToSearchSourceBuilder(SearchSourceBuilder searchSourceBuilder, boolean compoundUsed) {
+        in.extractToSearchSourceBuilder(searchSourceBuilder, compoundUsed);
+    }
+
+    @Override
+    public String getName() {
+        return in.getName();
+    }
+
+    @Override
+    protected void doToXContent(XContentBuilder builder, Params params) throws IOException {
+        in.doToXContent(builder, params);
+    }
+
+    @Override
+    protected boolean doEquals(Object o) {
+        // Handle the edge case where we need to unwrap the incoming retriever
+        if (o instanceof RetrieverBuilderWrapper<?> wrapper) {
+            return in.doEquals(wrapper.in);
+        } else {
+            return in.doEquals(o);
+        }
+    }
+
+    @Override
+    protected int doHashCode() {
+        return in.doHashCode();
+    }
+}

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

@@ -89,6 +89,8 @@ public void doToXContent(XContentBuilder builder, Params params) throws IOExcept
 
     @Override
     public boolean doEquals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
         TestRetrieverBuilder that = (TestRetrieverBuilder) o;
         return Objects.equals(value, that.value);
     }