def _create_bag_state(self, state_spec: userstate.StateSpec, encoded_namespace) \
-> userstate.AccumulatingRuntimeState:
if isinstance(state_spec, userstate.BagStateSpec):
bag_state = SynchronousBagRuntimeState(
self._state_handler,
state_key=self.get_bag_state_key(
state_spec.name, self._encoded_current_key, encoded_namespace),
value_coder=state_spec.coder)
return bag_state
else:
raise NotImplementedError(state_spec)
def _create_state(self, state_spec: userstate.StateSpec) -> userstate.AccumulatingRuntimeState:
if isinstance(state_spec, userstate.BagStateSpec):
bag_state = SynchronousBagRuntimeState(
self._state_handler,
state_key=beam_fn_api_pb2.StateKey(
bag_user_state=beam_fn_api_pb2.StateKey.BagUserState(
transform_id="",
user_state_id=state_spec.name,
key=self._encoded_current_key)),
value_coder=state_spec.coder)
return bag_state
else:
raise NotImplementedError(state_spec)