def _create_internal_map_state(self, name, encoded_namespace,
map_key_coder, map_value_coder, ttl_config,
cache_type):
# Currently the `beam_fn_api.proto` does not support MapState, so we use the
# the `MultimapSideInput` message to mark the state as a MapState for now.
from pyflink.fn_execution.flink_fn_execution_pb2 import StateDescriptor
state_proto = StateDescriptor()
state_proto.state_name = name
if ttl_config is not None:
state_proto.state_ttl_config.CopyFrom(ttl_config._to_proto())
state_key = beam_fn_api_pb2.StateKey(
multimap_side_input=beam_fn_api_pb2.StateKey.MultimapSideInput(
transform_id="",
window=encoded_namespace,
side_input_id=base64.b64encode(
state_proto.SerializeToString()),
key=self._encoded_current_key))
if cache_type == SynchronousKvRuntimeState.CacheType.DISABLE_CACHE:
write_cache_size = 0
else:
write_cache_size = self._map_state_write_cache_size
return InternalSynchronousMapRuntimeState(self._map_state_handler,
state_key, map_key_coder,
map_value_coder,
write_cache_size)
def get_bag_state_key(self, name, encoded_key, encoded_namespace,
ttl_config):
from pyflink.fn_execution.flink_fn_execution_pb2 import StateDescriptor
state_proto = StateDescriptor()
state_proto.state_name = name
if ttl_config is not None:
state_proto.state_ttl_config.CopyFrom(ttl_config._to_proto())
return beam_fn_api_pb2.StateKey(
bag_user_state=beam_fn_api_pb2.StateKey.BagUserState(
transform_id="",
window=encoded_namespace,
user_state_id=base64.b64encode(
state_proto.SerializeToString()),
key=encoded_key))
def _to_proto(self):
from pyflink.fn_execution.flink_fn_execution_pb2 import StateDescriptor
state_ttl_config = StateDescriptor.StateTTLConfig()
state_ttl_config.update_type = self._update_type._to_proto()
state_ttl_config.state_visibility = self._state_visibility._to_proto()
state_ttl_config.ttl_time_characteristic = self._ttl_time_characteristic._to_proto()
state_ttl_config.ttl = self._ttl.to_milliseconds()
state_ttl_config.cleanup_strategies.CopyFrom(self._cleanup_strategies._to_proto())
return state_ttl_config
