fix(search): implement queryByDefault annotation for SearchableRef (#…

…10603)
datahub-project · May 28, 2024 · 0fd451c · 0fd451c
1 parent 7eba5f2
commit 0fd451c
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 4 deletions.
diff --git a/...-registry/src/main/java/com/linkedin/metadata/models/SearchableRefFieldSpecExtractor.java b/...-registry/src/main/java/com/linkedin/metadata/models/SearchableRefFieldSpecExtractor.java
@@ -115,6 +115,7 @@ private void extractSearchableRefAnnotation(
             new SearchableRefAnnotation(
                 pathName,
                 annotation.getFieldType(),
+                annotation.isQueryByDefault(),
                 annotation.getBoostScore(),
                 annotation.getDepth(),
                 annotation.getRefType(),

diff --git a/...gistry/src/main/java/com/linkedin/metadata/models/annotation/SearchableRefAnnotation.java b/...gistry/src/main/java/com/linkedin/metadata/models/annotation/SearchableRefAnnotation.java
@@ -32,6 +32,8 @@ public class SearchableRefAnnotation {
   String fieldName;
   // Type of the field. Defines how the field is indexed and matched
   SearchableAnnotation.FieldType fieldType;
+  // Whether we should match the field for the default search query
+  boolean queryByDefault;
   // Boost multiplier to the match score. Matches on fields with higher boost score ranks higher
   double boostScore;
   // defines what depth should be explored of reference object
@@ -72,20 +74,35 @@ public static SearchableRefAnnotation fromPegasusAnnotationObject(
                   + "Mandatory input field refType defining the Entity Type is not provided",
               ANNOTATION_NAME, context));
     }
+    final Optional<Boolean> queryByDefault =
+        AnnotationUtils.getField(map, "queryByDefault", Boolean.class);
     final Optional<Integer> depth = AnnotationUtils.getField(map, "depth", Integer.class);
     final Optional<Double> boostScore = AnnotationUtils.getField(map, "boostScore", Double.class);
     final List<String> fieldNameAliases = getFieldNameAliases(map);
     final SearchableAnnotation.FieldType resolvedFieldType =
         getFieldType(fieldType, schemaDataType);
+
     return new SearchableRefAnnotation(
         fieldName.orElse(schemaFieldName),
         resolvedFieldType,
+        getQueryByDefault(queryByDefault, resolvedFieldType),
         boostScore.orElse(1.0),
         depth.orElse(2),
         refType.get(),
         fieldNameAliases);
   }
 
+  private static Boolean getQueryByDefault(
+      Optional<Boolean> maybeQueryByDefault, SearchableAnnotation.FieldType fieldType) {
+    if (!maybeQueryByDefault.isPresent()) {
+      if (DEFAULT_QUERY_FIELD_TYPES.contains(fieldType)) {
+        return Boolean.TRUE;
+      }
+      return Boolean.FALSE;
+    }
+    return maybeQueryByDefault.get();
+  }
+
   private static SearchableAnnotation.FieldType getFieldType(
       Optional<String> maybeFieldType, DataSchema.Type schemaDataType) {
     if (!maybeFieldType.isPresent()) {

diff --git a/...ain/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchQueryBuilder.java b/...ain/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchQueryBuilder.java
@@ -181,6 +181,13 @@ public Set<SearchFieldConfig> getStandardFields(
     return fields;
   }
 
+  /**
+   * Return query by default fields
+   *
+   * @param entityRegistry entity registry with search annotations
+   * @param entitySpec the entity spect
+   * @return set of queryByDefault field configurations
+   */
   @VisibleForTesting
   public Set<SearchFieldConfig> getFieldsFromEntitySpec(
       @Nonnull EntityRegistry entityRegistry, EntitySpec entitySpec) {
@@ -212,14 +219,20 @@ public Set<SearchFieldConfig> getFieldsFromEntitySpec(
 
     List<SearchableRefFieldSpec> searchableRefFieldSpecs = entitySpec.getSearchableRefFieldSpecs();
     for (SearchableRefFieldSpec refFieldSpec : searchableRefFieldSpecs) {
+      if (!refFieldSpec.getSearchableRefAnnotation().isQueryByDefault()) {
+        continue;
+      }
+
       int depth = refFieldSpec.getSearchableRefAnnotation().getDepth();
-      Set<SearchFieldConfig> searchFieldConfig =
-          SearchFieldConfig.detectSubFieldType(refFieldSpec, depth, entityRegistry);
-      fields.addAll(searchFieldConfig);
+      Set<SearchFieldConfig> searchFieldConfigs =
+          SearchFieldConfig.detectSubFieldType(refFieldSpec, depth, entityRegistry).stream()
+              .filter(SearchFieldConfig::isQueryByDefault)
+              .collect(Collectors.toSet());
+      fields.addAll(searchFieldConfigs);
 
       Map<String, SearchableAnnotation.FieldType> fieldTypeMap =
           getAllFieldTypeFromSearchableRef(refFieldSpec, depth, entityRegistry, "");
-      for (SearchFieldConfig fieldConfig : searchFieldConfig) {
+      for (SearchFieldConfig fieldConfig : searchFieldConfigs) {
         if (fieldConfig.hasDelimitedSubfield()) {
           fields.add(
               SearchFieldConfig.detectSubFieldType(

diff --git a/...a-io/src/test/java/com/linkedin/metadata/search/query/request/SearchQueryBuilderTest.java b/...a-io/src/test/java/com/linkedin/metadata/search/query/request/SearchQueryBuilderTest.java
@@ -541,4 +541,16 @@ public void testGetStandardFields() {
             .map(SearchFieldConfig::boost),
         Optional.of(2.0F));
   }
+
+  @Test
+  public void testStandardFieldsQueryByDefault() {
+    assertTrue(
+        TEST_BUILDER
+            .getStandardFields(
+                opContext.getEntityRegistry(),
+                opContext.getEntityRegistry().getEntitySpecs().values())
+            .stream()
+            .allMatch(SearchFieldConfig::isQueryByDefault),
+        "Expect all search fields to be queryByDefault.");
+  }
 }