org.apache.beam.sdk.state.StateSpec Java Examples

@Override
public <KeyT, ValueT> org.apache.beam.sdk.state.MapState<KeyT, ValueT> bindMap(
    String id,
    StateSpec<org.apache.beam.sdk.state.MapState<KeyT, ValueT>> spec,
    Coder<KeyT> mapKeyCoder,
    Coder<ValueT> mapValueCoder) {
  try {
    keyedStateBackend.getOrCreateKeyedState(
        StringSerializer.INSTANCE,
        new MapStateDescriptor<>(
            id,
            new CoderTypeSerializer<>(mapKeyCoder),
            new CoderTypeSerializer<>(mapValueCoder)));
  } catch (Exception e) {
    throw new RuntimeException(e);
  }
  return null;
}
 

private void verifySetStateUnsupported(PipelineOptions options) throws Exception {
  Pipeline p = Pipeline.create(options);
  p.apply(Create.of(KV.of(13, 42)))
      .apply(
          ParDo.of(
              new DoFn<KV<Integer, Integer>, Void>() {
                @StateId("fizzle")
                private final StateSpec<SetState<Void>> voidState = StateSpecs.set();

                @ProcessElement
                public void process() {}
              }));

  thrown.expectMessage("SetState");
  thrown.expect(UnsupportedOperationException.class);
  p.run();
}
 

private static DoFn<KV<String, Long>, Integer> timeSortedDoFn() {
  return new DoFn<KV<String, Long>, Integer>() {

    @StateId("last")
    private final StateSpec<ValueState<Long>> lastSpec = StateSpecs.value();

    @RequiresTimeSortedInput
    @ProcessElement
    public void process(
        @Element KV<String, Long> element,
        @StateId("last") ValueState<Long> last,
        OutputReceiver<Integer> output) {
      long lastVal = MoreObjects.firstNonNull(last.read(), element.getValue() - 1);
      last.write(element.getValue());
      output.output((int) (element.getValue() - lastVal));
    }
  };
}
 

@Test
public void testStateNotDeterministic() {
  final String stateId = "foo";

  // DoubleCoder is not deterministic, so this should crash
  DoFn<KV<Double, String>, Integer> fn =
      new DoFn<KV<Double, String>, Integer>() {

        @StateId(stateId)
        private final StateSpec<ValueState<Integer>> intState = StateSpecs.value();

        @ProcessElement
        public void processElement(
            ProcessContext c, @StateId(stateId) ValueState<Integer> state) {}
      };

  thrown.expect(IllegalArgumentException.class);
  thrown.expectMessage("state");
  thrown.expectMessage("deterministic");

  pipeline
      .apply(Create.of(KV.of(1.0, "hello"), KV.of(5.4, "goodbye"), KV.of(7.2, "hello again")))
      .apply(ParDo.of(fn));
}
 

@Test
public void testStateNotKeyed() {
  final String stateId = "foo";

  DoFn<String, Integer> fn =
      new DoFn<String, Integer>() {

        @StateId(stateId)
        private final StateSpec<ValueState<Integer>> intState = StateSpecs.value();

        @ProcessElement
        public void processElement(
            ProcessContext c, @StateId(stateId) ValueState<Integer> state) {}
      };

  thrown.expect(IllegalArgumentException.class);
  thrown.expectMessage("state");
  thrown.expectMessage("KvCoder");

  pipeline.apply(Create.of("hello", "goodbye", "hello again")).apply(ParDo.of(fn));
}
 

private static void populateDisplayData(
    DisplayData.Builder builder,
    DoFn<?, ?> fn,
    DisplayData.ItemSpec<? extends Class<?>> fnDisplayData) {
  builder.include("fn", fn).add(fnDisplayData);
  for (DoFnSignature.StateDeclaration stateDeclaration :
      DoFnSignatures.signatureForDoFn(fn).stateDeclarations().values()) {
    try {
      StateSpec<?> stateSpec = (StateSpec<?>) stateDeclaration.field().get(fn);
      builder.add(
          DisplayData.item("state_" + stateDeclaration.id(), stateDescription(stateSpec))
              .withLabel("State \"" + stateDeclaration.id() + "\""));
    } catch (IllegalAccessException e) {
      throw new RuntimeException(e);
    }
  }
}
 

public static StateSpec<?> getStateSpecOrThrow(
    StateDeclaration stateDeclaration, DoFn<?, ?> target) {
  try {
    Object fieldValue = stateDeclaration.field().get(target);
    checkState(
        fieldValue instanceof StateSpec,
        "Malformed %s class %s: state declaration field %s does not have type %s.",
        format(DoFn.class),
        target.getClass().getName(),
        stateDeclaration.field().getName(),
        StateSpec.class);

    return (StateSpec<?>) stateDeclaration.field().get(target);
  } catch (IllegalAccessException exc) {
    throw new RuntimeException(
        String.format(
            "Malformed %s class %s: state declaration field %s is not accessible.",
            format(DoFn.class), target.getClass().getName(), stateDeclaration.field().getName()));
  }
}
 

private void verifyMergingStatefulParDoRejected(PipelineOptions options) throws Exception {
  Pipeline p = Pipeline.create(options);

  p.apply(Create.of(KV.of(13, 42)))
      .apply(Window.into(Sessions.withGapDuration(Duration.millis(1))))
      .apply(
          ParDo.of(
              new DoFn<KV<Integer, Integer>, Void>() {
                @StateId("fizzle")
                private final StateSpec<ValueState<Void>> voidState = StateSpecs.value();

                @ProcessElement
                public void process() {}
              }));

  thrown.expectMessage("merging");
  thrown.expect(UnsupportedOperationException.class);
  p.run();
}
 

private void verifyMapStateUnsupported(PipelineOptions options) throws Exception {
  Pipeline p = Pipeline.create(options);
  p.apply(Create.of(KV.of(13, 42)))
      .apply(
          ParDo.of(
              new DoFn<KV<Integer, Integer>, Void>() {
                @StateId("fizzle")
                private final StateSpec<MapState<Void, Void>> voidState = StateSpecs.map();

                @ProcessElement
                public void process() {}
              }));

  thrown.expectMessage("MapState");
  thrown.expect(UnsupportedOperationException.class);
  p.run();
}
 

/** Tests that the generated {@link DoFnInvoker} passes the state parameter that it should. */
@Test
public void testDoFnWithState() throws Exception {
  ValueState<Integer> mockState = mock(ValueState.class);
  final String stateId = "my-state-id-here";
  when(mockArgumentProvider.state(stateId, false)).thenReturn(mockState);

  class MockFn extends DoFn<String, String> {
    @StateId(stateId)
    private final StateSpec<ValueState<Integer>> spec = StateSpecs.value(VarIntCoder.of());

    @ProcessElement
    public void processElement(ProcessContext c, @StateId(stateId) ValueState<Integer> valueState)
        throws Exception {}
  }

  MockFn fn = mock(MockFn.class);
  assertEquals(stop(), invokeProcessElement(fn));
  verify(fn).processElement(mockProcessContext, mockState);
}
 

@Test
public void testStateIdDuplicate() throws Exception {
  thrown.expect(IllegalArgumentException.class);
  thrown.expectMessage("Duplicate");
  thrown.expectMessage("StateId");
  thrown.expectMessage("my-id");
  thrown.expectMessage("myfield1");
  thrown.expectMessage("myfield2");
  thrown.expectMessage(not(mentionsTimers()));
  DoFnSignature sig =
      DoFnSignatures.getSignature(
          new DoFn<KV<String, Integer>, Long>() {
            @StateId("my-id")
            private final StateSpec<ValueState<Integer>> myfield1 =
                StateSpecs.value(VarIntCoder.of());

            @StateId("my-id")
            private final StateSpec<ValueState<Long>> myfield2 =
                StateSpecs.value(VarLongCoder.of());

            @ProcessElement
            public void foo(ProcessContext context) {}
          }.getClass());
}
 

@Test
public void testStateIdNonFinal() throws Exception {
  thrown.expect(IllegalArgumentException.class);
  thrown.expectMessage("State declarations must be final");
  thrown.expectMessage("Non-final field");
  thrown.expectMessage("myfield");
  thrown.expectMessage(not(mentionsTimers()));
  DoFnSignatures.getSignature(
      new DoFn<KV<String, Integer>, Long>() {
        @StateId("my-id")
        private StateSpec<ValueState<Integer>> myfield = StateSpecs.value(VarIntCoder.of());

        @ProcessElement
        public void foo(ProcessContext context) {}
      }.getClass());
}
 

@Test
public void testStateParameterAlwaysFetchNonReadableState() {
  thrown.expect(IllegalArgumentException.class);
  thrown.expectMessage("ReadableStates");
  DoFnSignatures.getSignature(
      new DoFn<KV<String, Integer>, Long>() {
        @StateId("my-id")
        private final StateSpec<MapState<Integer, Integer>> myfield =
            StateSpecs.map(VarIntCoder.of(), VarIntCoder.of());

        @ProcessElement
        public void myProcessElement(
            ProcessContext context,
            @AlwaysFetched @StateId("my-id") MapState<Integer, Integer> one) {}
      }.getClass());
}
 

private static DoFn<KV<String, Integer>, String> statefulConcat() {

    final String stateId = "sum";

    return new DoFn<KV<String, Integer>, String>() {

      @StateId(stateId)
      final StateSpec<ValueState<String>> stateSpec = StateSpecs.value();

      @ProcessElement
      @RequiresTimeSortedInput
      public void processElement(
          ProcessContext context, @StateId(stateId) ValueState<String> state) {
        String current = MoreObjects.firstNonNull(state.read(), "");
        if (!current.isEmpty()) {
          current += ":";
        }
        current += context.element().getValue();
        context.output(current);
        state.write(current);
      }
    };
  }
 

@Override
public void clearForWindow(W window) {
  for (Map.Entry<String, DoFnSignature.StateDeclaration> entry :
      signature.stateDeclarations().entrySet()) {
    try {
      StateSpec<?> spec = (StateSpec<?>) entry.getValue().field().get(fn);
      State state =
          stateInternals.state(
              StateNamespaces.window(windowCoder, window),
              StateTags.tagForSpec(entry.getKey(), (StateSpec) spec));
      state.clear();
    } catch (IllegalAccessException e) {
      throw new RuntimeException(e);
    }
  }
}
 

@Override
public <InputT, AccumT, OutputT>
    CombiningState<InputT, AccumT, OutputT> bindCombiningWithContext(
        String id,
        StateSpec<CombiningState<InputT, AccumT, OutputT>> spec,
        Coder<AccumT> accumCoder,
        CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn) {
  try {
    keyedStateBackend.getOrCreateKeyedState(
        StringSerializer.INSTANCE,
        new ValueStateDescriptor<>(id, new CoderTypeSerializer<>(accumCoder)));
  } catch (Exception e) {
    throw new RuntimeException(e);
  }
  return null;
}
 

private void earlyBindStateIfNeeded() throws IllegalArgumentException, IllegalAccessException {
  if (keyCoder != null) {
    if (doFn != null) {
      DoFnSignature signature = DoFnSignatures.getSignature(doFn.getClass());
      FlinkStateInternals.EarlyBinder earlyBinder =
          new FlinkStateInternals.EarlyBinder(getKeyedStateBackend());
      for (DoFnSignature.StateDeclaration value : signature.stateDeclarations().values()) {
        StateSpec<?> spec =
            (StateSpec<?>) signature.stateDeclarations().get(value.id()).field().get(doFn);
        spec.bind(value.id(), earlyBinder);
      }
      if (doFnRunner instanceof StatefulDoFnRunner) {
        ((StatefulDoFnRunner<InputT, OutputT, BoundedWindow>) doFnRunner)
            .getSystemStateTags()
            .forEach(tag -> tag.getSpec().bind(tag.getId(), earlyBinder));
      }
    }
  }
}
 

@Test
public void testStateSpecToFromProto() throws Exception {
  // Encode
  SdkComponents sdkComponents = SdkComponents.create();
  sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java"));
  RunnerApi.StateSpec stateSpecProto =
      ParDoTranslation.translateStateSpec(stateSpec, sdkComponents);

  // Decode
  RehydratedComponents rehydratedComponents =
      RehydratedComponents.forComponents(sdkComponents.toComponents());
  StateSpec<?> deserializedStateSpec =
      ParDoTranslation.fromProto(stateSpecProto, rehydratedComponents);

  assertThat(stateSpec, equalTo(deserializedStateSpec));
}
 

@Test
public void testSimpleStateIdNamedDoFn() throws Exception {
  class DoFnForTestSimpleStateIdNamedDoFn extends DoFn<KV<String, Integer>, Long> {
    @StateId("foo")
    private final StateSpec<ValueState<Integer>> bizzle = StateSpecs.value(VarIntCoder.of());

    @ProcessElement
    public void foo(ProcessContext context) {}
  }

  // Test classes at the bottom of the file
  DoFnSignature sig = DoFnSignatures.signatureForDoFn(new DoFnForTestSimpleStateIdNamedDoFn());

  assertThat(sig.stateDeclarations().size(), equalTo(1));
  DoFnSignature.StateDeclaration decl = sig.stateDeclarations().get("foo");

  assertThat(decl.id(), equalTo("foo"));
  assertThat(
      decl.field(), equalTo(DoFnForTestSimpleStateIdNamedDoFn.class.getDeclaredField("bizzle")));
  assertThat(
      decl.stateType(),
      Matchers.<TypeDescriptor<?>>equalTo(new TypeDescriptor<ValueState<Integer>>() {}));
}
 

@Test
public void testSimpleStateIdAnonymousDoFn() throws Exception {
  DoFnSignature sig =
      DoFnSignatures.getSignature(
          new DoFn<KV<String, Integer>, Long>() {
            @StateId("foo")
            private final StateSpec<ValueState<Integer>> bizzle =
                StateSpecs.value(VarIntCoder.of());

            @ProcessElement
            public void foo(ProcessContext context) {}
          }.getClass());

  assertThat(sig.stateDeclarations().size(), equalTo(1));
  DoFnSignature.StateDeclaration decl = sig.stateDeclarations().get("foo");

  assertThat(decl.id(), equalTo("foo"));
  assertThat(decl.field().getName(), equalTo("bizzle"));
  assertThat(
      decl.stateType(),
      Matchers.<TypeDescriptor<?>>equalTo(new TypeDescriptor<ValueState<Integer>>() {}));
}
 

@Test
public void testStateParameterWrongGenericType() throws Exception {
  thrown.expect(IllegalArgumentException.class);
  thrown.expectMessage("ValueState<String>");
  thrown.expectMessage("reference to");
  thrown.expectMessage("supertype");
  thrown.expectMessage("ValueState<Integer>");
  thrown.expectMessage("my-id");
  thrown.expectMessage("myProcessElement");
  thrown.expectMessage("index 1");
  thrown.expectMessage(not(mentionsTimers()));
  DoFnSignatures.getSignature(
      new DoFn<KV<String, Integer>, Long>() {
        @StateId("my-id")
        private final StateSpec<ValueState<Integer>> myfield = StateSpecs.value(VarIntCoder.of());

        @ProcessElement
        public void myProcessElement(
            ProcessContext context, @StateId("my-id") ValueState<String> stringState) {}
      }.getClass());
}
 

@Test
public void testStateParameterWrongStateType() throws Exception {
  thrown.expect(IllegalArgumentException.class);
  thrown.expectMessage("WatermarkHoldState");
  thrown.expectMessage("reference to");
  thrown.expectMessage("supertype");
  thrown.expectMessage("ValueState");
  thrown.expectMessage("my-id");
  thrown.expectMessage("myProcessElement");
  thrown.expectMessage("index 1");
  thrown.expectMessage(not(mentionsTimers()));
  DoFnSignatures.getSignature(
      new DoFn<KV<String, Integer>, Long>() {
        @StateId("my-id")
        private final StateSpec<ValueState<Integer>> myfield = StateSpecs.value(VarIntCoder.of());

        @ProcessElement
        public void myProcessElement(
            ProcessContext context, @StateId("my-id") WatermarkHoldState watermark) {}
      }.getClass());
}
 

@Test
public void testStateParameterDuplicate() throws Exception {
  thrown.expect(IllegalArgumentException.class);
  thrown.expectMessage("duplicate");
  thrown.expectMessage("my-id");
  thrown.expectMessage("myProcessElement");
  thrown.expectMessage("index 2");
  thrown.expectMessage(not(mentionsTimers()));
  DoFnSignatures.getSignature(
      new DoFn<KV<String, Integer>, Long>() {
        @StateId("my-id")
        private final StateSpec<ValueState<Integer>> myfield = StateSpecs.value(VarIntCoder.of());

        @ProcessElement
        public void myProcessElement(
            ProcessContext context,
            @StateId("my-id") ValueState<Integer> one,
            @StateId("my-id") ValueState<Integer> two) {}
      }.getClass());
}
 

@Test
public void testStateParameterAlwaysFetched() {
  thrown.expect(IllegalArgumentException.class);
  thrown.expectMessage("ReadableStates");
  DoFnSignature sig =
      DoFnSignatures.getSignature(
          new DoFn<KV<String, Integer>, Long>() {
            @StateId("my-id")
            private final StateSpec<MapState<Integer, Integer>> myfield =
                StateSpecs.map(VarIntCoder.of(), VarIntCoder.of());

            @ProcessElement
            public void myProcessElement(
                ProcessContext context,
                @AlwaysFetched @StateId("my-id") MapState<Integer, Integer> one) {}
          }.getClass());
  StateParameter stateParameter = (StateParameter) sig.processElement().extraParameters().get(1);
  assertTrue(stateParameter.alwaysFetched());
}
 

@Override
public <T> SetState<T> bindSet(String id, StateSpec<SetState<T>> spec, Coder<T> elemCoder) {
  try {
    keyedStateBackend.getOrCreateKeyedState(
        StringSerializer.INSTANCE,
        new MapStateDescriptor<>(
            id, new CoderTypeSerializer<>(elemCoder), VoidSerializer.INSTANCE));
  } catch (Exception e) {
    throw new RuntimeException(e);
  }
  return null;
}
 

/**
 * @deprecated The Fn API has no plans to implement WatermarkHoldState as of this writing and is
 *     waiting on resolution of BEAM-2535.
 */
@Override
@Deprecated
public WatermarkHoldState bindWatermark(
    String id, StateSpec<WatermarkHoldState> spec, TimestampCombiner timestampCombiner) {
  throw new UnsupportedOperationException("WatermarkHoldState is unsupported by the Fn API.");
}
 

@Override
public <InputT, AccumT, OutputT> CombiningState<InputT, AccumT, OutputT> bindCombining(
    String id,
    StateSpec<CombiningState<InputT, AccumT, OutputT>> spec,
    Coder<AccumT> accumCoder,
    Combine.CombineFn<InputT, AccumT, OutputT> combineFn) {
  try {
    keyedStateBackend.getOrCreateKeyedState(
        StringSerializer.INSTANCE,
        new ValueStateDescriptor<>(id, new CoderTypeSerializer<>(accumCoder)));
  } catch (Exception e) {
    throw new RuntimeException(e);
  }
  return null;
}
 

@Override
public <KeyT, ValueT> MapState<KeyT, ValueT> bindMap(
    String id,
    StateSpec<MapState<KeyT, ValueT>> spec,
    Coder<KeyT> mapKeyCoder,
    Coder<ValueT> mapValueCoder) {
  throw new UnsupportedOperationException("TODO: Add support for a map state to the Fn API.");
}
 

@Override
public WatermarkHoldState bindWatermark(
    String id, StateSpec<WatermarkHoldState> spec, TimestampCombiner timestampCombiner) {
  try {
    keyedStateBackend.getOrCreateKeyedState(
        VoidNamespaceSerializer.INSTANCE,
        new MapStateDescriptor<>(
            "watermark-holds",
            StringSerializer.INSTANCE,
            new CoderTypeSerializer<>(InstantCoder.of())));
  } catch (Exception e) {
    throw new RuntimeException(e);
  }
  return null;
}
 

@Override
public <ElementT, AccumT, ResultT>
    CombiningState<ElementT, AccumT, ResultT> bindCombiningWithContext(
        String id,
        StateSpec<CombiningState<ElementT, AccumT, ResultT>> spec,
        Coder<AccumT> accumCoder,
        CombineFnWithContext<ElementT, AccumT, ResultT> combineFn) {
  return (CombiningState<ElementT, AccumT, ResultT>)
      stateKeyObjectCache.computeIfAbsent(
          createBagUserStateKey(id),
          key ->
              bindCombining(
                  id,
                  spec,
                  accumCoder,
                  CombineFnUtil.bindContext(
                      combineFn,
                      new StateContext<BoundedWindow>() {
                        @Override
                        public PipelineOptions getPipelineOptions() {
                          return pipelineOptions;
                        }

                        @Override
                        public <T> T sideInput(PCollectionView<T> view) {
                          return get(view, currentWindowSupplier.get());
                        }

                        @Override
                        public BoundedWindow window() {
                          return currentWindowSupplier.get();
                        }
                      })));
}