diff --git a/server/src/main/java/org/elasticsearch/search/retriever/KnnRetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/KnnRetrieverBuilder.java index a44f509158c53..22ab15e6eb141 100644 --- a/server/src/main/java/org/elasticsearch/search/retriever/KnnRetrieverBuilder.java +++ b/server/src/main/java/org/elasticsearch/search/retriever/KnnRetrieverBuilder.java @@ -139,7 +139,7 @@ public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException { } @Override - public QueryBuilder originalQuery(QueryBuilder leadQuery) { + public QueryBuilder topDocsQuery(QueryBuilder leadQuery) { // TODO nested + inner_hits BoolQueryBuilder ret = new BoolQueryBuilder().must(leadQuery) .should(new ExactKnnQueryBuilder(knnSearchBuilder.getQueryVector(), knnSearchBuilder.getField())); diff --git a/server/src/main/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilder.java index 68541372b5157..2cbeae6f809ab 100644 --- a/server/src/main/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilder.java +++ b/server/src/main/java/org/elasticsearch/search/retriever/RankDocsRetrieverBuilder.java @@ -74,10 +74,10 @@ public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException { } @Override - public QueryBuilder originalQuery(QueryBuilder leadQuery) { + public QueryBuilder topDocsQuery(QueryBuilder leadQuery) { DisMaxQueryBuilder disMax = new DisMaxQueryBuilder().tieBreaker(0f); for (var source : sources) { - var query = source.originalQuery(leadQuery); + var query = source.topDocsQuery(leadQuery); if (query != null) { if (source.retrieverName != null) { query.queryName(source.retrieverName); @@ -105,7 +105,7 @@ public void extractToSearchSourceBuilder(SearchSourceBuilder searchSourceBuilder for (var preFilterQueryBuilder : preFilterQueryBuilders) { bq.filter(preFilterQueryBuilder); } - QueryBuilder originalQuery = originalQuery(rankQuery); + QueryBuilder originalQuery = topDocsQuery(rankQuery); if (originalQuery != null) { bq.should(originalQuery); } 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 eccefc2e5ea61..a3c069359596a 100644 --- a/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java +++ b/server/src/main/java/org/elasticsearch/search/retriever/RetrieverBuilder.java @@ -34,9 +34,10 @@ /** * A retriever represents an API element that returns an ordered list of top * documents. These can be obtained from a query, from another retriever, etc. - * Internally, a {@link RetrieverBuilder} is just a wrapper for other search - * elements that are extracted into a {@link SearchSourceBuilder}. The advantage - * retrievers have is in the API they appear as a tree-like structure enabling + * Internally, a {@link RetrieverBuilder} is first rewritten into its simplest + * form and then its elements are extracted into a {@link SearchSourceBuilder}. + * + * The advantage retrievers have is in the API they appear as a tree-like structure enabling * easier reasoning about what a search does. * * This is the base class for all other retrievers. This class does not support @@ -181,7 +182,7 @@ public List getPreFilterQueryBuilders() { } /** - * Determines if this retriever contains sub-retrievers that need to be rewritten into simpler forms. + * Determines if this retriever contains sub-retrievers that need to be executed prior to search. */ public boolean isCompound() { return false; @@ -193,10 +194,12 @@ public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException { } /** - * Returns the original {@link QueryBuilder} used to compute the top documents. - * @param leadQuery + * This function is called by compound {@link RetrieverBuilder} to return the original query that + * was used by this retriever to compute its top documents. + * + * @param leadQuery The query identifying the top documents of the parent retriever. */ - public abstract QueryBuilder originalQuery(QueryBuilder leadQuery); + public abstract QueryBuilder topDocsQuery(QueryBuilder leadQuery); /** * This method is called at the end of rewrite on the final retriever. diff --git a/server/src/main/java/org/elasticsearch/search/retriever/StandardRetrieverBuilder.java b/server/src/main/java/org/elasticsearch/search/retriever/StandardRetrieverBuilder.java index 6b61bee45c009..2412af905117d 100644 --- a/server/src/main/java/org/elasticsearch/search/retriever/StandardRetrieverBuilder.java +++ b/server/src/main/java/org/elasticsearch/search/retriever/StandardRetrieverBuilder.java @@ -134,7 +134,7 @@ public RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOException { } @Override - public QueryBuilder originalQuery(QueryBuilder leadQuery) { + public QueryBuilder topDocsQuery(QueryBuilder leadQuery) { /** * What actions should we take with {@link KnnVectorQueryBuilder} or {@link MultiTermQueryBuilder} when a * compound retriever executes the original queries? Our goal is to retain these queries in scenarios where @@ -142,7 +142,7 @@ public QueryBuilder originalQuery(QueryBuilder leadQuery) { * retrievers since they will be executed twice: once as a must clause at this level and a second time as a * should clause at the upper level (compound retriever). * Therefore, it would be beneficial to rewrite these queries at the upper level to focus solely on - * scoring/matching similar to what {@link RetrieverBuilder#originalQuery(QueryBuilder)} is doing. + * scoring/matching similar to what {@link RetrieverBuilder#topDocsQuery(QueryBuilder)} is doing. */ if (preFilterQueryBuilders.isEmpty()) { return queryBuilder; diff --git a/test/framework/src/main/java/org/elasticsearch/search/retriever/TestRetrieverBuilder.java b/test/framework/src/main/java/org/elasticsearch/search/retriever/TestRetrieverBuilder.java index 739494e20f8b3..84def9b2d0bbf 100644 --- a/test/framework/src/main/java/org/elasticsearch/search/retriever/TestRetrieverBuilder.java +++ b/test/framework/src/main/java/org/elasticsearch/search/retriever/TestRetrieverBuilder.java @@ -67,7 +67,7 @@ public TestRetrieverBuilder(String value) { } @Override - public QueryBuilder originalQuery(QueryBuilder leadQuery) { + public QueryBuilder topDocsQuery(QueryBuilder leadQuery) { throw new UnsupportedOperationException("only used for parsing tests"); } diff --git a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilder.java b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilder.java index e13e21f083081..5695c1c4ad9c4 100644 --- a/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilder.java +++ b/x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilder.java @@ -217,7 +217,7 @@ public void onFailure(Exception e) { } @Override - public QueryBuilder originalQuery(QueryBuilder leadQuery) { + public QueryBuilder topDocsQuery(QueryBuilder leadQuery) { throw new IllegalStateException(NAME + " cannot be nested"); }