Examples with PTransform org.apache.beam.model.pipeline.v1.RunnerApi.PTransform used on opensource projects

Example 21 with PTransform

use of org.apache.beam.model.pipeline.v1.RunnerApi.PTransform in project beam by apache.

the class BeamFnDataWriteRunnerTest method testReuseForMultipleBundles.

@Test
public void testReuseForMultipleBundles() throws Exception {
    AtomicReference<String> bundleId = new AtomicReference<>("0");
    String localInputId = "inputPC";
    RunnerApi.PTransform pTransform = RemoteGrpcPortWrite.writeToPort(localInputId, PORT_SPEC).toPTransform();
    List<WindowedValue<String>> output0 = new ArrayList<>();
    List<WindowedValue<String>> output1 = new ArrayList<>();
    Map<ApiServiceDescriptor, BeamFnDataOutboundAggregator> aggregators = new HashMap<>();
    BeamFnDataOutboundAggregator aggregator = createRecordingAggregator(ImmutableMap.of("0", output0, "1", output1), bundleId::get);
    aggregators.put(PORT_SPEC.getApiServiceDescriptor(), aggregator);
    PTransformRunnerFactoryTestContext context = PTransformRunnerFactoryTestContext.builder(TRANSFORM_ID, pTransform).beamFnDataClient(mockBeamFnDataClient).processBundleInstructionIdSupplier(bundleId::get).outboundAggregators(aggregators).pCollections(ImmutableMap.of(localInputId, RunnerApi.PCollection.newBuilder().setCoderId(ELEM_CODER_ID).build())).coders(COMPONENTS.getCodersMap()).windowingStrategies(COMPONENTS.getWindowingStrategiesMap()).build();
    new BeamFnDataWriteRunner.Factory<String>().createRunnerForPTransform(context);
    assertThat(context.getPCollectionConsumers().keySet(), containsInAnyOrder(localInputId));
    FnDataReceiver<Object> pCollectionConsumer = context.getPCollectionConsumer(localInputId);
    pCollectionConsumer.accept(valueInGlobalWindow("ABC"));
    pCollectionConsumer.accept(valueInGlobalWindow("DEF"));
    assertThat(output0, contains(valueInGlobalWindow("ABC"), valueInGlobalWindow("DEF")));
    output0.clear();
    // Process for bundle id 1
    bundleId.set("1");
    pCollectionConsumer.accept(valueInGlobalWindow("GHI"));
    pCollectionConsumer.accept(valueInGlobalWindow("JKL"));
    assertThat(output1, contains(valueInGlobalWindow("GHI"), valueInGlobalWindow("JKL")));
    verifyNoMoreInteractions(mockBeamFnDataClient);
}

Example 22 with PTransform

use of org.apache.beam.model.pipeline.v1.RunnerApi.PTransform in project beam by apache.

the class WindowMappingFnRunner method createMapFunctionForPTransform.

static <T, W1 extends BoundedWindow, W2 extends BoundedWindow> ThrowingFunction<KV<T, W1>, KV<T, W2>> createMapFunctionForPTransform(String ptransformId, PTransform pTransform) throws IOException {
    FunctionSpec windowMappingFnPayload = FunctionSpec.parseFrom(pTransform.getSpec().getPayload());
    WindowMappingFn<W2> windowMappingFn = (WindowMappingFn<W2>) PCollectionViewTranslation.windowMappingFnFromProto(windowMappingFnPayload);
    return (KV<T, W1> input) -> KV.of(input.getKey(), windowMappingFn.getSideInputWindow(input.getValue()));
}

Example 23 with PTransform

use of org.apache.beam.model.pipeline.v1.RunnerApi.PTransform in project beam by apache.

the class JavaClassLookupTransformProvider method findMappingConstructor.

private Constructor<PTransform<InputT, OutputT>> findMappingConstructor(Constructor<?>[] constructors, JavaClassLookupPayload payload) {
    Row constructorRow = decodeRow(payload.getConstructorSchema(), payload.getConstructorPayload());
    List<Constructor<?>> mappingConstructors = Arrays.stream(constructors).filter(c -> c.getParameterCount() == payload.getConstructorSchema().getFieldsCount()).filter(c -> parametersCompatible(c.getParameters(), constructorRow)).collect(Collectors.toList());
    if (mappingConstructors.size() == 0) {
        throw new RuntimeException("Could not find a matching constructor. When using field names, make sure they are " + "available in the compiled Java class.");
    } else if (mappingConstructors.size() != 1) {
        throw new RuntimeException("Expected to find a single mapping constructor but found " + mappingConstructors.size());
    }
    return (Constructor<PTransform<InputT, OutputT>>) mappingConstructors.get(0);
}

Example 24 with PTransform

use of org.apache.beam.model.pipeline.v1.RunnerApi.PTransform in project beam by apache.

the class RemoteGrpcPortReadTest method toFromPTransform.

@Test
public void toFromPTransform() throws InvalidProtocolBufferException {
    RemoteGrpcPort port = RemoteGrpcPort.newBuilder().setApiServiceDescriptor(ApiServiceDescriptor.newBuilder().setUrl("foo").setAuthentication(AuthenticationSpec.getDefaultInstance()).build()).build();
    RemoteGrpcPortRead read = RemoteGrpcPortRead.readFromPort(port, "myPort");
    PTransform ptransform = PTransform.parseFrom(read.toPTransform().toByteArray());
    RemoteGrpcPortRead serDeRead = RemoteGrpcPortRead.fromPTransform(ptransform);
    assertThat(serDeRead, equalTo(read));
    assertThat(serDeRead.getPort(), equalTo(read.getPort()));
    assertThat(serDeRead.toPTransform(), equalTo(ptransform));
}

Example 25 with PTransform

use of org.apache.beam.model.pipeline.v1.RunnerApi.PTransform in project beam by apache.

the class ParDoTranslation method getMainInputName.

/**
 * Returns the name of the main input of the ptransform.
 */
public static String getMainInputName(RunnerApi.PTransformOrBuilder ptransform) throws IOException {
    checkArgument(PAR_DO_TRANSFORM_URN.equals(ptransform.getSpec().getUrn()) || SPLITTABLE_PAIR_WITH_RESTRICTION_URN.equals(ptransform.getSpec().getUrn()) || SPLITTABLE_SPLIT_AND_SIZE_RESTRICTIONS_URN.equals(ptransform.getSpec().getUrn()) || SPLITTABLE_TRUNCATE_SIZED_RESTRICTION_URN.equals(ptransform.getSpec().getUrn()) || SPLITTABLE_PROCESS_ELEMENTS_URN.equals(ptransform.getSpec().getUrn()) || SPLITTABLE_PROCESS_SIZED_ELEMENTS_AND_RESTRICTIONS_URN.equals(ptransform.getSpec().getUrn()), "Unexpected payload type %s", ptransform.getSpec().getUrn());
    ParDoPayload payload = ParDoPayload.parseFrom(ptransform.getSpec().getPayload());
    return getMainInputName(ptransform, payload);
}