From c5030a0ebc09499a15704301703654cb34444547 Mon Sep 17 00:00:00 2001
From: Luke Cwik <lcwik@google.com>
Date: Mon, 18 Jul 2016 13:14:33 -0700
Subject: [PATCH] [BEAM-468] NullableCoder should not ask valueCoder
 isRegisterByteSizeObserverCheap when value is null

---
 .../apache/beam/sdk/coders/NullableCoder.java |  3 +++
 .../beam/sdk/coders/NullableCoderTest.java    | 26 +++++++++++--------
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/NullableCoder.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/NullableCoder.java
index 6fa7305de91a..cae101ab663e 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/NullableCoder.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/coders/NullableCoder.java
@@ -175,6 +175,9 @@ protected long getEncodedElementByteSize(@Nullable T value, Context context) thr
    */
   @Override
   public boolean isRegisterByteSizeObserverCheap(@Nullable T value, Context context) {
+    if (value == null) {
+      return true;
+    }
     return valueCoder.isRegisterByteSizeObserverCheap(value, context.nested());
   }
 }
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/coders/NullableCoderTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/coders/NullableCoderTest.java
index 1e669acf27b2..5bfbe05c1a46 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/coders/NullableCoderTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/coders/NullableCoderTest.java
@@ -90,29 +90,33 @@ public void testEncodingId() throws Exception {
 
   @Test
   public void testWireFormatEncode() throws Exception {
-      CoderProperties.coderEncodesBase64(TEST_CODER, TEST_VALUES, TEST_ENCODINGS);
+    CoderProperties.coderEncodesBase64(TEST_CODER, TEST_VALUES, TEST_ENCODINGS);
   }
 
   @Test
   public void testEncodedSize() throws Exception {
-      NullableCoder<Double> coder = NullableCoder.of(DoubleCoder.of());
-      assertEquals(1, coder.getEncodedElementByteSize(null, Coder.Context.OUTER));
-      assertEquals(9, coder.getEncodedElementByteSize(5.0, Coder.Context.OUTER));
+    NullableCoder<Double> coder = NullableCoder.of(DoubleCoder.of());
+    assertEquals(1, coder.getEncodedElementByteSize(null, Coder.Context.OUTER));
+    assertEquals(9, coder.getEncodedElementByteSize(5.0, Coder.Context.OUTER));
   }
 
   @Test
   public void testObserverIsCheap() throws Exception {
-      NullableCoder<Double> coder = NullableCoder.of(DoubleCoder.of());
-      assertTrue(coder.isRegisterByteSizeObserverCheap(null, Coder.Context.OUTER));
-      assertTrue(coder.isRegisterByteSizeObserverCheap(5.0, Coder.Context.OUTER));
+    NullableCoder<Double> coder = NullableCoder.of(DoubleCoder.of());
+    assertTrue(coder.isRegisterByteSizeObserverCheap(5.0, Coder.Context.OUTER));
   }
 
   @Test
   public void testObserverIsNotCheap() throws Exception {
-      NullableCoder<List<String>> coder = NullableCoder.of(ListCoder.of(StringUtf8Coder.of()));
-      assertFalse(coder.isRegisterByteSizeObserverCheap(null, Coder.Context.OUTER));
-      assertFalse(coder.isRegisterByteSizeObserverCheap(
-          ImmutableList.of("hi", "test"), Coder.Context.OUTER));
+    NullableCoder<List<String>> coder = NullableCoder.of(ListCoder.of(StringUtf8Coder.of()));
+    assertFalse(coder.isRegisterByteSizeObserverCheap(
+        ImmutableList.of("hi", "test"), Coder.Context.OUTER));
+  }
+
+  @Test
+  public void testObserverIsAlwaysCheapForNullValues() throws Exception {
+    NullableCoder<List<String>> coder = NullableCoder.of(ListCoder.of(StringUtf8Coder.of()));
+    assertTrue(coder.isRegisterByteSizeObserverCheap(null, Coder.Context.OUTER));
   }
 
   @Test