You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by av...@apache.org on 2017/04/01 07:28:35 UTC
[1/2] beam git commit: [BEAM-1337] Infer state coders
Repository: beam
Updated Branches:
refs/heads/master 03dce6dcc -> e31ca8b0d
[BEAM-1337] Infer state coders
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/42e690e8
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/42e690e8
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/42e690e8
Branch: refs/heads/master
Commit: 42e690e84a9f05d508f2528b1444b26ce031e080
Parents: 03dce6d
Author: Aviem Zur <av...@gmail.com>
Authored: Wed Mar 1 07:27:57 2017 +0200
Committer: Aviem Zur <av...@gmail.com>
Committed: Sat Apr 1 10:27:14 2017 +0300
----------------------------------------------------------------------
.../org/apache/beam/sdk/transforms/ParDo.java | 62 ++
.../apache/beam/sdk/util/state/StateSpec.java | 15 +
.../apache/beam/sdk/util/state/StateSpecs.java | 264 ++++++++-
.../apache/beam/sdk/transforms/ParDoTest.java | 578 +++++++++++++++++++
4 files changed, 902 insertions(+), 17 deletions(-)
----------------------------------------------------------------------
http://git-wip-us.apache.org/repos/asf/beam/blob/42e690e8/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java
index 664fbc3..3de845b 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/ParDo.java
@@ -22,6 +22,8 @@ import static com.google.common.base.Preconditions.checkArgument;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import java.io.Serializable;
+import java.lang.reflect.ParameterizedType;
+import java.lang.reflect.Type;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
@@ -29,6 +31,7 @@ import java.util.Map;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.CannotProvideCoderException;
import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.CoderRegistry;
import org.apache.beam.sdk.runners.PipelineRunner;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.display.DisplayData.Builder;
@@ -41,6 +44,7 @@ import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.WindowFn;
import org.apache.beam.sdk.util.NameUtils;
import org.apache.beam.sdk.util.SerializableUtils;
+import org.apache.beam.sdk.util.state.StateSpec;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.PCollectionView;
@@ -434,6 +438,59 @@ public class ParDo {
return DisplayData.item("fn", fn.getClass()).withLabel("Transform Function");
}
+ private static void finishSpecifyingStateSpecs(
+ DoFn<?, ?> fn,
+ CoderRegistry coderRegistry,
+ Coder<?> inputCoder) {
+ DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
+ Map<String, DoFnSignature.StateDeclaration> stateDeclarations = signature.stateDeclarations();
+ for (DoFnSignature.StateDeclaration stateDeclaration : stateDeclarations.values()) {
+ try {
+ StateSpec<?, ?> stateSpec = (StateSpec<?, ?>) stateDeclaration.field().get(fn);
+ stateSpec.offerCoders(codersForStateSpecTypes(stateDeclaration, coderRegistry, inputCoder));
+ stateSpec.finishSpecifying();
+ } catch (IllegalAccessException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ }
+
+ /**
+ * Try to provide coders for as many of the type arguments of given
+ * {@link DoFnSignature.StateDeclaration} as possible.
+ */
+ private static <InputT> Coder[] codersForStateSpecTypes(
+ DoFnSignature.StateDeclaration stateDeclaration,
+ CoderRegistry coderRegistry,
+ Coder<InputT> inputCoder) {
+ Type stateType = stateDeclaration.stateType().getType();
+
+ if (!(stateType instanceof ParameterizedType)) {
+ // No type arguments means no coders to infer.
+ return new Coder[0];
+ }
+
+ Type[] typeArguments = ((ParameterizedType) stateType).getActualTypeArguments();
+ Coder[] coders = new Coder[typeArguments.length];
+
+ for (int i = 0; i < typeArguments.length; i++) {
+ Type typeArgument = typeArguments[i];
+ TypeDescriptor<?> typeDescriptor = TypeDescriptor.of(typeArgument);
+ try {
+ coders[i] = coderRegistry.getDefaultCoder(typeDescriptor);
+ } catch (CannotProvideCoderException e) {
+ try {
+ coders[i] = coderRegistry.getDefaultCoder(
+ typeDescriptor, inputCoder.getEncodedTypeDescriptor(), inputCoder);
+ } catch (CannotProvideCoderException ignored) {
+ // Since not all type arguments will have a registered coder we ignore this exception.
+ }
+ }
+ }
+
+ return coders;
+ }
+
/**
* Perform common validations of the {@link DoFn} against the input {@link PCollection}, for
* example ensuring that the window type expected by the {@link DoFn} matches the window type of
@@ -557,6 +614,7 @@ public class ParDo {
@Override
public PCollection<OutputT> expand(PCollection<? extends InputT> input) {
+ finishSpecifyingStateSpecs(fn, input.getPipeline().getCoderRegistry(), input.getCoder());
TupleTag<OutputT> mainOutput = new TupleTag<>();
return input.apply(withOutputTags(mainOutput, TupleTagList.empty())).get(mainOutput);
}
@@ -681,6 +739,10 @@ public class ParDo {
public PCollectionTuple expand(PCollection<? extends InputT> input) {
// SplittableDoFn should be forbidden on the runner-side.
validateWindowType(input, fn);
+
+ // Use coder registry to determine coders for all StateSpec defined in the fn signature.
+ finishSpecifyingStateSpecs(fn, input.getPipeline().getCoderRegistry(), input.getCoder());
+
PCollectionTuple outputs = PCollectionTuple.ofPrimitiveOutputsInternal(
input.getPipeline(),
TupleTagList.of(mainOutputTag).and(sideOutputTags.getAll()),
http://git-wip-us.apache.org/repos/asf/beam/blob/42e690e8/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpec.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpec.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpec.java
index 4fdeefb..6b94c40 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpec.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpec.java
@@ -20,6 +20,7 @@ package org.apache.beam.sdk.util.state;
import java.io.Serializable;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.annotations.Experimental.Kind;
+import org.apache.beam.sdk.coders.Coder;
/**
* A specification of a persistent state cell. This includes information necessary to encode the
@@ -36,4 +37,18 @@ public interface StateSpec<K, StateT extends State> extends Serializable {
* Use the {@code binder} to create an instance of {@code StateT} appropriate for this address.
*/
StateT bind(String id, StateBinder<? extends K> binder);
+
+ /**
+ * Given {code coders} are inferred from type arguments defined for this class.
+ * Coders which are already set should take precedence over offered coders.
+ * @param coders Array of coders indexed by the type arguments order. Entries might be null if
+ * the coder could not be inferred.
+ */
+ void offerCoders(Coder[] coders);
+
+ /**
+ * Validates that this {@link StateSpec} has been specified correctly and finalizes it.
+ * Automatically invoked when the pipeline is built.
+ */
+ void finishSpecifying();
}
http://git-wip-us.apache.org/repos/asf/beam/blob/42e690e8/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpecs.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpecs.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpecs.java
index 8912993..6a8c80b 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpecs.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/util/state/StateSpecs.java
@@ -17,7 +17,10 @@
*/
package org.apache.beam.sdk.util.state;
+import static com.google.common.base.Preconditions.checkArgument;
+
import java.util.Objects;
+import javax.annotation.Nullable;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.annotations.Experimental.Kind;
import org.apache.beam.sdk.coders.CannotProvideCoderException;
@@ -44,7 +47,13 @@ public class StateSpecs {
private StateSpecs() {}
/** Create a simple state spec for values of type {@code T}. */
+ public static <T> StateSpec<Object, ValueState<T>> value() {
+ return new ValueStateSpec<>(null);
+ }
+
+ /** Create a simple state spec for values of type {@code T}. */
public static <T> StateSpec<Object, ValueState<T>> value(Coder<T> valueCoder) {
+ checkArgument(valueCoder != null, "valueCoder should not be null. Consider value() instead");
return new ValueStateSpec<>(valueCoder);
}
@@ -53,8 +62,21 @@ public class StateSpecs {
* {@code InputT}s into a single {@code OutputT}.
*/
public static <InputT, AccumT, OutputT>
+ StateSpec<Object, AccumulatorCombiningState<InputT, AccumT, OutputT>> combiningValue(
+ CombineFn<InputT, AccumT, OutputT> combineFn) {
+ return new CombiningValueStateSpec<InputT, AccumT, OutputT>(null, combineFn);
+ }
+
+ /**
+ * Create a state spec for values that use a {@link CombineFn} to automatically merge multiple
+ * {@code InputT}s into a single {@code OutputT}.
+ */
+ public static <InputT, AccumT, OutputT>
StateSpec<Object, AccumulatorCombiningState<InputT, AccumT, OutputT>> combiningValue(
Coder<AccumT> accumCoder, CombineFn<InputT, AccumT, OutputT> combineFn) {
+ checkArgument(accumCoder != null,
+ "accumCoder should not be null. "
+ + "Consider using combiningValue(CombineFn<> combineFn) instead.");
return combiningValueInternal(accumCoder, combineFn);
}
@@ -63,8 +85,21 @@ public class StateSpecs {
* multiple {@code InputT}s into a single {@code OutputT}.
*/
public static <K, InputT, AccumT, OutputT>
+ StateSpec<K, AccumulatorCombiningState<InputT, AccumT, OutputT>> keyedCombiningValue(
+ KeyedCombineFn<K, InputT, AccumT, OutputT> combineFn) {
+ return new KeyedCombiningValueStateSpec<K, InputT, AccumT, OutputT>(null, combineFn);
+ }
+
+ /**
+ * Create a state spec for values that use a {@link KeyedCombineFn} to automatically merge
+ * multiple {@code InputT}s into a single {@code OutputT}.
+ */
+ public static <K, InputT, AccumT, OutputT>
StateSpec<K, AccumulatorCombiningState<InputT, AccumT, OutputT>> keyedCombiningValue(
Coder<AccumT> accumCoder, KeyedCombineFn<K, InputT, AccumT, OutputT> combineFn) {
+ checkArgument(accumCoder != null,
+ "accumCoder should not be null. "
+ + "Consider using keyedCombiningValue(KeyedCombineFn<> combineFn) instead.");
return keyedCombiningValueInternal(accumCoder, combineFn);
}
@@ -73,10 +108,23 @@ public class StateSpecs {
* merge multiple {@code InputT}s into a single {@code OutputT}.
*/
public static <K, InputT, AccumT, OutputT>
+ StateSpec<K, AccumulatorCombiningState<InputT, AccumT, OutputT>>
+ keyedCombiningValueWithContext(KeyedCombineFnWithContext<K, InputT, AccumT, OutputT> combineFn) {
+ return new KeyedCombiningValueWithContextStateSpec<K, InputT, AccumT, OutputT>(null, combineFn);
+ }
+
+ /**
+ * Create a state spec for values that use a {@link KeyedCombineFnWithContext} to automatically
+ * merge multiple {@code InputT}s into a single {@code OutputT}.
+ */
+ public static <K, InputT, AccumT, OutputT>
StateSpec<K, AccumulatorCombiningState<InputT, AccumT, OutputT>>
keyedCombiningValueWithContext(
Coder<AccumT> accumCoder,
KeyedCombineFnWithContext<K, InputT, AccumT, OutputT> combineFn) {
+ checkArgument(accumCoder != null,
+ "accumCoder should not be null. Consider using "
+ + "keyedCombiningValueWithContext(KeyedCombineFnWithContext<> combineFn) instead.");
return new KeyedCombiningValueWithContextStateSpec<K, InputT, AccumT, OutputT>(
accumCoder, combineFn);
}
@@ -121,8 +169,23 @@ public class StateSpecs {
* Create a state spec that is optimized for adding values frequently, and occasionally retrieving
* all the values that have been added.
*/
+ public static <T> StateSpec<Object, BagState<T>> bag() {
+ return bag(null);
+ }
+
+ /**
+ * Create a state spec that is optimized for adding values frequently, and occasionally retrieving
+ * all the values that have been added.
+ */
public static <T> StateSpec<Object, BagState<T>> bag(Coder<T> elemCoder) {
- return new BagStateSpec<T>(elemCoder);
+ return new BagStateSpec<>(elemCoder);
+ }
+
+ /**
+ * Create a state spec that supporting for {@link java.util.Set} like access patterns.
+ */
+ public static <T> StateSpec<Object, SetState<T>> set() {
+ return set(null);
}
/**
@@ -135,6 +198,13 @@ public class StateSpecs {
/**
* Create a state spec that supporting for {@link java.util.Map} like access patterns.
*/
+ public static <K, V> StateSpec<Object, MapState<K, V>> map() {
+ return new MapStateSpec<>(null, null);
+ }
+
+ /**
+ * Create a state spec that supporting for {@link java.util.Map} like access patterns.
+ */
public static <K, V> StateSpec<Object, MapState<K, V>> map(Coder<K> keyCoder,
Coder<V> valueCoder) {
return new MapStateSpec<>(keyCoder, valueCoder);
@@ -174,9 +244,10 @@ public class StateSpecs {
*/
private static class ValueStateSpec<T> implements StateSpec<Object, ValueState<T>> {
- private final Coder<T> coder;
+ @Nullable
+ private Coder<T> coder;
- private ValueStateSpec(Coder<T> coder) {
+ private ValueStateSpec(@Nullable Coder<T> coder) {
this.coder = coder;
}
@@ -185,6 +256,25 @@ public class StateSpecs {
return visitor.bindValue(id, this, coder);
}
+ @SuppressWarnings("unchecked")
+ @Override
+ public void offerCoders(Coder[] coders) {
+ if (this.coder == null) {
+ if (coders[0] != null) {
+ this.coder = (Coder<T>) coders[0];
+ }
+ }
+ }
+
+ @Override public void finishSpecifying() {
+ if (coder == null) {
+ throw new IllegalStateException("Unable to infer a coder for ValueState and no Coder"
+ + " was specified. Please set a coder by either invoking"
+ + " StateSpecs.value(Coder<T> valueCoder) or by registering the coder in the"
+ + " Pipeline's CoderRegistry.");
+ }
+ }
+
@Override
public boolean equals(Object obj) {
if (obj == this) {
@@ -214,15 +304,32 @@ public class StateSpecs {
extends KeyedCombiningValueStateSpec<Object, InputT, AccumT, OutputT>
implements StateSpec<Object, AccumulatorCombiningState<InputT, AccumT, OutputT>> {
- private final Coder<AccumT> accumCoder;
+ @Nullable
+ private Coder<AccumT> accumCoder;
private final CombineFn<InputT, AccumT, OutputT> combineFn;
private CombiningValueStateSpec(
- Coder<AccumT> accumCoder, CombineFn<InputT, AccumT, OutputT> combineFn) {
+ @Nullable Coder<AccumT> accumCoder,
+ CombineFn<InputT, AccumT, OutputT> combineFn) {
super(accumCoder, combineFn.asKeyedFn());
this.combineFn = combineFn;
this.accumCoder = accumCoder;
}
+
+ @Override
+ protected Coder<AccumT> getAccumCoder() {
+ return accumCoder;
+ }
+
+ @SuppressWarnings("unchecked")
+ @Override
+ public void offerCoders(Coder[] coders) {
+ if (this.accumCoder == null) {
+ if (coders[1] != null) {
+ this.accumCoder = (Coder<AccumT>) coders[1];
+ }
+ }
+ }
}
/**
@@ -234,11 +341,13 @@ public class StateSpecs {
private static class KeyedCombiningValueWithContextStateSpec<K, InputT, AccumT, OutputT>
implements StateSpec<K, AccumulatorCombiningState<InputT, AccumT, OutputT>> {
- private final Coder<AccumT> accumCoder;
+ @Nullable
+ private Coder<AccumT> accumCoder;
private final KeyedCombineFnWithContext<K, InputT, AccumT, OutputT> combineFn;
protected KeyedCombiningValueWithContextStateSpec(
- Coder<AccumT> accumCoder, KeyedCombineFnWithContext<K, InputT, AccumT, OutputT> combineFn) {
+ @Nullable Coder<AccumT> accumCoder,
+ KeyedCombineFnWithContext<K, InputT, AccumT, OutputT> combineFn) {
this.combineFn = combineFn;
this.accumCoder = accumCoder;
}
@@ -249,6 +358,27 @@ public class StateSpecs {
return visitor.bindKeyedCombiningValueWithContext(id, this, accumCoder, combineFn);
}
+ @SuppressWarnings("unchecked")
+ @Override
+ public void offerCoders(Coder[] coders) {
+ if (this.accumCoder == null) {
+ if (coders[2] != null) {
+ this.accumCoder = (Coder<AccumT>) coders[2];
+ }
+ }
+ }
+
+ @Override public void finishSpecifying() {
+ if (accumCoder == null) {
+ throw new IllegalStateException("Unable to infer a coder for"
+ + " KeyedCombiningValueWithContextState and no Coder was specified."
+ + " Please set a coder by either invoking"
+ + " StateSpecs.keyedCombiningValue(Coder<AccumT> accumCoder,"
+ + " KeyedCombineFn<K, InputT, AccumT, OutputT> combineFn)"
+ + " or by registering the coder in the Pipeline's CoderRegistry.");
+ }
+ }
+
@Override
public boolean equals(Object obj) {
if (obj == this) {
@@ -282,19 +412,45 @@ public class StateSpecs {
private static class KeyedCombiningValueStateSpec<K, InputT, AccumT, OutputT>
implements StateSpec<K, AccumulatorCombiningState<InputT, AccumT, OutputT>> {
- private final Coder<AccumT> accumCoder;
+ @Nullable
+ private Coder<AccumT> accumCoder;
private final KeyedCombineFn<K, InputT, AccumT, OutputT> keyedCombineFn;
protected KeyedCombiningValueStateSpec(
- Coder<AccumT> accumCoder, KeyedCombineFn<K, InputT, AccumT, OutputT> keyedCombineFn) {
+ @Nullable Coder<AccumT> accumCoder,
+ KeyedCombineFn<K, InputT, AccumT, OutputT> keyedCombineFn) {
this.keyedCombineFn = keyedCombineFn;
this.accumCoder = accumCoder;
}
+ protected Coder<AccumT> getAccumCoder() {
+ return accumCoder;
+ }
+
@Override
public AccumulatorCombiningState<InputT, AccumT, OutputT> bind(
String id, StateBinder<? extends K> visitor) {
- return visitor.bindKeyedCombiningValue(id, this, accumCoder, keyedCombineFn);
+ return visitor.bindKeyedCombiningValue(id, this, getAccumCoder(), keyedCombineFn);
+ }
+
+ @SuppressWarnings("unchecked")
+ @Override
+ public void offerCoders(Coder[] coders) {
+ if (this.accumCoder == null) {
+ if (coders[2] != null) {
+ this.accumCoder = (Coder<AccumT>) coders[2];
+ }
+ }
+ }
+
+ @Override public void finishSpecifying() {
+ if (getAccumCoder() == null) {
+ throw new IllegalStateException("Unable to infer a coder for CombiningState and no"
+ + " Coder was specified. Please set a coder by either invoking"
+ + " StateSpecs.combiningValue(Coder<AccumT> accumCoder,"
+ + " CombineFn<InputT, AccumT, OutputT> combineFn)"
+ + " or by registering the coder in the Pipeline's CoderRegistry.");
+ }
}
@Override
@@ -330,9 +486,10 @@ public class StateSpecs {
*/
private static class BagStateSpec<T> implements StateSpec<Object, BagState<T>> {
- private final Coder<T> elemCoder;
+ @Nullable
+ private Coder<T> elemCoder;
- private BagStateSpec(Coder<T> elemCoder) {
+ private BagStateSpec(@Nullable Coder<T> elemCoder) {
this.elemCoder = elemCoder;
}
@@ -341,6 +498,25 @@ public class StateSpecs {
return visitor.bindBag(id, this, elemCoder);
}
+ @SuppressWarnings("unchecked")
+ @Override
+ public void offerCoders(Coder[] coders) {
+ if (this.elemCoder == null) {
+ if (coders[0] != null) {
+ this.elemCoder = (Coder<T>) coders[0];
+ }
+ }
+ }
+
+ @Override public void finishSpecifying() {
+ if (elemCoder == null) {
+ throw new IllegalStateException("Unable to infer a coder for BagState and no Coder"
+ + " was specified. Please set a coder by either invoking"
+ + " StateSpecs.bag(Coder<T> elemCoder) or by registering the coder in the"
+ + " Pipeline's CoderRegistry.");
+ }
+ }
+
@Override
public boolean equals(Object obj) {
if (obj == this) {
@@ -363,10 +539,12 @@ public class StateSpecs {
private static class MapStateSpec<K, V> implements StateSpec<Object, MapState<K, V>> {
- private final Coder<K> keyCoder;
- private final Coder<V> valueCoder;
+ @Nullable
+ private Coder<K> keyCoder;
+ @Nullable
+ private Coder<V> valueCoder;
- private MapStateSpec(Coder<K> keyCoder, Coder<V> valueCoder) {
+ private MapStateSpec(@Nullable Coder<K> keyCoder, @Nullable Coder<V> valueCoder) {
this.keyCoder = keyCoder;
this.valueCoder = valueCoder;
}
@@ -376,6 +554,30 @@ public class StateSpecs {
return visitor.bindMap(id, this, keyCoder, valueCoder);
}
+ @SuppressWarnings("unchecked")
+ @Override
+ public void offerCoders(Coder[] coders) {
+ if (this.keyCoder == null) {
+ if (coders[0] != null) {
+ this.keyCoder = (Coder<K>) coders[0];
+ }
+ }
+ if (this.valueCoder == null) {
+ if (coders[1] != null) {
+ this.valueCoder = (Coder<V>) coders[1];
+ }
+ }
+ }
+
+ @Override public void finishSpecifying() {
+ if (keyCoder == null || valueCoder == null) {
+ throw new IllegalStateException("Unable to infer a coder for MapState and no Coder"
+ + " was specified. Please set a coder by either invoking"
+ + " StateSpecs.map(Coder<K> keyCoder, Coder<V> valueCoder) or by registering the"
+ + " coder in the Pipeline's CoderRegistry.");
+ }
+ }
+
@Override
public boolean equals(Object obj) {
if (obj == this) {
@@ -404,9 +606,10 @@ public class StateSpecs {
*/
private static class SetStateSpec<T> implements StateSpec<Object, SetState<T>> {
- private final Coder<T> elemCoder;
+ @Nullable
+ private Coder<T> elemCoder;
- private SetStateSpec(Coder<T> elemCoder) {
+ private SetStateSpec(@Nullable Coder<T> elemCoder) {
this.elemCoder = elemCoder;
}
@@ -415,6 +618,25 @@ public class StateSpecs {
return visitor.bindSet(id, this, elemCoder);
}
+ @SuppressWarnings("unchecked")
+ @Override
+ public void offerCoders(Coder[] coders) {
+ if (this.elemCoder == null) {
+ if (coders[0] != null) {
+ this.elemCoder = (Coder<T>) coders[0];
+ }
+ }
+ }
+
+ @Override public void finishSpecifying() {
+ if (elemCoder == null) {
+ throw new IllegalStateException("Unable to infer a coder for SetState and no Coder"
+ + " was specified. Please set a coder by either invoking"
+ + " StateSpecs.set(Coder<T> elemCoder) or by registering the coder in the"
+ + " Pipeline's CoderRegistry.");
+ }
+ }
+
@Override
public boolean equals(Object obj) {
if (obj == this) {
@@ -461,6 +683,14 @@ public class StateSpecs {
}
@Override
+ public void offerCoders(Coder[] coders) {
+ }
+
+ @Override public void finishSpecifying() {
+ // Currently an empty implementation as there are no coders to validate.
+ }
+
+ @Override
public boolean equals(Object obj) {
if (obj == this) {
return true;
http://git-wip-us.apache.org/repos/asf/beam/blob/42e690e8/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
----------------------------------------------------------------------
diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
index cbbbe5f..4249a77 100644
--- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
+++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java
@@ -39,6 +39,7 @@ import static org.junit.Assert.assertThat;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.google.common.base.MoreObjects;
+import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Sets;
@@ -55,8 +56,12 @@ import java.util.Map;
import java.util.Set;
import org.apache.beam.sdk.Pipeline.PipelineExecutionException;
import org.apache.beam.sdk.coders.AtomicCoder;
+import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderException;
+import org.apache.beam.sdk.coders.CustomCoder;
import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.ListCoder;
+import org.apache.beam.sdk.coders.SetCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.coders.VarIntCoder;
import org.apache.beam.sdk.io.CountingInput;
@@ -1036,6 +1041,71 @@ public class ParDoTest implements Serializable {
}
}
+ private static class MyInteger implements Comparable<MyInteger> {
+ private final int value;
+
+ MyInteger(int value) {
+ this.value = value;
+ }
+
+ public int getValue() {
+ return value;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+
+ if (!(o instanceof MyInteger)) {
+ return false;
+ }
+
+ MyInteger myInteger = (MyInteger) o;
+
+ return value == myInteger.value;
+
+ }
+
+ @Override
+ public int hashCode() {
+ return value;
+ }
+
+ @Override
+ public int compareTo(MyInteger o) {
+ return Integer.compare(this.getValue(), o.getValue());
+ }
+
+ @Override
+ public String toString() {
+ return "MyInteger{" + "value=" + value + '}';
+ }
+ }
+
+ private static class MyIntegerCoder extends CustomCoder<MyInteger> {
+ private static final MyIntegerCoder INSTANCE = new MyIntegerCoder();
+
+ private final VarIntCoder delegate = VarIntCoder.of();
+
+ public static MyIntegerCoder of() {
+ return INSTANCE;
+ }
+
+ @Override
+ public void encode(MyInteger value, OutputStream outStream, Context context)
+ throws CoderException, IOException {
+ delegate.encode(value.getValue(), outStream, context);
+ }
+
+ @Override
+ public MyInteger decode(InputStream inStream, Context context) throws CoderException,
+ IOException {
+ return new MyInteger(delegate.decode(inStream, context));
+ }
+ }
+
/** PAssert "matcher" for expected output. */
static class HasExpectedOutput
implements SerializableFunction<Iterable<String>, Void>, Serializable {
@@ -1619,6 +1689,132 @@ public class ParDoTest implements Serializable {
@Test
@Category({ValidatesRunner.class, UsesStatefulParDo.class})
+ public void testValueStateCoderInference() {
+ final String stateId = "foo";
+ MyIntegerCoder myIntegerCoder = MyIntegerCoder.of();
+ pipeline.getCoderRegistry().registerCoder(MyInteger.class, myIntegerCoder);
+
+ DoFn<KV<String, Integer>, MyInteger> fn =
+ new DoFn<KV<String, Integer>, MyInteger>() {
+
+ @StateId(stateId)
+ private final StateSpec<Object, ValueState<MyInteger>> intState =
+ StateSpecs.value();
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c, @StateId(stateId) ValueState<MyInteger> state) {
+ MyInteger currentValue = MoreObjects.firstNonNull(state.read(), new MyInteger(0));
+ c.output(currentValue);
+ state.write(new MyInteger(currentValue.getValue() + 1));
+ }
+ };
+
+ PCollection<MyInteger> output =
+ pipeline.apply(Create.of(KV.of("hello", 42), KV.of("hello", 97), KV.of("hello", 84)))
+ .apply(ParDo.of(fn)).setCoder(myIntegerCoder);
+
+ PAssert.that(output).containsInAnyOrder(new MyInteger(0), new MyInteger(1), new MyInteger(2));
+ pipeline.run();
+ }
+
+ @Test
+ @Category({ValidatesRunner.class, UsesStatefulParDo.class})
+ public void testValueStateCoderInferenceFailure() throws Exception {
+ final String stateId = "foo";
+ MyIntegerCoder myIntegerCoder = MyIntegerCoder.of();
+
+ DoFn<KV<String, Integer>, MyInteger> fn =
+ new DoFn<KV<String, Integer>, MyInteger>() {
+
+ @StateId(stateId)
+ private final StateSpec<Object, ValueState<MyInteger>> intState =
+ StateSpecs.value();
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c, @StateId(stateId) ValueState<MyInteger> state) {
+ MyInteger currentValue = MoreObjects.firstNonNull(state.read(), new MyInteger(0));
+ c.output(currentValue);
+ state.write(new MyInteger(currentValue.getValue() + 1));
+ }
+ };
+
+ thrown.expect(RuntimeException.class);
+ thrown.expectMessage("Unable to infer a coder for ValueState and no Coder was specified.");
+
+ pipeline.apply(Create.of(KV.of("hello", 42), KV.of("hello", 97), KV.of("hello", 84)))
+ .apply(ParDo.of(fn)).setCoder(myIntegerCoder);
+
+ pipeline.run();
+ }
+
+ @Test
+ @Category({ValidatesRunner.class, UsesStatefulParDo.class})
+ public void testValueStateCoderInferenceFromInputCoder() {
+ final String stateId = "foo";
+ MyIntegerCoder myIntegerCoder = MyIntegerCoder.of();
+
+ DoFn<KV<String, MyInteger>, MyInteger> fn =
+ new DoFn<KV<String, MyInteger>, MyInteger>() {
+
+ @StateId(stateId)
+ private final StateSpec<Object, ValueState<MyInteger>> intState =
+ StateSpecs.value();
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c, @StateId(stateId) ValueState<MyInteger> state) {
+ MyInteger currentValue = MoreObjects.firstNonNull(state.read(), new MyInteger(0));
+ c.output(currentValue);
+ state.write(new MyInteger(currentValue.getValue() + 1));
+ }
+ };
+
+ pipeline
+ .apply(Create.of(KV.of("hello", new MyInteger(42)),
+ KV.of("hello", new MyInteger(97)), KV.of("hello", new MyInteger(84)))
+ .withCoder(KvCoder.of(StringUtf8Coder.of(), myIntegerCoder)))
+ .apply(ParDo.of(fn)).setCoder(myIntegerCoder);
+
+ pipeline.run();
+ }
+
+ @Test
+ @Category({ValidatesRunner.class, UsesStatefulParDo.class})
+ public void testCoderInferenceOfList() {
+ final String stateId = "foo";
+ MyIntegerCoder myIntegerCoder = MyIntegerCoder.of();
+ pipeline.getCoderRegistry().registerCoder(MyInteger.class, myIntegerCoder);
+
+ DoFn<KV<String, Integer>, List<MyInteger>> fn =
+ new DoFn<KV<String, Integer>, List<MyInteger>>() {
+
+ @StateId(stateId)
+ private final StateSpec<Object, ValueState<List<MyInteger>>> intState =
+ StateSpecs.value();
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c, @StateId(stateId) ValueState<List<MyInteger>> state) {
+ MyInteger myInteger = new MyInteger(c.element().getValue());
+ List<MyInteger> currentValue = state.read();
+ List<MyInteger> newValue = currentValue != null
+ ? ImmutableList.<MyInteger>builder().addAll(currentValue).add(myInteger).build()
+ : Collections.singletonList(myInteger);
+ c.output(newValue);
+ state.write(newValue);
+ }
+ };
+
+ pipeline.apply(Create.of(KV.of("hello", 42), KV.of("hello", 97), KV.of("hello", 84)))
+ .apply(ParDo.of(fn)).setCoder(ListCoder.of(myIntegerCoder));
+
+ pipeline.run();
+ }
+
+ @Test
+ @Category({ValidatesRunner.class, UsesStatefulParDo.class})
public void testValueStateFixedWindows() {
final String stateId = "foo";
@@ -1801,6 +1997,82 @@ public class ParDoTest implements Serializable {
}
@Test
+ @Category({ValidatesRunner.class, UsesStatefulParDo.class})
+ public void testBagStateCoderInference() {
+ final String stateId = "foo";
+ Coder<MyInteger> myIntegerCoder = MyIntegerCoder.of();
+ pipeline.getCoderRegistry().registerCoder(MyInteger.class, myIntegerCoder);
+
+ DoFn<KV<String, Integer>, List<MyInteger>> fn =
+ new DoFn<KV<String, Integer>, List<MyInteger>>() {
+
+ @StateId(stateId)
+ private final StateSpec<Object, BagState<MyInteger>> bufferState =
+ StateSpecs.bag();
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c, @StateId(stateId) BagState<MyInteger> state) {
+ Iterable<MyInteger> currentValue = state.read();
+ state.add(new MyInteger(c.element().getValue()));
+ if (Iterables.size(state.read()) >= 4) {
+ List<MyInteger> sorted = Lists.newArrayList(currentValue);
+ Collections.sort(sorted);
+ c.output(sorted);
+ }
+ }
+ };
+
+ PCollection<List<MyInteger>> output =
+ pipeline.apply(
+ Create.of(
+ KV.of("hello", 97), KV.of("hello", 42), KV.of("hello", 84), KV.of("hello", 12)))
+ .apply(ParDo.of(fn)).setCoder(ListCoder.of(myIntegerCoder));
+
+ PAssert.that(output).containsInAnyOrder(Lists.newArrayList(new MyInteger(12), new MyInteger(42),
+ new MyInteger(84), new MyInteger(97)));
+
+ pipeline.run();
+ }
+
+ @Test
+ @Category({ValidatesRunner.class, UsesStatefulParDo.class})
+ public void testBagStateCoderInferenceFailure() throws Exception {
+ final String stateId = "foo";
+ Coder<MyInteger> myIntegerCoder = MyIntegerCoder.of();
+
+ DoFn<KV<String, Integer>, List<MyInteger>> fn =
+ new DoFn<KV<String, Integer>, List<MyInteger>>() {
+
+ @StateId(stateId)
+ private final StateSpec<Object, BagState<MyInteger>> bufferState =
+ StateSpecs.bag();
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c, @StateId(stateId) BagState<MyInteger> state) {
+ Iterable<MyInteger> currentValue = state.read();
+ state.add(new MyInteger(c.element().getValue()));
+ if (Iterables.size(state.read()) >= 4) {
+ List<MyInteger> sorted = Lists.newArrayList(currentValue);
+ Collections.sort(sorted);
+ c.output(sorted);
+ }
+ }
+ };
+
+ thrown.expect(RuntimeException.class);
+ thrown.expectMessage("Unable to infer a coder for BagState and no Coder was specified.");
+
+ pipeline.apply(
+ Create.of(
+ KV.of("hello", 97), KV.of("hello", 42), KV.of("hello", 84), KV.of("hello", 12)))
+ .apply(ParDo.of(fn)).setCoder(ListCoder.of(myIntegerCoder));
+
+ pipeline.run();
+ }
+
+ @Test
@Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesSetState.class})
public void testSetState() {
final String stateId = "foo";
@@ -1843,6 +2115,93 @@ public class ParDoTest implements Serializable {
}
@Test
+ @Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesSetState.class})
+ public void testSetStateCoderInference() {
+ final String stateId = "foo";
+ final String countStateId = "count";
+ Coder<MyInteger> myIntegerCoder = MyIntegerCoder.of();
+ pipeline.getCoderRegistry().registerCoder(MyInteger.class, myIntegerCoder);
+
+ DoFn<KV<String, Integer>, Set<MyInteger>> fn =
+ new DoFn<KV<String, Integer>, Set<MyInteger>>() {
+
+ @StateId(stateId)
+ private final StateSpec<Object, SetState<MyInteger>> setState = StateSpecs.set();
+
+ @StateId(countStateId)
+ private final StateSpec<Object, AccumulatorCombiningState<Integer, int[], Integer>>
+ countState = StateSpecs.combiningValueFromInputInternal(VarIntCoder.of(),
+ Sum.ofIntegers());
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ @StateId(stateId) SetState<MyInteger> state,
+ @StateId(countStateId) AccumulatorCombiningState<Integer, int[], Integer> count) {
+ state.add(new MyInteger(c.element().getValue()));
+ count.add(1);
+ if (count.read() >= 4) {
+ Set<MyInteger> set = Sets.newHashSet(state.read());
+ c.output(set);
+ }
+ }
+ };
+
+ PCollection<Set<MyInteger>> output =
+ pipeline.apply(
+ Create.of(
+ KV.of("hello", 97), KV.of("hello", 42), KV.of("hello", 42), KV.of("hello", 12)))
+ .apply(ParDo.of(fn)).setCoder(SetCoder.of(myIntegerCoder));
+
+ PAssert.that(output).containsInAnyOrder(
+ Sets.newHashSet(new MyInteger(97), new MyInteger(42), new MyInteger(12)));
+ pipeline.run();
+ }
+
+ @Test
+ @Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesSetState.class})
+ public void testSetStateCoderInferenceFailure() throws Exception {
+ final String stateId = "foo";
+ final String countStateId = "count";
+ Coder<MyInteger> myIntegerCoder = MyIntegerCoder.of();
+
+ DoFn<KV<String, Integer>, Set<MyInteger>> fn =
+ new DoFn<KV<String, Integer>, Set<MyInteger>>() {
+
+ @StateId(stateId)
+ private final StateSpec<Object, SetState<MyInteger>> setState = StateSpecs.set();
+
+ @StateId(countStateId)
+ private final StateSpec<Object, AccumulatorCombiningState<Integer, int[], Integer>>
+ countState = StateSpecs.combiningValueFromInputInternal(VarIntCoder.of(),
+ Sum.ofIntegers());
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ @StateId(stateId) SetState<MyInteger> state,
+ @StateId(countStateId) AccumulatorCombiningState<Integer, int[], Integer> count) {
+ state.add(new MyInteger(c.element().getValue()));
+ count.add(1);
+ if (count.read() >= 4) {
+ Set<MyInteger> set = Sets.newHashSet(state.read());
+ c.output(set);
+ }
+ }
+ };
+
+ thrown.expect(RuntimeException.class);
+ thrown.expectMessage("Unable to infer a coder for SetState and no Coder was specified.");
+
+ pipeline.apply(
+ Create.of(
+ KV.of("hello", 97), KV.of("hello", 42), KV.of("hello", 42), KV.of("hello", 12)))
+ .apply(ParDo.of(fn)).setCoder(SetCoder.of(myIntegerCoder));
+
+ pipeline.run();
+ }
+
+ @Test
@Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesMapState.class})
public void testMapState() {
final String stateId = "foo";
@@ -1888,6 +2247,99 @@ public class ParDoTest implements Serializable {
}
@Test
+ @Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesMapState.class})
+ public void testMapStateCoderInference() {
+ final String stateId = "foo";
+ final String countStateId = "count";
+ Coder<MyInteger> myIntegerCoder = MyIntegerCoder.of();
+ pipeline.getCoderRegistry().registerCoder(MyInteger.class, myIntegerCoder);
+
+ DoFn<KV<String, KV<String, Integer>>, KV<String, MyInteger>> fn =
+ new DoFn<KV<String, KV<String, Integer>>, KV<String, MyInteger>>() {
+
+ @StateId(stateId)
+ private final StateSpec<Object, MapState<String, MyInteger>> mapState = StateSpecs.map();
+ @StateId(countStateId)
+ private final StateSpec<Object, AccumulatorCombiningState<Integer, int[], Integer>>
+ countState = StateSpecs.combiningValueFromInputInternal(VarIntCoder.of(),
+ Sum.ofIntegers());
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c, @StateId(stateId) MapState<String, MyInteger> state,
+ @StateId(countStateId) AccumulatorCombiningState<Integer, int[], Integer>
+ count) {
+ KV<String, Integer> value = c.element().getValue();
+ state.put(value.getKey(), new MyInteger(value.getValue()));
+ count.add(1);
+ if (count.read() >= 4) {
+ Iterable<Map.Entry<String, MyInteger>> iterate = state.iterate();
+ for (Map.Entry<String, MyInteger> entry : iterate) {
+ c.output(KV.of(entry.getKey(), entry.getValue()));
+ }
+ }
+ }
+ };
+
+ PCollection<KV<String, MyInteger>> output =
+ pipeline.apply(
+ Create.of(
+ KV.of("hello", KV.of("a", 97)), KV.of("hello", KV.of("b", 42)),
+ KV.of("hello", KV.of("b", 42)), KV.of("hello", KV.of("c", 12))))
+ .apply(ParDo.of(fn)).setCoder(KvCoder.of(StringUtf8Coder.of(), myIntegerCoder));
+
+ PAssert.that(output).containsInAnyOrder(KV.of("a", new MyInteger(97)),
+ KV.of("b", new MyInteger(42)), KV.of("c", new MyInteger(12)));
+ pipeline.run();
+ }
+
+ @Test
+ @Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesMapState.class})
+ public void testMapStateCoderInferenceFailure() throws Exception {
+ final String stateId = "foo";
+ final String countStateId = "count";
+ Coder<MyInteger> myIntegerCoder = MyIntegerCoder.of();
+
+ DoFn<KV<String, KV<String, Integer>>, KV<String, MyInteger>> fn =
+ new DoFn<KV<String, KV<String, Integer>>, KV<String, MyInteger>>() {
+
+ @StateId(stateId)
+ private final StateSpec<Object, MapState<String, MyInteger>> mapState = StateSpecs.map();
+ @StateId(countStateId)
+ private final StateSpec<Object, AccumulatorCombiningState<Integer, int[], Integer>>
+ countState = StateSpecs.combiningValueFromInputInternal(VarIntCoder.of(),
+ Sum.ofIntegers());
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c, @StateId(stateId) MapState<String, MyInteger> state,
+ @StateId(countStateId) AccumulatorCombiningState<Integer, int[], Integer>
+ count) {
+ KV<String, Integer> value = c.element().getValue();
+ state.put(value.getKey(), new MyInteger(value.getValue()));
+ count.add(1);
+ if (count.read() >= 4) {
+ Iterable<Map.Entry<String, MyInteger>> iterate = state.iterate();
+ for (Map.Entry<String, MyInteger> entry : iterate) {
+ c.output(KV.of(entry.getKey(), entry.getValue()));
+ }
+ }
+ }
+ };
+
+ thrown.expect(RuntimeException.class);
+ thrown.expectMessage("Unable to infer a coder for MapState and no Coder was specified.");
+
+ pipeline.apply(
+ Create.of(
+ KV.of("hello", KV.of("a", 97)), KV.of("hello", KV.of("b", 42)),
+ KV.of("hello", KV.of("b", 42)), KV.of("hello", KV.of("c", 12))))
+ .apply(ParDo.of(fn)).setCoder(KvCoder.of(StringUtf8Coder.of(), myIntegerCoder));
+
+ pipeline.run();
+ }
+
+ @Test
@Category({ValidatesRunner.class, UsesStatefulParDo.class})
public void testCombiningState() {
final String stateId = "foo";
@@ -1928,6 +2380,132 @@ public class ParDoTest implements Serializable {
@Test
@Category({ValidatesRunner.class, UsesStatefulParDo.class})
+ public void testCombiningStateCoderInference() {
+ pipeline.getCoderRegistry().registerCoder(MyInteger.class, MyIntegerCoder.of());
+
+ final String stateId = "foo";
+
+ DoFn<KV<String, Integer>, String> fn =
+ new DoFn<KV<String, Integer>, String>() {
+ private static final int EXPECTED_SUM = 16;
+
+ @StateId(stateId)
+ private final StateSpec<
+ Object, AccumulatorCombiningState<Integer, MyInteger, Integer>>
+ combiningState =
+ StateSpecs.combiningValue(new Combine.CombineFn<Integer, MyInteger, Integer>() {
+ @Override
+ public MyInteger createAccumulator() {
+ return new MyInteger(0);
+ }
+
+ @Override
+ public MyInteger addInput(MyInteger accumulator, Integer input) {
+ return new MyInteger(accumulator.getValue() + input);
+ }
+
+ @Override
+ public MyInteger mergeAccumulators(Iterable<MyInteger> accumulators) {
+ int newValue = 0;
+ for (MyInteger myInteger : accumulators) {
+ newValue += myInteger.getValue();
+ }
+ return new MyInteger(newValue);
+ }
+
+ @Override
+ public Integer extractOutput(MyInteger accumulator) {
+ return accumulator.getValue();
+ }
+ });
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ @StateId(stateId)
+ AccumulatorCombiningState<Integer, MyInteger, Integer> state) {
+ state.add(c.element().getValue());
+ Integer currentValue = state.read();
+ if (currentValue == EXPECTED_SUM) {
+ c.output("right on");
+ }
+ }
+ };
+
+ PCollection<String> output =
+ pipeline
+ .apply(Create.of(KV.of("hello", 3), KV.of("hello", 6), KV.of("hello", 7)))
+ .apply(ParDo.of(fn));
+
+ // There should only be one moment at which the average is exactly 16
+ PAssert.that(output).containsInAnyOrder("right on");
+ pipeline.run();
+ }
+
+ @Test
+ @Category({ValidatesRunner.class, UsesStatefulParDo.class})
+ public void testCombiningStateCoderInferenceFailure() throws Exception {
+ final String stateId = "foo";
+
+ DoFn<KV<String, Integer>, String> fn =
+ new DoFn<KV<String, Integer>, String>() {
+ private static final int EXPECTED_SUM = 16;
+
+ @StateId(stateId)
+ private final StateSpec<
+ Object, AccumulatorCombiningState<Integer, MyInteger, Integer>>
+ combiningState =
+ StateSpecs.combiningValue(new Combine.CombineFn<Integer, MyInteger, Integer>() {
+ @Override
+ public MyInteger createAccumulator() {
+ return new MyInteger(0);
+ }
+
+ @Override
+ public MyInteger addInput(MyInteger accumulator, Integer input) {
+ return new MyInteger(accumulator.getValue() + input);
+ }
+
+ @Override
+ public MyInteger mergeAccumulators(Iterable<MyInteger> accumulators) {
+ int newValue = 0;
+ for (MyInteger myInteger : accumulators) {
+ newValue += myInteger.getValue();
+ }
+ return new MyInteger(newValue);
+ }
+
+ @Override
+ public Integer extractOutput(MyInteger accumulator) {
+ return accumulator.getValue();
+ }
+ });
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext c,
+ @StateId(stateId)
+ AccumulatorCombiningState<Integer, MyInteger, Integer> state) {
+ state.add(c.element().getValue());
+ Integer currentValue = state.read();
+ if (currentValue == EXPECTED_SUM) {
+ c.output("right on");
+ }
+ }
+ };
+
+ thrown.expect(RuntimeException.class);
+ thrown.expectMessage("Unable to infer a coder for CombiningState and no Coder was specified.");
+
+ pipeline
+ .apply(Create.of(KV.of("hello", 3), KV.of("hello", 6), KV.of("hello", 7)))
+ .apply(ParDo.of(fn));
+
+ pipeline.run();
+ }
+
+ @Test
+ @Category({ValidatesRunner.class, UsesStatefulParDo.class})
public void testBagStateSideInput() {
final PCollectionView<List<Integer>> listView =
[2/2] beam git commit: This closes #2133
Posted by av...@apache.org.
This closes #2133
Project: http://git-wip-us.apache.org/repos/asf/beam/repo
Commit: http://git-wip-us.apache.org/repos/asf/beam/commit/e31ca8b0
Tree: http://git-wip-us.apache.org/repos/asf/beam/tree/e31ca8b0
Diff: http://git-wip-us.apache.org/repos/asf/beam/diff/e31ca8b0
Branch: refs/heads/master
Commit: e31ca8b0d05e47c2588d5db29c92bac49aa410da
Parents: 03dce6d 42e690e
Author: Aviem Zur <av...@gmail.com>
Authored: Sat Apr 1 10:28:25 2017 +0300
Committer: Aviem Zur <av...@gmail.com>
Committed: Sat Apr 1 10:28:25 2017 +0300
----------------------------------------------------------------------
.../org/apache/beam/sdk/transforms/ParDo.java | 62 ++
.../apache/beam/sdk/util/state/StateSpec.java | 15 +
.../apache/beam/sdk/util/state/StateSpecs.java | 264 ++++++++-
.../apache/beam/sdk/transforms/ParDoTest.java | 578 +++++++++++++++++++
4 files changed, 902 insertions(+), 17 deletions(-)
----------------------------------------------------------------------