| # |
| # 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. |
| # |
| |
| """Support for user state in the BundleBasedDirectRunner.""" |
| from __future__ import absolute_import |
| |
| import itertools |
| |
| from apache_beam.transforms import userstate |
| from apache_beam.transforms.trigger import _ListStateTag |
| from apache_beam.transforms.trigger import _SetStateTag |
| |
| |
| class DirectRuntimeState(userstate.RuntimeState): |
| def __init__(self, state_spec, state_tag, current_value_accessor): |
| self._state_spec = state_spec |
| self._state_tag = state_tag |
| self._current_value_accessor = current_value_accessor |
| |
| @staticmethod |
| def for_spec(state_spec, state_tag, current_value_accessor): |
| if isinstance(state_spec, userstate.BagStateSpec): |
| return BagRuntimeState(state_spec, state_tag, current_value_accessor) |
| elif isinstance(state_spec, userstate.CombiningValueStateSpec): |
| return CombiningValueRuntimeState(state_spec, state_tag, |
| current_value_accessor) |
| elif isinstance(state_spec, userstate.SetStateSpec): |
| return SetRuntimeState(state_spec, state_tag, current_value_accessor) |
| else: |
| raise ValueError('Invalid state spec: %s' % state_spec) |
| |
| def _encode(self, value): |
| return self._state_spec.coder.encode(value) |
| |
| def _decode(self, value): |
| return self._state_spec.coder.decode(value) |
| |
| |
| # Sentinel designating an unread value. |
| UNREAD_VALUE = object() |
| |
| |
| class BagRuntimeState(DirectRuntimeState, userstate.BagRuntimeState): |
| def __init__(self, state_spec, state_tag, current_value_accessor): |
| super(BagRuntimeState, self).__init__( |
| state_spec, state_tag, current_value_accessor) |
| self._cached_value = UNREAD_VALUE |
| self._cleared = False |
| self._new_values = [] |
| |
| def read(self): |
| if self._cached_value is UNREAD_VALUE: |
| self._cached_value = self._current_value_accessor() |
| if not self._cleared: |
| encoded_values = itertools.chain(self._cached_value, self._new_values) |
| else: |
| encoded_values = self._new_values |
| return (self._decode(v) for v in encoded_values) |
| |
| def add(self, value): |
| self._new_values.append(self._encode(value)) |
| |
| def clear(self): |
| self._cleared = True |
| self._cached_value = [] |
| self._new_values = [] |
| |
| |
| class SetRuntimeState(DirectRuntimeState, userstate.SetRuntimeState): |
| def __init__(self, state_spec, state_tag, current_value_accessor): |
| super(SetRuntimeState, self).__init__( |
| state_spec, state_tag, current_value_accessor) |
| self._current_accumulator = UNREAD_VALUE |
| self._modified = False |
| |
| def _read_initial_value(self): |
| if self._current_accumulator is UNREAD_VALUE: |
| self._current_accumulator = { |
| self._decode(a) for a in self._current_value_accessor() |
| } |
| |
| def read(self): |
| self._read_initial_value() |
| return self._current_accumulator |
| |
| def add(self, value): |
| self._read_initial_value() |
| self._modified = True |
| self._current_accumulator.add(value) |
| |
| def clear(self): |
| self._current_accumulator = set() |
| self._modified = True |
| |
| def is_modified(self): |
| return self._modified |
| |
| |
| class CombiningValueRuntimeState( |
| DirectRuntimeState, userstate.CombiningValueRuntimeState): |
| """Combining value state interface object passed to user code.""" |
| |
| def __init__(self, state_spec, state_tag, current_value_accessor): |
| super(CombiningValueRuntimeState, self).__init__( |
| state_spec, state_tag, current_value_accessor) |
| self._current_accumulator = UNREAD_VALUE |
| self._modified = False |
| self._combine_fn = state_spec.combine_fn |
| |
| def _read_initial_value(self): |
| if self._current_accumulator is UNREAD_VALUE: |
| existing_accumulators = list( |
| self._decode(a) for a in self._current_value_accessor()) |
| if existing_accumulators: |
| self._current_accumulator = self._combine_fn.merge_accumulators( |
| existing_accumulators) |
| else: |
| self._current_accumulator = self._combine_fn.create_accumulator() |
| |
| def read(self): |
| self._read_initial_value() |
| return self._combine_fn.extract_output(self._current_accumulator) |
| |
| def add(self, value): |
| self._read_initial_value() |
| self._modified = True |
| self._current_accumulator = self._combine_fn.add_input( |
| self._current_accumulator, value) |
| |
| def clear(self): |
| self._modified = True |
| self._current_accumulator = self._combine_fn.create_accumulator() |
| |
| |
| class DirectUserStateContext(userstate.UserStateContext): |
| """userstate.UserStateContext for the BundleBasedDirectRunner. |
| |
| The DirectUserStateContext buffers up updates that are to be committed |
| by the TransformEvaluator after running a DoFn. |
| """ |
| |
| def __init__(self, step_context, dofn, key_coder): |
| self.step_context = step_context |
| self.dofn = dofn |
| self.key_coder = key_coder |
| |
| self.all_state_specs, self.all_timer_specs = ( |
| userstate.get_dofn_specs(dofn)) |
| self.state_tags = {} |
| for state_spec in self.all_state_specs: |
| state_key = 'user/%s' % state_spec.name |
| if isinstance(state_spec, userstate.BagStateSpec): |
| state_tag = _ListStateTag(state_key) |
| elif isinstance(state_spec, userstate.CombiningValueStateSpec): |
| state_tag = _ListStateTag(state_key) |
| elif isinstance(state_spec, userstate.SetStateSpec): |
| state_tag = _SetStateTag(state_key) |
| else: |
| raise ValueError('Invalid state spec: %s' % state_spec) |
| self.state_tags[state_spec] = state_tag |
| |
| self.cached_states = {} |
| self.cached_timers = {} |
| |
| def get_timer(self, timer_spec, key, window): |
| assert timer_spec in self.all_timer_specs |
| encoded_key = self.key_coder.encode(key) |
| cache_key = (encoded_key, window, timer_spec) |
| if cache_key not in self.cached_timers: |
| self.cached_timers[cache_key] = userstate.RuntimeTimer(timer_spec) |
| return self.cached_timers[cache_key] |
| |
| def get_state(self, state_spec, key, window): |
| assert state_spec in self.all_state_specs |
| encoded_key = self.key_coder.encode(key) |
| cache_key = (encoded_key, window, state_spec) |
| if cache_key not in self.cached_states: |
| state_tag = self.state_tags[state_spec] |
| value_accessor = ( |
| lambda: self._get_underlying_state(state_spec, key, window)) |
| self.cached_states[cache_key] = DirectRuntimeState.for_spec( |
| state_spec, state_tag, value_accessor) |
| return self.cached_states[cache_key] |
| |
| def _get_underlying_state(self, state_spec, key, window): |
| state_tag = self.state_tags[state_spec] |
| encoded_key = self.key_coder.encode(key) |
| return (self.step_context.get_keyed_state(encoded_key) |
| .get_state(window, state_tag)) |
| |
| def commit(self): |
| # Commit state modifications. |
| for cache_key, runtime_state in self.cached_states.items(): |
| encoded_key, window, state_spec = cache_key |
| state = self.step_context.get_keyed_state(encoded_key) |
| state_tag = self.state_tags[state_spec] |
| if isinstance(state_spec, userstate.BagStateSpec): |
| if runtime_state._cleared: |
| state.clear_state(window, state_tag) |
| for new_value in runtime_state._new_values: |
| state.add_state(window, state_tag, new_value) |
| elif isinstance(state_spec, userstate.CombiningValueStateSpec): |
| if runtime_state._modified: |
| state.clear_state(window, state_tag) |
| state.add_state( |
| window, state_tag, |
| state_spec.coder.encode(runtime_state._current_accumulator)) |
| elif isinstance(state_spec, userstate.SetStateSpec): |
| if runtime_state.is_modified(): |
| state.clear_state(window, state_tag) |
| for new_value in runtime_state._current_accumulator: |
| state.add_state( |
| window, state_tag, state_spec.coder.encode(new_value)) |
| else: |
| raise ValueError('Invalid state spec: %s' % state_spec) |
| |
| # Commit new timers. |
| for cache_key, runtime_timer in self.cached_timers.items(): |
| encoded_key, window, timer_spec = cache_key |
| state = self.step_context.get_keyed_state(encoded_key) |
| timer_name = 'user/%s' % timer_spec.name |
| if runtime_timer._cleared: |
| state.clear_timer(window, timer_name, timer_spec.time_domain) |
| if runtime_timer._new_timestamp is not None: |
| # TODO(ccy): add corresponding watermark holds after the DirectRunner |
| # allows for keyed watermark holds. |
| state.set_timer(window, timer_name, timer_spec.time_domain, |
| runtime_timer._new_timestamp) |