feat(search): search query rewriter (datahub-project#11279)

dushayntAW · Sep 12, 2024 · d385076 · d385076
1 parent f23765e
commit d385076
Show file tree

Hide file tree

Showing 50 changed files with 1,917 additions and 126 deletions.
diff --git a/datahub-graphql-core/src/main/resources/search.graphql b/datahub-graphql-core/src/main/resources/search.graphql
@@ -550,6 +550,21 @@ enum FilterOperator {
   Represent the relation less than or equal to, e.g. ownerCount <= 3
   """
   LESS_THAN_OR_EQUAL_TO
+
+  """
+  Represent the relation: URN field any nested children in addition to the given URN
+  """
+  DESCENDANTS_INCL
+
+  """
+  Represent the relation: URN field matches any nested parent in addition to the given URN
+  """
+  ANCESTORS_INCL
+
+  """
+  Represent the relation: URN field matches any nested child or parent in addition to the given URN
+  """
+  RELATED_INCL
 }
 
 """

diff --git a/entity-registry/src/main/java/com/linkedin/metadata/aspect/GraphRetriever.java b/entity-registry/src/main/java/com/linkedin/metadata/aspect/GraphRetriever.java
@@ -5,6 +5,7 @@
 import com.linkedin.metadata.query.filter.RelationshipFilter;
 import com.linkedin.metadata.query.filter.SortCriterion;
 import java.util.List;
+import java.util.function.Function;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
@@ -40,4 +41,60 @@ RelatedEntitiesScrollResult scrollRelatedEntities(
       int count,
       @Nullable Long startTimeMillis,
       @Nullable Long endTimeMillis);
+
+  /**
+   * Consume graph edges
+   *
+   * @param consumer consumer function, return true to exit early
+   * @param sourceTypes
+   * @param sourceEntityFilter
+   * @param destinationTypes
+   * @param destinationEntityFilter
+   * @param relationshipTypes
+   * @param relationshipFilter
+   * @param sortCriteria
+   * @param count
+   * @param startTimeMillis
+   * @param endTimeMillis
+   */
+  default void consumeRelatedEntities(
+      @Nonnull Function<RelatedEntitiesScrollResult, Boolean> consumer,
+      @Nullable List<String> sourceTypes,
+      @Nonnull Filter sourceEntityFilter,
+      @Nullable List<String> destinationTypes,
+      @Nonnull Filter destinationEntityFilter,
+      @Nonnull List<String> relationshipTypes,
+      @Nonnull RelationshipFilter relationshipFilter,
+      @Nonnull List<SortCriterion> sortCriteria,
+      int count,
+      @Nullable Long startTimeMillis,
+      @Nullable Long endTimeMillis) {
+
+    String scrollId = null;
+    boolean exitCriteria = false;
+
+    while (!exitCriteria) {
+      RelatedEntitiesScrollResult result =
+          scrollRelatedEntities(
+              sourceTypes,
+              sourceEntityFilter,
+              destinationTypes,
+              destinationEntityFilter,
+              relationshipTypes,
+              relationshipFilter,
+              sortCriteria,
+              scrollId,
+              count,
+              startTimeMillis,
+              endTimeMillis);
+
+      exitCriteria = consumer.apply(result);
+
+      if (result == null || result.getEntities().isEmpty() || result.getScrollId() == null) {
+        exitCriteria = true;
+      } else {
+        scrollId = result.getScrollId();
+      }
+    }
+  }
 }
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java
@@ -23,6 +23,7 @@
 import com.linkedin.metadata.models.annotation.SearchableAnnotation;
 import com.linkedin.metadata.query.SearchFlags;
 import com.linkedin.metadata.query.filter.Filter;
+import com.linkedin.metadata.search.elasticsearch.query.filter.QueryFilterRewriteChain;
 import com.linkedin.metadata.search.elasticsearch.query.request.SearchRequestHandler;
 import com.linkedin.metadata.search.utils.ESUtils;
 import com.linkedin.metadata.search.utils.SearchUtils;
@@ -69,6 +70,7 @@ public class ESBrowseDAO {
   private final RestHighLevelClient client;
   @Nonnull private final SearchConfiguration searchConfiguration;
   @Nullable private final CustomSearchConfiguration customSearchConfiguration;
+  @Nonnull private final QueryFilterRewriteChain queryFilterRewriteChain;
 
   private static final String BROWSE_PATH = "browsePaths";
   private static final String BROWSE_PATH_DEPTH = "browsePaths.length";
@@ -607,7 +609,8 @@ private QueryBuilder buildQueryStringV2(
 
     EntitySpec entitySpec = opContext.getEntityRegistry().getEntitySpec(entityName);
     QueryBuilder query =
-        SearchRequestHandler.getBuilder(entitySpec, searchConfiguration, customSearchConfiguration)
+        SearchRequestHandler.getBuilder(
+                entitySpec, searchConfiguration, customSearchConfiguration, queryFilterRewriteChain)
             .getQuery(
                 finalOpContext,
                 input,
@@ -623,7 +626,7 @@ private QueryBuilder buildQueryStringV2(
 
     queryBuilder.filter(
         SearchRequestHandler.getFilterQuery(
-            finalOpContext, filter, entitySpec.getSearchableFieldTypes()));
+            finalOpContext, filter, entitySpec.getSearchableFieldTypes(), queryFilterRewriteChain));
 
     return queryBuilder;
   }
@@ -643,7 +646,11 @@ private QueryBuilder buildQueryStringBrowseAcrossEntities(
     final BoolQueryBuilder queryBuilder = QueryBuilders.boolQuery();
 
     QueryBuilder query =
-        SearchRequestHandler.getBuilder(entitySpecs, searchConfiguration, customSearchConfiguration)
+        SearchRequestHandler.getBuilder(
+                entitySpecs,
+                searchConfiguration,
+                customSearchConfiguration,
+                queryFilterRewriteChain)
             .getQuery(
                 finalOpContext,
                 input,
@@ -669,7 +676,8 @@ private QueryBuilder buildQueryStringBrowseAcrossEntities(
                       return set1;
                     }));
     queryBuilder.filter(
-        SearchRequestHandler.getFilterQuery(finalOpContext, filter, searchableFields));
+        SearchRequestHandler.getFilterQuery(
+            finalOpContext, filter, searchableFields, queryFilterRewriteChain));
 
     return queryBuilder;
   }

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java
@@ -20,6 +20,7 @@
 import com.linkedin.metadata.search.FilterValueArray;
 import com.linkedin.metadata.search.ScrollResult;
 import com.linkedin.metadata.search.SearchResult;
+import com.linkedin.metadata.search.elasticsearch.query.filter.QueryFilterRewriteChain;
 import com.linkedin.metadata.search.elasticsearch.query.request.AggregationQueryBuilder;
 import com.linkedin.metadata.search.elasticsearch.query.request.AutocompleteRequestHandler;
 import com.linkedin.metadata.search.elasticsearch.query.request.SearchAfterWrapper;
@@ -76,6 +77,7 @@ public class ESSearchDAO {
   private final String elasticSearchImplementation;
   @Nonnull private final SearchConfiguration searchConfiguration;
   @Nullable private final CustomSearchConfiguration customSearchConfiguration;
+  @Nonnull private final QueryFilterRewriteChain queryFilterRewriteChain;
 
   public long docCount(@Nonnull OperationContext opContext, @Nonnull String entityName) {
     return docCount(opContext, entityName, null);
@@ -88,7 +90,10 @@ public long docCount(
         new CountRequest(opContext.getSearchContext().getIndexConvention().getIndexName(entitySpec))
             .query(
                 SearchRequestHandler.getFilterQuery(
-                    opContext, filter, entitySpec.getSearchableFieldTypes()));
+                    opContext,
+                    filter,
+                    entitySpec.getSearchableFieldTypes(),
+                    queryFilterRewriteChain));
     try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "docCount").time()) {
       return client.count(countRequest, RequestOptions.DEFAULT).getCount();
     } catch (IOException e) {
@@ -115,7 +120,10 @@ private SearchResult executeAndExtract(
       return transformIndexIntoEntityName(
           opContext.getSearchContext().getIndexConvention(),
           SearchRequestHandler.getBuilder(
-                  entitySpec, searchConfiguration, customSearchConfiguration)
+                  entitySpec,
+                  searchConfiguration,
+                  customSearchConfiguration,
+                  queryFilterRewriteChain)
               .extractResult(opContext, searchResponse, filter, from, size));
     } catch (Exception e) {
       log.error("Search query failed", e);
@@ -212,7 +220,10 @@ private ScrollResult executeAndExtract(
       return transformIndexIntoEntityName(
           opContext.getSearchContext().getIndexConvention(),
           SearchRequestHandler.getBuilder(
-                  entitySpecs, searchConfiguration, customSearchConfiguration)
+                  entitySpecs,
+                  searchConfiguration,
+                  customSearchConfiguration,
+                  queryFilterRewriteChain)
               .extractScrollResult(
                   opContext, searchResponse, filter, keepAlive, size, supportsPointInTime()));
     } catch (Exception e) {
@@ -255,7 +266,11 @@ public SearchResult search(
     Filter transformedFilters = transformFilterForEntities(postFilters, indexConvention);
     // Step 1: construct the query
     final SearchRequest searchRequest =
-        SearchRequestHandler.getBuilder(entitySpecs, searchConfiguration, customSearchConfiguration)
+        SearchRequestHandler.getBuilder(
+                entitySpecs,
+                searchConfiguration,
+                customSearchConfiguration,
+                queryFilterRewriteChain)
             .getSearchRequest(
                 opContext, finalInput, transformedFilters, sortCriteria, from, size, facets);
     searchRequest.indices(
@@ -288,7 +303,8 @@ public SearchResult filter(
     EntitySpec entitySpec = opContext.getEntityRegistry().getEntitySpec(entityName);
     Filter transformedFilters = transformFilterForEntities(filters, indexConvention);
     final SearchRequest searchRequest =
-        SearchRequestHandler.getBuilder(entitySpec, searchConfiguration, customSearchConfiguration)
+        SearchRequestHandler.getBuilder(
+                entitySpec, searchConfiguration, customSearchConfiguration, queryFilterRewriteChain)
             .getFilterRequest(opContext, transformedFilters, sortCriteria, from, size);
 
     searchRequest.indices(indexConvention.getIndexName(entitySpec));
@@ -321,7 +337,8 @@ public AutoCompleteResult autoComplete(
       EntitySpec entitySpec = opContext.getEntityRegistry().getEntitySpec(entityName);
       IndexConvention indexConvention = opContext.getSearchContext().getIndexConvention();
       AutocompleteRequestHandler builder =
-          AutocompleteRequestHandler.getBuilder(entitySpec, customSearchConfiguration);
+          AutocompleteRequestHandler.getBuilder(
+              entitySpec, customSearchConfiguration, queryFilterRewriteChain);
       SearchRequest req =
           builder.getSearchRequest(
               opContext,
@@ -366,7 +383,11 @@ public Map<String, Long> aggregateByValue(
     }
     IndexConvention indexConvention = opContext.getSearchContext().getIndexConvention();
     final SearchRequest searchRequest =
-        SearchRequestHandler.getBuilder(entitySpecs, searchConfiguration, customSearchConfiguration)
+        SearchRequestHandler.getBuilder(
+                entitySpecs,
+                searchConfiguration,
+                customSearchConfiguration,
+                queryFilterRewriteChain)
             .getAggregationRequest(
                 opContext,
                 field,
@@ -481,7 +502,7 @@ private SearchRequest getScrollRequest(
     }
 
     return SearchRequestHandler.getBuilder(
-            entitySpecs, searchConfiguration, customSearchConfiguration)
+            entitySpecs, searchConfiguration, customSearchConfiguration, queryFilterRewriteChain)
         .getSearchRequest(
             opContext, finalInput, postFilters, sortCriteria, sort, pitId, keepAlive, size, facets);
   }