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

Example 16 with FunctionSpec

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

the class WindowingStrategyTranslation method toProto.

/**
 * Converts a {@link WindowingStrategy} into a {@link RunnerApi.WindowingStrategy}, registering
 * any components in the provided {@link SdkComponents}.
 */
public static RunnerApi.WindowingStrategy toProto(WindowingStrategy<?, ?> windowingStrategy, SdkComponents components) throws IOException {
    WindowFn<?, ?> windowFn = windowingStrategy.getWindowFn();
    FunctionSpec windowFnSpec = toProto(windowFn, components);
    String environmentId = Strings.isNullOrEmpty(windowingStrategy.getEnvironmentId()) ? components.getEnvironmentIdFor(ResourceHints.create()) : windowingStrategy.getEnvironmentId();
    RunnerApi.WindowingStrategy.Builder windowingStrategyProto = RunnerApi.WindowingStrategy.newBuilder().setOutputTime(toProto(windowingStrategy.getTimestampCombiner())).setAccumulationMode(toProto(windowingStrategy.getMode())).setClosingBehavior(toProto(windowingStrategy.getClosingBehavior())).setAllowedLateness(windowingStrategy.getAllowedLateness().getMillis()).setTrigger(TriggerTranslation.toProto(windowingStrategy.getTrigger())).setWindowFn(windowFnSpec).setAssignsToOneWindow(windowFn.assignsToOneWindow()).setMergeStatus(windowFn.isNonMerging() ? MergeStatus.Enum.NON_MERGING : (windowingStrategy.isAlreadyMerged() ? MergeStatus.Enum.ALREADY_MERGED : MergeStatus.Enum.NEEDS_MERGE)).setOnTimeBehavior(toProto(windowingStrategy.getOnTimeBehavior())).setWindowCoderId(components.registerCoder(windowFn.windowCoder())).setEnvironmentId(environmentId);
    return windowingStrategyProto.build();
}

Example 17 with FunctionSpec

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

the class CreatePCollectionViewTranslationTest method testExtractionDirectFromTransform.

@Test
public void testExtractionDirectFromTransform() throws Exception {
    SdkComponents components = SdkComponents.create();
    components.registerEnvironment(Environments.createDockerEnvironment("java"));
    components.registerPCollection(testPCollection);
    AppliedPTransform<?, ?, ?> appliedPTransform = AppliedPTransform.of("foo", PValues.expandInput(testPCollection), PValues.expandOutput(createViewTransform.getView()), createViewTransform, ResourceHints.create(), p);
    CreatePCollectionViewTranslation.getView((AppliedPTransform) appliedPTransform);
    FunctionSpec payload = PTransformTranslation.toProto(appliedPTransform, components).getSpec();
    // Checks that the payload is what it should be
    PCollectionView<?> deserializedView = (PCollectionView<?>) SerializableUtils.deserializeFromByteArray(payload.getPayload().toByteArray(), PCollectionView.class.getSimpleName());
    assertThat(deserializedView, Matchers.equalTo(createViewTransform.getView()));
}

Example 18 with FunctionSpec

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

the class JavaClassLookupTransformProvider method getTransform.

@Override
public PTransform<PInput, POutput> getTransform(FunctionSpec spec) {
    JavaClassLookupPayload payload;
    try {
        payload = JavaClassLookupPayload.parseFrom(spec.getPayload());
    } catch (InvalidProtocolBufferException e) {
        throw new IllegalArgumentException("Invalid payload type for URN " + getUrn(ExpansionMethods.Enum.JAVA_CLASS_LOOKUP), e);
    }
    String className = payload.getClassName();
    try {
        AllowedClass allowlistClass = allowList.getAllowedClass(className);
        Class<PTransform<InputT, OutputT>> transformClass = (Class<PTransform<InputT, OutputT>>) ReflectHelpers.findClassLoader().loadClass(className);
        PTransform<PInput, POutput> transform;
        Row constructorRow = decodeRow(payload.getConstructorSchema(), payload.getConstructorPayload());
        if (payload.getConstructorMethod().isEmpty()) {
            Constructor<?>[] constructors = transformClass.getConstructors();
            Constructor<PTransform<InputT, OutputT>> constructor = findMappingConstructor(constructors, payload);
            Object[] parameterValues = getParameterValues(constructor.getParameters(), constructorRow, constructor.getGenericParameterTypes());
            transform = (PTransform<PInput, POutput>) constructor.newInstance(parameterValues);
        } else {
            Method[] methods = transformClass.getMethods();
            Method method = findMappingConstructorMethod(methods, payload, allowlistClass);
            Object[] parameterValues = getParameterValues(method.getParameters(), constructorRow, method.getGenericParameterTypes());
            transform = (PTransform<PInput, POutput>) method.invoke(null, /* static */
            parameterValues);
        }
        return applyBuilderMethods(transform, payload, allowlistClass);
    } catch (ClassNotFoundException e) {
        throw new IllegalArgumentException("Could not find class " + className, e);
    } catch (InstantiationException | IllegalArgumentException | IllegalAccessException | InvocationTargetException e) {
        throw new IllegalArgumentException("Could not instantiate class " + className, e);
    }
}