Merge pull request apache#9 from seznam/simunek/addTypeHint/BEAM-4429

[BEAM-4429] Add type hint
cosmoskitten · Jun 21, 2018 · 529d05d · 529d05d
2 parents aa1af7a + cef3e44
commit 529d05d
Show file tree

Hide file tree

Showing 24 changed files with 548 additions and 140 deletions.
diff --git a/...core/client/functional/TypeHintAware.java → ...lient/functional/TypeDescriptorAware.java b/...core/client/functional/TypeHintAware.java → ...lient/functional/TypeDescriptorAware.java
@@ -19,17 +19,18 @@
 
 import java.io.Serializable;
 import org.apache.beam.sdk.extensions.euphoria.core.annotation.audience.Audience;
-import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeHint;
+import org.apache.beam.sdk.values.TypeDescriptor;
+
 
 /** @param <T> */
 @Audience(Audience.Type.INTERNAL)
-public interface TypeHintAware<T> extends Serializable {
+public interface TypeDescriptorAware<T> extends Serializable {
 
   /**
    * Retrieve type hint associated with this object. Mostly this represents type returned by
    * function.
    *
-   * @return {@link TypeHint} associated with this object
+   * @return {@link TypeDescriptor} associated with this object
    */
-  TypeHint<T> getTypeHint();
+  TypeDescriptor<T> getTypeDescriptor();
 }
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
@@ -35,11 +35,13 @@
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.base.StateAwareWindowWiseSingleInputOperator;
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.hint.OutputHint;
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.windowing.WindowingDesc;
+import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeAwareUnaryFunction;
 import org.apache.beam.sdk.extensions.euphoria.core.client.util.Pair;
 import org.apache.beam.sdk.extensions.euphoria.core.executor.graph.DAG;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.Trigger;
 import org.apache.beam.sdk.transforms.windowing.WindowFn;
+import org.apache.beam.sdk.values.TypeDescriptor;
 import org.apache.beam.sdk.values.WindowingStrategy;
 
 /**
@@ -157,6 +159,12 @@ public <K> WindowingBuilder<InputT, K> keyBy(UnaryFunction<InputT, K> keyExtract
       paramsCasted.keyExtractor = Objects.requireNonNull(keyExtractor);
       return new WindowingBuilder<>(paramsCasted);
     }
+
+    @Override
+    public <K> WindowingBuilder<InputT, K> keyBy(
+        UnaryFunction<InputT, K> keyExtractor, TypeDescriptor<K> typeHint) {
+      return keyBy(TypeAwareUnaryFunction.of(keyExtractor, typeHint));
+    }
   }
 
   /**

diff --git a/...e/src/main/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/FlatMap.java b/...e/src/main/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/FlatMap.java
@@ -33,6 +33,8 @@
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.base.Builders;
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.base.ElementWiseOperator;
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.hint.OutputHint;
+import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeAwareUnaryFunctor;
+import org.apache.beam.sdk.values.TypeDescriptor;
 
 /**
  * A transformation of a dataset from one type into another allowing user code to generate zero,
@@ -183,6 +185,11 @@ public <OutputT> EventTimeBuilder<InputT, OutputT> using(
         UnaryFunctor<InputT, OutputT> functor) {
       return new EventTimeBuilder<>(this, functor);
     }
+
+    public <OutputT> EventTimeBuilder<InputT, OutputT> using(
+        UnaryFunctor<InputT, OutputT> functor, TypeDescriptor<OutputT> outputTypeDescriptor) {
+      return using(TypeAwareUnaryFunctor.of(functor, outputTypeDescriptor));
+    }
   }
 
   /** TODO: complete javadoc. */

diff --git a/.../src/main/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/FullJoin.java b/.../src/main/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/FullJoin.java
@@ -27,7 +27,10 @@
 import org.apache.beam.sdk.extensions.euphoria.core.client.functional.UnaryFunction;
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.Join.BuilderParams;
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.Join.Type;
+import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeAwareBinaryFunctor;
+import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeAwareUnaryFunction;
 import org.apache.beam.sdk.transforms.windowing.WindowFn;
+import org.apache.beam.sdk.values.TypeDescriptor;
 
 /**
  * Full outer join of two input datasets producing single new dataset.
@@ -115,6 +118,13 @@ public <K> UsingBuilder<LeftT, RightT, K> by(
 
       return new UsingBuilder<>(paramsCasted);
     }
+
+    public <K> UsingBuilder<LeftT, RightT, K> by(
+        UnaryFunction<LeftT, K> leftKeyExtractor, UnaryFunction<RightT, K> rightKeyExtractor,
+        TypeDescriptor<K> keyTypeDescriptor) {
+      return by(TypeAwareUnaryFunction.of(leftKeyExtractor, keyTypeDescriptor),
+          TypeAwareUnaryFunction.of(rightKeyExtractor, keyTypeDescriptor));
+    }
   }
 
   /**
@@ -142,5 +152,23 @@ public <OutputT> Join.WindowingBuilder<LeftT, RightT, K, OutputT> using(
 
       return new Join.WindowingBuilder<>(paramsCasted);
     }
+
+    public <OutputT> Join.WindowingBuilder<LeftT, RightT, K, OutputT> using(
+        BinaryFunctor<Optional<LeftT>, Optional<RightT>, OutputT> joinFunc,
+        TypeDescriptor<OutputT> outputTypeDescriptor) {
+      Objects.requireNonNull(joinFunc);
+
+      @SuppressWarnings("unchecked")
+      BuilderParams<LeftT, RightT, K, OutputT, ?> paramsCasted =
+          (BuilderParams<LeftT, RightT, K, OutputT, ?>) params;
+
+      paramsCasted.joinFunc =
+          TypeAwareBinaryFunctor.of(
+              (left, right, context) ->
+                  joinFunc.apply(Optional.ofNullable(left), Optional.ofNullable(right), context),
+              outputTypeDescriptor);
+
+      return new Join.WindowingBuilder<>(paramsCasted);
+    }
   }
 }
diff --git a/...core/src/main/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/Join.java b/...core/src/main/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/Join.java
@@ -46,12 +46,15 @@
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.state.StateContext;
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.state.StorageProvider;
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.windowing.WindowingDesc;
+import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeAwareBinaryFunctor;
+import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeAwareUnaryFunction;
 import org.apache.beam.sdk.extensions.euphoria.core.client.util.Either;
 import org.apache.beam.sdk.extensions.euphoria.core.client.util.Pair;
 import org.apache.beam.sdk.extensions.euphoria.core.executor.graph.DAG;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.Trigger;
 import org.apache.beam.sdk.transforms.windowing.WindowFn;
+import org.apache.beam.sdk.values.TypeDescriptor;
 import org.apache.beam.sdk.values.WindowingStrategy;
 
 /**
@@ -307,6 +310,13 @@ public <K> UsingBuilder<LeftT, RightT, K> by(
       paramsCasted.rightKeyExtractor = Objects.requireNonNull(rightKeyExtractor);
       return new UsingBuilder<>(paramsCasted);
     }
+
+    public <K> UsingBuilder<LeftT, RightT, K> by(
+        UnaryFunction<LeftT, K> leftKeyExtractor, UnaryFunction<RightT, K> rightKeyExtractor,
+        TypeDescriptor<K> keyTypeDescriptor) {
+      return by(TypeAwareUnaryFunction.of(leftKeyExtractor, keyTypeDescriptor),
+          TypeAwareUnaryFunction.of(rightKeyExtractor, keyTypeDescriptor));
+    }
   }
 
   /**
@@ -333,6 +343,11 @@ public <OutputT> Join.WindowingBuilder<LeftT, RightT, K, OutputT> using(
       return new Join.WindowingBuilder<>(paramsCasted);
     }
 
+    public <OutputT> Join.WindowingBuilder<LeftT, RightT, K, OutputT> using(
+        BinaryFunctor<LeftT, RightT, OutputT> joinFunc,
+        TypeDescriptor<OutputT> outputTypeDescriptor) {
+      return using(TypeAwareBinaryFunctor.of(joinFunc, outputTypeDescriptor));
+    }
   }
 
   /**

diff --git a/.../src/main/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/LeftJoin.java b/.../src/main/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/LeftJoin.java
@@ -27,7 +27,10 @@
 import org.apache.beam.sdk.extensions.euphoria.core.client.functional.UnaryFunction;
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.Join.BuilderParams;
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.Join.Type;
+import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeAwareBinaryFunctor;
+import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeAwareUnaryFunction;
 import org.apache.beam.sdk.transforms.windowing.WindowFn;
+import org.apache.beam.sdk.values.TypeDescriptor;
 
 /**
  * Left outer join of two input datasets producing single new dataset.
@@ -125,6 +128,13 @@ public <K> UsingBuilder<LeftT, RightT, K> by(
 
       return new UsingBuilder<>(paramsCasted);
     }
+
+    public <K> UsingBuilder<LeftT, RightT, K> by(
+        UnaryFunction<LeftT, K> leftKeyExtractor, UnaryFunction<RightT, K> rightKeyExtractor,
+        TypeDescriptor<K> keyTypeDescriptor) {
+      return by(TypeAwareUnaryFunction.of(leftKeyExtractor, keyTypeDescriptor),
+          TypeAwareUnaryFunction.of(rightKeyExtractor, keyTypeDescriptor));
+    }
   }
 
   /**
@@ -152,5 +162,22 @@ public <OutputT> Join.WindowingBuilder<LeftT, RightT, K, OutputT> using(
 
       return new Join.WindowingBuilder<>(paramsCasted);
     }
+
+    public <OutputT> Join.WindowingBuilder<LeftT, RightT, K, OutputT> using(
+        BinaryFunctor<LeftT, Optional<RightT>, OutputT> joinFunc,
+        TypeDescriptor<OutputT> outputTypeDescriptor) {
+      Objects.requireNonNull(joinFunc);
+
+      @SuppressWarnings("unchecked")
+      BuilderParams<LeftT, RightT, K, OutputT, ?> paramsCasted =
+          (BuilderParams<LeftT, RightT, K, OutputT, ?>) params;
+
+      paramsCasted.joinFunc =
+          TypeAwareBinaryFunctor.of(
+              (left, right, context) -> joinFunc.apply(left, Optional.ofNullable(right), context),
+              outputTypeDescriptor);
+
+      return new Join.WindowingBuilder<>(paramsCasted);
+    }
   }
 }
diff --git a/...c/main/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/ReduceByKey.java b/...c/main/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/ReduceByKey.java
@@ -55,13 +55,13 @@
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.windowing.WindowingDesc;
 import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeAwareReduceFunctor;
 import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeAwareUnaryFunction;
-import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeHint;
 import org.apache.beam.sdk.extensions.euphoria.core.client.util.Pair;
 import org.apache.beam.sdk.extensions.euphoria.core.executor.graph.DAG;
 import org.apache.beam.sdk.extensions.euphoria.core.executor.util.SingleValueContext;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.Trigger;
 import org.apache.beam.sdk.transforms.windowing.WindowFn;
+import org.apache.beam.sdk.values.TypeDescriptor;
 import org.apache.beam.sdk.values.WindowingStrategy;
 
 /**
@@ -260,6 +260,13 @@ default <OutputT> WithSortedValuesBuilder<InputT, K, V, OutputT> reduceBy(
       return reduceBy((Stream<V> in, Collector<OutputT> ctx) -> ctx.collect(reducer.apply(in)));
     }
 
+    default <OutputT> WithSortedValuesBuilder<InputT, K, V, OutputT> reduceBy(
+        ReduceFunction<V, OutputT> reducer, TypeDescriptor<OutputT> outputTypeDescriptor) {
+      return reduceBy(
+          (Stream<V> in, Collector<OutputT> ctx) ->
+              ctx.collect(reducer.apply(in)), outputTypeDescriptor);
+    }
+
     /**
      * Define a function that reduces all values related to one key into one or more result objects.
      * The function is not combinable - i.e. partial results cannot be made up before shuffle. To
@@ -272,6 +279,11 @@ default <OutputT> WithSortedValuesBuilder<InputT, K, V, OutputT> reduceBy(
     <OutputT> WithSortedValuesBuilder<InputT, K, V, OutputT> reduceBy(
         ReduceFunctor<V, OutputT> reducer);
 
+    default <OutputT> WithSortedValuesBuilder<InputT, K, V, OutputT> reduceBy(
+        ReduceFunctor<V, OutputT> reducer, TypeDescriptor<OutputT> outputTypeDescriptor) {
+      return reduceBy(TypeAwareReduceFunctor.of(reducer, outputTypeDescriptor));
+    }
+
     /**
      * Define a function that reduces all values related to one key into one result object. The
      * function is combinable (associative and commutative) so it can be used to compute partial
@@ -285,7 +297,7 @@ default WindowByBuilder<InputT, K, V, V> combineBy(CombinableReduceFunction<V> r
     }
 
     default WindowByBuilder<InputT, K, V, V> combineBy(
-        CombinableReduceFunction<V> reducer, TypeHint<V> typeHint) {
+        CombinableReduceFunction<V> reducer, TypeDescriptor<V> typeHint) {
       return reduceBy(TypeAwareReduceFunctor.of(toReduceFunctor(reducer), typeHint));
     }
   }
@@ -349,15 +361,8 @@ public <K> ValueByReduceByBuilder<InputT, K> keyBy(UnaryFunction<InputT, K> keyE
 
     @Override
     public <K> ValueByReduceByBuilder<InputT, K> keyBy(
-        UnaryFunction<InputT, K> keyExtractor, TypeHint<K> typeHint) {
-
-      @SuppressWarnings("unchecked")
-      BuilderParams<InputT, K, ?, ?, ?> paramsCasted =
-          (BuilderParams<InputT, K, ?, ?, ?>) params;
-
-      paramsCasted.keyExtractor = Objects.requireNonNull(keyExtractor);
-
-      return new ValueByReduceByBuilder<>(paramsCasted);
+        UnaryFunction<InputT, K> keyExtractor, TypeDescriptor<K> typeHint) {
+      return keyBy(TypeAwareUnaryFunction.of(keyExtractor, typeHint));
     }
   }
 
@@ -393,7 +398,7 @@ public <V> ReduceByCombineByBuilder<InputT, K, V> valueBy(
     }
 
     public <V> ReduceByCombineByBuilder<InputT, K, V> valueBy(
-        UnaryFunction<InputT, V> valueExtractor, TypeHint<V> typeHint) {
+        UnaryFunction<InputT, V> valueExtractor, TypeDescriptor<V> typeHint) {
       return valueBy(TypeAwareUnaryFunction.of(valueExtractor, typeHint));
     }
 

diff --git a/...n/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/ReduceStateByKey.java b/...n/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/ReduceStateByKey.java
@@ -39,11 +39,11 @@
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.state.StateMerger;
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.windowing.WindowingDesc;
 import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeAwareUnaryFunction;
-import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeHint;
 import org.apache.beam.sdk.extensions.euphoria.core.client.util.Pair;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.Trigger;
 import org.apache.beam.sdk.transforms.windowing.WindowFn;
+import org.apache.beam.sdk.values.TypeDescriptor;
 import org.apache.beam.sdk.values.WindowingStrategy;
 
 /**
@@ -250,15 +250,14 @@ public <K> ValueByBuilder<InputT, K> keyBy(UnaryFunction<InputT, K> keyExtractor
     }
 
     public <K> ValueByBuilder<InputT, K> keyBy(
-        UnaryFunction<InputT, K> keyExtractor, TypeHint<K> typeHint) {
+        UnaryFunction<InputT, K> keyExtractor, TypeDescriptor<K> typeHint) {
 
-      TypeAwareUnaryFunction<InputT, K> typeAwareKeyExtractor =
-          TypeAwareUnaryFunction.of(Objects.requireNonNull(keyExtractor), typeHint);
-
-      BuilderParams<InputT, K, ?, ?, ?, ?> params = new BuilderParams<>(
-          name, input, typeAwareKeyExtractor);
+      return keyBy(TypeAwareUnaryFunction.of(keyExtractor, typeHint));
+    }
 
-      return new ValueByBuilder<>(params);
+    public <K> ValueByBuilder<InputT, K> keyBy(
+        UnaryFunction<InputT, K> keyExtractor, Class<K> typeHint) {
+      return keyBy(TypeAwareUnaryFunction.of(keyExtractor, TypeDescriptor.of(typeHint)));
     }
   }
 

diff --git a/...src/main/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/RightJoin.java b/...src/main/java/org/apache/beam/sdk/extensions/euphoria/core/client/operator/RightJoin.java
@@ -27,7 +27,10 @@
 import org.apache.beam.sdk.extensions.euphoria.core.client.functional.UnaryFunction;
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.Join.BuilderParams;
 import org.apache.beam.sdk.extensions.euphoria.core.client.operator.Join.Type;
+import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeAwareBinaryFunctor;
+import org.apache.beam.sdk.extensions.euphoria.core.client.type.TypeAwareUnaryFunction;
 import org.apache.beam.sdk.transforms.windowing.WindowFn;
+import org.apache.beam.sdk.values.TypeDescriptor;
 
 /**
  * Right outer join of two input datasets producing single new dataset.
@@ -119,6 +122,13 @@ public <K> UsingBuilder<LeftT, RightT, K> by(
       paramsCasted.rightKeyExtractor = Objects.requireNonNull(rightKeyExtractor);
       return new UsingBuilder<>(paramsCasted);
     }
+
+    public <K> UsingBuilder<LeftT, RightT, K> by(
+        UnaryFunction<LeftT, K> leftKeyExtractor, UnaryFunction<RightT, K> rightKeyExtractor,
+        TypeDescriptor<K> keyTypeDescriptor) {
+      return by(TypeAwareUnaryFunction.of(leftKeyExtractor, keyTypeDescriptor),
+          TypeAwareUnaryFunction.of(rightKeyExtractor, keyTypeDescriptor));
+    }
   }
 
   /**
@@ -128,6 +138,7 @@ public static class UsingBuilder<LeftT, RightT, K> {
 
     private final BuilderParams<LeftT, RightT, K, ?, ?> params;
 
+
     UsingBuilder(BuilderParams<LeftT, RightT, K, ?, ?> params) {
       this.params = params;
     }
@@ -146,5 +157,23 @@ public <OutputT> Join.WindowingBuilder<LeftT, RightT, K, OutputT> using(
 
       return new Join.WindowingBuilder<>(paramsCasted);
     }
+
+    public <OutputT> Join.WindowingBuilder<LeftT, RightT, K, OutputT> using(
+        BinaryFunctor<Optional<LeftT>, RightT, OutputT> joinFunc,
+        TypeDescriptor<OutputT> outputTypeDescriptor) {
+      Objects.requireNonNull(joinFunc);
+
+      @SuppressWarnings("unchecked")
+      BuilderParams<LeftT, RightT, K, OutputT, ?> paramsCasted =
+          (BuilderParams<LeftT, RightT, K, OutputT, ?>) params;
+
+      paramsCasted.joinFunc =
+          TypeAwareBinaryFunctor.of(
+              (left, right, context) -> joinFunc.apply(Optional.ofNullable(left), right, context),
+              outputTypeDescriptor);
+
+      return new Join.WindowingBuilder<>(paramsCasted);
+    }
+
   }
 }