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

mlrun/feature_store/retrieval/spark_merger.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.
@@ -12,17 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+
 import mlrun
 from mlrun.datastore.targets import get_offline_target
 
 from ...runtimes import RemoteSparkRuntime
 from ...runtimes.sparkjob.abstract import AbstractSparkRuntime
-from ..feature_vector import OfflineVectorResponse
 from .base import BaseMerger
 
 
 class SparkFeatureMerger(BaseMerger):
     engine = "spark"
+    support_offline = True
 
     def __init__(self, vector, **engine_args):
         super().__init__(vector, **engine_args)
@@ -33,109 +34,6 @@ class SparkFeatureMerger(BaseMerger):
     def to_spark_df(self, session, path):
         return session.read.load(path)
 
-    def _generate_vector(
-        self,
-        entity_rows,
-        entity_timestamp_column,
-        feature_set_objects,
-        feature_set_fields,
-        start_time=None,
-        end_time=None,
-        query=None,
-    ):
-        from pyspark.sql import SparkSession
-        from pyspark.sql.functions import col
-
-        if self.spark is None:
-            # create spark context
-            self.spark = SparkSession.builder.appName(
-                f"vector-merger-{self.vector.metadata.name}"
-            ).getOrCreate()
-
-        feature_sets = []
-        dfs = []
-
-        for name, columns in feature_set_fields.items():
-            feature_set = feature_set_objects[name]
-            feature_sets.append(feature_set)
-            column_names = [name for name, alias in columns]
-
-            if feature_set.spec.passthrough:
-                if not feature_set.spec.source:
-                    raise mlrun.errors.MLRunNotFoundError(
-                        f"passthrough feature set {name} with no source"
-                    )
-                source_kind = feature_set.spec.source.kind
-                source_path = feature_set.spec.source.path
-            else:
-                target = get_offline_target(feature_set)
-                if not target:
-                    raise mlrun.errors.MLRunInvalidArgumentError(
-                        f"feature set {name} does not have offline targets"
-                    )
-                source_kind = target.kind
-                source_path = target.get_target_path()
-
-            # handling case where there are multiple feature sets and user creates vector where
-            # entity_timestamp_column is from a specific feature set (can't be entity timestamp)
-            source_driver = mlrun.datastore.sources.source_kind_to_driver[source_kind]
-            if (
-                entity_timestamp_column in column_names
-                or feature_set.spec.timestamp_key == entity_timestamp_column
-            ):
-                source = source_driver(
-                    name=self.vector.metadata.name,
-                    path=source_path,
-                    time_field=entity_timestamp_column,
-                    start_time=start_time,
-                    end_time=end_time,
-                )
-            else:
-                source = source_driver(
-                    name=self.vector.metadata.name,
-                    path=source_path,
-                    time_field=entity_timestamp_column,
-                )
-
-            # add the index/key to selected columns
-            timestamp_key = feature_set.spec.timestamp_key
-
-            df = source.to_spark_df(
-                self.spark, named_view=self.named_view, time_field=timestamp_key
-            )
-
-            if timestamp_key and timestamp_key not in column_names:
-                columns.append((timestamp_key, None))
-            for entity in feature_set.spec.entities.keys():
-                if entity not in column_names:
-                    columns.append((entity, None))
-
-            # select requested columns and rename with alias where needed
-            df = df.select([col(name).alias(alias or name) for name, alias in columns])
-            dfs.append(df)
-            del df
-
-        # convert pandas entity_rows to spark DF if needed
-        if entity_rows is not None and not hasattr(entity_rows, "rdd"):
-            entity_rows = self.spark.createDataFrame(entity_rows)
-
-        # join the feature data frames
-        self.merge(entity_rows, entity_timestamp_column, feature_sets, dfs)
-
-        # filter joined data frame by the query param
-        if query:
-            self._result_df = self._result_df.filter(query)
-
-        self._result_df = self._result_df.drop(*self._drop_columns)
-
-        if self.vector.status.label_column:
-            self._result_df = self._result_df.dropna(
-                subset=[self.vector.status.label_column]
-            )
-
-        self._write_to_target()
-        return OfflineVectorResponse(self)
-
     def _unpersist_df(self, df):
         df.unpersist()
 
@@ -147,7 +45,6 @@ class SparkFeatureMerger(BaseMerger):
         featureset_df,
         left_keys: list,
         right_keys: list,
-        columns: list,
     ):
 
         """Perform an as of join between entity and featureset.
@@ -170,30 +67,32 @@ class SparkFeatureMerger(BaseMerger):
         from pyspark.sql.functions import col, monotonically_increasing_id, row_number
 
         entity_with_id = entity_df.withColumn("_row_nr", monotonically_increasing_id())
-        indexes = list(featureset.spec.entities.keys())
-
+        rename_right_keys = {}
+        for key in right_keys + [featureset.spec.timestamp_key]:
+            if key in entity_df.columns:
+                rename_right_keys[key] = f"ft__{key}"
         # get columns for projection
         projection = [
-            col(col_name).alias(
-                f"ft__{col_name}"
-                if col_name in indexes + [entity_timestamp_column]
-                else col_name
-            )
+            col(col_name).alias(rename_right_keys.get(col_name, col_name))
             for col_name in featureset_df.columns
         ]
 
         aliased_featureset_df = featureset_df.select(projection)
+        right_timestamp = rename_right_keys.get(
+            featureset.spec.timestamp_key, featureset.spec.timestamp_key
+        )
 
         # set join conditions
         join_cond = (
             entity_with_id[entity_timestamp_column]
-            >= aliased_featureset_df[f"ft__{entity_timestamp_column}"]
+            >= aliased_featureset_df[right_timestamp]
         )
 
         # join based on entities
-        for key in indexes:
+        for key_l, key_r in zip(left_keys, right_keys):
             join_cond = join_cond & (
-                entity_with_id[key] == aliased_featureset_df[f"ft__{key}"]
+                entity_with_id[key_l]
+                == aliased_featureset_df[rename_right_keys.get(key_r, key_r)]
             )
 
         conditional_join = entity_with_id.join(
@@ -201,20 +100,22 @@ class SparkFeatureMerger(BaseMerger):
         )
 
         window = Window.partitionBy("_row_nr").orderBy(
-            col(f"ft__{entity_timestamp_column}").desc(),
+            col(right_timestamp).desc(),
         )
         filter_most_recent_feature_timestamp = conditional_join.withColumn(
             "_rank", row_number().over(window)
         ).filter(col("_rank") == 1)
 
-        for key in indexes + [entity_timestamp_column]:
-            filter_most_recent_feature_timestamp = (
-                filter_most_recent_feature_timestamp.drop(
-                    aliased_featureset_df[f"ft__{key}"]
+        for key in right_keys + [featureset.spec.timestamp_key]:
+            if key in entity_df.columns + [entity_timestamp_column]:
+                filter_most_recent_feature_timestamp = (
+                    filter_most_recent_feature_timestamp.drop(
+                        aliased_featureset_df[f"ft__{key}"]
+                    )
                 )
-            )
-
-        return filter_most_recent_feature_timestamp.drop("_row_nr", "_rank")
+        return filter_most_recent_feature_timestamp.drop("_row_nr", "_rank").orderBy(
+            col(entity_timestamp_column)
+        )
 
     def _join(
         self,
@@ -224,7 +125,6 @@ class SparkFeatureMerger(BaseMerger):
         featureset_df,
         left_keys: list,
         right_keys: list,
-        columns: list,
     ):
 
         """
@@ -245,8 +145,19 @@ class SparkFeatureMerger(BaseMerger):
                 be prefixed with featureset_df name.
 
         """
-        indexes = list(featureset.spec.entities.keys())
-        merged_df = entity_df.join(featureset_df, on=indexes)
+        if left_keys != right_keys:
+            join_cond = [
+                entity_df[key_l] == featureset_df[key_r]
+                for key_l, key_r in zip(left_keys, right_keys)
+            ]
+        else:
+            join_cond = left_keys
+
+        merged_df = entity_df.join(
+            featureset_df,
+            join_cond,
+            how=self._join_type,
+        )
         return merged_df
 
     def get_df(self, to_pandas=True):
@@ -268,3 +179,94 @@ class SparkFeatureMerger(BaseMerger):
             return RemoteSparkRuntime.default_image
         else:
             raise mlrun.errors.MLRunInvalidArgumentError(f"Unsupported kind '{kind}'")
+
+    def _create_engine_env(self):
+        from pyspark.sql import SparkSession
+
+        if self.spark is None:
+            # create spark context
+            self.spark = SparkSession.builder.appName(
+                f"vector-merger-{self.vector.metadata.name}"
+            ).getOrCreate()
+
+    def _get_engine_df(
+        self,
+        feature_set,
+        feature_set_name,
+        column_names=None,
+        start_time=None,
+        end_time=None,
+        time_column=None,
+    ):
+        source_kwargs = {}
+        if feature_set.spec.passthrough:
+            if not feature_set.spec.source:
+                raise mlrun.errors.MLRunNotFoundError(
+                    f"passthrough feature set {feature_set_name} with no source"
+                )
+            source_kind = feature_set.spec.source.kind
+            source_path = feature_set.spec.source.path
+            source_kwargs.update(feature_set.spec.source.attributes)
+        else:
+            target = get_offline_target(feature_set)
+            if not target:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    f"feature set {feature_set_name} does not have offline targets"
+                )
+            source_kind = target.kind
+            source_path = target.get_target_path()
+
+        # handling case where there are multiple feature sets and user creates vector where
+        # entity_timestamp_column is from a specific feature set (can't be entity timestamp)
+        source_driver = mlrun.datastore.sources.source_kind_to_driver[source_kind]
+
+        source = source_driver(
+            name=self.vector.metadata.name,
+            path=source_path,
+            time_field=time_column,
+            start_time=start_time,
+            end_time=end_time,
+            **source_kwargs,
+        )
+
+        columns = column_names + [ent.name for ent in feature_set.spec.entities]
+        if (
+            feature_set.spec.timestamp_key
+            and feature_set.spec.timestamp_key not in columns
+        ):
+            columns.append(feature_set.spec.timestamp_key)
+
+        return source.to_spark_df(
+            self.spark,
+            named_view=self.named_view,
+            time_field=time_column,
+            columns=columns,
+        )
+
+    def _rename_columns_and_select(
+        self,
+        df,
+        rename_col_dict,
+        columns=None,
+    ):
+        from pyspark.sql.functions import col
+
+        return df.select(
+            [
+                col(name).alias(rename_col_dict.get(name, name))
+                for name in columns or rename_col_dict.keys()
+            ]
+        )
+
+    def _drop_columns_from_result(self):
+        self._result_df = self._result_df.drop(*self._drop_columns)
+
+    def _filter(self, query):
+        self._result_df = self._result_df.filter(query)
+
+    def _order_by(self, order_by_active):
+        from pyspark.sql.functions import col
+
+        self._result_df = self._result_df.orderBy(
+            *[col(col_name).asc_nulls_last() for col_name in order_by_active]
+        )

mlrun/feature_store/retrieval/storey_merger.py

@@ -0,0 +1,173 @@
+# Copyright 2018 Iguazio
+#
+# Licensed 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.
+import mlrun
+from mlrun.datastore.store_resources import ResourceCache
+from mlrun.datastore.targets import get_online_target
+from mlrun.serving.server import create_graph_server
+
+from ..feature_vector import OnlineVectorService
+from .base import BaseMerger
+
+
+class StoreyFeatureMerger(BaseMerger):
+    engine = "storey"
+    support_online = True
+
+    def __init__(self, vector, **engine_args):
+        super().__init__(vector, **engine_args)
+        self.impute_policy = engine_args.get("impute_policy")
+
+    def _generate_online_feature_vector_graph(
+        self,
+        entity_keys,
+        feature_set_fields,
+        feature_set_objects,
+        fixed_window_type,
+    ):
+        graph = self.vector.spec.graph.copy()
+        start_states, default_final_state, responders = graph.check_and_process_graph(
+            allow_empty=True
+        )
+        next = graph
+
+        fs_link_list = self._create_linked_relation_list(
+            feature_set_objects, feature_set_fields, entity_keys
+        )
+
+        all_columns = []
+        save_column = []
+        entity_keys = []
+        end_aliases = {}
+        for node in fs_link_list:
+            name = node.name
+            if name == self._entity_rows_node_name:
+                continue
+            featureset = feature_set_objects[name]
+            columns = feature_set_fields[name]
+            column_names = [name for name, alias in columns]
+            aliases = {name: alias for name, alias in columns if alias}
+            all_columns += [aliases.get(name, name) for name in column_names]
+            for col in node.data["save_cols"]:
+                if col not in column_names:
+                    column_names.append(col)
+                else:
+                    save_column.append(col)
+
+            entity_list = node.data["right_keys"] or list(
+                featureset.spec.entities.keys()
+            )
+            if not entity_keys:
+                # if entity_keys not provided by the user we will set it to be the entity of the first feature set
+                entity_keys = entity_list
+            end_aliases.update(
+                {
+                    k: v
+                    for k, v in zip(entity_list, node.data["left_keys"])
+                    if k != v and v in save_column
+                }
+            )
+            mapping = {
+                k: v for k, v in zip(node.data["left_keys"], entity_list) if k != v
+            }
+            if mapping:
+                next = next.to(
+                    "storey.Rename",
+                    f"rename-{name}",
+                    mapping=mapping,
+                )
+
+            next = next.to(
+                "storey.QueryByKey",
+                f"query-{name}",
+                features=column_names,
+                table=featureset.uri,
+                key_field=entity_list,
+                aliases=aliases,
+                fixed_window_type=fixed_window_type.to_qbk_fixed_window_type(),
+            )
+        if end_aliases:
+            # run if the user want to save a column that related to another entity
+            next = next.to(
+                "storey.Rename",
+                "rename-entity-to-features",
+                mapping=end_aliases,
+            )
+        for name in start_states:
+            next.set_next(name)
+
+        if not start_states:  # graph was empty
+            next.respond()
+        elif not responders and default_final_state:  # graph has clear state sequence
+            graph[default_final_state].respond()
+        elif not responders:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "the graph doesnt have an explicit final step to respond on"
+            )
+        return graph, all_columns, entity_keys
+
+    def init_online_vector_service(
+        self, entity_keys, fixed_window_type, update_stats=False
+    ):
+        try:
+            from storey import SyncEmitSource
+        except ImportError as exc:
+            raise ImportError(f"storey not installed, use pip install storey, {exc}")
+
+        feature_set_objects, feature_set_fields = self.vector.parse_features(
+            offline=False, update_stats=update_stats
+        )
+        if not feature_set_fields:
+            raise mlrun.errors.MLRunRuntimeError(
+                f"No features found for feature vector '{self.vector.metadata.name}'"
+            )
+        (
+            graph,
+            requested_columns,
+            entity_keys,
+        ) = self._generate_online_feature_vector_graph(
+            entity_keys,
+            feature_set_fields,
+            feature_set_objects,
+            fixed_window_type,
+        )
+        graph.set_flow_source(SyncEmitSource())
+        server = create_graph_server(graph=graph, parameters={})
+
+        cache = ResourceCache()
+        all_fs_entities = []
+        for featureset in feature_set_objects.values():
+            driver = get_online_target(featureset)
+            if not driver:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    f"resource {featureset.uri} does not have an online data target"
+                )
+            cache.cache_table(featureset.uri, driver.get_table_object())
+
+            for key in featureset.spec.entities.keys():
+                if key not in all_fs_entities:
+                    all_fs_entities.append(key)
+        server.init_states(context=None, namespace=None, resource_cache=cache)
+        server.init_object(None)
+
+        service = OnlineVectorService(
+            self.vector,
+            graph,
+            entity_keys,
+            all_fs_entities=all_fs_entities,
+            impute_policy=self.impute_policy,
+            requested_columns=requested_columns,
+        )
+        service.initialize()
+
+        return service