@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); }