mlrun 1.3.2rc1__py3-none-any.whl → 1.3.2rc2__py3-none-any.whl

mlrun/feature_store/retrieval/dask_merger.py

@@ -20,6 +20,7 @@ from dask.distributed import Client
 
 import mlrun
 
+from ..feature_vector import OfflineVectorResponse
 from .base import BaseMerger
 
 
@@ -31,6 +32,139 @@ class DaskFeatureMerger(BaseMerger):
         self.client = engine_args.get("dask_client")
         self._dask_cluster_uri = engine_args.get("dask_cluster_uri")
 
+    def _generate_vector(
+        self,
+        entity_rows,
+        entity_timestamp_column,
+        feature_set_objects,
+        feature_set_fields,
+        start_time=None,
+        end_time=None,
+        query=None,
+    ):
+        if "index" not in self._index_columns:
+            self._append_drop_column("index")
+
+        # init the dask client if needed
+        if not self.client:
+            if self._dask_cluster_uri:
+                function = mlrun.import_function(self._dask_cluster_uri)
+                self.client = function.client
+            else:
+                self.client = Client()
+
+        # load dataframes
+        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
+        all_columns = []
+
+        fs_link_list = self._create_linked_relation_list(
+            feature_set_objects, feature_set_fields
+        )
+
+        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]
+            column_names = [name for name, alias in columns]
+
+            for col in node.data["save_cols"]:
+                if col not in column_names:
+                    self._append_drop_column(col)
+            column_names += node.data["save_cols"]
+
+            df = feature_set.to_dataframe(
+                columns=column_names,
+                df_module=dd,
+                start_time=start_time,
+                end_time=end_time,
+                time_column=entity_timestamp_column,
+                index=False,
+            )
+
+            df = df.reset_index()
+            column_names += node.data["save_index"]
+            node.data["save_cols"] += node.data["save_index"]
+            entity_timestamp_column_list = (
+                [entity_timestamp_column]
+                if entity_timestamp_column
+                else feature_set.spec.timestamp_key
+            )
+            if entity_timestamp_column_list:
+                column_names += entity_timestamp_column_list
+                node.data["save_cols"] += entity_timestamp_column_list
+
+            df = df.persist()
+
+            # rename columns to be unique for each feature set
+            rename_col_dict = {
+                col: f"{col}_{name}"
+                for col in column_names
+                if col not in node.data["save_cols"]
+            }
+            df = df.rename(
+                columns=rename_col_dict,
+            )
+
+            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 = []
+            for col, alias in columns:
+                if col in rename_col_dict and alias:
+                    new_columns.append((rename_col_dict[col], alias))
+                elif col in rename_col_dict and not alias:
+                    new_columns.append((rename_col_dict[col], col))
+                else:
+                    new_columns.append((col, alias))
+            all_columns.append(new_columns)
+            self._update_alias(
+                dictionary={name: alias for name, alias in new_columns if alias}
+            )
+
+        self.merge(
+            entity_df=entity_rows,
+            entity_timestamp_column=entity_timestamp_column,
+            featuresets=feature_sets,
+            featureset_dfs=dfs,
+            keys=keys,
+            all_columns=all_columns,
+        )
+
+        self._result_df = self._result_df.drop(
+            columns=self._drop_columns, errors="ignore"
+        )
+
+        # renaming all columns according to self._alias
+        self._result_df = self._result_df.rename(
+            columns=self._alias,
+        )
+
+        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._result_df = self._result_df.query(query)
+
+        if self._drop_indexes:
+            self._result_df = self._reset_index(self._result_df)
+        else:
+            self._result_df = self._set_indexes(self._result_df)
+        self._write_to_target()
+
+        return OfflineVectorResponse(self)
+
     def _reset_index(self, df):
         to_drop = df.index.name is None
         df = df.reset_index(drop=to_drop)
@@ -44,13 +178,27 @@ class DaskFeatureMerger(BaseMerger):
         featureset_df,
         left_keys: list,
         right_keys: list,
+        columns: list,
     ):
 
+        entity_df = self._reset_index(entity_df)
+        entity_df = (
+            entity_df
+            if entity_timestamp_column not in entity_df
+            else entity_df.set_index(entity_timestamp_column, drop=True)
+        )
+        featureset_df = self._reset_index(featureset_df)
+        featureset_df = (
+            featureset_df
+            if entity_timestamp_column not in featureset_df
+            else featureset_df.set_index(entity_timestamp_column, drop=True)
+        )
+
         merged_df = merge_asof(
             entity_df,
             featureset_df,
-            left_on=entity_timestamp_column,
-            right_on=entity_timestamp_column,
+            left_index=True,
+            right_index=True,
             left_by=left_keys or None,
             right_by=right_keys or None,
             suffixes=("", f"_{featureset.metadata.name}_"),
@@ -69,6 +217,7 @@ class DaskFeatureMerger(BaseMerger):
         featureset_df,
         left_keys: list,
         right_keys: list,
+        columns: list,
     ):
 
         fs_name = featureset.metadata.name
@@ -92,56 +241,5 @@ class DaskFeatureMerger(BaseMerger):
 
     def get_df(self, to_pandas=True):
         if to_pandas and hasattr(self._result_df, "dask"):
-            df = self._result_df.compute()
-        else:
-            df = self._result_df
-        self._set_indexes(df)
-        return df
-
-    def _create_engine_env(self):
-        if "index" not in self._index_columns:
-            self._append_drop_column("index")
-
-        # init the dask client if needed
-        if not self.client:
-            if self._dask_cluster_uri:
-                function = mlrun.import_function(self._dask_cluster_uri)
-                self.client = function.client
-            else:
-                self.client = Client()
-
-    def _get_engine_df(
-        self,
-        feature_set,
-        feature_set_name,
-        column_names=None,
-        start_time=None,
-        end_time=None,
-        entity_timestamp_column=None,
-    ):
-        df = feature_set.to_dataframe(
-            columns=column_names,
-            df_module=dd,
-            start_time=start_time,
-            end_time=end_time,
-            time_column=entity_timestamp_column,
-            index=False,
-        )
-
-        return self._reset_index(df).persist()
-
-    def _rename_columns_and_select(self, df, rename_col_dict, columns=None):
-        return df.rename(
-            columns=rename_col_dict,
-        )
-
-    def _drop_columns_from_result(self):
-        self._result_df = self._result_df.drop(
-            columns=self._drop_columns, errors="ignore"
-        )
-
-    def _filter(self, query):
-        self._result_df = self._result_df.query(query)
-
-    def _order_by(self, order_by_active):
-        self._result_df.sort_values(by=order_by_active)
+            return self._result_df.compute()
+        return self._result_df

mlrun/feature_store/retrieval/job.py

@@ -39,7 +39,6 @@ def run_merge_job(
     with_indexes=None,
     query=None,
     join_type="inner",
-    order_by=None,
 ):
     name = vector.metadata.name
     if not target or not hasattr(target, "to_dict"):
@@ -104,7 +103,6 @@ def run_merge_job(
             "with_indexes": with_indexes,
             "query": query,
             "join_type": join_type,
-            "order_by": order_by,
             "engine_args": engine_args,
         },
         inputs={"entity_rows": entity_rows},
@@ -149,18 +147,12 @@ class RemoteVectorResponse:
         :param df_module: optional, py module used to create the DataFrame (e.g. pd, dd, cudf, ..)
         :param kwargs:    extended DataItem.as_df() args
         """
-
         file_format = kwargs.get("format")
         if not file_format:
             file_format = self.run.status.results["target"]["kind"]
-        df = mlrun.get_dataitem(self.target_uri).as_df(
+        return mlrun.get_dataitem(self.target_uri).as_df(
             columns=columns, df_module=df_module, format=file_format, **kwargs
         )
-        if self.vector.spec.with_indexes:
-            df.set_index(
-                list(self.vector.spec.entity_fields.keys()), inplace=True, drop=True
-            )
-        return df
 
     @property
     def target_uri(self):
@@ -174,8 +166,7 @@ import mlrun
 import mlrun.feature_store.retrieval
 from mlrun.datastore.targets import get_target_driver
 def merge_handler(context, vector_uri, target, entity_rows=None, 
-                  timestamp_column=None, drop_columns=None, with_indexes=None, query=None, join_type='inner', 
-                  engine_args=None, order_by=None):
+                  timestamp_column=None, drop_columns=None, with_indexes=None, query=None, join_type='inner', engine_args=None):
     vector = context.get_store_resource(vector_uri)
     store_target = get_target_driver(target, vector)
     entity_timestamp_column = timestamp_column or vector.spec.timestamp_field
@@ -185,7 +176,7 @@ def merge_handler(context, vector_uri, target, entity_rows=None,
     context.logger.info(f"starting vector merge task to {vector.uri}")
     merger = mlrun.feature_store.retrieval.{{{engine}}}(vector, **(engine_args or {}))
     merger.start(entity_rows, entity_timestamp_column, store_target, drop_columns, with_indexes=with_indexes, 
-                 query=query, join_type=join_type, order_by=order_by)
+                 query=query, join_type=join_type)
 
     target = vector.status.targets[store_target.name].to_dict()
     context.log_result('feature_vector', vector.uri)

mlrun/feature_store/retrieval/local_merger.py

@@ -16,6 +16,7 @@ import re
 
 import pandas as pd
 
+from ..feature_vector import OfflineVectorResponse
 from .base import BaseMerger
 
 
@@ -25,6 +26,133 @@ class LocalFeatureMerger(BaseMerger):
     def __init__(self, vector, **engine_args):
         super().__init__(vector, **engine_args)
 
+    def _generate_vector(
+        self,
+        entity_rows,
+        entity_timestamp_column,
+        feature_set_objects,
+        feature_set_fields,
+        start_time=None,
+        end_time=None,
+        query=None,
+    ):
+
+        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
+        all_columns = []
+
+        fs_link_list = self._create_linked_relation_list(
+            feature_set_objects, feature_set_fields
+        )
+
+        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]
+            column_names = [name for name, alias in columns]
+
+            for col in node.data["save_cols"]:
+                if col not in column_names:
+                    self._append_drop_column(col)
+            column_names += node.data["save_cols"]
+
+            # 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)
+            if (
+                entity_timestamp_column in column_names
+                or feature_set.spec.timestamp_key == entity_timestamp_column
+            ):
+                df = feature_set.to_dataframe(
+                    columns=column_names,
+                    start_time=start_time,
+                    end_time=end_time,
+                    time_column=entity_timestamp_column,
+                )
+            else:
+                df = feature_set.to_dataframe(
+                    columns=column_names,
+                    time_column=entity_timestamp_column,
+                )
+            if df.index.names[0]:
+                df.reset_index(inplace=True)
+            column_names += node.data["save_index"]
+            node.data["save_cols"] += node.data["save_index"]
+            entity_timestamp_column_list = (
+                [entity_timestamp_column]
+                if entity_timestamp_column
+                else feature_set.spec.timestamp_key
+            )
+            if entity_timestamp_column_list:
+                column_names += entity_timestamp_column_list
+                node.data["save_cols"] += entity_timestamp_column_list
+            # rename columns to be unique for each feature set
+            rename_col_dict = {
+                col: f"{col}_{name}"
+                for col in column_names
+                if col not in node.data["save_cols"]
+            }
+            df.rename(
+                columns=rename_col_dict,
+                inplace=True,
+            )
+
+            dfs.append(df)
+            keys.append([node.data["left_keys"], node.data["right_keys"]])
+
+            # update alias according to the unique column name
+            new_columns = []
+            for col, alias in columns:
+                if col in rename_col_dict and alias:
+                    new_columns.append((rename_col_dict[col], alias))
+                elif col in rename_col_dict and not alias:
+                    new_columns.append((rename_col_dict[col], col))
+                else:
+                    new_columns.append((col, alias))
+            all_columns.append(new_columns)
+            self._update_alias(
+                dictionary={name: alias for name, alias in new_columns if alias}
+            )
+
+        self.merge(
+            entity_df=entity_rows,
+            entity_timestamp_column=entity_timestamp_column,
+            featuresets=feature_sets,
+            featureset_dfs=dfs,
+            keys=keys,
+            all_columns=all_columns,
+        )
+
+        self._result_df.drop(columns=self._drop_columns, inplace=True, errors="ignore")
+
+        # renaming all columns according to self._alias
+        self._result_df.rename(
+            columns=self._alias,
+            inplace=True,
+        )
+        if self.vector.status.label_column:
+            self._result_df.dropna(
+                subset=[self.vector.status.label_column],
+                inplace=True,
+            )
+        # filter joined data frame by the query param
+        if query:
+            self._result_df.query(query, inplace=True)
+
+        if self._drop_indexes:
+            self._result_df.reset_index(drop=True, inplace=True)
+        else:
+            self._set_indexes(self._result_df)
+
+        self._write_to_target()
+
+        return OfflineVectorResponse(self)
+
     def _asof_join(
         self,
         entity_df,
@@ -33,6 +161,7 @@ class LocalFeatureMerger(BaseMerger):
         featureset_df,
         left_keys: list,
         right_keys: list,
+        columns: list,
     ):
 
         indexes = None
@@ -84,6 +213,7 @@ class LocalFeatureMerger(BaseMerger):
         featureset_df,
         left_keys: list,
         right_keys: list,
+        columns: list,
     ):
         fs_name = featureset.metadata.name
         merged_df = pd.merge(
@@ -98,51 +228,3 @@ class LocalFeatureMerger(BaseMerger):
             if re.findall(f"_{fs_name}_$", col):
                 self._append_drop_column(col)
         return merged_df
-
-    def _create_engine_env(self):
-        pass
-
-    def _get_engine_df(
-        self,
-        feature_set,
-        feature_set_name,
-        column_names=None,
-        start_time=None,
-        end_time=None,
-        entity_timestamp_column=None,
-    ):
-        # 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)
-        if (
-            entity_timestamp_column in column_names
-            or feature_set.spec.timestamp_key == entity_timestamp_column
-        ):
-            df = feature_set.to_dataframe(
-                columns=column_names,
-                start_time=start_time,
-                end_time=end_time,
-                time_column=entity_timestamp_column,
-            )
-        else:
-            df = feature_set.to_dataframe(
-                columns=column_names,
-                time_column=entity_timestamp_column,
-            )
-        if df.index.names[0]:
-            df.reset_index(inplace=True)
-        return df
-
-    def _rename_columns_and_select(self, df, rename_col_dict, columns=None):
-        df.rename(
-            columns=rename_col_dict,
-            inplace=True,
-        )
-
-    def _drop_columns_from_result(self):
-        self._result_df.drop(columns=self._drop_columns, inplace=True, errors="ignore")
-
-    def _filter(self, query):
-        self._result_df.query(query, inplace=True)
-
-    def _order_by(self, order_by_active):
-        self._result_df.sort_values(by=order_by_active, ignore_index=True, inplace=True)