| ################################################################################ |
| # 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 abc import abstractmethod, ABC |
| from typing import Generic, List, Iterable, Dict, Set |
| |
| from pyflink.common import Row |
| from pyflink.common.constants import MAX_LONG_VALUE |
| from pyflink.datastream.state import MapState |
| from pyflink.fn_execution.table.window_assigner import WindowAssigner, PanedWindowAssigner, \ |
| MergingWindowAssigner |
| from pyflink.fn_execution.table.window_context import Context, K, W |
| |
| |
| def join_row(left: List, right: List): |
| return Row(*(left + right)) |
| |
| |
| class InternalWindowProcessFunction(Generic[K, W], ABC): |
| """ |
| The internal interface for functions that process over grouped windows. |
| """ |
| |
| def __init__(self, |
| allowed_lateness: int, |
| window_assigner: WindowAssigner[W], |
| window_aggregator): |
| self._allowed_lateness = allowed_lateness |
| self._window_assigner = window_assigner |
| self._window_aggregator = window_aggregator |
| self._ctx = None # type: Context[K, W] |
| |
| def open(self, ctx: Context[K, W]): |
| self._ctx = ctx |
| self._window_assigner.open(ctx) |
| |
| def close(self): |
| pass |
| |
| def is_cleanup_time(self, window: W, time: int) -> bool: |
| return time == self._cleanup_time(window) |
| |
| def is_window_late(self, window: W) -> bool: |
| return self._window_assigner.is_event_time() and \ |
| self._cleanup_time(window) <= self._ctx.current_watermark() |
| |
| def _cleanup_time(self, window: W) -> int: |
| if self._window_assigner.is_event_time(): |
| cleanup_time = window.max_timestamp() + self._allowed_lateness |
| if cleanup_time >= window.max_timestamp(): |
| return cleanup_time |
| else: |
| return MAX_LONG_VALUE |
| else: |
| return window.max_timestamp() |
| |
| @abstractmethod |
| def assign_state_namespace(self, input_row: List, timestamp: int) -> List[W]: |
| """ |
| Assigns the input element into the state namespace which the input element should be |
| accumulated/retracted into. |
| |
| :param input_row: The input element |
| :param timestamp: The timestamp of the element or the processing time (depends on the type |
| of assigner) |
| :return: The state namespace. |
| """ |
| pass |
| |
| @abstractmethod |
| def assign_actual_windows(self, input_row: List, timestamp: int) -> List[W]: |
| """ |
| Assigns the input element into the actual windows which the {@link Trigger} should trigger |
| on. |
| |
| :param input_row: The input element |
| :param timestamp: The timestamp of the element or the processing time (depends on the type |
| of assigner) |
| :return: The actual windows |
| """ |
| pass |
| |
| @abstractmethod |
| def prepare_aggregate_accumulator_for_emit(self, window: W): |
| """ |
| Prepares the accumulator of the given window before emit the final result. The accumulator |
| is stored in the state or will be created if there is no corresponding accumulator in state. |
| |
| :param window: The window |
| """ |
| pass |
| |
| @abstractmethod |
| def clean_window_if_needed(self, window: W, current_time: int): |
| """ |
| Cleans the given window if needed. |
| |
| :param window: The window to cleanup |
| :param current_time: The current timestamp |
| """ |
| pass |
| |
| |
| class GeneralWindowProcessFunction(InternalWindowProcessFunction[K, W]): |
| """ |
| The general implementation of InternalWindowProcessFunction. The WindowAssigner should be a |
| regular assigner without implement PanedWindowAssigner or MergingWindowAssigner. |
| """ |
| |
| def __init__(self, |
| allowed_lateness: int, |
| window_assigner: WindowAssigner[W], |
| window_aggregator): |
| super(GeneralWindowProcessFunction, self).__init__( |
| allowed_lateness, window_assigner, window_aggregator) |
| self._reuse_affected_windows = None # type: List[W] |
| |
| def assign_state_namespace(self, input_row: List, timestamp: int) -> List[W]: |
| element_windows = self._window_assigner.assign_windows(input_row, timestamp) |
| self._reuse_affected_windows = [] |
| for window in element_windows: |
| if not self.is_window_late(window): |
| self._reuse_affected_windows.append(window) |
| return self._reuse_affected_windows |
| |
| def assign_actual_windows(self, input_row: List, timestamp: int) -> List[W]: |
| # actual windows is equal to affected window, reuse it |
| return self._reuse_affected_windows |
| |
| def prepare_aggregate_accumulator_for_emit(self, window: W): |
| acc = self._ctx.get_window_accumulators(window) |
| if acc is None: |
| acc = self._window_aggregator.create_accumulators() |
| self._window_aggregator.set_accumulators(window, acc) |
| |
| def clean_window_if_needed(self, window: W, current_time: int): |
| if self.is_cleanup_time(window, current_time): |
| self._ctx.clear_window_state(window) |
| self._window_aggregator.cleanup(window) |
| self._ctx.clear_trigger(window) |
| |
| |
| class PanedWindowProcessFunction(InternalWindowProcessFunction[K, W]): |
| """ |
| The implementation of InternalWindowProcessFunction for PanedWindowAssigner. |
| """ |
| |
| def __init__(self, |
| allowed_lateness: int, |
| window_assigner: PanedWindowAssigner[W], |
| window_aggregator): |
| super(PanedWindowProcessFunction, self).__init__( |
| allowed_lateness, window_assigner, window_aggregator) |
| self._window_assigner = window_assigner |
| |
| def assign_state_namespace(self, input_row: List, timestamp: int) -> List[W]: |
| pane = self._window_assigner.assign_pane(input_row, timestamp) |
| if not self._is_pane_late(pane): |
| return [pane] |
| else: |
| return [] |
| |
| def assign_actual_windows(self, input_row: List, timestamp: int) -> List[W]: |
| element_windows = self._window_assigner.assign_windows(input_row, timestamp) |
| actual_windows = [] |
| for window in element_windows: |
| if not self.is_window_late(window): |
| actual_windows.append(window) |
| return actual_windows |
| |
| def prepare_aggregate_accumulator_for_emit(self, window: W): |
| panes = self._window_assigner.split_into_panes(window) |
| acc = self._window_aggregator.create_accumulators() |
| # null namespace means use heap data views |
| self._window_aggregator.set_accumulators(None, acc) |
| for pane in panes: |
| pane_acc = self._ctx.get_window_accumulators(pane) |
| if pane_acc: |
| self._window_aggregator.merge(pane, pane_acc) |
| |
| def clean_window_if_needed(self, window: W, current_time: int): |
| if self.is_cleanup_time(window, current_time): |
| panes = self._window_assigner.split_into_panes(window) |
| for pane in panes: |
| last_window = self._window_assigner.get_last_window(pane) |
| if window == last_window: |
| self._ctx.clear_window_state(pane) |
| self._ctx.clear_trigger(window) |
| |
| def _is_pane_late(self, pane: W): |
| # whether the pane is late depends on the last window which the pane is belongs to is late |
| return self._window_assigner.is_event_time() and \ |
| self.is_window_late(self._window_assigner.get_last_window(pane)) |
| |
| |
| class MergeResultCollector(MergingWindowAssigner.MergeCallback): |
| |
| def __init__(self): |
| self.merge_results = {} # type: Dict[W, Iterable[W]] |
| |
| def merge(self, merge_result: W, to_be_merged: Iterable[W]): |
| self.merge_results[merge_result] = to_be_merged |
| |
| |
| class MergingWindowProcessFunction(InternalWindowProcessFunction[K, W]): |
| """ |
| The implementation of InternalWindowProcessFunction for MergingWindowAssigner. |
| """ |
| |
| def __init__(self, |
| allowed_lateness: int, |
| window_assigner: MergingWindowAssigner[W], |
| window_aggregator, |
| state_backend): |
| super(MergingWindowProcessFunction, self).__init__( |
| allowed_lateness, window_assigner, window_aggregator) |
| self._window_assigner = window_assigner |
| self._reuse_actual_windows = None # type: List |
| self._window_mapping = None # type: MapState |
| self._state_backend = state_backend |
| self._sorted_windows = None # type: List |
| |
| from pyflink.fn_execution.state_impl import LRUCache |
| |
| self._cached_sorted_windows = LRUCache(10000, None) |
| |
| def open(self, ctx: Context[K, W]): |
| super(MergingWindowProcessFunction, self).open(ctx) |
| self._window_mapping = self._state_backend.get_map_state( |
| 'session-window-mapping', |
| self._state_backend.namespace_coder, |
| self._state_backend.namespace_coder) |
| |
| def assign_state_namespace(self, input_row: List, timestamp: int) -> List[W]: |
| element_windows = self._window_assigner.assign_windows(input_row, timestamp) |
| self._initialize_cache(self._ctx.current_key()) |
| self._reuse_actual_windows = [] |
| for window in element_windows: |
| # adding the new window might result in a merge, in that case the actualWindow |
| # is the merged window and we work with that. If we don't merge then |
| # actualWindow == window |
| actual_window = self._add_window(window) |
| |
| # drop if the window is already late |
| if self.is_window_late(actual_window): |
| self._window_mapping.remove(actual_window) |
| self._sorted_windows.remove(actual_window) |
| else: |
| self._reuse_actual_windows.append(actual_window) |
| |
| affected_windows = [self._window_mapping.get(actual) |
| for actual in self._reuse_actual_windows] |
| return affected_windows |
| |
| def assign_actual_windows(self, input_row: List, timestamp: int) -> List[W]: |
| # the actual windows is calculated in assignStateNamespace |
| return self._reuse_actual_windows |
| |
| def prepare_aggregate_accumulator_for_emit(self, window: W): |
| state_window = self._window_mapping.get(window) |
| acc = self._ctx.get_window_accumulators(state_window) |
| if acc is None: |
| acc = self._window_aggregator.create_accumulators() |
| |
| self._window_aggregator.set_accumulators(state_window, acc) |
| |
| def clean_window_if_needed(self, window: W, current_time: int): |
| if self.is_cleanup_time(window, current_time): |
| self._ctx.clear_trigger(window) |
| state_window = self._window_mapping.get(window) |
| self._ctx.clear_window_state(state_window) |
| # retire expired window |
| self._initialize_cache(self._ctx.current_key()) |
| self._window_mapping.remove(window) |
| self._sorted_windows.remove(window) |
| |
| def _initialize_cache(self, key): |
| tuple_key = tuple(key) |
| self._sorted_windows = self._cached_sorted_windows.get(tuple_key) |
| if self._sorted_windows is None: |
| self._sorted_windows = [k for k in self._window_mapping] |
| self._sorted_windows.sort() |
| self._cached_sorted_windows.put(tuple_key, self._sorted_windows) |
| |
| def _add_window(self, new_window: W): |
| collector = MergeResultCollector() |
| self._window_assigner.merge_windows(new_window, self._sorted_windows, collector) |
| |
| result_window = new_window |
| is_new_window_merged = False |
| |
| # perform the merge |
| merge_results = collector.merge_results |
| for merge_result in merge_results: |
| merge_windows = merge_results[merge_result] # type: Set[W] |
| |
| # if our new window is in the merged windows make the merge result the result window |
| try: |
| merge_windows.remove(new_window) |
| is_new_window_merged = True |
| result_window = merge_result |
| except KeyError: |
| pass |
| |
| # if our new window is the same as a pre-existing window, nothing to do |
| if not merge_windows: |
| continue |
| |
| # pick any of the merged windows and choose that window's state window |
| # as the state window for the merge result |
| merged_state_namespace = self._window_mapping.get(iter(merge_windows).__next__()) |
| merged_state_windows = [] |
| |
| # figure out the state windows that we are merging |
| for merged_window in merge_windows: |
| res = self._window_mapping.get(merged_window) |
| if res is not None: |
| self._window_mapping.remove(merged_window) |
| self._sorted_windows.remove(merged_window) |
| # don't put the target state window into the merged windows |
| if res != merged_state_namespace: |
| merged_state_windows.append(res) |
| |
| self._window_mapping.put(merge_result, merged_state_namespace) |
| self._sorted_windows.append(merge_result) |
| self._sorted_windows.sort() |
| |
| # don't merge the new window itself, it never had any state associated with it |
| # i.e. if we are only merging one pre-existing window into itself |
| # without extending the pre-existing window |
| if not (len(merge_windows) == 1 and merge_result in merge_windows): |
| self._merge( |
| merge_result, merge_windows, merged_state_namespace, merged_state_windows) |
| |
| # the new window created a new, self-contained window without merging |
| if len(merge_results) == 0 or result_window == new_window and not is_new_window_merged: |
| self._window_mapping.put(result_window, result_window) |
| self._sorted_windows.append(result_window) |
| self._sorted_windows.sort() |
| |
| return result_window |
| |
| def _merge(self, merge_result: W, merge_windows: Set[W], state_window_result: W, |
| state_windows_tobe_merged: Iterable[W]): |
| self._ctx.on_merge(merge_result, state_windows_tobe_merged) |
| |
| # clear registered timers |
| for window in merge_windows: |
| self._ctx.clear_trigger(window) |
| self._ctx.delete_cleanup_timer(window) |
| |
| # merge the merged state windows into the newly resulting state window |
| if state_windows_tobe_merged: |
| target_acc = self._ctx.get_window_accumulators(state_window_result) |
| if target_acc is None: |
| target_acc = self._window_aggregator.create_accumulators() |
| self._window_aggregator.set_accumulators(state_window_result, target_acc) |
| for window in state_windows_tobe_merged: |
| acc = self._ctx.get_window_accumulators(window) |
| if acc is not None: |
| self._window_aggregator.merge(window, acc) |
| # clear merged window |
| self._ctx.clear_window_state(window) |
| target_acc = self._window_aggregator.get_accumulators() |
| self._ctx.set_window_accumulators(state_window_result, target_acc) |
