Examples with DataflowStepContext org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext used on opensource projects

Example 16 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class UserDistributionMonitoringInfoToCounterUpdateTransformer method validate.

private Optional<String> validate(MonitoringInfo monitoringInfo) {
    Optional<String> validatorResult = specValidator.validate(monitoringInfo);
    if (validatorResult.isPresent()) {
        return validatorResult;
    }
    String urn = monitoringInfo.getUrn();
    if (!urn.equals(Urns.USER_DISTRIBUTION_INT64)) {
        throw new RuntimeException(String.format("Received unexpected counter urn. Expected urn: %s, received: %s", Urns.USER_DISTRIBUTION_INT64, urn));
    }
    String type = monitoringInfo.getType();
    if (!type.equals(TypeUrns.DISTRIBUTION_INT64_TYPE)) {
        throw new RuntimeException(String.format("Received unexpected counter type. Expected type: %s, received: %s", TypeUrns.DISTRIBUTION_INT64_TYPE, type));
    }
    final String ptransform = monitoringInfo.getLabelsMap().get(MonitoringInfoConstants.Labels.PTRANSFORM);
    DataflowStepContext stepContext = transformIdMapping.get(ptransform);
    if (stepContext == null) {
        return Optional.of("Encountered user-counter MonitoringInfo with unknown ptransformId: " + monitoringInfo.toString());
    }
    return Optional.empty();
}

Example 17 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class UserDistributionMonitoringInfoToCounterUpdateTransformer method transform.

/**
 * Transforms user counter MonitoringInfo to relevant CounterUpdate.
 *
 * @return Relevant CounterUpdate or null if transformation failed.
 */
@Override
@Nullable
public CounterUpdate transform(MonitoringInfo monitoringInfo) {
    Optional<String> validationResult = validate(monitoringInfo);
    if (validationResult.isPresent()) {
        LOG.debug(validationResult.get());
        return null;
    }
    DistributionData data = decodeInt64Distribution(monitoringInfo.getPayload());
    Map<String, String> miLabels = monitoringInfo.getLabelsMap();
    final String ptransform = miLabels.get(MonitoringInfoConstants.Labels.PTRANSFORM);
    final String counterName = miLabels.get(MonitoringInfoConstants.Labels.NAME);
    final String counterNamespace = miLabels.get(MonitoringInfoConstants.Labels.NAMESPACE);
    CounterStructuredNameAndMetadata name = new CounterStructuredNameAndMetadata();
    DataflowStepContext stepContext = transformIdMapping.get(ptransform);
    name.setName(new CounterStructuredName().setOrigin(Origin.USER.toString()).setName(counterName).setOriginalStepName(stepContext.getNameContext().originalName()).setOriginNamespace(counterNamespace)).setMetadata(new CounterMetadata().setKind(Kind.DISTRIBUTION.toString()));
    return new CounterUpdate().setStructuredNameAndMetadata(name).setCumulative(true).setDistribution(new DistributionUpdate().setMax(DataflowCounterUpdateExtractor.longToSplitInt(data.max())).setMin(DataflowCounterUpdateExtractor.longToSplitInt(data.min())).setSum(DataflowCounterUpdateExtractor.longToSplitInt(data.sum())).setCount(DataflowCounterUpdateExtractor.longToSplitInt(data.count())));
}

Example 18 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class RegisterAndProcessBundleOperation method handleBagUserState.

private CompletionStage<BeamFnApi.StateResponse.Builder> handleBagUserState(StateRequest stateRequest) {
    StateKey.BagUserState bagUserStateKey = stateRequest.getStateKey().getBagUserState();
    DataflowStepContext userStepContext = ptransformIdToUserStepContext.get(bagUserStateKey.getTransformId());
    checkState(userStepContext != null, String.format("Unknown PTransform id '%s'", bagUserStateKey.getTransformId()));
    // TODO: We should not be required to hold onto a pointer to the bag states for the
    // user. InMemoryStateInternals assumes that the Java garbage collector does the clean-up work
    // but instead StateInternals should hold its own references and write out any data and
    // clear references when the MapTask within Dataflow completes like how WindmillStateInternals
    // works.
    BagState<ByteString> state = userStateData.computeIfAbsent(stateRequest.getStateKey(), unused -> userStepContext.stateInternals().state(// window.
    StateNamespaces.window(GlobalWindow.Coder.INSTANCE, GlobalWindow.INSTANCE), StateTags.bag(bagUserStateKey.getUserStateId(), ByteStringCoder.of())));
    switch(stateRequest.getRequestCase()) {
        case GET:
            return CompletableFuture.completedFuture(StateResponse.newBuilder().setGet(StateGetResponse.newBuilder().setData(concat(state.read()))));
        case APPEND:
            state.add(stateRequest.getAppend().getData());
            return CompletableFuture.completedFuture(StateResponse.newBuilder().setAppend(StateAppendResponse.getDefaultInstance()));
        case CLEAR:
            state.clear();
            return CompletableFuture.completedFuture(StateResponse.newBuilder().setClear(StateClearResponse.getDefaultInstance()));
        default:
            throw new IllegalArgumentException(String.format("Unknown request type %s", stateRequest.getRequestCase()));
    }
}

Example 19 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class BeamFnMapTaskExecutorFactory method createOperationTransformForRegisterFnNodes.

private Function<Node, Node> createOperationTransformForRegisterFnNodes(final IdGenerator idGenerator, final InstructionRequestHandler instructionRequestHandler, final StateDelegator beamFnStateDelegator, final String stageName, final DataflowExecutionContext<?> executionContext) {
    return new TypeSafeNodeFunction<RegisterRequestNode>(RegisterRequestNode.class) {

        @Override
        public Node typedApply(RegisterRequestNode input) {
            ImmutableMap.Builder<String, DataflowOperationContext> ptransformIdToOperationContextBuilder = ImmutableMap.builder();
            ImmutableMap.Builder<String, DataflowStepContext> ptransformIdToStepContext = ImmutableMap.builder();
            for (Map.Entry<String, NameContext> entry : input.getPTransformIdToPartialNameContextMap().entrySet()) {
                NameContext fullNameContext = NameContext.create(stageName, entry.getValue().originalName(), entry.getValue().systemName(), entry.getValue().userName());
                DataflowOperationContext operationContext = executionContext.createOperationContext(fullNameContext);
                ptransformIdToOperationContextBuilder.put(entry.getKey(), operationContext);
                ptransformIdToStepContext.put(entry.getKey(), executionContext.getStepContext(operationContext));
            }
            ImmutableMap.Builder<String, NameContext> pcollectionIdToNameContext = ImmutableMap.builder();
            for (Map.Entry<String, NameContext> entry : input.getPCollectionToPartialNameContextMap().entrySet()) {
                pcollectionIdToNameContext.put(entry.getKey(), NameContext.create(stageName, entry.getValue().originalName(), entry.getValue().systemName(), entry.getValue().userName()));
            }
            ImmutableMap<String, DataflowOperationContext> ptransformIdToOperationContexts = ptransformIdToOperationContextBuilder.build();
            ImmutableMap<String, SideInputReader> ptransformIdToSideInputReaders = buildPTransformIdToSideInputReadersMap(executionContext, input, ptransformIdToOperationContexts);
            ImmutableTable<String, String, PCollectionView<?>> ptransformIdToSideInputIdToPCollectionView = buildPTransformIdToSideInputIdToPCollectionView(input);
            return OperationNode.create(new RegisterAndProcessBundleOperation(idGenerator, instructionRequestHandler, beamFnStateDelegator, input.getRegisterRequest(), ptransformIdToOperationContexts, ptransformIdToStepContext.build(), ptransformIdToSideInputReaders, ptransformIdToSideInputIdToPCollectionView, pcollectionIdToNameContext.build(), // TODO: Set NameContext properly for these operations.
            executionContext.createOperationContext(NameContext.create(stageName, stageName, stageName, stageName))));
        }
    };
}

Example 20 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class StreamingPCollectionViewWriterDoFnFactoryTest method testConstruction.

@Test
public void testConstruction() throws Exception {
    DataflowOperationContext mockOperationContext = Mockito.mock(DataflowOperationContext.class);
    DataflowExecutionContext mockExecutionContext = Mockito.mock(DataflowExecutionContext.class);
    DataflowStepContext mockStepContext = Mockito.mock(StreamingModeExecutionContext.StepContext.class);
    when(mockExecutionContext.getStepContext(mockOperationContext)).thenReturn(mockStepContext);
    CloudObject coder = CloudObjects.asCloudObject(WindowedValue.getFullCoder(BigEndianIntegerCoder.of(), GlobalWindow.Coder.INSTANCE), /*sdkComponents=*/
    null);
    ParDoFn parDoFn = new StreamingPCollectionViewWriterDoFnFactory().create(null, /* pipeline options */
    CloudObject.fromSpec(ImmutableMap.of(PropertyNames.OBJECT_TYPE_NAME, "StreamingPCollectionViewWriterDoFn", PropertyNames.ENCODING, coder, WorkerPropertyNames.SIDE_INPUT_ID, "test-side-input-id")), null, /* side input infos */
    null, /* main output tag */
    null, /* output tag to receiver index */
    mockExecutionContext, mockOperationContext);
    assertThat(parDoFn, instanceOf(StreamingPCollectionViewWriterParDoFn.class));
}