mlrun 1.7.2rc3__py3-none-any.whl → 1.8.0__py3-none-any.whl

mlrun/feature_store/feature_vector.py

@@ -19,18 +19,18 @@ from datetime import datetime
 from enum import Enum
 from typing import Union
 
-import numpy as np
 import pandas as pd
 
 import mlrun
 
 from ..config import config as mlconf
 from ..datastore import get_store_uri
-from ..datastore.targets import get_offline_target
+from ..datastore.targets import BaseStoreTarget, get_offline_target
 from ..feature_store.common import (
     get_feature_set_by_uri,
     parse_feature_string,
     parse_project_name_from_feature_string,
+    verify_feature_vector_permissions,
 )
 from ..feature_store.feature_set import FeatureSet
 from ..features import Entity, Feature
@@ -47,6 +47,22 @@ from ..runtimes.function_reference import FunctionReference
 from ..serving.states import RootFlowStep
 from ..utils import StorePrefix
 from .common import RunConfig
+from .feature_vector_utils import JoinGraph, OnlineVectorService
+from .retrieval import get_merger, run_merge_job
+
+
+def _features_to_vector_and_check_permissions(features: "FeatureVector", update_stats):
+    vector = features
+    if not vector.metadata.name:
+        raise mlrun.errors.MLRunInvalidArgumentError(
+            "feature vector name must be specified"
+        )
+    verify_feature_vector_permissions(
+        vector, mlrun.common.schemas.AuthorizationAction.update
+    )
+
+    vector.save()
+    return vector
 
 
 class FeatureVectorSpec(ModelObj):
@@ -201,251 +217,6 @@ class FeatureVectorStatus(ModelObj):
         self._features = ObjectList.from_list(Feature, features)
 
 
-class JoinGraph(ModelObj):
-    """
-    explain here about the class
-    """
-
-    default_graph_name = "$__join_graph_fv__$"
-    first_join_type = "first"
-    _dict_fields = ["name", "first_feature_set", "steps"]
-
-    def __init__(
-        self,
-        name: str = None,
-        first_feature_set: Union[str, FeatureSet] = None,
-    ):
-        """
-        JoinGraph is a class that represents a graph of data joins between feature sets. It allows users to define
-        data joins step by step, specifying the join type for each step. The graph can be used to build a sequence of
-        joins that will be executed in order, allowing the creation of complex join operations between feature sets.
-
-
-        Example:
-        # Create a new JoinGraph and add steps for joining feature sets.
-        join_graph = JoinGraph(name="my_join_graph", first_feature_set="featureset1")
-        join_graph.inner("featureset2")
-        join_graph.left("featureset3", asof_join=True)
-
-
-        :param name:                (str, optional) The name of the join graph. If not provided,
-                                    a default name will be used.
-        :param first_feature_set:   (str or FeatureSet, optional) The first feature set to join. It can be
-                                    specified either as a string representing the name of the feature set or as a
-                                    FeatureSet object.
-        """
-        self.name = name or self.default_graph_name
-        self._steps: ObjectList = None
-        self._feature_sets = None
-        if first_feature_set:
-            self._start(first_feature_set)
-
-    def inner(self, other_operand: typing.Union[str, FeatureSet]):
-        """
-        Specifies an inner join with the given feature set
-
-        :param other_operand:       (str or FeatureSet) The name of the feature set or a FeatureSet object to join with.
-
-        :return:                    JoinGraph: The updated JoinGraph object with the specified inner join.
-        """
-        return self._join_operands(other_operand, "inner")
-
-    def outer(self, other_operand: typing.Union[str, FeatureSet]):
-        """
-        Specifies an outer join with the given feature set
-
-        :param other_operand:       (str or FeatureSet) The name of the feature set or a FeatureSet object to join with.
-        :return:                    JoinGraph: The updated JoinGraph object with the specified outer join.
-        """
-        return self._join_operands(other_operand, "outer")
-
-    def left(self, other_operand: typing.Union[str, FeatureSet], asof_join):
-        """
-        Specifies a left join with the given feature set
-
-        :param other_operand:       (str or FeatureSet) The name of the feature set or a FeatureSet object to join with.
-        :param asof_join:           (bool) A flag indicating whether to perform an as-of join.
-
-        :return:                    JoinGraph: The updated JoinGraph object with the specified left join.
-        """
-        return self._join_operands(other_operand, "left", asof_join=asof_join)
-
-    def right(self, other_operand: typing.Union[str, FeatureSet]):
-        """
-        Specifies a right join with the given feature set
-
-        :param other_operand:       (str or FeatureSet) The name of the feature set or a FeatureSet object to join with.
-
-        :return:                    JoinGraph: The updated JoinGraph object with the specified right join.
-        """
-        return self._join_operands(other_operand, "right")
-
-    def _join_operands(
-        self,
-        other_operand: typing.Union[str, FeatureSet],
-        join_type: str,
-        asof_join: bool = False,
-    ):
-        if isinstance(other_operand, FeatureSet):
-            other_operand = other_operand.metadata.name
-
-        first_key_num = len(self._steps.keys()) if self._steps else 0
-        left_last_step_name, left_all_feature_sets = (
-            self.last_step_name,
-            self.all_feature_sets_names,
-        )
-        is_first_fs = (
-            join_type == JoinGraph.first_join_type or left_all_feature_sets == self.name
-        )
-        # create_new_step
-        new_step = _JoinStep(
-            f"step_{first_key_num}",
-            left_last_step_name if not is_first_fs else "",
-            other_operand,
-            left_all_feature_sets if not is_first_fs else [],
-            other_operand,
-            join_type,
-            asof_join,
-        )
-
-        if self.steps is not None:
-            self.steps.update(new_step)
-        else:
-            self.steps = [new_step]
-        return self
-
-    def _start(self, other_operand: typing.Union[str, FeatureSet]):
-        return self._join_operands(other_operand, JoinGraph.first_join_type)
-
-    def _init_all_join_keys(
-        self, feature_set_objects, vector, entity_rows_keys: list[str] = None
-    ):
-        for step in self.steps:
-            step.init_join_keys(feature_set_objects, vector, entity_rows_keys)
-
-    @property
-    def all_feature_sets_names(self):
-        """
-         Returns a list of all feature set names included in the join graph.
-
-        :return:                    List[str]: A list of feature set names.
-        """
-        if self._steps:
-            return self._steps[-1].left_feature_set_names + [
-                self._steps[-1].right_feature_set_name
-            ]
-        else:
-            return self.name
-
-    @property
-    def last_step_name(self):
-        """
-        Returns the name of the last step in the join graph.
-
-        :return:                    str: The name of the last step.
-        """
-        if self._steps:
-            return self._steps[-1].name
-        else:
-            return self.name
-
-    @property
-    def steps(self):
-        """
-        Returns the list of join steps as ObjectList, which can be used to iterate over the steps
-        or access the properties of each step.
-        :return:                    ObjectList: The list of join steps.
-        """
-        return self._steps
-
-    @steps.setter
-    def steps(self, steps):
-        """
-         Setter for the steps property. It allows updating the join steps.
-
-        :param steps:               (List[_JoinStep]) The list of join steps.
-        """
-        self._steps = ObjectList.from_list(child_class=_JoinStep, children=steps)
-
-
-class _JoinStep(ModelObj):
-    def __init__(
-        self,
-        name: str = None,
-        left_step_name: str = None,
-        right_step_name: str = None,
-        left_feature_set_names: Union[str, list[str]] = None,
-        right_feature_set_name: str = None,
-        join_type: str = "inner",
-        asof_join: bool = False,
-    ):
-        self.name = name
-        self.left_step_name = left_step_name
-        self.right_step_name = right_step_name
-        self.left_feature_set_names = (
-            left_feature_set_names
-            if left_feature_set_names is None
-            or isinstance(left_feature_set_names, list)
-            else [left_feature_set_names]
-        )
-        self.right_feature_set_name = right_feature_set_name
-        self.join_type = join_type
-        self.asof_join = asof_join
-
-        self.left_keys = []
-        self.right_keys = []
-
-    def init_join_keys(
-        self,
-        feature_set_objects: ObjectList,
-        vector,
-        entity_rows_keys: list[str] = None,
-    ):
-        if feature_set_objects[self.right_feature_set_name].is_connectable_to_df(
-            entity_rows_keys
-        ):
-            self.left_keys, self.right_keys = [
-                list(
-                    feature_set_objects[
-                        self.right_feature_set_name
-                    ].spec.entities.keys()
-                )
-            ] * 2
-
-        if (
-            self.join_type == JoinGraph.first_join_type
-            or not self.left_feature_set_names
-        ):
-            self.join_type = (
-                "inner"
-                if self.join_type == JoinGraph.first_join_type
-                else self.join_type
-            )
-            return
-
-        for left_fset in self.left_feature_set_names:
-            current_left_keys = feature_set_objects[left_fset].extract_relation_keys(
-                feature_set_objects[self.right_feature_set_name],
-                vector.get_feature_set_relations(feature_set_objects[left_fset]),
-            )
-            current_right_keys = list(
-                feature_set_objects[self.right_feature_set_name].spec.entities.keys()
-            )
-            for i in range(len(current_left_keys)):
-                if (
-                    current_left_keys[i] not in self.left_keys
-                    and current_right_keys[i] not in self.right_keys
-                ):
-                    self.left_keys.append(current_left_keys[i])
-                    self.right_keys.append(current_right_keys[i])
-
-        if not self.left_keys:
-            raise mlrun.errors.MLRunRuntimeError(
-                f"{self.name} can't be preform due to undefined relation between "
-                f"{self.left_feature_set_names} to {self.right_feature_set_name}"
-            )
-
-
 class FixedWindowType(Enum):
     CurrentOpenWindow = 1
     LastClosedWindow = 2
@@ -479,7 +250,7 @@ class FeatureVector(ModelObj):
         description=None,
         with_indexes=None,
         join_graph: JoinGraph = None,
-        relations: dict[str, dict[str, Union[Entity, str]]] = None,
+        relations: typing.Optional[dict[str, dict[str, Union[Entity, str]]]] = None,
     ):
         """Feature vector, specify selected features, their metadata and material views
 
@@ -727,21 +498,21 @@ class FeatureVector(ModelObj):
     def get_offline_features(
         self,
         entity_rows=None,
-        entity_timestamp_column: str = None,
+        entity_timestamp_column: typing.Optional[str] = None,
         target: DataTargetBase = None,
         run_config: RunConfig = None,
-        drop_columns: list[str] = None,
-        start_time: Union[str, datetime] = None,
-        end_time: Union[str, datetime] = None,
+        drop_columns: typing.Optional[list[str]] = None,
+        start_time: typing.Optional[Union[str, datetime]] = None,
+        end_time: typing.Optional[Union[str, datetime]] = None,
         with_indexes: bool = False,
-        update_stats: bool = False,
-        engine: str = None,
-        engine_args: dict = None,
-        query: str = None,
-        order_by: Union[str, list[str]] = None,
-        spark_service: str = None,
-        timestamp_for_filtering: Union[str, dict[str, str]] = None,
-        additional_filters: list = None,
+        update_stats: bool = True,
+        engine: typing.Optional[str] = None,
+        engine_args: typing.Optional[dict] = None,
+        query: typing.Optional[str] = None,
+        order_by: typing.Optional[Union[str, list[str]]] = None,
+        spark_service: typing.Optional[str] = None,
+        timestamp_for_filtering: typing.Optional[Union[str, dict[str, str]]] = None,
+        additional_filters: typing.Optional[list] = None,
     ):
         """retrieve offline feature vector results
 
@@ -784,8 +555,9 @@ class FeatureVector(ModelObj):
                                         columns. This property can be specified also in the feature vector spec
                                         (feature_vector.spec.with_indexes)
                                         (default False)
-        :param update_stats:            update features statistics from the requested feature sets on the vector.
-                                        (default False).
+        :param update_stats:            When set to True (default), updates feature statistics from the requested
+                                        feature sets on the vector, which requires 'update' permissions. When set to
+                                        False, uses read-only operations that only require 'read' permissions.
         :param engine:                  processing engine kind ("local", "dask", or "spark")
         :param engine_args:             kwargs for the processing engine
         :param query:                   The query string used to filter rows on the output
@@ -806,35 +578,79 @@ class FeatureVector(ModelObj):
                             https://arrow.apache.org/docs/python/generated/pyarrow.parquet.ParquetDataset.html
 
         """
+        if entity_rows is None and entity_timestamp_column is not None:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "entity_timestamp_column param "
+                "can not be specified without entity_rows param"
+            )
+
+        if isinstance(target, BaseStoreTarget) and not target.support_pandas:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"get_offline_features does not support targets that do not support pandas engine."
+                f" Target kind: {target.kind}"
+            )
+
+        if update_stats:
+            feature_vector = _features_to_vector_and_check_permissions(
+                self, update_stats
+            )
+        else:
+            feature_vector = self
+            verify_feature_vector_permissions(
+                feature_vector, mlrun.common.schemas.AuthorizationAction.read
+            )
+
+        entity_timestamp_column = (
+            entity_timestamp_column or feature_vector.spec.timestamp_field
+        )
+
+        merger_engine = get_merger(engine)
+
+        if run_config and not run_config.local:
+            return run_merge_job(
+                feature_vector,
+                target,
+                merger_engine,
+                engine,
+                engine_args,
+                spark_service,
+                entity_rows,
+                entity_timestamp_column=entity_timestamp_column,
+                run_config=run_config,
+                drop_columns=drop_columns,
+                with_indexes=with_indexes,
+                query=query,
+                order_by=order_by,
+                start_time=start_time,
+                end_time=end_time,
+                timestamp_for_filtering=timestamp_for_filtering,
+                additional_filters=additional_filters,
+            )
 
-        return mlrun.feature_store.api._get_offline_features(
-            self,
+        merger = merger_engine(feature_vector, **(engine_args or {}))
+        return merger.start(
             entity_rows,
             entity_timestamp_column,
-            target,
-            run_config,
-            drop_columns,
-            start_time,
-            end_time,
-            with_indexes,
-            update_stats,
-            engine,
-            engine_args,
-            query,
-            order_by,
-            spark_service,
-            timestamp_for_filtering,
-            additional_filters,
+            target=target,
+            drop_columns=drop_columns,
+            start_time=start_time,
+            end_time=end_time,
+            timestamp_for_filtering=timestamp_for_filtering,
+            with_indexes=with_indexes,
+            update_stats=update_stats,
+            query=query,
+            order_by=order_by,
+            additional_filters=additional_filters,
         )
 
     def get_online_feature_service(
         self,
         run_config: RunConfig = None,
         fixed_window_type: FixedWindowType = FixedWindowType.LastClosedWindow,
-        impute_policy: dict = None,
+        impute_policy: typing.Optional[dict] = None,
         update_stats: bool = False,
-        entity_keys: list[str] = None,
-    ):
+        entity_keys: typing.Optional[list[str]] = None,
+    ) -> OnlineVectorService:
         """initialize and return online feature vector service api,
         returns :py:class:`~mlrun.feature_store.OnlineVectorService`
 
@@ -897,204 +713,14 @@ class FeatureVector(ModelObj):
         :return:                    Initialize the `OnlineVectorService`.
                                     Will be used in subclasses where `support_online=True`.
         """
-        return mlrun.feature_store.api._get_online_feature_service(
-            self,
-            run_config,
-            fixed_window_type,
-            impute_policy,
-            update_stats,
-            entity_keys,
-        )
-
-
-class OnlineVectorService:
-    """get_online_feature_service response object"""
-
-    def __init__(
-        self,
-        vector,
-        graph,
-        index_columns,
-        impute_policy: dict = None,
-        requested_columns: list[str] = None,
-    ):
-        self.vector = vector
-        self.impute_policy = impute_policy or {}
-
-        self._controller = graph.controller
-        self._index_columns = index_columns
-        self._impute_values = {}
-        self._requested_columns = requested_columns
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
-
-    def initialize(self):
-        """internal, init the feature service and prep the imputing logic"""
-        if not self.impute_policy:
-            return
-
-        impute_policy = copy(self.impute_policy)
-        vector = self.vector
-        feature_stats = vector.get_stats_table()
-        self._impute_values = {}
-
-        feature_keys = list(vector.status.features.keys())
-        if vector.status.label_column in feature_keys:
-            feature_keys.remove(vector.status.label_column)
-
-        if "*" in impute_policy:
-            value = impute_policy["*"]
-            del impute_policy["*"]
-
-            for name in feature_keys:
-                if name not in impute_policy:
-                    if isinstance(value, str) and value.startswith("$"):
-                        self._impute_values[name] = feature_stats.loc[name, value[1:]]
-                    else:
-                        self._impute_values[name] = value
-
-        for name, value in impute_policy.items():
-            if name not in feature_keys:
-                raise mlrun.errors.MLRunInvalidArgumentError(
-                    f"feature {name} in impute_policy but not in feature vector"
-                )
-            if isinstance(value, str) and value.startswith("$"):
-                self._impute_values[name] = feature_stats.loc[name, value[1:]]
-            else:
-                self._impute_values[name] = value
-
-    @property
-    def status(self):
-        """vector merger function status (ready, running, error)"""
-        return "ready"
-
-    def get(self, entity_rows: list[Union[dict, list]], as_list=False):
-        """get feature vector given the provided entity inputs
-
-        take a list of input vectors/rows and return a list of enriched feature vectors
-        each input and/or output vector can be a list of values or a dictionary of field names and values,
-        to return the vector as a list of values set the `as_list` to True.
-
-        if the input is a list of list (vs a list of dict), the values in the list will correspond to the
-        index/entity values, i.e. [["GOOG"], ["MSFT"]] means "GOOG" and "MSFT" are the index/entity fields.
-
-        example::
-
-            # accept list of dict, return list of dict
-            svc = fstore.get_online_feature_service(vector)
-            resp = svc.get([{"name": "joe"}, {"name": "mike"}])
-
-            # accept list of list, return list of list
-            svc = fstore.get_online_feature_service(vector, as_list=True)
-            resp = svc.get([["joe"], ["mike"]])
-
-        :param entity_rows:  list of list/dict with input entity data/rows
-        :param as_list:      return a list of list (list input is required by many ML frameworks)
-        """
-        results = []
-        futures = []
-        if isinstance(entity_rows, dict):
-            entity_rows = [entity_rows]
-
-        # validate we have valid input struct
-        if (
-            not entity_rows
-            or not isinstance(entity_rows, list)
-            or not isinstance(entity_rows[0], (list, dict))
-        ):
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                f"input data is of type {type(entity_rows)}. must be a list of lists or list of dicts"
-            )
+        feature_vector = _features_to_vector_and_check_permissions(self, True)
 
-        # if list of list, convert to dicts (with the index columns as the dict keys)
-        if isinstance(entity_rows[0], list):
-            if not self._index_columns or len(entity_rows[0]) != len(
-                self._index_columns
-            ):
-                raise mlrun.errors.MLRunInvalidArgumentError(
-                    "input list must be in the same size of the index_keys list"
-                )
-            index_range = range(len(self._index_columns))
-            entity_rows = [
-                {self._index_columns[i]: item[i] for i in index_range}
-                for item in entity_rows
-            ]
-
-        for row in entity_rows:
-            futures.append(self._controller.emit(row, return_awaitable_result=True))
-
-        for future in futures:
-            result = future.await_result()
-            data = result.body
-            if data:
-                actual_columns = data.keys()
-                if all([col in self._index_columns for col in actual_columns]):
-                    # didn't get any data from the graph
-                    results.append(None)
-                    continue
-                for column in self._requested_columns:
-                    if (
-                        column not in actual_columns
-                        and column != self.vector.status.label_column
-                    ):
-                        data[column] = None
-
-                if self._impute_values:
-                    for name in data.keys():
-                        v = data[name]
-                        if v is None or (
-                            isinstance(v, float) and (np.isinf(v) or np.isnan(v))
-                        ):
-                            data[name] = self._impute_values.get(name, v)
-                if not self.vector.spec.with_indexes:
-                    for name in self.vector.status.index_keys:
-                        data.pop(name, None)
-                if not any(data.values()):
-                    data = None
-
-            if as_list and data:
-                data = [
-                    data.get(key, None)
-                    for key in self._requested_columns
-                    if key != self.vector.status.label_column
-                ]
-            results.append(data)
-
-        return results
-
-    def close(self):
-        """terminate the async loop"""
-        self._controller.terminate()
-
-
-class OfflineVectorResponse:
-    """get_offline_features response object"""
-
-    def __init__(self, merger):
-        self._merger = merger
-        self.vector = merger.vector
+        engine_args = {"impute_policy": impute_policy}
+        merger_engine = get_merger("storey")
+        # todo: support remote service (using remote nuclio/mlrun function if run_config)
 
-    @property
-    def status(self):
-        """vector prep job status (ready, running, error)"""
-        return self._merger.get_status()
-
-    def to_dataframe(self, to_pandas=True):
-        """return result as dataframe"""
-        if self.status != "completed":
-            raise mlrun.errors.MLRunTaskNotReadyError(
-                "feature vector dataset is not ready"
-            )
-        return self._merger.get_df(to_pandas=to_pandas)
+        merger = merger_engine(feature_vector, **engine_args)
 
-    def to_parquet(self, target_path, **kw):
-        """return results as parquet file"""
-        return self._merger.to_parquet(target_path, **kw)
-
-    def to_csv(self, target_path, **kw):
-        """return results as csv file"""
-        return self._merger.to_csv(target_path, **kw)
+        return merger.init_online_vector_service(
+            entity_keys, fixed_window_type, update_stats=True
+        )