Implement support for weighted rrf #130658

mridula-s109 · 2025-07-04T23:39:23Z

Implement support for weighted RRF

Summary

This PR implements support for weighted RRF (Reciprocal Rank Fusion) retrievers, allowing users to specify custom weights for each sub-retriever within an RRF retriever configuration. This addresses a common customer request to customize the influence of different retrievers in the RRF scoring process.

Core Implementation

Enhanced RRFRetrieverBuilder: Extended to support both weighted and non-weighted retriever configurations
New RRFRetrieverComponent: Added weight validation and handling for individual retrievers
Backward Compatibility: Maintains support for existing RRF configurations without weights (default weight = 1.0)

Key Features

1. Weighted Retriever Support

Users can now specify weights for individual retrievers:

{
  "retriever": {
    "rrf": {
      "retrievers": [
        {
          "retriever": {
            "standard": { 
              "query": { "match": { "description": "pizza" } } 
            }
          },
          "weight": 0.7
        },
        {
          "retriever": {
            "knn": { 
              "field": "vector", 
              "query_vector": [1,2,3], 
              "k": 10 
            }
          },
          "weight": 0.3
        }
      ]
    }
  }
}

…ilder

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

github-actions · 2025-07-17T12:06:40Z

🔍 Preview links for changed docs

More links …

Mikep86

Nice progress, this is looking better. The hybrid object parsing is definitely a beast.

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

Mikep86 · 2025-07-24T19:57:46Z

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

+    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);
+        }
+    );


Where do we actually use this?

It is used in the static method:
public static RRFRetrieverComponent fromXContent(XContentParser parser, RetrieverParserContext context) which typically calls:
return PARSER.apply(parser, context);

I get that's where it's normally used, but where is it actually used in this implementation? RRFRetrieverComponen#fromXContent, in its current form, doesn't use it.

Mikep86 · 2025-07-24T20:14:03Z

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

+        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();
+                    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 {
+                    if (retriever != null) {
+                        throw new ParsingException(parser.getTokenLocation(), "only one retriever can be specified");
+                    }
+                    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 {


I understand this is complex, but is there an opportunity to use a ConstructingObjectParser here once we know this is a structured component? @ioanatia WDYT?

I am bit skeptical on that change, lets see what @ioanatia thinks as well.

Mikep86

A bit of cleanup left, but this is coming along 👍

Mikep86 · 2025-07-25T18:16:59Z

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

+                        retrievers.add(retriever.retrieverSource());
+                    }
+                    float[] weights = new float[retrievers.size()];
+                    Arrays.fill(weights, 1.0f);


We don't need to hard-code a weight of 1.0 here. The weight validator ensures that every WeightedRetrieverSource's weight is 1.0, so we can use that to populate the weights array.

Mikep86 · 2025-07-25T18:21:14Z

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

-            ).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)
-                rewritten = new RRFRetrieverBuilder(fieldsInnerRetrievers, rankWindowSize, rankConstant);
+                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);


Not a functional difference, but we could clean this up a bit and remove some duplicated logic:

Keep .stream().map(RetrieverSource::from).toList() to build a list of RetrieverSources. We need that anyways.

Use createDefaultWeights to create the weights array.

Mikep86 · 2025-07-25T18:27:36Z

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

+    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);
+        }
+    );


I get that's where it's normally used, but where is it actually used in this implementation? RRFRetrieverComponen#fromXContent, in its current form, doesn't use it.

Mikep86 · 2025-07-25T18:29:42Z

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

+        for (int i = 0; i < innerRetrievers.size(); i++) {
+            weights[i] = randomFloat();
+        }


Nit: we could populate the weights array in the while (retrieverCount > 0) loop

Mikep86 · 2025-07-25T18:35:24Z

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

+    public void testRRFRetrieverParsingWithDefaultWeights() throws IOException {
+        String restContent = """
+            {
+              "retriever": {
+                "rrf": {
+                  "retrievers": [
+                    {
+                      "test": {
+                        "value": "first"
+                      }
+                    },
+                    {
+                      "test": {
+                        "value": "second"
+                      }
+                    }
+                  ],
+                  "rank_window_size": 100,
+                  "rank_constant": 10,
+                  "min_score": 20.0,
+                  "_name": "foo_rrf"
+                }
+              }
+            }
+            """;
+        checkRRFRetrieverParsing(restContent);
+    }


How is this test functionally different than testRRFRetrieverParsing?

Mikep86 · 2025-07-25T18:36:27Z

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

+            }
+            """;
+
+        expectParsingException(negativeWeightContent, "weight] must be non-negative");


Nit: Missing a [ here

Mikep86 · 2025-07-25T18:41:25Z

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

+        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);
+        }


Nit: we could factor this duplicated code out. It doesn't necessarily need to be refactored into a named method, we could make a local BiConsumer (or CheckedBiConsumer) that takes the JSON string and the expected weight array.

Mikep86 · 2025-07-25T18:46:40Z

...in/rank-rrf/src/yamlRestTest/resources/rest-api-spec/test/rrf/320_rrf_weighted_retriever.yml

+  - match: { hits.hits.0._id: "1" }
+
+---
+"Weighted RRF retriever defaults to weight 1":


Very big nit here, but this test method is misleading. It doesn't actually check that that the default weight is one (testRRFRetrieverParsingSyntax does that). It only checks that weight is optional. Maybe change the test name to be more accurate?

Mikep86 · 2025-07-25T18:48:24Z

...in/rank-rrf/src/yamlRestTest/resources/rest-api-spec/test/rrf/320_rrf_weighted_retriever.yml

Could we add a test where we show that we can use weight to boost a document in the result set? In other words, show that by changing weight we can change the result order in an expected way.

mridula-s109 added 2 commits July 4, 2025 13:47

RRFRetrieverComponent added:

8be20f0

Modified parser, toXcontent and included component in the RetrieverBu…

a8f6487

…ilder

mridula-s109 marked this pull request as draft July 4, 2025 23:39

elasticsearchmachine added the v9.2.0 label Jul 4, 2025

[CI] Auto commit changes from spotless

e07c38d

mridula-s109 commented Jul 15, 2025

View reviewed changes

x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilder.java Outdated Show resolved Hide resolved

mridula-s109 commented Jul 15, 2025

View reviewed changes

x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilder.java Outdated Show resolved Hide resolved

mridula-s109 added 2 commits July 15, 2025 11:28

Resolved merge conflicts

33d3da4

Merge branch 'main' into SEARCH-1026-implement-support-for-weighted-rrf

5fb5568

Mikep86 self-requested a review July 15, 2025 13:32

mridula-s109 and others added 3 commits July 15, 2025 21:20

Fixed compile issues in tests

3ba149c

[CI] Auto commit changes from spotless

d5749f6

trying to resolve parse errros

7614936

mridula-s109 force-pushed the SEARCH-1026-implement-support-for-weighted-rrf branch from 7614936 to e7f0a90 Compare July 17, 2025 12:05

github-actions bot deployed to docs-preview July 17, 2025 12:06 View deployment

wip

a5d9e34

github-actions bot deployed to docs-preview July 17, 2025 12:19 View deployment

mridula-s109 force-pushed the SEARCH-1026-implement-support-for-weighted-rrf branch 4 times, most recently from 782e3ca to f1eede5 Compare July 17, 2025 16:55

github-actions bot had a problem deploying to docs-preview July 17, 2025 16:55 Failure

mridula-s109 force-pushed the SEARCH-1026-implement-support-for-weighted-rrf branch from f1eede5 to f5e1572 Compare July 17, 2025 16:56

github-actions bot deployed to docs-preview July 17, 2025 16:57 View deployment

mridula-s109 requested review from a team and removed request for a team July 17, 2025 16:59

github-actions bot deployed to docs-preview July 17, 2025 17:02 View deployment

Modified builder

0640099

mridula-s109 force-pushed the SEARCH-1026-implement-support-for-weighted-rrf branch from f1eede5 to 0640099 Compare July 17, 2025 17:05

mridula-s109 added 4 commits July 24, 2025 13:04

Fixed the relaxed constraints

fafb50f

Resolving issues

e535864

Resolved PR comments

78f8641

removed simplified rrf

02647b1

github-actions bot deployed to docs-preview July 24, 2025 12:05 View deployment

mridula-s109 added 3 commits July 24, 2025 15:18

changed the test file back to its original state

2010f3a

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

7433023

made parsing robust

a2bf4de

github-actions bot deployed to docs-preview July 24, 2025 15:28 View deployment

Merge branch 'main' into SEARCH-1026-implement-support-for-weighted-rrf

eebf577

github-actions bot deployed to docs-preview July 24, 2025 15:33 View deployment

Merge branch 'main' into SEARCH-1026-implement-support-for-weighted-rrf

0388abd

github-actions bot deployed to docs-preview July 24, 2025 15:45 View deployment

IT test reverted

74ed8db

github-actions bot deployed to docs-preview July 24, 2025 18:58 View deployment

mridula-s109 requested review from ioanatia and Mikep86 July 24, 2025 19:03

Mikep86 reviewed Jul 24, 2025

View reviewed changes

mridula-s109 added 3 commits July 25, 2025 10:29

Replaced the declareString array parser

6d7e8ff

Enforced weights as nonnull

f1e14ce

Fixed the weights null

fd30387

github-actions bot deployed to docs-preview July 25, 2025 09:58 View deployment

Empty weight shouldnt be serialised

3a82a28

github-actions bot deployed to docs-preview July 25, 2025 10:02 View deployment

[CI] Auto commit changes from spotless

77c14d3

github-actions bot deployed to docs-preview July 25, 2025 10:08 View deployment

Merge branch 'main' into SEARCH-1026-implement-support-for-weighted-rrf

45ca068

github-actions bot deployed to docs-preview July 25, 2025 10:10 View deployment

mridula-s109 requested a review from Mikep86 July 25, 2025 10:16

Mikep86 reviewed Jul 25, 2025

View reviewed changes

Implement support for weighted rrf #130658

Are you sure you want to change the base?

Implement support for weighted rrf #130658

Uh oh!

Conversation

mridula-s109 commented Jul 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Implement support for weighted RRF

Summary

Core Implementation

Key Features

1. Weighted Retriever Support

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Jul 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔍 Preview links for changed docs

Uh oh!

Mikep86 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mridula-s109 Jul 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Mikep86 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Mikep86 Jul 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mridula-s109 commented Jul 4, 2025 •

edited

Loading

github-actions bot commented Jul 17, 2025 •

edited

Loading

mridula-s109 Jul 25, 2025 •

edited

Loading

Mikep86 Jul 25, 2025 •

edited

Loading