@Test
public void testCreateNormalParDoFn() throws Exception {
String stringState = "some state";
long longState = 42L;
TestDoFn fn = new TestDoFn(stringState, longState);
String serializedFn =
StringUtils.byteArrayToJsonString(
SerializableUtils.serializeToByteArray(
new DoFnInfo(fn, WindowingStrategy.globalDefault())));
CloudObject cloudUserFn = CloudObject.forClassName("DoFn");
addString(cloudUserFn, "serialized_fn", serializedFn);
String tag = "output";
MultiOutputInfo multiOutputInfo = new MultiOutputInfo();
multiOutputInfo.setTag(tag);
List<MultiOutputInfo> multiOutputInfos = Arrays.asList(multiOutputInfo);
PipelineOptions options = PipelineOptionsFactory.create();
DataflowExecutionContext context = BatchModeExecutionContext.fromOptions(options);
CounterSet counters = new CounterSet();
StateSampler stateSampler = new StateSampler("test", counters.getAddCounterMutator());
ParDoFn parDoFn =
factory.create(
options,
cloudUserFn,
"name",
"transformName",
null,
multiOutputInfos,
1,
context,
counters.getAddCounterMutator(),
stateSampler);
// Test that the factory created the correct class
assertThat(parDoFn, instanceOf(NormalParDoFn.class));
// Test that the DoFnInfo reflects the one passed in
NormalParDoFn normalParDoFn = (NormalParDoFn) parDoFn;
DoFnInfo doFnInfo = normalParDoFn.getDoFnInfo();
DoFn actualDoFn = doFnInfo.getDoFn();
assertThat(actualDoFn, instanceOf(TestDoFn.class));
assertThat(doFnInfo.getWindowingStrategy().getWindowFn(), instanceOf(GlobalWindows.class));
assertThat(
doFnInfo.getWindowingStrategy().getTrigger().getSpec(), instanceOf(DefaultTrigger.class));
// Test that the deserialized user DoFn is as expected
TestDoFn actualTestDoFn = (TestDoFn) actualDoFn;
assertEquals(stringState, actualTestDoFn.stringState);
assertEquals(longState, actualTestDoFn.longState);
assertEquals(context, normalParDoFn.getExecutionContext());
}
@Override
public PCollectionTuple apply(PCollection<Integer> input) {
PCollection<Integer> sum = input.apply(Sum.integersGlobally());
// Fails here when attempting to construct a tuple with an unbound object.
return PCollectionTuple.of(sumTag, sum)
.and(
doneTag,
PCollection.<Void>createPrimitiveOutputInternal(
input.getPipeline(), WindowingStrategy.globalDefault(), input.isBounded()));
}
@Override
public PCollection<T> apply(PInput input) {
if (topic == null && subscription == null) {
throw new IllegalStateException(
"need to set either the topic or the subscription for "
+ "a PubsubIO.Read transform");
}
if (topic != null && subscription != null) {
throw new IllegalStateException(
"Can't set both the topic and the subscription for a " + "PubsubIO.Read transform");
}
return PCollection.<T>createPrimitiveOutputInternal(
input.getPipeline(), WindowingStrategy.globalDefault(), IsBounded.UNBOUNDED)
.setCoder(coder);
}
protected <T> WindowedValue<T> makeWindowedValue(
T output, Instant timestamp, Collection<? extends BoundedWindow> windows, PaneInfo pane) {
final Instant inputTimestamp = timestamp;
final WindowFn windowFn = windowingStrategy.getWindowFn();
if (timestamp == null) {
timestamp = BoundedWindow.TIMESTAMP_MIN_VALUE;
}
if (windows == null) {
try {
windows =
windowFn.assignWindows(
windowFn.new AssignContext() {
@Override
public Object element() {
throw new UnsupportedOperationException(
"WindowFn attempted to access input element when none was available"); // TODO: 12/16/15 aljoscha's comment in slack
}
@Override
public Instant timestamp() {
if (inputTimestamp == null) {
throw new UnsupportedOperationException(
"WindowFn attempted to access input timestamp when none was available");
}
return inputTimestamp;
}
@Override
public Collection<? extends BoundedWindow> windows() {
throw new UnsupportedOperationException(
"WindowFn attempted to access input windows when none were available");
}
});
} catch (Exception e) {
Throwables.propagateIfInstanceOf(e, UserCodeException.class);
throw new UserCodeException(e);
}
}
return WindowedValue.of(output, timestamp, windows, pane);
}
@Override
public PCollection<String> apply(PCollection<String> input) {
return PCollection.createPrimitiveOutputInternal(
input.getPipeline(), WindowingStrategy.globalDefault(), input.isBounded());
}
@SuppressWarnings({"rawtypes", "unchecked"})
public static GroupAlsoByWindowsParDoFn create(
PipelineOptions options,
CloudObject cloudUserFn,
String stepName,
@Nullable List<SideInputInfo> sideInputInfos,
@Nullable List<MultiOutputInfo> multiOutputInfos,
Integer numOutputs,
ExecutionContext executionContext,
CounterSet.AddCounterMutator addCounterMutator,
StateSampler sampler /* unused */)
throws Exception {
Object windowingStrategyObj;
byte[] encodedWindowingStrategy = getBytes(cloudUserFn, PropertyNames.SERIALIZED_FN);
if (encodedWindowingStrategy.length == 0) {
windowingStrategyObj = WindowingStrategy.globalDefault();
} else {
windowingStrategyObj =
SerializableUtils.deserializeFromByteArray(
encodedWindowingStrategy, "serialized windowing strategy");
if (!(windowingStrategyObj instanceof WindowingStrategy)) {
throw new Exception(
"unexpected kind of WindowingStrategy: " + windowingStrategyObj.getClass().getName());
}
}
WindowingStrategy windowingStrategy = (WindowingStrategy) windowingStrategyObj;
byte[] serializedCombineFn = getBytes(cloudUserFn, PropertyNames.COMBINE_FN, null);
KeyedCombineFn combineFn;
if (serializedCombineFn != null) {
Object combineFnObj =
SerializableUtils.deserializeFromByteArray(serializedCombineFn, "serialized combine fn");
if (!(combineFnObj instanceof KeyedCombineFn)) {
throw new Exception(
"unexpected kind of KeyedCombineFn: " + combineFnObj.getClass().getName());
}
combineFn = (KeyedCombineFn) combineFnObj;
} else {
combineFn = null;
}
Map<String, Object> inputCoderObject = getObject(cloudUserFn, PropertyNames.INPUT_CODER);
Coder inputCoder = Serializer.deserialize(inputCoderObject, Coder.class);
if (!(inputCoder instanceof WindowedValueCoder)) {
throw new Exception(
"Expected WindowedValueCoder for inputCoder, got: " + inputCoder.getClass().getName());
}
Coder elemCoder = ((WindowedValueCoder) inputCoder).getValueCoder();
if (!(elemCoder instanceof KvCoder)) {
throw new Exception(
"Expected KvCoder for inputCoder, got: " + elemCoder.getClass().getName());
}
KvCoder kvCoder = (KvCoder) elemCoder;
boolean isStreamingPipeline = false;
if (options instanceof StreamingOptions) {
isStreamingPipeline = ((StreamingOptions) options).isStreaming();
}
KeyedCombineFn maybeMergingCombineFn = null;
if (combineFn != null) {
String phase = getString(cloudUserFn, PropertyNames.PHASE, CombinePhase.ALL);
Preconditions.checkArgument(
phase.equals(CombinePhase.ALL) || phase.equals(CombinePhase.MERGE),
"Unexpected phase: " + phase);
if (phase.equals(CombinePhase.MERGE)) {
maybeMergingCombineFn = new MergingKeyedCombineFn(combineFn);
} else {
maybeMergingCombineFn = combineFn;
}
}
DoFnInfoFactory fnFactory;
final DoFn groupAlsoByWindowsDoFn =
getGroupAlsoByWindowsDoFn(
isStreamingPipeline, windowingStrategy, kvCoder, maybeMergingCombineFn);
fnFactory =
new DoFnInfoFactory() {
@Override
public DoFnInfo createDoFnInfo() {
return new DoFnInfo(groupAlsoByWindowsDoFn, null);
}
};
return new GroupAlsoByWindowsParDoFn(
options, fnFactory, stepName, executionContext, addCounterMutator);
}
@Override
public PCollection<KV<K, V>> apply(PInput input) {
return PCollection.createPrimitiveOutputInternal(
input.getPipeline(), WindowingStrategy.globalDefault(), PCollection.IsBounded.BOUNDED);
}
