beam-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From amits...@apache.org
Subject [05/23] beam git commit: Refactor translators according to new GroupAlsoByWindow implemenation for the Spark runnner.
Date Tue, 28 Feb 2017 22:35:12 GMT
Refactor translators according to new GroupAlsoByWindow implemenation for the Spark runnner.


Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/96abe4f0
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/96abe4f0
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/96abe4f0

Branch: refs/heads/master
Commit: 96abe4f08be12ac10dac39b55b8f8319a227b1ea
Parents: 8c37970
Author: Sela <ansela@paypal.com>
Authored: Fri Feb 17 01:19:23 2017 +0200
Committer: Sela <ansela@paypal.com>
Committed: Wed Mar 1 00:17:59 2017 +0200

----------------------------------------------------------------------
 .../SparkGroupAlsoByWindowViaWindowSet.java     |  14 +-
 .../spark/stateful/SparkStateInternals.java     |   4 +-
 .../spark/stateful/SparkTimerInternals.java     |   6 +-
 .../translation/GroupCombineFunctions.java      | 237 ++++++------------
 .../spark/translation/TransformTranslator.java  | 238 +++++++++++++------
 .../spark/translation/TranslationUtils.java     |  22 +-
 .../streaming/StreamingTransformTranslator.java | 163 ++++---------
 .../translation/streaming/UnboundedDataset.java |  12 +-
 .../beam/runners/spark/util/LateDataUtils.java  |   2 +-
 .../spark/util/UnsupportedSideInputReader.java  |   2 +-
 .../streaming/TrackStreamingSourcesTest.java    |   2 +-
 11 files changed, 314 insertions(+), 388 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/beam/blob/96abe4f0/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkGroupAlsoByWindowViaWindowSet.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkGroupAlsoByWindowViaWindowSet.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkGroupAlsoByWindowViaWindowSet.java
index 7902d7c..2fb4100 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkGroupAlsoByWindowViaWindowSet.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkGroupAlsoByWindowViaWindowSet.java
@@ -17,8 +17,8 @@
  */
 package org.apache.beam.runners.spark.stateful;
 
-import com.google.common.collect.Table;
 import com.google.common.collect.AbstractIterator;
+import com.google.common.collect.Table;
 import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -84,7 +84,8 @@ import scala.runtime.AbstractFunction1;
  * in the following steps.
  */
 public class SparkGroupAlsoByWindowViaWindowSet {
-  private static final Logger LOG = LoggerFactory.getLogger(SparkGroupAlsoByWindowViaWindowSet.class);
+  private static final Logger LOG = LoggerFactory.getLogger(
+      SparkGroupAlsoByWindowViaWindowSet.class);
 
   /**
    * A helper class that is essentially a {@link Serializable} {@link AbstractFunction1}.
@@ -101,7 +102,7 @@ public class SparkGroupAlsoByWindowViaWindowSet {
           final SparkRuntimeContext runtimeContext,
           final List<Integer> sourceIds) {
 
-    Long checkpointDuration =
+    long checkpointDurationMillis =
         runtimeContext.getPipelineOptions().as(SparkPipelineOptions.class)
             .getCheckpointDurationMillis();
 
@@ -271,8 +272,11 @@ public class SparkGroupAlsoByWindowViaWindowSet {
         return scala.collection.JavaConversions.asScalaIterator(outIter);
       }
     }, partitioner, true, JavaSparkContext$.MODULE$.<Tuple2<StateAndTimers,
-        List<WindowedValue<KV<K, Iterable<InputT>>>>>>fakeClassTag())
-            .checkpoint(new Duration(checkpointDuration));
+        List<WindowedValue<KV<K, Iterable<InputT>>>>>>fakeClassTag());
+
+    if (checkpointDurationMillis > 0) {
+      firedStream.checkpoint(new Duration(checkpointDurationMillis));
+    }
 
     // go back to Java now.
     JavaPairDStream<K, Tuple2<StateAndTimers, List<WindowedValue<KV<K, Iterable<InputT>>>>>>

http://git-wip-us.apache.org/repos/asf/beam/blob/96abe4f0/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkStateInternals.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkStateInternals.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkStateInternals.java
index e628d31..93b1f63 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkStateInternals.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkStateInternals.java
@@ -22,11 +22,11 @@ import com.google.common.collect.Table;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
-import org.apache.beam.runners.spark.coders.CoderHelpers;
 import org.apache.beam.runners.core.StateInternals;
 import org.apache.beam.runners.core.StateNamespace;
 import org.apache.beam.runners.core.StateTag;
 import org.apache.beam.runners.core.StateTag.StateBinder;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.InstantCoder;
 import org.apache.beam.sdk.coders.ListCoder;
@@ -399,4 +399,4 @@ class SparkStateInternals<K> implements StateInternals<K> {
       };
     }
   }
-}
\ No newline at end of file
+}

http://git-wip-us.apache.org/repos/asf/beam/blob/96abe4f0/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java
index 65225c5..4072240 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/stateful/SparkTimerInternals.java
@@ -27,9 +27,9 @@ import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import javax.annotation.Nullable;
-import org.apache.beam.runners.spark.util.GlobalWatermarkHolder.SparkWatermarks;
 import org.apache.beam.runners.core.StateNamespace;
 import org.apache.beam.runners.core.TimerInternals;
+import org.apache.beam.runners.spark.util.GlobalWatermarkHolder.SparkWatermarks;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.util.TimeDomain;
 import org.apache.spark.broadcast.Broadcast;
@@ -145,7 +145,7 @@ class SparkTimerInternals implements TimerInternals {
     return inputWatermark;
   }
 
-  /** Advances the watermark - since */
+  /** Advances the watermark. */
   public void advanceWatermark() {
     inputWatermark = highWatermark;
   }
@@ -170,4 +170,4 @@ class SparkTimerInternals implements TimerInternals {
     throw new UnsupportedOperationException("Deleting a timer by ID is not yet supported.");
   }
 
-}
\ No newline at end of file
+}

http://git-wip-us.apache.org/repos/asf/beam/blob/96abe4f0/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/GroupCombineFunctions.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/GroupCombineFunctions.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/GroupCombineFunctions.java
index 8a41b4e..1e879ce 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/GroupCombineFunctions.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/GroupCombineFunctions.java
@@ -18,32 +18,21 @@
 
 package org.apache.beam.runners.spark.translation;
 
-import com.google.common.collect.Lists;
-import java.util.Collections;
-import java.util.Map;
-import org.apache.beam.runners.core.SystemReduceFn;
-import org.apache.beam.runners.spark.aggregators.NamedAggregators;
+import static com.google.common.base.Preconditions.checkArgument;
+
 import org.apache.beam.runners.spark.coders.CoderHelpers;
 import org.apache.beam.runners.spark.util.ByteArray;
-import org.apache.beam.runners.spark.util.SideInputBroadcast;
-import org.apache.beam.sdk.coders.CannotProvideCoderException;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.IterableCoder;
 import org.apache.beam.sdk.coders.KvCoder;
-import org.apache.beam.sdk.transforms.CombineWithContext;
-import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder;
 import org.apache.beam.sdk.util.WindowingStrategy;
 import org.apache.beam.sdk.values.KV;
-import org.apache.beam.sdk.values.TupleTag;
-import org.apache.spark.Accumulator;
 import org.apache.spark.api.java.JavaPairRDD;
 import org.apache.spark.api.java.JavaRDD;
-import org.apache.spark.api.java.JavaSparkContext;
 import org.apache.spark.api.java.function.Function;
 import org.apache.spark.api.java.function.Function2;
-import org.apache.spark.api.java.function.PairFlatMapFunction;
-import scala.Tuple2;
 
 
 
@@ -53,113 +42,71 @@ import scala.Tuple2;
 public class GroupCombineFunctions {
 
   /**
-   * Apply {@link org.apache.beam.sdk.transforms.GroupByKey} to a Spark RDD.
+   * An implementation of
+   * {@link org.apache.beam.runners.core.GroupByKeyViaGroupByKeyOnly.GroupByKeyOnly}
+   * for the Spark runner.
    */
-  public static <K, V, W extends BoundedWindow> JavaRDD<WindowedValue<KV<K,
-      Iterable<V>>>> groupByKey(JavaRDD<WindowedValue<KV<K, V>>> rdd,
-                                Accumulator<NamedAggregators> accum,
-                                KvCoder<K, V> coder,
-                                SparkRuntimeContext runtimeContext,
-                                WindowingStrategy<?, W> windowingStrategy) {
-    //--- coders.
-    final Coder<K> keyCoder = coder.getKeyCoder();
-    final Coder<V> valueCoder = coder.getValueCoder();
-    final WindowedValue.WindowedValueCoder<V> wvCoder = WindowedValue.FullWindowedValueCoder.of(
-        valueCoder, windowingStrategy.getWindowFn().windowCoder());
+  public static <K, V> JavaRDD<WindowedValue<KV<K, Iterable<WindowedValue<V>>>>> groupByKeyOnly(
+      JavaRDD<WindowedValue<KV<K, V>>> rdd,
+      Coder<K> keyCoder,
+      WindowedValueCoder<V> wvCoder) {
 
-    //--- groupByKey.
     // Use coders to convert objects in the PCollection to byte arrays, so they
     // can be transferred over the network for the shuffle.
-    JavaRDD<WindowedValue<KV<K, Iterable<WindowedValue<V>>>>> groupedByKey =
-        rdd.map(new ReifyTimestampsAndWindowsFunction<K, V>())
-            .map(WindowingHelpers.<KV<K, WindowedValue<V>>>unwindowFunction())
-            .mapToPair(TranslationUtils.<K, WindowedValue<V>>toPairFunction())
-            .mapToPair(CoderHelpers.toByteFunction(keyCoder, wvCoder))
-            .groupByKey()
-            .mapToPair(CoderHelpers.fromByteFunctionIterable(keyCoder, wvCoder))
-            // empty windows are OK here, see GroupByKey#evaluateHelper in the SDK
-            .map(TranslationUtils.<K, Iterable<WindowedValue<V>>>fromPairFunction())
-            .map(WindowingHelpers.<KV<K, Iterable<WindowedValue<V>>>>windowFunction());
-
-    //--- now group also by window.
-    // GroupAlsoByWindow currently uses a dummy in-memory StateInternals
-    return groupedByKey.flatMap(
-        new SparkGroupAlsoByWindowViaOutputBufferFn<>(
-            windowingStrategy,
-            new TranslationUtils.InMemoryStateInternalsFactory<K>(),
-            SystemReduceFn.<K, V, W>buffering(valueCoder),
-            runtimeContext,
-            accum));
+    return rdd
+        .map(new ReifyTimestampsAndWindowsFunction<K, V>())
+        .map(WindowingHelpers.<KV<K, WindowedValue<V>>>unwindowFunction())
+        .mapToPair(TranslationUtils.<K, WindowedValue<V>>toPairFunction())
+        .mapToPair(CoderHelpers.toByteFunction(keyCoder, wvCoder))
+        .groupByKey()
+        .mapToPair(CoderHelpers.fromByteFunctionIterable(keyCoder, wvCoder))
+        .map(TranslationUtils.<K, Iterable<WindowedValue<V>>>fromPairFunction())
+        .map(WindowingHelpers.<KV<K, Iterable<WindowedValue<V>>>>windowFunction());
   }
 
   /**
    * Apply a composite {@link org.apache.beam.sdk.transforms.Combine.Globally} transformation.
    */
-  public static <InputT, AccumT, OutputT> JavaRDD<WindowedValue<OutputT>>
-  combineGlobally(JavaRDD<WindowedValue<InputT>> rdd,
-                  final CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn,
-                  final Coder<InputT> iCoder,
-                  final Coder<OutputT> oCoder,
-                  final SparkRuntimeContext runtimeContext,
-                  final WindowingStrategy<?, ?> windowingStrategy,
-                  final Map<TupleTag<?>, KV<WindowingStrategy<?, ?>, SideInputBroadcast<?>>>
-                      sideInputs,
-                  boolean hasDefault) {
-    // handle empty input RDD, which will natively skip the entire execution as Spark will not
-    // run on empty RDDs.
-    if (rdd.isEmpty()) {
-      JavaSparkContext jsc = new JavaSparkContext(rdd.context());
-      if (hasDefault) {
-        OutputT defaultValue = combineFn.defaultValue();
-        return jsc
-            .parallelize(Lists.newArrayList(CoderHelpers.toByteArray(defaultValue, oCoder)))
-            .map(CoderHelpers.fromByteFunction(oCoder))
-            .map(WindowingHelpers.<OutputT>windowFunction());
-      } else {
-        return jsc.emptyRDD();
-      }
-    }
-
-    //--- coders.
-    final Coder<AccumT> aCoder;
-    try {
-      aCoder = combineFn.getAccumulatorCoder(runtimeContext.getCoderRegistry(), iCoder);
-    } catch (CannotProvideCoderException e) {
-      throw new IllegalStateException("Could not determine coder for accumulator", e);
-    }
-    // windowed coders.
+  public static <InputT, AccumT> Iterable<WindowedValue<AccumT>> combineGlobally(
+      JavaRDD<WindowedValue<InputT>> rdd,
+      final SparkGlobalCombineFn<InputT, AccumT, ?> sparkCombineFn,
+      final Coder<InputT> iCoder,
+      final Coder<AccumT> aCoder,
+      final WindowingStrategy<?, ?> windowingStrategy) {
+    checkArgument(!rdd.isEmpty(), "CombineGlobally computation should be skipped for empty RDDs.");
+
+    // coders.
     final WindowedValue.FullWindowedValueCoder<InputT> wviCoder =
         WindowedValue.FullWindowedValueCoder.of(iCoder,
             windowingStrategy.getWindowFn().windowCoder());
     final WindowedValue.FullWindowedValueCoder<AccumT> wvaCoder =
         WindowedValue.FullWindowedValueCoder.of(aCoder,
             windowingStrategy.getWindowFn().windowCoder());
-    final WindowedValue.FullWindowedValueCoder<OutputT> wvoCoder =
-        WindowedValue.FullWindowedValueCoder.of(oCoder,
-            windowingStrategy.getWindowFn().windowCoder());
-
-    final SparkGlobalCombineFn<InputT, AccumT, OutputT> sparkCombineFn =
-        new SparkGlobalCombineFn<>(combineFn, runtimeContext, sideInputs, windowingStrategy);
     final IterableCoder<WindowedValue<AccumT>> iterAccumCoder = IterableCoder.of(wvaCoder);
 
-
     // Use coders to convert objects in the PCollection to byte arrays, so they
     // can be transferred over the network for the shuffle.
-    JavaRDD<byte[]> inRddBytes = rdd.map(CoderHelpers.toByteFunction(wviCoder));
-    /*AccumT*/ byte[] acc = inRddBytes.aggregate(
+    // for readability, we add comments with actual type next to byte[].
+    // to shorten line length, we use:
+    //---- WV: WindowedValue
+    //---- Iterable: Itr
+    //---- AccumT: A
+    //---- InputT: I
+    JavaRDD<byte[]> inputRDDBytes = rdd.map(CoderHelpers.toByteFunction(wviCoder));
+    /*Itr<WV<A>>*/ byte[] accumulatedBytes = inputRDDBytes.aggregate(
         CoderHelpers.toByteArray(sparkCombineFn.zeroValue(), iterAccumCoder),
-        new Function2</*AccumT*/ byte[], /*InputT*/ byte[], /*AccumT*/ byte[]>() {
+        new Function2</*A*/ byte[], /*I*/ byte[], /*A*/ byte[]>() {
           @Override
-          public /*AccumT*/ byte[] call(/*AccumT*/ byte[] ab, /*InputT*/ byte[] ib)
+          public /*Itr<WV<A>>*/ byte[] call(/*Itr<WV<A>>*/ byte[] ab, /*WV<I>*/ byte[] ib)
               throws Exception {
             Iterable<WindowedValue<AccumT>> a = CoderHelpers.fromByteArray(ab, iterAccumCoder);
             WindowedValue<InputT> i = CoderHelpers.fromByteArray(ib, wviCoder);
             return CoderHelpers.toByteArray(sparkCombineFn.seqOp(a, i), iterAccumCoder);
           }
         },
-        new Function2</*AccumT*/ byte[], /*AccumT*/ byte[], /*AccumT*/ byte[]>() {
+        new Function2</*Itr<WV<A>>>*/ byte[], /*Itr<WV<A>>>*/ byte[], /*Itr<WV<A>>>*/ byte[]>() {
           @Override
-          public /*AccumT*/ byte[] call(/*AccumT*/ byte[] a1b, /*AccumT*/ byte[] a2b)
+          public /*Itr<WV<A>>>*/ byte[] call(/*Itr<WV<A>>>*/ byte[] a1b, /*Itr<WV<A>>>*/ byte[] a2b)
               throws Exception {
             Iterable<WindowedValue<AccumT>> a1 = CoderHelpers.fromByteArray(a1b, iterAccumCoder);
             Iterable<WindowedValue<AccumT>> a2 = CoderHelpers.fromByteArray(a2b, iterAccumCoder);
@@ -168,10 +115,7 @@ public class GroupCombineFunctions {
           }
         }
     );
-    Iterable<WindowedValue<OutputT>> output =
-        sparkCombineFn.extractOutput(CoderHelpers.fromByteArray(acc, iterAccumCoder));
-    return new JavaSparkContext(rdd.context()).parallelize(
-        CoderHelpers.toByteArrays(output, wvoCoder)).map(CoderHelpers.fromByteFunction(wvoCoder));
+    return CoderHelpers.fromByteArray(accumulatedBytes, iterAccumCoder);
   }
 
   /**
@@ -183,31 +127,22 @@ public class GroupCombineFunctions {
    * For streaming, this will be called from within a serialized context
    * (DStream's transform callback), so passed arguments need to be Serializable.
    */
-  public static <K, InputT, AccumT, OutputT> JavaRDD<WindowedValue<KV<K, OutputT>>>
-  combinePerKey(JavaRDD<WindowedValue<KV<K, InputT>>> rdd,
-                final CombineWithContext.KeyedCombineFnWithContext<K, InputT, AccumT, OutputT>
-                    combineFn,
-                final KvCoder<K, InputT> inputCoder,
-                final SparkRuntimeContext runtimeContext,
-                final WindowingStrategy<?, ?> windowingStrategy,
-                final Map<TupleTag<?>, KV<WindowingStrategy<?, ?>, SideInputBroadcast<?>>>
-                    sideInputs) {
-    //--- coders.
-    final Coder<K> keyCoder = inputCoder.getKeyCoder();
-    final Coder<InputT> viCoder = inputCoder.getValueCoder();
-    final Coder<AccumT> vaCoder;
-    try {
-      vaCoder = combineFn.getAccumulatorCoder(runtimeContext.getCoderRegistry(), keyCoder, viCoder);
-    } catch (CannotProvideCoderException e) {
-      throw new IllegalStateException("Could not determine coder for accumulator", e);
-    }
-    // windowed coders.
+  public static <K, InputT, AccumT> JavaPairRDD<K, Iterable<WindowedValue<KV<K, AccumT>>>>
+      combinePerKey(
+          JavaRDD<WindowedValue<KV<K, InputT>>> rdd,
+          final SparkKeyedCombineFn<K, InputT, AccumT, ?> sparkCombineFn,
+          final Coder<K> keyCoder,
+          final Coder<InputT> iCoder,
+          final Coder<AccumT> aCoder,
+          final WindowingStrategy<?, ?> windowingStrategy) {
+    // coders.
     final WindowedValue.FullWindowedValueCoder<KV<K, InputT>> wkviCoder =
-        WindowedValue.FullWindowedValueCoder.of(KvCoder.of(keyCoder, viCoder),
+        WindowedValue.FullWindowedValueCoder.of(KvCoder.of(keyCoder, iCoder),
             windowingStrategy.getWindowFn().windowCoder());
     final WindowedValue.FullWindowedValueCoder<KV<K, AccumT>> wkvaCoder =
-        WindowedValue.FullWindowedValueCoder.of(KvCoder.of(keyCoder, vaCoder),
+        WindowedValue.FullWindowedValueCoder.of(KvCoder.of(keyCoder, aCoder),
             windowingStrategy.getWindowFn().windowCoder());
+    final IterableCoder<WindowedValue<KV<K, AccumT>>> iterAccumCoder = IterableCoder.of(wkvaCoder);
 
     // We need to duplicate K as both the key of the JavaPairRDD as well as inside the value,
     // since the functions passed to combineByKey don't receive the associated key of each
@@ -217,53 +152,46 @@ public class GroupCombineFunctions {
     // we won't need to duplicate the keys anymore.
     // Key has to bw windowed in order to group by window as well.
     JavaPairRDD<K, WindowedValue<KV<K, InputT>>> inRddDuplicatedKeyPair =
-        rdd.flatMapToPair(
-            new PairFlatMapFunction<WindowedValue<KV<K, InputT>>, K,
-                WindowedValue<KV<K, InputT>>>() {
-              @Override
-              public Iterable<Tuple2<K, WindowedValue<KV<K, InputT>>>>
-              call(WindowedValue<KV<K, InputT>> wkv) {
-                return Collections.singletonList(new Tuple2<>(wkv.getValue().getKey(), wkv));
-              }
-            });
-
-    final SparkKeyedCombineFn<K, InputT, AccumT, OutputT> sparkCombineFn =
-        new SparkKeyedCombineFn<>(combineFn, runtimeContext, sideInputs, windowingStrategy);
-    final IterableCoder<WindowedValue<KV<K, AccumT>>> iterAccumCoder = IterableCoder.of(wkvaCoder);
+        rdd.mapToPair(TranslationUtils.<K, InputT>toPairByKeyInWindowedValue());
 
     // Use coders to convert objects in the PCollection to byte arrays, so they
     // can be transferred over the network for the shuffle.
+    // for readability, we add comments with actual type next to byte[].
+    // to shorten line length, we use:
+    //---- WV: WindowedValue
+    //---- Iterable: Itr
+    //---- AccumT: A
+    //---- InputT: I
     JavaPairRDD<ByteArray, byte[]> inRddDuplicatedKeyPairBytes = inRddDuplicatedKeyPair
         .mapToPair(CoderHelpers.toByteFunction(keyCoder, wkviCoder));
 
-    // The output of combineByKey will be "AccumT" (accumulator)
-    // types rather than "OutputT" (final output types) since Combine.CombineFn
-    // only provides ways to merge VAs, and no way to merge VOs.
-    JavaPairRDD</*K*/ ByteArray, /*KV<K, AccumT>*/ byte[]> accumulatedBytes =
+    JavaPairRDD</*K*/ ByteArray, /*Itr<WV<KV<K, A>>>*/ byte[]> accumulatedBytes =
         inRddDuplicatedKeyPairBytes.combineByKey(
-        new Function</*KV<K, InputT>*/ byte[], /*KV<K, AccumT>*/ byte[]>() {
+        new Function</*WV<KV<K, I>>*/ byte[], /*Itr<WV<KV<K, A>>>*/ byte[]>() {
           @Override
-          public /*KV<K, AccumT>*/ byte[] call(/*KV<K, InputT>*/ byte[] input) {
+          public /*Itr<WV<KV<K, A>>>*/ byte[] call(/*WV<KV<K, I>>*/ byte[] input) {
             WindowedValue<KV<K, InputT>> wkvi = CoderHelpers.fromByteArray(input, wkviCoder);
             return CoderHelpers.toByteArray(sparkCombineFn.createCombiner(wkvi), iterAccumCoder);
           }
         },
-        new Function2</*KV<K, AccumT>*/ byte[], /*KV<K, InputT>*/ byte[],
-            /*KV<K, AccumT>*/ byte[]>() {
+        new Function2</*Itr<WV<KV<K, A>>>*/ byte[], /*WV<KV<K, I>>*/ byte[],
+            /*Itr<WV<KV<K, A>>>*/ byte[]>() {
           @Override
-          public /*KV<K, AccumT>*/ byte[] call(/*KV<K, AccumT>*/ byte[] acc,
-              /*KV<K, InputT>*/ byte[] input) {
+          public /*Itr<WV<KV<K, A>>>*/ byte[] call(
+              /*Itr<WV<KV<K, A>>>*/ byte[] acc,
+              /*WV<KV<K, I>>*/ byte[] input) {
             Iterable<WindowedValue<KV<K, AccumT>>> wkvas =
                 CoderHelpers.fromByteArray(acc, iterAccumCoder);
             WindowedValue<KV<K, InputT>> wkvi = CoderHelpers.fromByteArray(input, wkviCoder);
             return CoderHelpers.toByteArray(sparkCombineFn.mergeValue(wkvi, wkvas), iterAccumCoder);
           }
         },
-        new Function2</*KV<K, AccumT>*/ byte[], /*KV<K, AccumT>*/ byte[],
-            /*KV<K, AccumT>*/ byte[]>() {
+        new Function2</*Itr<WV<KV<K, A>>>*/ byte[], /*Itr<WV<KV<K, A>>>*/ byte[],
+            /*Itr<WV<KV<K, A>>>*/ byte[]>() {
           @Override
-          public /*KV<K, AccumT>*/ byte[] call(/*KV<K, AccumT>*/ byte[] acc1,
-              /*KV<K, AccumT>*/ byte[] acc2) {
+          public /*Itr<WV<KV<K, A>>>*/ byte[] call(
+              /*Itr<WV<KV<K, A>>>*/ byte[] acc1,
+              /*Itr<WV<KV<K, A>>>*/ byte[] acc2) {
             Iterable<WindowedValue<KV<K, AccumT>>> wkvas1 =
                 CoderHelpers.fromByteArray(acc1, iterAccumCoder);
             Iterable<WindowedValue<KV<K, AccumT>>> wkvas2 =
@@ -273,23 +201,6 @@ public class GroupCombineFunctions {
           }
         });
 
-    JavaPairRDD<K, WindowedValue<OutputT>> extracted = accumulatedBytes
-        .mapToPair(CoderHelpers.fromByteFunction(keyCoder, iterAccumCoder))
-        .flatMapValues(new Function<Iterable<WindowedValue<KV<K, AccumT>>>,
-            Iterable<WindowedValue<OutputT>>>() {
-              @Override
-              public Iterable<WindowedValue<OutputT>> call(
-                  Iterable<WindowedValue<KV<K, AccumT>>> accums) {
-                return sparkCombineFn.extractOutput(accums);
-              }
-            });
-    return extracted.map(TranslationUtils.<K, WindowedValue<OutputT>>fromPairFunction()).map(
-        new Function<KV<K, WindowedValue<OutputT>>, WindowedValue<KV<K, OutputT>>>() {
-          @Override
-          public WindowedValue<KV<K, OutputT>> call(KV<K, WindowedValue<OutputT>> kwvo)
-              throws Exception {
-            return kwvo.getValue().withValue(KV.of(kwvo.getKey(), kwvo.getValue().getValue()));
-          }
-        });
+    return accumulatedBytes.mapToPair(CoderHelpers.fromByteFunction(keyCoder, iterAccumCoder));
   }
 }

http://git-wip-us.apache.org/repos/asf/beam/blob/96abe4f0/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TransformTranslator.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TransformTranslator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TransformTranslator.java
index 14c14dc..a643651 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TransformTranslator.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TransformTranslator.java
@@ -27,6 +27,7 @@ import static org.apache.beam.runners.spark.io.hadoop.ShardNameBuilder.replaceSh
 import static org.apache.beam.runners.spark.translation.TranslationUtils.rejectSplittable;
 import static org.apache.beam.runners.spark.translation.TranslationUtils.rejectStateAndTimers;
 
+import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import java.io.IOException;
 import java.util.Collections;
@@ -35,6 +36,7 @@ import java.util.Map;
 import org.apache.avro.mapred.AvroKey;
 import org.apache.avro.mapreduce.AvroJob;
 import org.apache.avro.mapreduce.AvroKeyInputFormat;
+import org.apache.beam.runners.core.SystemReduceFn;
 import org.apache.beam.runners.spark.aggregators.NamedAggregators;
 import org.apache.beam.runners.spark.aggregators.SparkAggregators;
 import org.apache.beam.runners.spark.coders.CoderHelpers;
@@ -46,6 +48,7 @@ import org.apache.beam.runners.spark.io.hadoop.TemplatedTextOutputFormat;
 import org.apache.beam.runners.spark.metrics.MetricsAccumulator;
 import org.apache.beam.runners.spark.metrics.SparkMetricsContainer;
 import org.apache.beam.runners.spark.util.SideInputBroadcast;
+import org.apache.beam.sdk.coders.CannotProvideCoderException;
 import org.apache.beam.sdk.coders.Coder;
 import org.apache.beam.sdk.coders.KvCoder;
 import org.apache.beam.sdk.coders.StringUtf8Coder;
@@ -63,6 +66,7 @@ import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.View;
 import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
 import org.apache.beam.sdk.transforms.windowing.Window;
+import org.apache.beam.sdk.transforms.windowing.WindowFn;
 import org.apache.beam.sdk.util.CombineFnUtil;
 import org.apache.beam.sdk.util.WindowedValue;
 import org.apache.beam.sdk.util.WindowingStrategy;
@@ -119,101 +123,159 @@ public final class TransformTranslator {
     };
   }
 
-  private static <K, V> TransformEvaluator<GroupByKey<K, V>> groupByKey() {
+  private static <K, V, W extends BoundedWindow> TransformEvaluator<GroupByKey<K, V>> groupByKey() {
     return new TransformEvaluator<GroupByKey<K, V>>() {
       @Override
       public void evaluate(GroupByKey<K, V> transform, EvaluationContext context) {
         @SuppressWarnings("unchecked")
         JavaRDD<WindowedValue<KV<K, V>>> inRDD =
             ((BoundedDataset<KV<K, V>>) context.borrowDataset(transform)).getRDD();
-
         @SuppressWarnings("unchecked")
         final KvCoder<K, V> coder = (KvCoder<K, V>) context.getInput(transform).getCoder();
-
         final Accumulator<NamedAggregators> accum =
             SparkAggregators.getNamedAggregators(context.getSparkContext());
-
-        context.putDataset(
-            transform,
-            new BoundedDataset<>(
-                GroupCombineFunctions.groupByKey(
-                    inRDD,
-                    accum,
-                    coder,
-                    context.getRuntimeContext(),
-                    context.getInput(transform).getWindowingStrategy())));
+        @SuppressWarnings("unchecked")
+        final WindowingStrategy<?, W> windowingStrategy =
+            (WindowingStrategy<?, W>) context.getInput(transform).getWindowingStrategy();
+        @SuppressWarnings("unchecked")
+        final WindowFn<Object, W> windowFn = (WindowFn<Object, W>) windowingStrategy.getWindowFn();
+
+        //--- coders.
+        final Coder<K> keyCoder = coder.getKeyCoder();
+        final WindowedValue.WindowedValueCoder<V> wvCoder =
+            WindowedValue.FullWindowedValueCoder.of(coder.getValueCoder(), windowFn.windowCoder());
+
+        //--- group by key only.
+        JavaRDD<WindowedValue<KV<K, Iterable<WindowedValue<V>>>>> groupedByKey =
+            GroupCombineFunctions.groupByKeyOnly(inRDD, keyCoder, wvCoder);
+
+        //--- now group also by window.
+        // for batch, GroupAlsoByWindow uses an in-memory StateInternals.
+        JavaRDD<WindowedValue<KV<K, Iterable<V>>>> groupedAlsoByWindow = groupedByKey.flatMap(
+            new SparkGroupAlsoByWindowViaOutputBufferFn<>(
+                windowingStrategy,
+                new TranslationUtils.InMemoryStateInternalsFactory<K>(),
+                SystemReduceFn.<K, V, W>buffering(coder.getValueCoder()),
+                context.getRuntimeContext(),
+                accum));
+
+        context.putDataset(transform, new BoundedDataset<>(groupedAlsoByWindow));
       }
     };
   }
 
   private static <K, InputT, OutputT> TransformEvaluator<Combine.GroupedValues<K, InputT, OutputT>>
-  combineGrouped() {
-    return new TransformEvaluator<Combine.GroupedValues<K, InputT, OutputT>>() {
-      @Override
-      public void evaluate(Combine.GroupedValues<K, InputT, OutputT> transform,
-                           EvaluationContext context) {
-        // get the applied combine function.
-        PCollection<? extends KV<K, ? extends Iterable<InputT>>> input =
-            context.getInput(transform);
-        WindowingStrategy<?, ?> windowingStrategy = input.getWindowingStrategy();
-        @SuppressWarnings("unchecked")
-        CombineWithContext.KeyedCombineFnWithContext<K, InputT, ?, OutputT> fn =
-            (CombineWithContext.KeyedCombineFnWithContext<K, InputT, ?, OutputT>)
-                CombineFnUtil.toFnWithContext(transform.getFn());
-
-        @SuppressWarnings("unchecked")
-        JavaRDD<WindowedValue<KV<K, Iterable<InputT>>>> inRDD =
-            ((BoundedDataset<KV<K, Iterable<InputT>>>)
-                context.borrowDataset(transform)).getRDD();
-
-        SparkKeyedCombineFn<K, InputT, ?, OutputT> combineFnWithContext =
-            new SparkKeyedCombineFn<>(fn, context.getRuntimeContext(),
-                TranslationUtils.getSideInputs(transform.getSideInputs(), context),
-                windowingStrategy);
-        context.putDataset(transform, new BoundedDataset<>(inRDD.map(new TranslationUtils
-            .CombineGroupedValues<>(
-            combineFnWithContext))));
-      }
-    };
+      combineGrouped() {
+          return new TransformEvaluator<Combine.GroupedValues<K, InputT, OutputT>>() {
+            @Override
+            public void evaluate(
+                Combine.GroupedValues<K, InputT, OutputT> transform,
+                EvaluationContext context) {
+              @SuppressWarnings("unchecked")
+              CombineWithContext.KeyedCombineFnWithContext<K, InputT, ?, OutputT> combineFn =
+                  (CombineWithContext.KeyedCombineFnWithContext<K, InputT, ?, OutputT>)
+                      CombineFnUtil.toFnWithContext(transform.getFn());
+              final SparkKeyedCombineFn<K, InputT, ?, OutputT> sparkCombineFn =
+                  new SparkKeyedCombineFn<>(combineFn, context.getRuntimeContext(),
+                      TranslationUtils.getSideInputs(transform.getSideInputs(), context),
+                          context.getInput(transform).getWindowingStrategy());
+
+              @SuppressWarnings("unchecked")
+              JavaRDD<WindowedValue<KV<K, Iterable<InputT>>>> inRDD =
+                  ((BoundedDataset<KV<K, Iterable<InputT>>>) context.borrowDataset(transform))
+                      .getRDD();
+
+              JavaRDD<WindowedValue<KV<K, OutputT>>> outRDD = inRDD.map(
+                   new Function<WindowedValue<KV<K, Iterable<InputT>>>,
+                       WindowedValue<KV<K, OutputT>>>() {
+                         @Override
+                         public WindowedValue<KV<K, OutputT>> call(
+                             WindowedValue<KV<K, Iterable<InputT>>> in) throws Exception {
+                               return WindowedValue.of(
+                                   KV.of(in.getValue().getKey(), sparkCombineFn.apply(in)),
+                                   in.getTimestamp(),
+                                   in.getWindows(),
+                                   in.getPane());
+                             }
+                       });
+               context.putDataset(transform, new BoundedDataset<>(outRDD));
+            }
+          };
   }
 
   private static <InputT, AccumT, OutputT> TransformEvaluator<Combine.Globally<InputT, OutputT>>
-  combineGlobally() {
-    return new TransformEvaluator<Combine.Globally<InputT, OutputT>>() {
-
-      @Override
-      public void evaluate(Combine.Globally<InputT, OutputT> transform, EvaluationContext context) {
-        final PCollection<InputT> input = context.getInput(transform);
-        // serializable arguments to pass.
-        final Coder<InputT> iCoder = context.getInput(transform).getCoder();
-        final Coder<OutputT> oCoder = context.getOutput(transform).getCoder();
-        @SuppressWarnings("unchecked")
-        final CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn =
-            (CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT>)
-                CombineFnUtil.toFnWithContext(transform.getFn());
-        final WindowingStrategy<?, ?> windowingStrategy = input.getWindowingStrategy();
-        final SparkRuntimeContext runtimeContext = context.getRuntimeContext();
-        final Map<TupleTag<?>, KV<WindowingStrategy<?, ?>, SideInputBroadcast<?>>> sideInputs =
-            TranslationUtils.getSideInputs(transform.getSideInputs(), context);
-        final boolean hasDefault = transform.isInsertDefault();
+      combineGlobally() {
+        return new TransformEvaluator<Combine.Globally<InputT, OutputT>>() {
 
-        @SuppressWarnings("unchecked")
-        JavaRDD<WindowedValue<InputT>> inRdd =
-            ((BoundedDataset<InputT>) context.borrowDataset(transform)).getRDD();
-
-        context.putDataset(transform, new BoundedDataset<>(GroupCombineFunctions
-            .combineGlobally(inRdd, combineFn,
-                iCoder, oCoder, runtimeContext, windowingStrategy, sideInputs, hasDefault)));
-      }
-    };
+          @Override
+          public void evaluate(
+              Combine.Globally<InputT, OutputT> transform,
+              EvaluationContext context) {
+            final PCollection<InputT> input = context.getInput(transform);
+            final Coder<InputT> iCoder = context.getInput(transform).getCoder();
+            final Coder<OutputT> oCoder = context.getOutput(transform).getCoder();
+            final WindowingStrategy<?, ?> windowingStrategy = input.getWindowingStrategy();
+            @SuppressWarnings("unchecked")
+            final CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn =
+                (CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT>)
+                    CombineFnUtil.toFnWithContext(transform.getFn());
+            final WindowedValue.FullWindowedValueCoder<OutputT> wvoCoder =
+                WindowedValue.FullWindowedValueCoder.of(oCoder,
+                    windowingStrategy.getWindowFn().windowCoder());
+            final SparkRuntimeContext runtimeContext = context.getRuntimeContext();
+            final boolean hasDefault = transform.isInsertDefault();
+
+            final SparkGlobalCombineFn<InputT, AccumT, OutputT> sparkCombineFn =
+                new SparkGlobalCombineFn<>(
+                    combineFn,
+                    runtimeContext,
+                    TranslationUtils.getSideInputs(transform.getSideInputs(), context),
+                    windowingStrategy);
+            final Coder<AccumT> aCoder;
+            try {
+              aCoder = combineFn.getAccumulatorCoder(runtimeContext.getCoderRegistry(), iCoder);
+            } catch (CannotProvideCoderException e) {
+              throw new IllegalStateException("Could not determine coder for accumulator", e);
+            }
+
+            @SuppressWarnings("unchecked")
+            JavaRDD<WindowedValue<InputT>> inRdd =
+                ((BoundedDataset<InputT>) context.borrowDataset(transform)).getRDD();
+
+            JavaRDD<WindowedValue<OutputT>> outRdd;
+            // handle empty input RDD, which will naturally skip the entire execution
+            // as Spark will not run on empty RDDs.
+            if (inRdd.isEmpty()) {
+              JavaSparkContext jsc = new JavaSparkContext(inRdd.context());
+              if (hasDefault) {
+                OutputT defaultValue = combineFn.defaultValue();
+                outRdd = jsc
+                    .parallelize(Lists.newArrayList(CoderHelpers.toByteArray(defaultValue, oCoder)))
+                    .map(CoderHelpers.fromByteFunction(oCoder))
+                    .map(WindowingHelpers.<OutputT>windowFunction());
+              } else {
+                outRdd = jsc.emptyRDD();
+              }
+            } else {
+              Iterable<WindowedValue<AccumT>> accumulated = GroupCombineFunctions.combineGlobally(
+                  inRdd, sparkCombineFn, iCoder, aCoder, windowingStrategy);
+              Iterable<WindowedValue<OutputT>> output = sparkCombineFn.extractOutput(accumulated);
+              outRdd = context.getSparkContext()
+                  .parallelize(CoderHelpers.toByteArrays(output, wvoCoder))
+                  .map(CoderHelpers.fromByteFunction(wvoCoder));
+            }
+            context.putDataset(transform, new BoundedDataset<>(outRdd));
+          }
+        };
   }
 
   private static <K, InputT, AccumT, OutputT>
   TransformEvaluator<Combine.PerKey<K, InputT, OutputT>> combinePerKey() {
     return new TransformEvaluator<Combine.PerKey<K, InputT, OutputT>>() {
       @Override
-      public void evaluate(Combine.PerKey<K, InputT, OutputT> transform,
-                           EvaluationContext context) {
+      public void evaluate(
+          Combine.PerKey<K, InputT, OutputT> transform,
+          EvaluationContext context) {
         final PCollection<KV<K, InputT>> input = context.getInput(transform);
         // serializable arguments to pass.
         @SuppressWarnings("unchecked")
@@ -227,14 +289,44 @@ public final class TransformTranslator {
         final SparkRuntimeContext runtimeContext = context.getRuntimeContext();
         final Map<TupleTag<?>, KV<WindowingStrategy<?, ?>, SideInputBroadcast<?>>> sideInputs =
             TranslationUtils.getSideInputs(transform.getSideInputs(), context);
+        final SparkKeyedCombineFn<K, InputT, AccumT, OutputT> sparkCombineFn =
+            new SparkKeyedCombineFn<>(combineFn, runtimeContext, sideInputs, windowingStrategy);
+        final Coder<AccumT> vaCoder;
+        try {
+          vaCoder = combineFn.getAccumulatorCoder(runtimeContext.getCoderRegistry(),
+              inputCoder.getKeyCoder(), inputCoder.getValueCoder());
+        } catch (CannotProvideCoderException e) {
+          throw new IllegalStateException("Could not determine coder for accumulator", e);
+        }
 
         @SuppressWarnings("unchecked")
         JavaRDD<WindowedValue<KV<K, InputT>>> inRdd =
             ((BoundedDataset<KV<K, InputT>>) context.borrowDataset(transform)).getRDD();
 
-        context.putDataset(transform, new BoundedDataset<>(GroupCombineFunctions
-            .combinePerKey(inRdd, combineFn,
-                inputCoder, runtimeContext, windowingStrategy, sideInputs)));
+        JavaPairRDD<K, Iterable<WindowedValue<KV<K, AccumT>>>> accumulatePerKey =
+            GroupCombineFunctions.combinePerKey(inRdd, sparkCombineFn, inputCoder.getKeyCoder(),
+                inputCoder.getValueCoder(), vaCoder, windowingStrategy);
+
+        JavaRDD<WindowedValue<KV<K, OutputT>>> outRdd =
+            accumulatePerKey.flatMapValues(new Function<Iterable<WindowedValue<KV<K, AccumT>>>,
+                Iterable<WindowedValue<OutputT>>>() {
+                  @Override
+                  public Iterable<WindowedValue<OutputT>> call(
+                      Iterable<WindowedValue<KV<K, AccumT>>> iter) throws Exception {
+                        return sparkCombineFn.extractOutput(iter);
+                      }
+                }).map(TranslationUtils.<K, WindowedValue<OutputT>>fromPairFunction())
+                  .map(new Function<KV<K, WindowedValue<OutputT>>,
+                      WindowedValue<KV<K, OutputT>>>() {
+                        @Override
+                          public WindowedValue<KV<K, OutputT>> call(
+                              KV<K, WindowedValue<OutputT>> kv) throws Exception {
+                                WindowedValue<OutputT> wv = kv.getValue();
+                                return wv.withValue(KV.of(kv.getKey(), wv.getValue()));
+                              }
+                      });
+
+        context.putDataset(transform, new BoundedDataset<>(outRdd));
       }
     };
   }

http://git-wip-us.apache.org/repos/asf/beam/blob/96abe4f0/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TranslationUtils.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TranslationUtils.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TranslationUtils.java
index 7d83230..6b27436 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TranslationUtils.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TranslationUtils.java
@@ -21,7 +21,6 @@ package org.apache.beam.runners.spark.translation;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Maps;
 import java.io.Serializable;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import org.apache.beam.runners.core.InMemoryStateInternals;
@@ -42,11 +41,11 @@ import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollectionView;
 import org.apache.beam.sdk.values.TupleTag;
 import org.apache.spark.api.java.JavaSparkContext;
-import org.apache.spark.api.java.function.FlatMapFunction;
 import org.apache.spark.api.java.function.Function;
 import org.apache.spark.api.java.function.PairFunction;
 import org.apache.spark.streaming.api.java.JavaDStream;
 import org.apache.spark.streaming.api.java.JavaPairDStream;
+
 import scala.Tuple2;
 
 /**
@@ -148,20 +147,17 @@ public final class TranslationUtils {
     };
   }
 
-  /** A Flatmap iterator function, flattening iterators into their elements. */
-  public static <T> FlatMapFunction<Iterator<T>, T> flattenIter() {
-    return new FlatMapFunction<Iterator<T>, T>() {
-      @Override
-      public Iterable<T> call(final Iterator<T> t) throws Exception {
-        return new Iterable<T>() {
+  /** Extract key from a {@link WindowedValue} {@link KV} into a pair. */
+  public static <K, V> PairFunction<WindowedValue<KV<K, V>>, K, WindowedValue<KV<K, V>>>
+      toPairByKeyInWindowedValue() {
+        return new PairFunction<WindowedValue<KV<K, V>>, K, WindowedValue<KV<K, V>>>() {
           @Override
-          public Iterator<T> iterator() {
-            return t;
-          }
+          public Tuple2<K, WindowedValue<KV<K, V>>> call(
+              WindowedValue<KV<K, V>> windowedKv) throws Exception {
+                return new Tuple2<>(windowedKv.getValue().getKey(), windowedKv);
+              }
         };
       }
-    };
-  }
 
   /**
    * A utility class to filter {@link TupleTag}s.

http://git-wip-us.apache.org/repos/asf/beam/blob/96abe4f0/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingTransformTranslator.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingTransformTranslator.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingTransformTranslator.java
index 9451df7..e90b490 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingTransformTranslator.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingTransformTranslator.java
@@ -31,6 +31,7 @@ import org.apache.beam.runners.spark.aggregators.SparkAggregators;
 import org.apache.beam.runners.spark.io.ConsoleIO;
 import org.apache.beam.runners.spark.io.CreateStream;
 import org.apache.beam.runners.spark.io.SparkUnboundedSource;
+import org.apache.beam.runners.spark.stateful.SparkGroupAlsoByWindowViaWindowSet;
 import org.apache.beam.runners.spark.metrics.MetricsAccumulator;
 import org.apache.beam.runners.spark.metrics.SparkMetricsContainer;
 import org.apache.beam.runners.spark.translation.BoundedDataset;
@@ -148,7 +149,7 @@ final class StreamingTransformTranslator {
           Dataset dataset = context.borrowDataset(pcol);
           if (dataset instanceof UnboundedDataset) {
             UnboundedDataset<T> unboundedDataset = (UnboundedDataset<T>) dataset;
-            streamingSources.addAll(unboundedDataset.getStreamingSources());
+            streamingSources.addAll(unboundedDataset.getStreamSources());
             dStreams.add(unboundedDataset.getDStream());
           } else {
             rdds.add(((BoundedDataset<T>) dataset).getRDD());
@@ -205,7 +206,7 @@ final class StreamingTransformTranslator {
         //--- then we apply windowing to the elements
         if (TranslationUtils.skipAssignWindows(transform, context)) {
           context.putDataset(transform,
-              new UnboundedDataset<>(windowedDStream, unboundedDataset.getStreamingSources()));
+              new UnboundedDataset<>(windowedDStream, unboundedDataset.getStreamSources()));
         } else {
           JavaDStream<WindowedValue<T>> outStream = windowedDStream.transform(
               new Function<JavaRDD<WindowedValue<T>>, JavaRDD<WindowedValue<T>>>() {
@@ -215,42 +216,55 @@ final class StreamingTransformTranslator {
             }
           });
           context.putDataset(transform,
-              new UnboundedDataset<>(outStream, unboundedDataset.getStreamingSources()));
+              new UnboundedDataset<>(outStream, unboundedDataset.getStreamSources()));
         }
       }
     };
   }
 
-  private static <K, V> TransformEvaluator<GroupByKey<K, V>> groupByKey() {
+  private static <K, V, W extends BoundedWindow> TransformEvaluator<GroupByKey<K, V>> groupByKey() {
     return new TransformEvaluator<GroupByKey<K, V>>() {
       @Override
       public void evaluate(GroupByKey<K, V> transform, EvaluationContext context) {
-        @SuppressWarnings("unchecked")
-        UnboundedDataset<KV<K, V>> unboundedDataset =
-            ((UnboundedDataset<KV<K, V>>) context.borrowDataset(transform));
-        JavaDStream<WindowedValue<KV<K, V>>> dStream = unboundedDataset.getDStream();
-
+        @SuppressWarnings("unchecked") UnboundedDataset<KV<K, V>> inputDataset =
+            (UnboundedDataset<KV<K, V>>) context.borrowDataset(transform);
+        List<Integer> streamSources = inputDataset.getStreamSources();
+        JavaDStream<WindowedValue<KV<K, V>>> dStream = inputDataset.getDStream();
         @SuppressWarnings("unchecked")
         final KvCoder<K, V> coder = (KvCoder<K, V>) context.getInput(transform).getCoder();
-
         final SparkRuntimeContext runtimeContext = context.getRuntimeContext();
-        final WindowingStrategy<?, ?> windowingStrategy =
-            context.getInput(transform).getWindowingStrategy();
+        @SuppressWarnings("unchecked")
+        final WindowingStrategy<?, W> windowingStrategy =
+            (WindowingStrategy<?, W>) context.getInput(transform).getWindowingStrategy();
+        @SuppressWarnings("unchecked")
+        final WindowFn<Object, W> windowFn = (WindowFn<Object, W>) windowingStrategy.getWindowFn();
 
-        JavaDStream<WindowedValue<KV<K, Iterable<V>>>> outStream =
+        //--- coders.
+        final WindowedValue.WindowedValueCoder<V> wvCoder =
+            WindowedValue.FullWindowedValueCoder.of(coder.getValueCoder(), windowFn.windowCoder());
+
+        //--- group by key only.
+        JavaDStream<WindowedValue<KV<K, Iterable<WindowedValue<V>>>>> groupedByKeyStream =
             dStream.transform(new Function<JavaRDD<WindowedValue<KV<K, V>>>,
-                JavaRDD<WindowedValue<KV<K, Iterable<V>>>>>() {
-          @Override
-          public JavaRDD<WindowedValue<KV<K, Iterable<V>>>> call(
-              JavaRDD<WindowedValue<KV<K, V>>> rdd) throws Exception {
-            final Accumulator<NamedAggregators> accum =
-                SparkAggregators.getNamedAggregators(new JavaSparkContext(rdd.context()));
-            return GroupCombineFunctions.groupByKey(rdd, accum, coder, runtimeContext,
-                windowingStrategy);
-          }
-        });
-        context.putDataset(transform,
-            new UnboundedDataset<>(outStream, unboundedDataset.getStreamingSources()));
+                JavaRDD<WindowedValue<KV<K, Iterable<WindowedValue<V>>>>>>() {
+                  @Override
+                  public JavaRDD<WindowedValue<KV<K, Iterable<WindowedValue<V>>>>> call(
+                      JavaRDD<WindowedValue<KV<K, V>>> rdd) throws Exception {
+                        return GroupCombineFunctions.groupByKeyOnly(
+                            rdd, coder.getKeyCoder(), wvCoder);
+                      }
+                });
+
+        //--- now group also by window.
+        JavaDStream<WindowedValue<KV<K, Iterable<V>>>> outStream =
+            SparkGroupAlsoByWindowViaWindowSet.groupAlsoByWindow(
+                groupedByKeyStream,
+                coder.getValueCoder(),
+                windowingStrategy,
+                runtimeContext,
+                streamSources);
+
+        context.putDataset(transform, new UnboundedDataset<>(outStream, streamSources));
       }
     };
   }
@@ -296,96 +310,7 @@ final class StreamingTransformTranslator {
                 });
 
         context.putDataset(transform,
-            new UnboundedDataset<>(outStream, unboundedDataset.getStreamingSources()));
-      }
-    };
-  }
-
-  private static <InputT, AccumT, OutputT> TransformEvaluator<Combine.Globally<InputT, OutputT>>
-  combineGlobally() {
-    return new TransformEvaluator<Combine.Globally<InputT, OutputT>>() {
-
-      @Override
-      public void evaluate(
-          final Combine.Globally<InputT, OutputT> transform,
-          EvaluationContext context) {
-        final PCollection<InputT> input = context.getInput(transform);
-        // serializable arguments to pass.
-        final Coder<InputT> iCoder = context.getInput(transform).getCoder();
-        final Coder<OutputT> oCoder = context.getOutput(transform).getCoder();
-        @SuppressWarnings("unchecked")
-        final CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn =
-            (CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT>)
-                CombineFnUtil.toFnWithContext(transform.getFn());
-        final WindowingStrategy<?, ?> windowingStrategy = input.getWindowingStrategy();
-        final SparkRuntimeContext runtimeContext = context.getRuntimeContext();
-        final boolean hasDefault = transform.isInsertDefault();
-        final SparkPCollectionView pviews = context.getPViews();
-
-        @SuppressWarnings("unchecked")
-        UnboundedDataset<InputT> unboundedDataset =
-            ((UnboundedDataset<InputT>) context.borrowDataset(transform));
-        JavaDStream<WindowedValue<InputT>> dStream = unboundedDataset.getDStream();
-
-        JavaDStream<WindowedValue<OutputT>> outStream = dStream.transform(
-            new Function<JavaRDD<WindowedValue<InputT>>, JavaRDD<WindowedValue<OutputT>>>() {
-          @Override
-          public JavaRDD<WindowedValue<OutputT>> call(JavaRDD<WindowedValue<InputT>> rdd)
-              throws Exception {
-            final Map<TupleTag<?>, KV<WindowingStrategy<?, ?>, SideInputBroadcast<?>>> sideInputs =
-                TranslationUtils.getSideInputs(transform.getSideInputs(),
-                    JavaSparkContext.fromSparkContext(rdd.context()),
-                    pviews);
-            return GroupCombineFunctions.combineGlobally(rdd, combineFn, iCoder, oCoder,
-                runtimeContext, windowingStrategy, sideInputs, hasDefault);
-          }
-        });
-
-        context.putDataset(transform,
-            new UnboundedDataset<>(outStream, unboundedDataset.getStreamingSources()));
-      }
-    };
-  }
-
-  private static <K, InputT, AccumT, OutputT>
-  TransformEvaluator<Combine.PerKey<K, InputT, OutputT>> combinePerKey() {
-    return new TransformEvaluator<Combine.PerKey<K, InputT, OutputT>>() {
-      @Override
-      public void evaluate(final Combine.PerKey<K, InputT, OutputT> transform,
-                           final EvaluationContext context) {
-        final PCollection<KV<K, InputT>> input = context.getInput(transform);
-        // serializable arguments to pass.
-        final KvCoder<K, InputT> inputCoder =
-            (KvCoder<K, InputT>) context.getInput(transform).getCoder();
-        @SuppressWarnings("unchecked")
-        final CombineWithContext.KeyedCombineFnWithContext<K, InputT, AccumT, OutputT> combineFn =
-            (CombineWithContext.KeyedCombineFnWithContext<K, InputT, AccumT, OutputT>)
-                CombineFnUtil.toFnWithContext(transform.getFn());
-        final WindowingStrategy<?, ?> windowingStrategy = input.getWindowingStrategy();
-        final SparkRuntimeContext runtimeContext = context.getRuntimeContext();
-        final SparkPCollectionView pviews = context.getPViews();
-
-        @SuppressWarnings("unchecked")
-        UnboundedDataset<KV<K, InputT>> unboundedDataset =
-            ((UnboundedDataset<KV<K, InputT>>) context.borrowDataset(transform));
-        JavaDStream<WindowedValue<KV<K, InputT>>> dStream = unboundedDataset.getDStream();
-
-        JavaDStream<WindowedValue<KV<K, OutputT>>> outStream =
-            dStream.transform(new Function<JavaRDD<WindowedValue<KV<K, InputT>>>,
-                JavaRDD<WindowedValue<KV<K, OutputT>>>>() {
-          @Override
-          public JavaRDD<WindowedValue<KV<K, OutputT>>> call(
-              JavaRDD<WindowedValue<KV<K, InputT>>> rdd) throws Exception {
-            final Map<TupleTag<?>, KV<WindowingStrategy<?, ?>, SideInputBroadcast<?>>> sideInputs =
-                TranslationUtils.getSideInputs(transform.getSideInputs(),
-                    JavaSparkContext.fromSparkContext(rdd.context()),
-                    pviews);
-            return GroupCombineFunctions.combinePerKey(rdd, combineFn, inputCoder, runtimeContext,
-                windowingStrategy, sideInputs);
-          }
-        });
-        context.putDataset(transform,
-            new UnboundedDataset<>(outStream, unboundedDataset.getStreamingSources()));
+            new UnboundedDataset<>(outStream, unboundedDataset.getStreamSources()));
       }
     };
   }
@@ -431,7 +356,7 @@ final class StreamingTransformTranslator {
         });
 
         context.putDataset(transform,
-            new UnboundedDataset<>(outStream, unboundedDataset.getStreamingSources()));
+            new UnboundedDataset<>(outStream, unboundedDataset.getStreamSources()));
       }
     };
   }
@@ -486,7 +411,7 @@ final class StreamingTransformTranslator {
               (JavaDStream<WindowedValue<Object>>)
                   (JavaDStream<?>) TranslationUtils.dStreamValues(filtered);
           context.putDataset(e.getValue(),
-              new UnboundedDataset<>(values, unboundedDataset.getStreamingSources()));
+              new UnboundedDataset<>(values, unboundedDataset.getStreamSources()));
         }
       }
     };
@@ -499,8 +424,6 @@ final class StreamingTransformTranslator {
     EVALUATORS.put(Read.Unbounded.class, readUnbounded());
     EVALUATORS.put(GroupByKey.class, groupByKey());
     EVALUATORS.put(Combine.GroupedValues.class, combineGrouped());
-    EVALUATORS.put(Combine.Globally.class, combineGlobally());
-    EVALUATORS.put(Combine.PerKey.class, combinePerKey());
     EVALUATORS.put(ParDo.Bound.class, parDo());
     EVALUATORS.put(ParDo.BoundMulti.class, multiDo());
     EVALUATORS.put(ConsoleIO.Write.Unbound.class, print());
@@ -523,7 +446,7 @@ final class StreamingTransformTranslator {
     @Override
     public boolean hasTranslation(Class<? extends PTransform<?, ?>> clazz) {
       // streaming includes rdd/bounded transformations as well
-      return EVALUATORS.containsKey(clazz) || batchTranslator.hasTranslation(clazz);
+      return EVALUATORS.containsKey(clazz);
     }
 
     @Override

http://git-wip-us.apache.org/repos/asf/beam/blob/96abe4f0/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/UnboundedDataset.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/UnboundedDataset.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/UnboundedDataset.java
index 6f5fa93..8624f41 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/UnboundedDataset.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/UnboundedDataset.java
@@ -56,11 +56,11 @@ public class UnboundedDataset<T> implements Dataset {
   // should be greater > 1 in case of Flatten for example.
   // when using GlobalWatermarkHolder this information helps to take only the relevant watermarks
   // and reason about them accordingly.
-  private final List<Integer> streamingSources = new ArrayList<>();
+  private final List<Integer> streamSources = new ArrayList<>();
 
-  public UnboundedDataset(JavaDStream<WindowedValue<T>> dStream, List<Integer> streamingSources) {
+  public UnboundedDataset(JavaDStream<WindowedValue<T>> dStream, List<Integer> streamSources) {
     this.dStream = dStream;
-    this.streamingSources.addAll(streamingSources);
+    this.streamSources.addAll(streamSources);
   }
 
   public UnboundedDataset(Iterable<Iterable<T>> values, JavaStreamingContext jssc, Coder<T> coder) {
@@ -68,7 +68,7 @@ public class UnboundedDataset<T> implements Dataset {
     this.jssc = jssc;
     this.coder = coder;
     // QueuedStream will have a negative (decreasing) unique id.
-    this.streamingSources.add(queuedStreamIds.decrementAndGet());
+    this.streamSources.add(queuedStreamIds.decrementAndGet());
   }
 
   @VisibleForTesting
@@ -97,8 +97,8 @@ public class UnboundedDataset<T> implements Dataset {
     return dStream;
   }
 
-  public List<Integer> getStreamingSources() {
-    return streamingSources;
+  public List<Integer> getStreamSources() {
+    return streamSources;
   }
 
   public void cache() {

http://git-wip-us.apache.org/repos/asf/beam/blob/96abe4f0/runners/spark/src/main/java/org/apache/beam/runners/spark/util/LateDataUtils.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/LateDataUtils.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/LateDataUtils.java
index 96e6ee5..18689bd 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/LateDataUtils.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/LateDataUtils.java
@@ -89,4 +89,4 @@ public class LateDataUtils {
               }
             });
   }
-}
\ No newline at end of file
+}

http://git-wip-us.apache.org/repos/asf/beam/blob/96abe4f0/runners/spark/src/main/java/org/apache/beam/runners/spark/util/UnsupportedSideInputReader.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/UnsupportedSideInputReader.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/UnsupportedSideInputReader.java
index 6de7e86..96d889d 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/UnsupportedSideInputReader.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/UnsupportedSideInputReader.java
@@ -49,4 +49,4 @@ public class UnsupportedSideInputReader implements SideInputReader {
     throw new UnsupportedOperationException(
         String.format("%s does not support side inputs.", transformName));
   }
-}
\ No newline at end of file
+}

http://git-wip-us.apache.org/repos/asf/beam/blob/96abe4f0/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/TrackStreamingSourcesTest.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/TrackStreamingSourcesTest.java b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/TrackStreamingSourcesTest.java
index fbe5777..8449724 100644
--- a/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/TrackStreamingSourcesTest.java
+++ b/runners/spark/src/test/java/org/apache/beam/runners/spark/translation/streaming/TrackStreamingSourcesTest.java
@@ -157,7 +157,7 @@ public class TrackStreamingSourcesTest {
         ctxt.setCurrentTransform(appliedTransform);
         //noinspection unchecked
         Dataset dataset = ctxt.borrowDataset((PTransform<? extends PValue, ?>) transform);
-        assertSourceIds(((UnboundedDataset<?>) dataset).getStreamingSources());
+        assertSourceIds(((UnboundedDataset<?>) dataset).getStreamSources());
         ctxt.setCurrentTransform(null);
       } else {
         evaluator.visitPrimitiveTransform(node);


Mime
View raw message