Examples with SchemaElementParameter org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.SchemaElementParameter used on opensource projects

Example 1 with SchemaElementParameter

use of org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.SchemaElementParameter in project beam by apache.

the class ParDo method getDoFnSchemaInformation.

/**
 * Extract information on how the DoFn uses schemas. In particular, if the schema of an element
 * parameter does not match the input PCollection's schema, convert.
 */
@Internal
public static DoFnSchemaInformation getDoFnSchemaInformation(DoFn<?, ?> fn, PCollection<?> input) {
    DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
    DoFnSignature.ProcessElementMethod processElementMethod = signature.processElement();
    if (!processElementMethod.getSchemaElementParameters().isEmpty()) {
        if (!input.hasSchema()) {
            throw new IllegalArgumentException("Type of @Element must match the DoFn type" + input);
        }
    }
    SchemaRegistry schemaRegistry = input.getPipeline().getSchemaRegistry();
    DoFnSchemaInformation doFnSchemaInformation = DoFnSchemaInformation.create();
    for (SchemaElementParameter parameter : processElementMethod.getSchemaElementParameters()) {
        TypeDescriptor<?> elementT = parameter.elementT();
        FieldAccessDescriptor accessDescriptor = getFieldAccessDescriptorFromParameter(parameter.fieldAccessString(), input.getSchema(), signature.fieldAccessDeclarations(), fn);
        doFnSchemaInformation = doFnSchemaInformation.withFieldAccessDescriptor(accessDescriptor);
        Schema selectedSchema = SelectHelpers.getOutputSchema(input.getSchema(), accessDescriptor);
        ConvertHelpers.ConvertedSchemaInformation converted = ConvertHelpers.getConvertedSchemaInformation(selectedSchema, elementT, schemaRegistry);
        if (converted.outputSchemaCoder != null) {
            doFnSchemaInformation = doFnSchemaInformation.withSelectFromSchemaParameter((SchemaCoder<?>) input.getCoder(), accessDescriptor, selectedSchema, converted.outputSchemaCoder, converted.unboxedType != null);
        } else {
            // If the selected schema is a Row containing a single primitive type (which is the output
            // of Select when selecting a primitive), attempt to unbox it and match against the
            // parameter.
            checkArgument(converted.unboxedType != null);
            doFnSchemaInformation = doFnSchemaInformation.withUnboxPrimitiveParameter((SchemaCoder<?>) input.getCoder(), accessDescriptor, selectedSchema, elementT);
        }
    }
    for (DoFnSignature.Parameter p : processElementMethod.extraParameters()) {
        if (p instanceof ProcessContextParameter || p instanceof ElementParameter) {
            doFnSchemaInformation = doFnSchemaInformation.withFieldAccessDescriptor(FieldAccessDescriptor.withAllFields());
            break;
        }
    }
    return doFnSchemaInformation;
}

Example 2 with SchemaElementParameter

use of org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.SchemaElementParameter in project beam by apache.

the class DoFnSignatures method analyzeGetInitialWatermarkEstimatorStateMethod.

@VisibleForTesting
static DoFnSignature.GetInitialWatermarkEstimatorStateMethod analyzeGetInitialWatermarkEstimatorStateMethod(ErrorReporter errors, TypeDescriptor<? extends DoFn<?, ?>> fnT, Method m, TypeDescriptor<?> inputT, TypeDescriptor<?> outputT, FnAnalysisContext fnContext) {
    // Method is of the form:
    // @GetInitialWatermarkEstimatorState
    // WatermarkEstimatorStateT getInitialWatermarkEstimatorState(... parameters ...);
    Type[] params = m.getGenericParameterTypes();
    MethodAnalysisContext methodContext = MethodAnalysisContext.create();
    TypeDescriptor<? extends BoundedWindow> windowT = getWindowType(fnT, m);
    for (int i = 0; i < params.length; ++i) {
        Parameter extraParam = analyzeExtraParameter(errors, fnContext, methodContext, ParameterDescription.of(m, i, fnT.resolveType(params[i]), Arrays.asList(m.getParameterAnnotations()[i])), inputT, outputT);
        if (extraParam instanceof SchemaElementParameter) {
            errors.throwIllegalArgument("Schema @%s are not supported for @%s method. Found %s, did you mean to use %s?", format(DoFn.Element.class), format(DoFn.GetInitialWatermarkEstimatorState.class), format(((SchemaElementParameter) extraParam).elementT()), format(inputT));
        }
        methodContext.addParameter(extraParam);
    }
    for (Parameter parameter : methodContext.getExtraParameters()) {
        checkParameterOneOf(errors, parameter, ALLOWED_GET_INITIAL_WATERMARK_ESTIMATOR_STATE_PARAMETERS);
    }
    return DoFnSignature.GetInitialWatermarkEstimatorStateMethod.create(m, fnT.resolveType(m.getGenericReturnType()), windowT, methodContext.extraParameters);
}

Example 3 with SchemaElementParameter

use of org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.SchemaElementParameter in project beam by apache.

the class DoFnSignatures method analyzeSplitRestrictionMethod.

@VisibleForTesting
static DoFnSignature.SplitRestrictionMethod analyzeSplitRestrictionMethod(ErrorReporter errors, TypeDescriptor<? extends DoFn<?, ?>> fnT, Method m, TypeDescriptor<?> inputT, TypeDescriptor<?> outputT, TypeDescriptor<?> restrictionT, FnAnalysisContext fnContext) {
    // Method is of the form:
    // @SplitRestriction
    // void splitRestriction(... parameters ...);
    errors.checkArgument(void.class.equals(m.getReturnType()), "Must return void");
    Type[] params = m.getGenericParameterTypes();
    MethodAnalysisContext methodContext = MethodAnalysisContext.create();
    TypeDescriptor<? extends BoundedWindow> windowT = getWindowType(fnT, m);
    for (int i = 0; i < params.length; ++i) {
        Parameter extraParam = analyzeExtraParameter(errors, fnContext, methodContext, ParameterDescription.of(m, i, fnT.resolveType(params[i]), Arrays.asList(m.getParameterAnnotations()[i])), inputT, restrictionT);
        if (extraParam instanceof SchemaElementParameter) {
            errors.throwIllegalArgument("Schema @%s are not supported for @%s method. Found %s, did you mean to use %s?", format(DoFn.Element.class), format(DoFn.SplitRestriction.class), format(((SchemaElementParameter) extraParam).elementT()), format(inputT));
        } else if (extraParam instanceof RestrictionParameter) {
            errors.checkArgument(restrictionT.equals(((RestrictionParameter) extraParam).restrictionT()), "Uses restriction type %s, but @%s method uses restriction type %s", format(((RestrictionParameter) extraParam).restrictionT()), format(DoFn.GetInitialRestriction.class), format(restrictionT));
        }
        methodContext.addParameter(extraParam);
    }
    for (Parameter parameter : methodContext.getExtraParameters()) {
        checkParameterOneOf(errors, parameter, ALLOWED_SPLIT_RESTRICTION_PARAMETERS);
    }
    return DoFnSignature.SplitRestrictionMethod.create(m, windowT, methodContext.getExtraParameters());
}

Example 4 with SchemaElementParameter

use of org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.SchemaElementParameter in project beam by apache.

the class DoFnSignatures method analyzeProcessElementMethod.

@VisibleForTesting
static DoFnSignature.ProcessElementMethod analyzeProcessElementMethod(ErrorReporter errors, TypeDescriptor<? extends DoFn<?, ?>> fnClass, Method m, TypeDescriptor<?> inputT, TypeDescriptor<?> outputT, FnAnalysisContext fnContext) {
    errors.checkArgument(void.class.equals(m.getReturnType()) || DoFn.ProcessContinuation.class.equals(m.getReturnType()), "Must return void or %s", format(DoFn.ProcessContinuation.class));
    MethodAnalysisContext methodContext = MethodAnalysisContext.create();
    boolean requiresStableInput = m.isAnnotationPresent(DoFn.RequiresStableInput.class);
    boolean requiresTimeSortedInput = m.isAnnotationPresent(DoFn.RequiresTimeSortedInput.class);
    TypeDescriptor<? extends BoundedWindow> windowT = getWindowType(fnClass, m);
    Type[] params = m.getGenericParameterTypes();
    for (int i = 0; i < params.length; ++i) {
        Parameter extraParam = analyzeExtraParameter(errors.forMethod(DoFn.ProcessElement.class, m), fnContext, methodContext, ParameterDescription.of(m, i, fnClass.resolveType(params[i]), Arrays.asList(m.getParameterAnnotations()[i])), inputT, outputT);
        methodContext.addParameter(extraParam);
    }
    int schemaElementIndex = 0;
    for (int i = 0; i < methodContext.getExtraParameters().size(); ++i) {
        Parameter parameter = methodContext.getExtraParameters().get(i);
        if (parameter instanceof SchemaElementParameter) {
            SchemaElementParameter schemaParameter = (SchemaElementParameter) parameter;
            schemaParameter = schemaParameter.toBuilder().setIndex(schemaElementIndex).build();
            methodContext.setParameter(i, schemaParameter);
            ++schemaElementIndex;
        }
    }
    TypeDescriptor<?> trackerT = methodContext.findParameter(RestrictionTrackerParameter.class).map(p -> p.trackerT()).orElse(null);
    TypeDescriptor<?> watermarkEstimatorT = methodContext.findParameter(WatermarkEstimatorParameter.class).map(p -> p.estimatorT()).orElse(null);
    // The allowed parameters depend on whether this DoFn is splittable
    if (trackerT != null) {
        for (Parameter parameter : methodContext.getExtraParameters()) {
            checkParameterOneOf(errors, parameter, ALLOWED_SPLITTABLE_PROCESS_ELEMENT_PARAMETERS);
        }
    } else {
        for (Parameter parameter : methodContext.getExtraParameters()) {
            checkParameterOneOf(errors, parameter, ALLOWED_NON_SPLITTABLE_PROCESS_ELEMENT_PARAMETERS);
        }
    }
    return DoFnSignature.ProcessElementMethod.create(m, methodContext.getExtraParameters(), requiresStableInput, requiresTimeSortedInput, trackerT, watermarkEstimatorT, windowT, DoFn.ProcessContinuation.class.equals(m.getReturnType()));
}

Example 5 with SchemaElementParameter

use of org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.SchemaElementParameter in project beam by apache.

the class DoFnSignatures method analyzeGetInitialRestrictionMethod.

@VisibleForTesting
static DoFnSignature.GetInitialRestrictionMethod analyzeGetInitialRestrictionMethod(ErrorReporter errors, TypeDescriptor<? extends DoFn<?, ?>> fnT, Method m, TypeDescriptor<?> inputT, TypeDescriptor<?> outputT, FnAnalysisContext fnContext) {
    // Method is of the form:
    // @GetInitialRestriction
    // RestrictionT getInitialRestriction(... parameters ...);
    Type[] params = m.getGenericParameterTypes();
    MethodAnalysisContext methodContext = MethodAnalysisContext.create();
    TypeDescriptor<? extends BoundedWindow> windowT = getWindowType(fnT, m);
    for (int i = 0; i < params.length; ++i) {
        Parameter extraParam = analyzeExtraParameter(errors, fnContext, methodContext, ParameterDescription.of(m, i, fnT.resolveType(params[i]), Arrays.asList(m.getParameterAnnotations()[i])), inputT, outputT);
        if (extraParam instanceof SchemaElementParameter) {
            errors.throwIllegalArgument("Schema @%s are not supported for @%s method. Found %s, did you mean to use %s?", format(DoFn.Element.class), format(DoFn.GetInitialRestriction.class), format(((SchemaElementParameter) extraParam).elementT()), format(inputT));
        }
        methodContext.addParameter(extraParam);
    }
    for (Parameter parameter : methodContext.getExtraParameters()) {
        checkParameterOneOf(errors, parameter, ALLOWED_GET_INITIAL_RESTRICTION_PARAMETERS);
    }
    return DoFnSignature.GetInitialRestrictionMethod.create(m, fnT.resolveType(m.getGenericReturnType()), windowT, methodContext.extraParameters);
}