Examples with ElementCounter org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter used on opensource projects

Example 1 with ElementCounter

use of org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter in project beam by apache.

the class IntrinsicMapTaskExecutorFactory 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());
        }
    };
}

Example 2 with ElementCounter

use of org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter 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();
}

Example 3 with ElementCounter

use of org.apache.beam.runners.dataflow.worker.util.common.worker.ElementCounter 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());
        }
    };
}