use of org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver in project beam by apache.
the class SimpleParDoFn method reallyStartBundle.
private void reallyStartBundle() throws Exception {
checkState(fnRunner == null, "bundle already started (or not properly finished)");
OutputManager outputManager = new OutputManager() {
final Map<TupleTag<?>, OutputReceiver> undeclaredOutputs = new HashMap<>();
@Nullable
private Receiver getReceiverOrNull(TupleTag<?> tag) {
Integer receiverIndex = outputTupleTagsToReceiverIndices.get(tag);
if (receiverIndex != null) {
return receivers[receiverIndex];
} else {
return undeclaredOutputs.get(tag);
}
}
@Override
public <T> void output(TupleTag<T> tag, WindowedValue<T> output) {
outputsPerElementTracker.onOutput();
Receiver receiver = getReceiverOrNull(tag);
if (receiver == null) {
// A new undeclared output.
// TODO: plumb through the operationName, so that we can
// name implicit outputs after it.
String outputName = "implicit-" + tag.getId();
// TODO: plumb through the counter prefix, so we can
// make it available to the OutputReceiver class in case
// it wants to use it in naming output counterFactory. (It
// doesn't today.)
OutputReceiver undeclaredReceiver = new OutputReceiver();
ElementCounter outputCounter = new DataflowOutputCounter(outputName, counterFactory, stepContext.getNameContext());
undeclaredReceiver.addOutputCounter(outputCounter);
undeclaredOutputs.put(tag, undeclaredReceiver);
receiver = undeclaredReceiver;
}
try {
receiver.process(output);
} catch (RuntimeException | Error e) {
// via a chain of DoFn's.
throw e;
} catch (Exception e) {
// with other Receivers.
throw new RuntimeException(e);
}
}
};
fnInfo = (DoFnInfo) doFnInstanceManager.get();
fnSignature = DoFnSignatures.getSignature(fnInfo.getDoFn().getClass());
fnRunner = runnerFactory.createRunner(fnInfo.getDoFn(), options, mainOutputTag, sideOutputTags, fnInfo.getSideInputViews(), sideInputReader, fnInfo.getInputCoder(), fnInfo.getOutputCoders(), fnInfo.getWindowingStrategy(), stepContext, userStepContext, outputManager, doFnSchemaInformation, sideInputMapping);
fnRunner.startBundle();
}
use of org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver in project beam by apache.
the class RemoteGrpcPortReadOperationTest method setUp.
@Before
public void setUp() {
MockitoAnnotations.initMocks(this);
testReceiver = new TestOutputReceiver(CODER, NameContextsForTests.nameContextForTest());
operation = new RemoteGrpcPortReadOperation<>(beamFnDataService, TRANSFORM_ID, bundleIdSupplier, CODER, new OutputReceiver[] { testReceiver }, operationContext);
}
use of org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver in project beam by apache.
the class NodesTest method testOutputReceiverNode.
@Test
public void testOutputReceiverNode() {
OutputReceiver receiver = new OutputReceiver();
Coder<?> coder = StringUtf8Coder.of();
assertSame(receiver, OutputReceiverNode.create(receiver, coder, PCOLLECTION_ID).getOutputReceiver());
assertSame(coder, OutputReceiverNode.create(receiver, coder, PCOLLECTION_ID).getCoder());
assertNotEquals(OutputReceiverNode.create(receiver, coder, PCOLLECTION_ID), OutputReceiverNode.create(receiver, coder, PCOLLECTION_ID));
}
use of org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver in project beam by apache.
the class BeamFnMapTaskExecutorFactory method createOutputReceiversTransform.
/**
* Returns a function which can convert {@link InstructionOutput}s into {@link OutputReceiver}s.
*/
static Function<Node, Node> createOutputReceiversTransform(final String stageName, final CounterFactory counterFactory) {
return new TypeSafeNodeFunction<InstructionOutputNode>(InstructionOutputNode.class) {
@Override
public Node typedApply(InstructionOutputNode input) {
InstructionOutput cloudOutput = input.getInstructionOutput();
OutputReceiver outputReceiver = new OutputReceiver();
Coder<?> coder = CloudObjects.coderFromCloudObject(CloudObject.fromSpec(cloudOutput.getCodec()));
@SuppressWarnings("unchecked") ElementCounter outputCounter = new DataflowOutputCounter(cloudOutput.getName(), new ElementByteSizeObservableCoder<>(coder), counterFactory, NameContext.create(stageName, cloudOutput.getOriginalName(), cloudOutput.getSystemName(), cloudOutput.getName()));
outputReceiver.addOutputCounter(outputCounter);
return OutputReceiverNode.create(outputReceiver, coder, input.getPcollectionId());
}
};
}
use of org.apache.beam.runners.dataflow.worker.util.common.worker.OutputReceiver in project beam by apache.
the class BeamFnMapTaskExecutorFactory method createPartialGroupByKeyOperation.
<K> OperationNode createPartialGroupByKeyOperation(Network<Node, Edge> network, ParallelInstructionNode node, PipelineOptions options, DataflowExecutionContext<?> executionContext, DataflowOperationContext operationContext) throws Exception {
ParallelInstruction instruction = node.getParallelInstruction();
PartialGroupByKeyInstruction pgbk = instruction.getPartialGroupByKey();
OutputReceiver[] receivers = getOutputReceivers(network, node);
Coder<?> windowedCoder = CloudObjects.coderFromCloudObject(CloudObject.fromSpec(pgbk.getInputElementCodec()));
if (!(windowedCoder instanceof WindowedValueCoder)) {
throw new IllegalArgumentException(String.format("unexpected kind of input coder for PartialGroupByKeyOperation: %s", windowedCoder));
}
Coder<?> elemCoder = ((WindowedValueCoder<?>) windowedCoder).getValueCoder();
if (!(elemCoder instanceof KvCoder)) {
throw new IllegalArgumentException(String.format("unexpected kind of input element coder for PartialGroupByKeyOperation: %s", elemCoder));
}
@SuppressWarnings("unchecked") KvCoder<K, ?> keyedElementCoder = (KvCoder<K, ?>) elemCoder;
CloudObject cloudUserFn = pgbk.getValueCombiningFn() != null ? CloudObject.fromSpec(pgbk.getValueCombiningFn()) : null;
ParDoFn fn = PartialGroupByKeyParDoFns.create(options, keyedElementCoder, cloudUserFn, pgbk.getSideInputs(), Arrays.<Receiver>asList(receivers), executionContext, operationContext);
return OperationNode.create(new ParDoOperation(fn, receivers, operationContext));
}
Aggregations