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

The following examples show how to use org.apache.beam.sdk.state.StateSpec. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example #1
Source File: FlinkStateInternals.java    From beam with Apache License 2.0 6 votes vote down vote up
@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;
}
 
Example #2
Source File: DataflowRunnerTest.java    From beam with Apache License 2.0 6 votes vote down vote up
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();
}
 
Example #3
Source File: ParDoTest.java    From beam with Apache License 2.0 6 votes vote down vote up
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));
    }
  };
}
 
Example #4
Source File: ParDoTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@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));
}
 
Example #5
Source File: ParDoTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@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));
}
 
Example #6
Source File: ParDo.java    From beam with Apache License 2.0 6 votes vote down vote up
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);
    }
  }
}
 
Example #7
Source File: DoFnSignatures.java    From beam with Apache License 2.0 6 votes vote down vote up
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()));
  }
}
 
Example #8
Source File: DataflowRunnerTest.java    From beam with Apache License 2.0 6 votes vote down vote up
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();
}
 
Example #9
Source File: DataflowRunnerTest.java    From beam with Apache License 2.0 6 votes vote down vote up
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();
}
 
Example #10
Source File: DoFnInvokersTest.java    From beam with Apache License 2.0 6 votes vote down vote up
/** 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);
}
 
Example #11
Source File: DoFnSignaturesTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example #12
Source File: DoFnSignaturesTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example #13
Source File: DoFnSignaturesTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example #14
Source File: StatefulParDoEvaluatorFactoryTest.java    From beam with Apache License 2.0 6 votes vote down vote up
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);
      }
    };
  }
 
Example #15
Source File: StatefulDoFnRunner.java    From beam with Apache License 2.0 6 votes vote down vote up
@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);
    }
  }
}
 
Example #16
Source File: FlinkStateInternals.java    From beam with Apache License 2.0 6 votes vote down vote up
@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;
}
 
Example #17
Source File: DoFnOperator.java    From beam with Apache License 2.0 6 votes vote down vote up
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));
      }
    }
  }
}
 
Example #18
Source File: ParDoTranslationTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@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));
}
 
Example #19
Source File: DoFnSignaturesTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@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>>() {}));
}
 
Example #20
Source File: DoFnSignaturesTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@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>>() {}));
}
 
Example #21
Source File: DoFnSignaturesTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example #22
Source File: DoFnSignaturesTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example #23
Source File: DoFnSignaturesTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example #24
Source File: DoFnSignaturesTest.java    From beam with Apache License 2.0 6 votes vote down vote up
@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());
}
 
Example #25
Source File: FlinkStateInternals.java    From beam with Apache License 2.0 5 votes vote down vote up
@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;
}
 
Example #26
Source File: FnApiStateAccessor.java    From beam with Apache License 2.0 5 votes vote down vote up
/**
 * @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.");
}
 
Example #27
Source File: FlinkStateInternals.java    From beam with Apache License 2.0 5 votes vote down vote up
@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;
}
 
Example #28
Source File: FnApiStateAccessor.java    From beam with Apache License 2.0 5 votes vote down vote up
@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.");
}
 
Example #29
Source File: FlinkStateInternals.java    From beam with Apache License 2.0 5 votes vote down vote up
@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;
}
 
Example #30
Source File: FnApiStateAccessor.java    From beam with Apache License 2.0 5 votes vote down vote up
@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();
                        }
                      })));
}