[BEAM-4294] Join operator translator

sdks/java/extensions/euphoria/euphoria-beam/build.gradle

@@ -13,9 +13,10 @@ dependencies {
     shadow 'com.google.code.findbugs:annotations:3.0.1'
     testCompile project(':beam-sdks-java-extensions-euphoria-operator-testkit')
     testCompile project(':beam-sdks-java-extensions-euphoria-testing')
-    testCompile project(':beam-runners-direct-java')
+    //testCompile project(':beam-runners-direct-java')
     testCompile library.java.slf4j_api
     testCompile library.java.hamcrest_core
+    testCompile "org.apache.beam:beam-runners-direct-java:2.4.0"
 
 }
 

sdks/java/extensions/euphoria/euphoria-beam/src/main/java/cz/seznam/euphoria/beam/FlowTranslator.java

@@ -20,6 +20,7 @@
 import cz.seznam.euphoria.core.client.flow.Flow;
 import cz.seznam.euphoria.core.client.io.DataSink;
 import cz.seznam.euphoria.core.client.operator.FlatMap;
+import cz.seznam.euphoria.core.client.operator.Join;
 import cz.seznam.euphoria.core.client.operator.Operator;
 import cz.seznam.euphoria.core.client.operator.ReduceByKey;
 import cz.seznam.euphoria.core.client.operator.ReduceStateByKey;
@@ -52,6 +53,7 @@ class FlowTranslator {
     // extended operators
     translators.put(ReduceByKey.class, new ReduceByKeyTranslator());
     translators.put(ReduceStateByKey.class, new ReduceStateByKeyTranslator());
+    translators.put(Join.class, new JoinTranslator());
   }
 
   static Pipeline toPipeline(

sdks/java/extensions/euphoria/euphoria-beam/src/main/java/cz/seznam/euphoria/beam/JoinTranslator.java

@@ -0,0 +1,125 @@
+package cz.seznam.euphoria.beam;
+
+import cz.seznam.euphoria.beam.common.InputToKvDoFn;
+import cz.seznam.euphoria.beam.io.KryoCoder;
+import cz.seznam.euphoria.beam.join.FullJoinFn;
+import cz.seznam.euphoria.beam.join.InnerJoinFn;
+import cz.seznam.euphoria.beam.join.JoinFn;
+import cz.seznam.euphoria.beam.join.LeftOuterJoinFn;
+import cz.seznam.euphoria.beam.join.RightOuterJoinFn;
+import cz.seznam.euphoria.beam.window.WindowingUtils;
+import cz.seznam.euphoria.core.client.dataset.windowing.Window;
+import cz.seznam.euphoria.core.client.functional.BinaryFunctor;
+import cz.seznam.euphoria.core.client.functional.UnaryFunction;
+import cz.seznam.euphoria.core.client.operator.Join;
+import cz.seznam.euphoria.core.client.util.Pair;
+import java.util.List;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.join.CoGbkResult;
+import org.apache.beam.sdk.transforms.join.CoGroupByKey;
+import org.apache.beam.sdk.transforms.join.KeyedPCollectionTuple;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.TupleTag;
+
+
+/**
+ * {@link OperatorTranslator Translator } for Euphoria {@link Join} operator.
+ */
+public class JoinTranslator implements OperatorTranslator<Join> {
+
+  @Override
+  @SuppressWarnings("unchecked")
+  public PCollection<?> translate(Join operator, BeamExecutorContext context) {
+    return doTranslate(operator, context);
+  }
+
+
+  public <K, LeftT, RightT, OutputT, W extends Window<W>> PCollection<Pair<K, OutputT>>
+  doTranslate(Join<LeftT, RightT, K, OutputT, W> operator, BeamExecutorContext context) {
+
+    Coder<K> keyCoder = context.getCoder(operator.getLeftKeyExtractor());
+
+    // get input data-sets transformed to Pcollections<KV<K,LeftT/RightT>>
+    List<PCollection<Object>> inputs = context.getInputs(operator);
+
+    PCollection<KV<K, LeftT>> leftKvInput = getKVInputCollection(inputs.get(0),
+        operator.getLeftKeyExtractor(),
+        keyCoder, new KryoCoder<>(), "::extract-keys-left");
+
+    PCollection<KV<K, RightT>> rightKvInput = getKVInputCollection(inputs.get(1),
+        operator.getRightKeyExtractor(),
+        keyCoder, new KryoCoder<>(), "::extract-keys-right");
+
+    // and apply the same widowing on input Pcolections since the documentation states:
+    //'all of the PCollections you want to group must use the same
+    // windowing strategy and window sizing'
+    leftKvInput = WindowingUtils.applyWindowingIfSpecified(
+        operator, leftKvInput, context.getAllowedLateness(operator));
+    rightKvInput = WindowingUtils.applyWindowingIfSpecified(
+        operator, rightKvInput, context.getAllowedLateness(operator));
+
+    // GoGroupByKey collections
+    TupleTag<LeftT> leftTag = new TupleTag<>();
+    TupleTag<RightT> rightTag = new TupleTag<>();
+
+    PCollection<KV<K, CoGbkResult>> coGrouped = KeyedPCollectionTuple
+        .of(leftTag, leftKvInput)
+        .and(rightTag, rightKvInput)
+        .apply("::co-group-by-key", CoGroupByKey.create());
+
+    // Join
+    JoinFn<LeftT, RightT, K, OutputT> joinFn = chooseJoinFn(operator, leftTag, rightTag);
+
+    return coGrouped.apply(joinFn.getFnName(), ParDo.of(joinFn));
+  }
+
+  private <K, ValueT> PCollection<KV<K, ValueT>> getKVInputCollection(
+      PCollection<Object> inputPCollection,
+      UnaryFunction<ValueT, K> keyExtractor,
+      Coder<K> keyCoder, Coder<ValueT> valueCoder, String transformName) {
+
+    @SuppressWarnings("unchecked")
+    PCollection<ValueT> typedInput = (PCollection<ValueT>) inputPCollection;
+    typedInput.setCoder(valueCoder);
+
+    PCollection<KV<K, ValueT>> leftKvInput =
+        typedInput.apply(transformName, ParDo.of(new InputToKvDoFn<>(keyExtractor)));
+    leftKvInput.setCoder(KvCoder.of(keyCoder, valueCoder));
+
+    return leftKvInput;
+  }
+
+  private <K, LeftT, RightT, OutputT, W extends Window<W>> JoinFn<LeftT, RightT, K, OutputT>
+  chooseJoinFn(
+      Join<LeftT, RightT, K, OutputT, W> operator, TupleTag<LeftT> leftTag,
+      TupleTag<RightT> rightTag) {
+
+    JoinFn<LeftT, RightT, K, OutputT> joinFn;
+    BinaryFunctor<LeftT, RightT, OutputT> joiner = operator.getJoiner();
+
+    switch (operator.getType()) {
+      case INNER:
+        joinFn = new InnerJoinFn<>(joiner, leftTag, rightTag);
+        break;
+      case LEFT:
+        joinFn = new LeftOuterJoinFn<>(joiner, leftTag, rightTag);
+        break;
+      case RIGHT:
+        joinFn = new RightOuterJoinFn<>(joiner, leftTag, rightTag);
+        break;
+      case FULL:
+        joinFn = new FullJoinFn<>(joiner, leftTag, rightTag);
+        break;
+
+      default:
+        throw new UnsupportedOperationException(String.format(
+            "Cannot translate Euphoria '%s' operator to Beam transformations."
+                + " Given join type '%s' is not supported.",
+            Join.class.getSimpleName(), operator.getType()));
+    }
+    return joinFn;
+  }
+}

sdks/java/extensions/euphoria/euphoria-beam/src/main/java/cz/seznam/euphoria/beam/ReduceByKeyTranslator.java

@@ -15,7 +15,7 @@
  */
 package cz.seznam.euphoria.beam;
 
-import cz.seznam.euphoria.beam.window.BeamWindowFn;
+import cz.seznam.euphoria.beam.window.WindowingUtils;
 import cz.seznam.euphoria.core.client.accumulators.AccumulatorProvider;
 import cz.seznam.euphoria.core.client.dataset.windowing.Window;
 import cz.seznam.euphoria.core.client.functional.ReduceFunctor;
@@ -33,7 +33,6 @@
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.SerializableFunction;
 import org.apache.beam.sdk.transforms.SimpleFunction;
-import org.apache.beam.sdk.transforms.windowing.AfterWatermark;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
 
@@ -54,24 +53,8 @@ class ReduceByKeyTranslator implements OperatorTranslator<ReduceByKey> {
     final Coder<K> keyCoder = context.getCoder(keyExtractor);
     final Coder<V> valueCoder = context.getCoder(valueExtractor);
 
-    final PCollection<InputT> input;
-
-    // ~ apply windowing if specified
-    if (operator.getWindowing() == null) {
-      input = context.getInput(operator);
-    } else {
-      input =
-          context
-              .getInput(operator)
-              .apply(
-                  operator.getName() + "::windowing",
-                  org.apache.beam.sdk.transforms.windowing.Window.into(
-                      BeamWindowFn.wrap(operator.getWindowing()))
-                      // TODO: trigger
-                      .triggering(AfterWatermark.pastEndOfWindow())
-                      .discardingFiredPanes()
-                      .withAllowedLateness(context.getAllowedLateness(operator)));
-    }
+    final PCollection<InputT> input = WindowingUtils.applyWindowingIfSpecified(operator,
+        context.getInput(operator), context.getAllowedLateness(operator));
 
     // ~ create key & value extractor
     final MapElements<InputT, KV<K, V>> extractor =
@@ -120,8 +103,8 @@ public Pair<K, V> apply(KV<K, V> in) {
   private static <InputT, OutputT> SerializableFunction<Iterable<InputT>, InputT> asCombiner(
       ReduceFunctor<InputT, OutputT> reducer) {
 
-    @SuppressWarnings("unchecked")
-    final ReduceFunctor<InputT, InputT> combiner = (ReduceFunctor<InputT, InputT>) reducer;
+    @SuppressWarnings("unchecked") final ReduceFunctor<InputT, InputT> combiner =
+        (ReduceFunctor<InputT, InputT>) reducer;
     final SingleValueCollector<InputT> collector = new SingleValueCollector<>();
     return (Iterable<InputT> input) -> {
       combiner.apply(StreamSupport.stream(input.spliterator(), false), collector);

sdks/java/extensions/euphoria/euphoria-beam/src/main/java/cz/seznam/euphoria/beam/ReduceStateByKeyTranslator.java

@@ -63,6 +63,6 @@ public PCollection<?> translate(ReduceStateByKey operator, BeamExecutorContext c
           .setCoder(new KryoCoder<>());
     }
     */
-    throw new UnsupportedOperationException("Not supported yet");
+    throw new UnsupportedOperationException("ReduceStateByKy is not supported yet.");
   }
 }

sdks/java/extensions/euphoria/euphoria-beam/src/main/java/cz/seznam/euphoria/beam/common/InputToKvDoFn.java

@@ -0,0 +1,26 @@
+package cz.seznam.euphoria.beam.common;
+
+import cz.seznam.euphoria.core.client.functional.UnaryFunction;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.values.KV;
+
+/**
+ * {@link DoFn} which takes input elements and transforms them to {@link KV} using given key
+ * extractor.
+ */
+public class InputToKvDoFn<InputT, K> extends DoFn<InputT, KV<K, InputT>> {
+
+  private final UnaryFunction<InputT, K> keyExtractor;
+
+  public InputToKvDoFn(UnaryFunction<InputT, K> leftKeyExtractor) {
+    this.keyExtractor = leftKeyExtractor;
+  }
+
+  @ProcessElement
+  public void processElement(ProcessContext c) {
+    InputT element = c.element();
+    K key = keyExtractor.apply(element);
+    c.output(KV.of(key, element));
+  }
+
+}

sdks/java/extensions/euphoria/euphoria-beam/src/main/java/cz/seznam/euphoria/beam/common/package-info.java

@@ -0,0 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+/**
+ * A set of commonly used classes enabling some code reuse.
+ */
+package cz.seznam.euphoria.beam.common;

sdks/java/extensions/euphoria/euphoria-beam/src/main/java/cz/seznam/euphoria/beam/join/FullJoinFn.java

@@ -0,0 +1,60 @@
+package cz.seznam.euphoria.beam.join;
+
+import cz.seznam.euphoria.beam.SingleValueCollector;
+import cz.seznam.euphoria.core.client.functional.BinaryFunctor;
+import cz.seznam.euphoria.core.client.util.Pair;
+import org.apache.beam.sdk.transforms.join.CoGbkResult;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+
+/**
+ * Full join implementation of {@link JoinFn}.
+ */
+public class FullJoinFn<LeftT, RightT, K, OutputT> extends JoinFn<LeftT, RightT, K, OutputT> {
+
+  public FullJoinFn(BinaryFunctor<LeftT, RightT, OutputT> joiner, TupleTag<LeftT> leftTag,
+      TupleTag<RightT> rightTag) {
+    super(joiner, leftTag, rightTag);
+  }
+
+  @Override
+  public void processElement(ProcessContext c) {
+
+    KV<K, CoGbkResult> element = c.element();
+    CoGbkResult value = element.getValue();
+    K key = element.getKey();
+
+    Iterable<LeftT> leftSideIter = value.getAll(leftTag);
+    Iterable<RightT> rightSIdeIter = value.getAll(rightTag);
+
+    SingleValueCollector<OutputT> outCollector = new SingleValueCollector<>();
+
+    boolean leftHasValues = leftSideIter.iterator().hasNext();
+    boolean rightHasValues = rightSIdeIter.iterator().hasNext();
+
+    if (leftHasValues && rightHasValues) {
+      for (RightT rightValue : rightSIdeIter) {
+        for (LeftT leftValue : leftSideIter) {
+          joiner.apply(leftValue, rightValue, outCollector);
+          c.output(Pair.of(key, outCollector.get()));
+        }
+      }
+    } else if (leftHasValues && !rightHasValues) {
+      for (LeftT leftValue : leftSideIter) {
+        joiner.apply(leftValue, null, outCollector);
+        c.output(Pair.of(key, outCollector.get()));
+      }
+    } else if (!leftHasValues && rightHasValues) {
+      for (RightT rightValue : rightSIdeIter) {
+        joiner.apply(null, rightValue, outCollector);
+        c.output(Pair.of(key, outCollector.get()));
+      }
+    }
+
+  }
+
+  @Override
+  public String getFnName() {
+    return "::full-join";
+  }
+}

sdks/java/extensions/euphoria/euphoria-beam/src/main/java/cz/seznam/euphoria/beam/join/InnerJoinFn.java

@@ -0,0 +1,47 @@
+package cz.seznam.euphoria.beam.join;
+
+import cz.seznam.euphoria.beam.SingleValueCollector;
+import cz.seznam.euphoria.core.client.functional.BinaryFunctor;
+import cz.seznam.euphoria.core.client.util.Pair;
+import org.apache.beam.sdk.transforms.join.CoGbkResult;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.TupleTag;
+
+/**
+ * Inner join implementation of {@link JoinFn}.
+ */
+public class InnerJoinFn<LeftT, RightT, K, OutputT> extends JoinFn<LeftT, RightT, K, OutputT> {
+
+  public InnerJoinFn(
+      BinaryFunctor<LeftT, RightT, OutputT> functor,
+      TupleTag<LeftT> leftTag,
+      TupleTag<RightT> rightTag) {
+    super(functor, leftTag, rightTag);
+  }
+
+  @Override
+  public void processElement(ProcessContext c) {
+
+    KV<K, CoGbkResult> element = c.element();
+    CoGbkResult value = element.getValue();
+    K key = element.getKey();
+
+    Iterable<LeftT> leftSideIter = value.getAll(leftTag);
+    Iterable<RightT> rightSideIter = value.getAll(rightTag);
+
+    SingleValueCollector<OutputT> outCollector = new SingleValueCollector<>();
+
+    for (LeftT leftItem : leftSideIter) {
+      for (RightT rightItem : rightSideIter) {
+        joiner.apply(leftItem, rightItem, outCollector);
+        c.output(Pair.of(key, outCollector.get()));
+      }
+    }
+  }
+
+  @Override
+  public String getFnName() {
+    return "::inner-join";
+  }
+
+}