Fix data search for fieldType = OBJECT for non-string fields (#11524)

Co-authored-by: david-leifker <114954101+david-leifker@users.noreply.github.com>
datahub-project · Oct 8, 2024 · 5335454 · 5335454
1 parent 86d394c
commit 5335454
Show file tree

Hide file tree

Showing 3 changed files with 106 additions and 3 deletions.
diff --git a/...n/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java b/...n/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java
@@ -13,6 +13,8 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.linkedin.common.urn.Urn;
+import com.linkedin.data.schema.DataSchema;
+import com.linkedin.data.schema.MapDataSchema;
 import com.linkedin.data.template.DoubleMap;
 import com.linkedin.metadata.config.search.SearchConfiguration;
 import com.linkedin.metadata.config.search.custom.CustomSearchConfiguration;
@@ -662,8 +664,48 @@ private static Map<String, Set<SearchableAnnotation.FieldType>> buildSearchableF
                                       Collections.emptySet())))
                       .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
 
+              List<SearchableFieldSpec> objectFieldSpec =
+                  entitySpec.getSearchableFieldSpecs().stream()
+                      .filter(
+                          searchableFieldSpec ->
+                              searchableFieldSpec.getSearchableAnnotation().getFieldType()
+                                  == SearchableAnnotation.FieldType.OBJECT)
+                      .collect(Collectors.toList());
+
+              Map<String, Set<SearchableAnnotation.FieldType>> objectFieldTypes = new HashMap<>();
+
+              objectFieldSpec.forEach(
+                  fieldSpec -> {
+                    String fieldName = fieldSpec.getSearchableAnnotation().getFieldName();
+                    DataSchema.Type dataType =
+                        ((MapDataSchema) fieldSpec.getPegasusSchema()).getValues().getType();
+
+                    Set<SearchableAnnotation.FieldType> fieldType;
+
+                    switch (dataType) {
+                      case BOOLEAN:
+                        fieldType = Set.of(SearchableAnnotation.FieldType.BOOLEAN);
+                        break;
+                      case INT:
+                        fieldType = Set.of(SearchableAnnotation.FieldType.COUNT);
+                        break;
+                      case DOUBLE:
+                      case LONG:
+                      case FLOAT:
+                        fieldType = Set.of(SearchableAnnotation.FieldType.DOUBLE);
+                        break;
+                      default:
+                        fieldType = Set.of(SearchableAnnotation.FieldType.TEXT);
+                        break;
+                    }
+                    objectFieldTypes.put(fieldName, fieldType);
+                    annotationFieldTypes.remove(fieldName);
+                  });
+
               return Stream.concat(
-                  annotationFieldTypes.entrySet().stream(),
+                  Stream.concat(
+                      objectFieldTypes.entrySet().stream(),
+                      annotationFieldTypes.entrySet().stream()),
                   Stream.concat(
                       mappingFieldTypes.entrySet().stream(), aliasFieldTypes.entrySet().stream()));
             })

diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java b/metadata-io/src/main/java/com/linkedin/metadata/search/utils/ESUtils.java
@@ -759,7 +759,7 @@ private static Set<String> getFieldTypes(
           StructuredPropertyUtils.toElasticsearchFieldType(fieldName, aspectRetriever);
     } else {
       Set<SearchableAnnotation.FieldType> fieldTypes =
-          searchableFields.getOrDefault(fieldName, Collections.emptySet());
+          searchableFields.getOrDefault(fieldName.split("\\.")[0], Collections.emptySet());
       finalFieldTypes =
           fieldTypes.stream().map(ESUtils::getElasticTypeForFieldType).collect(Collectors.toSet());
     }
@@ -785,6 +785,7 @@ private static RangeQueryBuilder buildRangeQueryFromCriterion(
     // Determine criterion value, range query only accepts single value so take first value in
     // values if multiple
     String criterionValueString = criterion.getValues().get(0).trim();
+
     Object criterionValue;
     String documentFieldName;
     if (fieldTypes.contains(BOOLEAN_FIELD_TYPE)) {

diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/utils/ESUtilsTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/utils/ESUtilsTest.java
@@ -10,10 +10,13 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
+import com.google.common.collect.ImmutableList;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.data.template.SetMode;
+import com.linkedin.data.template.StringArray;
 import com.linkedin.entity.Aspect;
 import com.linkedin.metadata.aspect.AspectRetriever;
+import com.linkedin.metadata.models.annotation.SearchableAnnotation;
 import com.linkedin.metadata.query.filter.Condition;
 import com.linkedin.metadata.query.filter.Criterion;
 import com.linkedin.metadata.search.elasticsearch.query.filter.QueryFilterRewriteChain;
@@ -38,7 +41,7 @@ public class ESUtilsTest {
   private static AspectRetriever aspectRetrieverV1;
 
   @BeforeClass
-  public void setup() throws RemoteInvocationException, URISyntaxException {
+  public static void setup() throws RemoteInvocationException, URISyntaxException {
     Urn abFghTenUrn = Urn.createFromString("urn:li:structuredProperty:ab.fgh.ten");
 
     // legacy
@@ -835,4 +838,61 @@ public void testGetQueryBuilderFromStructPropExistsV1() {
             + "}";
     Assert.assertEquals(result.toString(), expected);
   }
+
+  @Test
+  public void testGetQueryBuilderForObjectFields() {
+    final Criterion singleValueCriterion =
+        new Criterion()
+            .setField("testObjectField.numericField")
+            .setCondition(Condition.EQUAL)
+            .setValues(new StringArray(ImmutableList.of("10")));
+
+    Map<String, Set<SearchableAnnotation.FieldType>> searchableFieldTypes = new HashMap<>();
+    searchableFieldTypes.put("testObjectField", Set.of(SearchableAnnotation.FieldType.DOUBLE));
+
+    QueryBuilder result =
+        ESUtils.getQueryBuilderFromCriterion(
+            singleValueCriterion,
+            false,
+            searchableFieldTypes,
+            mock(OperationContext.class),
+            QueryFilterRewriteChain.EMPTY);
+    String expected =
+        "{\n"
+            + "  \"terms\" : {\n"
+            + "    \"testObjectField.numericField\" : [\n"
+            + "      10.0\n"
+            + "    ],\n"
+            + "    \"boost\" : 1.0,\n"
+            + "    \"_name\" : \"testObjectField.numericField\"\n"
+            + "  }\n"
+            + "}";
+    Assert.assertEquals(result.toString(), expected);
+
+    final Criterion multiValueCriterion =
+        new Criterion()
+            .setField("testObjectField.numericField")
+            .setCondition(Condition.EQUAL)
+            .setValues(new StringArray(ImmutableList.of("10", "20")));
+
+    result =
+        ESUtils.getQueryBuilderFromCriterion(
+            multiValueCriterion,
+            false,
+            searchableFieldTypes,
+            mock(OperationContext.class),
+            QueryFilterRewriteChain.EMPTY);
+    expected =
+        "{\n"
+            + "  \"terms\" : {\n"
+            + "    \"testObjectField.numericField\" : [\n"
+            + "      10.0,\n"
+            + "      20.0\n"
+            + "    ],\n"
+            + "    \"boost\" : 1.0,\n"
+            + "    \"_name\" : \"testObjectField.numericField\"\n"
+            + "  }\n"
+            + "}";
+    Assert.assertEquals(result.toString(), expected);
+  }
 }