mlrun 1.6.0rc35__py3-none-any.whl → 1.7.0rc2__py3-none-any.whl

mlrun/feature_store/feature_vector.py

@@ -17,7 +17,7 @@ import typing
 from copy import copy
 from datetime import datetime
 from enum import Enum
-from typing import Dict, List, Union
+from typing import Union
 
 import numpy as np
 import pandas as pd
@@ -69,18 +69,16 @@ class FeatureVectorSpec(ModelObj):
         self._entity_fields: ObjectList = None
         self._entity_source: DataSource = None
         self._function: FunctionReference = None
-        self._relations: typing.Dict[str, ObjectDict] = None
+        self._relations: dict[str, ObjectDict] = None
         self._join_graph: JoinGraph = None
 
         self.description = description
-        self.features: List[str] = features or []
+        self.features: list[str] = features or []
         self.entity_source = entity_source
         self.entity_fields = entity_fields or []
         self.graph = graph
         self.join_graph = join_graph
-        self.relations: typing.Dict[str, typing.Dict[str, Union[Entity, str]]] = (
-            relations or {}
-        )
+        self.relations: dict[str, dict[str, Union[Entity, str]]] = relations or {}
         self.timestamp_field = timestamp_field
         self.label_feature = label_feature
         self.with_indexes = with_indexes
@@ -97,12 +95,12 @@ class FeatureVectorSpec(ModelObj):
         self._entity_source = self._verify_dict(source, "entity_source", DataSource)
 
     @property
-    def entity_fields(self) -> List[Feature]:
+    def entity_fields(self) -> list[Feature]:
         """the schema/metadata for the entity source fields"""
         return self._entity_fields
 
     @entity_fields.setter
-    def entity_fields(self, entity_fields: List[Feature]):
+    def entity_fields(self, entity_fields: list[Feature]):
         self._entity_fields = ObjectList.from_list(Feature, entity_fields)
 
     @property
@@ -125,14 +123,12 @@ class FeatureVectorSpec(ModelObj):
         self._function = self._verify_dict(function, "function", FunctionReference)
 
     @property
-    def relations(self) -> typing.Dict[str, ObjectDict]:
+    def relations(self) -> dict[str, ObjectDict]:
         """feature set relations dict"""
         return self._relations
 
     @relations.setter
-    def relations(
-        self, relations: typing.Dict[str, typing.Dict[str, Union[Entity, str]]]
-    ):
+    def relations(self, relations: dict[str, dict[str, Union[Entity, str]]]):
         temp_relations = {}
         for fs_name, relation in relations.items():
             for col, ent in relation.items():
@@ -179,29 +175,29 @@ class FeatureVectorStatus(ModelObj):
         self.stats = stats or {}
         self.index_keys = index_keys
         self.preview = preview or []
-        self.features: List[Feature] = features or []
+        self.features: list[Feature] = features or []
         self.run_uri = run_uri
         self.timestamp_key = timestamp_key
 
     @property
-    def targets(self) -> List[DataTarget]:
+    def targets(self) -> list[DataTarget]:
         """list of material storage targets + their status/path"""
         return self._targets
 
     @targets.setter
-    def targets(self, targets: List[DataTarget]):
+    def targets(self, targets: list[DataTarget]):
         self._targets = ObjectList.from_list(DataTarget, targets)
 
     def update_target(self, target: DataTarget):
         self._targets.update(target)
 
     @property
-    def features(self) -> List[Feature]:
+    def features(self) -> list[Feature]:
         """list of features (result of joining features from the source feature sets)"""
         return self._features
 
     @features.setter
-    def features(self, features: List[Feature]):
+    def features(self, features: list[Feature]):
         self._features = ObjectList.from_list(Feature, features)
 
 
@@ -378,7 +374,7 @@ class _JoinStep(ModelObj):
         name: str = None,
         left_step_name: str = None,
         right_step_name: str = None,
-        left_feature_set_names: Union[str, List[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,
@@ -388,7 +384,8 @@ class _JoinStep(ModelObj):
         self.right_step_name = right_step_name
         self.left_feature_set_names = (
             left_feature_set_names
-            if isinstance(left_feature_set_names, list)
+            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
@@ -402,7 +399,7 @@ class _JoinStep(ModelObj):
         self,
         feature_set_objects: ObjectList,
         vector,
-        entity_rows_keys: List[str] = None,
+        entity_rows_keys: list[str] = None,
     ):
         if feature_set_objects[self.right_feature_set_name].is_connectable_to_df(
             entity_rows_keys
@@ -482,7 +479,7 @@ class FeatureVector(ModelObj):
         description=None,
         with_indexes=None,
         join_graph: JoinGraph = None,
-        relations: typing.Dict[str, typing.Dict[str, Union[Entity, str]]] = None,
+        relations: dict[str, dict[str, Union[Entity, str]]] = None,
     ):
         """Feature vector, specify selected features, their metadata and material views
 
@@ -732,7 +729,7 @@ class FeatureVector(ModelObj):
         entity_timestamp_column: str = None,
         target: DataTargetBase = None,
         run_config: RunConfig = None,
-        drop_columns: List[str] = None,
+        drop_columns: list[str] = None,
         start_time: Union[str, datetime] = None,
         end_time: Union[str, datetime] = None,
         with_indexes: bool = False,
@@ -740,9 +737,9 @@ class FeatureVector(ModelObj):
         engine: str = None,
         engine_args: dict = None,
         query: str = None,
-        order_by: Union[str, List[str]] = None,
+        order_by: Union[str, list[str]] = None,
         spark_service: str = None,
-        timestamp_for_filtering: Union[str, Dict[str, str]] = None,
+        timestamp_for_filtering: Union[str, dict[str, str]] = None,
     ):
         """retrieve offline feature vector results
 
@@ -827,7 +824,7 @@ class FeatureVector(ModelObj):
         fixed_window_type: FixedWindowType = FixedWindowType.LastClosedWindow,
         impute_policy: dict = None,
         update_stats: bool = False,
-        entity_keys: List[str] = None,
+        entity_keys: list[str] = None,
     ):
         """initialize and return online feature vector service api,
         returns :py:class:`~mlrun.feature_store.OnlineVectorService`
@@ -910,7 +907,7 @@ class OnlineVectorService:
         graph,
         index_columns,
         impute_policy: dict = None,
-        requested_columns: List[str] = None,
+        requested_columns: list[str] = None,
     ):
         self.vector = vector
         self.impute_policy = impute_policy or {}
@@ -966,7 +963,7 @@ class OnlineVectorService:
         """vector merger function status (ready, running, error)"""
         return "ready"
 
-    def get(self, entity_rows: List[Union[dict, list]], as_list=False):
+    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

mlrun/feature_store/retrieval/base.py

@@ -540,8 +540,8 @@ class BaseMerger(abc.ABC):
             self,
             name: str,
             order: int,
-            left_keys: typing.List[str] = None,
-            right_keys: typing.List[str] = None,
+            left_keys: list[str] = None,
+            right_keys: list[str] = None,
         ):
             self.name = name
             self.left_keys = left_keys if left_keys is not None else []
@@ -750,8 +750,8 @@ class BaseMerger(abc.ABC):
     def _get_engine_df(
         self,
         feature_set: FeatureSet,
-        feature_set_name: typing.List[str],
-        column_names: typing.List[str] = None,
+        feature_set_name: list[str],
+        column_names: list[str] = None,
         start_time: typing.Union[str, datetime] = None,
         end_time: typing.Union[str, datetime] = None,
         time_column: typing.Optional[str] = None,
@@ -773,8 +773,8 @@ class BaseMerger(abc.ABC):
     def _rename_columns_and_select(
         self,
         df,
-        rename_col_dict: typing.Dict[str, str],
-        columns: typing.List[str] = None,
+        rename_col_dict: dict[str, str],
+        columns: list[str] = None,
     ):
         """
         rename the columns of the df according to rename_col_dict, and select only `columns` if it is not none
@@ -801,7 +801,7 @@ class BaseMerger(abc.ABC):
         """
         raise NotImplementedError
 
-    def _order_by(self, order_by_active: typing.List[str]):
+    def _order_by(self, order_by_active: list[str]):
         """
         Order by `order_by_active` along all axis.
 

mlrun/feature_store/retrieval/conversion.py

@@ -19,7 +19,7 @@ from collections import Counter
 # np.bool -> bool and np.object -> object fix backported from pyspark v3.3.3.
 
 
-class PandasConversionMixin(object):
+class PandasConversionMixin:
     """
     Min-in for the conversion from Spark to pandas. Currently, only :class:`DataFrame`
     can use this class.
@@ -108,9 +108,7 @@ class PandasConversionMixin(object):
                     )
 
                     # Rename columns to avoid duplicated column names.
-                    tmp_column_names = [
-                        "col_{}".format(i) for i in range(len(self.columns))
-                    ]
+                    tmp_column_names = [f"col_{i}" for i in range(len(self.columns))]
                     self_destruct = self.sql_ctx._conf.arrowPySparkSelfDestructEnabled()
                     batches = self.toDF(*tmp_column_names)._collect_as_arrow(
                         split_batches=self_destruct

mlrun/feature_store/steps.py

@@ -16,7 +16,7 @@ import math
 import re
 import uuid
 from collections import OrderedDict
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -92,8 +92,6 @@ class MLRunStep(MapClass):
 
 
 class FeaturesetValidator(StepToDict, MLRunStep):
-    """Validate feature values according to the feature set validation policy"""
-
     def __init__(self, featureset=None, columns=None, name=None, **kwargs):
         """Validate feature values according to the feature set validation policy
 
@@ -152,11 +150,9 @@ class FeaturesetValidator(StepToDict, MLRunStep):
 
 
 class MapValues(StepToDict, MLRunStep):
-    """Map column values to new values"""
-
     def __init__(
         self,
-        mapping: Dict[str, Dict[Union[str, int, bool], Any]],
+        mapping: dict[str, dict[Union[str, int, bool], Any]],
         with_original_features: bool = False,
         suffix: str = "mapped",
         **kwargs,
@@ -377,7 +373,7 @@ class Imputer(StepToDict, MLRunStep):
         self,
         method: str = "avg",
         default_value=None,
-        mapping: Dict[str, Any] = None,
+        mapping: dict[str, Any] = None,
         **kwargs,
     ):
         """Replace None values with default values
@@ -423,7 +419,7 @@ class Imputer(StepToDict, MLRunStep):
 
 
 class OneHotEncoder(StepToDict, MLRunStep):
-    def __init__(self, mapping: Dict[str, List[Union[int, str]]], **kwargs):
+    def __init__(self, mapping: dict[str, list[Union[int, str]]], **kwargs):
         """Create new binary fields, one per category (one hot encoded)
 
         example::
@@ -510,15 +506,13 @@ class OneHotEncoder(StepToDict, MLRunStep):
 
 
 class DateExtractor(StepToDict, MLRunStep):
-    """Date Extractor allows you to extract a date-time component"""
-
     def __init__(
         self,
-        parts: Union[Dict[str, str], List[str]],
+        parts: Union[dict[str, str], list[str]],
         timestamp_col: str = None,
         **kwargs,
     ):
-        """Date Extractor extract a date-time component into new columns
+        """Date Extractor extracts a date-time component into new columns
 
         The extracted date part will appear as `<timestamp_col>_<date_part>` feature.
 
@@ -629,8 +623,6 @@ class DateExtractor(StepToDict, MLRunStep):
 
 
 class SetEventMetadata(MapClass):
-    """Set the event metadata (id and key) from the event body"""
-
     def __init__(
         self,
         id_path: Optional[str] = None,
@@ -695,7 +687,7 @@ class SetEventMetadata(MapClass):
 
 
 class DropFeatures(StepToDict, MLRunStep):
-    def __init__(self, features: List[str], **kwargs):
+    def __init__(self, features: list[str], **kwargs):
         """Drop all the features from feature list
 
         :param features:    string list of the features names to drop

mlrun/features.py

@@ -14,7 +14,7 @@
 #
 import math
 import re
-from typing import Dict, List, Optional, Union
+from typing import Optional, Union
 
 from .data_types import ValueType, python_type_to_value_type
 from .errors import MLRunRuntimeError, err_to_str
@@ -44,7 +44,7 @@ class Entity(ModelObj):
         name: str = None,
         value_type: Union[ValueType, str] = None,
         description: str = None,
-        labels: Optional[Dict[str, str]] = None,
+        labels: Optional[dict[str, str]] = None,
     ):
         """data entity (index key)
 
@@ -65,8 +65,6 @@ class Entity(ModelObj):
 
 
 class Feature(ModelObj):
-    """data feature"""
-
     _dict_fields = [
         "name",
         "description",
@@ -82,13 +80,13 @@ class Feature(ModelObj):
     def __init__(
         self,
         value_type: Union[ValueType, str] = None,
-        dims: List[int] = None,
+        dims: list[int] = None,
         description: str = None,
         aggregate: bool = None,
         name: str = None,
         validator=None,
         default: str = None,
-        labels: Dict[str, str] = None,
+        labels: dict[str, str] = None,
     ):
         """data feature
 
@@ -445,7 +443,7 @@ class RegexValidator(Validator):
 
     @classmethod
     def from_dict(cls, struct=None, fields=None, deprecated_fields: dict = None):
-        new_obj = super(RegexValidator, cls).from_dict(
+        new_obj = super().from_dict(
             struct=struct, fields=fields, deprecated_fields=deprecated_fields
         )
         if hasattr(new_obj, "regex"):

mlrun/frameworks/_common/artifacts_library.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 #
 from abc import ABC, abstractmethod
-from typing import Dict, List, Type, Union
+from typing import Union
 
 import mlrun
 
@@ -39,12 +39,12 @@ class ArtifactsLibrary(ABC):
     @classmethod
     def get_plans(
         cls,
-        artifacts: Union[List[Plan], Dict[str, dict], List[str]] = None,
+        artifacts: Union[list[Plan], dict[str, dict], list[str]] = None,
         context: mlrun.MLClientCtx = None,
         include_default: bool = True,
         # custom_plans: dict = None, :param custom_plans: Custom user plans objects to initialize from.
         **default_kwargs,
-    ) -> List[Plan]:
+    ) -> list[Plan]:
         """
         Get plans for a run. The plans will be taken from the provided artifacts / configuration via code, from provided
         configuration via MLRun context and if the 'include_default' is True, from the framework artifact library's
@@ -97,7 +97,7 @@ class ArtifactsLibrary(ABC):
 
     @classmethod
     @abstractmethod
-    def default(cls, **kwargs) -> List[Plan]:
+    def default(cls, **kwargs) -> list[Plan]:
         """
         Get the default artifacts plans list of this framework's library.
 
@@ -106,7 +106,7 @@ class ArtifactsLibrary(ABC):
         pass
 
     @classmethod
-    def _get_library_plans(cls) -> Dict[str, Type[Plan]]:
+    def _get_library_plans(cls) -> dict[str, type[Plan]]:
         """
         Get all the supported plans in this library.
 
@@ -120,8 +120,8 @@ class ArtifactsLibrary(ABC):
 
     @staticmethod
     def _from_dict(
-        requested_plans: Dict[str, dict], available_plans: Dict[str, Type[Plan]]
-    ) -> List[Plan]:
+        requested_plans: dict[str, dict], available_plans: dict[str, type[Plan]]
+    ) -> list[Plan]:
         """
         Initialize a list of plans from a given configuration dictionary. The configuration is expected to be a
         dictionary of plans and their initialization parameters in the following format:
@@ -162,8 +162,8 @@ class ArtifactsLibrary(ABC):
 
     @staticmethod
     def _from_list(
-        requested_plans: List[str], available_plans: Dict[str, Type[Plan]]
-    ) -> List[Plan]:
+        requested_plans: list[str], available_plans: dict[str, type[Plan]]
+    ) -> list[Plan]:
         """
         Initialize a list of plans from a given configuration list. The configuration is expected to be a list of plans
         names to be initialized with their default configuration.

mlrun/frameworks/_common/mlrun_interface.py

@@ -17,7 +17,7 @@ import functools
 import inspect
 from abc import ABC
 from types import FunctionType, MethodType
-from typing import Any, Dict, Generic, List, Tuple, Union
+from typing import Any, Generic, Union
 
 from .utils import CommonTypes
 
@@ -173,7 +173,7 @@ class MLRunInterface(ABC, Generic[CommonTypes.MLRunInterfaceableType]):
     def _insert_properties(
         cls,
         obj: CommonTypes.MLRunInterfaceableType,
-        properties: Dict[str, Any] = None,
+        properties: dict[str, Any] = None,
     ):
         """
         Insert the properties of the interface to the object. The properties default values are being copied (not deep
@@ -238,7 +238,7 @@ class MLRunInterface(ABC, Generic[CommonTypes.MLRunInterfaceableType]):
 
     @classmethod
     def _replace_properties(
-        cls, obj: CommonTypes.MLRunInterfaceableType, properties: Dict[str, Any] = None
+        cls, obj: CommonTypes.MLRunInterfaceableType, properties: dict[str, Any] = None
     ):
         """
         Replace the properties of the given object according to the configuration in the MLRun interface.
@@ -282,7 +282,7 @@ class MLRunInterface(ABC, Generic[CommonTypes.MLRunInterfaceableType]):
 
     @classmethod
     def _replace_functions(
-        cls, obj: CommonTypes.MLRunInterfaceableType, functions: List[str] = None
+        cls, obj: CommonTypes.MLRunInterfaceableType, functions: list[str] = None
     ):
         """
         Replace the functions / methods of the given object according to the configuration in the MLRun interface.
@@ -420,7 +420,7 @@ class MLRunInterface(ABC, Generic[CommonTypes.MLRunInterfaceableType]):
         passed_args: tuple = None,
         passed_kwargs: dict = None,
         default_value: Any = None,
-    ) -> Tuple[Any, Union[str, int, None]]:
+    ) -> tuple[Any, Union[str, int, None]]:
         """
         Get a passed argument (from *args or **kwargs) to a function. If the argument was not found the default value
         will be returned. In addition, the keyword of the argument in `kwargs` or the index of the argument in `args`