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

mlrun/feature_store/retrieval/dask_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.
@@ -20,151 +20,18 @@ from dask.distributed import Client
 
 import mlrun
 
-from ..feature_vector import OfflineVectorResponse
 from .base import BaseMerger
 
 
 class DaskFeatureMerger(BaseMerger):
     engine = "dask"
+    support_offline = True
 
     def __init__(self, vector, **engine_args):
         super().__init__(vector, **engine_args)
         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)
@@ -178,27 +45,22 @@ class DaskFeatureMerger(BaseMerger):
         featureset_df,
         left_keys: list,
         right_keys: list,
-        columns: list,
     ):
+        def sort_partition(partition, timestamp):
+            return partition.sort_values(timestamp)
 
-        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)
+        entity_df = entity_df.map_partitions(
+            sort_partition, timestamp=entity_timestamp_column
         )
-        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)
+        featureset_df = featureset_df.map_partitions(
+            sort_partition, timestamp=featureset.spec.timestamp_key
         )
 
         merged_df = merge_asof(
             entity_df,
             featureset_df,
-            left_index=True,
-            right_index=True,
+            left_on=entity_timestamp_column,
+            right_on=featureset.spec.timestamp_key,
             left_by=left_keys or None,
             right_by=right_keys or None,
             suffixes=("", f"_{featureset.metadata.name}_"),
@@ -217,7 +79,6 @@ class DaskFeatureMerger(BaseMerger):
         featureset_df,
         left_keys: list,
         right_keys: list,
-        columns: list,
     ):
 
         fs_name = featureset.metadata.name
@@ -241,5 +102,56 @@ class DaskFeatureMerger(BaseMerger):
 
     def get_df(self, to_pandas=True):
         if to_pandas and hasattr(self._result_df, "dask"):
-            return self._result_df.compute()
-        return self._result_df
+            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,
+        time_column=None,
+    ):
+        df = feature_set.to_dataframe(
+            columns=column_names,
+            df_module=dd,
+            start_time=start_time,
+            end_time=end_time,
+            time_column=time_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)

mlrun/feature_store/retrieval/job.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.
@@ -33,12 +33,15 @@ def run_merge_job(
     engine_args: dict,
     spark_service: str = None,
     entity_rows=None,
-    timestamp_column=None,
+    entity_timestamp_column=None,
     run_config=None,
     drop_columns=None,
     with_indexes=None,
     query=None,
-    join_type="inner",
+    order_by=None,
+    start_time=None,
+    end_time=None,
+    timestamp_for_filtering=None,
 ):
     name = vector.metadata.name
     if not target or not hasattr(target, "to_dict"):
@@ -92,20 +95,27 @@ def run_merge_job(
         set_default_resources(
             function.spec.executor_resources, function.with_executor_requests
         )
+    if start_time and not isinstance(start_time, str):
+        start_time = start_time.isoformat()
+    if end_time and not isinstance(end_time, str):
+        end_time = end_time.isoformat()
 
     task = new_task(
         name=name,
         params={
             "vector_uri": vector.uri,
             "target": target.to_dict(),
-            "timestamp_column": timestamp_column,
+            "entity_timestamp_column": entity_timestamp_column,
             "drop_columns": drop_columns,
             "with_indexes": with_indexes,
             "query": query,
-            "join_type": join_type,
+            "order_by": order_by,
+            "start_time": start_time,
+            "end_time": end_time,
+            "timestamp_for_filtering": timestamp_for_filtering,
             "engine_args": engine_args,
         },
-        inputs={"entity_rows": entity_rows},
+        inputs={"entity_rows": entity_rows} if entity_rows is not None else {},
     )
     task.spec.secret_sources = run_config.secret_sources
     task.set_label("job-type", "feature-merge").set_label("feature-vector", vector.uri)
@@ -120,15 +130,16 @@ def run_merge_job(
         watch=run_config.watch,
     )
     logger.info(f"feature vector merge job started, run id = {run.uid()}")
-    return RemoteVectorResponse(vector, run)
+    return RemoteVectorResponse(vector, run, with_indexes)
 
 
 class RemoteVectorResponse:
     """get_offline_features response object"""
 
-    def __init__(self, vector, run):
+    def __init__(self, vector, run, with_indexes=False):
         self.run = run
         self.vector = vector
+        self.with_indexes = with_indexes or self.vector.spec.with_indexes
 
     @property
     def status(self):
@@ -147,12 +158,18 @@ 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"]
-        return mlrun.get_dataitem(self.target_uri).as_df(
+        df = mlrun.get_dataitem(self.target_uri).as_df(
             columns=columns, df_module=df_module, format=file_format, **kwargs
         )
+        if self.with_indexes:
+            df.set_index(
+                list(self.vector.spec.entity_fields.keys()), inplace=True, drop=True
+            )
+        return df
 
     @property
     def target_uri(self):
@@ -166,17 +183,18 @@ 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):
+                  entity_timestamp_column=None, drop_columns=None, with_indexes=None, query=None,
+                  engine_args=None, order_by=None, start_time=None, end_time=None, timestamp_for_filtering=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
     if entity_rows:
         entity_rows = entity_rows.as_df()
 
     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)
+                 query=query, order_by=order_by, start_time=start_time, end_time=end_time,
+                 timestamp_for_filtering=timestamp_for_filtering)
 
     target = vector.status.targets[store_target.name].to_dict()
     context.log_result('feature_vector', vector.uri)

mlrun/feature_store/retrieval/local_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.
@@ -16,143 +16,16 @@ import re
 
 import pandas as pd
 
-from ..feature_vector import OfflineVectorResponse
 from .base import BaseMerger
 
 
 class LocalFeatureMerger(BaseMerger):
     engine = "local"
+    support_offline = True
 
     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,
@@ -161,7 +34,6 @@ class LocalFeatureMerger(BaseMerger):
         featureset_df,
         left_keys: list,
         right_keys: list,
-        columns: list,
     ):
 
         indexes = None
@@ -176,7 +48,7 @@ class LocalFeatureMerger(BaseMerger):
             featureset_df[featureset.spec.timestamp_key]
         )
         entity_df.sort_values(by=entity_timestamp_column, inplace=True)
-        featureset_df.sort_values(by=entity_timestamp_column, inplace=True)
+        featureset_df.sort_values(by=featureset.spec.timestamp_key, inplace=True)
 
         merged_df = pd.merge_asof(
             entity_df,
@@ -191,7 +63,6 @@ class LocalFeatureMerger(BaseMerger):
         for col in merged_df.columns:
             if re.findall(f"_{featureset.metadata.name}_$", col):
                 self._append_drop_column(col)
-
         # Undo indexing tricks for asof merge
         # to return the correct indexes and not
         # overload `index` columns
@@ -213,7 +84,6 @@ class LocalFeatureMerger(BaseMerger):
         featureset_df,
         left_keys: list,
         right_keys: list,
-        columns: list,
     ):
         fs_name = featureset.metadata.name
         merged_df = pd.merge(
@@ -228,3 +98,40 @@ 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,
+        time_column=None,
+    ):
+        df = feature_set.to_dataframe(
+            columns=column_names,
+            start_time=start_time,
+            end_time=end_time,
+            time_column=time_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)