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

mlrun/feature_store/retrieval/base.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.
@@ -13,16 +13,30 @@
 # limitations under the License.
 #
 import abc
+import typing
+from datetime import datetime
+
+import dask.dataframe as dd
+import pandas as pd
 
 import mlrun
 from mlrun.datastore.targets import CSVTarget, ParquetTarget
+from mlrun.feature_store.feature_set import FeatureSet
+from mlrun.feature_store.feature_vector import Feature
 
-from ...utils import logger
+from ...utils import logger, str_to_timestamp
+from ..feature_vector import OfflineVectorResponse
 
 
 class BaseMerger(abc.ABC):
     """abstract feature merger class"""
 
+    # In order to be an online merger, the merger should implement `init_online_vector_service` function.
+    support_online = False
+
+    # In order to be an offline merger, the merger should implement
+    # `_order_by`, `_filter`, `_drop_columns_from_result`, `_rename_columns_and_select`, `_get_engine_df` functions.
+    support_offline = False
     engine = None
 
     def __init__(self, vector, **engine_args):
@@ -36,6 +50,8 @@ class BaseMerger(abc.ABC):
         self._drop_indexes = True
         self._target = None
         self._alias = dict()
+        self._origin_alias = dict()
+        self._entity_rows_node_name = "__mlrun__$entity_rows$"
 
     def _append_drop_column(self, key):
         if key and key not in self._drop_columns:
@@ -67,22 +83,19 @@ class BaseMerger(abc.ABC):
         drop_columns=None,
         start_time=None,
         end_time=None,
+        timestamp_for_filtering=None,
         with_indexes=None,
         update_stats=None,
         query=None,
-        join_type="inner",
+        order_by=None,
     ):
         self._target = target
-        self._join_type = join_type
 
         # calculate the index columns and columns we need to drop
         self._drop_columns = drop_columns or self._drop_columns
         if self.vector.spec.with_indexes or with_indexes:
             self._drop_indexes = False
 
-        if entity_timestamp_column and self._drop_indexes:
-            self._append_drop_column(entity_timestamp_column)
-
         # retrieve the feature set objects/fields needed for the vector
         feature_set_objects, feature_set_fields = self.vector.parse_features(
             update_stats=update_stats
@@ -96,23 +109,34 @@ class BaseMerger(abc.ABC):
             # update the feature vector objects with refreshed stats
             self.vector.save()
 
+        if self._drop_indexes and entity_timestamp_column:
+            self._append_drop_column(entity_timestamp_column)
+
         for feature_set in feature_set_objects.values():
-            if not entity_timestamp_column and self._drop_indexes:
+            if self._drop_indexes:
                 self._append_drop_column(feature_set.spec.timestamp_key)
             for key in feature_set.spec.entities.keys():
                 self._append_index(key)
 
-        return self._generate_vector(
+        start_time = str_to_timestamp(start_time)
+        end_time = str_to_timestamp(end_time)
+        if start_time and not end_time:
+            # if end_time is not specified set it to now()
+            end_time = pd.Timestamp.now()
+
+        return self._generate_offline_vector(
             entity_rows,
             entity_timestamp_column,
             feature_set_objects=feature_set_objects,
             feature_set_fields=feature_set_fields,
             start_time=start_time,
             end_time=end_time,
+            timestamp_for_filtering=timestamp_for_filtering,
             query=query,
+            order_by=order_by,
         )
 
-    def _write_to_target(self):
+    def _write_to_offline_target(self):
         if self._target:
             is_persistent_vector = self.vector.metadata.name is not None
             if not self._target.path and not is_persistent_vector:
@@ -125,6 +149,14 @@ class BaseMerger(abc.ABC):
                 target_status = self._target.update_resource_status("ready", size=size)
                 logger.info(f"wrote target: {target_status}")
                 self.vector.save()
+        if not self._drop_indexes:
+            self.vector.spec.entity_fields = [
+                Feature(name=feature, value_type=self._result_df[feature].dtype)
+                if self._result_df[feature].dtype.name != "object"
+                else Feature(name=feature, value_type="str")
+                for feature in self._index_columns
+            ]
+            self.vector.save()
 
     def _set_indexes(self, df):
         if self._index_columns and not self._drop_indexes:
@@ -134,29 +166,16 @@ class BaseMerger(abc.ABC):
                     if index not in df.columns:
                         index_columns_missing.append(index)
                 if not index_columns_missing:
-                    if self.engine == "local" or self.engine == "spark":
-                        df.set_index(self._index_columns, inplace=True)
-                    elif self.engine == "dask":
-                        if len(self._index_columns) == 1:
-                            return df.set_index(self._index_columns[0])
-                        elif len(self._index_columns) != 1:
-                            return self._reset_index(self._result_df)
-                        else:
-                            logger.info(
-                                "The entities will stay as columns because "
-                                "Dask dataframe does not yet support multi-indexes"
-                            )
-                            return self._result_df
+                    df.set_index(self._index_columns, inplace=True)
                 else:
                     logger.warn(
                         f"Can't set index, not all index columns found: {index_columns_missing}. "
                         f"It is possible that column was already indexed."
                     )
-            else:
-                return df
+        else:
+            df.reset_index(drop=True, inplace=True)
 
-    @abc.abstractmethod
-    def _generate_vector(
+    def _generate_offline_vector(
         self,
         entity_rows,
         entity_timestamp_column,
@@ -164,9 +183,203 @@ class BaseMerger(abc.ABC):
         feature_set_fields,
         start_time=None,
         end_time=None,
+        timestamp_for_filtering=None,
         query=None,
+        order_by=None,
     ):
-        raise NotImplementedError("_generate_vector() operation not supported in class")
+        self._create_engine_env()
+
+        feature_sets = []
+        dfs = []
+        keys = (
+            []
+        )  # the struct of key is [[[],[]], ..] So that each record indicates which way the corresponding
+        # featureset is connected to the previous one, and within each record the left keys are indicated in index 0
+        # and the right keys in index 1, this keys will be the keys that will be used in this join
+
+        fs_link_list = self._create_linked_relation_list(
+            feature_set_objects, feature_set_fields
+        )
+
+        filtered = False
+        for node in fs_link_list:
+            name = node.name
+            feature_set = feature_set_objects[name]
+            feature_sets.append(feature_set)
+            columns = feature_set_fields[name]
+            self._origin_alias.update({name: alias for name, alias in columns})
+            column_names = [name for name, _ in columns]
+
+            for column in node.data["save_cols"]:
+                if column not in column_names:
+                    column_names.append(column)
+                    if column not in self._index_columns:
+                        self._append_drop_column(column)
+
+            if isinstance(timestamp_for_filtering, dict):
+                time_column = timestamp_for_filtering.get(
+                    name, feature_set.spec.timestamp_key
+                )
+            elif isinstance(timestamp_for_filtering, str):
+                time_column = timestamp_for_filtering
+            else:
+                time_column = feature_set.spec.timestamp_key
+
+            if time_column != feature_set.spec.timestamp_key and time_column not in [
+                feature.name for feature in feature_set.spec.features
+            ]:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    f"Feature set `{name}` "
+                    f"does not have a column named `{time_column}` to filter on."
+                )
+
+            if self._drop_indexes:
+                self._append_drop_column(time_column)
+            if (start_time or end_time) and time_column:
+                filtered = True
+
+            df = self._get_engine_df(
+                feature_set,
+                name,
+                column_names,
+                start_time if time_column else None,
+                end_time if time_column else None,
+                time_column,
+            )
+
+            column_names += node.data["save_index"]
+            node.data["save_cols"] += node.data["save_index"]
+            fs_entities_and_timestamp = list(feature_set.spec.entities.keys())
+            if feature_set.spec.timestamp_key:
+                column_names.append(feature_set.spec.timestamp_key)
+                node.data["save_cols"].append(feature_set.spec.timestamp_key)
+                fs_entities_and_timestamp.append(feature_set.spec.timestamp_key)
+
+            # rename columns to be unique for each feature set and select if needed
+            rename_col_dict = {
+                column: f"{column}_{name}"
+                for column in column_names
+                if column not in node.data["save_cols"]
+            }
+            df_temp = self._rename_columns_and_select(
+                df,
+                rename_col_dict,
+                columns=list(set(column_names + fs_entities_and_timestamp)),
+            )
+
+            if df_temp is not None:
+                df = df_temp
+                del df_temp
+
+            dfs.append(df)
+            del df
+
+            keys.append([node.data["left_keys"], node.data["right_keys"]])
+
+            # update alias according to the unique column name
+            new_columns = []
+            if not self._drop_indexes:
+                new_columns.extend([(ind, ind) for ind in fs_entities_and_timestamp])
+            for column, alias in columns:
+                if column in rename_col_dict:
+                    new_columns.append((rename_col_dict[column], alias or column))
+                else:
+                    new_columns.append((column, alias))
+            self._update_alias(dictionary={name: alias for name, alias in new_columns})
+
+        # None of the feature sets was filtered as required
+        if not filtered and (start_time or end_time):
+            raise mlrun.errors.MLRunRuntimeError(
+                "start_time and end_time can only be provided in conjunction with "
+                "a timestamp column, or when the at least one feature_set has a timestamp key"
+            )
+        # convert pandas entity_rows to spark\dask DF if needed
+        if (
+            entity_rows is not None
+            and not hasattr(entity_rows, "rdd")
+            and self.engine == "spark"
+        ):
+            entity_rows = self.spark.createDataFrame(entity_rows)
+        elif (
+            entity_rows is not None
+            and not hasattr(entity_rows, "dask")
+            and self.engine == "dask"
+        ):
+            entity_rows = dd.from_pandas(
+                entity_rows, npartitions=len(entity_rows.columns)
+            )
+
+        # join the feature data frames
+        result_timestamp = self.merge(
+            entity_df=entity_rows,
+            entity_timestamp_column=entity_timestamp_column
+            if entity_rows is not None
+            else None,
+            featuresets=feature_sets,
+            featureset_dfs=dfs,
+            keys=keys,
+        )
+
+        all_columns = None
+        if not self._drop_indexes and result_timestamp:
+            if result_timestamp not in self._alias.values():
+                self._update_alias(key=result_timestamp, val=result_timestamp)
+            all_columns = list(self._alias.keys())
+
+        df_temp = self._rename_columns_and_select(
+            self._result_df, self._alias, columns=all_columns
+        )
+        if df_temp is not None:
+            self._result_df = df_temp
+            del df_temp
+
+        df_temp = self._drop_columns_from_result()
+        if df_temp is not None:
+            self._result_df = df_temp
+            del df_temp
+
+        if self.vector.status.label_column:
+            self._result_df = self._result_df.dropna(
+                subset=[self.vector.status.label_column]
+            )
+        # filter joined data frame by the query param
+        if query:
+            self._filter(query)
+
+        if order_by:
+            if isinstance(order_by, str):
+                order_by = [order_by]
+            order_by_active = [
+                order_col
+                if order_col in self._result_df.columns
+                else self._origin_alias.get(order_col, None)
+                for order_col in order_by
+            ]
+            if None in order_by_active:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    f"Result dataframe contains {self._result_df.columns} "
+                    f"columns and can't order by {order_by}"
+                )
+            self._order_by(order_by_active)
+
+        self._write_to_offline_target()
+        return OfflineVectorResponse(self)
+
+    def init_online_vector_service(
+        self, entity_keys, fixed_window_type, update_stats=False
+    ):
+        """
+        initialize the `OnlineVectorService`
+
+        :param entity_keys:         list of the feature_vector indexes.
+        :param fixed_window_type:   determines how to query the fixed window values which were previously
+                                    inserted by ingest
+        :param update_stats:        update features statistics from the requested feature sets on the vector.
+                                    Default: False.
+
+        :return:                    `OnlineVectorService`
+        """
+        raise NotImplementedError
 
     def _unpersist_df(self, df):
         pass
@@ -178,7 +391,6 @@ class BaseMerger(abc.ABC):
         featuresets: list,
         featureset_dfs: list,
         keys: list = None,
-        all_columns: list = None,
     ):
         """join the entities and feature set features into a result dataframe"""
         merged_df = entity_df
@@ -190,10 +402,6 @@ class BaseMerger(abc.ABC):
             else:
                 # keys can be multiple keys on each side of the join
                 keys = [[[], []]] * len(featureset_dfs)
-            if all_columns is not None:
-                all_columns.pop(0)
-            else:
-                all_columns = [[]] * len(featureset_dfs)
             entity_timestamp_column = (
                 entity_timestamp_column or featureset.spec.timestamp_key
             )
@@ -203,16 +411,9 @@ class BaseMerger(abc.ABC):
             # and it can join only by the entities of the first `featureset`
             keys[0][0] = keys[0][1] = list(featuresets[0].spec.entities.keys())
 
-        for featureset, featureset_df, lr_key, columns in zip(
-            featuresets, featureset_dfs, keys, all_columns
-        ):
-            if featureset.spec.timestamp_key:
+        for featureset, featureset_df, lr_key in zip(featuresets, featureset_dfs, keys):
+            if featureset.spec.timestamp_key and entity_timestamp_column:
                 merge_func = self._asof_join
-                if self._join_type != "inner":
-                    logger.warn(
-                        "Merge all the features with as_of_join and don't "
-                        "take into account the join_type that was given"
-                    )
             else:
                 merge_func = self._join
 
@@ -223,7 +424,9 @@ class BaseMerger(abc.ABC):
                 featureset_df,
                 lr_key[0],
                 lr_key[1],
-                columns,
+            )
+            entity_timestamp_column = (
+                entity_timestamp_column or featureset.spec.timestamp_key
             )
 
             # unpersist as required by the implementation (e.g. spark) and delete references
@@ -232,8 +435,8 @@ class BaseMerger(abc.ABC):
             del featureset_df
 
         self._result_df = merged_df
+        return entity_timestamp_column
 
-    @abc.abstractmethod
     def _asof_join(
         self,
         entity_df,
@@ -242,11 +445,9 @@ class BaseMerger(abc.ABC):
         featureset_df,
         left_keys: list,
         right_keys: list,
-        columns: list,
     ):
         raise NotImplementedError("_asof_join() operation not implemented in class")
 
-    @abc.abstractmethod
     def _join(
         self,
         entity_df,
@@ -255,7 +456,6 @@ class BaseMerger(abc.ABC):
         featureset_df,
         left_keys: list,
         right_keys: list,
-        columns: list,
     ):
         raise NotImplementedError("_join() operation not implemented in class")
 
@@ -267,6 +467,7 @@ class BaseMerger(abc.ABC):
 
     def get_df(self, to_pandas=True):
         """return the result as a dataframe (pandas by default)"""
+        self._set_indexes(self._result_df)
         return self._result_df
 
     def to_parquet(self, target_path, **kw):
@@ -293,6 +494,9 @@ class BaseMerger(abc.ABC):
         def __eq__(self, other):
             return self.name == other.name
 
+        def __copy__(self):
+            return BaseMerger._Node(self.name, self.order, self.data.copy())
+
     class _LinkedList:
         def __init__(self, head=None):
             self.head = head
@@ -313,6 +517,19 @@ class BaseMerger(abc.ABC):
                 yield node
                 node = node.next
 
+        def __copy__(self):
+            ll = BaseMerger._LinkedList()
+            prev_node = None
+            for node in self:
+                new_node = node.__copy__()
+                if ll.head is None:
+                    ll.head = new_node
+                else:
+                    prev_node.next = new_node
+                prev_node = new_node
+            ll.len = self.len
+            return ll
+
         def add_first(self, node):
             node.next = self.head
             self.head = node
@@ -325,7 +542,9 @@ class BaseMerger(abc.ABC):
             for current_node in self:
                 pass
             current_node.next = node
-            self.len += 1
+            while node:
+                self.len += 1
+                node = node.next
 
         def add_after(self, target_node, new_node):
             new_node.next = target_node.next
@@ -346,7 +565,9 @@ class BaseMerger(abc.ABC):
             node = self.find_node(other_head.name)
             if node is None:
                 return
-            node.data["save_cols"] += other_head.data["save_cols"]
+            for col in other_head.data["save_cols"]:
+                if col not in node.data["save_cols"]:
+                    node.data["save_cols"].append(col)
             for other_node in other_iter:
                 if self.find_node(other_node.name) is None:
                     while node is not None and other_node.order > node.order:
@@ -357,10 +578,11 @@ class BaseMerger(abc.ABC):
                         self.add_last(other_node)
                     node = other_node
 
-    @staticmethod
-    def _create_linked_relation_list(feature_set_objects, feature_set_fields):
+    def _create_linked_relation_list(
+        self, feature_set_objects, feature_set_fields, entity_rows_keys=None
+    ):
         feature_set_names = list(feature_set_fields.keys())
-        if len(feature_set_names) == 1:
+        if len(feature_set_names) == 1 and not entity_rows_keys:
             return BaseMerger._LinkedList(
                 head=BaseMerger._Node(
                     name=feature_set_names[0],
@@ -420,10 +642,9 @@ class BaseMerger(abc.ABC):
                 )
             )
 
-            # checking if feature_set have relation with feature_set_in
-            relation_wise = all(curr_col_relation_list)
-
-            if relation_wise:
+            if all(
+                curr_col_relation_list
+            ):  # checking if feature_set have relation with feature_set_in
                 # add to the link list feature set according to the defined relation
                 linked_list_relation.add_last(
                     BaseMerger._Node(
@@ -437,8 +658,8 @@ class BaseMerger(abc.ABC):
                         order=name_in_order,
                     )
                 )
-                linked_list_relation.head.data["save_cols"].append(
-                    *curr_col_relation_list
+                linked_list_relation.head.data["save_cols"].extend(
+                    curr_col_relation_list
                 )
             elif name_in_order > head_order and sorted(
                 feature_set_in_entity_list_names
@@ -460,26 +681,127 @@ class BaseMerger(abc.ABC):
                 linked_list_relation.head.data["save_index"] = keys
             return linked_list_relation
 
+        def _build_entity_rows_relation(entity_rows_relation, fs_name, fs_order):
+            feature_set_entity_list = feature_set_entity_list_dict[fs_name]
+            feature_set_entity_list_names = list(feature_set_entity_list.keys())
+
+            if all([ent in entity_rows_keys for ent in feature_set_entity_list_names]):
+                # add to the link list feature set according to indexes match,
+                # only if all entities in the feature set exist in the entity rows
+                keys = feature_set_entity_list_names
+                entity_rows_relation.add_last(
+                    BaseMerger._Node(
+                        fs_name,
+                        data={
+                            "left_keys": keys,
+                            "right_keys": keys,
+                            "save_cols": [],
+                            "save_index": keys,
+                        },
+                        order=fs_order,
+                    )
+                )
+                entity_rows_relation.head.data["save_index"] = keys
+
+        if entity_rows_keys is not None:
+            entity_rows_linked_relation = _create_relation(
+                self._entity_rows_node_name, -1
+            )
+            relation_linked_lists.append(entity_rows_linked_relation)
+            linked_list_len_goal = len(feature_set_objects) + 1
+        else:
+            entity_rows_linked_relation = None
+            linked_list_len_goal = len(feature_set_objects)
+
         for i, name in enumerate(feature_set_names):
             linked_relation = _create_relation(name, i)
+            if entity_rows_linked_relation is not None:
+                _build_entity_rows_relation(entity_rows_linked_relation, name, i)
             for j, name_in in enumerate(feature_set_names):
                 if name != name_in:
                     linked_relation = _build_relation(name_in, j, linked_relation, i)
             relation_linked_lists.append(linked_relation)
 
         # concat all the link lists to one, for the merging process
-        link_list_iter = iter(relation_linked_lists)
-        return_relation = next(link_list_iter)
-        for relation_list in link_list_iter:
-            return_relation.concat(relation_list)
-        if return_relation.len != len(feature_set_objects):
-            raise mlrun.errors.MLRunRuntimeError("Failed to merge")
+        for i in range(len(relation_linked_lists)):
+            return_relation = relation_linked_lists[i].__copy__()
+            for relation_list in relation_linked_lists:
+                return_relation.concat(relation_list)
+            if return_relation.len == linked_list_len_goal:
+                return return_relation
 
-        return return_relation
+        raise mlrun.errors.MLRunRuntimeError("Failed to merge")
 
-    @classmethod
     def get_default_image(cls, kind):
         return mlrun.mlconf.feature_store.default_job_image
 
     def _reset_index(self, _result_df):
         raise NotImplementedError
+
+    def _create_engine_env(self):
+        """
+        initialize engine env if needed
+        """
+        raise NotImplementedError
+
+    def _get_engine_df(
+        self,
+        feature_set: FeatureSet,
+        feature_set_name: typing.List[str],
+        column_names: typing.List[str] = None,
+        start_time: typing.Union[str, datetime] = None,
+        end_time: typing.Union[str, datetime] = None,
+        time_column: typing.Optional[str] = None,
+    ):
+        """
+        Return the feature_set data frame according to the args
+
+        :param feature_set:             current feature_set to extract from the data frame
+        :param feature_set_name:        the name of the current feature_set
+        :param column_names:            list of columns to select (if not all)
+        :param start_time:              filter by start time
+        :param end_time:                filter by end time
+        :param time_column:             specify the time column name to filter on
+
+        :return: Data frame of the current engine
+        """
+        raise NotImplementedError
+
+    def _rename_columns_and_select(
+        self,
+        df,
+        rename_col_dict: typing.Dict[str, str],
+        columns: typing.List[str] = None,
+    ):
+        """
+        rename the columns of the df according to rename_col_dict, and select only `columns` if it is not none
+
+        :param df:              the data frame to change
+        :param rename_col_dict: the renaming dictionary - {<current_column_name>: <new_column_name>, ...}
+        :param columns:         list of columns to select (if not all)
+
+        :return: the data frame after the transformation or None if the transformation were preformed inplace
+        """
+        raise NotImplementedError
+
+    def _drop_columns_from_result(self):
+        """
+        drop `self._drop_columns` from `self._result_df`
+        """
+        raise NotImplementedError
+
+    def _filter(self, query: str):
+        """
+        filter `self._result_df` by `query`
+
+        :param query: The query string used to filter rows
+        """
+        raise NotImplementedError
+
+    def _order_by(self, order_by_active: typing.List[str]):
+        """
+        Order by `order_by_active` along all axis.
+
+        :param order_by_active: list of names to sort by.
+        """
+        raise NotImplementedError