@Override
public CountSum<NumT> decode(InputStream inStream, Coder.Context context)
throws CoderException, IOException {
Coder.Context nestedContext = context.nested();
return new CountSum<>(
LONG_CODER.decode(inStream, nestedContext), DOUBLE_CODER.decode(inStream, nestedContext));
}
@Override
public void encode(CountSum<NumT> value, OutputStream outStream, Coder.Context context)
throws CoderException, IOException {
Coder.Context nestedContext = context.nested();
LONG_CODER.encode(value.count, outStream, nestedContext);
DOUBLE_CODER.encode(value.sum, outStream, nestedContext);
}
<T> void runTestRead(T[] expected, Coder<T> coder) throws Exception {
File tmpFile = tmpFolder.newFile("file.txt");
String filename = tmpFile.getPath();
try (PrintStream writer = new PrintStream(new FileOutputStream(tmpFile))) {
for (T elem : expected) {
byte[] encodedElem = CoderUtils.encodeToByteArray(coder, elem);
String line = new String(encodedElem);
writer.println(line);
}
}
Pipeline p = TestPipeline.create();
TextIO.Read.Bound<T> read;
if (coder.equals(StringUtf8Coder.of())) {
TextIO.Read.Bound<String> readStrings = TextIO.Read.from(filename);
// T==String
read = (TextIO.Read.Bound<T>) readStrings;
} else {
read = TextIO.Read.from(filename).withCoder(coder);
}
PCollection<T> output = p.apply(read);
DataflowAssert.that(output).containsInAnyOrder(expected);
p.run();
}
<T> void runTestWrite(T[] elems, Coder<T> coder) throws Exception {
File tmpFile = tmpFolder.newFile("file.txt");
String filename = tmpFile.getPath();
Pipeline p = TestPipeline.create();
PCollection<T> input = p.apply(Create.of(Arrays.asList(elems)).withCoder(coder));
TextIO.Write.Bound<T> write;
if (coder.equals(StringUtf8Coder.of())) {
TextIO.Write.Bound<String> writeStrings = TextIO.Write.to(filename).withoutSharding();
// T==String
write = (TextIO.Write.Bound<T>) writeStrings;
} else {
write = TextIO.Write.to(filename).withCoder(coder).withoutSharding();
}
input.apply(write);
p.run();
List<String> actual = new ArrayList<>();
try (BufferedReader reader = new BufferedReader(new FileReader(tmpFile))) {
for (; ; ) {
String line = reader.readLine();
if (line == null) {
break;
}
actual.add(line);
}
}
String[] expected = new String[elems.length];
for (int i = 0; i < elems.length; i++) {
T elem = elems[i];
byte[] encodedElem = CoderUtils.encodeToByteArray(coder, elem);
String line = new String(encodedElem);
expected[i] = line;
}
assertThat(actual, containsInAnyOrder(expected));
}
@Override
public void flatMap(IN value, Collector<WindowedValue<OUT>> out) throws Exception {
@SuppressWarnings("unchecked")
OUT voidValue = (OUT) VoidCoderTypeSerializer.VoidValue.INSTANCE;
for (byte[] element : elements) {
ByteArrayInputStream bai = new ByteArrayInputStream(element);
OUT outValue = coder.decode(bai, Coder.Context.OUTER);
if (outValue == null) {
out.collect(
WindowedValue.of(voidValue, Instant.now(), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING));
} else {
out.collect(
WindowedValue.of(outValue, Instant.now(), GlobalWindow.INSTANCE, PaneInfo.NO_FIRING));
}
}
out.close();
}
@Override
public TimestampedValue<T> decode(InputStream inStream, Context context) throws IOException {
T value = valueCoder.decode(inStream, context.nested());
Instant timestamp = InstantCoder.of().decode(inStream, context);
return TimestampedValue.of(value, timestamp);
}
@Override
public void encode(TimestampedValue<T> windowedElem, OutputStream outStream, Context context)
throws IOException {
valueCoder.encode(windowedElem.getValue(), outStream, context.nested());
InstantCoder.of().encode(windowedElem.getTimestamp(), outStream, context);
}
private void checkRegistry(Class<? extends Mutation> mutationClass)
throws CannotProvideCoderException {
Coder<? extends Mutation> coder = registry.getCoder(TypeDescriptor.of(mutationClass));
assertNotNull(coder);
assertEquals(HBaseMutationCoder.class, coder.getClass());
}
@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);
}
private <X> X decode(Coder<X> coder, InputStream input) throws IOException {
return coder.decode(input, Coder.Context.OUTER);
}