mlrun 1.4.0rc25__py3-none-any.whl → 1.5.0rc2__py3-none-any.whl

mlrun/feature_store/retrieval/base.py

@@ -16,13 +16,12 @@ 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 mlrun.feature_store.feature_vector import Feature, JoinGraph
 
 from ...utils import logger, str_to_timestamp
 from ..feature_vector import OfflineVectorResponse
@@ -42,6 +41,7 @@ class BaseMerger(abc.ABC):
     def __init__(self, vector, **engine_args):
         self._relation = dict()
         self._join_type = "inner"
+        self._default_join_type = "default_join"
         self.vector = vector
 
         self._result_df = None
@@ -196,21 +196,34 @@ class BaseMerger(abc.ABC):
         )  # 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
+        join_types = []
 
-        fs_link_list = self._create_linked_relation_list(
-            feature_set_objects, feature_set_fields
+        entity_rows_keys = (
+            list(entity_rows.columns) if entity_rows is not None else None
         )
+        join_graph = self._get_graph(
+            feature_set_objects, feature_set_fields, entity_rows_keys
+        )
+        if entity_rows_keys:
+            entity_rows = self._convert_entity_rows_to_engine_df(entity_rows)
+            dfs.append(entity_rows)
+            keys.append([[], []])
+            feature_sets.append(None)
+            join_types.append(None)
 
         filtered = False
-        for node in fs_link_list:
-            name = node.name
+        for step in join_graph.steps:
+            name = step.right_feature_set_name
             feature_set = feature_set_objects[name]
+            saved_columns_for_relation = list(
+                self.vector.get_feature_set_relations(feature_set).keys()
+            )
             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"]:
+            for column in saved_columns_for_relation:
                 if column not in column_names:
                     column_names.append(column)
                     if column not in self._index_columns:
@@ -247,19 +260,19 @@ class BaseMerger(abc.ABC):
                 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())
+            column_names += fs_entities_and_timestamp
+            saved_columns_for_relation += fs_entities_and_timestamp
             if feature_set.spec.timestamp_key:
                 column_names.append(feature_set.spec.timestamp_key)
-                node.data["save_cols"].append(feature_set.spec.timestamp_key)
+                saved_columns_for_relation.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"]
+                if column not in saved_columns_for_relation
             }
             df_temp = self._rename_columns_and_select(
                 df,
@@ -274,7 +287,8 @@ class BaseMerger(abc.ABC):
             dfs.append(df)
             del df
 
-            keys.append([node.data["left_keys"], node.data["right_keys"]])
+            keys.append([step.left_keys, step.right_keys])
+            join_types.append([step.join_type, step.asof_join])
 
             # update alias according to the unique column name
             new_columns = []
@@ -293,31 +307,13 @@ class BaseMerger(abc.ABC):
                 "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,
+            entity_timestamp_column=entity_timestamp_column,
             featuresets=feature_sets,
             featureset_dfs=dfs,
             keys=keys,
+            join_types=join_types,
         )
 
         all_columns = None
@@ -386,41 +382,46 @@ class BaseMerger(abc.ABC):
 
     def merge(
         self,
-        entity_df,
         entity_timestamp_column: str,
         featuresets: list,
         featureset_dfs: list,
         keys: list = None,
+        join_types: list = None,
     ):
         """join the entities and feature set features into a result dataframe"""
-        merged_df = entity_df
-        if entity_df is None and featureset_dfs:
-            merged_df = featureset_dfs.pop(0)
-            featureset = featuresets.pop(0)
-            if keys is not None:
-                keys.pop(0)
-            else:
-                # keys can be multiple keys on each side of the join
-                keys = [[[], []]] * len(featureset_dfs)
-            entity_timestamp_column = (
-                entity_timestamp_column or featureset.spec.timestamp_key
-            )
-        elif entity_df is not None and featureset_dfs:
-            # when `entity_rows` passed to `get_offline_features`
-            # keys[0] mention the way that `entity_rows`  joins to the first `featureset`
-            # 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 in zip(featuresets, featureset_dfs, keys):
-            if featureset.spec.timestamp_key and entity_timestamp_column:
+
+        merged_df = featureset_dfs.pop(0)
+        featureset = featuresets.pop(0)
+        keys.pop(0)
+        join_types.pop(0)
+
+        if not entity_timestamp_column and featureset:
+            entity_timestamp_column = featureset.spec.timestamp_key
+
+        for featureset, featureset_df, lr_key, join_type in zip(
+            featuresets, featureset_dfs, keys, join_types
+        ):
+            join_type, as_of = join_type
+            if (
+                featureset.spec.timestamp_key
+                and entity_timestamp_column
+                and join_type == self._default_join_type
+            ):
                 merge_func = self._asof_join
-            else:
+            elif join_type == self._default_join_type:
+                merge_func = self._join
+            elif join_type != self._default_join_type and not as_of:
+                self._join_type = join_type
                 merge_func = self._join
+            else:
+                self._join_type = join_type
+                merge_func = self._asof_join
 
             merged_df = merge_func(
                 merged_df,
                 entity_timestamp_column,
-                featureset,
+                featureset.metadata.name,
+                featureset.spec.timestamp_key,
                 featureset_df,
                 lr_key[0],
                 lr_key[1],
@@ -441,7 +442,8 @@ class BaseMerger(abc.ABC):
         self,
         entity_df,
         entity_timestamp_column: str,
-        featureset,
+        featureset_name: str,
+        featureset_timstamp: str,
         featureset_df,
         left_keys: list,
         right_keys: list,
@@ -452,7 +454,8 @@ class BaseMerger(abc.ABC):
         self,
         entity_df,
         entity_timestamp_column: str,
-        featureset,
+        featureset_name: str,
+        featureset_timestamp: str,
         featureset_df,
         left_keys: list,
         right_keys: list,
@@ -480,10 +483,42 @@ class BaseMerger(abc.ABC):
         size = CSVTarget(path=target_path).write_dataframe(self._result_df, **kw)
         return size
 
+    def _get_graph(
+        self, feature_set_objects, feature_set_fields, entity_rows_keys=None
+    ):
+        join_graph = self.vector.spec.join_graph
+        if not join_graph:
+            fs_link_list = self._create_linked_relation_list(
+                feature_set_objects, feature_set_fields, entity_rows_keys
+            )
+            join_graph = None
+            for i, node in enumerate(fs_link_list):
+                if node.name != self._entity_rows_node_name and join_graph is None:
+                    join_graph = JoinGraph(first_feature_set=node.name)
+                elif node.name == self._entity_rows_node_name:
+                    continue
+                else:
+                    join_graph.inner(other_operand=node.name)
+
+                last_step = join_graph.steps[-1]
+                last_step.join_type = self._default_join_type
+                last_step.left_keys = node.left_keys
+                last_step.right_keys = node.right_keys
+        else:
+            join_graph._init_all_join_keys(feature_set_objects, self.vector)
+        return join_graph
+
     class _Node:
-        def __init__(self, name: str, order: int, data=None):
+        def __init__(
+            self,
+            name: str,
+            order: int,
+            left_keys: typing.List[str] = None,
+            right_keys: typing.List[str] = None,
+        ):
             self.name = name
-            self.data = data
+            self.left_keys = left_keys if left_keys is not None else []
+            self.right_keys = right_keys if right_keys is not None else []
             # order of this feature_set in the original list
             self.order = order
             self.next = None
@@ -495,7 +530,9 @@ class BaseMerger(abc.ABC):
             return self.name == other.name
 
         def __copy__(self):
-            return BaseMerger._Node(self.name, self.order, self.data.copy())
+            return BaseMerger._Node(
+                self.name, self.order, self.left_keys, self.right_keys
+            )
 
     class _LinkedList:
         def __init__(self, head=None):
@@ -565,9 +602,6 @@ class BaseMerger(abc.ABC):
             node = self.find_node(other_head.name)
             if node is None:
                 return
-            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:
@@ -587,24 +621,24 @@ class BaseMerger(abc.ABC):
                 head=BaseMerger._Node(
                     name=feature_set_names[0],
                     order=0,
-                    data={
-                        "left_keys": [],
-                        "right_keys": [],
-                        "save_cols": [],
-                        "save_index": [],
-                    },
                 )
             )
         relation_linked_lists = []
         feature_set_entity_list_dict = {
             name: feature_set_objects[name].spec.entities for name in feature_set_names
         }
-        entity_relation_val_list = {
-            name: list(feature_set_objects[name].spec.relations.values())
+        relation_val_list = {
+            name: list(
+                self.vector.get_feature_set_relations(
+                    feature_set_objects[name]
+                ).values()
+            )
             for name in feature_set_names
         }
-        entity_relation_key_list = {
-            name: list(feature_set_objects[name].spec.relations.keys())
+        relation_key_list = {
+            name: list(
+                self.vector.get_feature_set_relations(feature_set_objects[name]).keys()
+            )
             for name in feature_set_names
         }
 
@@ -612,12 +646,6 @@ class BaseMerger(abc.ABC):
             relations = BaseMerger._LinkedList()
             main_node = BaseMerger._Node(
                 name,
-                data={
-                    "left_keys": [],
-                    "right_keys": [],
-                    "save_cols": [],
-                    "save_index": [],
-                },
                 order=order,
             )
             relations.add_first(main_node)
@@ -629,8 +657,8 @@ class BaseMerger(abc.ABC):
             name_head = linked_list_relation.head.name
             feature_set_in_entity_list = feature_set_entity_list_dict[fs_name_in]
             feature_set_in_entity_list_names = list(feature_set_in_entity_list.keys())
-            entity_relation_list = entity_relation_val_list[name_head]
-            col_relation_list = entity_relation_key_list[name_head]
+            entity_relation_list = relation_val_list[name_head]
+            col_relation_list = relation_key_list[name_head]
             curr_col_relation_list = list(
                 map(
                     lambda ent: (
@@ -649,18 +677,11 @@ class BaseMerger(abc.ABC):
                 linked_list_relation.add_last(
                     BaseMerger._Node(
                         fs_name_in,
-                        data={
-                            "left_keys": curr_col_relation_list,
-                            "right_keys": feature_set_in_entity_list_names,
-                            "save_cols": [],
-                            "save_index": [],
-                        },
+                        left_keys=curr_col_relation_list,
+                        right_keys=feature_set_in_entity_list_names,
                         order=name_in_order,
                     )
                 )
-                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
             ) == sorted(feature_set_entity_list_dict[name_head].keys()):
@@ -669,16 +690,11 @@ class BaseMerger(abc.ABC):
                 linked_list_relation.add_last(
                     BaseMerger._Node(
                         fs_name_in,
-                        data={
-                            "left_keys": keys,
-                            "right_keys": keys,
-                            "save_cols": [],
-                            "save_index": keys,
-                        },
+                        left_keys=keys,
+                        right_keys=keys,
                         order=name_in_order,
                     )
                 )
-                linked_list_relation.head.data["save_index"] = keys
             return linked_list_relation
 
         def _build_entity_rows_relation(entity_rows_relation, fs_name, fs_order):
@@ -692,16 +708,11 @@ class BaseMerger(abc.ABC):
                 entity_rows_relation.add_last(
                     BaseMerger._Node(
                         fs_name,
-                        data={
-                            "left_keys": keys,
-                            "right_keys": keys,
-                            "save_cols": [],
-                            "save_index": keys,
-                        },
+                        left_keys=keys,
+                        right_keys=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(
@@ -805,3 +816,6 @@ class BaseMerger(abc.ABC):
         :param order_by_active: list of names to sort by.
         """
         raise NotImplementedError
+
+    def _convert_entity_rows_to_engine_df(self, entity_rows):
+        raise NotImplementedError

mlrun/feature_store/retrieval/dask_merger.py

@@ -41,8 +41,9 @@ class DaskFeatureMerger(BaseMerger):
         self,
         entity_df,
         entity_timestamp_column: str,
-        featureset,
-        featureset_df,
+        featureset_name: str,
+        featureset_timestamp: str,
+        featureset_df: list,
         left_keys: list,
         right_keys: list,
     ):
@@ -53,20 +54,20 @@ class DaskFeatureMerger(BaseMerger):
             sort_partition, timestamp=entity_timestamp_column
         )
         featureset_df = featureset_df.map_partitions(
-            sort_partition, timestamp=featureset.spec.timestamp_key
+            sort_partition, timestamp=featureset_timestamp
         )
 
         merged_df = merge_asof(
             entity_df,
             featureset_df,
             left_on=entity_timestamp_column,
-            right_on=featureset.spec.timestamp_key,
+            right_on=featureset_timestamp,
             left_by=left_keys or None,
             right_by=right_keys or None,
-            suffixes=("", f"_{featureset.metadata.name}_"),
+            suffixes=("", f"_{featureset_name}_"),
         )
         for col in merged_df.columns:
-            if re.findall(f"_{featureset.metadata.name}_$", col):
+            if re.findall(f"_{featureset_name}_$", col):
                 self._append_drop_column(col)
 
         return merged_df
@@ -75,23 +76,23 @@ class DaskFeatureMerger(BaseMerger):
         self,
         entity_df,
         entity_timestamp_column: str,
-        featureset,
+        featureset_name,
+        featureset_timestamp,
         featureset_df,
         left_keys: list,
         right_keys: list,
     ):
 
-        fs_name = featureset.metadata.name
         merged_df = merge(
             entity_df,
             featureset_df,
             how=self._join_type,
             left_on=left_keys,
             right_on=right_keys,
-            suffixes=("", f"_{fs_name}_"),
+            suffixes=("", f"_{featureset_name}_"),
         )
         for col in merged_df.columns:
-            if re.findall(f"_{fs_name}_$", col):
+            if re.findall(f"_{featureset_name}_$", col):
                 self._append_drop_column(col)
         return merged_df
 
@@ -155,3 +156,9 @@ class DaskFeatureMerger(BaseMerger):
 
     def _order_by(self, order_by_active):
         self._result_df.sort_values(by=order_by_active)
+
+    def _convert_entity_rows_to_engine_df(self, entity_rows):
+        if entity_rows is not None and not hasattr(entity_rows, "dask"):
+            return dd.from_pandas(entity_rows, npartitions=len(entity_rows.columns))
+
+        return entity_rows

mlrun/feature_store/retrieval/job.py

@@ -62,9 +62,12 @@ def run_merge_job(
     function = run_config.to_function(kind, merger.get_default_image(kind))
 
     # Avoid overriding a handler that was provided by the user
-    # The user shouldn't have to provide a handler, but we leave this option open just in case
     if not run_config.handler:
         function.with_code(body=default_code)
+    else:
+        raise mlrun.errors.MLRunInvalidArgumentError(
+            "get_offline_features does not support run_config with a handler"
+        )
 
     function.metadata.project = vector.metadata.project
     function.metadata.name = function.metadata.name or name

mlrun/feature_store/retrieval/local_merger.py

@@ -30,45 +30,42 @@ class LocalFeatureMerger(BaseMerger):
         self,
         entity_df,
         entity_timestamp_column: str,
-        featureset,
-        featureset_df,
+        featureset_name,
+        featureset_timstamp,
+        featureset_df: list,
         left_keys: list,
         right_keys: list,
     ):
 
-        indexes = None
-        if not right_keys:
-            indexes = list(featureset.spec.entities.keys())
         index_col_not_in_entity = "index" not in entity_df.columns
         index_col_not_in_featureset = "index" not in featureset_df.columns
         entity_df[entity_timestamp_column] = pd.to_datetime(
             entity_df[entity_timestamp_column]
         )
-        featureset_df[featureset.spec.timestamp_key] = pd.to_datetime(
-            featureset_df[featureset.spec.timestamp_key]
+        featureset_df[featureset_timstamp] = pd.to_datetime(
+            featureset_df[featureset_timstamp]
         )
         entity_df.sort_values(by=entity_timestamp_column, inplace=True)
-        featureset_df.sort_values(by=featureset.spec.timestamp_key, inplace=True)
+        featureset_df.sort_values(by=featureset_timstamp, inplace=True)
 
         merged_df = pd.merge_asof(
             entity_df,
             featureset_df,
             left_on=entity_timestamp_column,
-            right_on=featureset.spec.timestamp_key,
-            by=indexes,
+            right_on=featureset_timstamp,
             left_by=left_keys or None,
             right_by=right_keys or None,
-            suffixes=("", f"_{featureset.metadata.name}_"),
+            suffixes=("", f"_{featureset_name}_"),
         )
         for col in merged_df.columns:
-            if re.findall(f"_{featureset.metadata.name}_$", col):
+            if re.findall(f"_{featureset_name}_$", col):
                 self._append_drop_column(col)
         # Undo indexing tricks for asof merge
         # to return the correct indexes and not
         # overload `index` columns
         if (
-            indexes
-            and "index" not in indexes
+            "index" not in left_keys
+            and "index" not in right_keys
             and index_col_not_in_entity
             and index_col_not_in_featureset
             and "index" in merged_df.columns
@@ -80,22 +77,22 @@ class LocalFeatureMerger(BaseMerger):
         self,
         entity_df,
         entity_timestamp_column: str,
-        featureset,
+        featureset_name,
+        featureset_timestamp,
         featureset_df,
         left_keys: list,
         right_keys: list,
     ):
-        fs_name = featureset.metadata.name
         merged_df = pd.merge(
             entity_df,
             featureset_df,
             how=self._join_type,
             left_on=left_keys,
             right_on=right_keys,
-            suffixes=("", f"_{fs_name}_"),
+            suffixes=("", f"_{featureset_name}_"),
         )
         for col in merged_df.columns:
-            if re.findall(f"_{fs_name}_$", col):
+            if re.findall(f"_{featureset_name}_$", col):
                 self._append_drop_column(col)
         return merged_df
 
@@ -135,3 +132,6 @@ class LocalFeatureMerger(BaseMerger):
 
     def _order_by(self, order_by_active):
         self._result_df.sort_values(by=order_by_active, ignore_index=True, inplace=True)
+
+    def _convert_entity_rows_to_engine_df(self, entity_rows):
+        return entity_rows