Disallow query rule retriever as a sub retriever

kderusso · Oct 18, 2024 · 5588c89 · 5588c89
1 parent 7aff304
commit 5588c89
Show file tree

Hide file tree

Showing 4 changed files with 105 additions and 0 deletions.
diff --git a/server/src/main/java/org/elasticsearch/search/retriever/CompoundRetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/CompoundRetrieverBuilder.java
@@ -50,6 +50,15 @@ public record RetrieverSource(RetrieverBuilder retriever, SearchSourceBuilder so
     protected final List<RetrieverSource> innerRetrievers;
 
     protected CompoundRetrieverBuilder(List<RetrieverSource> innerRetrievers, int rankWindowSize) {
+
+        for (RetrieverSource innerRetriever : innerRetrievers) {
+            if (innerRetriever.retriever.isAllowedAsChildRetriever() == false) {
+                throw new IllegalArgumentException(
+                    "[" + innerRetriever.retriever.getName() + "] must be a standalone or top level retriever only"
+                );
+            }
+        }
+
         this.rankWindowSize = rankWindowSize;
         this.innerRetrievers = innerRetrievers;
     }

diff --git a/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java
@@ -182,6 +182,14 @@ public boolean isCompound() {
         return false;
     }
 
+    /**
+     * Determines if this retriever is allowed as a child retriever, or if it must be used
+     * as a top level or standalone retriever only
+     */
+    public boolean isAllowedAsChildRetriever() {
+        return true;
+    }
+
     protected RankDoc[] rankDocs = null;
 
     public RetrieverBuilder() {}

diff --git a/...rc/yamlRestTest/resources/rest-api-spec/test/entsearch/rules/80_query_rules_retriever.yml b/...rc/yamlRestTest/resources/rest-api-spec/test/entsearch/rules/80_query_rules_retriever.yml
@@ -127,3 +127,86 @@ setup:
   - match: { hits.hits.0._id: foo }
   - match: { hits.hits.1._id: foo2 }
 
+
+---
+"query rules retriever combined with rrf and pagination":
+
+  - do:
+      search:
+        index: test-index1
+        body:
+          size: 1
+          from: 1
+          retriever:
+            rule:
+              match_criteria:
+                foo: foo
+                bar: bar
+              ruleset_ids:
+                test-ruleset
+              retriever:
+                rrf:
+                  retrievers: [
+                    {
+                      standard: {
+                        query: {
+                          query_string: {
+                            query: bar
+                          }
+                        }
+                      }
+                    },
+                    {
+                      standard: {
+                        query: {
+                          query_string: {
+                            query: baz
+                          }
+                        }
+                      }
+                    }
+                  ]
+
+  - match: { hits.total.value: 4 }
+  - length: { hits.hits: 1 }
+  - match: { hits.hits.0._id: foo2 }
+
+---
+"query rules will error if defined as a sub-retriever":
+
+  - do:
+      catch: /\[rule\] must be a standalone or top level retriever only/
+      search:
+        index: test-index1
+        body:
+          retriever:
+            rrf:
+              retrievers: [
+                {
+                  standard: {
+                    query: {
+                      query_string: {
+                        query: bar
+                      }
+                    }
+                  }
+                },
+                {
+                  rule: {
+                    match_criteria: {
+                      foo: foo,
+                      bar: bar
+                    },
+                    ruleset_ids: test-ruleset,
+                    retriever: {
+                      standard: {
+                        query: {
+                          query_string: {
+                            query: baz
+                          }
+                        }
+                      }
+                    }
+                  }
+                }
+              ]
diff --git a/...n/java/org/elasticsearch/xpack/application/rules/retriever/QueryRuleRetrieverBuilder.java b/...n/java/org/elasticsearch/xpack/application/rules/retriever/QueryRuleRetrieverBuilder.java
@@ -120,6 +120,11 @@ public String getName() {
         return NAME;
     }
 
+    @Override
+    public boolean isAllowedAsChildRetriever() {
+        return false;
+    }
+
     @Override
     protected SearchSourceBuilder createSearchSourceBuilder(PointInTimeBuilder pit, RetrieverBuilder retrieverBuilder) {
         var ret = super.createSearchSourceBuilder(pit, retrieverBuilder);