Modified parser, toXcontent and included component in the RetrieverBuilder

Author: mridula-s109

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

@@ -37,7 +37,9 @@
 import java.util.Map;
 import java.util.Objects;
 
+import static org.elasticsearch.action.ValidateActions.addValidationError;
 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
@@ -57,33 +59,63 @@ 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<Object> rawRetrievers = args[0] == null ? List.of() : (List<Object>) 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);
+            List<RetrieverSource> innerRetrievers = new ArrayList<>(rawRetrievers.size());
+            float[] weights = new float[rawRetrievers.size()];
+            
+            int weightIndex = 0;
+            for (Object retrieverOrComponent : rawRetrievers) {
+                if (retrieverOrComponent instanceof RRFRetrieverComponent component) {
+                    innerRetrievers.add(RetrieverSource.from(component.retriever));
+                    weights[weightIndex++] = component.weight;
+                } else {
+                    RetrieverBuilder bareRetriever = (RetrieverBuilder) retrieverOrComponent;
+                    innerRetrievers.add(RetrieverSource.from(bareRetriever));
+                    weights[weightIndex++] = RRFRetrieverComponent.DEFAULT_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.declareObjectArray(optionalConstructorArg(),
+            (p, c) -> {
+                List<Object> list = new ArrayList<>();
+                while (p.nextToken() != XContentParser.Token.END_ARRAY) {
+                    if (p.currentToken() == XContentParser.Token.START_OBJECT &&
+                        p.nextToken() == XContentParser.Token.FIELD_NAME &&
+                        RRFRetrieverComponent.RETRIEVER_FIELD.match(p.currentName(), p.getDeprecationHandler())) {
+                        // Handle wrapped retriever with weight
+                        list.add(RRFRetrieverComponent.fromXContent(p, c));
+                    } else {
+                        // Handle bare retriever (legacy format)
+                        String name = p.currentName();
+                        RetrieverBuilder retrieverBuilder = p.namedObject(RetrieverBuilder.class, name, c);
+                        c.trackRetrieverUsage(retrieverBuilder.getName());
+                        p.nextToken();
+                        list.add(retrieverBuilder);
+                    }
+                }
+                return list;
+            },
+            RETRIEVERS_FIELD
+        );
+        
         PARSER.declareStringArray(optionalConstructorArg(), FIELDS_FIELD);
         PARSER.declareString(optionalConstructorArg(), QUERY_FIELD);
         PARSER.declareInt(optionalConstructorArg(), RANK_WINDOW_SIZE_FIELD);
@@ -103,21 +135,30 @@ 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<RetrieverSource> 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;
+        this.weights = weights;
     }
 
     public int rankConstant() {
@@ -137,6 +178,14 @@ public ActionRequestValidationException validate(
         boolean allowPartialSearchResults
     ) {
         validationException = super.validate(source, validationException, isScroll, allowPartialSearchResults);
+
+        if (this.weights != null) {
+            for (float weight : this.weights) {
+                if (weight < 0) {
+                    validationException = addValidationError("[weight] must be non-negative, found [" + weight + "]", validationException);
+                }
+            }
+        }
         return MultiFieldsInnerRetrieverUtils.validateParams(
             innerRetrievers,
             fields,
@@ -151,7 +200,7 @@ 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 +232,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
@@ -233,29 +282,34 @@ protected RetrieverBuilder doRewrite(QueryRewriteContext ctx) {
                 );
             }
 
-            List<RetrieverSource> fieldsInnerRetrievers = MultiFieldsInnerRetrieverUtils.generateInnerRetrievers(
+            List<RetrieverBuilder> fieldsInnerRetrievers = MultiFieldsInnerRetrieverUtils.generateInnerRetrievers(
                 fields,
                 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()];
+                    int i = 0;
+                    for(var retriever: r) {
+                        retrievers.add(retriever.retrieverSource());
+                        weights[i++] = retriever.weight();
+                    }
+                    return new RRFRetrieverBuilder(retrievers, null, null, rankWindowSize, rankConstant, weights);
                 },
                 w -> {
-                    if (w != 1.0f) {
+                    if (w < 0) {
                         throw new IllegalArgumentException(
-                            "[" + NAME + "] does not support per-field weights in [" + FIELDS_FIELD.getPreferredName() + "]"
+                            "[" + NAME + "] per-field weights must be non-negative"
                         );
                     }
                 }
-            ).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);
+                RRFRetrieverBuilder g = (RRFRetrieverBuilder) fieldsInnerRetrievers.get(0);
+                rewritten = new RRFRetrieverBuilder(g.innerRetrievers, null, null, rankWindowSize, rankConstant, g.weights);
                 rewritten.getPreFilterQueryBuilders().addAll(preFilterQueryBuilders);
             } else {
                 // Inner retriever list can be empty when using an index wildcard pattern that doesn't match any indices
@@ -274,21 +328,22 @@ public boolean doEquals(Object o) {
         return super.doEquals(o)
             && Objects.equals(fields, that.fields)
             && Objects.equals(query, that.query)
-            && rankConstant == that.rankConstant;
+            && rankConstant == that.rankConstant
+            && Arrays.equals(weights, that.weights);
     }
 
     @Override
     public int doHashCode() {
-        return Objects.hash(super.doHashCode(), fields, query, rankConstant);
+        return Objects.hash(super.doHashCode(), fields, query, rankConstant, Arrays.hashCode(weights));
     }
 
     @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++) {
+                new RRFRetrieverComponent(innerRetrievers.get(i).retriever(), this.weights[i]).toXContent(builder, params);
             }
             builder.endArray();
         }

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

@@ -20,8 +20,8 @@ public class RRFRetrieverComponent implements ToXContentObject {
 
     static final float DEFAULT_WEIGHT = 1f;
 
-    RetrieverBuilder retriever;
-    float weight;
+    final RetrieverBuilder retriever;
+    final float weight;
 
     public RRFRetrieverComponent(RetrieverBuilder retrieverBuilder, Float weight) {
         assert retrieverBuilder != null;