[BEAM-121] DisplayData API tweaks

runners/google-cloud-dataflow-java/src/main/java/com/google/cloud/dataflow/sdk/runners/DataflowPipelineTranslator.java

@@ -81,7 +81,6 @@
 import com.google.cloud.dataflow.sdk.values.TypedPValue;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Strings;
-import com.google.common.collect.Lists;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -730,12 +729,8 @@ private void addOutput(String name, PValue value, Coder<?> valueCoder) {
     }
 
     private void addDisplayData(String name, DisplayData displayData) {
-      List<Map<String, Object>> serializedItems = Lists.newArrayList();
-      for (DisplayData.Item item : displayData.items()) {
-        serializedItems.add(MAPPER.convertValue(item, Map.class));
-      }
-
-      addList(getProperties(), name, serializedItems);
+      List<Map<String, Object>> list = MAPPER.convertValue(displayData, List.class);
+      addList(getProperties(), name, list);
     }
 
     @Override

runners/google-cloud-dataflow-java/src/test/java/com/google/cloud/dataflow/sdk/runners/DataflowPipelineTranslatorTest.java

@@ -818,8 +818,8 @@ public void processElement(ProcessContext c) throws Exception {
       @Override
       public void populateDisplayData(DisplayData.Builder builder) {
         builder
-                .add("foo", "bar")
-                .add("foo2", DataflowPipelineTranslatorTest.class)
+            .add("foo", "bar")
+            .add("foo2", DataflowPipelineTranslatorTest.class)
                 .withLabel("Test Class")
                 .withLinkUrl("http://www.google.com");
       }
@@ -833,7 +833,7 @@ public void processElement(ProcessContext c) throws Exception {
 
       @Override
       public void populateDisplayData(DisplayData.Builder builder) {
-        builder.add("foo3", "barge");
+        builder.add("foo3", 1234);
       }
     };
 
@@ -876,11 +876,11 @@ public void populateDisplayData(DisplayData.Builder builder) {
     );
 
     ImmutableList expectedFn2DisplayData = ImmutableList.of(
-            ImmutableMap.<String, String>builder()
+            ImmutableMap.<String, Object>builder()
                     .put("namespace", fn2.getClass().getName())
                     .put("key", "foo3")
-                    .put("type", "STRING")
-                    .put("value", "barge")
+                    .put("type", "INTEGER")
+                    .put("value", 1234L)
                     .build()
     );
 

sdks/java/core/src/main/java/com/google/cloud/dataflow/sdk/transforms/display/DisplayData.java

@@ -29,6 +29,7 @@
 
 import com.fasterxml.jackson.annotation.JsonGetter;
 import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonValue;
 
 import org.apache.avro.reflect.Nullable;
 import org.joda.time.Duration;
@@ -100,6 +101,7 @@ public static Type inferType(@Nullable Object value) {
     return Type.tryInferFrom(value);
   }
 
+  @JsonValue
   public Collection<Item> items() {
     return entries.values();
   }
@@ -175,6 +177,13 @@ public interface Builder {
      */
     ItemBuilder add(String key, long value);
 
+    /**
+     * Register the given numeric display data if the value is not null.
+     *
+     * @see DisplayData.Builder#add(String, long)
+     */
+    ItemBuilder addIfNotNull(String key, @Nullable Long value);
+
     /**
      * Register the given numeric display data if the value is different than the specified default.
      *
@@ -189,6 +198,13 @@ public interface Builder {
      */
     ItemBuilder add(String key, double value);
 
+    /**
+     * Register the given floating point display data if the value is not null.
+     *
+     * @see DisplayData.Builder#add(String, double)
+     */
+    ItemBuilder addIfNotNull(String key, @Nullable Double value);
+
     /**
      * Register the given floating point display data if the value is different than the specified
      * default.
@@ -204,6 +220,13 @@ public interface Builder {
      */
     ItemBuilder add(String key, boolean value);
 
+    /**
+     * Register the given boolean display data if the value is not null.
+     *
+     * @see DisplayData.Builder#add(String, boolean)
+     */
+    ItemBuilder addIfNotNull(String key, @Nullable Boolean value);
+
     /**
      * Register the given boolean display data if the value is different than the specified default.
      *
@@ -286,6 +309,7 @@ ItemBuilder addIfNotDefault(
    * transform or component.
    *
    * @throws ClassCastException if the value cannot be safely cast to the specified type.
+   *
    * @see DisplayData#inferType(Object)
    */
     ItemBuilder add(String key, Type type, Object value);
@@ -332,8 +356,8 @@ public static class Item {
     private final String key;
     private final String ns;
     private final Type type;
-    private final String value;
-    private final String shortValue;
+    private final Object value;
+    private final Object shortValue;
     private final String label;
     private final String url;
 
@@ -348,8 +372,8 @@ private Item(
         String namespace,
         String key,
         Type type,
-        String value,
-        String shortValue,
+        Object value,
+        Object shortValue,
         String url,
         String label) {
       this.ns = namespace;
@@ -384,7 +408,7 @@ public Type getType() {
      * Retrieve the value of the metadata item.
      */
     @JsonGetter("value")
-    public String getValue() {
+    public Object getValue() {
       return value;
     }
 
@@ -398,7 +422,7 @@ public String getValue() {
     @JsonGetter("shortValue")
     @JsonInclude(JsonInclude.Include.NON_NULL)
     @Nullable
-    public String getShortValue() {
+    public Object getShortValue() {
       return shortValue;
     }
 
@@ -540,48 +564,65 @@ public enum Type {
     STRING {
       @Override
       FormattedItemValue format(Object value) {
-        return new FormattedItemValue(value.toString());
+        return new FormattedItemValue(checkType(value, String.class, STRING));
       }
     },
     INTEGER {
       @Override
       FormattedItemValue format(Object value) {
-        Number number = (Number) value;
-        return new FormattedItemValue(Long.toString(number.longValue()));
+        if (value instanceof Integer) {
+          long l = ((Integer) value).longValue();
+          return format(l);
+        }
+
+        return new FormattedItemValue(checkType(value, Long.class, INTEGER));
       }
     },
     FLOAT {
       @Override
       FormattedItemValue format(Object value) {
-        return new FormattedItemValue(Double.toString((Double) value));
+        return new FormattedItemValue(checkType(value, Number.class, FLOAT));
       }
     },
     BOOLEAN() {
       @Override
       FormattedItemValue format(Object value) {
-        return new FormattedItemValue(Boolean.toString((boolean) value));
+        return new FormattedItemValue(checkType(value, Boolean.class, BOOLEAN));
       }
     },
     TIMESTAMP() {
       @Override
       FormattedItemValue format(Object value) {
-        return new FormattedItemValue((TIMESTAMP_FORMATTER.print((Instant) value)));
+        Instant instant = checkType(value, Instant.class, TIMESTAMP);
+        return new FormattedItemValue((TIMESTAMP_FORMATTER.print(instant)));
       }
     },
     DURATION {
       @Override
       FormattedItemValue format(Object value) {
-        return new FormattedItemValue(Long.toString(((Duration) value).getMillis()));
+        Duration duration = checkType(value, Duration.class, DURATION);
+        return new FormattedItemValue(duration.getMillis());
       }
     },
     JAVA_CLASS {
       @Override
       FormattedItemValue format(Object value) {
-        Class<?> clazz = (Class<?>) value;
+        Class<?> clazz = checkType(value, Class.class, JAVA_CLASS);
         return new FormattedItemValue(clazz.getName(), clazz.getSimpleName());
       }
     };
 
+    private static <T> T checkType(Object value, Class<T> clazz, DisplayData.Type expectedType) {
+      if (!clazz.isAssignableFrom(value.getClass())) {
+        throw new ClassCastException(String.format(
+            "Value is not valid for DisplayData type %s: %s", expectedType, value));
+      }
+
+      @SuppressWarnings("unchecked") // type checked above.
+      T typedValue = (T) value;
+      return typedValue;
+    }
+
     /**
      * Format the display metadata value into a long string representation, and optionally
      * a shorter representation for display.
@@ -592,7 +633,6 @@ FormattedItemValue format(Object value) {
 
     @Nullable
     private static Type tryInferFrom(@Nullable Object value) {
-      Type type;
       if (value instanceof Integer || value instanceof Long) {
         return INTEGER;
       } else if (value instanceof Double || value instanceof Float) {
@@ -614,23 +654,23 @@ private static Type tryInferFrom(@Nullable Object value) {
   }
 
   static class FormattedItemValue {
-    private final String shortValue;
-    private final String longValue;
+    private final Object shortValue;
+    private final Object longValue;
 
-    private FormattedItemValue(String longValue) {
+    private FormattedItemValue(Object longValue) {
       this(longValue, null);
     }
 
-    private FormattedItemValue(String longValue, String shortValue) {
+    private FormattedItemValue(Object longValue, Object shortValue) {
       this.longValue = longValue;
       this.shortValue = shortValue;
     }
 
-    String getLongValue() {
+    Object getLongValue() {
       return this.longValue;
     }
 
-    String getShortValue() {
+    Object getShortValue() {
       return this.shortValue;
     }
   }
@@ -700,29 +740,44 @@ public ItemBuilder add(String key, long value) {
       return addItemIf(true, key, Type.INTEGER, value);
     }
 
+    @Override
+    public ItemBuilder addIfNotNull(String key, @Nullable Long value) {
+      return addItemIf(value != null, key, Type.INTEGER, value);
+    }
+
     @Override
     public ItemBuilder addIfNotDefault(String key, long value, long defaultValue) {
-      return addItemIf(value != defaultValue, key, Type.INTEGER, value);
+      return addItemIf(!Objects.equals(value, defaultValue), key, Type.INTEGER, value);
     }
 
     @Override
     public ItemBuilder add(String key, double value) {
       return addItemIf(true, key, Type.FLOAT, value);
     }
 
+    @Override
+    public ItemBuilder addIfNotNull(String key, @Nullable Double value) {
+      return addItemIf(value != null, key, Type.FLOAT, value);
+    }
+
     @Override
     public ItemBuilder addIfNotDefault(String key, double value, double defaultValue) {
-      return addItemIf(value != defaultValue, key, Type.FLOAT, value);
+      return addItemIf(!Objects.equals(value, defaultValue), key, Type.FLOAT, value);
     }
 
     @Override
     public ItemBuilder add(String key, boolean value) {
       return addItemIf(true, key, Type.BOOLEAN, value);
     }
 
+    @Override
+    public ItemBuilder addIfNotNull(String key, @Nullable Boolean value) {
+      return addItemIf(value != null, key, Type.BOOLEAN, value);
+    }
+
     @Override
     public ItemBuilder addIfNotDefault(String key, boolean value, boolean defaultValue) {
-      return addItemIf(value != defaultValue, key, Type.BOOLEAN, value);
+      return addItemIf(!Objects.equals(value, defaultValue), key, Type.BOOLEAN, value);
     }
 
     @Override

sdks/java/core/src/test/java/com/google/cloud/dataflow/sdk/transforms/display/DisplayDataMatchers.java

@@ -301,20 +301,20 @@ protected DisplayData.Type featureValueOf(DisplayData.Item actual) {
    * value.
    */
 
-  public static Matcher<DisplayData.Item> hasValue(String value) {
+  public static Matcher<DisplayData.Item> hasValue(Object value) {
     return hasValue(Matchers.is(value));
   }
 
   /**
    * Creates a matcher that matches if the examined {@link DisplayData.Item} contains a value
    * matching the specified value matcher.
    */
-  public static Matcher<DisplayData.Item> hasValue(Matcher<String> valueMatcher) {
-    return new FeatureMatcher<DisplayData.Item, String>(
+  public static <T> Matcher<DisplayData.Item> hasValue(Matcher<T> valueMatcher) {
+    return new FeatureMatcher<DisplayData.Item, T>(
             valueMatcher, "with value", "value") {
       @Override
-      protected String featureValueOf(DisplayData.Item actual) {
-        return actual.getValue();
+      protected T featureValueOf(DisplayData.Item actual) {
+        return (T) actual.getValue();
       }
     };
   }