fix(adv search): two advanced search fixes (#6252)

* two advanced search fixes

* another approach

* adding unit tests

* fixing checkstyle issue

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ResolverUtils.java

@@ -135,15 +135,25 @@ public static Filter buildFilter(@Nullable List<FacetFilterInput> andFilters, @N
     public static Criterion criterionFromFilter(final FacetFilterInput filter) {
         Criterion result = new Criterion();
         result.setField(getFilterField(filter.getField()));
-        if (filter.getValues() != null) {
+
+        // `value` is deprecated in place of `values`- this is to support old query patterns. If values is provided,
+        // this statement will be skipped
+        if (filter.getValues() == null && filter.getValue() != null) {
+            result.setValues(new StringArray(filter.getValue()));
+            result.setValue(filter.getValue());
+        } else if (filter.getValues() != null) {
             result.setValues(new StringArray(filter.getValues()));
             if (!filter.getValues().isEmpty()) {
                 result.setValue(filter.getValues().get(0));
             } else {
                 result.setValue("");
             }
+        } else {
+            result.setValues(new StringArray());
+            result.setValue("");
         }
 
+
         if (filter.getCondition() != null) {
             result.setCondition(Condition.valueOf(filter.getCondition().toString()));
         } else {

datahub-graphql-core/src/main/resources/search.graphql

@@ -177,10 +177,16 @@ input FacetFilterInput {
   """
   field: String!
 
+  """
+  Value of the field to filter by. Deprecated in favor of `values`, which should accept a single element array for a
+  value
+  """
+  value: String @deprecated(reason: "Prefer `values` for single elements")
+
   """
   Values, one of which the intended field should match.
   """
-  values: [String!]!
+  values: [String!]
 
   """
   If the filter should or should not be matched

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ResolverUtilsTest.java

@@ -0,0 +1,69 @@
+package com.linkedin.datahub.graphql.resolvers;
+
+import com.google.common.collect.ImmutableList;
+import com.linkedin.data.template.StringArray;
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.TestUtils;
+import com.linkedin.datahub.graphql.generated.FacetFilterInput;
+import com.linkedin.datahub.graphql.generated.FilterOperator;
+import com.linkedin.metadata.query.filter.Condition;
+import com.linkedin.metadata.query.filter.Criterion;
+import graphql.schema.DataFetchingEnvironment;
+import junit.framework.TestCase;
+import org.testng.annotations.Test;
+import org.mockito.Mockito;
+
+import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.*;
+
+
+public class ResolverUtilsTest extends TestCase {
+
+  @Test
+  public void testCriterionFromFilter() throws Exception {
+    final DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class);
+    final QueryContext mockAllowContext = TestUtils.getMockAllowContext();
+    Mockito.when(mockEnv.getContext()).thenReturn(mockAllowContext);
+
+    // this is the expected path
+    Criterion valuesCriterion = criterionFromFilter(
+        new FacetFilterInput(
+            "tags",
+            null,
+            ImmutableList.of("urn:li:tag:abc", "urn:li:tag:def"),
+            false,
+            FilterOperator.EQUAL
+        )
+    );
+    assertEquals(valuesCriterion, new Criterion().setValue("urn:li:tag:abc").setValues(
+        new StringArray(ImmutableList.of("urn:li:tag:abc", "urn:li:tag:def"))
+    ).setNegated(false).setCondition(Condition.EQUAL).setField("tags.keyword"));
+
+    // this is the legacy pathway
+    Criterion valueCriterion = criterionFromFilter(
+        new FacetFilterInput(
+            "tags",
+            "urn:li:tag:abc",
+            null,
+            true,
+            FilterOperator.EQUAL
+        )
+    );
+    assertEquals(valueCriterion, new Criterion().setValue("urn:li:tag:abc").setValues(
+        new StringArray(ImmutableList.of("urn:li:tag:abc"))
+    ).setNegated(true).setCondition(Condition.EQUAL).setField("tags.keyword"));
+
+    // check that both being null doesn't cause a NPE. this should never happen except via API interaction
+    Criterion doubleNullCriterion = criterionFromFilter(
+        new FacetFilterInput(
+            "tags",
+            null,
+            null,
+            true,
+            FilterOperator.EQUAL
+        )
+    );
+    assertEquals(doubleNullCriterion, new Criterion().setValue("").setValues(
+        new StringArray(ImmutableList.of())
+    ).setNegated(true).setCondition(Condition.EQUAL).setField("tags.keyword"));
+  }
+}

datahub-web-react/src/app/entity/shared/components/styled/search/EmbeddedListSearch.tsx

@@ -100,7 +100,7 @@ export const EmbeddedListSearch = ({
     const finalFilters = (fixedFilter && [...filtersWithoutEntities, fixedFilter]) || filtersWithoutEntities;
     const entityFilters: Array<EntityType> = filters
         .filter((filter) => filter.field === ENTITY_FILTER_NAME)
-        .flatMap((filter) => filter.values.map((value) => value?.toUpperCase() as EntityType));
+        .flatMap((filter) => filter.values?.map((value) => value?.toUpperCase() as EntityType) || []);
 
     const [showFilters, setShowFilters] = useState(defaultShowFilters || false);
     const [isSelectMode, setIsSelectMode] = useState(false);

datahub-web-react/src/app/entity/shared/components/styled/search/SearchSelect.tsx

@@ -71,7 +71,7 @@ export const SearchSelect = ({ fixedEntityTypes, placeholderText, selectedEntiti
     );
     const entityFilters: Array<EntityType> = filters
         .filter((filter) => filter.field === ENTITY_FILTER_NAME)
-        .flatMap((filter) => filter.values.map((value) => value.toUpperCase() as EntityType));
+        .flatMap((filter) => filter.values?.map((value) => value.toUpperCase() as EntityType) || []);
     const finalEntityTypes = (entityFilters.length > 0 && entityFilters) || fixedEntityTypes || [];
 
     // Execute search

datahub-web-react/src/app/entity/shared/tabs/Lineage/ImpactAnalysis.tsx

@@ -33,7 +33,7 @@ export const ImpactAnalysis = ({ urn, direction }: Props) => {
     );
     const entityFilters: Array<EntityType> = filters
         .filter((filter) => filter.field === ENTITY_FILTER_NAME)
-        .flatMap((filter) => filter.values.map((value) => value.toUpperCase() as EntityType));
+        .flatMap((filter) => filter.values?.map((value) => value.toUpperCase() as EntityType) || []);
 
     const { data, loading } = useSearchAcrossLineageQuery({
         variables: {

datahub-web-react/src/app/search/SearchPage.tsx

@@ -39,7 +39,7 @@ export const SearchPage = () => {
     );
     const entityFilters: Array<EntityType> = filters
         .filter((filter) => filter.field === ENTITY_FILTER_NAME)
-        .flatMap((filter) => filter.values.map((value) => value?.toUpperCase() as EntityType));
+        .flatMap((filter) => filter.values?.map((value) => value?.toUpperCase() as EntityType) || []);
 
     const [numResultsPerPage, setNumResultsPerPage] = useState(SearchCfg.RESULTS_PER_PAGE);
     const [isSelectMode, setIsSelectMode] = useState(false);

datahub-web-react/src/app/search/SimpleSearchFilter.tsx

@@ -59,7 +59,7 @@ export const SimpleSearchFilter = ({ facet, selectedFilters, onFilterSelect, def
     const [expanded, setExpanded] = useState(false);
 
     const isFacetSelected = (field, value) => {
-        return selectedFilters.find((f) => f.field === field && f.values.includes(value)) !== undefined;
+        return selectedFilters.find((f) => f.field === field && f.values?.includes(value)) !== undefined;
     };
 
     // Aggregations filtered for count > 0 or selected = true

datahub-web-react/src/app/search/SimpleSearchFilters.tsx

@@ -51,10 +51,10 @@ export const SimpleSearchFilters = ({ facets, selectedFilters, onFilterSelect, l
             : selectedFilters
                   .map((filter) =>
                       filter.field === field
-                          ? { ...filter, values: filter.values.filter((val) => val !== value) }
+                          ? { ...filter, values: filter.values?.filter((val) => val !== value) }
                           : filter,
                   )
-                  .filter((filter) => filter.field !== field || !(filter.values.length === 0));
+                  .filter((filter) => filter.field !== field || !(filter.values?.length === 0));
         setCachedProps({ ...cachedProps, selectedFilters: newFilters });
         onFilterSelect(newFilters);
     };

datahub-web-react/src/app/search/utils/filtersToQueryStringParams.ts

@@ -12,7 +12,7 @@ function reduceFiltersToCombineDegreeFilters(acc: FacetFilterInput[], filter: Fa
     if (filter.field === DEGREE_FILTER && acc.filter((f) => f.field === DEGREE_FILTER).length > 0) {
         // instead of appending this new degree filter, combine it with the previous one and continue
         return acc.map((f) =>
-            f.field === DEGREE_FILTER ? { ...f, values: [...f.values, ...filter.values] } : f,
+            f.field === DEGREE_FILTER ? { ...f, values: [...(f.values || []), ...(filter.values || [])] } : f,
         ) as FacetFilterInput[];
     }
     return [...acc, filter] as FacetFilterInput[];

metadata-io/src/main/java/com/linkedin/metadata/search/aggregator/AllEntitiesSearchAggregator.java

@@ -25,6 +25,8 @@
 import com.linkedin.util.Pair;
 import io.opentelemetry.extension.annotations.WithSpan;
 import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -125,7 +127,7 @@ public SearchResult search(@Nonnull List<String> entities, @Nonnull String input
     finalAggregations.put("entity", new AggregationMetadata().setName("entity")
         .setDisplayName("Type")
         .setAggregations(new LongMap(numResultsPerEntity))
-        .setFilterValues(new FilterValueArray(SearchUtil.convertToFilters(numResultsPerEntity))));
+        .setFilterValues(new FilterValueArray(SearchUtil.convertToFilters(numResultsPerEntity, Collections.emptySet()))));
 
     // 4. Rank results across entities
     List<SearchEntity> rankedResult = _searchRanker.rank(matchedResults);
@@ -184,6 +186,8 @@ private Map<String, AggregationMetadata> trimMergedAggregations(Map<String, Aggr
    */
   private FilterValueArray trimFilterValues(FilterValueArray original) {
     if (original.size() > _maxAggregationValueCount) {
+      // sort so that values that appear in the filter appear first
+      original.sort(Comparator.comparingInt(val -> val.hasFiltered() && val.isFiltered() ? 0 : 1));
       return new FilterValueArray(
           original.subList(0, _maxAggregationValueCount)
       );

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

@@ -72,7 +72,6 @@ public class SearchRequestHandler {
 
   private static final String URN_FILTER = "urn";
   private static final int DEFAULT_MAX_TERM_BUCKET_SIZE = 20;
-
   private final EntitySpec _entitySpec;
   private final Set<String> _facetFields;
   private final Set<String> _defaultQueryFieldNames;
@@ -358,7 +357,7 @@ private List<AggregationMetadata> extractAggregationMetadata(@Nonnull SearchResp
       final AggregationMetadata aggregationMetadata = new AggregationMetadata().setName(entry.getKey())
           .setDisplayName(_filtersToDisplayName.get(entry.getKey()))
           .setAggregations(new LongMap(oneTermAggResult))
-          .setFilterValues(new FilterValueArray(SearchUtil.convertToFilters(oneTermAggResult)));
+          .setFilterValues(new FilterValueArray(SearchUtil.convertToFilters(oneTermAggResult, Collections.emptySet())));
       aggregationMetadataList.add(aggregationMetadata);
     }
 
@@ -476,7 +475,7 @@ private void addCriterionFiltersToAggregationMetadata(
           finalFacetField,
           _filtersToDisplayName.getOrDefault(finalFacetField, finalFacetField),
           new LongMap(criterion.getValues().stream().collect(Collectors.toMap(i -> i, i -> 0L))),
-          new FilterValueArray(criterion.getValues().stream().map(value -> createFilterValue(value, 0L)).collect(
+          new FilterValueArray(criterion.getValues().stream().map(value -> createFilterValue(value, 0L, true)).collect(
               Collectors.toList())))
       );
     }
@@ -489,7 +488,7 @@ private void addMissingAggregationValueToAggregationMetadata(@Nonnull final Stri
     ) {
       // No aggregation found for filtered value -- inject one!
       originalMetadata.getAggregations().put(value, 0L);
-      originalMetadata.getFilterValues().add(createFilterValue(value, 0L));
+      originalMetadata.getFilterValues().add(createFilterValue(value, 0L, true));
     }
   }
 

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

@@ -23,6 +23,7 @@
 import java.util.Collections;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -136,10 +137,16 @@ public static AggregationMetadata merge(AggregationMetadata one, AggregationMeta
     Map<String, Long> mergedMap =
         Stream.concat(one.getAggregations().entrySet().stream(), two.getAggregations().entrySet().stream())
             .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue, Long::sum));
+
+    // we want to make sure the values that were used in the filter are prioritized to appear in the response aggregation
+    Set<String> filteredValues = Stream.concat(one.getFilterValues().stream(), two.getFilterValues().stream()).filter(val -> val.isFiltered()).map(
+        val -> val.getValue()
+    ).collect(Collectors.toSet());
+
     return one.clone()
         .setDisplayName(two.getDisplayName() != two.getName() ? two.getDisplayName() : one.getDisplayName())
         .setAggregations(new LongMap(mergedMap))
-        .setFilterValues(new FilterValueArray(SearchUtil.convertToFilters(mergedMap)));
+        .setFilterValues(new FilterValueArray(SearchUtil.convertToFilters(mergedMap, filteredValues)));
   }
 
   public static ListResult toListResult(final SearchResult searchResult) {

metadata-models/src/main/pegasus/com/linkedin/metadata/search/AggregationMetadata.pdl

@@ -22,5 +22,7 @@ record AggregationMetadata {
       value: string
       entity: optional Urn
       facetCount: long
+      // indicates that the FilterValue is part of the search request
+      filtered: optional boolean
     }]
-}
+}

metadata-service/restli-api/src/main/snapshot/com.linkedin.entity.entities.snapshot.json

@@ -5460,6 +5460,10 @@
           }, {
             "name" : "facetCount",
             "type" : "long"
+          }, {
+            "name" : "filtered",
+            "type" : "boolean",
+            "optional" : true
           } ]
         }
       }

metadata-utils/src/main/java/com/linkedin/metadata/utils/SearchUtil.java

@@ -6,6 +6,7 @@
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 import lombok.extern.slf4j.Slf4j;
 
@@ -17,14 +18,18 @@ public class SearchUtil {
   private SearchUtil() {
   }
 
-  public static List<FilterValue> convertToFilters(Map<String, Long> aggregations) {
+  /*
+   * @param aggregations the aggregations coming back from elasticsearch combined with the filters from the search request
+   * @param filteredValues the set of values provided by the search request
+   */
+  public static List<FilterValue> convertToFilters(Map<String, Long> aggregations, Set<String> filteredValues) {
     return aggregations.entrySet().stream().map(entry -> {
-      return createFilterValue(entry.getKey(), entry.getValue());
+      return createFilterValue(entry.getKey(), entry.getValue(), filteredValues.contains(entry.getKey()));
     }).sorted(Comparator.comparingLong(value -> -value.getFacetCount())).collect(Collectors.toList());
   }
 
-  public static FilterValue createFilterValue(String value, Long facetCount) {
-    FilterValue result = new FilterValue().setValue(value).setFacetCount(facetCount);
+  public static FilterValue createFilterValue(String value, Long facetCount, Boolean isFilteredOn) {
+    FilterValue result = new FilterValue().setValue(value).setFacetCount(facetCount).setFiltered(isFilteredOn);
     if (value.startsWith(URN_PREFIX)) {
       try {
         result.setEntity(Urn.createFromString(value));

metadata-utils/src/test/java/com/linkedin/metadata/utils/SearchUtilTest.java

@@ -0,0 +1,45 @@
+package com.linkedin.metadata.utils;
+
+import com.google.common.collect.ImmutableSet;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.metadata.search.FilterValue;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.*;
+
+
+/**
+ * Tests the capabilities of {@link EntityKeyUtils}
+ */
+public class SearchUtilTest {
+
+  @Test
+  public void testConvertToFilters() throws Exception {
+    Map<String, Long> aggregations = new HashMap<>();
+    aggregations.put("urn:li:tag:abc", 3L);
+    aggregations.put("urn:li:tag:def", 0L);
+
+    Set<String> filteredValues = ImmutableSet.of("urn:li:tag:def");
+
+    List<FilterValue> filters =
+        SearchUtil.convertToFilters(aggregations, filteredValues);
+
+    assertEquals(filters.get(0), new FilterValue()
+        .setFiltered(false)
+        .setValue("urn:li:tag:abc")
+        .setEntity(Urn.createFromString("urn:li:tag:abc"))
+        .setFacetCount(3L)
+    );
+
+    assertEquals(filters.get(1), new FilterValue()
+        .setFiltered(true)
+        .setValue("urn:li:tag:def")
+        .setEntity(Urn.createFromString("urn:li:tag:def"))
+        .setFacetCount(0L)
+    );
+  }
+}