From 5588c899e885568cca9f731f579e2ecebcd46d77 Mon Sep 17 00:00:00 2001
From: Kathleen DeRusso <kathleen.derusso@elastic.co>
Date: Fri, 18 Oct 2024 16:42:09 -0400
Subject: [PATCH] Disallow query rule retriever as a sub retriever

---
 .../retriever/CompoundRetrieverBuilder.java   |  9 ++
 .../search/retriever/RetrieverBuilder.java    |  8 ++
 .../rules/80_query_rules_retriever.yml        | 83 +++++++++++++++++++
 .../retriever/QueryRuleRetrieverBuilder.java  |  5 ++
 4 files changed, 105 insertions(+)

diff --git a/server/src/main/java/org/elasticsearch/search/retriever/CompoundRetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/CompoundRetrieverBuilder.java
index 9c938fe6b1aab..46f61d4cdc114 100644
--- 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
index fa8d2ee1cbefc..e72cf44c443a0 100644
--- 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/x-pack/plugin/ent-search/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/entsearch/rules/80_query_rules_retriever.yml b/x-pack/plugin/ent-search/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/entsearch/rules/80_query_rules_retriever.yml
index 5ac8143cdfb41..4dd8d9b84fbf6 100644
--- a/x-pack/plugin/ent-search/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/entsearch/rules/80_query_rules_retriever.yml
+++ b/x-pack/plugin/ent-search/qa/rest/src/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/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/rules/retriever/QueryRuleRetrieverBuilder.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/rules/retriever/QueryRuleRetrieverBuilder.java
index 768e7157ed475..e5ac22043f429 100644
--- a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/rules/retriever/QueryRuleRetrieverBuilder.java
+++ b/x-pack/plugin/ent-search/src/main/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);