feat(search): Allow aggregating on facets that are not explicitly part of default filter set (#8540)

Author: jjoyce0510

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

@@ -21,13 +21,15 @@
 public class AggregationQueryBuilder {
 
   private final SearchConfiguration _configs;
-  private final Set<String> _facetFields;
+  private final Set<String> _defaultFacetFields;
+  private final Set<String> _allFacetFields;
 
   public AggregationQueryBuilder(
       @Nonnull final SearchConfiguration configs,
       @Nonnull final List<SearchableAnnotation> annotations) {
     this._configs = Objects.requireNonNull(configs, "configs must not be null");
-    this._facetFields = getFacetFields(annotations);
+    this._defaultFacetFields = getDefaultFacetFields(annotations);
+    this._allFacetFields = getAllFacetFields(annotations);
   }
 
   /**
@@ -44,27 +46,36 @@ public List<AggregationBuilder> getAggregations(@Nullable List<String> facets) {
     final Set<String> facetsToAggregate;
     if (facets != null) {
       facets.stream().filter(f -> !isValidAggregate(f)).forEach(facet -> {
-        log.warn(String.format("Provided facet for search filter aggregations that doesn't exist. Provided: %s; Available: %s", facet, _facetFields));
+        log.warn(String.format("Requested facet for search filter aggregations that isn't part of the default filters. Provided: %s; Available: %s", facet,
+            _defaultFacetFields));
       });
       facetsToAggregate = facets.stream().filter(this::isValidAggregate).collect(Collectors.toSet());
     } else {
-      facetsToAggregate = _facetFields;
+      facetsToAggregate = _defaultFacetFields;
     }
     return facetsToAggregate.stream().map(this::facetToAggregationBuilder).collect(Collectors.toList());
   }
 
 
-  private Set<String> getFacetFields(final List<SearchableAnnotation> annotations) {
+  private Set<String> getDefaultFacetFields(final List<SearchableAnnotation> annotations) {
     Set<String> facets = annotations.stream()
-        .flatMap(annotation -> getFacetFieldsFromAnnotation(annotation).stream())
+        .flatMap(annotation -> getDefaultFacetFieldsFromAnnotation(annotation).stream())
+        .collect(Collectors.toSet());
+    facets.add(INDEX_VIRTUAL_FIELD);
+    return facets;
+  }
+
+  private Set<String> getAllFacetFields(final List<SearchableAnnotation> annotations) {
+    Set<String> facets = annotations.stream()
+        .flatMap(annotation -> getAllFacetFieldsFromAnnotation(annotation).stream())
         .collect(Collectors.toSet());
     facets.add(INDEX_VIRTUAL_FIELD);
     return facets;
   }
 
   private boolean isValidAggregate(final String inputFacet) {
     Set<String> facets = Set.of(inputFacet.split(AGGREGATION_SEPARATOR_CHAR));
-    return facets.size() > 0 && _facetFields.containsAll(facets);
+    return facets.size() > 0 && _allFacetFields.containsAll(facets);
   }
 
   private AggregationBuilder facetToAggregationBuilder(final String inputFacet) {
@@ -97,7 +108,7 @@ private String getAggregationField(final String facet) {
     return ESUtils.toKeywordField(facet, false);
   }
 
-  List<String> getFacetFieldsFromAnnotation(final SearchableAnnotation annotation) {
+  List<String> getDefaultFacetFieldsFromAnnotation(final SearchableAnnotation annotation) {
     final List<String> facetsFromAnnotation = new ArrayList<>();
     if (annotation.isAddToFilters()) {
       facetsFromAnnotation.add(annotation.getFieldName());
@@ -107,4 +118,13 @@ List<String> getFacetFieldsFromAnnotation(final SearchableAnnotation annotation)
     }
     return facetsFromAnnotation;
   }
-}
+
+  List<String> getAllFacetFieldsFromAnnotation(final SearchableAnnotation annotation) {
+    final List<String> facetsFromAnnotation = new ArrayList<>();
+    facetsFromAnnotation.add(annotation.getFieldName());
+    if (annotation.getHasValuesFieldName().isPresent()) {
+      facetsFromAnnotation.add(annotation.getHasValuesFieldName().get());
+    }
+    return facetsFromAnnotation;
+  }
+}

metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilderTest.java

@@ -1,11 +1,14 @@
 package com.linkedin.metadata.search.elasticsearch.query.request;
 
+import com.google.common.collect.ImmutableSet;
 import com.linkedin.metadata.config.search.SearchConfiguration;
 import com.google.common.collect.ImmutableList;
 import com.linkedin.metadata.models.annotation.SearchableAnnotation;
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
 import org.elasticsearch.search.aggregations.AggregationBuilder;
 import org.testng.Assert;
 import org.testng.annotations.Test;
@@ -14,7 +17,7 @@
 public class AggregationQueryBuilderTest {
 
   @Test
-  public void testGetAggregationsHasFields() {
+  public void testGetDefaultAggregationsHasFields() {
 
     SearchableAnnotation annotation = new SearchableAnnotation(
         "test",
@@ -43,7 +46,7 @@ public void testGetAggregationsHasFields() {
   }
 
   @Test
-  public void testGetAggregationsFields() {
+  public void testGetDefaultAggregationsFields() {
 
     SearchableAnnotation annotation = new SearchableAnnotation(
         "test",
@@ -70,4 +73,58 @@ public void testGetAggregationsFields() {
 
     Assert.assertTrue(aggs.stream().anyMatch(agg -> agg.getName().equals("test")));
   }
+
+  @Test
+  public void testGetSpecificAggregationsHasFields() {
+
+    SearchableAnnotation annotation1 = new SearchableAnnotation(
+        "test1",
+        SearchableAnnotation.FieldType.KEYWORD,
+        true,
+        true,
+        false,
+        false,
+        Optional.empty(),
+        Optional.of("Has Test"),
+        1.0,
+        Optional.of("hasTest1"),
+        Optional.empty(),
+        Collections.emptyMap()
+    );
+
+    SearchableAnnotation annotation2 = new SearchableAnnotation(
+        "test2",
+        SearchableAnnotation.FieldType.KEYWORD,
+        true,
+        true,
+        false,
+        false,
+        Optional.of("Test Filter"),
+        Optional.empty(),
+        1.0,
+        Optional.empty(),
+        Optional.empty(),
+        Collections.emptyMap()
+    );
+
+    SearchConfiguration config = new SearchConfiguration();
+    config.setMaxTermBucketSize(25);
+
+    AggregationQueryBuilder builder = new AggregationQueryBuilder(
+        config, ImmutableList.of(annotation1, annotation2));
+
+    // Case 1: Ask for fields that should exist.
+    List<AggregationBuilder> aggs = builder.getAggregations(
+        ImmutableList.of("test1", "test2", "hasTest1")
+    );
+    Assert.assertEquals(aggs.size(), 3);
+    Set<String> facets = aggs.stream().map(AggregationBuilder::getName).collect(Collectors.toSet());
+    Assert.assertEquals(ImmutableSet.of("test1", "test2", "hasTest1"), facets);
+
+    // Case 2: Ask for fields that should NOT exist.
+    aggs = builder.getAggregations(
+        ImmutableList.of("hasTest2")
+    );
+    Assert.assertEquals(aggs.size(), 0);
+  }
 }