feat:(entity-registry): add ability to search for float and double (d…

…atahub-project#10324)
fjmacagno · Apr 24, 2024 · 437b7a1 · 437b7a1
1 parent e64229b
commit 437b7a1
Show file tree

Hide file tree

Showing 6 changed files with 23 additions and 12 deletions.
diff --git a/...-registry/src/main/java/com/linkedin/metadata/models/annotation/SearchableAnnotation.java b/...-registry/src/main/java/com/linkedin/metadata/models/annotation/SearchableAnnotation.java
@@ -153,10 +153,12 @@ private static FieldType getFieldType(
   private static FieldType getDefaultFieldType(DataSchema.Type schemaDataType) {
     switch (schemaDataType) {
       case INT:
-      case FLOAT:
         return FieldType.COUNT;
       case MAP:
         return FieldType.KEYWORD;
+      case FLOAT:
+      case DOUBLE:
+        return FieldType.DOUBLE;
       default:
         return FieldType.TEXT;
     }

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
@@ -98,10 +98,12 @@ private static SearchableAnnotation.FieldType getDefaultFieldType(
       DataSchema.Type schemaDataType) {
     switch (schemaDataType) {
       case INT:
-      case FLOAT:
         return SearchableAnnotation.FieldType.COUNT;
       case MAP:
         return SearchableAnnotation.FieldType.KEYWORD;
+      case FLOAT:
+      case DOUBLE:
+        return SearchableAnnotation.FieldType.DOUBLE;
       default:
         return SearchableAnnotation.FieldType.TEXT;
     }

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
@@ -314,6 +314,10 @@ private Optional<JsonNode> getNodeForValue(
         return Optional.of(JsonNodeFactory.instance.numberNode((Integer) fieldValue));
       case LONG:
         return Optional.of(JsonNodeFactory.instance.numberNode((Long) fieldValue));
+      case FLOAT:
+        return Optional.of(JsonNodeFactory.instance.numberNode((Float) fieldValue));
+      case DOUBLE:
+        return Optional.of(JsonNodeFactory.instance.numberNode((Double) fieldValue));
         // By default run toString
       default:
         String value = fieldValue.toString();

diff --git a/metadata-io/src/test/java/com/linkedin/metadata/TestEntityUtil.java b/metadata-io/src/test/java/com/linkedin/metadata/TestEntityUtil.java
@@ -69,6 +69,7 @@ public static TestEntityInfo getTestEntityInfo(Urn urn) {
                 "123",
                 "longValue",
                 "0123456789")));
+    testEntityInfo.setDoubleField(100.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
@@ -60,6 +60,7 @@ public void testTransform() throws IOException {
     assertTrue(result.isPresent());
     ObjectNode parsedJson = (ObjectNode) OBJECT_MAPPER.readTree(result.get());
     assertEquals(parsedJson.get("urn").asText(), snapshot.getUrn().toString());
+    assertEquals(parsedJson.get("doubleField").asDouble(), 100.456);
     assertEquals(parsedJson.get("keyPart1").asText(), "key");
     assertFalse(parsedJson.has("keyPart2"));
     assertEquals(parsedJson.get("keyPart3").asText(), "VALUE_1");
@@ -106,6 +107,7 @@ public void testTransformForDelete() throws IOException {
     parsedJson.get("nestedIntegerField").getNodeType().equals(JsonNodeType.NULL);
     parsedJson.get("feature1").getNodeType().equals(JsonNodeType.NULL);
     parsedJson.get("feature2").getNodeType().equals(JsonNodeType.NULL);
+    parsedJson.get("doubleField").getNodeType().equals(JsonNodeType.NULL);
   }
 
   @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
@@ -92,15 +92,15 @@ record TestEntityInfo includes CustomProperties {
   }
   esObjectField: optional map[string, string]
 
-    @Searchable = {
-      "fieldName": "doubleField",
-      "fieldType": "DOUBLE"
-    }
-    doubleField: optional double
+  @Searchable = {
+     "fieldName": "doubleField",
+     "fieldType": "DOUBLE"
+  }
+  doubleField: optional double
 
-    @Searchable = {
-      "fieldName": "removed",
-      "fieldType": "BOOLEAN"
-    }
-    removed: optional boolean
+  @Searchable = {
+     "fieldName": "removed",
+     "fieldType": "BOOLEAN"
+  }
+  removed: optional boolean
 }