Aggregation optimizations. (#144)

* split aggregateStoreElement to query and aggregate, where the latter is optimized

* lint

* typo + minor performance improvment

* AggregationValue optimizations.

* Replace min and max with if.

* Further optimizations.

* Optimization.

* Rename method.

* raise NotImplementedError

* Reuse AggregationValues to avoid allocations.

* More use of reset. Fix reset.

* More use of reset.

* More AggregationValue reuse.

* Remove dict key redundancy.

* Optimization.

* Optimization.

* Simplification/optimization

* Optimize _set_value.

* Remove max_value from reset().

Co-authored-by: Tal Neiman <Taln@iguazio.com>
Co-authored-by: Gal Topper <galt@iguazio.com>

Author: 3 people

storey/aggregations.py

@@ -1,5 +1,4 @@
 import asyncio
-import copy
 from datetime import datetime
 import time
 import re
@@ -174,10 +173,9 @@ async def _emit_event(self, key, event):
             emitted_attr_name = self._aliases.get(col, None) or col
             if col in self._table._get_static_attrs(key):
                 features[emitted_attr_name] = self._table._get_static_attrs(key)[col]
-        new_event = copy.copy(event)
-        new_event.key = key
-        new_event.body = features
-        await self._do_downstream(new_event)
+        event.key = key
+        event.body = features
+        await self._do_downstream(event)
 
     # Emit multiple events for every key in the store with the current time
     async def _emit_all_events(self, timestamp):

storey/dataframe.py

@@ -1,5 +1,3 @@
-import copy
-
 import pandas as pd
 
 from .flow import _termination_obj, Flow
@@ -78,6 +76,5 @@ async def _do(self, event):
             df = pd.DataFrame(event.body, columns=self._columns)
             if self._index:
                 df.set_index(self._index, inplace=True)
-            new_event = copy.copy(event)
-            new_event.body = df
-            return await self._do_downstream(new_event)
+            event.body = df
+            return await self._do_downstream(event)

storey/drivers.py

@@ -252,21 +252,21 @@ def _build_conditioned_feature_store_request(self, aggregation_element, pending=
                         if not initialized_attributes.get(array_attribute_name, 0) == expected_time:
                             initialized_attributes[array_attribute_name] = expected_time
                             init_expression = f'{array_attribute_name}=if_else(({get_array_time_expr}<{expected_time_expr}),' \
-                                f"init_array({bucket.total_number_of_buckets},'double'," \
-                                f'{aggregation_value.get_default_value()}),{array_attribute_name})'
+                                              f"init_array({bucket.total_number_of_buckets},'double'," \
+                                              f'{aggregation_value.default_value}),{array_attribute_name})'
                             expressions.append(init_expression)
 
                         arr_at_index = f'{array_attribute_name}[{index_to_update}]'
                         update_array_expression = f'{arr_at_index}=if_else(({get_array_time_expr}>{expected_time_expr}),{arr_at_index},' \
-                            f'{self._get_update_expression_by_aggregation(arr_at_index, aggregation_value)})'
+                                                  f'{aggregation_value.get_update_expression(arr_at_index)})'
 
                         expressions.append(update_array_expression)
 
                         # Separating time attribute updates, so that they will be executed in the end and only once per feature name.
                         if array_time_attribute_name not in times_update_expressions:
                             times_update_expressions[array_time_attribute_name] = f'{array_time_attribute_name}=' \
-                                f'if_else(({get_array_time_expr}<{expected_time_expr}),' \
-                                f'{expected_time_expr},{array_time_attribute_name})'
+                                                                                  f'if_else(({get_array_time_expr}<{expected_time_expr}),' \
+                                                                                  f'{expected_time_expr},{array_time_attribute_name})'
 
         expressions.extend(times_update_expressions.values())
 
@@ -304,7 +304,7 @@ def _build_simplified_feature_store_request(self, aggregation_element):
                             if not initialized_attributes.get(array_attribute_name, 0) == expected_time:
                                 initialized_attributes[array_attribute_name] = expected_time
                                 expressions.append(f"{array_attribute_name}=init_array({bucket.total_number_of_buckets},'double',"
-                                                   f'{aggregation_value.get_default_value()})')
+                                                   f'{aggregation_value.default_value})')
                             if array_time_attribute_name not in times_update_expressions:
                                 times_update_expressions[array_time_attribute_name] = \
                                     f'{array_time_attribute_name}={expected_time_expr}'
@@ -313,8 +313,7 @@ def _build_simplified_feature_store_request(self, aggregation_element):
                         # Updating the specific cells
                         if cached_time <= expected_time:
                             arr_at_index = f'{array_attribute_name}[{index_to_update}]'
-                            expressions.append(
-                                f'{arr_at_index}={self._get_update_expression_by_aggregation(arr_at_index, aggregation_value)}')
+                            expressions.append(f'{arr_at_index}={aggregation_value.get_update_expression(arr_at_index)}')
 
         # Separating time attribute updates, so that they will be executed in the end and only once per feature name.
         expressions.extend(times_update_expressions.values())
@@ -325,20 +324,6 @@ def _build_simplified_feature_store_request(self, aggregation_element):
             aggregation_element.aggregation_buckets[name].storage_specific_cache[attribute_name] = new_time
         return expressions, pending_updates
 
-    @staticmethod
-    def _get_update_expression_by_aggregation(old, aggregation):
-        value = aggregation.get_value()[1]
-        if aggregation.aggregation == 'max':
-            return f'max({old}, {value})'
-        elif aggregation.aggregation == 'min':
-            return f'min({old}, {value})'
-        elif aggregation.aggregation == 'last':
-            return f'{value}'
-        elif aggregation.aggregation == 'first':
-            return f'if_else(({old} == {aggregation.get_default_value()}), {value}, {old})'
-        else:
-            return f'{old}+{value}'
-
     @staticmethod
     def _convert_python_obj_to_expression_value(value):
         if isinstance(value, str):