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

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";
    }
  };
}
 

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

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

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

@Override
public void close() throws Exception {
  try {
    metricContainer.registerMetricsForPipelineResult();
    Optional.ofNullable(doFnInvoker).ifPresent(DoFnInvoker::invokeTeardown);
  } finally {
    Workarounds.deleteStaticCaches();
  }
}
 

@Override
public void close() {
  Optional.ofNullable(doFnInvoker).ifPresent(DoFnInvoker::invokeTeardown);
}
 

/** Returns the invoker for a {@link DoFn}. */
<InputT, OutputT> DoFnInvoker<InputT, OutputT> invokerFor(
    DoFn<InputT, OutputT> fn, Context context);
 

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

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

/**
 * 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);