def get_state(self, state_descriptor: ValueStateDescriptor) -> ValueState:
if self._keyed_state_backend:
return self._keyed_state_backend.get_value_state(
state_descriptor.name,
from_type_info(state_descriptor.type_info))
else:
raise Exception(
"This state is only accessible by functions executed on a KeyedStream."
)
def get_reducing_state(
self, state_descriptor: ReducingStateDescriptor) -> ReducingState:
if self._keyed_state_backend:
return self._keyed_state_backend.get_reducing_state(
state_descriptor.get_name(),
from_type_info(state_descriptor.type_info),
state_descriptor.get_reduce_function())
else:
raise Exception(
"This state is only accessible by functions executed on a KeyedStream."
)
def get_list_state(self,
state_descriptor: ListStateDescriptor) -> ListState:
if self._keyed_state_backend:
array_coder = from_type_info(
state_descriptor.type_info) # type: GenericArrayCoder
return self._keyed_state_backend.get_list_state(
state_descriptor.name, array_coder._elem_coder)
else:
raise Exception(
"This state is only accessible by functions executed on a KeyedStream."
)
def get_map_state(self, state_descriptor: MapStateDescriptor) -> MapState:
if self._keyed_state_backend:
map_coder = from_type_info(
state_descriptor.type_info) # type: MapCoder
key_coder = map_coder._key_coder
value_coder = map_coder._value_coder
return self._keyed_state_backend.get_map_state(
state_descriptor.name, key_coder, value_coder)
else:
raise Exception(
"This state is only accessible by functions executed on a KeyedStream."
)
def get_aggregating_state(
self,
state_descriptor: AggregatingStateDescriptor) -> AggregatingState:
if self._keyed_state_backend:
return self._keyed_state_backend.get_aggregating_state(
state_descriptor.get_name(),
from_type_info(state_descriptor.type_info),
state_descriptor.get_agg_function(),
state_descriptor._ttl_config)
else:
raise Exception(
"This state is only accessible by functions executed on a KeyedStream."
)
def get_partitioned_state(self,
state_descriptor: StateDescriptor) -> State:
if isinstance(state_descriptor, ValueStateDescriptor):
state = self._state_backend.get_value_state(
state_descriptor.name,
from_type_info(state_descriptor.type_info))
elif isinstance(state_descriptor, ListStateDescriptor):
array_coder = from_type_info(
state_descriptor.type_info) # type: GenericArrayCoder
state = self._state_backend.get_list_state(state_descriptor.name,
array_coder._elem_coder)
elif isinstance(state_descriptor, MapStateDescriptor):
map_coder = from_type_info(
state_descriptor.type_info) # type: MapCoder
key_coder = map_coder._key_coder
value_coder = map_coder._value_coder
state = self._state_backend.get_map_state(state_descriptor.name,
key_coder, value_coder)
else:
raise Exception("Unknown supported StateDescriptor %s" %
state_descriptor)
state.set_current_namespace(self.window)
return state
