[BEAM-3900] various TypeDescriptor bugfixes

seznam · Oct 4, 2018 · 49a6c18 · 49a6c18
1 parent 5875038
commit 49a6c18
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 4 deletions.
diff --git a/...rc/main/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/CountByKey.java b/...rc/main/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/CountByKey.java
@@ -266,7 +266,7 @@ private CountByKey(
   public Dataset<KV<KeyT, Long>> expand(List<Dataset<InputT>> inputs) {
     return ReduceByKey.named(getName().orElse(null))
         .of(Iterables.getOnlyElement(inputs))
-        .keyBy(getKeyExtractor())
+        .keyBy(getKeyExtractor(), getKeyType().orElse(null))
         .valueBy(v -> 1L, TypeDescriptors.longs())
         .combineBy(Sums.ofLongs())
         .applyIf(

diff --git a/...c/main/java/org/apache/beam/sdk/extensions/euphoria/core/translate/OperatorTransform.java b/...c/main/java/org/apache/beam/sdk/extensions/euphoria/core/translate/OperatorTransform.java
@@ -17,6 +17,7 @@
  */
 package org.apache.beam.sdk.extensions.euphoria.core.translate;
 
+import com.google.common.base.Preconditions;
 import java.util.List;
 import java.util.Optional;
 import java.util.stream.Collectors;
@@ -51,6 +52,8 @@ public static <InputT, OutputT, OperatorT extends Operator<OutputT>> Dataset<Out
           inputList.apply(
               operator.getName().orElseGet(() -> operator.getClass().getName()),
               new OperatorTransform<>(operator, maybeTranslator.get()));
+      Preconditions.checkState(
+          output.getTypeDescriptor() != null, "Translator should always return typed PCollection.");
       return Dataset.of(output, operator);
     }
 

diff --git a/...in/java/org/apache/beam/sdk/extensions/euphoria/core/translate/ReduceByKeyTranslator.java b/...in/java/org/apache/beam/sdk/extensions/euphoria/core/translate/ReduceByKeyTranslator.java
@@ -86,7 +86,11 @@ public PCollection<KV<KeyT, OutputT>> translate(
               Combine.perKey(asCombiner(reducer, accumulators, operator.getName().orElse(null))));
       @SuppressWarnings("unchecked")
       final PCollection<KV<KeyT, OutputT>> casted = (PCollection) combined;
-      return casted;
+      return casted.setTypeDescriptor(
+          operator
+              .getOutputType()
+              .orElseThrow(
+                  () -> new IllegalStateException("Unable to infer output type descriptor.")));
     }
 
     return extracted

diff --git a/...src/main/java/org/apache/beam/sdk/extensions/euphoria/core/translate/UnionTranslator.java b/...src/main/java/org/apache/beam/sdk/extensions/euphoria/core/translate/UnionTranslator.java
@@ -21,15 +21,17 @@
 import org.apache.beam.sdk.transforms.Flatten;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PCollectionList;
+import org.apache.beam.sdk.values.TypeDescriptor;
 
 /** Euphoria to Beam translation of {@link Union} operator. */
 class UnionTranslator<InputT> implements OperatorTranslator<InputT, InputT, Union<InputT>> {
 
   @Override
   public PCollection<InputT> translate(Union<InputT> operator, PCollectionList<InputT> inputs) {
+    final TypeDescriptor<InputT> outputType = operator.getOutputType().orElse(null);
     return operator
         .getName()
-        .map(name -> inputs.apply(name, Flatten.pCollections()))
-        .orElseGet(() -> inputs.apply(Flatten.pCollections()));
+        .map(name -> inputs.apply(name, Flatten.pCollections()).setTypeDescriptor(outputType))
+        .orElseGet(() -> inputs.apply(Flatten.pCollections()).setTypeDescriptor(outputType));
   }
 }