Implement support for weighted rrf

docs/changelog/130658.yaml

@@ -1,5 +1,5 @@
 pr: 130658
-summary: Implement support for weighted rrf
+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/rrf/RRFRetrieverBuilder.java

@@ -57,7 +57,6 @@ public final class RRFRetrieverBuilder extends CompoundRetrieverBuilder<RRFRetri
     public static final ParseField RANK_CONSTANT_FIELD = new ParseField("rank_constant");
     public static final ParseField FIELDS_FIELD = new ParseField("fields");
     public static final ParseField QUERY_FIELD = new ParseField("query");
-    public static final ParseField WEIGHTS_FIELD = new ParseField("weights");
 
     public static final int DEFAULT_RANK_CONSTANT = 60;
 
@@ -92,7 +91,6 @@ public final class RRFRetrieverBuilder extends CompoundRetrieverBuilder<RRFRetri
         PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), QUERY_FIELD);
         PARSER.declareInt(ConstructingObjectParser.optionalConstructorArg(), RANK_WINDOW_SIZE_FIELD);
         PARSER.declareInt(ConstructingObjectParser.optionalConstructorArg(), RANK_CONSTANT_FIELD);
-        PARSER.declareFloatArray(ConstructingObjectParser.optionalConstructorArg(), WEIGHTS_FIELD);
         RetrieverBuilder.declareBaseParserFields(PARSER);
     }
 
@@ -265,17 +263,18 @@ protected RetrieverBuilder doRewrite(QueryRewriteContext ctx) {
                 );
             }
 
-            List<RetrieverBuilder> fieldsInnerRetrievers = MultiFieldsInnerRetrieverUtils.generateInnerRetrievers(
+            List<RetrieverSource> fieldsInnerRetrievers = MultiFieldsInnerRetrieverUtils.generateInnerRetrievers(
                 fields,
                 query,
                 localIndicesMetadata.values(),
                 r -> {
                     List<RetrieverSource> retrievers = new ArrayList<>(r.size());
-                    for (var retriever : r) {
+                    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();
                     }
-                    float[] weights = new float[retrievers.size()];
-                    Arrays.fill(weights, 1.0f);
                     return new RRFRetrieverBuilder(retrievers, null, null, rankWindowSize, rankConstant, weights);
                 },
                 w -> {
@@ -285,20 +284,13 @@ protected RetrieverBuilder doRewrite(QueryRewriteContext ctx) {
                         );
                     }
                 }
-            );
+            ).stream().map(RetrieverSource::from).toList();
 
             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)
-                int size = fieldsInnerRetrievers.size();
-                List<RetrieverSource> sources = new ArrayList<>(size);
-                float[] weights = new float[size];
-                Arrays.fill(weights, RRFRetrieverComponent.DEFAULT_WEIGHT);
-                for (int i = 0; i < size; i++) {
-                    sources.add(RetrieverSource.from(fieldsInnerRetrievers.get(i)));
-                    weights[i] = RRFRetrieverComponent.DEFAULT_WEIGHT;
-                }
-                rewritten = new RRFRetrieverBuilder(sources, null, null, rankWindowSize, rankConstant, weights);
+                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
@@ -333,13 +325,6 @@ public void doToXContent(XContentBuilder builder, Params params) throws IOExcept
 
         builder.field(RANK_WINDOW_SIZE_FIELD.getPreferredName(), rankWindowSize);
         builder.field(RANK_CONSTANT_FIELD.getPreferredName(), rankConstant);
-        if (weights.length > 0) {
-            builder.startArray(WEIGHTS_FIELD.getPreferredName());
-            for (float weight : weights) {
-                builder.value(weight);
-            }
-            builder.endArray();
-        }
     }
 
     // ---- FOR TESTING XCONTENT PARSING ----

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

@@ -10,7 +10,6 @@
 import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.search.retriever.RetrieverBuilder;
 import org.elasticsearch.search.retriever.RetrieverParserContext;
-import org.elasticsearch.xcontent.ConstructingObjectParser;
 import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
@@ -19,9 +18,6 @@
 import java.io.IOException;
 import java.util.Objects;
 
-import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg;
-import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg;
-
 public class RRFRetrieverComponent implements ToXContentObject {
 
     public static final ParseField RETRIEVER_FIELD = new ParseField("retriever");
@@ -56,26 +52,6 @@ public XContentBuilder toXContent(XContentBuilder builder, ToXContentObject.Para
         return builder;
     }
 
-    @SuppressWarnings("unchecked")
-    static final ConstructingObjectParser<RRFRetrieverComponent, RetrieverParserContext> PARSER = new ConstructingObjectParser<>(
-        "rrf_component",
-        false,
-        (args, context) -> {
-            RetrieverBuilder retrieverBuilder = (RetrieverBuilder) args[0];
-            Float weight = (Float) args[1];
-            return new RRFRetrieverComponent(retrieverBuilder, weight);
-        }
-    );
-
-    static {
-        PARSER.declareNamedObject(constructorArg(), (p, c, n) -> {
-            RetrieverBuilder innerRetriever = p.namedObject(RetrieverBuilder.class, n, c);
-            c.trackRetrieverUsage(innerRetriever.getName());
-            return innerRetriever;
-        }, RETRIEVER_FIELD);
-        PARSER.declareFloat(optionalConstructorArg(), WEIGHT_FIELD);
-    }
-
     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());
@@ -106,6 +82,9 @@ public static RRFRetrieverComponent fromXContent(XContentParser parser, Retrieve
                         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);
@@ -118,9 +97,6 @@ public static RRFRetrieverComponent fromXContent(XContentParser parser, Retrieve
                     parser.nextToken();
                     weight = parser.floatValue();
                 } else {
-                    if (retriever != null) {
-                        throw new ParsingException(parser.getTokenLocation(), "only one retriever can be specified");
-                    }
                     throw new ParsingException(
                         parser.getTokenLocation(),
                         "unknown field [{}], expected [{}] or [{}]",
@@ -139,7 +115,8 @@ public static RRFRetrieverComponent fromXContent(XContentParser parser, Retrieve
         } else {
             RetrieverBuilder retriever = parser.namedObject(RetrieverBuilder.class, firstFieldName, context);
             context.trackRetrieverUsage(retriever.getName());
-            while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
+            if (parser.nextToken() != XContentParser.Token.END_OBJECT) {
+                throw new ParsingException(parser.getTokenLocation(), "unknown field [{}] after retriever", parser.currentName());
             }
             return new RRFRetrieverComponent(retriever, DEFAULT_WEIGHT);
         }

x-pack/plugin/rank-rrf/src/test/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilderParsingTests.java

@@ -57,14 +57,13 @@ public static RRFRetrieverBuilder createRandomRRFRetrieverBuilder() {
 
         int retrieverCount = randomIntBetween(2, 50);
         List<CompoundRetrieverBuilder.RetrieverSource> innerRetrievers = new ArrayList<>(retrieverCount);
+        float[] weights = new float[retrieverCount];
+        int i = 0;
         while (retrieverCount > 0) {
             innerRetrievers.add(CompoundRetrieverBuilder.RetrieverSource.from(TestRetrieverBuilder.createRandomTestRetrieverBuilder()));
+            weights[i++] = randomFloat();
             --retrieverCount;
         }
-        float[] weights = new float[innerRetrievers.size()];
-        for (int i = 0; i < innerRetrievers.size(); i++) {
-            weights[i] = randomFloat();
-        }
 
         return new RRFRetrieverBuilder(innerRetrievers, fields, query, rankWindowSize, rankConstant, weights);
     }
@@ -274,7 +273,7 @@ public void testRRFRetrieverComponentErrorCases() throws IOException {
             }
             """;
 
-        expectParsingException(multipleRetrieversContent, "only one retriever can be specified");
+        expectParsingException(multipleRetrieversContent, "unknown field [standard], expected [retriever] or [weight]");
 
         // Test case 2: Weight without retriever
         String weightOnlyContent = """
@@ -336,7 +335,28 @@ public void testRRFRetrieverComponentErrorCases() throws IOException {
             }
             """;
 
-        expectParsingException(negativeWeightContent, "weight] must be non-negative");
+        expectParsingException(negativeWeightContent, "[weight] must be non-negative");
+
+        // Test case 5: Retriever as non-object
+        String retrieverAsStringContent = """
+            {
+              "retriever": {
+                "rrf": {
+                  "retrievers": [
+                    {
+                      "retriever": "not_an_object"
+                    }
+                  ],
+                  "rank_window_size": 100,
+                  "rank_constant": 10,
+                  "min_score": 20.0,
+                  "_name": "foo_rrf"
+                }
+              }
+            }
+            """;
+
+        expectParsingException(retrieverAsStringContent, "retriever must be an object");
     }
 
     private void expectParsingException(String restContent, String expectedMessageFragment) throws IOException {

x-pack/plugin/rank-rrf/src/test/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilderTests.java

@@ -39,6 +39,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.BiConsumer;
 
 import static org.elasticsearch.search.rank.RankBuilder.DEFAULT_RANK_WINDOW_SIZE;
 import static org.hamcrest.Matchers.instanceOf;
@@ -86,6 +87,17 @@ public void testRetrieverExtractionErrors() throws IOException {
     }
 
     public void testRRFRetrieverParsingSyntax() throws IOException {
+        BiConsumer<String, float[]> testCase = (json, expectedWeights) -> {
+            try (XContentParser parser = createParser(JsonXContent.jsonXContent, json)) {
+                SearchSourceBuilder ssb = new SearchSourceBuilder().parseXContent(parser, true, nf -> true);
+                assertThat(ssb.retriever(), instanceOf(RRFRetrieverBuilder.class));
+                RRFRetrieverBuilder rrf = (RRFRetrieverBuilder) ssb.retriever();
+                assertArrayEquals(expectedWeights, rrf.weights(), 0.001f);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        };
+
         String legacyJson = """
             {
               "retriever": {
@@ -98,12 +110,7 @@ public void testRRFRetrieverParsingSyntax() throws IOException {
               }
             }
             """;
-        try (XContentParser parser = createParser(JsonXContent.jsonXContent, legacyJson)) {
-            SearchSourceBuilder ssb = new SearchSourceBuilder().parseXContent(parser, true, nf -> true);
-            assertThat(ssb.retriever(), instanceOf(RRFRetrieverBuilder.class));
-            RRFRetrieverBuilder rrf = (RRFRetrieverBuilder) ssb.retriever();
-            assertArrayEquals(new float[] { 1.0f, 1.0f }, rrf.weights(), 0.001f);
-        }
+        testCase.accept(legacyJson, new float[] { 1.0f, 1.0f });
 
         String weightedJson = """
             {
@@ -117,12 +124,7 @@ public void testRRFRetrieverParsingSyntax() throws IOException {
               }
             }
             """;
-        try (XContentParser parser = createParser(JsonXContent.jsonXContent, weightedJson)) {
-            SearchSourceBuilder ssb = new SearchSourceBuilder().parseXContent(parser, true, nf -> true);
-            assertThat(ssb.retriever(), instanceOf(RRFRetrieverBuilder.class));
-            RRFRetrieverBuilder rrf = (RRFRetrieverBuilder) ssb.retriever();
-            assertArrayEquals(new float[] { 2.5f, 0.5f }, rrf.weights(), 0.001f);
-        }
+        testCase.accept(weightedJson, new float[] { 2.5f, 0.5f });
 
         String mixedJson = """
             {
@@ -136,12 +138,7 @@ public void testRRFRetrieverParsingSyntax() throws IOException {
               }
             }
             """;
-        try (XContentParser parser = createParser(JsonXContent.jsonXContent, mixedJson)) {
-            SearchSourceBuilder ssb = new SearchSourceBuilder().parseXContent(parser, true, nf -> true);
-            assertThat(ssb.retriever(), instanceOf(RRFRetrieverBuilder.class));
-            RRFRetrieverBuilder rrf = (RRFRetrieverBuilder) ssb.retriever();
-            assertArrayEquals(new float[] { 1.0f, 0.6f }, rrf.weights(), 0.001f);
-        }
+        testCase.accept(mixedJson, new float[] { 1.0f, 0.6f });
     }
 
     public void testMultiFieldsParamsRewrite() {

x-pack/plugin/rank-rrf/src/yamlRestTest/resources/rest-api-spec/test/rrf/320_rrf_weighted_retriever.yml

@@ -2,7 +2,7 @@ setup:
   - requires:
       cluster_features: [ "rrf_retriever.weighted_support" ]
       reason: "RRF retriever Weighted support"
-      test_runner_features: [ "contains" ]
+      test_runner_features: [ "contains", "close_to" ]
   - do:
       indices.create:
         index: restaurants
@@ -58,7 +58,7 @@ setup:
   - match: { hits.hits.0._id: "1" }
 
 ---
-"Weighted RRF retriever defaults to weight 1":
+"Weighted RRF retriever allows optional weight field":
   - do:
       search:
         index: restaurants
@@ -80,6 +80,37 @@ setup:
   - match: { hits.total.value: 3 }
   - match: { hits.hits.0._id: "1" }
 
+---
+"Weighted RRF retriever changes result order":
+  - do:
+      search:
+        index: restaurants
+        body:
+          retriever:
+            rrf:
+              retrievers:
+                - retriever:
+                    standard:
+                      query:
+                        match:
+                          description: "pizza"
+                  weight: 0.1
+                - retriever:
+                    standard:
+                      query:
+                        match:
+                          description: "burgers"
+                  weight: 0.9
+  - match: { hits.total.value: 2 }
+  - match: { hits.hits.0._id: "2" }
+  - match: { hits.hits.1._id: "1" }
+  # Document 2: matches "burgers" with weight 0.9
+  # RRF score = 1/(60+1) * 0.9 = 0.01475
+  - close_to: {hits.hits.0._score: {value:  0.01475, error: 0.0001}}
+  # Document 1: matches "pizza" with weight 0.1
+  # RRF score = 1/(60+1) * 0.1 = 0.00164
+  - close_to: {hits.hits.1._score: {value:  0.00164, error: 0.0001}}
+
 ---
 "Weighted RRF retriever errors on negative weight":
   - do:
@@ -104,4 +135,5 @@ setup:
                           description: "pizza"
                   weight: 0.7
   - match: { error.type: "x_content_parse_exception" }
-  - contains: { error.reason: "failed to parse field" }
+  - contains: { error.caused_by.reason: "[weight] must be non-negative, found [-0.5]" }
+