dianfu commented on a change in pull request #13462: URL: https://github.com/apache/flink/pull/13462#discussion_r493998338
########## File path: flink-python/pyflink/common/state.py ########## @@ -0,0 +1,64 @@ +################################################################################ +# 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 ABC, abstractmethod + +from typing import TypeVar, Generic, Iterable, List, Iterator + +T = TypeVar('T') +K = TypeVar('K') +V = TypeVar('V') + + Review comment: Please add the Python docs for these classes and methods. ########## File path: flink-python/pyflink/fn_execution/beam/beam_coder_impl_slow.py ########## @@ -46,8 +46,7 @@ def __init__(self, field_coders): self._remaining_bits_num = (self._field_count + ROW_KIND_BIT_SIZE) % 8 self.null_mask_search_table = self.generate_null_mask_search_table() self.null_byte_search_table = (0x80, 0x40, 0x20, 0x10, 0x08, 0x04, 0x02, 0x01) - self.row_kind_search_table = \ - [i << (8 - ROW_KIND_BIT_SIZE) for i in range(2 ** ROW_KIND_BIT_SIZE)] + self.row_kind_search_table = [0x00, 0x80, 0x40, 0xC0] Review comment: what's the purpose of this change? ########## File path: flink-python/pyflink/fn_execution/beam/beam_coders.py ########## @@ -97,6 +97,80 @@ def __hash__(self): return hash(self._table_function_row_coder) +class BeamAggregateFunctionInputRowCoder(FastCoder): + """ + Coder for Table Function Row. Review comment: Correct the doc ########## File path: flink-python/pyflink/fn_execution/beam/beam_operations_fast.pyx ########## @@ -233,3 +339,28 @@ def _create_user_defined_function_operation(factory, transform_proto, consumers, factory.counter_factory, factory.state_sampler, consumers) + +def _create_stateful_user_defined_function_operation(factory, transform_proto, consumers, Review comment: Could we reuse _create_user_defined_function_operation without adding a new method? ########## File path: flink-python/pyflink/fn_execution/operation_utils.py ########## @@ -170,3 +181,485 @@ def _next_constant_num(): constant_value_name = 'c%s' % _next_constant_num() return constant_value_name, parsed_constant_value + + +class Count1AggFunction(AggregateFunction): + + def get_value(self, accumulator): + return accumulator[0] + + def create_accumulator(self): + return [0] + + def accumulate(self, accumulator, *args): + accumulator[0] += 1 + + def retract(self, accumulator, *args): + accumulator[0] -= 1 + + def merge(self, accumulator, accumulators): + for acc in accumulators: + accumulator[0] += acc[0] + + +def join_row(left: Row, right: Row): + fields = [] + for value in left: + fields.append(value) + for value in right: + fields.append(value) + return Row(*fields) + + +class RowKeySelector(object): + + def __init__(self, grouping): + self.grouping = grouping + + def get_key(self, data: Row): + return Row(*[data[i] for i in self.grouping]) + + +class StateDataViewStore(object): + + def __init__(self, function_context): + self.function_context = function_context + + def get_runtime_context(self): + return self.function_context + + +class AggsHandleFunction(ABC): + """ + The base class for handling aggregate functions. + """ + + @abstractmethod + def open(self, state_data_view_store): + """ + Initialization method for the function. It is called before the actual working methods. + + :param state_data_view_store: The object used to manage the DataView. + """ + pass + + @abstractmethod + def accumulate(self, input_data: Row): + """ + Accumulates the input values to the accumulators. + + :param input_data: Input values bundled in a row. + """ + pass + + @abstractmethod + def retract(self, input_data: Row): + """ + Retracts the input values from the accumulators. + + :param input_data: Input values bundled in a row. + """ + + @abstractmethod + def merge(self, accumulators: Row): + """ + Merges the other accumulators into current accumulators. + + :param accumulators: The other row of accumulators. + """ + pass + + @abstractmethod + def set_accumulators(self, accumulators: Row): + """ + Set the current accumulators (saved in a row) which contains the current aggregated results. + + In streaming: accumulators are store in the state, we need to restore aggregate buffers from + state. + + In batch: accumulators are store in the dict, we need to restore aggregate buffers from Review comment: ```suggestion In batch: accumulators are stored in the dict, we need to restore aggregate buffers from ``` ########## File path: flink-python/src/main/java/org/apache/flink/table/runtime/arrow/ArrowUtils.java ########## @@ -749,7 +749,7 @@ private static boolean isAppendOnlyTable(Table table) throws Exception { OutputConversionModifyOperation.UpdateMode.APPEND); tableEnv.getPlanner().translate(Collections.singletonList(modifyOperation)); } catch (Throwable t) { - if (t.getMessage().contains("doesn't support consuming update changes") || Review comment: Unnecessary change? ########## File path: flink-python/pyflink/fn_execution/beam/beam_coder_impl_slow.py ########## @@ -187,6 +189,35 @@ def __repr__(self): return 'TableFunctionRowCoderImpl[%s]' % repr(self._flatten_row_coder) +class AggregateFunctionRowCoderImpl(StreamCoderImpl): + """ + The aggregate function row coder impl is similar to the table function row coder + (one line in, multiple line out). The only differences are that this row coder will + encode row kind information into the output row and would not encode the + end message. + """ + + def __init__(self, flatten_row_coder): + self._flatten_row_coder = flatten_row_coder + self.data_out_stream = create_OutputStream() Review comment: ```suggestion self._data_out_stream = create_OutputStream() ``` ########## File path: flink-python/pyflink/fn_execution/beam/beam_operations_slow.py ########## @@ -172,6 +183,94 @@ def generate_func(self, udfs): return lambda it: map(mapper, it), user_defined_funcs +class StatefulFunctionOperation(StatelessFunctionOperation): + + def __init__(self, name, spec, counter_factory, sampler, consumers, keyed_state_backend): + self.keyed_state_backend = keyed_state_backend + super(StatefulFunctionOperation, self).__init__( + name, spec, counter_factory, sampler, consumers) + + def finish(self): + super().finish() + with self.scoped_finish_state: + if self.keyed_state_backend: + self.keyed_state_backend.commit() + + def reset(self): + super().reset() + if self.keyed_state_backend: + self.keyed_state_backend.reset() + + +TRIGGER_TIMER = 1 + + +class StreamGroupAggregateOperation(StatefulFunctionOperation): + + def __init__(self, name, spec, counter_factory, sampler, consumers, keyed_state_backend): + self.generate_update_before = spec.serialized_fn.generate_update_before + self.grouping = [i for i in spec.serialized_fn.grouping] + self.group_agg_function = None + self.index_of_count_star = spec.serialized_fn.index_of_count_star + self.state_cache_size = spec.serialized_fn.state_cache_size + self.state_cleaning_enabled = spec.serialized_fn.state_cleaning_enabled + super(StreamGroupAggregateOperation, self).__init__( + name, spec, counter_factory, sampler, consumers, keyed_state_backend) + + def open_func(self): + self.group_agg_function.open(FunctionContext(self.base_metric_group)) + + def generate_func(self, udfs): + user_defined_aggs = [] + input_offsets = [] + for i in range(len(udfs)): + if i != self.index_of_count_star: + user_defined_agg, input_offset = self._extract_user_defined_agg_info(udfs[i]) + else: + user_defined_agg = Count1AggFunction() + input_offset = [] + user_defined_aggs.append(user_defined_agg) + input_offsets.append(input_offset) + aggs_handler_function = SimpleAggsHandleFunction( + user_defined_aggs, input_offsets, self.index_of_count_star) + key_selector = RowKeySelector(self.grouping) + self.group_agg_function = GroupAggFunction( + aggs_handler_function, + key_selector, + self.keyed_state_backend, + self.generate_update_before, + self.state_cleaning_enabled, + self.index_of_count_star) + return lambda it: map(self.call_timer_or_process_element, it), [] + + def call_timer_or_process_element(self, input_data: Tuple[int, Row, int, Row]): Review comment: What about named it as process_element_or_timer? ########## File path: flink-python/pyflink/fn_execution/beam/beam_operations_slow.py ########## @@ -172,6 +183,94 @@ def generate_func(self, udfs): return lambda it: map(mapper, it), user_defined_funcs +class StatefulFunctionOperation(StatelessFunctionOperation): + + def __init__(self, name, spec, counter_factory, sampler, consumers, keyed_state_backend): + self.keyed_state_backend = keyed_state_backend + super(StatefulFunctionOperation, self).__init__( + name, spec, counter_factory, sampler, consumers) + + def finish(self): + super().finish() + with self.scoped_finish_state: + if self.keyed_state_backend: + self.keyed_state_backend.commit() + + def reset(self): + super().reset() + if self.keyed_state_backend: + self.keyed_state_backend.reset() + + +TRIGGER_TIMER = 1 + + +class StreamGroupAggregateOperation(StatefulFunctionOperation): + + def __init__(self, name, spec, counter_factory, sampler, consumers, keyed_state_backend): + self.generate_update_before = spec.serialized_fn.generate_update_before + self.grouping = [i for i in spec.serialized_fn.grouping] + self.group_agg_function = None + self.index_of_count_star = spec.serialized_fn.index_of_count_star + self.state_cache_size = spec.serialized_fn.state_cache_size + self.state_cleaning_enabled = spec.serialized_fn.state_cleaning_enabled + super(StreamGroupAggregateOperation, self).__init__( + name, spec, counter_factory, sampler, consumers, keyed_state_backend) + + def open_func(self): + self.group_agg_function.open(FunctionContext(self.base_metric_group)) + + def generate_func(self, udfs): + user_defined_aggs = [] + input_offsets = [] + for i in range(len(udfs)): + if i != self.index_of_count_star: + user_defined_agg, input_offset = self._extract_user_defined_agg_info(udfs[i]) + else: + user_defined_agg = Count1AggFunction() + input_offset = [] + user_defined_aggs.append(user_defined_agg) + input_offsets.append(input_offset) + aggs_handler_function = SimpleAggsHandleFunction( + user_defined_aggs, input_offsets, self.index_of_count_star) + key_selector = RowKeySelector(self.grouping) + self.group_agg_function = GroupAggFunction( + aggs_handler_function, + key_selector, + self.keyed_state_backend, + self.generate_update_before, + self.state_cleaning_enabled, + self.index_of_count_star) + return lambda it: map(self.call_timer_or_process_element, it), [] + + def call_timer_or_process_element(self, input_data: Tuple[int, Row, int, Row]): + # the structure of the input data: + # [element_type, element(for process_element), timestamp(for timer), key(for timer)] + # all the fields are nullable except the "element_type" + if input_data[0] != TRIGGER_TIMER: + return self.group_agg_function.process_element(input_data[1]) + else: + self.group_agg_function.on_timer(input_data[3]) + return [] + + def teardown(self): + if self.group_agg_function is not None: + self.group_agg_function.close() + super().teardown() + + def _extract_user_defined_agg_info(self, user_defined_function_proto): Review comment: ```suggestion def _extract_user_defined_aggregate_function(self, user_defined_function_proto): ``` ########## File path: flink-python/pyflink/fn_execution/beam/beam_coders.py ########## @@ -97,6 +97,80 @@ def __hash__(self): return hash(self._table_function_row_coder) +class BeamAggregateFunctionInputRowCoder(FastCoder): + """ + Coder for Table Function Row. + """ + + def __init__(self, aggregate_function_row_coder): + self._aggregate_function_row_coder = aggregate_function_row_coder + + def _create_impl(self): + return self._aggregate_function_row_coder.get_impl() + + def get_impl(self): + return BeamCoderImpl(self._create_impl()) + + def to_type_hint(self): + return typehints.List + + @Coder.register_urn(coders.FLINK_AGGREGATE_FUNCTION_INPUT_SCHEMA_CODER_URN, + flink_fn_execution_pb2.Schema) + def _pickle_from_runner_api_parameter(schema_proto, unused_components, unused_context): + return BeamAggregateFunctionInputRowCoder( + coders.AggregateFunctionInputRowCoder.from_schema_proto(schema_proto)) + + def __repr__(self): + return 'BeamAggregateFunctionInputRowCoder[%s]' % repr(self._aggregate_function_row_coder) + + def __eq__(self, other): + return (self.__class__ == other.__class__ + and self._aggregate_function_row_coder == other._table_function_row_coder) + + def __ne__(self, other): + return not self == other + + def __hash__(self): + return hash(self._aggregate_function_row_coder) + + +class BeamAggregateFunctionOutputRowCoder(FastCoder): + """ + Coder for Table Function Row. Review comment: ditto ########## File path: flink-python/pyflink/fn_execution/beam/beam_operations_slow.py ########## @@ -172,6 +183,94 @@ def generate_func(self, udfs): return lambda it: map(mapper, it), user_defined_funcs +class StatefulFunctionOperation(StatelessFunctionOperation): + + def __init__(self, name, spec, counter_factory, sampler, consumers, keyed_state_backend): + self.keyed_state_backend = keyed_state_backend + super(StatefulFunctionOperation, self).__init__( + name, spec, counter_factory, sampler, consumers) + + def finish(self): + super().finish() + with self.scoped_finish_state: + if self.keyed_state_backend: + self.keyed_state_backend.commit() + + def reset(self): + super().reset() + if self.keyed_state_backend: + self.keyed_state_backend.reset() + + +TRIGGER_TIMER = 1 + + +class StreamGroupAggregateOperation(StatefulFunctionOperation): + + def __init__(self, name, spec, counter_factory, sampler, consumers, keyed_state_backend): + self.generate_update_before = spec.serialized_fn.generate_update_before + self.grouping = [i for i in spec.serialized_fn.grouping] + self.group_agg_function = None + self.index_of_count_star = spec.serialized_fn.index_of_count_star Review comment: Why we need to handle count(*) specially? I would suggest adding some description about this. ########## File path: flink-python/pyflink/fn_execution/beam/beam_coder_impl_slow.py ########## @@ -187,6 +189,35 @@ def __repr__(self): return 'TableFunctionRowCoderImpl[%s]' % repr(self._flatten_row_coder) +class AggregateFunctionRowCoderImpl(StreamCoderImpl): + """ + The aggregate function row coder impl is similar to the table function row coder + (one line in, multiple line out). The only differences are that this row coder will Review comment: Could you add some description about why there are multiple outputs for one input for AggregateFunction? ########## File path: flink-python/pyflink/fn_execution/operation_utils.py ########## @@ -170,3 +181,485 @@ def _next_constant_num(): constant_value_name = 'c%s' % _next_constant_num() return constant_value_name, parsed_constant_value + + +class Count1AggFunction(AggregateFunction): + + def get_value(self, accumulator): + return accumulator[0] + + def create_accumulator(self): + return [0] + + def accumulate(self, accumulator, *args): + accumulator[0] += 1 + + def retract(self, accumulator, *args): + accumulator[0] -= 1 + + def merge(self, accumulator, accumulators): + for acc in accumulators: + accumulator[0] += acc[0] + + +def join_row(left: Row, right: Row): + fields = [] + for value in left: + fields.append(value) + for value in right: + fields.append(value) + return Row(*fields) + + +class RowKeySelector(object): + + def __init__(self, grouping): + self.grouping = grouping + + def get_key(self, data: Row): + return Row(*[data[i] for i in self.grouping]) + + +class StateDataViewStore(object): + + def __init__(self, function_context): + self.function_context = function_context + + def get_runtime_context(self): + return self.function_context + + +class AggsHandleFunction(ABC): + """ + The base class for handling aggregate functions. + """ + + @abstractmethod + def open(self, state_data_view_store): + """ + Initialization method for the function. It is called before the actual working methods. + + :param state_data_view_store: The object used to manage the DataView. + """ + pass + + @abstractmethod + def accumulate(self, input_data: Row): + """ + Accumulates the input values to the accumulators. + + :param input_data: Input values bundled in a row. + """ + pass + + @abstractmethod + def retract(self, input_data: Row): + """ + Retracts the input values from the accumulators. + + :param input_data: Input values bundled in a row. + """ + + @abstractmethod + def merge(self, accumulators: Row): + """ + Merges the other accumulators into current accumulators. + + :param accumulators: The other row of accumulators. + """ + pass + + @abstractmethod + def set_accumulators(self, accumulators: Row): + """ + Set the current accumulators (saved in a row) which contains the current aggregated results. + + In streaming: accumulators are store in the state, we need to restore aggregate buffers from Review comment: ```suggestion In streaming: accumulators are store in the state, we need to restore aggregate buffers from ``` ```suggestion In streaming: accumulators are stored in the state, we need to restore aggregate buffers from ``` ########## File path: flink-python/pyflink/fn_execution/beam/beam_operations_slow.py ########## @@ -172,6 +183,94 @@ def generate_func(self, udfs): return lambda it: map(mapper, it), user_defined_funcs +class StatefulFunctionOperation(StatelessFunctionOperation): + + def __init__(self, name, spec, counter_factory, sampler, consumers, keyed_state_backend): + self.keyed_state_backend = keyed_state_backend + super(StatefulFunctionOperation, self).__init__( + name, spec, counter_factory, sampler, consumers) + + def finish(self): + super().finish() + with self.scoped_finish_state: + if self.keyed_state_backend: + self.keyed_state_backend.commit() + + def reset(self): + super().reset() + if self.keyed_state_backend: + self.keyed_state_backend.reset() + + +TRIGGER_TIMER = 1 + + +class StreamGroupAggregateOperation(StatefulFunctionOperation): + + def __init__(self, name, spec, counter_factory, sampler, consumers, keyed_state_backend): + self.generate_update_before = spec.serialized_fn.generate_update_before + self.grouping = [i for i in spec.serialized_fn.grouping] + self.group_agg_function = None + self.index_of_count_star = spec.serialized_fn.index_of_count_star + self.state_cache_size = spec.serialized_fn.state_cache_size + self.state_cleaning_enabled = spec.serialized_fn.state_cleaning_enabled + super(StreamGroupAggregateOperation, self).__init__( + name, spec, counter_factory, sampler, consumers, keyed_state_backend) + + def open_func(self): + self.group_agg_function.open(FunctionContext(self.base_metric_group)) + + def generate_func(self, udfs): + user_defined_aggs = [] + input_offsets = [] + for i in range(len(udfs)): + if i != self.index_of_count_star: + user_defined_agg, input_offset = self._extract_user_defined_agg_info(udfs[i]) + else: + user_defined_agg = Count1AggFunction() + input_offset = [] + user_defined_aggs.append(user_defined_agg) + input_offsets.append(input_offset) + aggs_handler_function = SimpleAggsHandleFunction( + user_defined_aggs, input_offsets, self.index_of_count_star) + key_selector = RowKeySelector(self.grouping) + self.group_agg_function = GroupAggFunction( + aggs_handler_function, + key_selector, + self.keyed_state_backend, + self.generate_update_before, + self.state_cleaning_enabled, + self.index_of_count_star) + return lambda it: map(self.call_timer_or_process_element, it), [] + + def call_timer_or_process_element(self, input_data: Tuple[int, Row, int, Row]): + # the structure of the input data: + # [element_type, element(for process_element), timestamp(for timer), key(for timer)] + # all the fields are nullable except the "element_type" + if input_data[0] != TRIGGER_TIMER: + return self.group_agg_function.process_element(input_data[1]) + else: + self.group_agg_function.on_timer(input_data[3]) + return [] + + def teardown(self): + if self.group_agg_function is not None: + self.group_agg_function.close() + super().teardown() + + def _extract_user_defined_agg_info(self, user_defined_function_proto): + user_defined_agg = cloudpickle.loads(user_defined_function_proto.payload) + assert isinstance(user_defined_agg, AggregateFunction) + inputs = self._extract_user_defined_function_args(user_defined_function_proto.inputs) + return user_defined_agg, inputs + + def _extract_user_defined_function_args(self, args): Review comment: Merge the implementation into _extract_user_defined_agg_info as the logic is very simple and so I think it isn't necessary to extract this into a separate method? ########## File path: flink-python/pyflink/fn_execution/operation_utils.py ########## @@ -170,3 +181,485 @@ def _next_constant_num(): constant_value_name = 'c%s' % _next_constant_num() return constant_value_name, parsed_constant_value + + +class Count1AggFunction(AggregateFunction): Review comment: It doesn't make sense to simply place these classes to a utils module. Could we find a meaningful module name for these classes? ########## File path: flink-python/pyflink/table/table_environment.py ########## @@ -1606,6 +1612,17 @@ def _before_execute(self): self._add_jars_to_j_env_config(jars_key) self._add_jars_to_j_env_config(classpaths_key) + def _wrap_aggregate_function_if_needed(self, function): + if isinstance(function, (AggregateFunction, UserDefinedAggregateFunctionWrapper)): + if not self._is_blink_planner: + raise Exception("The Python UDAF is only supported on the blink planner") Review comment: ```suggestion raise Exception("Python UDAF is only supported in blink planner") ``` ########## File path: flink-python/pyflink/fn_execution/operation_utils.py ########## @@ -170,3 +181,485 @@ def _next_constant_num(): constant_value_name = 'c%s' % _next_constant_num() return constant_value_name, parsed_constant_value + + +class Count1AggFunction(AggregateFunction): + + def get_value(self, accumulator): + return accumulator[0] + + def create_accumulator(self): + return [0] + + def accumulate(self, accumulator, *args): + accumulator[0] += 1 + + def retract(self, accumulator, *args): + accumulator[0] -= 1 + + def merge(self, accumulator, accumulators): + for acc in accumulators: + accumulator[0] += acc[0] + + +def join_row(left: Row, right: Row): + fields = [] + for value in left: + fields.append(value) + for value in right: + fields.append(value) + return Row(*fields) + + +class RowKeySelector(object): + + def __init__(self, grouping): + self.grouping = grouping + + def get_key(self, data: Row): + return Row(*[data[i] for i in self.grouping]) + + +class StateDataViewStore(object): + + def __init__(self, function_context): + self.function_context = function_context + + def get_runtime_context(self): + return self.function_context + + +class AggsHandleFunction(ABC): + """ + The base class for handling aggregate functions. + """ + + @abstractmethod + def open(self, state_data_view_store): + """ + Initialization method for the function. It is called before the actual working methods. + + :param state_data_view_store: The object used to manage the DataView. + """ + pass + + @abstractmethod + def accumulate(self, input_data: Row): + """ + Accumulates the input values to the accumulators. + + :param input_data: Input values bundled in a row. + """ + pass + + @abstractmethod + def retract(self, input_data: Row): + """ + Retracts the input values from the accumulators. + + :param input_data: Input values bundled in a row. + """ + + @abstractmethod + def merge(self, accumulators: Row): + """ + Merges the other accumulators into current accumulators. + + :param accumulators: The other row of accumulators. + """ + pass + + @abstractmethod + def set_accumulators(self, accumulators: Row): + """ + Set the current accumulators (saved in a row) which contains the current aggregated results. + + In streaming: accumulators are store in the state, we need to restore aggregate buffers from + state. + + In batch: accumulators are store in the dict, we need to restore aggregate buffers from + dict. + + :param accumulators: Current accumulators. + """ + pass + + @abstractmethod + def get_accumulators(self) -> Row: + """ + Gets the current accumulators (saved in a row) which contains the current + aggregated results. + + :return: The current accumulators. + """ + pass + + @abstractmethod + def create_accumulators(self) -> Row: + """ + Initializes the accumulators and save them to a accumulators row. Review comment: ```suggestion Initializes the accumulators and save them to an accumulators row. ``` ########## File path: flink-python/pyflink/fn_execution/operation_utils.py ########## @@ -170,3 +181,485 @@ def _next_constant_num(): constant_value_name = 'c%s' % _next_constant_num() return constant_value_name, parsed_constant_value + + +class Count1AggFunction(AggregateFunction): + + def get_value(self, accumulator): + return accumulator[0] + + def create_accumulator(self): + return [0] + + def accumulate(self, accumulator, *args): + accumulator[0] += 1 + + def retract(self, accumulator, *args): + accumulator[0] -= 1 + + def merge(self, accumulator, accumulators): + for acc in accumulators: + accumulator[0] += acc[0] + + +def join_row(left: Row, right: Row): + fields = [] + for value in left: + fields.append(value) + for value in right: + fields.append(value) + return Row(*fields) + + +class RowKeySelector(object): + + def __init__(self, grouping): + self.grouping = grouping + + def get_key(self, data: Row): + return Row(*[data[i] for i in self.grouping]) + + +class StateDataViewStore(object): + + def __init__(self, function_context): + self.function_context = function_context + + def get_runtime_context(self): + return self.function_context + + +class AggsHandleFunction(ABC): + """ + The base class for handling aggregate functions. + """ + + @abstractmethod + def open(self, state_data_view_store): + """ + Initialization method for the function. It is called before the actual working methods. + + :param state_data_view_store: The object used to manage the DataView. + """ + pass + + @abstractmethod + def accumulate(self, input_data: Row): + """ + Accumulates the input values to the accumulators. + + :param input_data: Input values bundled in a row. + """ + pass + + @abstractmethod + def retract(self, input_data: Row): + """ + Retracts the input values from the accumulators. + + :param input_data: Input values bundled in a row. + """ + + @abstractmethod + def merge(self, accumulators: Row): + """ + Merges the other accumulators into current accumulators. + + :param accumulators: The other row of accumulators. + """ + pass + + @abstractmethod + def set_accumulators(self, accumulators: Row): + """ + Set the current accumulators (saved in a row) which contains the current aggregated results. + + In streaming: accumulators are store in the state, we need to restore aggregate buffers from + state. + + In batch: accumulators are store in the dict, we need to restore aggregate buffers from + dict. + + :param accumulators: Current accumulators. + """ + pass + + @abstractmethod + def get_accumulators(self) -> Row: + """ + Gets the current accumulators (saved in a row) which contains the current + aggregated results. + + :return: The current accumulators. + """ + pass + + @abstractmethod + def create_accumulators(self) -> Row: + """ + Initializes the accumulators and save them to a accumulators row. + + :return: A row of accumulators which contains the aggregated results. + """ + pass + + @abstractmethod + def cleanup(self): + """ + Cleanup for the retired accumulators state. + """ + pass + + @abstractmethod + def get_value(self) -> Row: + """ + Gets the result of the aggregation from the current accumulators. + + :return: The final result (saved in a row) of the current accumulators. + """ + pass + + @abstractmethod + def close(self): + """ + Tear-down method for this function. It can be used for clean up work. + By default, this method does nothing. + """ Review comment: Missing `pass`? ########## File path: flink-python/pyflink/fn_execution/operation_utils.py ########## @@ -170,3 +181,485 @@ def _next_constant_num(): constant_value_name = 'c%s' % _next_constant_num() return constant_value_name, parsed_constant_value + + +class Count1AggFunction(AggregateFunction): + + def get_value(self, accumulator): + return accumulator[0] + + def create_accumulator(self): + return [0] + + def accumulate(self, accumulator, *args): + accumulator[0] += 1 + + def retract(self, accumulator, *args): + accumulator[0] -= 1 + + def merge(self, accumulator, accumulators): + for acc in accumulators: + accumulator[0] += acc[0] + + +def join_row(left: Row, right: Row): + fields = [] + for value in left: + fields.append(value) + for value in right: + fields.append(value) + return Row(*fields) + + +class RowKeySelector(object): + + def __init__(self, grouping): + self.grouping = grouping + + def get_key(self, data: Row): + return Row(*[data[i] for i in self.grouping]) + + +class StateDataViewStore(object): + + def __init__(self, function_context): + self.function_context = function_context + + def get_runtime_context(self): + return self.function_context + + +class AggsHandleFunction(ABC): + """ + The base class for handling aggregate functions. + """ + + @abstractmethod + def open(self, state_data_view_store): + """ + Initialization method for the function. It is called before the actual working methods. + + :param state_data_view_store: The object used to manage the DataView. + """ + pass + + @abstractmethod + def accumulate(self, input_data: Row): + """ + Accumulates the input values to the accumulators. + + :param input_data: Input values bundled in a row. + """ + pass + + @abstractmethod + def retract(self, input_data: Row): + """ + Retracts the input values from the accumulators. + + :param input_data: Input values bundled in a row. + """ + + @abstractmethod + def merge(self, accumulators: Row): + """ + Merges the other accumulators into current accumulators. + + :param accumulators: The other row of accumulators. + """ + pass + + @abstractmethod + def set_accumulators(self, accumulators: Row): + """ + Set the current accumulators (saved in a row) which contains the current aggregated results. + + In streaming: accumulators are store in the state, we need to restore aggregate buffers from + state. + + In batch: accumulators are store in the dict, we need to restore aggregate buffers from + dict. + + :param accumulators: Current accumulators. + """ + pass + + @abstractmethod + def get_accumulators(self) -> Row: + """ + Gets the current accumulators (saved in a row) which contains the current + aggregated results. + + :return: The current accumulators. + """ + pass + + @abstractmethod + def create_accumulators(self) -> Row: + """ + Initializes the accumulators and save them to a accumulators row. + + :return: A row of accumulators which contains the aggregated results. + """ + pass + + @abstractmethod + def cleanup(self): + """ + Cleanup for the retired accumulators state. + """ + pass + + @abstractmethod + def get_value(self) -> Row: + """ + Gets the result of the aggregation from the current accumulators. + + :return: The final result (saved in a row) of the current accumulators. + """ + pass + + @abstractmethod + def close(self): + """ + Tear-down method for this function. It can be used for clean up work. + By default, this method does nothing. + """ + + +class SimpleAggsHandleFunction(AggsHandleFunction): + + def __init__(self, + udfs: List[AggregateFunction], + args_offsets_list: List[List[int]], + index_of_count_star: int): + self._udfs = udfs + self._args_offsets_list = args_offsets_list + self._accumulators = None # type: Row + self._get_value_indexes = [i for i in range(len(udfs))] + if index_of_count_star >= 0: + self._get_value_indexes.remove(index_of_count_star) + + def open(self, state_data_view_store): + for udf in self._udfs: + udf.open(state_data_view_store.get_runtime_context()) + + def accumulate(self, input_data: Row): + for i in range(len(self._udfs)): + args_offset = self._args_offsets_list[i] + args = [input_data[offset] for offset in args_offset] + self._udfs[i].accumulate(self._accumulators[i], *args) + + def retract(self, input_data: Row): + for i in range(len(self._udfs)): + args_offset = self._args_offsets_list[i] + args = [input_data[offset] for offset in args_offset] + self._udfs[i].retract(self._accumulators[i], *args) + + def merge(self, accumulators: Row): + for i in range(len(self._udfs)): + self._udfs[i].merge(self._accumulators[i], [accumulators[i]]) + + def set_accumulators(self, accumulators: Row): + self._accumulators = accumulators + + def get_accumulators(self): + return self._accumulators + + def create_accumulators(self): + return Row(*[udf.create_accumulator() for udf in self._udfs]) + + def cleanup(self): + # Due to DataView is unsupported currently, there is nothing to do in this method. + pass + + def get_value(self): + return Row(*[self._udfs[i].get_value(self._accumulators[i]) + for i in self._get_value_indexes]) + + def close(self): + for udf in self._udfs: + udf.close() + + +class RecordCounter(ABC): + + @abstractmethod + def record_count_is_zero(self, acc): + pass + + @staticmethod + def of(index_of_count_star): + if index_of_count_star >= 0: + return RetractionRecordCounter(index_of_count_star) + else: + return AccumulationRecordCounter() + + +class AccumulationRecordCounter(RecordCounter): + + def record_count_is_zero(self, acc): + return acc is None + + +class RetractionRecordCounter(RecordCounter): + + def __init__(self, index_of_count_star): + self._index_of_count_star = index_of_count_star + + def record_count_is_zero(self, acc): + return acc is None or acc[self._index_of_count_star][0] == 0 + + +class LRUCache(object): + + def __init__(self, max_entries, default_entry): + self._max_entries = max_entries + self._default_entry = default_entry + self._cache = collections.OrderedDict() + self._on_evict = None + + def get(self, key): + value = self._cache.pop(key, self._default_entry) + if value != self._default_entry: + self._cache[key] = value + return value + + def put(self, key, value): + self._cache[key] = value + while len(self._cache) > self._max_entries: + name, value = self._cache.popitem(last=False) + if self._on_evict is not None: + self._on_evict(name, value) + + def evict(self, key): + value = self._cache.pop(key, self._default_entry) + if self._on_evict is not None: + self._on_evict(key, value) + + def evict_all(self): + if self._on_evict is not None: + for item in self._cache.items(): + self._on_evict(*item) + self._cache.clear() + + def set_on_evict(self, func): + self._on_evict = func + + def __len__(self): + return len(self._cache) + + +class SynchronousListRuntimeState(ListState): + + def __init__(self, internal_state: SynchronousBagRuntimeState): + self._internal_state = internal_state + + def add(self, v): + self._internal_state.add(v) + + def get(self): + return self._internal_state.read() + + def add_all(self, values): + self._internal_state._added_elements.extend(values) + + def update(self, values): + self.clear() + self.add_all(values) + + def clear(self): + self._internal_state.clear() + + +class SynchronousValueRuntimeState(ValueState): + + def __init__(self, internal_state: SynchronousBagRuntimeState): + self._internal_state = internal_state + + def value(self): + for i in self._internal_state.read(): + return i + return None + + def update(self, value) -> None: + self._internal_state.clear() + self._internal_state.add(value) + + def clear(self) -> None: + self._internal_state.clear() + + +class RemoteKeyedStateBackend(object): Review comment: Is there any test case covering this class? ########## File path: flink-python/pyflink/fn_execution/beam/beam_operations_slow.py ########## @@ -172,6 +183,94 @@ def generate_func(self, udfs): return lambda it: map(mapper, it), user_defined_funcs +class StatefulFunctionOperation(StatelessFunctionOperation): + + def __init__(self, name, spec, counter_factory, sampler, consumers, keyed_state_backend): + self.keyed_state_backend = keyed_state_backend + super(StatefulFunctionOperation, self).__init__( + name, spec, counter_factory, sampler, consumers) + + def finish(self): + super().finish() + with self.scoped_finish_state: + if self.keyed_state_backend: + self.keyed_state_backend.commit() + + def reset(self): + super().reset() + if self.keyed_state_backend: + self.keyed_state_backend.reset() + + +TRIGGER_TIMER = 1 + + +class StreamGroupAggregateOperation(StatefulFunctionOperation): + + def __init__(self, name, spec, counter_factory, sampler, consumers, keyed_state_backend): + self.generate_update_before = spec.serialized_fn.generate_update_before + self.grouping = [i for i in spec.serialized_fn.grouping] + self.group_agg_function = None + self.index_of_count_star = spec.serialized_fn.index_of_count_star + self.state_cache_size = spec.serialized_fn.state_cache_size + self.state_cleaning_enabled = spec.serialized_fn.state_cleaning_enabled + super(StreamGroupAggregateOperation, self).__init__( + name, spec, counter_factory, sampler, consumers, keyed_state_backend) + + def open_func(self): + self.group_agg_function.open(FunctionContext(self.base_metric_group)) + + def generate_func(self, udfs): + user_defined_aggs = [] + input_offsets = [] + for i in range(len(udfs)): + if i != self.index_of_count_star: + user_defined_agg, input_offset = self._extract_user_defined_agg_info(udfs[i]) + else: + user_defined_agg = Count1AggFunction() + input_offset = [] + user_defined_aggs.append(user_defined_agg) + input_offsets.append(input_offset) + aggs_handler_function = SimpleAggsHandleFunction( + user_defined_aggs, input_offsets, self.index_of_count_star) + key_selector = RowKeySelector(self.grouping) + self.group_agg_function = GroupAggFunction( + aggs_handler_function, + key_selector, + self.keyed_state_backend, + self.generate_update_before, + self.state_cleaning_enabled, + self.index_of_count_star) + return lambda it: map(self.call_timer_or_process_element, it), [] + + def call_timer_or_process_element(self, input_data: Tuple[int, Row, int, Row]): + # the structure of the input data: + # [element_type, element(for process_element), timestamp(for timer), key(for timer)] + # all the fields are nullable except the "element_type" + if input_data[0] != TRIGGER_TIMER: + return self.group_agg_function.process_element(input_data[1]) + else: + self.group_agg_function.on_timer(input_data[3]) + return [] + + def teardown(self): + if self.group_agg_function is not None: + self.group_agg_function.close() + super().teardown() + + def _extract_user_defined_agg_info(self, user_defined_function_proto): Review comment: Move to operation_utils and then could share it between beam_operations_fast and beam_operations_slow. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected]
