Examples with DirectStepContext org.apache.beam.runners.direct.DirectExecutionContext.DirectStepContext used on opensource projects

Example 1 with DirectStepContext

use of org.apache.beam.runners.direct.DirectExecutionContext.DirectStepContext in project beam by apache.

the class ParDoEvaluatorTest method createEvaluator.

private ParDoEvaluator<Integer> createEvaluator(PCollectionView<Integer> singletonView, RecorderFn fn, PCollection<Integer> input, PCollection<Integer> output) {
    when(evaluationContext.createSideInputReader(ImmutableList.of(singletonView))).thenReturn(new ReadyInGlobalWindowReader());
    DirectExecutionContext executionContext = mock(DirectExecutionContext.class);
    DirectStepContext stepContext = mock(DirectStepContext.class);
    when(executionContext.getStepContext(Mockito.any(String.class))).thenReturn(stepContext);
    when(stepContext.getTimerUpdate()).thenReturn(TimerUpdate.empty());
    when(evaluationContext.getExecutionContext(Mockito.any(AppliedPTransform.class), Mockito.any(StructuralKey.class))).thenReturn(executionContext);
    DirectGraphs.performDirectOverrides(p);
    @SuppressWarnings("unchecked") AppliedPTransform<PCollection<Integer>, ?, ?> transform = (AppliedPTransform<PCollection<Integer>, ?, ?>) DirectGraphs.getProducer(output);
    return ParDoEvaluator.create(evaluationContext, PipelineOptionsFactory.create(), stepContext, transform, input.getCoder(), input.getWindowingStrategy(), fn, null, /* key */
    ImmutableList.of(singletonView), mainOutputTag, additionalOutputTags, ImmutableMap.of(mainOutputTag, output), DoFnSchemaInformation.create(), Collections.emptyMap(), ParDoEvaluator.defaultRunnerFactory());
}

Example 2 with DirectStepContext

use of org.apache.beam.runners.direct.DirectExecutionContext.DirectStepContext in project beam by apache.

the class ParDoEvaluatorFactory method createEvaluator.

/**
   * Creates an evaluator for an arbitrary {@link AppliedPTransform} node, with the pieces of the
   * {@link ParDo} unpacked.
   *
   * <p>This can thus be invoked regardless of whether the types in the {@link AppliedPTransform}
   * correspond with the type in the unpacked {@link DoFn}, side inputs, and output tags.
   */
@SuppressWarnings({ "unchecked", "rawtypes" })
DoFnLifecycleManagerRemovingTransformEvaluator<InputT> createEvaluator(AppliedPTransform<PCollection<InputT>, PCollectionTuple, ?> application, StructuralKey<?> inputBundleKey, DoFn<InputT, OutputT> doFn, List<PCollectionView<?>> sideInputs, TupleTag<OutputT> mainOutputTag, List<TupleTag<?>> additionalOutputTags) throws Exception {
    String stepName = evaluationContext.getStepName(application);
    DirectStepContext stepContext = evaluationContext.getExecutionContext(application, inputBundleKey).getStepContext(stepName);
    DoFnLifecycleManager fnManager = fnClones.getUnchecked(doFn);
    return DoFnLifecycleManagerRemovingTransformEvaluator.wrapping(createParDoEvaluator(application, inputBundleKey, sideInputs, mainOutputTag, additionalOutputTags, stepContext, fnManager.<InputT, OutputT>get(), fnManager), fnManager);
}

Example 3 with DirectStepContext

use of org.apache.beam.runners.direct.DirectExecutionContext.DirectStepContext in project beam by apache.

the class EvaluationContextTest method getExecutionContextSameStepSameKeyState.

@Test
public void getExecutionContextSameStepSameKeyState() {
    DirectExecutionContext fooContext = context.getExecutionContext(impulseProducer, StructuralKey.of("foo", StringUtf8Coder.of()));
    StateTag<BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of());
    DirectStepContext stepContext = fooContext.getStepContext("s1");
    stepContext.stateInternals().state(StateNamespaces.global(), intBag).add(1);
    context.handleResult(ImmutableListBundleFactory.create().createKeyedBundle(StructuralKey.of("foo", StringUtf8Coder.of()), impulse).commit(Instant.now()), ImmutableList.of(), StepTransformResult.withoutHold(impulseProducer).withState(stepContext.commitState()).build());
    DirectExecutionContext secondFooContext = context.getExecutionContext(impulseProducer, StructuralKey.of("foo", StringUtf8Coder.of()));
    assertThat(secondFooContext.getStepContext("s1").stateInternals().state(StateNamespaces.global(), intBag).read(), contains(1));
}

Example 4 with DirectStepContext

use of org.apache.beam.runners.direct.DirectExecutionContext.DirectStepContext in project beam by apache.

the class ParDoEvaluatorFactory method createEvaluator.

/**
 * Creates an evaluator for an arbitrary {@link AppliedPTransform} node, with the pieces of the
 * {@link ParDo} unpacked.
 *
 * <p>This can thus be invoked regardless of whether the types in the {@link AppliedPTransform}
 * correspond with the type in the unpacked {@link DoFn}, side inputs, and output tags.
 */
@SuppressWarnings({ "unchecked", "rawtypes" })
DoFnLifecycleManagerRemovingTransformEvaluator<InputT> createEvaluator(AppliedPTransform<PCollection<InputT>, PCollectionTuple, ?> application, PCollection<InputT> mainInput, StructuralKey<?> inputBundleKey, List<PCollectionView<?>> sideInputs, TupleTag<OutputT> mainOutputTag, List<TupleTag<?>> additionalOutputTags, DoFnSchemaInformation doFnSchemaInformation, Map<String, PCollectionView<?>> sideInputMapping) throws Exception {
    String stepName = evaluationContext.getStepName(application);
    DirectStepContext stepContext = evaluationContext.getExecutionContext(application, inputBundleKey).getStepContext(stepName);
    DoFnLifecycleManager fnManager = fnClones.getUnchecked(application);
    return DoFnLifecycleManagerRemovingTransformEvaluator.wrapping(createParDoEvaluator(application, inputBundleKey, mainInput, sideInputs, mainOutputTag, additionalOutputTags, stepContext, fnManager.get(), doFnSchemaInformation, sideInputMapping, fnManager), fnManager);
}

Example 5 with DirectStepContext

use of org.apache.beam.runners.direct.DirectExecutionContext.DirectStepContext in project beam by apache.

the class StatefulParDoEvaluatorFactory method createEvaluator.

@SuppressWarnings({ "unchecked", "rawtypes" })
private TransformEvaluator<KeyedWorkItem<K, KV<K, InputT>>> createEvaluator(AppliedPTransform<PCollection<? extends KeyedWorkItem<K, KV<K, InputT>>>, PCollectionTuple, StatefulParDo<K, InputT, OutputT>> application, CommittedBundle<KeyedWorkItem<K, KV<K, InputT>>> inputBundle) throws Exception {
    DoFnLifecycleManagerRemovingTransformEvaluator<KV<K, InputT>> delegateEvaluator = delegateFactory.createEvaluator((AppliedPTransform) application, (PCollection) inputBundle.getPCollection(), inputBundle.getKey(), application.getTransform().getSideInputs(), application.getTransform().getMainOutputTag(), application.getTransform().getAdditionalOutputTags().getAll(), application.getTransform().getSchemaInformation(), application.getTransform().getSideInputMapping());
    DirectStepContext stepContext = evaluationContext.getExecutionContext(application, inputBundle.getKey()).getStepContext(evaluationContext.getStepName(application));
    stepContext.stateInternals().commit();
    return new StatefulParDoEvaluator<>(delegateEvaluator, stepContext);
}