Java Code Examples for org.apache.beam.sdk.transforms.reflect.DoFnInvoker

The following examples show how to use org.apache.beam.sdk.transforms.reflect.DoFnInvoker. These examples are extracted from open source projects. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source Project: beam   Source File: SplittableParDoViaKeyedWorkItems.java    License: Apache License 2.0 6 votes vote down vote up
private DoFnInvoker.ArgumentProvider<InputT, OutputT> wrapContextAsStartBundle(
    final StartBundleContext baseContext) {
  return new BaseArgumentProvider<InputT, OutputT>() {
    @Override
    public DoFn<InputT, OutputT>.StartBundleContext startBundleContext(
        DoFn<InputT, OutputT> doFn) {
      return fn.new StartBundleContext() {
        @Override
        public PipelineOptions getPipelineOptions() {
          return baseContext.getPipelineOptions();
        }
      };
    }

    @Override
    public PipelineOptions pipelineOptions() {
      return baseContext.getPipelineOptions();
    }

    @Override
    public String getErrorContext() {
      return "SplittableParDoViaKeyedWorkItems/StartBundle";
    }
  };
}
 
Example 2
Source Project: beam   Source File: FlinkDoFnFunction.java    License: Apache License 2.0 6 votes vote down vote up
@Override
public void close() throws Exception {
  Exception suppressed = null;
  try {
    if (bundleStarted && !exceptionThrownInFlatMap) {
      doFnRunner.finishBundle();
    }
  } catch (Exception e) {
    // Suppress exception, so we can properly teardown DoFn.
    suppressed = e;
  }
  try {
    metricContainer.registerMetricsForPipelineResult();
    Optional.ofNullable(doFnInvoker).ifPresent(DoFnInvoker::invokeTeardown);
    if (suppressed != null) {
      throw suppressed;
    }
  } finally {
    Workarounds.deleteStaticCaches();
  }
}
 
Example 3
Source Project: nemo   Source File: DoTransform.java    License: Apache License 2.0 5 votes vote down vote up
@Override
public void onData(final Iterator<I> elements, final String srcVertexId) {
  final StartBundleContext startBundleContext = new StartBundleContext(doFn, serializedOptions);
  final FinishBundleContext finishBundleContext = new FinishBundleContext(doFn, outputCollector, serializedOptions);
  final ProcessContext processContext = new ProcessContext(doFn, outputCollector, sideInputs, serializedOptions);
  final DoFnInvoker invoker = DoFnInvokers.invokerFor(doFn);
  invoker.invokeSetup();
  invoker.invokeStartBundle(startBundleContext);
  elements.forEachRemaining(element -> { // No need to check for input index, since it is always 0 for DoTransform
    processContext.setElement(element);
    invoker.invokeProcessElement(processContext);
  });
  invoker.invokeFinishBundle(finishBundleContext);
  invoker.invokeTeardown();
}
 
Example 4
Source Project: beam   Source File: DoFnOperator.java    License: Apache License 2.0 5 votes vote down vote up
@Override
public void dispose() throws Exception {
  try {
    Optional.ofNullable(flinkMetricContainer)
        .ifPresent(FlinkMetricContainer::registerMetricsForPipelineResult);
    Optional.ofNullable(checkFinishBundleTimer).ifPresent(timer -> timer.cancel(true));
    Workarounds.deleteStaticCaches();
    Optional.ofNullable(doFnInvoker).ifPresent(DoFnInvoker::invokeTeardown);
  } finally {
    // This releases all task's resources. We need to call this last
    // to ensure that state, timers, or output buffers can still be
    // accessed during finishing the bundle.
    super.dispose();
  }
}
 
Example 5
Source Project: beam   Source File: FlinkStatefulDoFnFunction.java    License: Apache License 2.0 5 votes vote down vote up
@Override
public void close() throws Exception {
  try {
    metricContainer.registerMetricsForPipelineResult();
    Optional.ofNullable(doFnInvoker).ifPresent(DoFnInvoker::invokeTeardown);
  } finally {
    Workarounds.deleteStaticCaches();
  }
}
 
Example 6
Source Project: beam   Source File: DoFnFunction.java    License: Apache License 2.0 4 votes vote down vote up
@Override
public void close() {
  Optional.ofNullable(doFnInvoker).ifPresent(DoFnInvoker::invokeTeardown);
}
 
Example 7
Source Project: beam   Source File: SamzaDoFnInvokerRegistrar.java    License: Apache License 2.0 4 votes vote down vote up
/** Returns the invoker for a {@link DoFn}. */
<InputT, OutputT> DoFnInvoker<InputT, OutputT> invokerFor(
    DoFn<InputT, OutputT> fn, Context context);
 
Example 8
Source Project: beam   Source File: DataflowPipelineTranslator.java    License: Apache License 2.0 4 votes vote down vote up
private <InputT, OutputT> void translateMultiHelper(
    ParDo.MultiOutput<InputT, OutputT> transform, TranslationContext context) {
  StepTranslationContext stepContext = context.addStep(transform, "ParallelDo");
  DoFnSchemaInformation doFnSchemaInformation;
  doFnSchemaInformation =
      ParDoTranslation.getSchemaInformation(context.getCurrentTransform());
  Map<String, PCollectionView<?>> sideInputMapping =
      ParDoTranslation.getSideInputMapping(context.getCurrentTransform());
  Map<TupleTag<?>, Coder<?>> outputCoders =
      context.getOutputs(transform).entrySet().stream()
          .collect(
              Collectors.toMap(
                  Map.Entry::getKey, e -> ((PCollection) e.getValue()).getCoder()));
  translateInputs(
      stepContext,
      context.getInput(transform),
      transform.getSideInputs().values(),
      context);
  translateOutputs(context.getOutputs(transform), stepContext);
  String ptransformId =
      context.getSdkComponents().getPTransformIdOrThrow(context.getCurrentTransform());
  translateFn(
      stepContext,
      ptransformId,
      transform.getFn(),
      context.getInput(transform).getWindowingStrategy(),
      transform.getSideInputs().values(),
      context.getInput(transform).getCoder(),
      context,
      transform.getMainOutputTag(),
      outputCoders,
      doFnSchemaInformation,
      sideInputMapping);

  // TODO: Move this logic into translateFn once the legacy ProcessKeyedElements is
  // removed.
  if (context.isFnApi()) {
    DoFnSignature signature = DoFnSignatures.signatureForDoFn(transform.getFn());
    if (signature.processElement().isSplittable()) {
      DoFnInvoker<?, ?> doFnInvoker = DoFnInvokers.invokerFor(transform.getFn());
      Coder<?> restrictionAndWatermarkStateCoder =
          KvCoder.of(
              doFnInvoker.invokeGetRestrictionCoder(
                  context.getInput(transform).getPipeline().getCoderRegistry()),
              doFnInvoker.invokeGetWatermarkEstimatorStateCoder(
                  context.getInput(transform).getPipeline().getCoderRegistry()));
      stepContext.addInput(
          PropertyNames.RESTRICTION_ENCODING,
          translateCoder(restrictionAndWatermarkStateCoder, context));
    }
  }
}
 
Example 9
Source Project: beam   Source File: DataflowPipelineTranslator.java    License: Apache License 2.0 4 votes vote down vote up
private <InputT, OutputT> void translateSingleHelper(
    ParDoSingle<InputT, OutputT> transform, TranslationContext context) {

  DoFnSchemaInformation doFnSchemaInformation;
  doFnSchemaInformation =
      ParDoTranslation.getSchemaInformation(context.getCurrentTransform());
  Map<String, PCollectionView<?>> sideInputMapping =
      ParDoTranslation.getSideInputMapping(context.getCurrentTransform());
  StepTranslationContext stepContext = context.addStep(transform, "ParallelDo");
  Map<TupleTag<?>, Coder<?>> outputCoders =
      context.getOutputs(transform).entrySet().stream()
          .collect(
              Collectors.toMap(
                  Map.Entry::getKey, e -> ((PCollection) e.getValue()).getCoder()));

  translateInputs(
      stepContext,
      context.getInput(transform),
      transform.getSideInputs().values(),
      context);
  stepContext.addOutput(
      transform.getMainOutputTag().getId(), context.getOutput(transform));
  String ptransformId =
      context.getSdkComponents().getPTransformIdOrThrow(context.getCurrentTransform());
  translateFn(
      stepContext,
      ptransformId,
      transform.getFn(),
      context.getInput(transform).getWindowingStrategy(),
      transform.getSideInputs().values(),
      context.getInput(transform).getCoder(),
      context,
      transform.getMainOutputTag(),
      outputCoders,
      doFnSchemaInformation,
      sideInputMapping);

  // TODO: Move this logic into translateFn once the legacy ProcessKeyedElements is
  // removed.
  if (context.isFnApi()) {
    DoFnSignature signature = DoFnSignatures.signatureForDoFn(transform.getFn());
    if (signature.processElement().isSplittable()) {
      DoFnInvoker<?, ?> doFnInvoker = DoFnInvokers.invokerFor(transform.getFn());
      Coder<?> restrictionAndWatermarkStateCoder =
          KvCoder.of(
              doFnInvoker.invokeGetRestrictionCoder(
                  context.getInput(transform).getPipeline().getCoderRegistry()),
              doFnInvoker.invokeGetWatermarkEstimatorStateCoder(
                  context.getInput(transform).getPipeline().getCoderRegistry()));
      stepContext.addInput(
          PropertyNames.RESTRICTION_ENCODING,
          translateCoder(restrictionAndWatermarkStateCoder, context));
    }
  }
}
 
Example 10
Source Project: beam   Source File: SplittableProcessElementInvoker.java    License: Apache License 2.0 2 votes vote down vote up
/**
 * Invokes the {@link DoFn.ProcessElement} method using the given {@link DoFnInvoker} for the
 * original {@link DoFn}, on the given element and with the given {@link RestrictionTracker}.
 *
 * @return Information on how to resume the call: residual restriction, a {@link
 *     DoFn.ProcessContinuation}, and a future output watermark.
 */
public abstract Result invokeProcessElement(
    DoFnInvoker<InputT, OutputT> invoker,
    WindowedValue<InputT> element,
    RestrictionTracker<RestrictionT, PositionT> tracker,
    WatermarkEstimator<WatermarkEstimatorStateT> watermarkEstimator);