| ################################################################################ |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| ################################################################################ |
| |
| from pemja import findClass |
| |
| from pyflink.datastream.state import (ValueStateDescriptor, ListStateDescriptor, MapStateDescriptor, |
| StateDescriptor, ReducingStateDescriptor, |
| AggregatingStateDescriptor) |
| from pyflink.fn_execution.datastream.embedded.state_impl import (ValueStateImpl, ListStateImpl, |
| MapStateImpl, ReducingStateImpl, |
| AggregatingStateImpl) |
| from pyflink.fn_execution.embedded.converters import from_type_info |
| from pyflink.fn_execution.embedded.java_utils import to_java_state_descriptor |
| |
| JVoidNamespace = findClass('org.apache.flink.runtime.state.VoidNamespace') |
| JVoidNamespaceSerializer = findClass('org.apache.flink.runtime.state.VoidNamespaceSerializer') |
| |
| JVoidNamespace_INSTANCE = JVoidNamespace.INSTANCE |
| JVoidNamespaceSerializer_INSTANCE = JVoidNamespaceSerializer.INSTANCE |
| |
| |
| class KeyedStateBackend(object): |
| def __init__(self, |
| function_context, |
| keyed_state_backend, |
| window_serializer=JVoidNamespaceSerializer_INSTANCE, |
| window_converter=None): |
| self._function_context = function_context |
| self._keyed_state_backend = keyed_state_backend |
| self._window_serializer = window_serializer |
| self._window_converter = window_converter |
| |
| def get_current_key(self): |
| return self._function_context.get_current_key() |
| |
| def get_value_state(self, state_descriptor: ValueStateDescriptor) -> ValueStateImpl: |
| return ValueStateImpl( |
| self._get_or_create_keyed_state(state_descriptor), |
| from_type_info(state_descriptor.type_info), |
| self._window_converter) |
| |
| def get_list_state(self, state_descriptor: ListStateDescriptor) -> ListStateImpl: |
| return ListStateImpl( |
| self._get_or_create_keyed_state(state_descriptor), |
| from_type_info(state_descriptor.type_info), |
| self._window_converter) |
| |
| def get_map_state(self, state_descriptor: MapStateDescriptor) -> MapStateImpl: |
| return MapStateImpl( |
| self._get_or_create_keyed_state(state_descriptor), |
| from_type_info(state_descriptor.type_info), |
| self._window_converter) |
| |
| def get_reducing_state(self, state_descriptor: ReducingStateDescriptor): |
| return ReducingStateImpl( |
| self._get_or_create_keyed_state(state_descriptor), |
| from_type_info(state_descriptor.type_info), |
| state_descriptor.get_reduce_function(), |
| self._window_converter) |
| |
| def get_aggregating_state(self, state_descriptor: AggregatingStateDescriptor): |
| return AggregatingStateImpl( |
| self._get_or_create_keyed_state(state_descriptor), |
| from_type_info(state_descriptor.type_info), |
| state_descriptor.get_agg_function(), |
| self._window_converter) |
| |
| def _get_or_create_keyed_state(self, state_descriptor: StateDescriptor): |
| return self._keyed_state_backend.getPartitionedState( |
| JVoidNamespace_INSTANCE, |
| self._window_serializer, |
| to_java_state_descriptor(state_descriptor)) |
