Make NestedHelper a utility class (#118071)

Noticed instantiating these instances taking a visible and unexpected amount of CPU in profiles (probably from bootstrapping the lambda/callsite for the predicate). This fixes the logic to effectively disappear from profiling and makes it easier to reason about as well by removing the indirect use of the search context and just explicitly passing it around. No need to instantiate instances of this thing either, escape analysis probably isn't able to remove it because of the recursive instance method calls.
elastic · Dec 5, 2024 · 5d1bca3 · 5d1bca3
1 parent 9d35053
commit 5d1bca3
Show file tree

Hide file tree

Showing 7 changed files with 185 additions and 179 deletions.
diff --git a/server/src/main/java/org/elasticsearch/index/query/NestedQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/NestedQueryBuilder.java
@@ -321,8 +321,7 @@ public static <E extends Exception> Query toQuery(
 
         // ToParentBlockJoinQuery requires that the inner query only matches documents
         // in its child space
-        NestedHelper nestedHelper = new NestedHelper(context.nestedLookup(), context::isFieldMapped);
-        if (nestedHelper.mightMatchNonNestedDocs(innerQuery, path)) {
+        if (NestedHelper.mightMatchNonNestedDocs(innerQuery, path, context)) {
             innerQuery = Queries.filtered(innerQuery, mapper.nestedTypeFilter());
         }
 

diff --git a/server/src/main/java/org/elasticsearch/index/search/NestedHelper.java b/server/src/main/java/org/elasticsearch/index/search/NestedHelper.java
@@ -21,61 +21,53 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TermInSetQuery;
 import org.apache.lucene.search.TermQuery;
-import org.elasticsearch.index.mapper.NestedLookup;
 import org.elasticsearch.index.mapper.NestedObjectMapper;
-
-import java.util.function.Predicate;
+import org.elasticsearch.index.query.SearchExecutionContext;
 
 /** Utility class to filter parent and children clauses when building nested
  * queries. */
 public final class NestedHelper {
 
-    private final NestedLookup nestedLookup;
-    private final Predicate<String> isMappedFieldPredicate;
-
-    public NestedHelper(NestedLookup nestedLookup, Predicate<String> isMappedFieldPredicate) {
-        this.nestedLookup = nestedLookup;
-        this.isMappedFieldPredicate = isMappedFieldPredicate;
-    }
+    private NestedHelper() {}
 
     /** Returns true if the given query might match nested documents. */
-    public boolean mightMatchNestedDocs(Query query) {
+    public static boolean mightMatchNestedDocs(Query query, SearchExecutionContext searchExecutionContext) {
         if (query instanceof ConstantScoreQuery) {
-            return mightMatchNestedDocs(((ConstantScoreQuery) query).getQuery());
+            return mightMatchNestedDocs(((ConstantScoreQuery) query).getQuery(), searchExecutionContext);
         } else if (query instanceof BoostQuery) {
-            return mightMatchNestedDocs(((BoostQuery) query).getQuery());
+            return mightMatchNestedDocs(((BoostQuery) query).getQuery(), searchExecutionContext);
         } else if (query instanceof MatchAllDocsQuery) {
             return true;
         } else if (query instanceof MatchNoDocsQuery) {
             return false;
         } else if (query instanceof TermQuery) {
             // We only handle term(s) queries and range queries, which should already
             // cover a high majority of use-cases
-            return mightMatchNestedDocs(((TermQuery) query).getTerm().field());
+            return mightMatchNestedDocs(((TermQuery) query).getTerm().field(), searchExecutionContext);
         } else if (query instanceof TermInSetQuery tis) {
             if (tis.getTermsCount() > 0) {
-                return mightMatchNestedDocs(tis.getField());
+                return mightMatchNestedDocs(tis.getField(), searchExecutionContext);
             } else {
                 return false;
             }
         } else if (query instanceof PointRangeQuery) {
-            return mightMatchNestedDocs(((PointRangeQuery) query).getField());
+            return mightMatchNestedDocs(((PointRangeQuery) query).getField(), searchExecutionContext);
         } else if (query instanceof IndexOrDocValuesQuery) {
-            return mightMatchNestedDocs(((IndexOrDocValuesQuery) query).getIndexQuery());
+            return mightMatchNestedDocs(((IndexOrDocValuesQuery) query).getIndexQuery(), searchExecutionContext);
         } else if (query instanceof final BooleanQuery bq) {
             final boolean hasRequiredClauses = bq.clauses().stream().anyMatch(BooleanClause::isRequired);
             if (hasRequiredClauses) {
                 return bq.clauses()
                     .stream()
                     .filter(BooleanClause::isRequired)
                     .map(BooleanClause::query)
-                    .allMatch(this::mightMatchNestedDocs);
+                    .allMatch(f -> mightMatchNestedDocs(f, searchExecutionContext));
             } else {
                 return bq.clauses()
                     .stream()
                     .filter(c -> c.occur() == Occur.SHOULD)
                     .map(BooleanClause::query)
-                    .anyMatch(this::mightMatchNestedDocs);
+                    .anyMatch(f -> mightMatchNestedDocs(f, searchExecutionContext));
             }
         } else if (query instanceof ESToParentBlockJoinQuery) {
             return ((ESToParentBlockJoinQuery) query).getPath() != null;
@@ -85,7 +77,7 @@ public boolean mightMatchNestedDocs(Query query) {
     }
 
     /** Returns true if a query on the given field might match nested documents. */
-    boolean mightMatchNestedDocs(String field) {
+    private static boolean mightMatchNestedDocs(String field, SearchExecutionContext searchExecutionContext) {
         if (field.startsWith("_")) {
             // meta field. Every meta field behaves differently, eg. nested
             // documents have the same _uid as their parent, put their path in
@@ -94,50 +86,50 @@ boolean mightMatchNestedDocs(String field) {
             // we might add a nested filter when it is nor required.
             return true;
         }
-        if (isMappedFieldPredicate.test(field) == false) {
+        if (searchExecutionContext.isFieldMapped(field) == false) {
             // field does not exist
             return false;
         }
-        return nestedLookup.getNestedParent(field) != null;
+        return searchExecutionContext.nestedLookup().getNestedParent(field) != null;
     }
 
     /** Returns true if the given query might match parent documents or documents
      *  that are nested under a different path. */
-    public boolean mightMatchNonNestedDocs(Query query, String nestedPath) {
+    public static boolean mightMatchNonNestedDocs(Query query, String nestedPath, SearchExecutionContext searchExecutionContext) {
         if (query instanceof ConstantScoreQuery) {
-            return mightMatchNonNestedDocs(((ConstantScoreQuery) query).getQuery(), nestedPath);
+            return mightMatchNonNestedDocs(((ConstantScoreQuery) query).getQuery(), nestedPath, searchExecutionContext);
         } else if (query instanceof BoostQuery) {
-            return mightMatchNonNestedDocs(((BoostQuery) query).getQuery(), nestedPath);
+            return mightMatchNonNestedDocs(((BoostQuery) query).getQuery(), nestedPath, searchExecutionContext);
         } else if (query instanceof MatchAllDocsQuery) {
             return true;
         } else if (query instanceof MatchNoDocsQuery) {
             return false;
         } else if (query instanceof TermQuery) {
-            return mightMatchNonNestedDocs(((TermQuery) query).getTerm().field(), nestedPath);
+            return mightMatchNonNestedDocs(searchExecutionContext, ((TermQuery) query).getTerm().field(), nestedPath);
         } else if (query instanceof TermInSetQuery tis) {
             if (tis.getTermsCount() > 0) {
-                return mightMatchNonNestedDocs(tis.getField(), nestedPath);
+                return mightMatchNonNestedDocs(searchExecutionContext, tis.getField(), nestedPath);
             } else {
                 return false;
             }
         } else if (query instanceof PointRangeQuery) {
-            return mightMatchNonNestedDocs(((PointRangeQuery) query).getField(), nestedPath);
+            return mightMatchNonNestedDocs(searchExecutionContext, ((PointRangeQuery) query).getField(), nestedPath);
         } else if (query instanceof IndexOrDocValuesQuery) {
-            return mightMatchNonNestedDocs(((IndexOrDocValuesQuery) query).getIndexQuery(), nestedPath);
+            return mightMatchNonNestedDocs(((IndexOrDocValuesQuery) query).getIndexQuery(), nestedPath, searchExecutionContext);
         } else if (query instanceof final BooleanQuery bq) {
             final boolean hasRequiredClauses = bq.clauses().stream().anyMatch(BooleanClause::isRequired);
             if (hasRequiredClauses) {
                 return bq.clauses()
                     .stream()
                     .filter(BooleanClause::isRequired)
                     .map(BooleanClause::query)
-                    .allMatch(q -> mightMatchNonNestedDocs(q, nestedPath));
+                    .allMatch(q -> mightMatchNonNestedDocs(q, nestedPath, searchExecutionContext));
             } else {
                 return bq.clauses()
                     .stream()
                     .filter(c -> c.occur() == Occur.SHOULD)
                     .map(BooleanClause::query)
-                    .anyMatch(q -> mightMatchNonNestedDocs(q, nestedPath));
+                    .anyMatch(q -> mightMatchNonNestedDocs(q, nestedPath, searchExecutionContext));
             }
         } else {
             return true;
@@ -146,7 +138,7 @@ public boolean mightMatchNonNestedDocs(Query query, String nestedPath) {
 
     /** Returns true if a query on the given field might match parent documents
      *  or documents that are nested under a different path. */
-    boolean mightMatchNonNestedDocs(String field, String nestedPath) {
+    private static boolean mightMatchNonNestedDocs(SearchExecutionContext searchExecutionContext, String field, String nestedPath) {
         if (field.startsWith("_")) {
             // meta field. Every meta field behaves differently, eg. nested
             // documents have the same _uid as their parent, put their path in
@@ -155,9 +147,10 @@ boolean mightMatchNonNestedDocs(String field, String nestedPath) {
             // we might add a nested filter when it is nor required.
             return true;
         }
-        if (isMappedFieldPredicate.test(field) == false) {
+        if (searchExecutionContext.isFieldMapped(field) == false) {
             return false;
         }
+        var nestedLookup = searchExecutionContext.nestedLookup();
         String nestedParent = nestedLookup.getNestedParent(field);
         if (nestedParent == null || nestedParent.startsWith(nestedPath) == false) {
             // the field is not a sub field of the nested path

diff --git a/server/src/main/java/org/elasticsearch/search/DefaultSearchContext.java b/server/src/main/java/org/elasticsearch/search/DefaultSearchContext.java
@@ -444,10 +444,9 @@ public void preProcess() {
     public Query buildFilteredQuery(Query query) {
         List<Query> filters = new ArrayList<>();
         NestedLookup nestedLookup = searchExecutionContext.nestedLookup();
-        NestedHelper nestedHelper = new NestedHelper(nestedLookup, searchExecutionContext::isFieldMapped);
         if (nestedLookup != NestedLookup.EMPTY
-            && nestedHelper.mightMatchNestedDocs(query)
-            && (aliasFilter == null || nestedHelper.mightMatchNestedDocs(aliasFilter))) {
+            && NestedHelper.mightMatchNestedDocs(query, searchExecutionContext)
+            && (aliasFilter == null || NestedHelper.mightMatchNestedDocs(aliasFilter, searchExecutionContext))) {
             filters.add(Queries.newNonNestedFilter(searchExecutionContext.indexVersionCreated()));
         }
 

diff --git a/server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java b/server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java
@@ -481,10 +481,9 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
             }
             parentBitSet = context.bitsetFilter(parentFilter);
             if (filterQuery != null) {
-                NestedHelper nestedHelper = new NestedHelper(context.nestedLookup(), context::isFieldMapped);
                 // We treat the provided filter as a filter over PARENT documents, so if it might match nested documents
                 // we need to adjust it.
-                if (nestedHelper.mightMatchNestedDocs(filterQuery)) {
+                if (NestedHelper.mightMatchNestedDocs(filterQuery, context)) {
                     // Ensure that the query only returns parent documents matching `filterQuery`
                     filterQuery = Queries.filtered(filterQuery, parentFilter);
                 }