feat(search): support filtering on count type searchable fields for equality

entity-registry/src/main/java/com/linkedin/metadata/models/EntitySpec.java

@@ -3,8 +3,11 @@
 import com.linkedin.data.schema.RecordDataSchema;
 import com.linkedin.data.schema.TyperefDataSchema;
 import com.linkedin.metadata.models.annotation.EntityAnnotation;
+import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 /** A specification of a DataHub Entity */
@@ -36,6 +39,18 @@ default List<SearchableFieldSpec> getSearchableFieldSpecs() {
         .collect(Collectors.toList());
   }
 
+  default Map<String, Set<SearchableFieldSpec>> getSearchableFieldSpecMap() {
+    return getSearchableFieldSpecs().stream()
+        .collect(
+            Collectors.toMap(
+                searchableFieldSpec -> searchableFieldSpec.getSearchableAnnotation().getFieldName(),
+                searchableFieldSpec -> new HashSet<>(Collections.singleton(searchableFieldSpec)),
+                (set1, set2) -> {
+                  set1.addAll(set2);
+                  return set1;
+                }));
+  }
+
   default List<SearchScoreFieldSpec> getSearchScoreFieldSpecs() {
     return getAspectSpecs().stream()
         .map(AspectSpec::getSearchScoreFieldSpecs)

entity-registry/src/main/java/com/linkedin/metadata/models/EntitySpecBuilder.java

@@ -248,9 +248,9 @@ public AspectSpec buildAspectSpec(
       // Extract SearchScore Field Specs
       final SearchScoreFieldSpecExtractor searchScoreFieldSpecExtractor =
           new SearchScoreFieldSpecExtractor();
-      final DataSchemaRichContextTraverser searcScoreFieldSpecTraverser =
+      final DataSchemaRichContextTraverser searchScoreFieldSpecTraverser =
           new DataSchemaRichContextTraverser(searchScoreFieldSpecExtractor);
-      searcScoreFieldSpecTraverser.traverse(processedSearchScoreResult.getResultSchema());
+      searchScoreFieldSpecTraverser.traverse(processedSearchScoreResult.getResultSchema());
 
       final SchemaAnnotationProcessor.SchemaAnnotationProcessResult processedRelationshipResult =
           SchemaAnnotationProcessor.process(

entity-registry/src/main/java/com/linkedin/metadata/models/registry/ConfigEntityRegistry.java

@@ -91,7 +91,7 @@ private static Pair<Path, Path> getFileAndClassPath(String entityRegistryRoot)
               .filter(Files::isRegularFile)
               .filter(f -> f.endsWith("entity-registry.yml") || f.endsWith("entity-registry.yaml"))
               .collect(Collectors.toList());
-      if (yamlFiles.size() == 0) {
+      if (yamlFiles.isEmpty()) {
         throw new EntityRegistryException(
             String.format(
                 "Did not find an entity registry (entity_registry.yaml/yml) under %s",

entity-registry/src/main/java/com/linkedin/metadata/models/registry/MergedEntityRegistry.java

@@ -58,7 +58,7 @@ private void validateEntitySpec(EntitySpec entitySpec, final ValidationResult va
       validationResult.setValid(false);
       validationResult
           .getValidationFailures()
-          .add(String.format("Key aspect is missing in entity {}", entitySpec.getName()));
+          .add(String.format("Key aspect is missing in entity %s", entitySpec.getName()));
     }
   }
 
@@ -86,7 +86,7 @@ public MergedEntityRegistry apply(EntityRegistry patchEntityRegistry)
     }
 
     // Merge Event Specs
-    if (patchEntityRegistry.getEventSpecs().size() > 0) {
+    if (!patchEntityRegistry.getEventSpecs().isEmpty()) {
       eventNameToSpec.putAll(patchEntityRegistry.getEventSpecs());
     }
     // TODO: Validate that the entity registries don't have conflicts among each other
@@ -116,19 +116,18 @@ private void checkMergeable(
     if (existingEntitySpec != null) {
       existingEntitySpec
           .getAspectSpecMap()
-          .entrySet()
           .forEach(
-              aspectSpecEntry -> {
-                if (newEntitySpec.hasAspect(aspectSpecEntry.getKey())) {
+              (key, value) -> {
+                if (newEntitySpec.hasAspect(key)) {
                   CompatibilityResult result =
                       CompatibilityChecker.checkCompatibility(
-                          aspectSpecEntry.getValue().getPegasusSchema(),
-                          newEntitySpec.getAspectSpec(aspectSpecEntry.getKey()).getPegasusSchema(),
+                          value.getPegasusSchema(),
+                          newEntitySpec.getAspectSpec(key).getPegasusSchema(),
                           new CompatibilityOptions());
                   if (result.isError()) {
                     log.error(
                         "{} schema is not compatible with previous schema due to {}",
-                        aspectSpecEntry.getKey(),
+                        key,
                         result.getMessages());
                     // we want to continue processing all aspects to collect all failures
                     validationResult.setValid(false);
@@ -137,11 +136,11 @@ private void checkMergeable(
                         .add(
                             String.format(
                                 "%s schema is not compatible with previous schema due to %s",
-                                aspectSpecEntry.getKey(), result.getMessages()));
+                                key, result.getMessages()));
                   } else {
                     log.info(
                         "{} schema is compatible with previous schema due to {}",
-                        aspectSpecEntry.getKey(),
+                        key,
                         result.getMessages());
                   }
                 }
@@ -222,7 +221,7 @@ public PluginFactory getPluginFactory() {
 
   @Setter
   @Getter
-  private class ValidationResult {
+  private static class ValidationResult {
     boolean valid = true;
     List<String> validationFailures = new ArrayList<>();
   }

entity-registry/src/main/java/com/linkedin/metadata/models/registry/PatchEntityRegistry.java

@@ -71,19 +71,17 @@ public class PatchEntityRegistry implements EntityRegistry {
   @Override
   public String toString() {
     StringBuilder sb = new StringBuilder("PatchEntityRegistry[" + "identifier=" + identifier + ';');
-    entityNameToSpec.entrySet().stream()
-        .forEach(
-            entry ->
-                sb.append("[entityName=")
-                    .append(entry.getKey())
-                    .append(";aspects=[")
-                    .append(
-                        entry.getValue().getAspectSpecs().stream()
-                            .map(spec -> spec.getName())
-                            .collect(Collectors.joining(",")))
-                    .append("]]"));
-    eventNameToSpec.entrySet().stream()
-        .forEach(entry -> sb.append("[eventName=").append(entry.getKey()).append("]"));
+    entityNameToSpec.forEach(
+        (key1, value1) ->
+            sb.append("[entityName=")
+                .append(key1)
+                .append(";aspects=[")
+                .append(
+                    value1.getAspectSpecs().stream()
+                        .map(AspectSpec::getName)
+                        .collect(Collectors.joining(",")))
+                .append("]]"));
+    eventNameToSpec.forEach((key, value) -> sb.append("[eventName=").append(key).append("]"));
     return sb.toString();
   }
 
@@ -119,7 +117,7 @@ private static Pair<Path, Path> getFileAndClassPath(String entityRegistryRoot)
               .filter(Files::isRegularFile)
               .filter(f -> f.endsWith("entity-registry.yml") || f.endsWith("entity-registry.yaml"))
               .collect(Collectors.toList());
-      if (yamlFiles.size() == 0) {
+      if (yamlFiles.isEmpty()) {
         throw new EntityRegistryException(
             String.format(
                 "Did not find an entity registry (entity-registry.yaml/yml) under %s",
@@ -175,7 +173,7 @@ private PatchEntityRegistry(
       entities = OBJECT_MAPPER.readValue(configFileStream, Entities.class);
       this.pluginFactory = PluginFactory.withCustomClasspath(entities.getPlugins(), classLoaders);
     } catch (IOException e) {
-      e.printStackTrace();
+      log.error("Unable to read Patch configuration.", e);
       throw new IllegalArgumentException(
           String.format(
               "Error while reading config file in path %s: %s", configFileStream, e.getMessage()));

entity-registry/src/main/java/com/linkedin/metadata/models/registry/SnapshotEntityRegistry.java

@@ -120,7 +120,7 @@ public AspectTemplateEngine getAspectTemplateEngine() {
   }
 
   @Override
-  public EventSpec getEventSpec(final String ignored) {
+  public EventSpec getEventSpec(@Nonnull final String ignored) {
     return null;
   }
 

entity-registry/src/test/java/com/linkedin/metadata/models/EntitySpecBuilderTest.java

@@ -189,7 +189,7 @@ private void validateTestEntityInfo(final AspectSpec testEntityInfo) {
         testEntityInfo.getPegasusSchema().getFullName());
 
     // Assert on Searchable Fields
-    assertEquals(testEntityInfo.getSearchableFieldSpecs().size(), 11);
+    assertEquals(testEntityInfo.getSearchableFieldSpecs().size(), 12);
     assertEquals(
         "customProperties",
         testEntityInfo
@@ -340,6 +340,20 @@ private void validateTestEntityInfo(final AspectSpec testEntityInfo) {
             .get(new PathSpec("doubleField").toString())
             .getSearchableAnnotation()
             .getFieldType());
+    assertEquals(
+        "removed",
+        testEntityInfo
+            .getSearchableFieldSpecMap()
+            .get(new PathSpec("removed").toString())
+            .getSearchableAnnotation()
+            .getFieldName());
+    assertEquals(
+        SearchableAnnotation.FieldType.BOOLEAN,
+        testEntityInfo
+            .getSearchableFieldSpecMap()
+            .get(new PathSpec("removed").toString())
+            .getSearchableAnnotation()
+            .getFieldType());
 
     // Assert on Relationship Fields
     assertEquals(4, testEntityInfo.getRelationshipFieldSpecs().size());

metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESBrowseDAO.java

@@ -19,6 +19,7 @@
 import com.linkedin.metadata.config.search.SearchConfiguration;
 import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration;
 import com.linkedin.metadata.models.EntitySpec;
+import com.linkedin.metadata.models.SearchableFieldSpec;
 import com.linkedin.metadata.models.registry.EntityRegistry;
 import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.search.elasticsearch.query.request.SearchRequestHandler;
@@ -33,6 +34,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -554,7 +556,8 @@ private QueryBuilder buildQueryStringV2(
 
     queryBuilder.filter(QueryBuilders.rangeQuery(BROWSE_PATH_V2_DEPTH).gt(browseDepthVal));
 
-    queryBuilder.filter(SearchRequestHandler.getFilterQuery(filter));
+    queryBuilder.filter(
+        SearchRequestHandler.getFilterQuery(filter, entitySpec.getSearchableFieldSpecMap()));
 
     return queryBuilder;
   }
@@ -580,7 +583,18 @@ private QueryBuilder buildQueryStringBrowseAcrossEntities(
 
     queryBuilder.filter(QueryBuilders.rangeQuery(BROWSE_PATH_V2_DEPTH).gt(browseDepthVal));
 
-    queryBuilder.filter(SearchRequestHandler.getFilterQuery(filter));
+    Map<String, Set<SearchableFieldSpec>> searchableFields =
+        entitySpecs.stream()
+            .flatMap(entitySpec -> entitySpec.getSearchableFieldSpecMap().entrySet().stream())
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    Map.Entry::getValue,
+                    (set1, set2) -> {
+                      set1.addAll(set2);
+                      return set1;
+                    }));
+    queryBuilder.filter(SearchRequestHandler.getFilterQuery(filter, searchableFields));
 
     return queryBuilder;
   }

metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/ESSearchDAO.java

@@ -78,7 +78,8 @@ public long docCount(@Nonnull String entityName) {
     EntitySpec entitySpec = entityRegistry.getEntitySpec(entityName);
     CountRequest countRequest =
         new CountRequest(indexConvention.getIndexName(entitySpec))
-            .query(SearchRequestHandler.getFilterQuery(null));
+            .query(
+                SearchRequestHandler.getFilterQuery(null, entitySpec.getSearchableFieldSpecMap()));
     try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "docCount").time()) {
       return client.count(countRequest, RequestOptions.DEFAULT).getCount();
     } catch (IOException e) {
@@ -315,9 +316,17 @@ public Map<String, Long> aggregateByValue(
       @Nonnull String field,
       @Nullable Filter requestParams,
       int limit) {
+    List<EntitySpec> entitySpecs;
+    if (entityNames == null || entityNames.isEmpty()) {
+      entitySpecs = new ArrayList<>(entityRegistry.getEntitySpecs().values());
+    } else {
+      entitySpecs =
+          entityNames.stream().map(entityRegistry::getEntitySpec).collect(Collectors.toList());
+    }
     final SearchRequest searchRequest =
-        SearchRequestHandler.getAggregationRequest(
-            field, transformFilterForEntities(requestParams, indexConvention), limit);
+        SearchRequestHandler.getBuilder(entitySpecs, searchConfiguration, customSearchConfiguration)
+            .getAggregationRequest(
+                field, transformFilterForEntities(requestParams, indexConvention), limit);
     if (entityNames == null) {
       String indexName = indexConvention.getAllEntityIndicesPattern();
       searchRequest.indices(indexName);

metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AutocompleteRequestHandler.java

@@ -14,6 +14,7 @@
 import com.linkedin.metadata.query.filter.Filter;
 import com.linkedin.metadata.search.utils.ESUtils;
 import java.net.URISyntaxException;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
@@ -40,19 +41,33 @@
 public class AutocompleteRequestHandler {
 
   private final List<String> _defaultAutocompleteFields;
+  private final Map<String, Set<SearchableFieldSpec>> searchableFields;
 
   private static final Map<EntitySpec, AutocompleteRequestHandler>
       AUTOCOMPLETE_QUERY_BUILDER_BY_ENTITY_NAME = new ConcurrentHashMap<>();
 
   public AutocompleteRequestHandler(@Nonnull EntitySpec entitySpec) {
+    List<SearchableFieldSpec> fieldSpecs = entitySpec.getSearchableFieldSpecs();
     _defaultAutocompleteFields =
         Stream.concat(
-                entitySpec.getSearchableFieldSpecs().stream()
+                fieldSpecs.stream()
                     .map(SearchableFieldSpec::getSearchableAnnotation)
                     .filter(SearchableAnnotation::isEnableAutocomplete)
                     .map(SearchableAnnotation::getFieldName),
                 Stream.of("urn"))
             .collect(Collectors.toList());
+    searchableFields =
+        fieldSpecs.stream()
+            .collect(
+                Collectors.toMap(
+                    searchableFieldSpec ->
+                        searchableFieldSpec.getSearchableAnnotation().getFieldName(),
+                    searchableFieldSpec ->
+                        new HashSet<>(Collections.singleton(searchableFieldSpec)),
+                    (set1, set2) -> {
+                      set1.addAll(set2);
+                      return set1;
+                    }));
   }
 
   public static AutocompleteRequestHandler getBuilder(@Nonnull EntitySpec entitySpec) {
@@ -66,7 +81,7 @@ public SearchRequest getSearchRequest(
     SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
     searchSourceBuilder.size(limit);
     searchSourceBuilder.query(getQuery(input, field));
-    searchSourceBuilder.postFilter(ESUtils.buildFilterQuery(filter, false));
+    searchSourceBuilder.postFilter(ESUtils.buildFilterQuery(filter, false, searchableFields));
     searchSourceBuilder.highlighter(getHighlights(field));
     searchRequest.source(searchSourceBuilder);
     return searchRequest;

metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java

@@ -97,6 +97,7 @@ public class SearchRequestHandler {
   private final SearchConfiguration _configs;
   private final SearchQueryBuilder _searchQueryBuilder;
   private final AggregationQueryBuilder _aggregationQueryBuilder;
+  private final Map<String, Set<SearchableFieldSpec>> searchableFields;
 
   private SearchRequestHandler(
       @Nonnull EntitySpec entitySpec,
@@ -121,6 +122,17 @@ private SearchRequestHandler(
     _searchQueryBuilder = new SearchQueryBuilder(configs, customSearchConfiguration);
     _aggregationQueryBuilder = new AggregationQueryBuilder(configs, annotations);
     _configs = configs;
+    searchableFields =
+        _entitySpecs.stream()
+            .flatMap(entitySpec -> entitySpec.getSearchableFieldSpecMap().entrySet().stream())
+            .collect(
+                Collectors.toMap(
+                    Map.Entry::getKey,
+                    Map.Entry::getValue,
+                    (set1, set2) -> {
+                      set1.addAll(set2);
+                      return set1;
+                    }));
   }
 
   public static SearchRequestHandler getBuilder(
@@ -169,8 +181,13 @@ private BinaryOperator<String> mapMerger() {
     };
   }
 
-  public static BoolQueryBuilder getFilterQuery(@Nullable Filter filter) {
-    BoolQueryBuilder filterQuery = ESUtils.buildFilterQuery(filter, false);
+  public BoolQueryBuilder getFilterQuery(@Nullable Filter filter) {
+    return getFilterQuery(filter, searchableFields);
+  }
+
+  public static BoolQueryBuilder getFilterQuery(
+      @Nullable Filter filter, Map<String, Set<SearchableFieldSpec>> searchableFields) {
+    BoolQueryBuilder filterQuery = ESUtils.buildFilterQuery(filter, false, searchableFields);
 
     return filterSoftDeletedByDefault(filter, filterQuery);
   }
@@ -354,7 +371,7 @@ public SearchRequest getFilterRequest(
    * @return {@link SearchRequest} that contains the aggregation query
    */
   @Nonnull
-  public static SearchRequest getAggregationRequest(
+  public SearchRequest getAggregationRequest(
       @Nonnull String field, @Nullable Filter filter, int limit) {
     SearchRequest searchRequest = new SearchRequest();
     BoolQueryBuilder filterQuery = getFilterQuery(filter);