feat(structuredproperties): aggregration fix & docs (#10780)

datahub-project · Jul 16, 2024 · 1d298b7 · 1d298b7
1 parent b08c3a9
commit 1d298b7
Show file tree

Hide file tree

Showing 3 changed files with 105 additions and 15 deletions.
diff --git a/docs/api/tutorials/structured-properties.md b/docs/api/tutorials/structured-properties.md
@@ -20,7 +20,7 @@ This guide will show you how to execute the following actions with structured pr
 - Add structured properties to a dataset
 - Patch structured properties (add / remove / update a single property)
 - Update structured property with breaking schema changes
-- Search using structured properties
+- Search & aggregations using structured properties
 
 ## Prerequisites
 
@@ -1107,7 +1107,7 @@ schema changes will remove previously written data.
 Breaking schema changes are implemented by setting a version string within the Structured Property definition. This
 version must be in the following format: `yyyyMMddhhmmss`, i.e. `20240614080000`
 
-:::IMPORTANT NOTES
+:::note IMPORTANT NOTES
 Old values will not be retrieve-able after the new Structured Property definition is applied. 
 
 The old values will be subject to deletion asynchronously (future work).
@@ -1234,7 +1234,7 @@ Example Response:
 </TabItem>
 </Tabs>
 
-## Structured Properties & Search
+## Structured Properties - Search & Aggregation
 
 Currently Structured Properties can be used to filter search results. This currently excludes fulltext search.
 
@@ -1525,5 +1525,73 @@ Example Response:
 }
 ```
 
+</TabItem>
+</Tabs>
+
+### Structured Property Aggregations
+
+Structured properties can also be used in GraphQL's aggregation queries using the same naming convention outlined above 
+for search filter field names. There are currently no aggregation endpoints for OpenAPI.
+
+<Tabs>
+<TabItem value="GraphQL" label="GraphQL" default>
+
+Aggregation Query:
+
+```graphql
+query {
+  aggregateAcrossEntities(
+    input: {
+      types: [], 
+      facets: [
+        "structuredProperties.io.acryl.privacy.retentionTime02",
+        "structuredProperties.io.acryl.privacy.retentionTime"], 
+      query: "*", 
+      orFilters: [], 
+      searchFlags: {maxAggValues: 100}
+    }) {
+  facets {
+    field
+      aggregations {
+        value
+        count
+      }
+    }
+  }
+}
+```
+
+Example Response:
+
+```json
+{
+  "data": {
+    "aggregateAcrossEntities": {
+      "facets": [
+        {
+          "field": "structuredProperties.io.acryl.privacy.retentionTime02",
+          "aggregations": [
+            {
+              "value": "bar2",
+              "count": 1
+            }
+          ]
+        },
+        {
+          "field": "structuredProperties.io.acryl.privacy.retentionTime",
+          "aggregations": [
+            {
+              "value": "60.0",
+              "count": 1
+            }
+          ]
+        }
+      ]
+    }
+  },
+  "extensions": {}
+}
+```
+
 </TabItem>
 </Tabs>
diff --git a/entity-registry/src/main/java/com/linkedin/metadata/models/StructuredPropertyUtils.java b/entity-registry/src/main/java/com/linkedin/metadata/models/StructuredPropertyUtils.java
@@ -80,6 +80,34 @@ public static String getValueTypeId(@Nullable final Urn valueType) {
     }
   }
 
+  /**
+   * Given a structured property input field or facet name, return a valid structured property facet
+   * name
+   *
+   * @param fieldOrFacetName input name
+   * @param aspectRetriever aspect retriever
+   * @return guranteed facet name
+   */
+  public static Optional<String> toStructuredPropertyFacetName(
+      @Nonnull String fieldOrFacetName, @Nullable AspectRetriever aspectRetriever) {
+    return lookupDefinitionFromFilterOrFacetName(fieldOrFacetName, aspectRetriever)
+        .map(
+            urnDefinition -> {
+              switch (getLogicalValueType(urnDefinition.getSecond())) {
+                case DATE:
+                case NUMBER:
+                  return STRUCTURED_PROPERTY_MAPPING_FIELD_PREFIX
+                      + StructuredPropertyUtils.toElasticsearchFieldName(
+                          urnDefinition.getFirst(), urnDefinition.getSecond());
+                default:
+                  return STRUCTURED_PROPERTY_MAPPING_FIELD_PREFIX
+                      + StructuredPropertyUtils.toElasticsearchFieldName(
+                          urnDefinition.getFirst(), urnDefinition.getSecond())
+                      + ".keyword";
+              }
+            });
+  }
+
   /**
    * Lookup structured property definition given the name used for the field in APIs such as a
    * search filter or aggregation query facet name.

diff --git a/...ava/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilder.java b/...ava/com/linkedin/metadata/search/elasticsearch/query/request/AggregationQueryBuilder.java
@@ -1,14 +1,14 @@
 package com.linkedin.metadata.search.elasticsearch.query.request;
 
 import static com.linkedin.metadata.Constants.*;
+import static com.linkedin.metadata.models.StructuredPropertyUtils.toStructuredPropertyFacetName;
 import static com.linkedin.metadata.search.utils.ESUtils.toParentField;
 import static com.linkedin.metadata.utils.SearchUtil.*;
 
 import com.linkedin.data.template.LongMap;
 import com.linkedin.metadata.aspect.AspectRetriever;
 import com.linkedin.metadata.config.search.SearchConfiguration;
 import com.linkedin.metadata.models.EntitySpec;
-import com.linkedin.metadata.models.StructuredPropertyUtils;
 import com.linkedin.metadata.models.annotation.SearchableAnnotation;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterion;
 import com.linkedin.metadata.query.filter.ConjunctiveCriterionArray;
@@ -136,15 +136,7 @@ private AggregationBuilder facetToAggregationBuilder(
             opContext.getSearchContext().getSearchFlags().getMaxAggValues(),
             configs.getMaxTermBucketSize());
     for (int i = facets.size() - 1; i >= 0; i--) {
-      String facet =
-          StructuredPropertyUtils.lookupDefinitionFromFilterOrFacetName(
-                  facets.get(i), opContext.getAspectRetriever())
-              .map(
-                  urnDefinition ->
-                      STRUCTURED_PROPERTY_MAPPING_FIELD_PREFIX
-                          + StructuredPropertyUtils.toElasticsearchFieldName(
-                              urnDefinition.getFirst(), urnDefinition.getSecond()))
-              .orElse(facets.get(i));
+      String facet = facets.get(i);
 
       AggregationBuilder aggBuilder;
       if (facet.contains(AGGREGATION_SPECIAL_TYPE_DELIMITER)) {
@@ -190,8 +182,10 @@ private String getAggregationField(
       // Boolean hasX field, not a keyword field. Return the name of the original facet.
       return facet;
     }
-    // Otherwise assume that this field is of keyword type.
-    return ESUtils.toKeywordField(facet, false, aspectRetriever);
+    // intercept structured property if it exists
+    return toStructuredPropertyFacetName(facet, aspectRetriever)
+        // Otherwise assume that this field is of keyword type.
+        .orElse(ESUtils.toKeywordField(facet, false, aspectRetriever));
   }
 
   List<String> getDefaultFacetFieldsFromAnnotation(final SearchableAnnotation annotation) {