public StateInternalsStateCleaner( DoFn<?, ?> fn, StateInternals stateInternals, Coder<W> windowCoder) { this.fn = fn; this.signature = DoFnSignatures.getSignature(fn.getClass()); this.stateInternals = stateInternals; this.windowCoder = windowCoder; }
/** @return the {@link DoFnInvoker} for the given {@link DoFn}. */ @SuppressWarnings({"unchecked", "rawtypes"}) public <InputT, OutputT> DoFnInvoker<InputT, OutputT> newByteBuddyInvoker( DoFn<InputT, OutputT> fn) { return newByteBuddyInvoker(DoFnSignatures.getSignature((Class) fn.getClass()), fn); }
private static void rejectSplittable(DoFn<?, ?> doFn) { DoFnSignature signature = DoFnSignatures.getSignature(doFn.getClass()); if (signature.processElement().isSplittable()) { throw new UnsupportedOperationException( String.format( "%s does not currently support splittable DoFn: %s", FlinkRunner.class.getSimpleName(), doFn)); } }
private static <InputT, OutputT> void verifyFnIsStateful(DoFn<InputT, OutputT> fn) { DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass()); // It is still correct to use this without state or timers, but a bad idea. // Since it is internal it should never be used wrong, so it is OK to crash. checkState( signature.usesState() || signature.usesTimers(), "%s used for %s that does not use state or timers.", BatchStatefulParDoOverrides.class.getSimpleName(), ParDo.class.getSimpleName()); } }
@Test public void testIdentityFnApplied() throws Exception { DoFnSignatures.getSignature(new IdentityFn<String>() {}.getClass()); } }
@Test public void testOnWindowExpirationNoParam() { DoFnSignature sig = DoFnSignatures.getSignature( new DoFn<String, String>() { @ProcessElement public void process(ProcessContext c) {} @OnWindowExpiration public void bar() {} }.getClass()); assertThat(sig.onWindowExpiration().extraParameters().size(), equalTo(0)); }
@Test public void testSplittableMissingRequiredMethods() throws Exception { class BadFn extends DoFn<Integer, String> { @ProcessElement public void process( ProcessContext context, RestrictionTracker<SomeRestriction, Void> tracker) {} } thrown.expectMessage( "Splittable, but does not define the following required methods: " + "[@GetInitialRestriction, @NewTracker]"); DoFnSignatures.getSignature(BadFn.class); }
@Test public void testUnsplittableButDeclaresBounded() throws Exception { @BoundedPerElement class SomeFn extends DoFn<Integer, String> { @ProcessElement public void process(ProcessContext context) {} } thrown.expectMessage("Non-splittable, but annotated as @Bounded"); DoFnSignatures.getSignature(SomeFn.class); }
@Test public void testNoProcessElement() throws Exception { thrown.expect(IllegalArgumentException.class); thrown.expectMessage("No method annotated with @ProcessElement found"); thrown.expectMessage(getClass().getName() + "$"); DoFnSignatures.getSignature(new DoFn<String, String>() {}.getClass()); }
@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 testWrongTimestampType() throws Exception { thrown.expect(IllegalArgumentException.class); thrown.expectMessage("@Timestamp argument must have type org.joda.time.Instant"); DoFnSignature sig = DoFnSignatures.getSignature( new DoFn<String, String>() { @ProcessElement public void process(@Timestamp String timestamp) {} }.getClass()); }
@Test public void testWrongOutputReceiverType() throws Exception { thrown.expect(IllegalArgumentException.class); thrown.expectMessage("OutputReceiver should be parameterized by java.lang.String"); DoFnSignature sig = DoFnSignatures.getSignature( new DoFn<String, String>() { @ProcessElement public void process(OutputReceiver<Integer> receiver) {} }.getClass()); }
@Test public void testMissingFieldAccess() { thrown.expect(IllegalArgumentException.class); thrown.expectMessage("No FieldAccessDescriptor defined."); DoFnSignature sig = DoFnSignatures.getSignature( new DoFn<String, String>() { @ProcessElement public void process(@FieldAccess("foo") Row row) {} }.getClass()); }
@Test public void testGoodStateParameterSuperclassStateType() throws Exception { DoFnSignatures.getSignature( new DoFn<KV<String, Integer>, Long>() { @StateId("my-id") private final StateSpec<CombiningState<Integer, int[], Integer>> state = StateSpecs.combining(Sum.ofIntegers()); @ProcessElement public void myProcessElement( ProcessContext context, @StateId("my-id") GroupingState<Integer, Integer> groupingState) {} }.getClass()); }
@Test public void testSplittableBoundednessInferredFromReturnValue() throws Exception { assertEquals( PCollection.IsBounded.BOUNDED, DoFnSignatures.getSignature(BaseFnWithoutContinuation.class).isBoundedPerElement()); assertEquals( PCollection.IsBounded.UNBOUNDED, DoFnSignatures.getSignature(BaseFnWithContinuation.class).isBoundedPerElement()); }
@Test public void testRowParameterWithoutFieldAccess() { DoFnSignature sig = DoFnSignatures.getSignature( new DoFn<String, String>() { @ProcessElement public void process(@Element Row row) {} }.getClass()); assertThat(sig.processElement().getRowParameter(), notNullValue()); }
@Test public void testUnsplittableIsBounded() throws Exception { class UnsplittableFn extends DoFn<Integer, String> { @ProcessElement public void process(ProcessContext context) {} } assertEquals( PCollection.IsBounded.BOUNDED, DoFnSignatures.getSignature(UnsplittableFn.class).isBoundedPerElement()); }
@Test public void testStateParameterNoAnnotation() throws Exception { thrown.expect(IllegalArgumentException.class); thrown.expectMessage("missing StateId annotation"); thrown.expectMessage("myProcessElement"); thrown.expectMessage("index 1"); thrown.expectMessage(not(mentionsTimers())); DoFnSignatures.getSignature( new DoFn<KV<String, Integer>, Long>() { @ProcessElement public void myProcessElement(ProcessContext context, ValueState<Integer> noAnnotation) {} }.getClass()); }
@Test public void testDeclOfStateUsedInSuperclass() throws Exception { thrown.expect(IllegalArgumentException.class); thrown.expectMessage("process"); thrown.expectMessage("declared in a different class"); thrown.expectMessage(DoFnUsingState.STATE_ID); DoFnSignatures.getSignature( new DoFnUsingState() { @StateId(DoFnUsingState.STATE_ID) private final StateSpec<ValueState<Integer>> spec = StateSpecs.value(VarIntCoder.of()); }.getClass()); }
@Test public void testBadTypeVariables() throws Exception { thrown.expect(IllegalArgumentException.class); thrown.expectMessage("DoFn<InputT, InputT>.ProcessContext"); thrown.expectMessage("must have type"); thrown.expectMessage("DoFn<InputT, OutputT>.ProcessContext"); DoFnSignatures.getSignature(BadTypeVariables.class); }