[Feature store] Support EmitPolicy in aggregations (#903)

mlrun/feature_store/feature_set.py

@@ -11,9 +11,10 @@
 # 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 typing import List
+from typing import List, Optional
 
 import pandas as pd
+from storey import EmitPolicy
 
 import mlrun
 
@@ -351,6 +352,7 @@ def add_aggregation(
         state_name=None,
         after=None,
         before=None,
+        emit_policy: Optional[EmitPolicy] = None,
     ):
         """add feature aggregation rule
 
@@ -366,6 +368,8 @@ def add_aggregation(
         :param state_name: optional, graph state name
         :param after:      optional, after which graph state it runs
         :param before:     optional, comes before graph state
+        :param emit_policy optional. Define emit policy of the aggregations. For example EmitAfterMaxEvent (will emit
+                            the Nth event). The default behaviour is emitting every event
         """
         aggregation = FeatureAggregation(
             name, column, operations, windows, period
@@ -384,14 +388,20 @@ def upsert_feature(name):
             aggregations = state.class_args.get("aggregates", [])
             aggregations.append(aggregation)
             state.class_args["aggregates"] = aggregations
+            if emit_policy:
+                state.class_args["emit_policy"] = emit_policy
         else:
+            class_args = {}
+            if emit_policy:
+                class_args["emit_policy"] = emit_policy
             state = graph.add_step(
                 name=state_name,
                 after=after or previous_step,
                 before=before,
                 class_name="storey.AggregateByKey",
                 aggregates=[aggregation],
                 table=".",
+                **class_args,
             )
 
         for operation in operations:

requirements.txt

@@ -51,4 +51,4 @@ fsspec~=0.9.0
 v3iofs~=0.1.5
 # 3.4 and above failed builidng in some images - see https://github.com/pyca/cryptography/issues/5771
 cryptography~=3.3.2
-storey~=0.4.6; python_version >= '3.7'
+storey~=0.4.8; python_version >= '3.7'

tests/system/feature_store/test_feature_store.py

@@ -5,7 +5,7 @@
 
 import pandas as pd
 import pytest
-from storey import MapClass
+from storey import EmitAfterMaxEvent, MapClass
 
 import mlrun
 import mlrun.feature_store as fs
@@ -396,6 +396,7 @@ def test_multiple_entities(self):
             operations=["sum", "max"],
             windows=["1h"],
             period="10m",
+            emit_policy=EmitAfterMaxEvent(1),
         )
         fs.infer_metadata(
             data_set,