/** The {@link OutputReceiverParameter} for a main output, or null if there is none. */ @Nullable public OutputReceiverParameter getMainOutputReceiver() { Optional<Parameter> parameter = extraParameters() .stream() .filter(Predicates.instanceOf(OutputReceiverParameter.class)::apply) .findFirst(); return parameter.isPresent() ? ((OutputReceiverParameter) parameter.get()) : null; }
/** * Whether this {@link DoFn} reads a schema {@link PCollection} type as a {@link * org.apache.beam.sdk.values.Row} object. */ @Nullable public RowParameter getRowParameter() { Optional<Parameter> parameter = extraParameters() .stream() .filter(Predicates.instanceOf(RowParameter.class)::apply) .findFirst(); return parameter.isPresent() ? ((RowParameter) parameter.get()) : null; }
/** * Whether this {@link DoFn} is <a href="https://s.apache.org/splittable-do-fn">splittable</a>. */ public boolean isSplittable() { return extraParameters() .stream() .anyMatch(Predicates.instanceOf(RestrictionTrackerParameter.class)::apply); } }
private DoFnTester(DoFn<InputT, OutputT> origFn) { this.origFn = origFn; DoFnSignature signature = DoFnSignatures.signatureForDoFn(origFn); for (DoFnSignature.Parameter param : signature.processElement().extraParameters()) { param.match( new DoFnSignature.Parameter.Cases.WithDefault<Void>() {
}.getClass()); assertThat(sig.processElement().extraParameters().size(), equalTo(2)); .extraParameters() .get(1) .match(
@Test public void testBasicDoFnAllParameters() throws Exception { DoFnSignature sig = DoFnSignatures.getSignature( new DoFn<String, String>() { @ProcessElement public void process( @Element String element, @Timestamp Instant timestamp, BoundedWindow window, PaneInfo paneInfo, OutputReceiver<String> receiver, PipelineOptions options) {} }.getClass()); assertThat(sig.processElement().extraParameters().size(), equalTo(6)); assertThat(sig.processElement().extraParameters().get(0), instanceOf(ElementParameter.class)); assertThat(sig.processElement().extraParameters().get(1), instanceOf(TimestampParameter.class)); assertThat(sig.processElement().extraParameters().get(2), instanceOf(WindowParameter.class)); assertThat(sig.processElement().extraParameters().get(3), instanceOf(PaneInfoParameter.class)); assertThat( sig.processElement().extraParameters().get(4), instanceOf(OutputReceiverParameter.class)); assertThat( sig.processElement().extraParameters().get(5), instanceOf(PipelineOptionsParameter.class)); }
@Test public void testHasRestrictionTracker() throws Exception { DoFnSignature.ProcessElementMethod signature = analyzeProcessElementMethod( new AnonymousMethod() { private void method( DoFn<Integer, String>.ProcessContext context, SomeRestrictionTracker tracker) {} }); assertTrue(signature.isSplittable()); assertTrue( signature .extraParameters() .stream() .anyMatch( Predicates.instanceOf(DoFnSignature.Parameter.RestrictionTrackerParameter.class) ::apply)); assertEquals(SomeRestrictionTracker.class, signature.trackerT().getRawType()); }
/** * Whether this {@link DoFn} observes - directly or indirectly - the window that an element * resides in. * * <p>{@link State} and {@link Timer} parameters indirectly observe the window, because they are * each scoped to a single window. */ public boolean observesWindow() { return extraParameters() .stream() .anyMatch( Predicates.or( Predicates.instanceOf(WindowParameter.class), Predicates.instanceOf(TimerParameter.class), Predicates.instanceOf(StateParameter.class)) ::apply); }
@Test public void testDeclAndUsageOfStateInSuperclass() throws Exception { class DoFnOverridingAbstractStateUse extends DoFnDeclaringStateAndAbstractUse { @Override public void processWithState(ProcessContext c, ValueState<String> state) {} } DoFnSignature sig = DoFnSignatures.getSignature(new DoFnOverridingAbstractStateUse().getClass()); assertThat(sig.stateDeclarations().size(), equalTo(1)); assertThat(sig.processElement().extraParameters().size(), equalTo(2)); DoFnSignature.StateDeclaration decl = sig.stateDeclarations().get(DoFnOverridingAbstractStateUse.STATE_ID); StateParameter stateParam = (StateParameter) sig.processElement().extraParameters().get(1); assertThat( decl.field(), equalTo(DoFnDeclaringStateAndAbstractUse.class.getDeclaredField("myStateSpec"))); // The method we pull out is the superclass method; this is what allows validation to remain // simple. The later invokeDynamic instruction causes it to invoke the actual implementation. assertThat(stateParam.referent(), equalTo(decl)); }
@Test public void testDeclAndUsageOfTimerInSuperclass() throws Exception { DoFnSignature sig = DoFnSignatures.getSignature(new DoFnOverridingAbstractTimerUse().getClass()); assertThat(sig.timerDeclarations().size(), equalTo(1)); assertThat(sig.processElement().extraParameters().size(), equalTo(2)); DoFnSignature.TimerDeclaration decl = sig.timerDeclarations().get(DoFnOverridingAbstractTimerUse.TIMER_ID); TimerParameter timerParam = (TimerParameter) sig.processElement().extraParameters().get(1); assertThat( decl.field(), equalTo(DoFnDeclaringTimerAndAbstractUse.class.getDeclaredField("myTimerSpec"))); // The method we pull out is the superclass method; this is what allows validation to remain // simple. The later invokeDynamic instruction causes it to invoke the actual implementation. assertThat(timerParam.referent(), equalTo(decl)); }
@Override public List<RunnerApi.Parameter> translateParameters() { return ParDoTranslation.translateParameters( signature.processElement().extraParameters()); }
@Test public void testPipelineOptionsParameter() throws Exception { DoFnSignature sig = DoFnSignatures.getSignature( new DoFn<String, String>() { @ProcessElement public void process(ProcessContext c, PipelineOptions options) {} }.getClass()); assertThat( sig.processElement().extraParameters(), Matchers.hasItem(instanceOf(Parameter.PipelineOptionsParameter.class))); }
@Test public void testBasicDoFnProcessContext() throws Exception { DoFnSignature sig = DoFnSignatures.getSignature( new DoFn<String, String>() { @ProcessElement public void process(ProcessContext c) {} }.getClass()); assertThat(sig.processElement().extraParameters().size(), equalTo(1)); assertThat( sig.processElement().extraParameters().get(0), instanceOf(ProcessContextParameter.class)); }
@Override public List<Parameter> translateParameters() { return ParDoTranslation.translateParameters( signature.processElement().extraParameters()); }
@Test public void testBasicDoFnMultiOutputReceiver() throws Exception { DoFnSignature sig = DoFnSignatures.getSignature( new DoFn<String, String>() { @ProcessElement public void process(MultiOutputReceiver receiver) {} }.getClass()); assertThat(sig.processElement().extraParameters().size(), equalTo(1)); assertThat( sig.processElement().extraParameters().get(0), instanceOf(TaggedOutputReceiverParameter.class)); }
@Override public List<RunnerApi.Parameter> translateParameters() { return ParDoTranslation.translateParameters( signature.processElement().extraParameters()); }