From abd5dd91ec8dc3a5f093638b67c74952a242f786 Mon Sep 17 00:00:00 2001
From: Thomas Groh <tgroh@google.com>
Date: Tue, 25 Apr 2017 17:39:38 -0700
Subject: [PATCH] Add getElementCoders to UnionCoder

This is an explicit method on UnionCoder that is completely unrelated to
the serialization and deserialization of the Coder.
---
 .../apache/beam/sdk/transforms/join/CoGbkResult.java  |  2 +-
 .../apache/beam/sdk/transforms/join/UnionCoder.java   |  4 ++++
 .../beam/sdk/transforms/join/UnionCoderTest.java      | 11 +++++++++++
 3 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/join/CoGbkResult.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/join/CoGbkResult.java
index 83e178e68d0e..6c62cbe73cc0 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/join/CoGbkResult.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/join/CoGbkResult.java
@@ -269,7 +269,7 @@ public CoGbkResult decode(
 
     @SuppressWarnings("rawtypes")
     private IterableCoder tagListCoder(int unionTag) {
-      return IterableCoder.of(unionCoder.getComponents().get(unionTag));
+      return IterableCoder.of(unionCoder.getElementCoders().get(unionTag));
     }
 
     @Override
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/join/UnionCoder.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/join/UnionCoder.java
index 2beff5730d55..f411cd1333b4 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/join/UnionCoder.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/join/UnionCoder.java
@@ -91,6 +91,10 @@ public List<? extends Coder<?>> getComponents() {
     return elementCoders;
   }
 
+  public List<? extends Coder<?>> getElementCoders() {
+    return elementCoders;
+  }
+
   /**
    * Since this coder uses elementCoders.get(index) and coders that are known to run in constant
    * time, we defer the return value to that coder.
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/join/UnionCoderTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/join/UnionCoderTest.java
index fe494959a42c..8c8c3f95e0da 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/join/UnionCoderTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/join/UnionCoderTest.java
@@ -17,11 +17,14 @@
  */
 package org.apache.beam.sdk.transforms.join;
 
+import static org.junit.Assert.assertThat;
+
 import com.google.common.collect.ImmutableList;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.DoubleCoder;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
 import org.apache.beam.sdk.testing.CoderProperties;
+import org.hamcrest.Matchers;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -31,6 +34,14 @@
  */
 @RunWith(JUnit4.class)
 public class UnionCoderTest {
+  @Test
+  public void testGetElementCoders() {
+    UnionCoder unionCoder =
+        UnionCoder.of(ImmutableList.<Coder<?>>of(StringUtf8Coder.of(), DoubleCoder.of()));
+    assertThat(
+        unionCoder.getElementCoders().get(0), Matchers.<Coder<?>>equalTo(StringUtf8Coder.of()));
+    assertThat(unionCoder.getElementCoders().get(1), Matchers.<Coder<?>>equalTo(DoubleCoder.of()));
+  }
 
   @Test
   public void testCoderIsSerializable() {