Search in sources :

Example 1 with TimerFamilyDeclaration

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

the class DoFnSignatures method analyzeTimerFamilyDeclarations.

private static ImmutableMap<String, TimerFamilyDeclaration> analyzeTimerFamilyDeclarations(ErrorReporter errors, Class<?> fnClazz) {
    Map<String, TimerFamilyDeclaration> declarations = new HashMap<>();
    for (Field field : ReflectHelpers.declaredFieldsWithAnnotation(DoFn.TimerFamily.class, fnClazz, DoFn.class)) {
        // TimerSpec fields may generally be private, but will be accessed via the signature
        field.setAccessible(true);
        String id = TimerFamilyDeclaration.PREFIX + field.getAnnotation(DoFn.TimerFamily.class).value();
        validateTimerFamilyField(errors, declarations, id, field);
        declarations.put(id, TimerFamilyDeclaration.create(id, field));
    }
    return ImmutableMap.copyOf(declarations);
}
Also used : TimerFamilyDeclaration(org.apache.beam.sdk.transforms.reflect.DoFnSignature.TimerFamilyDeclaration) Field(java.lang.reflect.Field) DoFn(org.apache.beam.sdk.transforms.DoFn) HashMap(java.util.HashMap) LinkedHashMap(java.util.LinkedHashMap) FormatString(com.google.errorprone.annotations.FormatString)

Example 2 with TimerFamilyDeclaration

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

the class DoFnSignatures method analyzeExtraParameter.

private static Parameter analyzeExtraParameter(ErrorReporter methodErrors, FnAnalysisContext fnContext, MethodAnalysisContext methodContext, ParameterDescription param, TypeDescriptor<?> inputT, TypeDescriptor<?> outputT) {
    TypeDescriptor<?> expectedProcessContextT = doFnProcessContextTypeOf(inputT, outputT);
    TypeDescriptor<?> expectedStartBundleContextT = doFnStartBundleContextTypeOf(inputT, outputT);
    TypeDescriptor<?> expectedFinishBundleContextT = doFnFinishBundleContextTypeOf(inputT, outputT);
    TypeDescriptor<?> expectedOnTimerContextT = doFnOnTimerContextTypeOf(inputT, outputT);
    TypeDescriptor<?> expectedOnWindowExpirationContextT = doFnOnWindowExpirationContextTypeOf(inputT, outputT);
    TypeDescriptor<?> paramT = param.getType();
    Class<?> rawType = paramT.getRawType();
    ErrorReporter paramErrors = methodErrors.forParameter(param);
    String fieldAccessString = getFieldAccessId(param.getAnnotations());
    if (fieldAccessString != null) {
        return Parameter.schemaElementParameter(paramT, fieldAccessString, param.getIndex());
    } else if (hasAnnotation(DoFn.Element.class, param.getAnnotations())) {
        return paramT.equals(inputT) ? Parameter.elementParameter(paramT) : Parameter.schemaElementParameter(paramT, null, param.getIndex());
    } else if (hasAnnotation(DoFn.Restriction.class, param.getAnnotations())) {
        return Parameter.restrictionParameter(paramT);
    } else if (hasAnnotation(DoFn.WatermarkEstimatorState.class, param.getAnnotations())) {
        return Parameter.watermarkEstimatorState(paramT);
    } else if (hasAnnotation(DoFn.Timestamp.class, param.getAnnotations())) {
        methodErrors.checkArgument(rawType.equals(Instant.class), "@Timestamp argument must have type org.joda.time.Instant.");
        return Parameter.timestampParameter();
    } else if (hasAnnotation(DoFn.Key.class, param.getAnnotations())) {
        methodErrors.checkArgument(KV.class.equals(inputT.getRawType()), "@Key argument is expected to be use with input element of type KV.");
        Type keyType = ((ParameterizedType) inputT.getType()).getActualTypeArguments()[0];
        methodErrors.checkArgument(TypeDescriptor.of(keyType).equals(paramT), "@Key argument is expected to be type of %s, but found %s.", keyType, rawType);
        return Parameter.keyT(paramT);
    } else if (rawType.equals(TimeDomain.class)) {
        return Parameter.timeDomainParameter();
    } else if (hasAnnotation(DoFn.SideInput.class, param.getAnnotations())) {
        String sideInputId = getSideInputId(param.getAnnotations());
        paramErrors.checkArgument(sideInputId != null, "%s missing %s annotation", sideInputId, format(SideInput.class));
        return Parameter.sideInputParameter(paramT, sideInputId);
    } else if (rawType.equals(PaneInfo.class)) {
        return Parameter.paneInfoParameter();
    } else if (rawType.equals(DoFn.BundleFinalizer.class)) {
        return Parameter.bundleFinalizer();
    } else if (rawType.equals(DoFn.ProcessContext.class)) {
        paramErrors.checkArgument(paramT.equals(expectedProcessContextT), "ProcessContext argument must have type %s", format(expectedProcessContextT));
        return Parameter.processContext();
    } else if (rawType.equals(DoFn.StartBundleContext.class)) {
        paramErrors.checkArgument(paramT.equals(expectedStartBundleContextT), "StartBundleContext argument must have type %s", format(expectedProcessContextT));
        return Parameter.startBundleContext();
    } else if (rawType.equals(DoFn.FinishBundleContext.class)) {
        paramErrors.checkArgument(paramT.equals(expectedFinishBundleContextT), "FinishBundleContext argument must have type %s", format(expectedProcessContextT));
        return Parameter.finishBundleContext();
    } else if (rawType.equals(DoFn.OnTimerContext.class)) {
        paramErrors.checkArgument(paramT.equals(expectedOnTimerContextT), "OnTimerContext argument must have type %s", format(expectedOnTimerContextT));
        return Parameter.onTimerContext();
    } else if (rawType.equals(DoFn.OnWindowExpirationContext.class)) {
        paramErrors.checkArgument(paramT.equals(expectedOnWindowExpirationContextT), "OnWindowExpirationContext argument must have type %s", format(expectedOnWindowExpirationContextT));
        return Parameter.onWindowExpirationContext();
    } else if (BoundedWindow.class.isAssignableFrom(rawType)) {
        methodErrors.checkArgument(!methodContext.hasParameter(WindowParameter.class), "Multiple %s parameters", format(BoundedWindow.class));
        return Parameter.boundedWindow((TypeDescriptor<? extends BoundedWindow>) paramT);
    } else if (rawType.equals(OutputReceiver.class)) {
        // It's a schema row receiver if it's an OutputReceiver<Row> _and_ the output type is not
        // already Row.
        boolean schemaRowReceiver = paramT.equals(outputReceiverTypeOf(TypeDescriptor.of(Row.class))) && !outputT.equals(TypeDescriptor.of(Row.class));
        if (!schemaRowReceiver) {
            TypeDescriptor<?> expectedReceiverT = outputReceiverTypeOf(outputT);
            paramErrors.checkArgument(paramT.equals(expectedReceiverT), "OutputReceiver should be parameterized by %s", outputT);
        }
        return Parameter.outputReceiverParameter(schemaRowReceiver);
    } else if (rawType.equals(MultiOutputReceiver.class)) {
        return Parameter.taggedOutputReceiverParameter();
    } else if (PipelineOptions.class.equals(rawType)) {
        methodErrors.checkArgument(!methodContext.hasParameter(PipelineOptionsParameter.class), "Multiple %s parameters", format(PipelineOptions.class));
        return Parameter.pipelineOptions();
    } else if (RestrictionTracker.class.isAssignableFrom(rawType)) {
        methodErrors.checkArgument(!methodContext.hasParameter(RestrictionTrackerParameter.class), "Multiple %s parameters", format(RestrictionTracker.class));
        return Parameter.restrictionTracker(paramT);
    } else if (WatermarkEstimator.class.isAssignableFrom(rawType)) {
        methodErrors.checkArgument(!methodContext.hasParameter(WatermarkEstimatorParameter.class), "Multiple %s parameters", format(WatermarkEstimator.class));
        return Parameter.watermarkEstimator(paramT);
    } else if (rawType.equals(Timer.class)) {
        // m.getParameters() is not available until Java 8
        String id = getTimerId(param.getAnnotations());
        paramErrors.checkArgument(id != null, "%s missing %s annotation", format(Timer.class), format(TimerId.class));
        paramErrors.checkArgument(!methodContext.getTimerParameters().containsKey(id), "duplicate %s: \"%s\"", format(TimerId.class), id);
        TimerDeclaration timerDecl = fnContext.getTimerDeclarations().get(id);
        paramErrors.checkArgument(timerDecl != null, "reference to undeclared %s: \"%s\"", format(TimerId.class), id);
        paramErrors.checkArgument(timerDecl.field().getDeclaringClass().equals(getDeclaringClass(param.getMethod())), "%s %s declared in a different class %s." + " Timers may be referenced only in the lexical scope where they are declared.", format(TimerId.class), id, timerDecl.field().getDeclaringClass().getName());
        return Parameter.timerParameter(timerDecl);
    } else if (hasAnnotation(DoFn.TimerId.class, param.getAnnotations())) {
        boolean isValidTimerIdForTimerFamily = fnContext.getTimerFamilyDeclarations().size() > 0 && rawType.equals(String.class);
        paramErrors.checkArgument(isValidTimerIdForTimerFamily, "%s not allowed here", format(DoFn.TimerId.class));
        return Parameter.timerIdParameter();
    } else if (rawType.equals(TimerMap.class)) {
        String id = getTimerFamilyId(param.getAnnotations());
        paramErrors.checkArgument(id != null, "%s missing %s annotation", format(TimerMap.class), format(DoFn.TimerFamily.class));
        paramErrors.checkArgument(!methodContext.getTimerFamilyParameters().containsKey(id), "duplicate %s: \"%s\"", format(DoFn.TimerFamily.class), id);
        TimerFamilyDeclaration timerDecl = fnContext.getTimerFamilyDeclarations().get(id);
        paramErrors.checkArgument(timerDecl != null, "reference to undeclared %s: \"%s\"", format(DoFn.TimerFamily.class), id);
        paramErrors.checkArgument(timerDecl.field().getDeclaringClass().equals(getDeclaringClass(param.getMethod())), "%s %s declared in a different class %s." + " Timers may be referenced only in the lexical scope where they are declared.", format(DoFn.TimerFamily.class), id, timerDecl.field().getDeclaringClass().getName());
        return Parameter.timerFamilyParameter(timerDecl);
    } else if (State.class.isAssignableFrom(rawType)) {
        // m.getParameters() is not available until Java 8
        String id = getStateId(param.getAnnotations());
        paramErrors.checkArgument(id != null, "missing %s annotation", format(DoFn.StateId.class));
        paramErrors.checkArgument(!methodContext.getStateParameters().containsKey(id), "duplicate %s: \"%s\"", format(DoFn.StateId.class), id);
        // By static typing this is already a well-formed State subclass
        TypeDescriptor<? extends State> stateType = (TypeDescriptor<? extends State>) param.getType();
        StateDeclaration stateDecl = fnContext.getStateDeclarations().get(id);
        paramErrors.checkArgument(stateDecl != null, "reference to undeclared %s: \"%s\"", format(DoFn.StateId.class), id);
        paramErrors.checkArgument(stateDecl.stateType().isSubtypeOf(stateType), "data type of reference to %s %s must be a supertype of %s", format(StateId.class), id, format(stateDecl.stateType()));
        paramErrors.checkArgument(stateDecl.field().getDeclaringClass().equals(getDeclaringClass(param.getMethod())), "%s %s declared in a different class %s." + " State may be referenced only in the class where it is declared.", format(StateId.class), id, stateDecl.field().getDeclaringClass().getName());
        boolean alwaysFetched = getStateAlwaysFetched(param.getAnnotations());
        if (alwaysFetched) {
            paramErrors.checkArgument(ReadableState.class.isAssignableFrom(rawType), "@AlwaysFetched can only be used on ReadableStates. It cannot be used on %s", format(stateDecl.stateType()));
        }
        return Parameter.stateParameter(stateDecl, alwaysFetched);
    } else {
        paramErrors.throwIllegalArgument("%s is not a valid context parameter.", format(paramT));
        // Unreachable
        return null;
    }
}
Also used : WindowParameter(org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.WindowParameter) RestrictionTracker(org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker) TimerFamilyDeclaration(org.apache.beam.sdk.transforms.reflect.DoFnSignature.TimerFamilyDeclaration) TimerId(org.apache.beam.sdk.transforms.DoFn.TimerId) StateId(org.apache.beam.sdk.transforms.DoFn.StateId) FormatString(com.google.errorprone.annotations.FormatString) TimerDeclaration(org.apache.beam.sdk.transforms.reflect.DoFnSignature.TimerDeclaration) MultiOutputReceiver(org.apache.beam.sdk.transforms.DoFn.MultiOutputReceiver) ParameterizedType(java.lang.reflect.ParameterizedType) RestrictionTrackerParameter(org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.RestrictionTrackerParameter) HasDefaultWatermarkEstimator(org.apache.beam.sdk.transforms.splittabledofn.HasDefaultWatermarkEstimator) WatermarkEstimator(org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimator) ManualWatermarkEstimator(org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator) BoundedWindow(org.apache.beam.sdk.transforms.windowing.BoundedWindow) SideInput(org.apache.beam.sdk.transforms.DoFn.SideInput) TimerMap(org.apache.beam.sdk.state.TimerMap) Instant(org.joda.time.Instant) ReadableState(org.apache.beam.sdk.state.ReadableState) Type(java.lang.reflect.Type) ParameterizedType(java.lang.reflect.ParameterizedType) DoFn(org.apache.beam.sdk.transforms.DoFn) TypeDescriptor(org.apache.beam.sdk.values.TypeDescriptor) Timer(org.apache.beam.sdk.state.Timer) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) ValueState(org.apache.beam.sdk.state.ValueState) OrderedListState(org.apache.beam.sdk.state.OrderedListState) WatermarkHoldState(org.apache.beam.sdk.state.WatermarkHoldState) SetState(org.apache.beam.sdk.state.SetState) MapState(org.apache.beam.sdk.state.MapState) State(org.apache.beam.sdk.state.State) BagState(org.apache.beam.sdk.state.BagState) Preconditions.checkState(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkState) ReadableState(org.apache.beam.sdk.state.ReadableState) Row(org.apache.beam.sdk.values.Row) StateDeclaration(org.apache.beam.sdk.transforms.reflect.DoFnSignature.StateDeclaration)

Example 3 with TimerFamilyDeclaration

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

the class DoFnSignatures method parseSignature.

/**
 * Analyzes a given {@link DoFn} class and extracts its {@link DoFnSignature}.
 */
private static DoFnSignature parseSignature(Class<? extends DoFn<?, ?>> fnClass) {
    DoFnSignature.Builder signatureBuilder = DoFnSignature.builder();
    ErrorReporter errors = new ErrorReporter(null, fnClass.getName());
    errors.checkArgument(DoFn.class.isAssignableFrom(fnClass), "Must be subtype of DoFn");
    signatureBuilder.setFnClass(fnClass);
    TypeDescriptor<? extends DoFn<?, ?>> fnT = TypeDescriptor.of(fnClass);
    // Extract the input and output type, and whether the fn is bounded.
    TypeDescriptor<?> inputT = null;
    TypeDescriptor<?> outputT = null;
    for (TypeDescriptor<?> supertype : fnT.getTypes()) {
        if (!supertype.getRawType().equals(DoFn.class)) {
            continue;
        }
        Type[] args = ((ParameterizedType) supertype.getType()).getActualTypeArguments();
        inputT = TypeDescriptor.of(args[0]);
        outputT = TypeDescriptor.of(args[1]);
    }
    errors.checkNotNull(inputT, "Unable to determine input type");
    // Find the state and timer declarations in advance of validating
    // method parameter lists
    FnAnalysisContext fnContext = FnAnalysisContext.create();
    fnContext.addStateDeclarations(analyzeStateDeclarations(errors, fnClass).values());
    fnContext.addTimerDeclarations(analyzeTimerDeclarations(errors, fnClass).values());
    fnContext.addTimerFamilyDeclarations(analyzeTimerFamilyDeclarations(errors, fnClass).values());
    fnContext.addFieldAccessDeclarations(analyzeFieldAccessDeclaration(errors, fnClass).values());
    Method processElementMethod = findAnnotatedMethod(errors, DoFn.ProcessElement.class, fnClass, true);
    Method startBundleMethod = findAnnotatedMethod(errors, DoFn.StartBundle.class, fnClass, false);
    Method finishBundleMethod = findAnnotatedMethod(errors, DoFn.FinishBundle.class, fnClass, false);
    Method setupMethod = findAnnotatedMethod(errors, DoFn.Setup.class, fnClass, false);
    Method teardownMethod = findAnnotatedMethod(errors, DoFn.Teardown.class, fnClass, false);
    Method onWindowExpirationMethod = findAnnotatedMethod(errors, DoFn.OnWindowExpiration.class, fnClass, false);
    Method getInitialRestrictionMethod = findAnnotatedMethod(errors, DoFn.GetInitialRestriction.class, fnClass, false);
    Method splitRestrictionMethod = findAnnotatedMethod(errors, DoFn.SplitRestriction.class, fnClass, false);
    Method truncateRestrictionMethod = findAnnotatedMethod(errors, TruncateRestriction.class, fnClass, false);
    Method getRestrictionCoderMethod = findAnnotatedMethod(errors, DoFn.GetRestrictionCoder.class, fnClass, false);
    Method newTrackerMethod = findAnnotatedMethod(errors, DoFn.NewTracker.class, fnClass, false);
    Method getSizeMethod = findAnnotatedMethod(errors, DoFn.GetSize.class, fnClass, false);
    Method getWatermarkEstimatorStateCoderMethod = findAnnotatedMethod(errors, DoFn.GetWatermarkEstimatorStateCoder.class, fnClass, false);
    Method getInitialWatermarkEstimatorStateMethod = findAnnotatedMethod(errors, DoFn.GetInitialWatermarkEstimatorState.class, fnClass, false);
    Method newWatermarkEstimatorMethod = findAnnotatedMethod(errors, DoFn.NewWatermarkEstimator.class, fnClass, false);
    Collection<Method> onTimerMethods = ReflectHelpers.declaredMethodsWithAnnotation(DoFn.OnTimer.class, fnClass, DoFn.class);
    HashMap<String, DoFnSignature.OnTimerMethod> onTimerMethodMap = Maps.newHashMapWithExpectedSize(onTimerMethods.size());
    for (Method onTimerMethod : onTimerMethods) {
        String id = TimerDeclaration.PREFIX + onTimerMethod.getAnnotation(DoFn.OnTimer.class).value();
        errors.checkArgument(fnContext.getTimerDeclarations().containsKey(id), "Callback %s is for undeclared timer %s", onTimerMethod, id);
        TimerDeclaration timerDecl = fnContext.getTimerDeclarations().get(id);
        errors.checkArgument(timerDecl.field().getDeclaringClass().equals(getDeclaringClass(onTimerMethod)), "Callback %s is for timer %s declared in a different class %s." + " Timer callbacks must be declared in the same lexical scope as their timer", onTimerMethod, id, timerDecl.field().getDeclaringClass().getCanonicalName());
        onTimerMethodMap.put(id, analyzeOnTimerMethod(errors, fnT, onTimerMethod, id, inputT, outputT, fnContext));
    }
    signatureBuilder.setOnTimerMethods(onTimerMethodMap);
    // Check for TimerFamily
    Collection<Method> onTimerFamilyMethods = ReflectHelpers.declaredMethodsWithAnnotation(DoFn.OnTimerFamily.class, fnClass, DoFn.class);
    HashMap<String, DoFnSignature.OnTimerFamilyMethod> onTimerFamilyMethodMap = Maps.newHashMapWithExpectedSize(onTimerFamilyMethods.size());
    for (Method onTimerFamilyMethod : onTimerFamilyMethods) {
        String id = TimerFamilyDeclaration.PREFIX + onTimerFamilyMethod.getAnnotation(DoFn.OnTimerFamily.class).value();
        errors.checkArgument(fnContext.getTimerFamilyDeclarations().containsKey(id), "Callback %s is for undeclared timerFamily %s", onTimerFamilyMethod, id);
        TimerFamilyDeclaration timerDecl = fnContext.getTimerFamilyDeclarations().get(id);
        errors.checkArgument(timerDecl.field().getDeclaringClass().equals(getDeclaringClass(onTimerFamilyMethod)), "Callback %s is for timerFamily %s declared in a different class %s." + " TimerFamily callbacks must be declared in the same lexical scope as their timer", onTimerFamilyMethod, id, timerDecl.field().getDeclaringClass().getCanonicalName());
        onTimerFamilyMethodMap.put(id, analyzeOnTimerFamilyMethod(errors, fnT, onTimerFamilyMethod, id, inputT, outputT, fnContext));
    }
    signatureBuilder.setOnTimerFamilyMethods(onTimerFamilyMethodMap);
    // those timers used in methods, once method parameter lists support timers.
    for (TimerDeclaration decl : fnContext.getTimerDeclarations().values()) {
        errors.checkArgument(onTimerMethodMap.containsKey(decl.id()), "No callback registered via %s for timer %s", format(DoFn.OnTimer.class), decl.id());
    }
    for (TimerFamilyDeclaration decl : fnContext.getTimerFamilyDeclarations().values()) {
        errors.checkArgument(onTimerFamilyMethodMap.containsKey(decl.id()), "No callback registered via %s for timerFamily %s", format(DoFn.OnTimerFamily.class), decl.id());
    }
    ErrorReporter processElementErrors = errors.forMethod(DoFn.ProcessElement.class, processElementMethod);
    DoFnSignature.ProcessElementMethod processElement = analyzeProcessElementMethod(processElementErrors, fnT, processElementMethod, inputT, outputT, fnContext);
    signatureBuilder.setProcessElement(processElement);
    if (startBundleMethod != null) {
        ErrorReporter startBundleErrors = errors.forMethod(DoFn.StartBundle.class, startBundleMethod);
        signatureBuilder.setStartBundle(analyzeStartBundleMethod(startBundleErrors, fnT, startBundleMethod, inputT, outputT, fnContext));
    }
    if (finishBundleMethod != null) {
        ErrorReporter finishBundleErrors = errors.forMethod(DoFn.FinishBundle.class, finishBundleMethod);
        signatureBuilder.setFinishBundle(analyzeFinishBundleMethod(finishBundleErrors, fnT, finishBundleMethod, inputT, outputT, fnContext));
    }
    if (setupMethod != null) {
        ErrorReporter setupErrors = errors.forMethod(DoFn.Setup.class, setupMethod);
        signatureBuilder.setSetup(analyzeSetupMethod(setupErrors, fnT, setupMethod, inputT, outputT, fnContext));
    }
    if (teardownMethod != null) {
        signatureBuilder.setTeardown(analyzeShutdownMethod(errors.forMethod(DoFn.Teardown.class, teardownMethod), teardownMethod));
    }
    if (onWindowExpirationMethod != null) {
        signatureBuilder.setOnWindowExpiration(analyzeOnWindowExpirationMethod(errors, fnT, onWindowExpirationMethod, inputT, outputT, fnContext));
    }
    if (processElement.isSplittable()) {
        ErrorReporter getInitialRestrictionErrors = errors.forMethod(DoFn.GetInitialRestriction.class, getInitialRestrictionMethod);
        getInitialRestrictionErrors.checkNotNull(getInitialRestrictionMethod, "Splittable, but does not define the required @%s method.", DoFnSignatures.format(DoFn.GetInitialRestriction.class));
        GetInitialRestrictionMethod initialRestrictionMethod = analyzeGetInitialRestrictionMethod(errors.forMethod(DoFn.GetInitialRestriction.class, getInitialRestrictionMethod), fnT, getInitialRestrictionMethod, inputT, outputT, fnContext);
        signatureBuilder.setGetInitialRestriction(initialRestrictionMethod);
        TypeDescriptor<?> restrictionT = initialRestrictionMethod.restrictionT();
        TypeDescriptor<?> watermarkEstimatorStateT = TypeDescriptors.voids();
        if (getInitialWatermarkEstimatorStateMethod != null) {
            GetInitialWatermarkEstimatorStateMethod initialWatermarkEstimatorStateMethod = analyzeGetInitialWatermarkEstimatorStateMethod(errors.forMethod(DoFn.GetInitialWatermarkEstimatorState.class, getInitialWatermarkEstimatorStateMethod), fnT, getInitialWatermarkEstimatorStateMethod, inputT, outputT, fnContext);
            watermarkEstimatorStateT = initialWatermarkEstimatorStateMethod.watermarkEstimatorStateT();
            signatureBuilder.setGetInitialWatermarkEstimatorState(initialWatermarkEstimatorStateMethod);
        }
        if (newTrackerMethod != null) {
            signatureBuilder.setNewTracker(analyzeNewTrackerMethod(errors.forMethod(DoFn.NewTracker.class, newTrackerMethod), fnT, newTrackerMethod, inputT, outputT, restrictionT, fnContext));
        } else {
            errors.forMethod(DoFn.NewTracker.class, null).checkArgument(restrictionT.isSubtypeOf(TypeDescriptor.of(HasDefaultTracker.class)), "Splittable, either @%s method must be defined or %s must implement %s.", format(DoFn.NewTracker.class), format(restrictionT), format(HasDefaultTracker.class));
        }
        if (splitRestrictionMethod != null) {
            signatureBuilder.setSplitRestriction(analyzeSplitRestrictionMethod(errors.forMethod(DoFn.SplitRestriction.class, splitRestrictionMethod), fnT, splitRestrictionMethod, inputT, outputT, restrictionT, fnContext));
        }
        if (truncateRestrictionMethod != null) {
            signatureBuilder.setTruncateRestriction(analyzeTruncateRestrictionMethod(errors.forMethod(TruncateRestriction.class, truncateRestrictionMethod), fnT, truncateRestrictionMethod, inputT, restrictionT, fnContext));
        }
        if (getSizeMethod != null) {
            signatureBuilder.setGetSize(analyzeGetSizeMethod(errors.forMethod(DoFn.GetSize.class, getSizeMethod), fnT, getSizeMethod, inputT, outputT, restrictionT, fnContext));
        }
        if (getRestrictionCoderMethod != null) {
            signatureBuilder.setGetRestrictionCoder(analyzeGetRestrictionCoderMethod(errors.forMethod(DoFn.GetRestrictionCoder.class, getRestrictionCoderMethod), fnT, getRestrictionCoderMethod));
        }
        if (getWatermarkEstimatorStateCoderMethod != null) {
            signatureBuilder.setGetWatermarkEstimatorStateCoder(analyzeGetWatermarkEstimatorStateCoderMethod(errors.forMethod(DoFn.GetWatermarkEstimatorStateCoder.class, getWatermarkEstimatorStateCoderMethod), fnT, getWatermarkEstimatorStateCoderMethod));
        }
        if (newWatermarkEstimatorMethod != null) {
            signatureBuilder.setNewWatermarkEstimator(analyzeNewWatermarkEstimatorMethod(errors.forMethod(DoFn.NewWatermarkEstimator.class, newWatermarkEstimatorMethod), fnT, newWatermarkEstimatorMethod, inputT, outputT, restrictionT, watermarkEstimatorStateT, fnContext));
        } else if (getInitialWatermarkEstimatorStateMethod != null) {
            errors.forMethod(DoFn.NewWatermarkEstimator.class, null).checkArgument(watermarkEstimatorStateT.isSubtypeOf(TypeDescriptor.of(HasDefaultWatermarkEstimator.class)), "Splittable, either @%s method must be defined or %s must implement %s.", format(DoFn.NewWatermarkEstimator.class), format(watermarkEstimatorStateT), format(HasDefaultWatermarkEstimator.class));
        }
    } else {
        // Validate that none of the splittable DoFn only methods have been declared.
        List<String> forbiddenMethods = new ArrayList<>();
        if (getInitialRestrictionMethod != null) {
            forbiddenMethods.add("@" + format(DoFn.GetInitialRestriction.class));
        }
        if (splitRestrictionMethod != null) {
            forbiddenMethods.add("@" + format(DoFn.SplitRestriction.class));
        }
        if (truncateRestrictionMethod != null) {
            forbiddenMethods.add("@" + format(TruncateRestriction.class));
        }
        if (newTrackerMethod != null) {
            forbiddenMethods.add("@" + format(DoFn.NewTracker.class));
        }
        if (getRestrictionCoderMethod != null) {
            forbiddenMethods.add("@" + format(DoFn.GetRestrictionCoder.class));
        }
        if (getSizeMethod != null) {
            forbiddenMethods.add("@" + format(DoFn.GetSize.class));
        }
        if (getInitialWatermarkEstimatorStateMethod != null) {
            forbiddenMethods.add("@" + format(DoFn.GetInitialWatermarkEstimatorState.class));
        }
        if (getWatermarkEstimatorStateCoderMethod != null) {
            forbiddenMethods.add("@" + format(DoFn.GetWatermarkEstimatorStateCoder.class));
        }
        if (newWatermarkEstimatorMethod != null) {
            forbiddenMethods.add("@" + format(DoFn.NewWatermarkEstimator.class));
        }
        errors.checkArgument(forbiddenMethods.isEmpty(), "Non-splittable, but defines methods: %s", forbiddenMethods);
    }
    signatureBuilder.setIsBoundedPerElement(inferBoundedness(fnT, processElement, errors));
    signatureBuilder.setStateDeclarations(fnContext.getStateDeclarations());
    signatureBuilder.setTimerDeclarations(fnContext.getTimerDeclarations());
    signatureBuilder.setTimerFamilyDeclarations(fnContext.getTimerFamilyDeclarations());
    signatureBuilder.setFieldAccessDeclarations(fnContext.getFieldAccessDeclarations());
    DoFnSignature signature = signatureBuilder.build();
    // Additional validation for splittable DoFn's.
    if (processElement.isSplittable()) {
        verifySplittableMethods(signature, errors);
    }
    return signature;
}
Also used : TimerFamilyDeclaration(org.apache.beam.sdk.transforms.reflect.DoFnSignature.TimerFamilyDeclaration) ArrayList(java.util.ArrayList) FormatString(com.google.errorprone.annotations.FormatString) TimerDeclaration(org.apache.beam.sdk.transforms.reflect.DoFnSignature.TimerDeclaration) ParameterizedType(java.lang.reflect.ParameterizedType) GetInitialWatermarkEstimatorStateMethod(org.apache.beam.sdk.transforms.reflect.DoFnSignature.GetInitialWatermarkEstimatorStateMethod) HasDefaultTracker(org.apache.beam.sdk.transforms.splittabledofn.HasDefaultTracker) Method(java.lang.reflect.Method) GetInitialRestrictionMethod(org.apache.beam.sdk.transforms.reflect.DoFnSignature.GetInitialRestrictionMethod) FormatMethod(com.google.errorprone.annotations.FormatMethod) GetInitialWatermarkEstimatorStateMethod(org.apache.beam.sdk.transforms.reflect.DoFnSignature.GetInitialWatermarkEstimatorStateMethod) HasDefaultWatermarkEstimator(org.apache.beam.sdk.transforms.splittabledofn.HasDefaultWatermarkEstimator) Type(java.lang.reflect.Type) ParameterizedType(java.lang.reflect.ParameterizedType) DoFn(org.apache.beam.sdk.transforms.DoFn) GetInitialRestrictionMethod(org.apache.beam.sdk.transforms.reflect.DoFnSignature.GetInitialRestrictionMethod)

Aggregations

FormatString (com.google.errorprone.annotations.FormatString)3 DoFn (org.apache.beam.sdk.transforms.DoFn)3 TimerFamilyDeclaration (org.apache.beam.sdk.transforms.reflect.DoFnSignature.TimerFamilyDeclaration)3 ParameterizedType (java.lang.reflect.ParameterizedType)2 Type (java.lang.reflect.Type)2 TimerDeclaration (org.apache.beam.sdk.transforms.reflect.DoFnSignature.TimerDeclaration)2 HasDefaultWatermarkEstimator (org.apache.beam.sdk.transforms.splittabledofn.HasDefaultWatermarkEstimator)2 FormatMethod (com.google.errorprone.annotations.FormatMethod)1 Field (java.lang.reflect.Field)1 Method (java.lang.reflect.Method)1 ArrayList (java.util.ArrayList)1 HashMap (java.util.HashMap)1 LinkedHashMap (java.util.LinkedHashMap)1 PipelineOptions (org.apache.beam.sdk.options.PipelineOptions)1 BagState (org.apache.beam.sdk.state.BagState)1 MapState (org.apache.beam.sdk.state.MapState)1 OrderedListState (org.apache.beam.sdk.state.OrderedListState)1 ReadableState (org.apache.beam.sdk.state.ReadableState)1 SetState (org.apache.beam.sdk.state.SetState)1 State (org.apache.beam.sdk.state.State)1