You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by am...@apache.org on 2016/09/21 17:25:07 UTC
[2/4] incubator-beam git commit: [BEAM-610] Enable spark's
checkpointing mechanism for driver-failure recovery in streaming.
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/0feb6499/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 8341c6d..1a0511f 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
@@ -19,39 +19,32 @@
package org.apache.beam.runners.spark.translation;
+import static com.google.common.base.Preconditions.checkState;
import static org.apache.beam.runners.spark.io.hadoop.ShardNameBuilder.getOutputDirectory;
import static org.apache.beam.runners.spark.io.hadoop.ShardNameBuilder.getOutputFilePrefix;
import static org.apache.beam.runners.spark.io.hadoop.ShardNameBuilder.getOutputFileTemplate;
import static org.apache.beam.runners.spark.io.hadoop.ShardNameBuilder.replaceShardCount;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import java.io.IOException;
-import java.io.Serializable;
-import java.lang.reflect.Field;
-import java.util.Arrays;
import java.util.Collections;
-import java.util.List;
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.AssignWindowsDoFn;
-import org.apache.beam.runners.core.GroupAlsoByWindowsViaOutputBufferDoFn;
import org.apache.beam.runners.core.GroupByKeyViaGroupByKeyOnly.GroupAlsoByWindow;
import org.apache.beam.runners.core.GroupByKeyViaGroupByKeyOnly.GroupByKeyOnly;
-import org.apache.beam.runners.core.SystemReduceFn;
+import org.apache.beam.runners.spark.aggregators.AccumulatorSingleton;
+import org.apache.beam.runners.spark.aggregators.NamedAggregators;
import org.apache.beam.runners.spark.coders.CoderHelpers;
import org.apache.beam.runners.spark.io.hadoop.HadoopIO;
import org.apache.beam.runners.spark.io.hadoop.ShardNameTemplateHelper;
import org.apache.beam.runners.spark.io.hadoop.TemplatedAvroKeyOutputFormat;
import org.apache.beam.runners.spark.io.hadoop.TemplatedTextOutputFormat;
import org.apache.beam.runners.spark.util.BroadcastHelper;
-import org.apache.beam.runners.spark.util.ByteArray;
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.io.AvroIO;
import org.apache.beam.sdk.io.TextIO;
@@ -63,36 +56,30 @@ import org.apache.beam.sdk.transforms.PTransform;
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.GlobalWindows;
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
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.util.state.InMemoryStateInternals;
-import org.apache.beam.sdk.util.state.StateInternals;
-import org.apache.beam.sdk.util.state.StateInternalsFactory;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
import org.apache.beam.sdk.values.PCollectionTuple;
-import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.io.NullWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
+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.JavaRDDLike;
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 org.apache.spark.api.java.function.PairFunction;
+
import scala.Tuple2;
+
/**
* Supports translation between a Beam transform, and Spark's operations on RDDs.
*/
@@ -101,31 +88,6 @@ public final class TransformTranslator {
private TransformTranslator() {
}
- /**
- * Getter of the field.
- */
- public static class FieldGetter {
- private final Map<String, Field> fields;
-
- public FieldGetter(Class<?> clazz) {
- this.fields = Maps.newHashMap();
- for (Field f : clazz.getDeclaredFields()) {
- f.setAccessible(true);
- this.fields.put(f.getName(), f);
- }
- }
-
- public <T> T get(String fieldname, Object value) {
- try {
- @SuppressWarnings("unchecked")
- T fieldValue = (T) fields.get(fieldname).get(value);
- return fieldValue;
- } catch (IllegalAccessException e) {
- throw new IllegalStateException(e);
- }
- }
- }
-
private static <T> TransformEvaluator<Flatten.FlattenPCollectionList<T>> flattenPColl() {
return new TransformEvaluator<Flatten.FlattenPCollectionList<T>>() {
@SuppressWarnings("unchecked")
@@ -142,28 +104,18 @@ public final class TransformTranslator {
};
}
- private static <K, V> TransformEvaluator<GroupByKeyOnly<K, V>> gbk() {
+ private static <K, V> TransformEvaluator<GroupByKeyOnly<K, V>> gbko() {
return new TransformEvaluator<GroupByKeyOnly<K, V>>() {
@Override
public void evaluate(GroupByKeyOnly<K, V> transform, EvaluationContext context) {
@SuppressWarnings("unchecked")
- JavaRDDLike<WindowedValue<KV<K, V>>, ?> inRDD =
- (JavaRDDLike<WindowedValue<KV<K, V>>, ?>) context.getInputRDD(transform);
+ JavaRDD<WindowedValue<KV<K, V>>> inRDD =
+ (JavaRDD<WindowedValue<KV<K, V>>>) context.getInputRDD(transform);
+
@SuppressWarnings("unchecked")
- KvCoder<K, V> coder = (KvCoder<K, V>) context.getInput(transform).getCoder();
- Coder<K> keyCoder = coder.getKeyCoder();
- Coder<V> valueCoder = coder.getValueCoder();
-
- // Use coders to convert objects in the PCollection to byte arrays, so they
- // can be transferred over the network for the shuffle.
- JavaRDDLike<WindowedValue<KV<K, Iterable<V>>>, ?> outRDD = fromPair(
- toPair(inRDD.map(WindowingHelpers.<KV<K, V>>unwindowFunction()))
- .mapToPair(CoderHelpers.toByteFunction(keyCoder, valueCoder))
- .groupByKey()
- .mapToPair(CoderHelpers.fromByteFunctionIterable(keyCoder, valueCoder)))
- // empty windows are OK here, see GroupByKey#evaluateHelper in the SDK
- .map(WindowingHelpers.<KV<K, Iterable<V>>>windowFunction());
- context.setOutputRDD(transform, outRDD);
+ final KvCoder<K, V> coder = (KvCoder<K, V>) context.getInput(transform).getCoder();
+
+ context.setOutputRDD(transform, GroupCombineFunctions.groupByKeyOnly(inRDD, coder));
}
};
}
@@ -174,81 +126,52 @@ public final class TransformTranslator {
@Override
public void evaluate(GroupAlsoByWindow<K, V> transform, EvaluationContext context) {
@SuppressWarnings("unchecked")
- JavaRDDLike<WindowedValue<KV<K, Iterable<WindowedValue<V>>>>, ?> inRDD =
- (JavaRDDLike<WindowedValue<KV<K, Iterable<WindowedValue<V>>>>, ?>)
+ JavaRDD<WindowedValue<KV<K, Iterable<WindowedValue<V>>>>> inRDD =
+ (JavaRDD<WindowedValue<KV<K, Iterable<WindowedValue<V>>>>>)
context.getInputRDD(transform);
- Coder<KV<K, Iterable<WindowedValue<V>>>> inputCoder =
- context.getInput(transform).getCoder();
- Coder<K> keyCoder = transform.getKeyCoder(inputCoder);
- Coder<V> valueCoder = transform.getValueCoder(inputCoder);
-
@SuppressWarnings("unchecked")
- KvCoder<K, Iterable<WindowedValue<V>>> inputKvCoder =
+ final KvCoder<K, Iterable<WindowedValue<V>>> inputKvCoder =
(KvCoder<K, Iterable<WindowedValue<V>>>) context.getInput(transform).getCoder();
- Coder<Iterable<WindowedValue<V>>> inputValueCoder = inputKvCoder.getValueCoder();
-
- IterableCoder<WindowedValue<V>> inputIterableValueCoder =
- (IterableCoder<WindowedValue<V>>) inputValueCoder;
- Coder<WindowedValue<V>> inputIterableElementCoder = inputIterableValueCoder.getElemCoder();
- WindowedValueCoder<V> inputIterableWindowedValueCoder =
- (WindowedValueCoder<V>) inputIterableElementCoder;
- Coder<V> inputIterableElementValueCoder = inputIterableWindowedValueCoder.getValueCoder();
+ final Accumulator<NamedAggregators> accum =
+ AccumulatorSingleton.getInstance(context.getSparkContext());
- @SuppressWarnings("unchecked")
- WindowingStrategy<?, W> windowingStrategy =
- (WindowingStrategy<?, W>) transform.getWindowingStrategy();
-
- OldDoFn<KV<K, Iterable<WindowedValue<V>>>, KV<K, Iterable<V>>> gabwDoFn =
- new GroupAlsoByWindowsViaOutputBufferDoFn<K, V, Iterable<V>, W>(
- windowingStrategy,
- new InMemoryStateInternalsFactory<K>(),
- SystemReduceFn.<K, V, W>buffering(inputIterableElementValueCoder));
-
- // GroupAlsoByWindow current uses a dummy in-memory StateInternals
- JavaRDDLike<WindowedValue<KV<K, Iterable<V>>>, ?> outRDD =
- inRDD.mapPartitions(
- new DoFnFunction<KV<K, Iterable<WindowedValue<V>>>, KV<K, Iterable<V>>>(
- gabwDoFn, context.getRuntimeContext(), null));
-
- context.setOutputRDD(transform, outRDD);
+ context.setOutputRDD(transform, GroupCombineFunctions.groupAlsoByWindow(inRDD, transform,
+ context.getRuntimeContext(), accum, inputKvCoder));
}
};
}
- private static final FieldGetter GROUPED_FG = new FieldGetter(Combine.GroupedValues.class);
-
private static <K, InputT, OutputT> TransformEvaluator<Combine.GroupedValues<K, InputT, OutputT>>
grouped() {
return new TransformEvaluator<Combine.GroupedValues<K, InputT, OutputT>>() {
@Override
public void evaluate(Combine.GroupedValues<K, InputT, OutputT> transform,
EvaluationContext context) {
- Combine.KeyedCombineFn<K, InputT, ?, OutputT> keyed = GROUPED_FG.get("fn", transform);
@SuppressWarnings("unchecked")
JavaRDDLike<WindowedValue<KV<K, Iterable<InputT>>>, ?> inRDD =
- (JavaRDDLike<WindowedValue<KV<K, Iterable<InputT>>>, ?>) context.getInputRDD(transform);
- context.setOutputRDD(transform,
- inRDD.map(new KVFunction<>(keyed)));
+ (JavaRDDLike<WindowedValue<KV<K, Iterable<InputT>>>, ?>)
+ context.getInputRDD(transform);
+ context.setOutputRDD(transform, inRDD.map(
+ new TranslationUtils.CombineGroupedValues<>(transform)));
}
};
}
- private static final FieldGetter COMBINE_GLOBALLY_FG = new FieldGetter(Combine.Globally.class);
-
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 Combine.CombineFn<InputT, AccumT, OutputT> globally =
- COMBINE_GLOBALLY_FG.get("fn", transform);
+ @SuppressWarnings("unchecked")
+ JavaRDD<WindowedValue<InputT>> inRdd =
+ (JavaRDD<WindowedValue<InputT>>) context.getInputRDD(transform);
@SuppressWarnings("unchecked")
- JavaRDDLike<WindowedValue<InputT>, ?> inRdd =
- (JavaRDDLike<WindowedValue<InputT>, ?>) context.getInputRDD(transform);
+ final Combine.CombineFn<InputT, AccumT, OutputT> globally =
+ (Combine.CombineFn<InputT, AccumT, OutputT>) transform.getFn();
final Coder<InputT> iCoder = context.getInput(transform).getCoder();
final Coder<AccumT> aCoder;
@@ -259,61 +182,26 @@ public final class TransformTranslator {
throw new IllegalStateException("Could not determine coder for accumulator", e);
}
- // 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 = inRdd
- .map(WindowingHelpers.<InputT>unwindowFunction())
- .map(CoderHelpers.toByteFunction(iCoder));
-
- /*AccumT*/ byte[] acc = inRddBytes.aggregate(
- CoderHelpers.toByteArray(globally.createAccumulator(), aCoder),
- new Function2</*AccumT*/ byte[], /*InputT*/ byte[], /*AccumT*/ byte[]>() {
- @Override
- public /*AccumT*/ byte[] call(/*AccumT*/ byte[] ab, /*InputT*/ byte[] ib)
- throws Exception {
- AccumT a = CoderHelpers.fromByteArray(ab, aCoder);
- InputT i = CoderHelpers.fromByteArray(ib, iCoder);
- return CoderHelpers.toByteArray(globally.addInput(a, i), aCoder);
- }
- },
- new Function2</*AccumT*/ byte[], /*AccumT*/ byte[], /*AccumT*/ byte[]>() {
- @Override
- public /*AccumT*/ byte[] call(/*AccumT*/ byte[] a1b, /*AccumT*/ byte[] a2b)
- throws Exception {
- AccumT a1 = CoderHelpers.fromByteArray(a1b, aCoder);
- AccumT a2 = CoderHelpers.fromByteArray(a2b, aCoder);
- // don't use Guava's ImmutableList.of as values may be null
- List<AccumT> accumulators = Collections.unmodifiableList(Arrays.asList(a1, a2));
- AccumT merged = globally.mergeAccumulators(accumulators);
- return CoderHelpers.toByteArray(merged, aCoder);
- }
- }
- );
- OutputT output = globally.extractOutput(CoderHelpers.fromByteArray(acc, aCoder));
-
- Coder<OutputT> coder = context.getOutput(transform).getCoder();
+ final Coder<OutputT> oCoder = context.getOutput(transform).getCoder();
JavaRDD<byte[]> outRdd = context.getSparkContext().parallelize(
// don't use Guava's ImmutableList.of as output may be null
- CoderHelpers.toByteArrays(Collections.singleton(output), coder));
- context.setOutputRDD(transform, outRdd.map(CoderHelpers.fromByteFunction(coder))
+ CoderHelpers.toByteArrays(Collections.singleton(
+ GroupCombineFunctions.combineGlobally(inRdd, globally, iCoder, aCoder)), oCoder));
+ context.setOutputRDD(transform, outRdd.map(CoderHelpers.fromByteFunction(oCoder))
.map(WindowingHelpers.<OutputT>windowFunction()));
}
};
}
- private static final FieldGetter COMBINE_PERKEY_FG = new FieldGetter(Combine.PerKey.class);
-
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) {
- final Combine.KeyedCombineFn<K, InputT, AccumT, OutputT> keyed =
- COMBINE_PERKEY_FG.get("fn", transform);
+ public void evaluate(Combine.PerKey<K, InputT, OutputT> transform,
+ EvaluationContext context) {
@SuppressWarnings("unchecked")
- JavaRDDLike<WindowedValue<KV<K, InputT>>, ?> inRdd =
- (JavaRDDLike<WindowedValue<KV<K, InputT>>, ?>) context.getInputRDD(transform);
+ final Combine.KeyedCombineFn<K, InputT, AccumT, OutputT> keyed =
+ (Combine.KeyedCombineFn<K, InputT, AccumT, OutputT>) transform.getFn();
@SuppressWarnings("unchecked")
KvCoder<K, InputT> inputCoder = (KvCoder<K, InputT>)
@@ -329,214 +217,66 @@ public final class TransformTranslator {
}
Coder<KV<K, InputT>> kviCoder = KvCoder.of(keyCoder, viCoder);
Coder<KV<K, AccumT>> kvaCoder = KvCoder.of(keyCoder, vaCoder);
-
- // 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
- // value, and we need to map back into methods in Combine.KeyedCombineFn, which each
- // require the key in addition to the InputT's and AccumT's being merged/accumulated.
- // Once Spark provides a way to include keys in the arguments of combine/merge functions,
- // we won't need to duplicate the keys anymore.
-
- // Key has to bw windowed in order to group by window as well
- JavaPairRDD<WindowedValue<K>, WindowedValue<KV<K, InputT>>> inRddDuplicatedKeyPair =
- inRdd.flatMapToPair(
- new PairFlatMapFunction<WindowedValue<KV<K, InputT>>, WindowedValue<K>,
- WindowedValue<KV<K, InputT>>>() {
- @Override
- public Iterable<Tuple2<WindowedValue<K>,
- WindowedValue<KV<K, InputT>>>>
- call(WindowedValue<KV<K, InputT>> kv) {
- List<Tuple2<WindowedValue<K>,
- WindowedValue<KV<K, InputT>>>> tuple2s =
- Lists.newArrayListWithCapacity(kv.getWindows().size());
- for (BoundedWindow boundedWindow: kv.getWindows()) {
- WindowedValue<K> wk = WindowedValue.of(kv.getValue().getKey(),
- boundedWindow.maxTimestamp(), boundedWindow, kv.getPane());
- tuple2s.add(new Tuple2<>(wk, kv));
- }
- return tuple2s;
- }
- });
//-- windowed coders
final WindowedValue.FullWindowedValueCoder<K> wkCoder =
- WindowedValue.FullWindowedValueCoder.of(keyCoder,
+ WindowedValue.FullWindowedValueCoder.of(keyCoder,
context.getInput(transform).getWindowingStrategy().getWindowFn().windowCoder());
final WindowedValue.FullWindowedValueCoder<KV<K, InputT>> wkviCoder =
- WindowedValue.FullWindowedValueCoder.of(kviCoder,
+ WindowedValue.FullWindowedValueCoder.of(kviCoder,
context.getInput(transform).getWindowingStrategy().getWindowFn().windowCoder());
final WindowedValue.FullWindowedValueCoder<KV<K, AccumT>> wkvaCoder =
- WindowedValue.FullWindowedValueCoder.of(kvaCoder,
+ WindowedValue.FullWindowedValueCoder.of(kvaCoder,
context.getInput(transform).getWindowingStrategy().getWindowFn().windowCoder());
- // Use coders to convert objects in the PCollection to byte arrays, so they
- // can be transferred over the network for the shuffle.
- JavaPairRDD<ByteArray, byte[]> inRddDuplicatedKeyPairBytes = inRddDuplicatedKeyPair
- .mapToPair(CoderHelpers.toByteFunction(wkCoder, 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 =
- inRddDuplicatedKeyPairBytes.combineByKey(
- new Function</*KV<K, InputT>*/ byte[], /*KV<K, AccumT>*/ byte[]>() {
- @Override
- public /*KV<K, AccumT>*/ byte[] call(/*KV<K, InputT>*/ byte[] input) {
- WindowedValue<KV<K, InputT>> wkvi =
- CoderHelpers.fromByteArray(input, wkviCoder);
- AccumT va = keyed.createAccumulator(wkvi.getValue().getKey());
- va = keyed.addInput(wkvi.getValue().getKey(), va, wkvi.getValue().getValue());
- WindowedValue<KV<K, AccumT>> wkva =
- WindowedValue.of(KV.of(wkvi.getValue().getKey(), va), wkvi.getTimestamp(),
- wkvi.getWindows(), wkvi.getPane());
- return CoderHelpers.toByteArray(wkva, wkvaCoder);
- }
- },
- new Function2</*KV<K, AccumT>*/ byte[],
- /*KV<K, InputT>*/ byte[],
- /*KV<K, AccumT>*/ byte[]>() {
- @Override
- public /*KV<K, AccumT>*/ byte[] call(/*KV<K, AccumT>*/ byte[] acc,
- /*KV<K, InputT>*/ byte[] input) {
- WindowedValue<KV<K, AccumT>> wkva =
- CoderHelpers.fromByteArray(acc, wkvaCoder);
- WindowedValue<KV<K, InputT>> wkvi =
- CoderHelpers.fromByteArray(input, wkviCoder);
- AccumT va =
- keyed.addInput(wkva.getValue().getKey(), wkva.getValue().getValue(),
- wkvi.getValue().getValue());
- wkva = WindowedValue.of(KV.of(wkva.getValue().getKey(), va), wkva.getTimestamp(),
- wkva.getWindows(), wkva.getPane());
- return CoderHelpers.toByteArray(wkva, wkvaCoder);
- }
- },
- new Function2</*KV<K, AccumT>*/ byte[],
- /*KV<K, AccumT>*/ byte[],
- /*KV<K, AccumT>*/ byte[]>() {
- @Override
- public /*KV<K, AccumT>*/ byte[] call(/*KV<K, AccumT>*/ byte[] acc1,
- /*KV<K, AccumT>*/ byte[] acc2) {
- WindowedValue<KV<K, AccumT>> wkva1 =
- CoderHelpers.fromByteArray(acc1, wkvaCoder);
- WindowedValue<KV<K, AccumT>> wkva2 =
- CoderHelpers.fromByteArray(acc2, wkvaCoder);
- AccumT va = keyed.mergeAccumulators(wkva1.getValue().getKey(),
- // don't use Guava's ImmutableList.of as values may be null
- Collections.unmodifiableList(Arrays.asList(wkva1.getValue().getValue(),
- wkva2.getValue().getValue())));
- WindowedValue<KV<K, AccumT>> wkva =
- WindowedValue.of(KV.of(wkva1.getValue().getKey(),
- va), wkva1.getTimestamp(), wkva1.getWindows(), wkva1.getPane());
- return CoderHelpers.toByteArray(wkva, wkvaCoder);
- }
- });
-
- JavaPairRDD<WindowedValue<K>, WindowedValue<OutputT>> extracted = accumulatedBytes
- .mapToPair(CoderHelpers.fromByteFunction(wkCoder, wkvaCoder))
- .mapValues(
- new Function<WindowedValue<KV<K, AccumT>>, WindowedValue<OutputT>>() {
- @Override
- public WindowedValue<OutputT> call(WindowedValue<KV<K, AccumT>> acc) {
- return WindowedValue.of(keyed.extractOutput(acc.getValue().getKey(),
- acc.getValue().getValue()), acc.getTimestamp(),
- acc.getWindows(), acc.getPane());
- }
- });
+ @SuppressWarnings("unchecked")
+ JavaRDD<WindowedValue<KV<K, InputT>>> inRdd =
+ (JavaRDD<WindowedValue<KV<K, InputT>>>) context.getInputRDD(transform);
- context.setOutputRDD(transform,
- fromPair(extracted)
- .map(new Function<KV<WindowedValue<K>, WindowedValue<OutputT>>,
- WindowedValue<KV<K, OutputT>>>() {
- @Override
- public WindowedValue<KV<K, OutputT>> call(KV<WindowedValue<K>,
- WindowedValue<OutputT>> kwvo)
- throws Exception {
- WindowedValue<OutputT> wvo = kwvo.getValue();
- KV<K, OutputT> kvo = KV.of(kwvo.getKey().getValue(), wvo.getValue());
- return WindowedValue.of(kvo, wvo.getTimestamp(), wvo.getWindows(), wvo.getPane());
- }
- }));
+ context.setOutputRDD(transform, GroupCombineFunctions.combinePerKey(inRdd, keyed, wkCoder,
+ wkviCoder, wkvaCoder));
}
};
}
- private static final class KVFunction<K, InputT, OutputT>
- implements Function<WindowedValue<KV<K, Iterable<InputT>>>,
- WindowedValue<KV<K, OutputT>>> {
- private final Combine.KeyedCombineFn<K, InputT, ?, OutputT> keyed;
-
- KVFunction(Combine.KeyedCombineFn<K, InputT, ?, OutputT> keyed) {
- this.keyed = keyed;
- }
-
- @Override
- public WindowedValue<KV<K, OutputT>> call(WindowedValue<KV<K,
- Iterable<InputT>>> windowedKv)
- throws Exception {
- KV<K, Iterable<InputT>> kv = windowedKv.getValue();
- return WindowedValue.of(KV.of(kv.getKey(), keyed.apply(kv.getKey(), kv.getValue())),
- windowedKv.getTimestamp(), windowedKv.getWindows(), windowedKv.getPane());
- }
- }
-
- private static <K, V> JavaPairRDD<K, V> toPair(JavaRDDLike<KV<K, V>, ?> rdd) {
- return rdd.mapToPair(new PairFunction<KV<K, V>, K, V>() {
- @Override
- public Tuple2<K, V> call(KV<K, V> kv) {
- return new Tuple2<>(kv.getKey(), kv.getValue());
- }
- });
- }
-
- private static <K, V> JavaRDDLike<KV<K, V>, ?> fromPair(JavaPairRDD<K, V> rdd) {
- return rdd.map(new Function<Tuple2<K, V>, KV<K, V>>() {
- @Override
- public KV<K, V> call(Tuple2<K, V> t2) {
- return KV.of(t2._1(), t2._2());
- }
- });
- }
-
private static <InputT, OutputT> TransformEvaluator<ParDo.Bound<InputT, OutputT>> parDo() {
return new TransformEvaluator<ParDo.Bound<InputT, OutputT>>() {
@Override
public void evaluate(ParDo.Bound<InputT, OutputT> transform, EvaluationContext context) {
- DoFnFunction<InputT, OutputT> dofn =
- new DoFnFunction<>(transform.getFn(),
- context.getRuntimeContext(),
- getSideInputs(transform.getSideInputs(), context));
@SuppressWarnings("unchecked")
JavaRDDLike<WindowedValue<InputT>, ?> inRDD =
(JavaRDDLike<WindowedValue<InputT>, ?>) context.getInputRDD(transform);
- context.setOutputRDD(transform, inRDD.mapPartitions(dofn));
+ Accumulator<NamedAggregators> accum =
+ AccumulatorSingleton.getInstance(context.getSparkContext());
+ Map<TupleTag<?>, BroadcastHelper<?>> sideInputs =
+ TranslationUtils.getSideInputs(transform.getSideInputs(), context);
+ context.setOutputRDD(transform,
+ inRDD.mapPartitions(new DoFnFunction<>(accum, transform.getFn(),
+ context.getRuntimeContext(), sideInputs)));
}
};
}
- private static final FieldGetter MULTIDO_FG = new FieldGetter(ParDo.BoundMulti.class);
-
- private static <InputT, OutputT> TransformEvaluator<ParDo.BoundMulti<InputT, OutputT>> multiDo() {
+ private static <InputT, OutputT> TransformEvaluator<ParDo.BoundMulti<InputT, OutputT>>
+ multiDo() {
return new TransformEvaluator<ParDo.BoundMulti<InputT, OutputT>>() {
@Override
public void evaluate(ParDo.BoundMulti<InputT, OutputT> transform, EvaluationContext context) {
- TupleTag<OutputT> mainOutputTag = MULTIDO_FG.get("mainOutputTag", transform);
- MultiDoFnFunction<InputT, OutputT> multifn = new MultiDoFnFunction<>(
- transform.getFn(),
- context.getRuntimeContext(),
- mainOutputTag,
- getSideInputs(transform.getSideInputs(), context));
-
@SuppressWarnings("unchecked")
JavaRDDLike<WindowedValue<InputT>, ?> inRDD =
(JavaRDDLike<WindowedValue<InputT>, ?>) context.getInputRDD(transform);
+ Accumulator<NamedAggregators> accum =
+ AccumulatorSingleton.getInstance(context.getSparkContext());
JavaPairRDD<TupleTag<?>, WindowedValue<?>> all = inRDD
- .mapPartitionsToPair(multifn)
- .cache();
-
+ .mapPartitionsToPair(
+ new MultiDoFnFunction<>(accum, transform.getFn(), context.getRuntimeContext(),
+ transform.getMainOutputTag(), TranslationUtils.getSideInputs(
+ transform.getSideInputs(), context)))
+ .cache();
PCollectionTuple pct = context.getOutput(transform);
for (Map.Entry<TupleTag<?>, PCollection<?>> e : pct.getAll().entrySet()) {
@SuppressWarnings("unchecked")
JavaPairRDD<TupleTag<?>, WindowedValue<?>> filtered =
- all.filter(new TupleTagFilter(e.getKey()));
+ all.filter(new TranslationUtils.TupleTagFilter(e.getKey()));
@SuppressWarnings("unchecked")
// Object is the best we can do since different outputs can have different tags
JavaRDD<WindowedValue<Object>> values =
@@ -753,22 +493,17 @@ public final class TransformTranslator {
JavaRDDLike<WindowedValue<T>, ?> inRDD =
(JavaRDDLike<WindowedValue<T>, ?>) context.getInputRDD(transform);
- @SuppressWarnings("unchecked")
- WindowFn<? super T, W> windowFn = (WindowFn<? super T, W>) transform.getWindowFn();
-
- // Avoid running assign windows if both source and destination are global window
- // or if the user has not specified the WindowFn (meaning they are just messing
- // with triggering or allowed lateness)
- if (windowFn == null
- || (context.getInput(transform).getWindowingStrategy().getWindowFn()
- instanceof GlobalWindows
- && windowFn instanceof GlobalWindows)) {
+ if (TranslationUtils.skipAssignWindows(transform, context)) {
context.setOutputRDD(transform, inRDD);
} else {
+ @SuppressWarnings("unchecked")
+ WindowFn<? super T, W> windowFn = (WindowFn<? super T, W>) transform.getWindowFn();
OldDoFn<T, T> addWindowsDoFn = new AssignWindowsDoFn<>(windowFn);
- DoFnFunction<T, T> dofn =
- new DoFnFunction<>(addWindowsDoFn, context.getRuntimeContext(), null);
- context.setOutputRDD(transform, inRDD.mapPartitions(dofn));
+ Accumulator<NamedAggregators> accum =
+ AccumulatorSingleton.getInstance(context.getSparkContext());
+ context.setOutputRDD(transform,
+ inRDD.mapPartitions(new DoFnFunction<>(accum, addWindowsDoFn,
+ context.getRuntimeContext(), null)));
}
}
};
@@ -822,42 +557,6 @@ public final class TransformTranslator {
};
}
- private static final class TupleTagFilter<V>
- implements Function<Tuple2<TupleTag<V>, WindowedValue<?>>, Boolean> {
-
- private final TupleTag<V> tag;
-
- private TupleTagFilter(TupleTag<V> tag) {
- this.tag = tag;
- }
-
- @Override
- public Boolean call(Tuple2<TupleTag<V>, WindowedValue<?>> input) {
- return tag.equals(input._1());
- }
- }
-
- private static Map<TupleTag<?>, BroadcastHelper<?>> getSideInputs(
- List<PCollectionView<?>> views,
- EvaluationContext context) {
- if (views == null) {
- return ImmutableMap.of();
- } else {
- Map<TupleTag<?>, BroadcastHelper<?>> sideInputs = Maps.newHashMap();
- for (PCollectionView<?> view : views) {
- Iterable<? extends WindowedValue<?>> collectionView = context.getPCollectionView(view);
- Coder<Iterable<WindowedValue<?>>> coderInternal = view.getCoderInternal();
- @SuppressWarnings("unchecked")
- BroadcastHelper<?> helper =
- BroadcastHelper.create((Iterable<WindowedValue<?>>) collectionView, coderInternal);
- //broadcast side inputs
- helper.broadcast(context.getSparkContext());
- sideInputs.put(view.getTagInternal(), helper);
- }
- return sideInputs;
- }
- }
-
private static final Map<Class<? extends PTransform>, TransformEvaluator<?>> EVALUATORS = Maps
.newHashMap();
@@ -870,7 +569,7 @@ public final class TransformTranslator {
EVALUATORS.put(HadoopIO.Write.Bound.class, writeHadoop());
EVALUATORS.put(ParDo.Bound.class, parDo());
EVALUATORS.put(ParDo.BoundMulti.class, multiDo());
- EVALUATORS.put(GroupByKeyOnly.class, gbk());
+ EVALUATORS.put(GroupByKeyOnly.class, gbko());
EVALUATORS.put(GroupAlsoByWindow.class, gabw());
EVALUATORS.put(Combine.GroupedValues.class, grouped());
EVALUATORS.put(Combine.Globally.class, combineGlobally());
@@ -883,17 +582,6 @@ public final class TransformTranslator {
EVALUATORS.put(Window.Bound.class, window());
}
- public static <TransformT extends PTransform<?, ?>> TransformEvaluator<TransformT>
- getTransformEvaluator(Class<TransformT> clazz) {
- @SuppressWarnings("unchecked")
- TransformEvaluator<TransformT> transform =
- (TransformEvaluator<TransformT>) EVALUATORS.get(clazz);
- if (transform == null) {
- throw new IllegalStateException("No TransformEvaluator registered for " + clazz);
- }
- return transform;
- }
-
/**
* Translator matches Beam transformation with the appropriate evaluator.
*/
@@ -905,17 +593,20 @@ public final class TransformTranslator {
}
@Override
- public <TransformT extends PTransform<?, ?>> TransformEvaluator<TransformT> translate(
- Class<TransformT> clazz) {
- return getTransformEvaluator(clazz);
+ public <TransformT extends PTransform<?, ?>> TransformEvaluator<TransformT>
+ translateBounded (Class<TransformT> clazz) {
+ @SuppressWarnings("unchecked") TransformEvaluator<TransformT> transformEvaluator =
+ (TransformEvaluator<TransformT>) EVALUATORS.get(clazz);
+ checkState(transformEvaluator != null,
+ "No TransformEvaluator registered for BOUNDED transform %s", clazz);
+ return transformEvaluator;
}
- }
- private static class InMemoryStateInternalsFactory<K> implements StateInternalsFactory<K>,
- Serializable {
@Override
- public StateInternals<K> stateInternalsForKey(K key) {
- return InMemoryStateInternals.forKey(key);
+ public <TransformT extends PTransform<?, ?>> TransformEvaluator<TransformT>
+ translateUnbounded(Class<TransformT> clazz) {
+ throw new IllegalStateException("TransformTranslator used in a batch pipeline only "
+ + "supports BOUNDED transforms.");
}
}
}
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/0feb6499/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
new file mode 100644
index 0000000..9b156fe
--- /dev/null
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/TranslationUtils.java
@@ -0,0 +1,195 @@
+/*
+ * 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.
+ */
+
+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.List;
+import java.util.Map;
+import org.apache.beam.runners.spark.util.BroadcastHelper;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.transforms.Combine;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindows;
+import org.apache.beam.sdk.transforms.windowing.Window;
+import org.apache.beam.sdk.transforms.windowing.WindowFn;
+import org.apache.beam.sdk.util.WindowedValue;
+import org.apache.beam.sdk.util.state.InMemoryStateInternals;
+import org.apache.beam.sdk.util.state.StateInternals;
+import org.apache.beam.sdk.util.state.StateInternalsFactory;
+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.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;
+
+/**
+ * A set of utilities to help translating Beam transformations into Spark transformations.
+ */
+public final class TranslationUtils {
+
+ private TranslationUtils() {
+ }
+
+ /**
+ * In-memory state internals factory.
+ *
+ * @param <K> State key type.
+ */
+ static class InMemoryStateInternalsFactory<K> implements StateInternalsFactory<K>,
+ Serializable {
+ @Override
+ public StateInternals<K> stateInternalsForKey(K key) {
+ return InMemoryStateInternals.forKey(key);
+ }
+ }
+
+ /**
+ * A {@link Combine.GroupedValues} function applied to grouped KVs.
+ *
+ * @param <K> Grouped key type.
+ * @param <InputT> Grouped values type.
+ * @param <OutputT> Output type.
+ */
+ public static class CombineGroupedValues<K, InputT, OutputT> implements
+ Function<WindowedValue<KV<K, Iterable<InputT>>>, WindowedValue<KV<K, OutputT>>> {
+ private final Combine.KeyedCombineFn<K, InputT, ?, OutputT> keyed;
+
+ public CombineGroupedValues(Combine.GroupedValues<K, InputT, OutputT> transform) {
+ //noinspection unchecked
+ keyed = (Combine.KeyedCombineFn<K, InputT, ?, OutputT>) transform.getFn();
+ }
+
+ @Override
+ public WindowedValue<KV<K, OutputT>> call(WindowedValue<KV<K, Iterable<InputT>>> windowedKv)
+ throws Exception {
+ KV<K, Iterable<InputT>> kv = windowedKv.getValue();
+ return WindowedValue.of(KV.of(kv.getKey(), keyed.apply(kv.getKey(), kv.getValue())),
+ windowedKv.getTimestamp(), windowedKv.getWindows(), windowedKv.getPane());
+ }
+ }
+
+ /**
+ * Checks if the window transformation should be applied or skipped.
+ *
+ * <p>
+ * Avoid running assign windows if both source and destination are global window
+ * or if the user has not specified the WindowFn (meaning they are just messing
+ * with triggering or allowed lateness).
+ * </p>
+ *
+ * @param transform The {@link Window.Bound} transformation.
+ * @param context The {@link EvaluationContext}.
+ * @param <T> PCollection type.
+ * @param <W> {@link BoundedWindow} type.
+ * @return if to apply the transformation.
+ */
+ public static <T, W extends BoundedWindow> boolean
+ skipAssignWindows(Window.Bound<T> transform, EvaluationContext context) {
+ @SuppressWarnings("unchecked")
+ WindowFn<? super T, W> windowFn = (WindowFn<? super T, W>) transform.getWindowFn();
+ return windowFn == null
+ || (context.getInput(transform).getWindowingStrategy().getWindowFn()
+ instanceof GlobalWindows
+ && windowFn instanceof GlobalWindows);
+ }
+
+ /** Transform a pair stream into a value stream. */
+ public static <T1, T2> JavaDStream<T2> dStreamValues(JavaPairDStream<T1, T2> pairDStream) {
+ return pairDStream.map(new Function<Tuple2<T1, T2>, T2>() {
+ @Override
+ public T2 call(Tuple2<T1, T2> v1) throws Exception {
+ return v1._2();
+ }
+ });
+ }
+
+ /** {@link KV} to pair function. */
+ static <K, V> PairFunction<KV<K, V>, K, V> toPairFunction() {
+ return new PairFunction<KV<K, V>, K, V>() {
+ @Override
+ public Tuple2<K, V> call(KV<K, V> kv) {
+ return new Tuple2<>(kv.getKey(), kv.getValue());
+ }
+ };
+ }
+
+ /** A pair to {@link KV} function . */
+ static <K, V> Function<Tuple2<K, V>, KV<K, V>> fromPairFunction() {
+ return new Function<Tuple2<K, V>, KV<K, V>>() {
+ @Override
+ public KV<K, V> call(Tuple2<K, V> t2) {
+ return KV.of(t2._1(), t2._2());
+ }
+ };
+ }
+
+ /**
+ * A utility class to filter {@link TupleTag}s.
+ *
+ * @param <V> TupleTag type.
+ */
+ public static final class TupleTagFilter<V>
+ implements Function<Tuple2<TupleTag<V>, WindowedValue<?>>, Boolean> {
+
+ private final TupleTag<V> tag;
+
+ public TupleTagFilter(TupleTag<V> tag) {
+ this.tag = tag;
+ }
+
+ @Override
+ public Boolean call(Tuple2<TupleTag<V>, WindowedValue<?>> input) {
+ return tag.equals(input._1());
+ }
+ }
+
+ /***
+ * Create SideInputs as Broadcast variables.
+ *
+ * @param views The {@link PCollectionView}s.
+ * @param context The {@link EvaluationContext}.
+ * @return a map of tagged {@link BroadcastHelper}s.
+ */
+ public static Map<TupleTag<?>, BroadcastHelper<?>> getSideInputs(List<PCollectionView<?>> views,
+ EvaluationContext context) {
+ if (views == null) {
+ return ImmutableMap.of();
+ } else {
+ Map<TupleTag<?>, BroadcastHelper<?>> sideInputs = Maps.newHashMap();
+ for (PCollectionView<?> view : views) {
+ Iterable<? extends WindowedValue<?>> collectionView = context.getPCollectionView(view);
+ Coder<Iterable<WindowedValue<?>>> coderInternal = view.getCoderInternal();
+ @SuppressWarnings("unchecked")
+ BroadcastHelper<?> helper =
+ BroadcastHelper.create((Iterable<WindowedValue<?>>) collectionView, coderInternal);
+ //broadcast side inputs
+ helper.broadcast(context.getSparkContext());
+ sideInputs.put(view.getTagInternal(), helper);
+ }
+ return sideInputs;
+ }
+ }
+
+}
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/0feb6499/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/SparkRunnerStreamingContextFactory.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/SparkRunnerStreamingContextFactory.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/SparkRunnerStreamingContextFactory.java
new file mode 100644
index 0000000..b7a407c
--- /dev/null
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/SparkRunnerStreamingContextFactory.java
@@ -0,0 +1,98 @@
+/*
+ * 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.
+ */
+
+package org.apache.beam.runners.spark.translation.streaming;
+
+import com.google.common.base.Predicates;
+import com.google.common.collect.Iterables;
+import java.net.MalformedURLException;
+import java.net.URL;
+import java.util.Arrays;
+import org.apache.beam.runners.spark.SparkPipelineOptions;
+import org.apache.beam.runners.spark.SparkRunner;
+import org.apache.beam.runners.spark.translation.SparkContextFactory;
+import org.apache.beam.runners.spark.translation.SparkPipelineTranslator;
+import org.apache.beam.runners.spark.translation.TransformTranslator;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.streaming.Duration;
+import org.apache.spark.streaming.api.java.JavaStreamingContext;
+import org.apache.spark.streaming.api.java.JavaStreamingContextFactory;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * A {@link JavaStreamingContext} factory for resilience.
+ * @see <a href="https://spark.apache.org/docs/1.6.2/streaming-programming-guide.html#how-to-configure-checkpointing">how-to-configure-checkpointing</a>
+ */
+public class SparkRunnerStreamingContextFactory implements JavaStreamingContextFactory {
+ private static final Logger LOG =
+ LoggerFactory.getLogger(SparkRunnerStreamingContextFactory.class);
+ private static final Iterable<String> KNOWN_RELIABLE_FS = Arrays.asList("hdfs", "s3", "gs");
+
+ private final Pipeline pipeline;
+ private final SparkPipelineOptions options;
+
+ public SparkRunnerStreamingContextFactory(Pipeline pipeline, SparkPipelineOptions options) {
+ this.pipeline = pipeline;
+ this.options = options;
+ }
+
+ private StreamingEvaluationContext ctxt;
+
+ @Override
+ public JavaStreamingContext create() {
+ LOG.info("Creating a new Spark Streaming Context");
+
+ SparkPipelineTranslator translator = new StreamingTransformTranslator.Translator(
+ new TransformTranslator.Translator());
+ Duration batchDuration = new Duration(options.getBatchIntervalMillis());
+ LOG.info("Setting Spark streaming batchDuration to {} msec", batchDuration.milliseconds());
+
+ JavaSparkContext jsc = SparkContextFactory.getSparkContext(options);
+ JavaStreamingContext jssc = new JavaStreamingContext(jsc, batchDuration);
+ ctxt = new StreamingEvaluationContext(jsc, pipeline, jssc,
+ options.getTimeout());
+ pipeline.traverseTopologically(new SparkRunner.Evaluator(translator, ctxt));
+ ctxt.computeOutputs();
+
+ // set checkpoint dir.
+ String checkpointDir = options.getCheckpointDir();
+ LOG.info("Checkpoint dir set to: {}", checkpointDir);
+ try {
+ // validate checkpoint dir and warn if not of a known durable filesystem.
+ URL checkpointDirUrl = new URL(checkpointDir);
+ if (!Iterables.any(KNOWN_RELIABLE_FS, Predicates.equalTo(checkpointDirUrl.getProtocol()))) {
+ LOG.warn("Checkpoint dir URL {} does not match a reliable filesystem, in case of failures "
+ + "this job may not recover properly or even at all.", checkpointDirUrl);
+ }
+ } catch (MalformedURLException e) {
+ throw new RuntimeException("Failed to form checkpoint dir URL. CheckpointDir should be in "
+ + "the form of hdfs:///path/to/dir or other reliable fs protocol, "
+ + "or file:///path/to/dir for local mode.", e);
+ }
+ jssc.checkpoint(checkpointDir);
+
+ return jssc;
+ }
+
+ public StreamingEvaluationContext getCtxt() {
+ return ctxt;
+ }
+}
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/0feb6499/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingEvaluationContext.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingEvaluationContext.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingEvaluationContext.java
index 2e4da44..5a43c55 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingEvaluationContext.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/translation/streaming/StreamingEvaluationContext.java
@@ -18,14 +18,18 @@
package org.apache.beam.runners.spark.translation.streaming;
+import com.google.common.collect.Iterables;
+
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Queue;
import java.util.Set;
import java.util.concurrent.LinkedBlockingQueue;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
import org.apache.beam.runners.spark.translation.EvaluationContext;
import org.apache.beam.runners.spark.translation.SparkRuntimeContext;
+import org.apache.beam.runners.spark.translation.WindowingHelpers;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.transforms.AppliedPTransform;
@@ -82,11 +86,17 @@ public class StreamingEvaluationContext extends EvaluationContext {
@SuppressWarnings("unchecked")
JavaDStream<WindowedValue<T>> getDStream() {
if (dStream == null) {
- // create the DStream from values
+ WindowedValue.ValueOnlyWindowedValueCoder<T> windowCoder =
+ WindowedValue.getValueOnlyCoder(coder);
+ // create the DStream from queue
Queue<JavaRDD<WindowedValue<T>>> rddQueue = new LinkedBlockingQueue<>();
for (Iterable<T> v : values) {
- setOutputRDDFromValues(currentTransform.getTransform(), v, coder);
- rddQueue.offer((JavaRDD<WindowedValue<T>>) getOutputRDD(currentTransform.getTransform()));
+ Iterable<WindowedValue<T>> windowedValues =
+ Iterables.transform(v, WindowingHelpers.<T>windowValueFunction());
+ JavaRDD<WindowedValue<T>> rdd = getSparkContext().parallelize(
+ CoderHelpers.toByteArrays(windowedValues, windowCoder)).map(
+ CoderHelpers.fromByteFunction(windowCoder));
+ rddQueue.offer(rdd);
}
// create dstream from queue, one at a time, no defaults
// mainly for unit test so no reason to have this configurable
@@ -102,7 +112,10 @@ public class StreamingEvaluationContext extends EvaluationContext {
}
<T> void setStream(PTransform<?, ?> transform, JavaDStream<WindowedValue<T>> dStream) {
- PValue pvalue = (PValue) getOutput(transform);
+ setStream((PValue) getOutput(transform), dStream);
+ }
+
+ <T> void setStream(PValue pvalue, JavaDStream<WindowedValue<T>> dStream) {
DStreamHolder<T> dStreamHolder = new DStreamHolder<>(dStream);
pstreams.put(pvalue, dStreamHolder);
leafStreams.add(dStreamHolder);
@@ -110,6 +123,10 @@ public class StreamingEvaluationContext extends EvaluationContext {
boolean hasStream(PTransform<?, ?> transform) {
PValue pvalue = (PValue) getInput(transform);
+ return hasStream(pvalue);
+ }
+
+ boolean hasStream(PValue pvalue) {
return pstreams.containsKey(pvalue);
}
@@ -141,19 +158,23 @@ public class StreamingEvaluationContext extends EvaluationContext {
@Override
public void computeOutputs() {
+ super.computeOutputs(); // in case the pipeline contains bounded branches as well.
for (DStreamHolder<?> streamHolder : leafStreams) {
computeOutput(streamHolder);
- }
+ } // force a DStream action
}
private static <T> void computeOutput(DStreamHolder<T> streamHolder) {
- streamHolder.getDStream().foreachRDD(new VoidFunction<JavaRDD<WindowedValue<T>>>() {
+ JavaDStream<WindowedValue<T>> dStream = streamHolder.getDStream();
+ // cache in DStream level not RDD
+ // because there could be a difference in StorageLevel if the DStream is windowed.
+ dStream.dstream().cache();
+ dStream.foreachRDD(new VoidFunction<JavaRDD<WindowedValue<T>>>() {
@Override
public void call(JavaRDD<WindowedValue<T>> rdd) throws Exception {
- rdd.rdd().cache();
rdd.count();
}
- }); // force a DStream action
+ });
}
@Override
@@ -163,8 +184,9 @@ public class StreamingEvaluationContext extends EvaluationContext {
} else {
jssc.awaitTermination();
}
- //TODO: stop gracefully ?
- jssc.stop(false, false);
+ // stop streaming context gracefully, so checkpointing (and other computations) get to
+ // finish before shutdown.
+ jssc.stop(false, true);
state = State.DONE;
super.close();
}
@@ -197,7 +219,7 @@ public class StreamingEvaluationContext extends EvaluationContext {
}
@Override
- protected void setCurrentTransform(AppliedPTransform<?, ?, ?> transform) {
+ public void setCurrentTransform(AppliedPTransform<?, ?, ?> transform) {
super.setCurrentTransform(transform);
}
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/0feb6499/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 c55be3d..64ddc57 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
@@ -17,53 +17,68 @@
*/
package org.apache.beam.runners.spark.translation.streaming;
-import com.google.common.collect.Lists;
+import static com.google.common.base.Preconditions.checkState;
+
import com.google.common.collect.Maps;
-import com.google.common.collect.Sets;
-import com.google.common.reflect.TypeToken;
-import java.lang.reflect.ParameterizedType;
-import java.lang.reflect.Type;
+import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Set;
import kafka.serializer.Decoder;
import org.apache.beam.runners.core.AssignWindowsDoFn;
+import org.apache.beam.runners.core.GroupByKeyViaGroupByKeyOnly;
+import org.apache.beam.runners.core.GroupByKeyViaGroupByKeyOnly.GroupAlsoByWindow;
+import org.apache.beam.runners.core.GroupByKeyViaGroupByKeyOnly.GroupByKeyOnly;
+import org.apache.beam.runners.spark.aggregators.AccumulatorSingleton;
+import org.apache.beam.runners.spark.aggregators.NamedAggregators;
+import org.apache.beam.runners.spark.coders.CoderHelpers;
import org.apache.beam.runners.spark.io.ConsoleIO;
import org.apache.beam.runners.spark.io.CreateStream;
import org.apache.beam.runners.spark.io.KafkaIO;
-import org.apache.beam.runners.spark.io.hadoop.HadoopIO;
import org.apache.beam.runners.spark.translation.DoFnFunction;
import org.apache.beam.runners.spark.translation.EvaluationContext;
+import org.apache.beam.runners.spark.translation.GroupCombineFunctions;
+import org.apache.beam.runners.spark.translation.MultiDoFnFunction;
import org.apache.beam.runners.spark.translation.SparkPipelineTranslator;
+import org.apache.beam.runners.spark.translation.SparkRuntimeContext;
import org.apache.beam.runners.spark.translation.TransformEvaluator;
+import org.apache.beam.runners.spark.translation.TranslationUtils;
import org.apache.beam.runners.spark.translation.WindowingHelpers;
+import org.apache.beam.runners.spark.util.BroadcastHelper;
+import org.apache.beam.sdk.coders.CannotProvideCoderException;
import org.apache.beam.sdk.coders.Coder;
-import org.apache.beam.sdk.io.AvroIO;
-import org.apache.beam.sdk.io.TextIO;
-import org.apache.beam.sdk.transforms.AppliedPTransform;
-import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.transforms.Combine;
import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.transforms.OldDoFn;
import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.FixedWindows;
import org.apache.beam.sdk.transforms.windowing.SlidingWindows;
import org.apache.beam.sdk.transforms.windowing.Window;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;
-import org.apache.beam.sdk.values.PDone;
+import org.apache.beam.sdk.values.PCollectionTuple;
+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.VoidFunction;
import org.apache.spark.streaming.Duration;
import org.apache.spark.streaming.Durations;
import org.apache.spark.streaming.api.java.JavaDStream;
import org.apache.spark.streaming.api.java.JavaDStreamLike;
+import org.apache.spark.streaming.api.java.JavaPairDStream;
import org.apache.spark.streaming.api.java.JavaPairInputDStream;
import org.apache.spark.streaming.api.java.JavaStreamingContext;
import org.apache.spark.streaming.kafka.KafkaUtils;
+
import scala.Tuple2;
@@ -114,19 +129,6 @@ public final class StreamingTransformTranslator {
};
}
- private static <T> TransformEvaluator<Create.Values<T>> create() {
- return new TransformEvaluator<Create.Values<T>>() {
- @SuppressWarnings("unchecked")
- @Override
- public void evaluate(Create.Values<T> transform, EvaluationContext context) {
- StreamingEvaluationContext sec = (StreamingEvaluationContext) context;
- Iterable<T> elems = transform.getElements();
- Coder<T> coder = sec.getOutput(transform).getCoder();
- sec.setDStreamFromQueue(transform, Collections.singletonList(elems), coder);
- }
- };
- }
-
private static <T> TransformEvaluator<CreateStream.QueuedValues<T>> createFromQueue() {
return new TransformEvaluator<CreateStream.QueuedValues<T>>() {
@Override
@@ -146,173 +148,325 @@ public final class StreamingTransformTranslator {
public void evaluate(Flatten.FlattenPCollectionList<T> transform, EvaluationContext context) {
StreamingEvaluationContext sec = (StreamingEvaluationContext) context;
PCollectionList<T> pcs = sec.getInput(transform);
- JavaDStream<WindowedValue<T>> first =
- (JavaDStream<WindowedValue<T>>) sec.getStream(pcs.get(0));
- List<JavaDStream<WindowedValue<T>>> rest = Lists.newArrayListWithCapacity(pcs.size() - 1);
- for (int i = 1; i < pcs.size(); i++) {
- rest.add((JavaDStream<WindowedValue<T>>) sec.getStream(pcs.get(i)));
+ // since this is a streaming pipeline, at least one of the PCollections to "flatten" are
+ // unbounded, meaning it represents a DStream.
+ // So we could end up with an unbounded unified DStream.
+ final List<JavaRDD<WindowedValue<T>>> rdds = new ArrayList<>();
+ final List<JavaDStream<WindowedValue<T>>> dStreams = new ArrayList<>();
+ for (PCollection<T> pcol: pcs.getAll()) {
+ if (sec.hasStream(pcol)) {
+ dStreams.add((JavaDStream<WindowedValue<T>>) sec.getStream(pcol));
+ } else {
+ rdds.add((JavaRDD<WindowedValue<T>>) context.getRDD(pcol));
+ }
+ }
+ // start by unifying streams into a single stream.
+ JavaDStream<WindowedValue<T>> unifiedStreams =
+ sec.getStreamingContext().union(dStreams.remove(0), dStreams);
+ // now unify in RDDs.
+ if (rdds.size() > 0) {
+ JavaDStream<WindowedValue<T>> joined = unifiedStreams.transform(
+ new Function<JavaRDD<WindowedValue<T>>, JavaRDD<WindowedValue<T>>>() {
+ @Override
+ public JavaRDD<WindowedValue<T>> call(JavaRDD<WindowedValue<T>> streamRdd)
+ throws Exception {
+ return new JavaSparkContext(streamRdd.context()).union(streamRdd, rdds);
+ }
+ });
+ sec.setStream(transform, joined);
+ } else {
+ sec.setStream(transform, unifiedStreams);
}
- JavaDStream<WindowedValue<T>> dstream = sec.getStreamingContext().union(first, rest);
- sec.setStream(transform, dstream);
}
};
}
- private static <TransformT extends PTransform<?, ?>> TransformEvaluator<TransformT> rddTransform(
- final SparkPipelineTranslator rddTranslator) {
- return new TransformEvaluator<TransformT>() {
- @SuppressWarnings("unchecked")
+ private static <T, W extends BoundedWindow> TransformEvaluator<Window.Bound<T>> window() {
+ return new TransformEvaluator<Window.Bound<T>>() {
@Override
- public void evaluate(TransformT transform, EvaluationContext context) {
- TransformEvaluator<TransformT> rddEvaluator =
- rddTranslator.translate((Class<TransformT>) transform.getClass());
-
+ public void evaluate(Window.Bound<T> transform, EvaluationContext context) {
StreamingEvaluationContext sec = (StreamingEvaluationContext) context;
- if (sec.hasStream(transform)) {
- JavaDStreamLike<WindowedValue<Object>, ?, JavaRDD<WindowedValue<Object>>> dStream =
- (JavaDStreamLike<WindowedValue<Object>, ?, JavaRDD<WindowedValue<Object>>>)
- sec.getStream(transform);
-
- sec.setStream(transform, dStream
- .transform(new RDDTransform<>(sec, rddEvaluator, transform)));
+ @SuppressWarnings("unchecked")
+ WindowFn<? super T, W> windowFn = (WindowFn<? super T, W>) transform.getWindowFn();
+ @SuppressWarnings("unchecked")
+ JavaDStream<WindowedValue<T>> dStream =
+ (JavaDStream<WindowedValue<T>>) sec.getStream(transform);
+ if (windowFn instanceof FixedWindows) {
+ Duration windowDuration = Durations.milliseconds(((FixedWindows) windowFn).getSize()
+ .getMillis());
+ sec.setStream(transform, dStream.window(windowDuration));
+ } else if (windowFn instanceof SlidingWindows) {
+ Duration windowDuration = Durations.milliseconds(((SlidingWindows) windowFn).getSize()
+ .getMillis());
+ Duration slideDuration = Durations.milliseconds(((SlidingWindows) windowFn).getPeriod()
+ .getMillis());
+ sec.setStream(transform, dStream.window(windowDuration, slideDuration));
+ }
+ //--- then we apply windowing to the elements
+ @SuppressWarnings("unchecked")
+ JavaDStream<WindowedValue<T>> dStream2 =
+ (JavaDStream<WindowedValue<T>>) sec.getStream(transform);
+ if (TranslationUtils.skipAssignWindows(transform, context)) {
+ sec.setStream(transform, dStream2);
} else {
- // if the transformation requires direct access to RDD (not in stream)
- // this is used for "fake" transformations like with PAssert
- rddEvaluator.evaluate(transform, context);
+ final OldDoFn<T, T> addWindowsDoFn = new AssignWindowsDoFn<>(windowFn);
+ final SparkRuntimeContext runtimeContext = sec.getRuntimeContext();
+ JavaDStream<WindowedValue<T>> outStream = dStream2.transform(
+ new Function<JavaRDD<WindowedValue<T>>, JavaRDD<WindowedValue<T>>>() {
+ @Override
+ public JavaRDD<WindowedValue<T>> call(JavaRDD<WindowedValue<T>> rdd) throws Exception {
+ final Accumulator<NamedAggregators> accum =
+ AccumulatorSingleton.getInstance(new JavaSparkContext(rdd.context()));
+ return rdd.mapPartitions(
+ new DoFnFunction<>(accum, addWindowsDoFn, runtimeContext, null));
+ }
+ });
+ sec.setStream(transform, outStream);
}
}
};
}
- /**
- * RDD transform function If the transformation function doesn't have an input, create a fake one
- * as an empty RDD.
- *
- * @param <TransformT> PTransform type
- */
- private static final class RDDTransform<TransformT extends PTransform<?, ?>>
- implements Function<JavaRDD<WindowedValue<Object>>, JavaRDD<WindowedValue<Object>>> {
-
- private final StreamingEvaluationContext context;
- private final AppliedPTransform<?, ?, ?> appliedPTransform;
- private final TransformEvaluator<TransformT> rddEvaluator;
- private final TransformT transform;
-
-
- private RDDTransform(StreamingEvaluationContext context,
- TransformEvaluator<TransformT> rddEvaluator,
- TransformT transform) {
- this.context = context;
- this.appliedPTransform = context.getCurrentTransform();
- this.rddEvaluator = rddEvaluator;
- this.transform = transform;
- }
+ private static <K, V> TransformEvaluator<GroupByKeyOnly<K, V>> gbko() {
+ return new TransformEvaluator<GroupByKeyOnly<K, V>>() {
+ @Override
+ public void evaluate(GroupByKeyOnly<K, V> transform, EvaluationContext context) {
+ StreamingEvaluationContext sec = (StreamingEvaluationContext) context;
- @Override
- @SuppressWarnings("unchecked")
- public JavaRDD<WindowedValue<Object>>
- call(JavaRDD<WindowedValue<Object>> rdd) throws Exception {
- AppliedPTransform<?, ?, ?> existingAPT = context.getCurrentTransform();
- context.setCurrentTransform(appliedPTransform);
- context.setInputRDD(transform, rdd);
- rddEvaluator.evaluate(transform, context);
- if (!context.hasOutputRDD(transform)) {
- // fake RDD as output
- context.setOutputRDD(transform,
- context.getSparkContext().<WindowedValue<Object>>emptyRDD());
+ @SuppressWarnings("unchecked")
+ JavaDStream<WindowedValue<KV<K, V>>> dStream =
+ (JavaDStream<WindowedValue<KV<K, V>>>) sec.getStream(transform);
+
+ @SuppressWarnings("unchecked")
+ final KvCoder<K, V> coder = (KvCoder<K, V>) sec.getInput(transform).getCoder();
+
+ JavaDStream<WindowedValue<KV<K, Iterable<V>>>> outStream =
+ 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 {
+ return GroupCombineFunctions.groupByKeyOnly(rdd, coder);
+ }
+ });
+ sec.setStream(transform, outStream);
}
- JavaRDD<WindowedValue<Object>> outRDD =
- (JavaRDD<WindowedValue<Object>>) context.getOutputRDD(transform);
- context.setCurrentTransform(existingAPT);
- return outRDD;
- }
+ };
}
- @SuppressWarnings("unchecked")
- private static <TransformT extends PTransform<?, ?>> TransformEvaluator<TransformT> foreachRDD(
- final SparkPipelineTranslator rddTranslator) {
- return new TransformEvaluator<TransformT>() {
+ private static <K, V, W extends BoundedWindow>
+ TransformEvaluator<GroupAlsoByWindow<K, V>> gabw() {
+ return new TransformEvaluator<GroupAlsoByWindow<K, V>>() {
@Override
- public void evaluate(TransformT transform, EvaluationContext context) {
- TransformEvaluator<TransformT> rddEvaluator =
- rddTranslator.translate((Class<TransformT>) transform.getClass());
+ public void evaluate(final GroupAlsoByWindow<K, V> transform, EvaluationContext context) {
+ final StreamingEvaluationContext sec = (StreamingEvaluationContext) context;
+ final SparkRuntimeContext runtimeContext = sec.getRuntimeContext();
+ @SuppressWarnings("unchecked")
+ JavaDStream<WindowedValue<KV<K, Iterable<WindowedValue<V>>>>> dStream =
+ (JavaDStream<WindowedValue<KV<K, Iterable<WindowedValue<V>>>>>)
+ sec.getStream(transform);
- StreamingEvaluationContext sec = (StreamingEvaluationContext) context;
- if (sec.hasStream(transform)) {
- JavaDStreamLike<WindowedValue<Object>, ?, JavaRDD<WindowedValue<Object>>> dStream =
- (JavaDStreamLike<WindowedValue<Object>, ?, JavaRDD<WindowedValue<Object>>>)
- sec.getStream(transform);
+ @SuppressWarnings("unchecked")
+ final KvCoder<K, Iterable<WindowedValue<V>>> inputKvCoder =
+ (KvCoder<K, Iterable<WindowedValue<V>>>) sec.getInput(transform).getCoder();
+
+ JavaDStream<WindowedValue<KV<K, Iterable<V>>>> outStream =
+ dStream.transform(new Function<JavaRDD<WindowedValue<KV<K,
+ Iterable<WindowedValue<V>>>>>, JavaRDD<WindowedValue<KV<K, Iterable<V>>>>>() {
+ @Override
+ public JavaRDD<WindowedValue<KV<K, Iterable<V>>>> call(JavaRDD<WindowedValue<KV<K,
+ Iterable<WindowedValue<V>>>>> rdd) throws Exception {
+ final Accumulator<NamedAggregators> accum =
+ AccumulatorSingleton.getInstance(new JavaSparkContext(rdd.context()));
+ return GroupCombineFunctions.groupAlsoByWindow(rdd, transform, runtimeContext,
+ accum, inputKvCoder);
+ }
+ });
+ sec.setStream(transform, outStream);
+ }
+ };
+ }
- dStream.foreachRDD(new RDDOutputOperator<>(sec, rddEvaluator, transform));
- } else {
- rddEvaluator.evaluate(transform, context);
- }
+ private static <K, InputT, OutputT> TransformEvaluator<Combine.GroupedValues<K, InputT, OutputT>>
+ grouped() {
+ return new TransformEvaluator<Combine.GroupedValues<K, InputT, OutputT>>() {
+ @Override
+ public void evaluate(Combine.GroupedValues<K, InputT, OutputT> transform,
+ EvaluationContext context) {
+ StreamingEvaluationContext sec = (StreamingEvaluationContext) context;
+ @SuppressWarnings("unchecked")
+ JavaDStream<WindowedValue<KV<K, Iterable<InputT>>>> dStream =
+ (JavaDStream<WindowedValue<KV<K, Iterable<InputT>>>>) sec.getStream(transform);
+ sec.setStream(transform, dStream.map(
+ new TranslationUtils.CombineGroupedValues<>(transform)));
}
};
}
- /**
- * RDD output function.
- *
- * @param <TransformT> PTransform type
- */
- private static final class RDDOutputOperator<TransformT extends PTransform<?, ?>>
- implements VoidFunction<JavaRDD<WindowedValue<Object>>> {
+ private static <InputT, AccumT, OutputT> TransformEvaluator<Combine.Globally<InputT, OutputT>>
+ combineGlobally() {
+ return new TransformEvaluator<Combine.Globally<InputT, OutputT>>() {
- private final StreamingEvaluationContext context;
- private final AppliedPTransform<?, ?, ?> appliedPTransform;
- private final TransformEvaluator<TransformT> rddEvaluator;
- private final TransformT transform;
+ @Override
+ public void evaluate(Combine.Globally<InputT, OutputT> transform, EvaluationContext context) {
+ StreamingEvaluationContext sec = (StreamingEvaluationContext) context;
+ @SuppressWarnings("unchecked")
+ final Combine.CombineFn<InputT, AccumT, OutputT> globally =
+ (Combine.CombineFn<InputT, AccumT, OutputT>) transform.getFn();
+ @SuppressWarnings("unchecked")
+ JavaDStream<WindowedValue<InputT>> dStream =
+ (JavaDStream<WindowedValue<InputT>>) sec.getStream(transform);
+
+ final Coder<InputT> iCoder = sec.getInput(transform).getCoder();
+ final Coder<OutputT> oCoder = sec.getOutput(transform).getCoder();
+ final Coder<AccumT> aCoder;
+ try {
+ aCoder = globally.getAccumulatorCoder(sec.getPipeline().getCoderRegistry(), iCoder);
+ } catch (CannotProvideCoderException e) {
+ throw new IllegalStateException("Could not determine coder for accumulator", e);
+ }
- private RDDOutputOperator(StreamingEvaluationContext context,
- TransformEvaluator<TransformT> rddEvaluator, TransformT transform) {
- this.context = context;
- this.appliedPTransform = context.getCurrentTransform();
- this.rddEvaluator = rddEvaluator;
- this.transform = transform;
- }
+ 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 {
+ JavaRDD<byte[]> outRdd = new JavaSparkContext(rdd.context()).parallelize(
+ // don't use Guava's ImmutableList.of as output may be null
+ CoderHelpers.toByteArrays(Collections.singleton(
+ GroupCombineFunctions.combineGlobally(rdd, globally, iCoder, aCoder)), oCoder));
+ return outRdd.map(CoderHelpers.fromByteFunction(oCoder)).map(
+ WindowingHelpers.<OutputT>windowFunction());
+ }
+ });
- @Override
- @SuppressWarnings("unchecked")
- public void call(JavaRDD<WindowedValue<Object>> rdd) throws Exception {
- AppliedPTransform<?, ?, ?> existingAPT = context.getCurrentTransform();
- context.setCurrentTransform(appliedPTransform);
- context.setInputRDD(transform, rdd);
- rddEvaluator.evaluate(transform, context);
- context.setCurrentTransform(existingAPT);
- }
+ sec.setStream(transform, outStream);
+ }
+ };
}
- private static <T> TransformEvaluator<Window.Bound<T>> window() {
- return new TransformEvaluator<Window.Bound<T>>() {
+ 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(Window.Bound<T> transform, EvaluationContext context) {
+ public void evaluate(Combine.PerKey<K, InputT, OutputT>
+ transform, EvaluationContext context) {
StreamingEvaluationContext sec = (StreamingEvaluationContext) context;
- WindowFn<? super T, ?> windowFn = transform.getWindowFn();
@SuppressWarnings("unchecked")
- JavaDStream<WindowedValue<T>> dStream =
- (JavaDStream<WindowedValue<T>>) sec.getStream(transform);
- if (windowFn instanceof FixedWindows) {
- Duration windowDuration = Durations.milliseconds(((FixedWindows) windowFn).getSize()
- .getMillis());
- sec.setStream(transform, dStream.window(windowDuration));
- } else if (windowFn instanceof SlidingWindows) {
- Duration windowDuration = Durations.milliseconds(((SlidingWindows) windowFn).getSize()
- .getMillis());
- Duration slideDuration = Durations.milliseconds(((SlidingWindows) windowFn).getPeriod()
- .getMillis());
- sec.setStream(transform, dStream.window(windowDuration, slideDuration));
+ final Combine.KeyedCombineFn<K, InputT, AccumT, OutputT> keyed =
+ (Combine.KeyedCombineFn<K, InputT, AccumT, OutputT>) transform.getFn();
+ @SuppressWarnings("unchecked")
+ JavaDStream<WindowedValue<KV<K, InputT>>> dStream =
+ (JavaDStream<WindowedValue<KV<K, InputT>>>) sec.getStream(transform);
+
+ @SuppressWarnings("unchecked")
+ KvCoder<K, InputT> inputCoder = (KvCoder<K, InputT>) sec.getInput(transform).getCoder();
+ Coder<K> keyCoder = inputCoder.getKeyCoder();
+ Coder<InputT> viCoder = inputCoder.getValueCoder();
+ Coder<AccumT> vaCoder;
+ try {
+ vaCoder = keyed.getAccumulatorCoder(
+ context.getPipeline().getCoderRegistry(), keyCoder, viCoder);
+ } catch (CannotProvideCoderException e) {
+ throw new IllegalStateException("Could not determine coder for accumulator", e);
}
- //--- then we apply windowing to the elements
- OldDoFn<T, T> addWindowsDoFn = new AssignWindowsDoFn<>(windowFn);
- DoFnFunction<T, T> dofn = new DoFnFunction<>(addWindowsDoFn,
- ((StreamingEvaluationContext) context).getRuntimeContext(), null);
+ Coder<KV<K, InputT>> kviCoder = KvCoder.of(keyCoder, viCoder);
+ Coder<KV<K, AccumT>> kvaCoder = KvCoder.of(keyCoder, vaCoder);
+ //-- windowed coders
+ final WindowedValue.FullWindowedValueCoder<K> wkCoder =
+ WindowedValue.FullWindowedValueCoder.of(keyCoder,
+ sec.getInput(transform).getWindowingStrategy().getWindowFn().windowCoder());
+ final WindowedValue.FullWindowedValueCoder<KV<K, InputT>> wkviCoder =
+ WindowedValue.FullWindowedValueCoder.of(kviCoder,
+ sec.getInput(transform).getWindowingStrategy().getWindowFn().windowCoder());
+ final WindowedValue.FullWindowedValueCoder<KV<K, AccumT>> wkvaCoder =
+ WindowedValue.FullWindowedValueCoder.of(kvaCoder,
+ sec.getInput(transform).getWindowingStrategy().getWindowFn().windowCoder());
+
+ 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 {
+ return GroupCombineFunctions.combinePerKey(rdd, keyed, wkCoder, wkviCoder, wkvaCoder);
+ }
+ });
+
+ sec.setStream(transform, outStream);
+ }
+ };
+ }
+
+ private static <InputT, OutputT> TransformEvaluator<ParDo.Bound<InputT, OutputT>> parDo() {
+ return new TransformEvaluator<ParDo.Bound<InputT, OutputT>>() {
+ @Override
+ public void evaluate(final ParDo.Bound<InputT, OutputT> transform,
+ final EvaluationContext context) {
+ final StreamingEvaluationContext sec = (StreamingEvaluationContext) context;
+ final SparkRuntimeContext runtimeContext = sec.getRuntimeContext();
+ final Map<TupleTag<?>, BroadcastHelper<?>> sideInputs =
+ TranslationUtils.getSideInputs(transform.getSideInputs(), context);
@SuppressWarnings("unchecked")
- JavaDStreamLike<WindowedValue<T>, ?, JavaRDD<WindowedValue<T>>> dstream =
- (JavaDStreamLike<WindowedValue<T>, ?, JavaRDD<WindowedValue<T>>>)
- sec.getStream(transform);
- sec.setStream(transform, dstream.mapPartitions(dofn));
+ JavaDStream<WindowedValue<InputT>> dStream =
+ (JavaDStream<WindowedValue<InputT>>) sec.getStream(transform);
+
+ 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 Accumulator<NamedAggregators> accum =
+ AccumulatorSingleton.getInstance(new JavaSparkContext(rdd.context()));
+ return rdd.mapPartitions(
+ new DoFnFunction<>(accum, transform.getFn(), runtimeContext, sideInputs));
+ }
+ });
+
+ sec.setStream(transform, outStream);
+ }
+ };
+ }
+
+ private static <InputT, OutputT> TransformEvaluator<ParDo.BoundMulti<InputT, OutputT>>
+ multiDo() {
+ return new TransformEvaluator<ParDo.BoundMulti<InputT, OutputT>>() {
+ @Override
+ public void evaluate(final ParDo.BoundMulti<InputT, OutputT> transform,
+ final EvaluationContext context) {
+ final StreamingEvaluationContext sec = (StreamingEvaluationContext) context;
+ final SparkRuntimeContext runtimeContext = sec.getRuntimeContext();
+ final Map<TupleTag<?>, BroadcastHelper<?>> sideInputs =
+ TranslationUtils.getSideInputs(transform.getSideInputs(), context);
+ @SuppressWarnings("unchecked")
+ JavaDStream<WindowedValue<InputT>> dStream =
+ (JavaDStream<WindowedValue<InputT>>) sec.getStream(transform);
+ JavaPairDStream<TupleTag<?>, WindowedValue<?>> all = dStream.transformToPair(
+ new Function<JavaRDD<WindowedValue<InputT>>,
+ JavaPairRDD<TupleTag<?>, WindowedValue<?>>>() {
+ @Override
+ public JavaPairRDD<TupleTag<?>, WindowedValue<?>> call(
+ JavaRDD<WindowedValue<InputT>> rdd) throws Exception {
+ final Accumulator<NamedAggregators> accum =
+ AccumulatorSingleton.getInstance(new JavaSparkContext(rdd.context()));
+ return rdd.mapPartitionsToPair(new MultiDoFnFunction<>(accum, transform.getFn(),
+ runtimeContext, transform.getMainOutputTag(), sideInputs));
+ }
+ }).cache();
+ PCollectionTuple pct = sec.getOutput(transform);
+ for (Map.Entry<TupleTag<?>, PCollection<?>> e : pct.getAll().entrySet()) {
+ @SuppressWarnings("unchecked")
+ JavaPairDStream<TupleTag<?>, WindowedValue<?>> filtered =
+ all.filter(new TranslationUtils.TupleTagFilter(e.getKey()));
+ @SuppressWarnings("unchecked")
+ // Object is the best we can do since different outputs can have different tags
+ JavaDStream<WindowedValue<Object>> values =
+ (JavaDStream<WindowedValue<Object>>)
+ (JavaDStream<?>) TranslationUtils.dStreamValues(filtered);
+ sec.setStream(e.getValue(), values);
+ }
}
};
}
@@ -321,79 +475,54 @@ public final class StreamingTransformTranslator {
.newHashMap();
static {
+ EVALUATORS.put(GroupByKeyViaGroupByKeyOnly.GroupByKeyOnly.class, gbko());
+ EVALUATORS.put(GroupByKeyViaGroupByKeyOnly.GroupAlsoByWindow.class, gabw());
+ EVALUATORS.put(Combine.GroupedValues.class, grouped());
+ 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());
EVALUATORS.put(CreateStream.QueuedValues.class, createFromQueue());
- EVALUATORS.put(Create.Values.class, create());
EVALUATORS.put(KafkaIO.Read.Unbound.class, kafka());
EVALUATORS.put(Window.Bound.class, window());
EVALUATORS.put(Flatten.FlattenPCollectionList.class, flattenPColl());
}
- private static final Set<Class<? extends PTransform>> UNSUPPORTED_EVALUATORS = Sets
- .newHashSet();
-
- static {
- //TODO - add support for the following
- UNSUPPORTED_EVALUATORS.add(TextIO.Read.Bound.class);
- UNSUPPORTED_EVALUATORS.add(TextIO.Write.Bound.class);
- UNSUPPORTED_EVALUATORS.add(AvroIO.Read.Bound.class);
- UNSUPPORTED_EVALUATORS.add(AvroIO.Write.Bound.class);
- UNSUPPORTED_EVALUATORS.add(HadoopIO.Read.Bound.class);
- UNSUPPORTED_EVALUATORS.add(HadoopIO.Write.Bound.class);
- }
-
- @SuppressWarnings("unchecked")
- private static <TransformT extends PTransform<?, ?>> TransformEvaluator<TransformT>
- getTransformEvaluator(Class<TransformT> clazz, SparkPipelineTranslator rddTranslator) {
- TransformEvaluator<TransformT> transform =
- (TransformEvaluator<TransformT>) EVALUATORS.get(clazz);
- if (transform == null) {
- if (UNSUPPORTED_EVALUATORS.contains(clazz)) {
- throw new UnsupportedOperationException("Beam transformation " + clazz
- .getCanonicalName()
- + " is currently unsupported by the Spark streaming pipeline");
- }
- // DStream transformations will transform an RDD into another RDD
- // Actions will create output
- // In Beam it depends on the PTransform's Input and Output class
- Class<?> pTOutputClazz = getPTransformOutputClazz(clazz);
- if (PDone.class.equals(pTOutputClazz)) {
- return foreachRDD(rddTranslator);
- } else {
- return rddTransform(rddTranslator);
- }
- }
- return transform;
- }
-
- private static <TransformT extends PTransform<?, ?>> Class<?>
- getPTransformOutputClazz(Class<TransformT> clazz) {
- Type[] types = ((ParameterizedType) clazz.getGenericSuperclass()).getActualTypeArguments();
- return TypeToken.of(clazz).resolveType(types[1]).getRawType();
- }
-
/**
- * Translator matches Beam transformation with the appropriate Spark streaming evaluator.
- * rddTranslator uses Spark evaluators in transform/foreachRDD to evaluate the transformation
+ * Translator matches Beam transformation with the appropriate evaluator.
*/
public static class Translator implements SparkPipelineTranslator {
- private final SparkPipelineTranslator rddTranslator;
+ private final SparkPipelineTranslator batchTranslator;
- public Translator(SparkPipelineTranslator rddTranslator) {
- this.rddTranslator = rddTranslator;
+ Translator(SparkPipelineTranslator batchTranslator) {
+ this.batchTranslator = batchTranslator;
}
@Override
public boolean hasTranslation(Class<? extends PTransform<?, ?>> clazz) {
- // streaming includes rdd transformations as well
- return EVALUATORS.containsKey(clazz) || rddTranslator.hasTranslation(clazz);
+ // streaming includes rdd/bounded transformations as well
+ return EVALUATORS.containsKey(clazz) || batchTranslator.hasTranslation(clazz);
+ }
+
+ @Override
+ public <TransformT extends PTransform<?, ?>> TransformEvaluator<TransformT>
+ translateBounded(Class<TransformT> clazz) {
+ TransformEvaluator<TransformT> transformEvaluator = batchTranslator.translateBounded(clazz);
+ checkState(transformEvaluator != null,
+ "No TransformEvaluator registered for BOUNDED transform %s", clazz);
+ return transformEvaluator;
}
@Override
public <TransformT extends PTransform<?, ?>> TransformEvaluator<TransformT>
- translate(Class<TransformT> clazz) {
- return getTransformEvaluator(clazz, rddTranslator);
+ translateUnbounded(Class<TransformT> clazz) {
+ @SuppressWarnings("unchecked") TransformEvaluator<TransformT> transformEvaluator =
+ (TransformEvaluator<TransformT>) EVALUATORS.get(clazz);
+ checkState(transformEvaluator != null,
+ "No TransformEvaluator registered for for UNBOUNDED transform %s", clazz);
+ return transformEvaluator;
}
}
}
http://git-wip-us.apache.org/repos/asf/incubator-beam/blob/0feb6499/runners/spark/src/main/java/org/apache/beam/runners/spark/util/BroadcastHelper.java
----------------------------------------------------------------------
diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/BroadcastHelper.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/BroadcastHelper.java
index 5c13b80..0e742eb 100644
--- a/runners/spark/src/main/java/org/apache/beam/runners/spark/util/BroadcastHelper.java
+++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/util/BroadcastHelper.java
@@ -52,8 +52,12 @@ public abstract class BroadcastHelper<T> implements Serializable {
public abstract T getValue();
+ public abstract boolean isBroadcasted();
+
public abstract void broadcast(JavaSparkContext jsc);
+ public abstract void unregister();
+
/**
* A {@link BroadcastHelper} that relies on the underlying
* Spark serialization (Kryo) to broadcast values. This is appropriate when
@@ -77,9 +81,20 @@ public abstract class BroadcastHelper<T> implements Serializable {
}
@Override
+ public boolean isBroadcasted() {
+ return bcast != null;
+ }
+
+ @Override
public void broadcast(JavaSparkContext jsc) {
this.bcast = jsc.broadcast(value);
}
+
+ @Override
+ public void unregister() {
+ this.bcast.destroy();
+ this.bcast = null;
+ }
}
/**
@@ -107,10 +122,21 @@ public abstract class BroadcastHelper<T> implements Serializable {
}
@Override
+ public boolean isBroadcasted() {
+ return bcast != null;
+ }
+
+ @Override
public void broadcast(JavaSparkContext jsc) {
this.bcast = jsc.broadcast(CoderHelpers.toByteArray(value, coder));
}
+ @Override
+ public void unregister() {
+ this.bcast.destroy();
+ this.bcast = null;
+ }
+
private T deserialize() {
T val;
try {