From e380f7d317bad20ccba7464fe9366d400f32e352 Mon Sep 17 00:00:00 2001 From: John Joyce Date: Tue, 1 Aug 2023 15:31:24 -0700 Subject: [PATCH] feat(search): Allow aggregating on facets that are not explicitly part of default filter set (#8540) --- .../request/AggregationQueryBuilder.java | 38 +++++++++--- .../request/AggregationQueryBuilderTest.java | 61 ++++++++++++++++++- 2 files changed, 88 insertions(+), 11 deletions(-) diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilder.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilder.java index 2e080747a899d..d02aa960b4868 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilder.java +++ b/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 _facetFields; + private final Set _defaultFacetFields; + private final Set _allFacetFields; public AggregationQueryBuilder( @Nonnull final SearchConfiguration configs, @Nonnull final List annotations) { this._configs = Objects.requireNonNull(configs, "configs must not be null"); - this._facetFields = getFacetFields(annotations); + this._defaultFacetFields = getDefaultFacetFields(annotations); + this._allFacetFields = getAllFacetFields(annotations); } /** @@ -44,19 +46,28 @@ public List getAggregations(@Nullable List facets) { final Set 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 getFacetFields(final List annotations) { + private Set getDefaultFacetFields(final List annotations) { Set 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 getAllFacetFields(final List annotations) { + Set facets = annotations.stream() + .flatMap(annotation -> getAllFacetFieldsFromAnnotation(annotation).stream()) .collect(Collectors.toSet()); facets.add(INDEX_VIRTUAL_FIELD); return facets; @@ -64,7 +75,7 @@ private Set getFacetFields(final List annotations) private boolean isValidAggregate(final String inputFacet) { Set 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 getFacetFieldsFromAnnotation(final SearchableAnnotation annotation) { + List getDefaultFacetFieldsFromAnnotation(final SearchableAnnotation annotation) { final List facetsFromAnnotation = new ArrayList<>(); if (annotation.isAddToFilters()) { facetsFromAnnotation.add(annotation.getFieldName()); @@ -107,4 +118,13 @@ List getFacetFieldsFromAnnotation(final SearchableAnnotation annotation) } return facetsFromAnnotation; } -} + + List getAllFacetFieldsFromAnnotation(final SearchableAnnotation annotation) { + final List facetsFromAnnotation = new ArrayList<>(); + facetsFromAnnotation.add(annotation.getFieldName()); + if (annotation.getHasValuesFieldName().isPresent()) { + facetsFromAnnotation.add(annotation.getHasValuesFieldName().get()); + } + return facetsFromAnnotation; + } +} \ No newline at end of file diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilderTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilderTest.java index 76160eb29af4e..10b4ee42b1a71 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilderTest.java +++ b/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 aggs = builder.getAggregations( + ImmutableList.of("test1", "test2", "hasTest1") + ); + Assert.assertEquals(aggs.size(), 3); + Set 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); + } }