fix: support for non-string types in object fields (#11066)

Co-authored-by: david-leifker <114954101+david-leifker@users.noreply.github.com>
Co-authored-by: milindgupta <milindgupta9@github.com>
Co-authored-by: Harshal Sheth <hsheth2@gmail.com>
Co-authored-by: Hyejin Yoon <0327jane@gmail.com>
Co-authored-by: milindgupta9 <133847413+milindgupta9@users.noreply.github.com>

entity-registry/src/test/java/com/linkedin/metadata/models/EntitySpecBuilderTest.java

@@ -189,7 +189,7 @@ public class EntitySpecBuilderTest {
         testEntityInfo.getPegasusSchema().getFullName());
 
     // Assert on Searchable Fields
-    assertEquals(testEntityInfo.getSearchableFieldSpecs().size(), 12);
+    assertEquals(testEntityInfo.getSearchableFieldSpecs().size(), 17);
     assertEquals(
         "customProperties",
         testEntityInfo

metadata-io/src/main/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformer.java

@@ -14,6 +14,7 @@ import com.linkedin.common.AuditStamp;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.data.DataMap;
 import com.linkedin.data.schema.DataSchema;
+import com.linkedin.data.schema.MapDataSchema;
 import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.entity.Aspect;
 import com.linkedin.events.metadata.ChangeType;
@@ -290,8 +291,39 @@ public class SearchDocumentTransformer {
                 String key = keyValues[0], value = "";
                 if (keyValues.length > 1) {
                   value = keyValues[1];
+                  if (((MapDataSchema) fieldSpec.getPegasusSchema())
+                      .getValues()
+                      .getType()
+                      .equals(DataSchema.Type.BOOLEAN)) {
+                    dictDoc.set(
+                        key, JsonNodeFactory.instance.booleanNode(Boolean.parseBoolean(value)));
+                  } else if (((MapDataSchema) fieldSpec.getPegasusSchema())
+                      .getValues()
+                      .getType()
+                      .equals(DataSchema.Type.INT)) {
+                    dictDoc.set(key, JsonNodeFactory.instance.numberNode(Integer.parseInt(value)));
+                  } else if (((MapDataSchema) fieldSpec.getPegasusSchema())
+                      .getValues()
+                      .getType()
+                      .equals(DataSchema.Type.DOUBLE)) {
+                    dictDoc.set(
+                        key, JsonNodeFactory.instance.numberNode(Double.parseDouble(value)));
+                  } else if (((MapDataSchema) fieldSpec.getPegasusSchema())
+                      .getValues()
+                      .getType()
+                      .equals(DataSchema.Type.LONG)) {
+                    dictDoc.set(key, JsonNodeFactory.instance.numberNode(Long.parseLong(value)));
+                  } else if (((MapDataSchema) fieldSpec.getPegasusSchema())
+                      .getValues()
+                      .getType()
+                      .equals(DataSchema.Type.FLOAT)) {
+                    dictDoc.set(key, JsonNodeFactory.instance.numberNode(Float.parseFloat(value)));
+                  } else {
+                    dictDoc.put(key, value);
+                  }
+                } else {
+                  dictDoc.put(key, value);
                 }
-                dictDoc.put(key, value);
               });
       searchDocument.set(fieldName, dictDoc);
     } else if (!fieldValues.isEmpty()) {

metadata-io/src/test/java/com/linkedin/metadata/TestEntityUtil.java

@@ -18,6 +18,11 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.linkedin.common.urn.TestEntityUrn;
 import com.linkedin.common.urn.Urn;
+import com.linkedin.data.template.BooleanMap;
+import com.linkedin.data.template.DoubleMap;
+import com.linkedin.data.template.FloatMap;
+import com.linkedin.data.template.IntegerMap;
+import com.linkedin.data.template.LongMap;
 import com.linkedin.data.template.StringArray;
 import com.linkedin.data.template.StringMap;
 
@@ -72,6 +77,13 @@ public class TestEntityUtil {
                 "longValue",
                 "0123456789")));
     testEntityInfo.setDoubleField(100.456);
+    testEntityInfo.setEsObjectFieldBoolean(
+        new BooleanMap(ImmutableMap.of("key1", true, "key2", false)));
+    testEntityInfo.setEsObjectFieldLong(new LongMap(ImmutableMap.of("key1", 1L, "key2", 2L)));
+    testEntityInfo.setEsObjectFieldFloat(new FloatMap(ImmutableMap.of("key1", 1.0f, "key2", 2.0f)));
+    testEntityInfo.setEsObjectFieldDouble(new DoubleMap(ImmutableMap.of("key1", 1.2, "key2", 2.4)));
+    testEntityInfo.setEsObjectFieldInteger(
+        new IntegerMap(ImmutableMap.of("key1", 123, "key2", 456)));
     return testEntityInfo;
   }
 

metadata-io/src/test/java/com/linkedin/metadata/search/indexbuilder/MappingsBuilderTest.java

@@ -31,7 +31,7 @@ public class MappingsBuilderTest {
     Map<String, Object> result = MappingsBuilder.getMappings(TestEntitySpecBuilder.getSpec());
     assertEquals(result.size(), 1);
     Map<String, Object> properties = (Map<String, Object>) result.get("properties");
-    assertEquals(properties.size(), 22);
+    assertEquals(properties.size(), 27);
     assertEquals(
         properties.get("urn"),
         ImmutableMap.of(

metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchQueryBuilderTest.java

@@ -111,19 +111,22 @@ public class SearchQueryBuilderTest extends AbstractTestNGSpringContextTests {
     assertEquals(keywordQuery.value(), "testQuery");
     assertEquals(keywordQuery.analyzer(), "keyword");
     Map<String, Float> keywordFields = keywordQuery.fields();
-    assertEquals(keywordFields.size(), 9);
+    assertEquals(keywordFields.size(), 14);
-    assertEquals(
+
-        keywordFields,
+    assertEquals(keywordFields.get("urn"), 10);
-        Map.of(
+    assertEquals(keywordFields.get("textArrayField"), 1);
-            "urn", 10.f,
+    assertEquals(keywordFields.get("customProperties"), 1);
-            "textArrayField", 1.0f,
+    assertEquals(keywordFields.get("wordGramField"), 1);
-            "customProperties", 1.0f,
+    assertEquals(keywordFields.get("nestedArrayArrayField"), 1);
-            "wordGramField", 1.0f,
+    assertEquals(keywordFields.get("textFieldOverride"), 1);
-            "nestedArrayArrayField", 1.0f,
+    assertEquals(keywordFields.get("nestedArrayStringField"), 1);
-            "textFieldOverride", 1.0f,
+    assertEquals(keywordFields.get("keyPart1"), 10);
-            "nestedArrayStringField", 1.0f,
+    assertEquals(keywordFields.get("esObjectField"), 1);
-            "keyPart1", 10.0f,
+    assertEquals(keywordFields.get("esObjectFieldFloat"), 1);
-            "esObjectField", 1.0f));
+    assertEquals(keywordFields.get("esObjectFieldDouble"), 1);
+    assertEquals(keywordFields.get("esObjectFieldLong"), 1);
+    assertEquals(keywordFields.get("esObjectFieldInteger"), 1);
+    assertEquals(keywordFields.get("esObjectFieldBoolean"), 1);
 
     SimpleQueryStringBuilder urnComponentQuery =
         (SimpleQueryStringBuilder) analyzerGroupQuery.should().get(1);
@@ -174,7 +177,7 @@ public class SearchQueryBuilderTest extends AbstractTestNGSpringContextTests {
                 })
             .collect(Collectors.toList());
 
-    assertEquals(prefixFieldWeights.size(), 29);
+    assertEquals(prefixFieldWeights.size(), 39);
 
     List.of(
             Pair.of("urn", 100.0f),
@@ -209,7 +212,7 @@ public class SearchQueryBuilderTest extends AbstractTestNGSpringContextTests {
     assertEquals(keywordQuery.queryString(), "testQuery");
     assertNull(keywordQuery.analyzer());
     Map<String, Float> keywordFields = keywordQuery.fields();
-    assertEquals(keywordFields.size(), 22);
+    assertEquals(keywordFields.size(), 27);
     assertEquals(keywordFields.get("keyPart1").floatValue(), 10.0f);
     assertFalse(keywordFields.containsKey("keyPart3"));
     assertEquals(keywordFields.get("textFieldOverride").floatValue(), 1.0f);
@@ -376,7 +379,7 @@ public class SearchQueryBuilderTest extends AbstractTestNGSpringContextTests {
     Set<SearchFieldConfig> fieldConfigs =
         TEST_CUSTOM_BUILDER.getStandardFields(
             mock(EntityRegistry.class), ImmutableList.of(TestEntitySpecBuilder.getSpec()));
-    assertEquals(fieldConfigs.size(), 22);
+    assertEquals(fieldConfigs.size(), 27);
     assertEquals(
         fieldConfigs.stream().map(SearchFieldConfig::fieldName).collect(Collectors.toSet()),
         Set.of(
@@ -401,7 +404,12 @@ public class SearchQueryBuilderTest extends AbstractTestNGSpringContextTests {
             "textFieldOverride.delimited",
             "urn",
             "wordGramField.wordGrams2",
-            "customProperties.delimited")); // customProperties.delimited Saas only
+            "customProperties.delimited",
+            "esObjectFieldBoolean",
+            "esObjectFieldInteger",
+            "esObjectFieldDouble",
+            "esObjectFieldFloat",
+            "esObjectFieldLong")); // customProperties.delimited Saas only
 
     assertEquals(
         fieldConfigs.stream()
@@ -487,7 +495,7 @@ public class SearchQueryBuilderTest extends AbstractTestNGSpringContextTests {
             ImmutableList.of(TestEntitySpecBuilder.getSpec(), mockEntitySpec));
     // Same 22 from the original entity + newFieldNotInOriginal + 3 word gram fields from the
     // textFieldOverride
-    assertEquals(fieldConfigs.size(), 27);
+    assertEquals(fieldConfigs.size(), 32);
     assertEquals(
         fieldConfigs.stream().map(SearchFieldConfig::fieldName).collect(Collectors.toSet()),
         Set.of(
@@ -517,7 +525,12 @@ public class SearchQueryBuilderTest extends AbstractTestNGSpringContextTests {
             "textFieldOverride.wordGrams2",
             "textFieldOverride.wordGrams3",
             "textFieldOverride.wordGrams4",
-            "customProperties.delimited"));
+            "customProperties.delimited",
+            "esObjectFieldBoolean",
+            "esObjectFieldInteger",
+            "esObjectFieldDouble",
+            "esObjectFieldFloat",
+            "esObjectFieldLong"));
 
     // Field which only exists in first one: Should be the same
     assertEquals(

metadata-io/src/test/java/com/linkedin/metadata/search/query/request/SearchRequestHandlerTest.java

@@ -228,7 +228,7 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests {
         highlightBuilder.fields().stream()
             .map(HighlightBuilder.Field::name)
             .collect(Collectors.toList());
-    assertEquals(fields.size(), 22);
+    assertEquals(fields.size(), 32);
     List<String> highlightableFields =
         ImmutableList.of(
             "keyPart1",
@@ -240,7 +240,12 @@ public class SearchRequestHandlerTest extends AbstractTestNGSpringContextTests {
             "nestedArrayArrayField",
             "customProperties",
             "esObjectField",
-            "wordGramField");
+            "wordGramField",
+            "esObjectFieldLong",
+            "esObjectFieldBoolean",
+            "esObjectFieldFloat",
+            "esObjectFieldDouble",
+            "esObjectFieldInteger");
     highlightableFields.forEach(
         field -> {
           assertTrue(fields.contains(field), "Missing: " + field);

metadata-io/src/test/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformerTest.java

@@ -85,7 +85,52 @@ public class SearchDocumentTransformerTest {
     assertEquals(parsedJson.get("feature2").asInt(), 1);
     JsonNode browsePathV2 = (JsonNode) parsedJson.get("browsePathV2");
     assertEquals(browsePathV2.asText(), "␟levelOne␟levelTwo");
+    assertEquals(
+        parsedJson.get("esObjectFieldBoolean").get("key1").getNodeType(),
+        JsonNodeFactory.instance.booleanNode(true).getNodeType());
+    assertEquals(
+        parsedJson.get("esObjectFieldLong").get("key1").getNodeType(),
+        JsonNodeFactory.instance.numberNode(1L).getNodeType());
+    assertEquals(
+        parsedJson.get("esObjectFieldFloat").get("key2").getNodeType(),
+        JsonNodeFactory.instance.numberNode(2.0f).getNodeType());
+    assertEquals(
+        parsedJson.get("esObjectFieldDouble").get("key1").getNodeType(),
+        JsonNodeFactory.instance.numberNode(1.2).getNodeType());
+    assertEquals(
+        parsedJson.get("esObjectFieldInteger").get("key2").getNodeType(),
+        JsonNodeFactory.instance.numberNode(456).getNodeType());
+    assertEquals(
+        parsedJson.get("esObjectFieldBoolean").get("key2").getNodeType(),
+        JsonNodeFactory.instance.booleanNode(false).getNodeType());
+    assertEquals(
+        parsedJson.get("esObjectFieldLong").get("key2").getNodeType(),
+        JsonNodeFactory.instance.numberNode(2L).getNodeType());
+    assertEquals(
+        parsedJson.get("esObjectFieldFloat").get("key1").getNodeType(),
+        JsonNodeFactory.instance.numberNode(1.0f).getNodeType());
+    assertEquals(
+        parsedJson.get("esObjectFieldDouble").get("key2").getNodeType(),
+        JsonNodeFactory.instance.numberNode(2.4).getNodeType());
+    assertEquals(
+        parsedJson.get("esObjectFieldInteger").get("key1").getNodeType(),
+        JsonNodeFactory.instance.numberNode(123).getNodeType());
     assertEquals(parsedJson.get("esObjectField").get("key3").asText(), "");
+    assertEquals(
+        parsedJson.get("esObjectFieldBoolean").get("key2").getNodeType(),
+        JsonNodeFactory.instance.booleanNode(false).getNodeType());
+    assertEquals(
+        parsedJson.get("esObjectFieldLong").get("key2").getNodeType(),
+        JsonNodeFactory.instance.numberNode(2L).getNodeType());
+    assertEquals(
+        parsedJson.get("esObjectFieldFloat").get("key1").getNodeType(),
+        JsonNodeFactory.instance.numberNode(1.0f).getNodeType());
+    assertEquals(
+        parsedJson.get("esObjectFieldDouble").get("key2").getNodeType(),
+        JsonNodeFactory.instance.numberNode(2.4).getNodeType());
+    assertEquals(
+        parsedJson.get("esObjectFieldInteger").get("key1").getNodeType(),
+        JsonNodeFactory.instance.numberNode(123).getNodeType());
   }
 
   @Test

test-models/src/main/pegasus/com/datahub/test/TestEntityInfo.pdl

@@ -103,4 +103,50 @@ record TestEntityInfo includes CustomProperties {
      "fieldType": "BOOLEAN"
   }
   removed: optional boolean
-}
+
+   @Searchable = {
+      "/*": {
+        "name": "esObjectFieldLong",
+        "fieldType": "OBJECT",
+        "queryByDefault": true
+     }
+  }
+  esObjectFieldLong: optional map[string, long]
+
+   @Searchable = {
+      "/*": {
+          "name": "esObjectFieldBoolean",
+          "fieldType": "OBJECT",
+          "queryByDefault": true
+      }
+  }
+  esObjectFieldBoolean: optional map[string, boolean]
+
+  @Searchable = {
+      "/*": {
+        "name": "esObjectFieldFloat",
+        "fieldType": "OBJECT",
+        "queryByDefault": true
+    }
+  }
+  esObjectFieldFloat: optional map[string, float]
+
+  @Searchable = {
+       "/*": {
+        "name": "esObjectFieldDouble",
+        "fieldType": "OBJECT",
+        "queryByDefault": true
+     }
+  }
+  esObjectFieldDouble: optional map[string, double]
+
+  @Searchable = {
+      "/*": {
+        "name": "esObjectFieldInteger",
+        "fieldType": "OBJECT",
+        "queryByDefault": true
+      }
+  }
+  esObjectFieldInteger: optional map[string, int]
+
+}