Implement support for weighted rrf (#130658)

* RRFRetrieverComponent added:

* Modified parser, toXcontent and included component in the RetrieverBuilder

* [CI] Auto commit changes from spotless

* Resolved merge conflicts

* Fixed compile issues in tests

* [CI] Auto commit changes from spotless

* trying to resolve parse errros

* wip

* Modified builder

* [CI] Auto commit changes from spotless

* Removed unnecessary code

* Fixed import

* Enhanced tests

* Fixed the failing tests

* Yaml tests were added

* Added cluster features to it

* Fixed spotless

* Update docs/changelog/130658.yaml

* Fixed the relaxed constraints

* Resolving issues

* Resolved PR comments

* removed simplified rrf

* changed the test file back to its original state

* Resolved comments to have ahelper method and the test case to use it

* made parsing robust

* IT test reverted

* Replaced the declareString array parser

* Enforced weights as nonnull

* Fixed the weights null

* Empty weight shouldnt be serialised

* [CI] Auto commit changes from spotless

* removed the hard coding

* Cleanup and optimised the code flow

* Fixed the comments

* [CI] Auto commit changes from spotless

* optimised test

* Added additional test

* addressed the commentS

* Update docs/changelog/130658.yaml

Co-authored-by: Liam Thompson <[email protected]>

* Explicit check for retriever object

* Resolved PR comments

* Fixed the error message

---------

Co-authored-by: elasticsearchmachine <[email protected]>
Co-authored-by: Ioana Tagirta <[email protected]>
Co-authored-by: Liam Thompson <[email protected]>

Author: 4 people

docs/changelog/130658.yaml

@@ -0,0 +1,5 @@
+pr: 130658
+summary: Add support for weighted RRF in retrievers
+area: Relevance
+type: enhancement
+issues: []

x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/RankRRFFeatures.java

@@ -36,7 +36,8 @@ public Set<NodeFeature> getTestFeatures() {
             LINEAR_RETRIEVER_L2_NORM,
             LINEAR_RETRIEVER_MINSCORE_FIX,
             LinearRetrieverBuilder.MULTI_FIELDS_QUERY_FORMAT_SUPPORT,
-            RRFRetrieverBuilder.MULTI_FIELDS_QUERY_FORMAT_SUPPORT
+            RRFRetrieverBuilder.MULTI_FIELDS_QUERY_FORMAT_SUPPORT,
+            RRFRetrieverBuilder.WEIGHTED_SUPPORT
         );
     }
 }

x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilder.java

@@ -20,6 +20,7 @@
 import org.elasticsearch.search.rank.RankBuilder;
 import org.elasticsearch.search.rank.RankDoc;
 import org.elasticsearch.search.retriever.CompoundRetrieverBuilder;
+import org.elasticsearch.search.retriever.CompoundRetrieverBuilder.RetrieverSource;
 import org.elasticsearch.search.retriever.RetrieverBuilder;
 import org.elasticsearch.search.retriever.RetrieverParserContext;
 import org.elasticsearch.search.retriever.StandardRetrieverBuilder;
@@ -37,7 +38,7 @@
 import java.util.Map;
 import java.util.Objects;
 
-import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg;
+import static org.elasticsearch.xpack.rank.rrf.RRFRetrieverComponent.DEFAULT_WEIGHT;
 
 /**
  * An rrf retriever is used to represent an rrf rank element, but
@@ -48,6 +49,7 @@
  */
 public final class RRFRetrieverBuilder extends CompoundRetrieverBuilder<RRFRetrieverBuilder> {
     public static final NodeFeature MULTI_FIELDS_QUERY_FORMAT_SUPPORT = new NodeFeature("rrf_retriever.multi_fields_query_format_support");
+    public static final NodeFeature WEIGHTED_SUPPORT = new NodeFeature("rrf_retriever.weighted_support");
 
     public static final String NAME = "rrf";
 
@@ -57,37 +59,38 @@ public final class RRFRetrieverBuilder extends CompoundRetrieverBuilder<RRFRetri
     public static final ParseField QUERY_FIELD = new ParseField("query");
 
     public static final int DEFAULT_RANK_CONSTANT = 60;
+
+    private final float[] weights;
+
     @SuppressWarnings("unchecked")
     static final ConstructingObjectParser<RRFRetrieverBuilder, RetrieverParserContext> PARSER = new ConstructingObjectParser<>(
         NAME,
         false,
         args -> {
-            List<RetrieverBuilder> childRetrievers = (List<RetrieverBuilder>) args[0];
+            List<RRFRetrieverComponent> retrieverComponents = args[0] == null ? List.of() : (List<RRFRetrieverComponent>) args[0];
             List<String> fields = (List<String>) args[1];
             String query = (String) args[2];
             int rankWindowSize = args[3] == null ? RankBuilder.DEFAULT_RANK_WINDOW_SIZE : (int) args[3];
             int rankConstant = args[4] == null ? DEFAULT_RANK_CONSTANT : (int) args[4];
 
-            List<RetrieverSource> innerRetrievers = childRetrievers != null
-                ? childRetrievers.stream().map(RetrieverSource::from).toList()
-                : List.of();
-            return new RRFRetrieverBuilder(innerRetrievers, fields, query, rankWindowSize, rankConstant);
+            int n = retrieverComponents.size();
+            List<RetrieverSource> innerRetrievers = new ArrayList<>(n);
+            float[] weights = new float[n];
+            for (int i = 0; i < n; i++) {
+                RRFRetrieverComponent component = retrieverComponents.get(i);
+                innerRetrievers.add(RetrieverSource.from(component.retriever()));
+                weights[i] = component.weight();
+            }
+            return new RRFRetrieverBuilder(innerRetrievers, fields, query, rankWindowSize, rankConstant, weights);
         }
     );
 
     static {
-        PARSER.declareObjectArray(optionalConstructorArg(), (p, c) -> {
-            p.nextToken();
-            String name = p.currentName();
-            RetrieverBuilder retrieverBuilder = p.namedObject(RetrieverBuilder.class, name, c);
-            c.trackRetrieverUsage(retrieverBuilder.getName());
-            p.nextToken();
-            return retrieverBuilder;
-        }, RETRIEVERS_FIELD);
-        PARSER.declareStringArray(optionalConstructorArg(), FIELDS_FIELD);
-        PARSER.declareString(optionalConstructorArg(), QUERY_FIELD);
-        PARSER.declareInt(optionalConstructorArg(), RANK_WINDOW_SIZE_FIELD);
-        PARSER.declareInt(optionalConstructorArg(), RANK_CONSTANT_FIELD);
+        PARSER.declareObjectArray(ConstructingObjectParser.optionalConstructorArg(), RRFRetrieverComponent::fromXContent, RETRIEVERS_FIELD);
+        PARSER.declareStringArray(ConstructingObjectParser.optionalConstructorArg(), FIELDS_FIELD);
+        PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), QUERY_FIELD);
+        PARSER.declareInt(ConstructingObjectParser.optionalConstructorArg(), RANK_WINDOW_SIZE_FIELD);
+        PARSER.declareInt(ConstructingObjectParser.optionalConstructorArg(), RANK_CONSTANT_FIELD);
         RetrieverBuilder.declareBaseParserFields(PARSER);
     }
 
@@ -103,27 +106,46 @@ public static RRFRetrieverBuilder fromXContent(XContentParser parser, RetrieverP
     private final int rankConstant;
 
     public RRFRetrieverBuilder(List<RetrieverSource> childRetrievers, int rankWindowSize, int rankConstant) {
-        this(childRetrievers, null, null, rankWindowSize, rankConstant);
+        this(childRetrievers, null, null, rankWindowSize, rankConstant, createDefaultWeights(childRetrievers));
+    }
+
+    private static float[] createDefaultWeights(List<?> retrievers) {
+        int size = retrievers == null ? 0 : retrievers.size();
+        float[] defaultWeights = new float[size];
+        Arrays.fill(defaultWeights, DEFAULT_WEIGHT);
+        return defaultWeights;
     }
 
     public RRFRetrieverBuilder(
         List<RetrieverSource> childRetrievers,
         List<String> fields,
         String query,
         int rankWindowSize,
-        int rankConstant
+        int rankConstant,
+        float[] weights
     ) {
         // Use a mutable list for childRetrievers so that we can use addChild
         super(childRetrievers == null ? new ArrayList<>() : new ArrayList<>(childRetrievers), rankWindowSize);
         this.fields = fields == null ? null : List.copyOf(fields);
         this.query = query;
         this.rankConstant = rankConstant;
+        Objects.requireNonNull(weights, "weights must not be null");
+        if (weights.length != innerRetrievers.size()) {
+            throw new IllegalArgumentException(
+                "weights array length [" + weights.length + "] must match retrievers count [" + innerRetrievers.size() + "]"
+            );
+        }
+        this.weights = weights;
     }
 
     public int rankConstant() {
         return rankConstant;
     }
 
+    public float[] weights() {
+        return weights;
+    }
+
     @Override
     public String getName() {
         return NAME;
@@ -137,6 +159,7 @@ public ActionRequestValidationException validate(
         boolean allowPartialSearchResults
     ) {
         validationException = super.validate(source, validationException, isScroll, allowPartialSearchResults);
+
         return MultiFieldsInnerRetrieverUtils.validateParams(
             innerRetrievers,
             fields,
@@ -151,7 +174,14 @@ public ActionRequestValidationException validate(
 
     @Override
     protected RRFRetrieverBuilder clone(List<RetrieverSource> newRetrievers, List<QueryBuilder> newPreFilterQueryBuilders) {
-        RRFRetrieverBuilder clone = new RRFRetrieverBuilder(newRetrievers, this.fields, this.query, this.rankWindowSize, this.rankConstant);
+        RRFRetrieverBuilder clone = new RRFRetrieverBuilder(
+            newRetrievers,
+            this.fields,
+            this.query,
+            this.rankWindowSize,
+            this.rankConstant,
+            this.weights
+        );
         clone.preFilterQueryBuilders = newPreFilterQueryBuilders;
         clone.retrieverName = retrieverName;
         return clone;
@@ -183,7 +213,7 @@ protected RRFRankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults
 
                     // calculate the current rrf score for this document
                     // later used to sort and covert to a rank
-                    value.score += 1.0f / (rankConstant + frank);
+                    value.score += this.weights[findex] * (1.0f / (rankConstant + frank));
 
                     if (explain && value.positions != null && value.scores != null) {
                         // record the position for each query
@@ -238,10 +268,14 @@ protected RetrieverBuilder doRewrite(QueryRewriteContext ctx) {
                 query,
                 localIndicesMetadata.values(),
                 r -> {
-                    List<RetrieverSource> retrievers = r.stream()
-                        .map(MultiFieldsInnerRetrieverUtils.WeightedRetrieverSource::retrieverSource)
-                        .toList();
-                    return new RRFRetrieverBuilder(retrievers, rankWindowSize, rankConstant);
+                    List<RetrieverSource> retrievers = new ArrayList<>(r.size());
+                    float[] weights = new float[r.size()];
+                    for (int i = 0; i < r.size(); i++) {
+                        var retriever = r.get(i);
+                        retrievers.add(retriever.retrieverSource());
+                        weights[i] = retriever.weight();
+                    }
+                    return new RRFRetrieverBuilder(retrievers, null, null, rankWindowSize, rankConstant, weights);
                 },
                 w -> {
                     if (w != 1.0f) {
@@ -255,7 +289,8 @@ protected RetrieverBuilder doRewrite(QueryRewriteContext ctx) {
             if (fieldsInnerRetrievers.isEmpty() == false) {
                 // TODO: This is a incomplete solution as it does not address other incomplete copy issues
                 // (such as dropping the retriever name and min score)
-                rewritten = new RRFRetrieverBuilder(fieldsInnerRetrievers, rankWindowSize, rankConstant);
+                float[] weights = createDefaultWeights(fieldsInnerRetrievers);
+                rewritten = new RRFRetrieverBuilder(fieldsInnerRetrievers, null, null, rankWindowSize, rankConstant, weights);
                 rewritten.getPreFilterQueryBuilders().addAll(preFilterQueryBuilders);
             } else {
                 // Inner retriever list can be empty when using an index wildcard pattern that doesn't match any indices
@@ -266,29 +301,13 @@ protected RetrieverBuilder doRewrite(QueryRewriteContext ctx) {
         return rewritten;
     }
 
-    // ---- FOR TESTING XCONTENT PARSING ----
-
-    @Override
-    public boolean doEquals(Object o) {
-        RRFRetrieverBuilder that = (RRFRetrieverBuilder) o;
-        return super.doEquals(o)
-            && Objects.equals(fields, that.fields)
-            && Objects.equals(query, that.query)
-            && rankConstant == that.rankConstant;
-    }
-
-    @Override
-    public int doHashCode() {
-        return Objects.hash(super.doHashCode(), fields, query, rankConstant);
-    }
-
     @Override
     public void doToXContent(XContentBuilder builder, Params params) throws IOException {
         if (innerRetrievers.isEmpty() == false) {
             builder.startArray(RETRIEVERS_FIELD.getPreferredName());
-
-            for (var entry : innerRetrievers) {
-                entry.retriever().toXContent(builder, params);
+            for (int i = 0; i < innerRetrievers.size(); i++) {
+                RRFRetrieverComponent component = new RRFRetrieverComponent(innerRetrievers.get(i).retriever(), weights[i]);
+                component.toXContent(builder, params);
             }
             builder.endArray();
         }
@@ -307,4 +326,20 @@ public void doToXContent(XContentBuilder builder, Params params) throws IOExcept
         builder.field(RANK_WINDOW_SIZE_FIELD.getPreferredName(), rankWindowSize);
         builder.field(RANK_CONSTANT_FIELD.getPreferredName(), rankConstant);
     }
+
+    // ---- FOR TESTING XCONTENT PARSING ----
+    @Override
+    public boolean doEquals(Object o) {
+        RRFRetrieverBuilder that = (RRFRetrieverBuilder) o;
+        return super.doEquals(o)
+            && Objects.equals(fields, that.fields)
+            && Objects.equals(query, that.query)
+            && rankConstant == that.rankConstant
+            && Arrays.equals(weights, that.weights);
+    }
+
+    @Override
+    public int doHashCode() {
+        return Objects.hash(super.doHashCode(), fields, query, rankConstant, Arrays.hashCode(weights));
+    }
 }

x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverComponent.java

@@ -0,0 +1,124 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.rank.rrf;
+
+import org.elasticsearch.common.ParsingException;
+import org.elasticsearch.search.retriever.RetrieverBuilder;
+import org.elasticsearch.search.retriever.RetrieverParserContext;
+import org.elasticsearch.xcontent.ParseField;
+import org.elasticsearch.xcontent.ToXContentObject;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentParser;
+
+import java.io.IOException;
+import java.util.Objects;
+
+public class RRFRetrieverComponent implements ToXContentObject {
+
+    public static final ParseField RETRIEVER_FIELD = new ParseField("retriever");
+    public static final ParseField WEIGHT_FIELD = new ParseField("weight");
+    static final float DEFAULT_WEIGHT = 1f;
+
+    final RetrieverBuilder retriever;
+    final float weight;
+
+    public RRFRetrieverComponent(RetrieverBuilder retrieverBuilder, Float weight) {
+        this.retriever = Objects.requireNonNull(retrieverBuilder, "retrieverBuilder must not be null");
+        this.weight = weight == null ? DEFAULT_WEIGHT : weight;
+        if (this.weight < 0) {
+            throw new IllegalArgumentException("[weight] must be non-negative, found [" + this.weight + "]");
+        }
+    }
+
+    public RetrieverBuilder retriever() {
+        return retriever;
+    }
+
+    public float weight() {
+        return weight;
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, ToXContentObject.Params params) throws IOException {
+        builder.startObject();
+        builder.field(RETRIEVER_FIELD.getPreferredName(), retriever);
+        builder.field(WEIGHT_FIELD.getPreferredName(), weight);
+        builder.endObject();
+        return builder;
+    }
+
+    public static RRFRetrieverComponent fromXContent(XContentParser parser, RetrieverParserContext context) throws IOException {
+        if (parser.currentToken() != XContentParser.Token.START_OBJECT) {
+            throw new ParsingException(parser.getTokenLocation(), "expected object but found [{}]", parser.currentToken());
+        }
+
+        // Peek at the first field to determine the format
+        XContentParser.Token token = parser.nextToken();
+        if (token == XContentParser.Token.END_OBJECT) {
+            throw new ParsingException(parser.getTokenLocation(), "retriever component must contain a retriever");
+        }
+        if (token != XContentParser.Token.FIELD_NAME) {
+            throw new ParsingException(parser.getTokenLocation(), "expected field name but found [{}]", token);
+        }
+
+        String firstFieldName = parser.currentName();
+
+        // Check if this is a structured component (starts with "retriever" or "weight")
+        if (RETRIEVER_FIELD.match(firstFieldName, parser.getDeprecationHandler())
+            || WEIGHT_FIELD.match(firstFieldName, parser.getDeprecationHandler())) {
+            // This is a structured component - parse manually
+            RetrieverBuilder retriever = null;
+            Float weight = null;
+
+            do {
+                String fieldName = parser.currentName();
+                if (RETRIEVER_FIELD.match(fieldName, parser.getDeprecationHandler())) {
+                    if (retriever != null) {
+                        throw new ParsingException(parser.getTokenLocation(), "only one retriever can be specified");
+                    }
+                    parser.nextToken();
+                    if (parser.currentToken() != XContentParser.Token.START_OBJECT) {
+                        throw new ParsingException(parser.getTokenLocation(), "retriever must be an object");
+                    }
+                    parser.nextToken();
+                    String retrieverType = parser.currentName();
+                    retriever = parser.namedObject(RetrieverBuilder.class, retrieverType, context);
+                    context.trackRetrieverUsage(retriever.getName());
+                    parser.nextToken();
+                } else if (WEIGHT_FIELD.match(fieldName, parser.getDeprecationHandler())) {
+                    if (weight != null) {
+                        throw new ParsingException(parser.getTokenLocation(), "[weight] field can only be specified once");
+                    }
+                    parser.nextToken();
+                    weight = parser.floatValue();
+                } else {
+                    throw new ParsingException(
+                        parser.getTokenLocation(),
+                        "unknown field [{}], expected [{}] or [{}]",
+                        fieldName,
+                        RETRIEVER_FIELD.getPreferredName(),
+                        WEIGHT_FIELD.getPreferredName()
+                    );
+                }
+            } while (parser.nextToken() == XContentParser.Token.FIELD_NAME);
+
+            if (retriever == null) {
+                throw new ParsingException(parser.getTokenLocation(), "retriever component must contain a retriever");
+            }
+
+            return new RRFRetrieverComponent(retriever, weight);
+        } else {
+            RetrieverBuilder retriever = parser.namedObject(RetrieverBuilder.class, firstFieldName, context);
+            context.trackRetrieverUsage(retriever.getName());
+            if (parser.nextToken() != XContentParser.Token.END_OBJECT) {
+                throw new ParsingException(parser.getTokenLocation(), "unknown field [{}] after retriever", parser.currentName());
+            }
+            return new RRFRetrieverComponent(retriever, DEFAULT_WEIGHT);
+        }
+    }
+}