fix: support for non-string types in object fields

datahub-project · Aug 1, 2024 · da12f1a · da12f1a
1 parent 2ab43f3
commit da12f1a
Show file tree

Hide file tree

Showing 4 changed files with 94 additions and 11 deletions.
diff --git a/...-io/src/main/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformer.java b/...-io/src/main/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformer.java
@@ -14,6 +14,7 @@
 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;
@@ -281,17 +282,33 @@ public void setSearchableValue(
               });
       searchDocument.set(fieldName, dictDoc);
     } else if (valueType == DataSchema.Type.MAP) {
-      ObjectNode dictDoc = JsonNodeFactory.instance.objectNode();
-      fieldValues
-          .subList(0, Math.min(fieldValues.size(), maxObjectKeys))
-          .forEach(
-              fieldValue -> {
+        ObjectNode dictDoc = JsonNodeFactory.instance.objectNode();
+        fieldValues
+        .subList(0, Math.min(fieldValues.size(), maxObjectKeys))
+        .forEach(
+            fieldValue -> {
                 String[] keyValues = fieldValue.toString().split("=");
-                String key = keyValues[0];
-                String value = keyValues[1];
-                dictDoc.put(key, value);
-              });
-      searchDocument.set(fieldName, dictDoc);
+                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);
+                }
+            });
+        searchDocument.set(fieldName, dictDoc);
     } else if (!fieldValues.isEmpty()) {
       getNodeForValue(valueType, fieldValues.get(0), fieldType)
           .ifPresent(node -> searchDocument.set(fieldName, node));

diff --git a/metadata-io/src/test/java/com/linkedin/metadata/TestEntityUtil.java b/metadata-io/src/test/java/com/linkedin/metadata/TestEntityUtil.java
@@ -20,6 +20,11 @@
 import com.linkedin.common.urn.Urn;
 import com.linkedin.data.template.StringArray;
 import com.linkedin.data.template.StringMap;
+import com.linkedin.data.template.BooleanMap;
+import com.linkedin.data.template.IntegerMap;
+import com.linkedin.data.template.LongMap;
+import com.linkedin.data.template.FloatMap;
+import com.linkedin.data.template.DoubleMap;
 
 public class TestEntityUtil {
   private TestEntityUtil() {}
@@ -70,6 +75,16 @@ public static TestEntityInfo getTestEntityInfo(Urn urn) {
                 "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;
   }
 

diff --git a/...src/test/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformerTest.java b/...src/test/java/com/linkedin/metadata/search/transformer/SearchDocumentTransformerTest.java
@@ -85,6 +85,11 @@ public void testTransform() throws IOException {
     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());
   }
 
   @Test

diff --git a/test-models/src/main/pegasus/com/datahub/test/TestEntityInfo.pdl b/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]
+
+}