Examples with NameContext org.apache.beam.runners.dataflow.worker.counters.NameContext used on opensource projects

Example 11 with NameContext

use of org.apache.beam.runners.dataflow.worker.counters.NameContext in project beam by apache.

the class IntrinsicMapTaskExecutorFactory method createOperationTransformForParallelInstructionNodes.

/**
 * Creates an {@link Operation} from the given {@link ParallelInstruction} definition using the
 * provided {@link ReaderFactory}.
 */
Function<Node, Node> createOperationTransformForParallelInstructionNodes(final String stageName, final Network<Node, Edge> network, final PipelineOptions options, final ReaderFactory readerFactory, final SinkFactory sinkFactory, final DataflowExecutionContext<?> executionContext) {
    return new TypeSafeNodeFunction<ParallelInstructionNode>(ParallelInstructionNode.class) {

        @Override
        public Node typedApply(ParallelInstructionNode node) {
            ParallelInstruction instruction = node.getParallelInstruction();
            NameContext nameContext = NameContext.create(stageName, instruction.getOriginalName(), instruction.getSystemName(), instruction.getName());
            try {
                DataflowOperationContext context = executionContext.createOperationContext(nameContext);
                if (instruction.getRead() != null) {
                    return createReadOperation(network, node, options, readerFactory, executionContext, context);
                } else if (instruction.getWrite() != null) {
                    return createWriteOperation(node, options, sinkFactory, executionContext, context);
                } else if (instruction.getParDo() != null) {
                    return createParDoOperation(network, node, options, executionContext, context);
                } else if (instruction.getPartialGroupByKey() != null) {
                    return createPartialGroupByKeyOperation(network, node, options, executionContext, context);
                } else if (instruction.getFlatten() != null) {
                    return createFlattenOperation(network, node, context);
                } else {
                    throw new IllegalArgumentException(String.format("Unexpected instruction: %s", instruction));
                }
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
    };
}

Example 12 with NameContext

use of org.apache.beam.runners.dataflow.worker.counters.NameContext in project beam by apache.

the class BeamFnMapTaskExecutorTest method generateDataflowStepContext.

/**
 * Generates bare minumum DataflowStepContext to use for testing.
 *
 * @param valuesPrefix prefix for all types of names that are specified in DataflowStepContext.
 * @return new instance of DataflowStepContext
 */
private DataflowStepContext generateDataflowStepContext(String valuesPrefix) {
    NameContext nc = new NameContext() {

        @Override
        @Nullable
        public String stageName() {
            return valuesPrefix + "Stage";
        }

        @Override
        @Nullable
        public String originalName() {
            return valuesPrefix + "OriginalName";
        }

        @Override
        @Nullable
        public String systemName() {
            return valuesPrefix + "SystemName";
        }

        @Override
        @Nullable
        public String userName() {
            return valuesPrefix + "UserName";
        }
    };
    DataflowStepContext dsc = new DataflowStepContext(nc) {

        @Override
        @Nullable
        public <W extends BoundedWindow> TimerData getNextFiredTimer(Coder<W> windowCoder) {
            return null;
        }

        @Override
        public <W extends BoundedWindow> void setStateCleanupTimer(String timerId, W window, Coder<W> windowCoder, Instant cleanupTime, Instant cleanupOutputTimestamp) {
        }

        @Override
        public DataflowStepContext namespacedToUser() {
            return this;
        }

        @Override
        public StateInternals stateInternals() {
            return null;
        }

        @Override
        public TimerInternals timerInternals() {
            return null;
        }
    };
    return dsc;
}

Example 13 with NameContext

use of org.apache.beam.runners.dataflow.worker.counters.NameContext in project beam by apache.

the class UserMonitoringInfoToCounterUpdateTransformerTest method testTransformReturnsValidCounterUpdateWhenValidUserMonitoringInfoReceived.

@Test
public void testTransformReturnsValidCounterUpdateWhenValidUserMonitoringInfoReceived() throws Exception {
    Map<String, DataflowStepContext> stepContextMapping = new HashMap<>();
    NameContext nc = NameContext.create("anyStageName", "anyOriginalName", "anySystemName", "anyUserName");
    DataflowStepContext dsc = mock(DataflowStepContext.class);
    when(dsc.getNameContext()).thenReturn(nc);
    stepContextMapping.put("anyValue", dsc);
    MonitoringInfo monitoringInfo = MonitoringInfo.newBuilder().setUrn(Urns.USER_SUM_INT64).setType(TypeUrns.SUM_INT64_TYPE).putLabels(MonitoringInfoConstants.Labels.NAME, "anyName").putLabels(MonitoringInfoConstants.Labels.NAMESPACE, "anyNamespace").putLabels(MonitoringInfoConstants.Labels.PTRANSFORM, "anyValue").setPayload(encodeInt64Counter(1L)).build();
    UserMonitoringInfoToCounterUpdateTransformer testObject = new UserMonitoringInfoToCounterUpdateTransformer(mockSpecValidator, stepContextMapping);
    when(mockSpecValidator.validate(any())).thenReturn(Optional.empty());
    CounterUpdate result = testObject.transform(monitoringInfo);
    assertNotNull(result);
    assertEqualsAsJson("{cumulative:true, integer:{highBits:0, lowBits:1}, " + "structuredNameAndMetadata:{metadata:{kind:'SUM'}, " + "name:{name:'anyName', origin:'USER', originNamespace:'anyNamespace', " + "originalStepName:'anyOriginalName'}}}", result);
}

Example 14 with NameContext

use of org.apache.beam.runners.dataflow.worker.counters.NameContext in project beam by apache.

the class UserDistributionMonitoringInfoToCounterUpdateTransformerTest method testTransformReturnsValidCounterUpdateWhenValidUserMonitoringInfoReceived.

@Test
public void testTransformReturnsValidCounterUpdateWhenValidUserMonitoringInfoReceived() throws Exception {
    Map<String, DataflowStepContext> stepContextMapping = new HashMap<>();
    NameContext nc = NameContext.create("anyStageName", "anyOriginalName", "anySystemName", "anyUserName");
    DataflowStepContext dsc = mock(DataflowStepContext.class);
    when(dsc.getNameContext()).thenReturn(nc);
    stepContextMapping.put("anyValue", dsc);
    MonitoringInfo monitoringInfo = MonitoringInfo.newBuilder().setUrn(Urns.USER_DISTRIBUTION_INT64).setType(TypeUrns.DISTRIBUTION_INT64_TYPE).putLabels(MonitoringInfoConstants.Labels.NAME, "anyName").putLabels(MonitoringInfoConstants.Labels.NAMESPACE, "anyNamespace").putLabels(MonitoringInfoConstants.Labels.PTRANSFORM, "anyValue").setPayload(encodeInt64Distribution(DistributionData.create(2L, /* sum */
    1L, /* count */
    3L, /* min */
    4L))).build();
    UserDistributionMonitoringInfoToCounterUpdateTransformer testObject = new UserDistributionMonitoringInfoToCounterUpdateTransformer(mockSpecValidator, stepContextMapping);
    when(mockSpecValidator.validate(any())).thenReturn(Optional.empty());
    CounterUpdate result = testObject.transform(monitoringInfo);
    assertNotNull(result);
    assertEqualsAsJson("{cumulative:true, distribution:{count:{highBits:0, lowBits:1}, " + "max:{highBits:0, lowBits:4}, min:{highBits:0, lowBits:3}, " + "sum:{highBits:0, lowBits:2}}, " + "structuredNameAndMetadata:{metadata:{kind:'DISTRIBUTION'}, " + "name:{name:'anyName', origin:'USER', originNamespace:'anyNamespace', " + "originalStepName:'anyOriginalName'}}}", result);
}

Example 15 with NameContext

use of org.apache.beam.runners.dataflow.worker.counters.NameContext in project beam by apache.

the class NodesTest method testFetchReadySideInputsAndFilterBlockedStreamingSideInputsNode.

@Test
public void testFetchReadySideInputsAndFilterBlockedStreamingSideInputsNode() {
    WindowingStrategy windowingStrategy = WindowingStrategy.globalDefault();
    Map<PCollectionView<?>, RunnerApi.FunctionSpec> pcollectionViewsToWindowMappingFns = ImmutableMap.of(mock(PCollectionView.class), FunctionSpec.newBuilder().setUrn("beam:test:urn:1.0").build());
    NameContext nameContext = NameContextsForTests.nameContextForTest();
    assertSame(FetchAndFilterStreamingSideInputsNode.create(windowingStrategy, pcollectionViewsToWindowMappingFns, nameContext).getWindowingStrategy(), windowingStrategy);
    assertSame(FetchAndFilterStreamingSideInputsNode.create(windowingStrategy, pcollectionViewsToWindowMappingFns, nameContext).getPCollectionViewsToWindowMappingFns(), pcollectionViewsToWindowMappingFns);
    assertSame(FetchAndFilterStreamingSideInputsNode.create(windowingStrategy, pcollectionViewsToWindowMappingFns, nameContext).getNameContext(), nameContext);
}