fix(field-extraction): supporting string union extraction in record u…

…tils (#125)

dao-api/src/main/java/com/linkedin/metadata/dao/utils/RecordUtils.java

@@ -413,18 +413,26 @@ private static Object invokeMethod(@Nonnull RecordTemplate record, @Nonnull Stri
   /**
    * Helper method for referencing array of RecordTemplate objects. Referencing a particular index or range of indices of an array is not supported.
    *
-   * @param reference {@link AbstractArrayTemplate} corresponding to array of {@link RecordTemplate} which needs to be referenced
+   * @param reference {@link AbstractArrayTemplate} corresponding to array of {@link RecordTemplate} or {@link UnionTemplate} which needs to be referenced
    * @param ps {@link PathSpec} for the entire path inside the array that needs to be referenced
    * @return {@link List} of objects from the array, referenced using the PathSpec
    */
   @Nonnull
   @SuppressWarnings("rawtypes")
-  private static List<Object> getReferenceForAbstractArray(@Nonnull AbstractArrayTemplate<RecordTemplate> reference, @Nonnull PathSpec ps) {
+  private static List<Object> getReferenceForAbstractArray(@Nonnull AbstractArrayTemplate<?> reference, @Nonnull PathSpec ps) {
     if (!reference.isEmpty()) {
-      return Arrays.stream((reference).toArray())
-          .map(x -> getFieldValue(((RecordTemplate) x), ps))
-          .flatMap(o -> o.map(Stream::of).orElseGet(Stream::empty))
-          .collect(Collectors.toList());
+      return Arrays.stream((reference).toArray()).map(x -> {
+        if (x instanceof RecordTemplate) {
+          return getFieldValue(((RecordTemplate) x), ps);
+        }
+        if (x instanceof UnionTemplate) {
+          if (ps.getPathComponents().size() != 1) {
+            throw new InvalidSchemaException("The currently selected member isn't a union of primitives: " + ps);
+          }
+          return Optional.ofNullable(((DataMap) ((UnionTemplate) x).data()).get(ps.getPathComponents().get(0)));
+        }
+        throw new InvalidSchemaException("The currently selected member is neither a RecordTemplate or a Union: " + ps);
+      }).flatMap(o -> o.map(Stream::of).orElseGet(Stream::empty)).collect(Collectors.toList());
     }
     return Collections.emptyList();
   }
@@ -463,7 +471,8 @@ public static Optional<Object> getFieldValue(@Nonnull RecordTemplate recordTempl
 
   /**
    * Given a {@link RecordTemplate} and {@link com.linkedin.data.schema.PathSpec} this will return value of the path from the record.
-   * This handles only RecordTemplate, fields of which can be primitive types, typeRefs, arrays of primitive types or array of records.
+   * This handles only RecordTemplate, fields of which can be primitive types, typeRefs, arrays of primitive types, arrays of records
+   * or arrays of primitive union types.
    * Fetching of values in a RecordTemplate where the field has a default value will return the field default value.
    * Referencing field corresponding to a particular index or range of indices of an array is not supported.
    * Fields corresponding to 1) multi-dimensional array 2) UnionTemplate 3) AbstractMapTemplate 4) FixedTemplate are currently not supported.
@@ -492,7 +501,7 @@ public static Optional<Object> getFieldValue(@Nonnull RecordTemplate recordTempl
         }
       } else if (reference instanceof AbstractArrayTemplate) {
         return Optional.of(getReferenceForAbstractArray(
-            (AbstractArrayTemplate<RecordTemplate>) reference, new PathSpec(ps.getPathComponents().subList(i, pathSize))));
+            (AbstractArrayTemplate<?>) reference, new PathSpec(ps.getPathComponents().subList(i, pathSize))));
       } else {
         throw new UnsupportedOperationException(String.format("Failed at extracting %s (%s from %s)", part, ps, recordTemplate));
       }

dao-api/src/test/java/com/linkedin/metadata/dao/utils/RecordUtilsTest.java

@@ -15,6 +15,8 @@
 import com.linkedin.testing.EntitySnapshot;
 import com.linkedin.testing.EntityValueArray;
 import com.linkedin.testing.MixedRecord;
+import com.linkedin.testing.StringUnion;
+import com.linkedin.testing.StringUnionArray;
 import com.linkedin.testing.singleaspectentity.EntityValue;
 import com.linkedin.testing.urn.FooUrn;
 import java.io.IOException;
@@ -373,6 +375,41 @@ public void testGetFieldValueArray() {
     assertFalse(o7.isPresent());
   }
 
+  @Test(description = "Test getFieldValue() when RecordTemplate has field of type array of primitive unions")
+  public void testGetFieldValueArrayOfPrimitiveUnions() {
+
+    // case 1: array of unions of strings
+    final MixedRecord mixedRecord1 =
+        new MixedRecord().setUnionArray(new StringUnionArray(Arrays.asList(
+            StringUnion.create("val1"),
+            StringUnion.create("val2"),
+            StringUnion.create("val3"),
+            StringUnion.create("val4")
+        )));
+
+    PathSpec ps1 = MixedRecord.fields().unionArray();
+    Object o1 = RecordUtils.getFieldValue(mixedRecord1, ps1).get();
+
+    PathSpec ps2 = MixedRecord.fields().unionArray().items();
+    Object o2 = RecordUtils.getFieldValue(mixedRecord1, ps2).get();
+
+    assertEquals(o1, new StringUnionArray(Arrays.asList(
+        StringUnion.create("val1"),
+        StringUnion.create("val2"),
+        StringUnion.create("val3"),
+        StringUnion.create("val4")
+    )));
+    assertEquals(ps1.toString(), "/unionArray");
+
+    assertEquals(o2, new StringUnionArray(Arrays.asList(
+        StringUnion.create("val1"),
+        StringUnion.create("val2"),
+        StringUnion.create("val3"),
+        StringUnion.create("val4")
+    )));
+    assertEquals(ps2.toString(), "/unionArray/*");
+  }
+
   @Test
   public void testCapitalizeFirst() {
     String s = "field1";

testing/test-models/src/main/pegasus/com/linkedin/testing/MixedRecord.pdl

@@ -45,6 +45,11 @@ record MixedRecord {
    */
   nestedRecordArray: optional array[EntityValue]
 
+  /**
+   * For unit tests
+   */
+  unionArray: optional array[StringUnion]
+
   /**
    * For unit tests
    */
@@ -59,4 +64,4 @@ record MixedRecord {
    * For unit tests
    */
   recordTypeRef: optional typeref AspectFooTypeRef = AspectFoo
-}
+}

testing/test-models/src/main/pegasus/com/linkedin/testing/StringUnion.pdl

@@ -0,0 +1,6 @@
+namespace com.linkedin.testing
+
+/**
+ * For unit testing
+ */
+typeref StringUnion = union[string]