query optimizer

Author: david-leifker

metadata-io/src/main/java/com/linkedin/metadata/graph/elastic/ESGraphQueryDAO.java

@@ -3,6 +3,7 @@
 import static com.linkedin.metadata.aspect.models.graph.Edge.*;
 import static com.linkedin.metadata.graph.elastic.ElasticSearchGraphService.*;
 import static com.linkedin.metadata.search.utils.ESUtils.applyResultLimit;
+import static com.linkedin.metadata.search.utils.ESUtils.queryOptimize;
 
 import com.datahub.util.exception.ESQueryException;
 import com.google.common.annotations.VisibleForTesting;
@@ -178,7 +179,7 @@ private static BoolQueryBuilder getAggregationFilter(
     BoolQueryBuilder subFilter = QueryBuilders.boolQuery();
     TermQueryBuilder relationshipTypeTerm =
         QueryBuilders.termQuery(RELATIONSHIP_TYPE, pair.getValue().getType()).caseInsensitive(true);
-    subFilter.must(relationshipTypeTerm);
+    subFilter.filter(relationshipTypeTerm);
 
     String sourceType;
     String destinationType;
@@ -192,10 +193,10 @@ private static BoolQueryBuilder getAggregationFilter(
 
     TermQueryBuilder sourceTypeTerm =
         QueryBuilders.termQuery(SOURCE_TYPE, sourceType).caseInsensitive(true);
-    subFilter.must(sourceTypeTerm);
+    subFilter.filter(sourceTypeTerm);
     TermQueryBuilder destinationTypeTerm =
         QueryBuilders.termQuery(DESTINATION_TYPE, destinationType).caseInsensitive(true);
-    subFilter.must(destinationTypeTerm);
+    subFilter.filter(destinationTypeTerm);
     return subFilter;
   }
 
@@ -261,6 +262,7 @@ public static BoolQueryBuilder buildQuery(
                   relationshipQuery.should(
                       QueryBuilders.termQuery(RELATIONSHIP_TYPE, relationshipType)));
       relationshipQuery.minimumShouldMatch(1);
+
       finalQuery.filter(relationshipQuery);
     }
 
@@ -1140,6 +1142,10 @@ private List<LineageRelationship> relationshipsGroupQuery(
 
     Set<Urn> entityUrnSet = new HashSet<>(entityUrns);
 
+    if (config.getSearch().getGraph().isQueryOptimization()) {
+      queryOptimize(baseQuery, false);
+    }
+
     // Get search responses as a stream with pagination
     Stream<SearchResponse> responseStream =
         executeGroupByLineageSearchQuery(
@@ -1167,16 +1173,16 @@ private List<LineageRelationship> relationshipsGroupQuery(
   private static BoolQueryBuilder getOutGoingEdgeQuery(
       @Nonnull Set<Urn> urns, @Nonnull Set<EdgeInfo> outgoingEdges) {
     BoolQueryBuilder outgoingEdgeQuery = QueryBuilders.boolQuery();
-    outgoingEdgeQuery.must(buildUrnFilters(urns, SOURCE));
-    outgoingEdgeQuery.must(buildEdgeFilters(outgoingEdges));
+    outgoingEdgeQuery.filter(buildUrnFilters(urns, SOURCE));
+    outgoingEdgeQuery.filter(buildEdgeFilters(outgoingEdges));
     return outgoingEdgeQuery;
   }
 
   private static BoolQueryBuilder getIncomingEdgeQuery(
       @Nonnull Set<Urn> urns, Set<EdgeInfo> incomingEdges) {
     BoolQueryBuilder incomingEdgeQuery = QueryBuilders.boolQuery();
-    incomingEdgeQuery.must(buildUrnFilters(urns, DESTINATION));
-    incomingEdgeQuery.must(buildEdgeFilters(incomingEdges));
+    incomingEdgeQuery.filter(buildUrnFilters(urns, DESTINATION));
+    incomingEdgeQuery.filter(buildEdgeFilters(incomingEdges));
     return incomingEdgeQuery;
   }
 

metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java

@@ -28,6 +28,7 @@
 import com.linkedin.metadata.search.elasticsearch.query.filter.QueryFilterRewriterContext;
 import com.linkedin.metadata.utils.CriterionUtils;
 import io.datahubproject.metadata.context.OperationContext;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
@@ -935,4 +936,317 @@ public static int applyResultLimit(@Nonnull ElasticSearchConfiguration config, i
     }
     return count;
   }
+
+  /**
+   * Performs a single optimization pass on the query tree.
+   *
+   * @param query The QueryBuilder to optimize
+   * @param considerScore If true, moves to must() for scoring; if false, moves to filter()
+   * @return true if any optimization was performed, false otherwise
+   */
+  static boolean optimizePass(QueryBuilder query, boolean considerScore) {
+    if (!(query instanceof BoolQueryBuilder)) {
+      return false;
+    }
+
+    BoolQueryBuilder boolQuery = (BoolQueryBuilder) query;
+    boolean changed = false;
+
+    // First, recursively optimize all nested queries
+    // Must clauses
+    for (int i = 0; i < boolQuery.must().size(); i++) {
+      changed |= optimizePass(boolQuery.must().get(i), considerScore);
+    }
+
+    // Filter clauses - always use considerScore=false for filters
+    for (int i = 0; i < boolQuery.filter().size(); i++) {
+      changed |= optimizePass(boolQuery.filter().get(i), false);
+    }
+
+    // Should clauses
+    for (int i = 0; i < boolQuery.should().size(); i++) {
+      changed |= optimizePass(boolQuery.should().get(i), considerScore);
+    }
+
+    // MustNot clauses - always use considerScore=false
+    for (int i = 0; i < boolQuery.mustNot().size(); i++) {
+      changed |= optimizePass(boolQuery.mustNot().get(i), false);
+    }
+
+    // After optimizing children, check if this query itself can be optimized
+
+    // Optimization 1: Convert single should clause with minimumShouldMatch=1
+    if (isOptimizableShould(boolQuery)) {
+      // Get the should clause
+      QueryBuilder shouldClause = boolQuery.should().get(0);
+
+      // Remove the should clause
+      boolQuery.should().clear();
+
+      // Reset minimumShouldMatch since there are no more should clauses
+      boolQuery.minimumShouldMatch(null);
+
+      // Add to appropriate clause type
+      if (considerScore) {
+        boolQuery.must(shouldClause);
+      } else {
+        boolQuery.filter(shouldClause);
+      }
+
+      changed = true;
+    }
+
+    // Optimization 2: Flatten nested bool queries with only filter clauses
+    if (canFlattenFilters(boolQuery)) {
+      changed |= flattenFilters(boolQuery);
+    }
+
+    // Note: We don't handle unwrapping here because it requires replacing the query object
+    // which can't be done in-place. Unwrapping is handled at the top level in queryOptimize.
+
+    return changed;
+  }
+
+  /**
+   * Checks if the BoolQueryBuilder has a single should clause with minimumShouldMatch=1
+   *
+   * @param query The BoolQueryBuilder to check
+   * @return true if the query can be optimized, false otherwise
+   */
+  static boolean isOptimizableShould(BoolQueryBuilder query) {
+    // Check if there's exactly one should clause
+    if (query.should().size() != 1) {
+      return false;
+    }
+
+    // Check if minimumShouldMatch is set to "1"
+    String minShouldMatch = query.minimumShouldMatch();
+    if (minShouldMatch == null) {
+      return false;
+    }
+
+    // Handle different formats of minimumShouldMatch
+    // It could be "1", "1%", or other formats
+    if (minShouldMatch.equals("1") || minShouldMatch.equals("100%")) {
+      return true;
+    }
+
+    // Check if it's a percentage format like "100%"
+    if (minShouldMatch.endsWith("%")) {
+      try {
+        int percentage = Integer.parseInt(minShouldMatch.substring(0, minShouldMatch.length() - 1));
+        // With only 1 should clause, 100% means that 1 clause must match
+        return percentage == 100;
+      } catch (NumberFormatException e) {
+        return false;
+      }
+    }
+
+    return false;
+  }
+
+  /**
+   * Checks if a bool query can have its filter clauses flattened. This is possible when it has
+   * filter clauses containing other bool queries that only have filter clauses.
+   *
+   * @param query The BoolQueryBuilder to check
+   * @return true if filters can be flattened
+   */
+  static boolean canFlattenFilters(BoolQueryBuilder query) {
+    // Check each filter clause
+    for (QueryBuilder filter : query.filter()) {
+      if (filter instanceof BoolQueryBuilder) {
+        BoolQueryBuilder nestedBool = (BoolQueryBuilder) filter;
+        // If the nested bool query has only filter clauses (no must, should, mustNot),
+        // then we can flatten it
+        if (nestedBool.must().isEmpty()
+            && nestedBool.should().isEmpty()
+            && nestedBool.mustNot().isEmpty()
+            && !nestedBool.filter().isEmpty()) {
+          return true;
+        }
+      }
+    }
+    return false;
+  }
+
+  /**
+   * Flattens nested filter clauses in a bool query. Extracts filter clauses from nested bool
+   * queries that only contain filters and adds them directly to the parent.
+   *
+   * @param query The BoolQueryBuilder to flatten
+   * @return true if any flattening occurred
+   */
+  private static boolean flattenFilters(BoolQueryBuilder query) {
+    boolean changed = false;
+    List<QueryBuilder> filtersToAdd = new ArrayList<>();
+    List<QueryBuilder> filtersToRemove = new ArrayList<>();
+
+    // Identify filters that can be flattened
+    for (QueryBuilder filter : query.filter()) {
+      if (filter instanceof BoolQueryBuilder) {
+        BoolQueryBuilder nestedBool = (BoolQueryBuilder) filter;
+        // If the nested bool query has only filter clauses, extract them
+        if (nestedBool.must().isEmpty()
+            && nestedBool.should().isEmpty()
+            && nestedBool.mustNot().isEmpty()
+            && !nestedBool.filter().isEmpty()) {
+
+          // Mark for removal
+          filtersToRemove.add(filter);
+
+          // Extract all filters from the nested bool query
+          filtersToAdd.addAll(nestedBool.filter());
+
+          changed = true;
+        }
+      }
+    }
+
+    // Apply the changes
+    if (changed) {
+      // Remove the nested bool queries
+      for (QueryBuilder toRemove : filtersToRemove) {
+        query.filter().remove(toRemove);
+      }
+
+      // Add the extracted filters directly
+      for (QueryBuilder toAdd : filtersToAdd) {
+        query.filter(toAdd);
+      }
+    }
+
+    return changed;
+  }
+
+  /**
+   * Checks if a bool query can be unwrapped (replaced by its single clause). This is possible when
+   * the bool query has exactly one clause total.
+   *
+   * @param query The BoolQueryBuilder to check
+   * @return true if the query can be unwrapped
+   */
+  static boolean canUnwrap(BoolQueryBuilder query) {
+    int totalClauses =
+        query.must().size()
+            + query.filter().size()
+            + query.should().size()
+            + query.mustNot().size();
+
+    // Can unwrap if there's exactly one clause and no minimumShouldMatch constraint
+    return totalClauses == 1 && query.minimumShouldMatch() == null;
+  }
+
+  /**
+   * Fully optimizes a query by running optimization passes until no more changes occur.
+   *
+   * @param query The QueryBuilder to optimize
+   * @param considerScore If true, preserves scoring; if false, optimizes for filtering
+   * @return The optimized query (may be a different instance if unwrapping occurred)
+   */
+  public static QueryBuilder queryOptimize(QueryBuilder query, boolean considerScore) {
+    if (query == null) {
+      return null;
+    }
+
+    // For non-bool queries, return as-is
+    if (!(query instanceof BoolQueryBuilder)) {
+      return query;
+    }
+
+    BoolQueryBuilder boolQuery = (BoolQueryBuilder) query;
+
+    // Keep optimizing until no more changes
+    boolean changed;
+    int iterations = 0;
+    int maxIterations = 100; // Safety limit to prevent infinite loops
+
+    do {
+      changed = optimizePass(boolQuery, considerScore);
+      iterations++;
+
+      if (iterations >= maxIterations) {
+        // Log warning or throw exception in production code
+        break;
+      }
+    } while (changed);
+
+    // After all optimization passes, check if we can unwrap at the top level
+    if (canUnwrap(boolQuery)) {
+      // Return the single clause directly
+      if (!boolQuery.must().isEmpty()) {
+        return queryOptimize(boolQuery.must().get(0), considerScore);
+      } else if (!boolQuery.filter().isEmpty()) {
+        return queryOptimize(boolQuery.filter().get(0), false);
+      } else if (!boolQuery.should().isEmpty()) {
+        return queryOptimize(boolQuery.should().get(0), considerScore);
+      } else if (!boolQuery.mustNot().isEmpty()) {
+        // mustNot can't stand alone, keep it wrapped
+        return boolQuery;
+      }
+    }
+
+    return boolQuery;
+  }
+
+  /**
+   * Recursively optimizes and unwraps nested bool queries where possible. This is a helper method
+   * that can be used to handle unwrapping during optimization.
+   *
+   * @param query The QueryBuilder to process
+   * @param parent The parent BoolQueryBuilder (null if this is the root)
+   * @param clauseType The type of clause this query is in the parent (must, filter, should,
+   *     mustNot)
+   * @param index The index of this query in the parent's clause list
+   * @return true if any changes were made
+   */
+  private static boolean optimizeAndUnwrapNested(
+      QueryBuilder query, BoolQueryBuilder parent, String clauseType, int index) {
+    if (!(query instanceof BoolQueryBuilder)) {
+      return false;
+    }
+
+    BoolQueryBuilder boolQuery = (BoolQueryBuilder) query;
+
+    // First optimize the query
+    boolean changed =
+        optimizePass(boolQuery, "must".equals(clauseType) || "should".equals(clauseType));
+
+    // If this query can be unwrapped and we have a parent, unwrap it
+    if (parent != null && canUnwrap(boolQuery)) {
+      QueryBuilder unwrapped = null;
+
+      if (!boolQuery.must().isEmpty()) {
+        unwrapped = boolQuery.must().get(0);
+      } else if (!boolQuery.filter().isEmpty()) {
+        unwrapped = boolQuery.filter().get(0);
+      } else if (!boolQuery.should().isEmpty()) {
+        unwrapped = boolQuery.should().get(0);
+      } else if (!boolQuery.mustNot().isEmpty()) {
+        // mustNot needs special handling - can't unwrap
+        return changed;
+      }
+
+      if (unwrapped != null) {
+        // Replace in parent
+        switch (clauseType) {
+          case "must":
+            parent.must().set(index, unwrapped);
+            break;
+          case "filter":
+            parent.filter().set(index, unwrapped);
+            break;
+          case "should":
+            parent.should().set(index, unwrapped);
+            break;
+          case "mustNot":
+            parent.mustNot().set(index, unwrapped);
+            break;
+        }
+        changed = true;
+      }
+    }
+
+    return changed;
+  }
 }