mlrun 1.3.3rc1__py3-none-any.whl → 1.4.0__py3-none-any.whl

mlrun/feature_store/feature_set.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -19,7 +19,7 @@ import pandas as pd
 from storey import EmitEveryEvent, EmitPolicy
 
 import mlrun
-import mlrun.api.schemas
+import mlrun.common.schemas
 
 from ..config import config as mlconf
 from ..datastore import get_store_uri
@@ -118,7 +118,7 @@ class FeatureSetSpec(ModelObj):
         self.owner = owner
         self.description = description
         self.entities: List[Union[Entity, str]] = entities or []
-        self.relations: Dict[str, Entity] = relations or {}
+        self.relations: Dict[str, Union[Entity, str]] = relations or {}
         self.features: List[Feature] = features or []
         self.partition_keys = partition_keys or []
         self.timestamp_key = timestamp_key
@@ -131,6 +131,7 @@ class FeatureSetSpec(ModelObj):
         self.engine = engine
         self.output_path = output_path or mlconf.artifact_path
         self.passthrough = passthrough
+        self.with_default_targets = True
 
     @property
     def entities(self) -> List[Entity]:
@@ -185,7 +186,8 @@ class FeatureSetSpec(ModelObj):
     @engine.setter
     def engine(self, engine: str):
         engine_list = ["pandas", "spark", "storey"]
-        if engine and engine not in engine_list:
+        engine = engine if engine else "storey"
+        if engine not in engine_list:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 f"engine must be one of {','.join(engine_list)}"
             )
@@ -232,6 +234,9 @@ class FeatureSetSpec(ModelObj):
 
     @relations.setter
     def relations(self, relations: Dict[str, Entity]):
+        for col, ent in relations.items():
+            if isinstance(ent, str):
+                relations[col] = Entity(ent)
         self._relations = ObjectDict.from_dict({"entity": Entity}, relations, "entity")
 
     def require_processing(self):
@@ -313,7 +318,7 @@ def emit_policy_to_dict(policy: EmitPolicy):
 class FeatureSet(ModelObj):
     """Feature set object, defines a set of features and their data pipeline"""
 
-    kind = mlrun.api.schemas.ObjectKind.feature_set.value
+    kind = mlrun.common.schemas.ObjectKind.feature_set.value
     _dict_fields = ["kind", "metadata", "spec", "status"]
 
     def __init__(
@@ -324,7 +329,7 @@ class FeatureSet(ModelObj):
         timestamp_key: str = None,
         engine: str = None,
         label_column: str = None,
-        relations: Dict[str, Entity] = None,
+        relations: Dict[str, Union[Entity, str]] = None,
         passthrough: bool = None,
     ):
         """Feature set object, defines a set of features and their data pipeline
@@ -372,6 +377,7 @@ class FeatureSet(ModelObj):
         self.status = None
         self._last_state = ""
         self._aggregations = {}
+        self.set_targets()
 
     @property
     def spec(self) -> FeatureSetSpec:
@@ -470,10 +476,25 @@ class FeatureSet(ModelObj):
             )
         targets = targets or []
         if with_defaults:
-            targets.extend(get_default_targets())
+            self.spec.with_default_targets = True
+            targets.extend(get_default_targets(offline_only=self.spec.passthrough))
+        else:
+            self.spec.with_default_targets = False
 
-        validate_target_list(targets=targets)
+        self.spec.targets = []
+        self.__set_targets_add_targets_helper(targets)
+
+        if default_final_step:
+            self.spec.graph.final_step = default_final_step
+
+    def __set_targets_add_targets_helper(self, targets):
+        """
+        Add the desired target list
 
+        :param targets: list of target type names ('csv', 'nosql', ..) or target objects
+                         CSVTarget(), ParquetTarget(), NoSqlTarget(), StreamTarget(), ..
+        """
+        validate_target_list(targets=targets)
         for target in targets:
             kind = target.kind if hasattr(target, "kind") else target
             if kind not in TargetTypes.all():
@@ -485,8 +506,6 @@ class FeatureSet(ModelObj):
                     target, name=str(target), partitioned=(target == "parquet")
                 )
             self.spec.targets.update(target)
-        if default_final_step:
-            self.spec.graph.final_step = default_final_step
 
     def validate_steps(self, namespace):
         if not self.spec:
@@ -522,7 +541,7 @@ class FeatureSet(ModelObj):
         :param silent: Fail silently if target doesn't exist in featureset status"""
 
         verify_feature_set_permissions(
-            self, mlrun.api.schemas.AuthorizationAction.delete
+            self, mlrun.common.schemas.AuthorizationAction.delete
         )
 
         purge_targets = self._reload_and_get_status_targets(
@@ -923,7 +942,18 @@ class FeatureSet(ModelObj):
                 raise mlrun.errors.MLRunNotFoundError(
                     "passthrough feature set {self.metadata.name} with no source"
                 )
-            return self.spec.source.to_dataframe()
+            df = self.spec.source.to_dataframe(
+                columns=columns,
+                start_time=start_time
+                or pd.Timestamp.min,  # overwrite `source.start_time` when the source is schedule.
+                end_time=end_time or pd.Timestamp.max,
+                time_field=time_column,
+                **kwargs,
+            )
+            # to_dataframe() can sometimes return an iterator of dataframes instead of one dataframe
+            if not isinstance(df, pd.DataFrame):
+                df = pd.concat(df)
+            return df
 
         target = get_offline_target(self, name=target_name)
         if not target:

mlrun/feature_store/feature_vector.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -154,7 +154,7 @@ class FeatureVectorStatus(ModelObj):
 class FeatureVector(ModelObj):
     """Feature vector, specify selected features, their metadata and material views"""
 
-    kind = mlrun.api.schemas.ObjectKind.feature_vector.value
+    kind = mlrun.common.schemas.ObjectKind.feature_vector.value
     _dict_fields = ["kind", "metadata", "spec", "status"]
 
     def __init__(
@@ -384,13 +384,23 @@ class FeatureVector(ModelObj):
 class OnlineVectorService:
     """get_online_feature_service response object"""
 
-    def __init__(self, vector, graph, index_columns, impute_policy: dict = None):
+    def __init__(
+        self,
+        vector,
+        graph,
+        index_columns,
+        all_fs_entities: List[str] = None,
+        impute_policy: dict = None,
+        requested_columns: List[str] = None,
+    ):
         self.vector = vector
         self.impute_policy = impute_policy or {}
 
         self._controller = graph.controller
         self._index_columns = index_columns
+        self._all_fs_entities = all_fs_entities
         self._impute_values = {}
+        self._requested_columns = requested_columns
 
     def __enter__(self):
         return self
@@ -493,38 +503,39 @@ class OnlineVectorService:
         for row in entity_rows:
             futures.append(self._controller.emit(row, return_awaitable_result=True))
 
-        requested_columns = list(self.vector.status.features.keys())
-        aliases = self.vector.get_feature_aliases()
-        for i, column in enumerate(requested_columns):
-            requested_columns[i] = aliases.get(column, column)
-
         for future in futures:
             result = future.await_result()
             data = result.body
-            for key in self._index_columns:
-                if data and key in data:
-                    del data[key]
-            if not data:
-                data = None
-            else:
+            if data:
                 actual_columns = data.keys()
-                for column in requested_columns:
+                if all([col in self._index_columns for col in actual_columns]):
+                    # didn't get any data from the graph
+                    results.append(None)
+                    continue
+                for column in self._requested_columns:
                     if (
                         column not in actual_columns
                         and column != self.vector.status.label_column
                     ):
                         data[column] = None
 
-            if self._impute_values and data:
-                for name in data.keys():
-                    v = data[name]
-                    if v is None or (type(v) == float and (np.isinf(v) or np.isnan(v))):
-                        data[name] = self._impute_values.get(name, v)
+                if self._impute_values:
+                    for name in data.keys():
+                        v = data[name]
+                        if v is None or (
+                            type(v) == float and (np.isinf(v) or np.isnan(v))
+                        ):
+                            data[name] = self._impute_values.get(name, v)
+                if not self.vector.spec.with_indexes:
+                    for name in self._all_fs_entities:
+                        data.pop(name, None)
+                if not any(data.values()):
+                    data = None
 
             if as_list and data:
                 data = [
                     data.get(key, None)
-                    for key in requested_columns
+                    for key in self._requested_columns
                     if key != self.vector.status.label_column
                 ]
             results.append(data)

mlrun/feature_store/ingestion.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -89,7 +89,7 @@ def init_featureset_graph(
     key_fields = entity_columns if entity_columns else None
 
     sizes = [0] * len(targets)
-    data_result = None
+    result_dfs = []
     total_rows = 0
     targets = [get_target_driver(target, featureset) for target in targets]
     if featureset.spec.passthrough:
@@ -100,11 +100,11 @@ def init_featureset_graph(
             # set the entities to be the indexes of the df
             event.body = entities_to_index(featureset, event.body)
 
-        data = server.run(event, get_body=True)
-        if data is not None:
+        df = server.run(event, get_body=True)
+        if df is not None:
             for i, target in enumerate(targets):
                 size = target.write_dataframe(
-                    data,
+                    df,
                     key_column=key_fields,
                     timestamp_key=featureset.spec.timestamp_key,
                     chunk_id=chunk_id,
@@ -112,21 +112,18 @@ def init_featureset_graph(
                 if size:
                     sizes[i] += size
         chunk_id += 1
-        if data_result is None:
-            # in case of multiple chunks only return the first chunk (last may be too small)
-            data_result = data
-        total_rows += data.shape[0]
+        result_dfs.append(df)
+        total_rows += df.shape[0]
         if rows_limit and total_rows >= rows_limit:
             break
 
-    # todo: fire termination event if iterator
-
     for i, target in enumerate(targets):
         target_status = target.update_resource_status("ready", size=sizes[i])
         if verbose:
             logger.info(f"wrote target: {target_status}")
 
-    return data_result
+    result_df = pd.concat(result_dfs)
+    return result_df.head(rows_limit)
 
 
 def featureset_initializer(server):

mlrun/feature_store/retrieval/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -16,13 +16,14 @@ import mlrun.errors
 from .dask_merger import DaskFeatureMerger
 from .job import run_merge_job  # noqa
 from .local_merger import LocalFeatureMerger
-from .online import init_feature_vector_graph  # noqa
 from .spark_merger import SparkFeatureMerger
+from .storey_merger import StoreyFeatureMerger
 
 mergers = {
     "local": LocalFeatureMerger,
     "dask": DaskFeatureMerger,
     "spark": SparkFeatureMerger,
+    "storey": StoreyFeatureMerger,
 }