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

The following examples show how to use org.apache.beam.sdk.transforms.reflect.DoFnInvoker. 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 File: SplittableParDoViaKeyedWorkItems.java    From beam with 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 File: FlinkDoFnFunction.java    From beam with 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 File: DoTransform.java    From nemo with 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 File: DoFnOperator.java    From beam with 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 File: FlinkStatefulDoFnFunction.java    From beam with 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 File: DoFnFunction.java    From beam with Apache License 2.0 4 votes vote down vote up
@Override
public void close() {
  Optional.ofNullable(doFnInvoker).ifPresent(DoFnInvoker::invokeTeardown);
}
 
Example #7
Source File: SamzaDoFnInvokerRegistrar.java    From beam with 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 File: DataflowPipelineTranslator.java    From beam with 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 File: DataflowPipelineTranslator.java    From beam with 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 File: SplittableProcessElementInvoker.java    From beam with 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);