Add downstream validation of the rank_window_size in the compound ret…

…riever
jimczi · Dec 13, 2024 · 768d716 · 768d716
1 parent 5966d1c
commit 768d716
Show file tree

Hide file tree

Showing 3 changed files with 119 additions and 0 deletions.
diff --git a/...src/yamlRestTest/resources/rest-api-spec/test/search.retrievers/30_rescorer_retriever.yml b/...src/yamlRestTest/resources/rest-api-spec/test/search.retrievers/30_rescorer_retriever.yml
@@ -198,3 +198,28 @@ setup:
   - match: { hits.hits.1._score: 8.0 }
   - match: { hits.hits.2._id: "6" }
   - match: { hits.hits.2._score: 5.0 }
+
+---
+"Rescorer retriever and invalid window size":
+  - do:
+      catch: "/\\[rescorer\\] requires \\[rank_window_size: 5\\] be greater than or equal to \\[size: 10\\]/"
+      search:
+        index: test
+        body:
+          retriever:
+            rescorer:
+              rescore:
+                window_size: 5
+                query:
+                  rescore_query:
+                    rank_feature:
+                      field: "features.second_stage"
+                      linear: { }
+                  query_weight: 0
+              retriever:
+                standard:
+                  query:
+                    rank_feature:
+                      field: "features.first_stage"
+                      linear: { }
+          size: 10
diff --git a/server/src/main/java/org/elasticsearch/search/retriever/CompoundRetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/CompoundRetrieverBuilder.java
@@ -36,6 +36,7 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Locale;
 import java.util.Objects;
 
 import static org.elasticsearch.action.ValidateActions.addValidationError;
@@ -231,6 +232,14 @@ public ActionRequestValidationException validate(
         }
         for (RetrieverSource innerRetriever : innerRetrievers) {
             validationException = innerRetriever.retriever().validate(source, validationException, isScroll, allowPartialSearchResults);
+            if (innerRetriever.retriever() instanceof CompoundRetrieverBuilder<?> compoundChild) {
+                String errorMessage = String.format(
+                        Locale.ROOT,
+                        "[%s] requires [rank_window_size: %d] to be smaller than or equal to its sub retriever's %s [rank_window_size: %d]",
+                        this.getName(), rankWindowSize, compoundChild.getName(), compoundChild.rankWindowSize
+                );
+                validationException = addValidationError(errorMessage, validationException);
+            }
         }
         return validationException;
     }

diff --git a/...ugin/rank-rrf/src/yamlRestTest/resources/rest-api-spec/test/rrf/900_rrf_with_rescorer.yml b/...ugin/rank-rrf/src/yamlRestTest/resources/rest-api-spec/test/rrf/900_rrf_with_rescorer.yml
@@ -322,3 +322,88 @@ setup:
   - match:  { aggregations.1.buckets.1.key: 0 }
   - match:  { aggregations.1.buckets.1.doc_count: 4 }
 
+---
+"RRF with rescorer retriever and invalid window size":
+  - do:
+      catch: "/\\[rescorer\\] requires \\[rank_window_size: 5\\] be greater than or equal to \\[size: 10\\]/"
+      search:
+        index: test
+        body:
+          retriever:
+            rescorer:
+              rescore:
+                window_size: 5
+                query:
+                  rescore_query:
+                    rank_feature:
+                      field: "features.final_score"
+                      linear: { }
+                  query_weight: 0
+              retriever:
+                rrf:
+                  rank_window_size: 5
+                  retrievers: [
+                    {
+                      standard: {
+                        query: {
+                          rank_feature: {
+                            field: "features.first_query",
+                            linear: { }
+                          }
+                        }
+                      }
+                    },
+                    {
+                      standard: {
+                        query: {
+                          rank_feature: {
+                            field: "features.second_query",
+                            linear: { }
+                          }
+                        }
+                      }
+                    }
+                  ]
+          size: 10
+
+  - do:
+      catch: "/\\[rescorer\\] requires \\[rank_window_size: 10\\] to be smaller than or equal to its sub retriever's rrf \\[rank_window_size: 5\\]/"
+      search:
+        index: test
+        body:
+          retriever:
+            rescorer:
+              rescore:
+                window_size: 10
+                query:
+                  rescore_query:
+                    rank_feature:
+                      field: "features.final_score"
+                      linear: { }
+                  query_weight: 0
+              retriever:
+                rrf:
+                  rank_window_size: 5
+                  retrievers: [
+                    {
+                      standard: {
+                        query: {
+                          rank_feature: {
+                            field: "features.first_query",
+                            linear: { }
+                          }
+                        }
+                      }
+                    },
+                    {
+                      standard: {
+                        query: {
+                          rank_feature: {
+                            field: "features.second_query",
+                            linear: { }
+                          }
+                        }
+                      }
+                    }
+                  ]
+          size: 5