From 7f9edc752963ce08383170c761de45fe44fd97dd Mon Sep 17 00:00:00 2001
From: Indy Prentice <indy@Indys-MacBook-Pro.local>
Date: Tue, 25 Jul 2023 11:41:41 -0300
Subject: [PATCH] Also apply entity-specific search analyzers to scroll API

---
 .../metadata/search/SearchService.java        | 25 +++++++++++--------
 1 file changed, 15 insertions(+), 10 deletions(-)

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java b/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java
index d2739fc813b45..41c8de3c553db 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java
@@ -62,7 +62,7 @@ public Map<String, Long> docCountPerEntity(@Nonnull List<String> entityNames) {
   public SearchResult search(@Nonnull List<String> entityNames, @Nonnull String input, @Nullable Filter postFilters,
       @Nullable SortCriterion sortCriterion, int from, int size, @Nullable SearchFlags searchFlags) {
     SearchResult result =
-        _cachingEntitySearchService.search(entityNames, input, postFilters, sortCriterion, from, size, searchFlags, null);
+        _cachingEntitySearchService.search(getEntitiesToSearch(entityNames), input, postFilters, sortCriterion, from, size, searchFlags, null);
 
     try {
       return result.copy().setEntities(new SearchEntityArray(_searchRanker.rank(result.getEntities())));
@@ -109,14 +109,7 @@ public SearchResult searchAcrossEntities(@Nonnull List<String> entities, @Nonnul
       facets = new ArrayList<>(facets);
       facets.add(INDEX_VIRTUAL_FIELD);
     }
-    List<String> nonEmptyEntities;
-    List<String> lowercaseEntities = entities.stream().map(String::toLowerCase).collect(Collectors.toList());
-    try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getNonEmptyEntities").time()) {
-      nonEmptyEntities = _entityDocCountCache.getNonEmptyEntities();
-    }
-    if (!entities.isEmpty()) {
-      nonEmptyEntities = nonEmptyEntities.stream().filter(lowercaseEntities::contains).collect(Collectors.toList());
-    }
+    List<String> nonEmptyEntities = getEntitiesToSearch(entities);
     SearchResult result = _cachingEntitySearchService.search(nonEmptyEntities, input, postFilters, sortCriterion, from, size, searchFlags, facets);
     if (facets == null || facets.contains("entity") || facets.contains("_entityType")) {
       Optional<AggregationMetadata> entityTypeAgg = result.getMetadata().getAggregations().stream().filter(
@@ -145,6 +138,18 @@ public SearchResult searchAcrossEntities(@Nonnull List<String> entities, @Nonnul
     return result;
   }
 
+  private List<String> getEntitiesToSearch(@Nonnull List<String> inputEntities) {
+    List<String> nonEmptyEntities;
+    List<String> lowercaseEntities = inputEntities.stream().map(String::toLowerCase).collect(Collectors.toList());
+    try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getNonEmptyEntities").time()) {
+      nonEmptyEntities = _entityDocCountCache.getNonEmptyEntities();
+    }
+    if (!inputEntities.isEmpty()) {
+      nonEmptyEntities = nonEmptyEntities.stream().filter(lowercaseEntities::contains).collect(Collectors.toList());
+    }
+    return nonEmptyEntities;
+  }
+
   /**
    * Gets a list of documents that match given search request across multiple entities. The results are aggregated and filters are applied to the
    * search hits and not the aggregation results.
@@ -165,6 +170,6 @@ public ScrollResult scrollAcrossEntities(@Nonnull List<String> entities, @Nonnul
     log.debug(String.format(
         "Searching Search documents entities: %s, input: %s, postFilters: %s, sortCriterion: %s, from: %s, size: %s",
         entities, input, postFilters, sortCriterion, scrollId, size));
-    return _cachingEntitySearchService.scroll(entities, input, postFilters, sortCriterion, scrollId, keepAlive, size, searchFlags);
+    return _cachingEntitySearchService.scroll(getEntitiesToSearch(entities), input, postFilters, sortCriterion, scrollId, keepAlive, size, searchFlags);
   }
 }