frogml-core 0.0.1__py3-none-any.whl

frogml/exceptions/quiet_error.py

@@ -0,0 +1,22 @@
+import sys
+
+
+class QuietError(Exception):
+    # All who inherit me shall not traceback, but be spoken of cleanly
+    def __init__(self, message):
+        self.message = None
+
+    def __str__(self):
+        return f"\033[91m{self.message}\033[0m"
+
+
+def quiet_hook(kind, message, traceback):
+    if issubclass(kind, QuietError):
+        print("{0}".format(message))  # Only print Error Type and Message
+    else:
+        sys.__excepthook__(
+            kind, message, traceback
+        )  # Print Error Type, Message and Traceback
+
+
+sys.excepthook = quiet_hook

frogml/feature_store/_common/artifact_utils.py

@@ -0,0 +1,57 @@
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union
+
+from frogml.clients.feature_store.management_client import FeatureRegistryClient
+from frogml.feature_store._common.packaging import upload_artifact
+
+if TYPE_CHECKING:
+    from frogml.feature_store.data_sources.streaming.kafka.deserialization import (
+        Deserializer,
+    )
+    from frogml.feature_store.feature_sets.transformations.transformations import (
+        BaseTransformation,
+    )
+
+
+@dataclass
+class ArtifactSpec:
+    """
+    Dataclass for holding artifacts which will be uploaded
+    """
+
+    artifact_name: str
+    root_module_path: Path
+    artifact_object: Union["BaseTransformation", "Deserializer"]
+    callables: List[Callable] = field(default_factory=lambda: list())
+    suffix: str = ""
+
+
+class ArtifactsUploader:
+    @staticmethod
+    def upload(artifact: ArtifactSpec) -> str:
+        return upload_artifact(
+            feature_store_object_name=artifact.artifact_name,
+            feature_store_object_name_suffix=artifact.suffix,
+            functions=artifact.callables,
+            feature_module_dir=artifact.root_module_path,
+            features_manager_client=FeatureRegistryClient(),
+            artifact_object=artifact.artifact_object,
+        )
+
+    @staticmethod
+    def get_artifact_spec(
+        transformation: "BaseTransformation",
+        featureset_name: str,
+        __instance_module_path__: str,
+    ) -> Optional[ArtifactSpec]:
+        transformation_functions: Optional[
+            List[Callable[..., Any]]
+        ] = transformation.get_functions()
+        if transformation_functions is not None and transformation_functions:
+            return ArtifactSpec(
+                artifact_name=featureset_name,
+                root_module_path=Path(__instance_module_path__).parent,
+                artifact_object=transformation,
+                callables=transformation_functions,
+            )

frogml/feature_store/_common/feature_set_utils.py

@@ -0,0 +1,261 @@
+import time
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+from _frogml_proto.qwak.build.v1.build_pb2 import BatchFeature as ProtoBatchFeature
+from _frogml_proto.qwak.build.v1.build_pb2 import BatchFeatureV1 as ProtoBatchFeatureV1
+from _frogml_proto.qwak.build.v1.build_pb2 import Feature as ProtoFeature
+from _frogml_proto.qwak.build.v1.build_pb2 import (
+    StreamingAggregationFeature as ProtoStreamingAggregationFeature,
+)
+from _frogml_proto.qwak.build.v1.build_pb2 import (
+    StreamingFeature as ProtoStreamingFeature,
+)
+from _frogml_proto.qwak.build.v1.build_pb2 import (
+    StreamingFeatureV1 as ProtoStreamingFeatureV1,
+)
+from _frogml_proto.qwak.feature_store.entities.entity_pb2 import (
+    EntitySpec as ProtoEntitySpec,
+)
+from _frogml_proto.qwak.feature_store.features.feature_set_types_pb2 import (
+    BatchFeatureSet as ProtoBatchFeatureSet,
+)
+from _frogml_proto.qwak.feature_store.features.feature_set_types_pb2 import (
+    BatchFeatureSetV1 as ProtoBatchFeatureSetV1,
+)
+from _frogml_proto.qwak.feature_store.features.feature_set_types_pb2 import (
+    FeatureSetType as ProtoFeatureSetType,
+)
+from _frogml_proto.qwak.feature_store.features.feature_set_types_pb2 import (
+    StreamingAggregationFeatureSet as ProtoStreamingAggregationFeatureSet,
+)
+from _frogml_proto.qwak.feature_store.features.feature_set_types_pb2 import (
+    StreamingFeatureSetV1 as ProtoStreamingFeatureSetV1,
+)
+from frogml.clients.feature_store import FeatureRegistryClient
+from frogml.exceptions import FrogmlException
+from frogml.model.schema_entities import BaseFeature, Entity, FeatureStoreInput
+
+_INTERNAL_KEY_PREFIX: str = "_frogml_key_"
+
+
+@dataclass(unsafe_hash=True)
+class BatchFeature(BaseFeature):
+    entity: Optional[Entity] = None
+
+    def to_proto(self):
+        return ProtoFeature(
+            batch_feature=ProtoBatchFeature(
+                name=self.name, entity=self.entity.to_proto()
+            )
+        )
+
+
+@dataclass(unsafe_hash=True)
+class BatchFeatureV1(BaseFeature):
+    entity: Optional[Entity] = None
+
+    def to_proto(self):
+        return ProtoFeature(
+            batch_feature_v1=ProtoBatchFeatureV1(
+                name=self.name, entity=self.entity.to_proto()
+            )
+        )
+
+
+@dataclass(unsafe_hash=True)
+class StreamingFeature(BaseFeature):
+    entity: Entity
+
+    def to_proto(self):
+        return ProtoFeature(
+            streaming_feature=ProtoStreamingFeature(
+                name=self.name, entity=self.entity.to_proto()
+            )
+        )
+
+
+@dataclass(unsafe_hash=True)
+class StreamingFeatureV1(BaseFeature):
+    entity: Optional[Entity] = None
+
+    def to_proto(self):
+        return ProtoFeature(
+            streaming_feature_v1=ProtoStreamingFeatureV1(
+                name=self.name, entity=self.entity.to_proto()
+            )
+        )
+
+
+@dataclass(unsafe_hash=True)
+class StreamingAggregationFeature(BaseFeature):
+    entity: Optional[Entity] = None
+
+    def to_proto(self):
+        return ProtoFeature(
+            streaming_aggregation_feature=ProtoStreamingAggregationFeature(
+                name=self.name, entity=self.entity.to_proto()
+            )
+        )
+
+
+@dataclass()
+class FeatureSetInfo:
+    entity_spec: ProtoEntitySpec
+    feature_set_type: ProtoFeatureSetType
+    feature_version: int
+
+
+def get_feature_set_info(
+    feature_manager_client: FeatureRegistryClient, feature_set_name: str
+) -> FeatureSetInfo:
+    """
+    Get the entities by the feature set name and feature type
+    Args:
+        feature_manager_client: feature manager client for the grpc request
+        feature_set_name: the required feature set name
+    Returns: tuple of entity spec, type and feature version
+
+    """
+    feature_set_response = feature_manager_client.get_feature_set_by_name(
+        feature_set_name
+    )
+    if not feature_set_response:
+        raise FrogmlException(f"Featureset '{feature_set_name}' does not exist")
+
+    featureset_type: ProtoFeatureSetType = get_feature_type(
+        feature_set_response.feature_set.feature_set_definition.feature_set_spec.feature_set_type
+    )
+    featureset_version = _get_featureset_version(featureset_type)
+
+    return FeatureSetInfo(
+        entity_spec=feature_set_response.feature_set.feature_set_definition.feature_set_spec.entity.entity_spec,
+        feature_set_type=featureset_type,
+        feature_version=featureset_version,
+    )
+
+
+def get_feature_type(feature_set_type: ProtoFeatureSetType):
+    return getattr(feature_set_type, feature_set_type.WhichOneof("set_type"))
+
+
+def _get_featureset_version(feature_set_type: ProtoFeatureSetType) -> int:
+    """
+    Get Featureset version. Return 0 if none version is set
+    """
+    if hasattr(feature_set_type, "qwak_internal_protocol_version"):
+        return feature_set_type.qwak_internal_protocol_version
+    return 0
+
+
+def get_typed_feature(
+    feature: FeatureStoreInput,
+    feature_type: ProtoFeatureSetType,
+    featureset_version: int,
+) -> BaseFeature:
+    """
+    convert InputFeature to the relevant type
+    Args:
+        feature: Input feature to cast to the correct type
+        feature_type: the feature type as it registered
+        featureset_version: the version of the featureset
+
+    Return:
+        BaseFeature with the correct type
+    """
+    if isinstance(feature_type, ProtoBatchFeatureSet):
+        return BatchFeature(name=feature.name, entity=feature.entity)
+    elif isinstance(feature_type, ProtoBatchFeatureSetV1) and featureset_version == 0:
+        return BatchFeature(name=feature.name, entity=feature.entity)
+    elif isinstance(feature_type, ProtoBatchFeatureSetV1) and featureset_version != 0:
+        return BatchFeatureV1(name=feature.name, entity=feature.entity)
+    elif (
+        isinstance(feature_type, ProtoStreamingFeatureSetV1) and featureset_version == 0
+    ):
+        return StreamingFeature(name=feature.name, entity=feature.entity)
+    elif (
+        isinstance(feature_type, ProtoStreamingFeatureSetV1) and featureset_version != 0
+    ):
+        return StreamingFeatureV1(name=feature.name, entity=feature.entity)
+    elif isinstance(feature_type, ProtoStreamingAggregationFeatureSet):
+        return StreamingAggregationFeature(name=feature.name, entity=feature.entity)
+    else:
+        raise ValueError(
+            f"Feature set type {feature_type} with protocol version {featureset_version} is not supported for extraction"
+        )
+
+
+def get_env_to_featuresets_mapping(
+    feature_manager_client: FeatureRegistryClient,
+) -> Dict[str, Dict[str, FeatureSetInfo]]:
+    """
+    Get mapping of the account's environment name to feature sets and converts each FeatureSet to FeatureSetInfo
+    :param feature_manager_client:
+    :return: Dict[str, List[FeatureSetInfo]] where the key is the environment name and the value is a list of FeatureSetInfo
+    """
+    env_to_featuresets_mapping = feature_manager_client.get_env_to_featuresets_mapping()
+
+    # Convert the mapping to a dictionary of type Dict[str, List[FeatureSetInfo]]
+    result = {}
+    for (
+        env_name,
+        feature_sets_infos,
+    ) in env_to_featuresets_mapping.env_to_feature_set_mapping.items():
+        env_name = env_name.lower()
+        result[env_name] = {}
+        for feature_set in feature_sets_infos.feature_sets:
+            featureset_type: ProtoFeatureSetType = get_feature_type(
+                feature_set.feature_set_definition.feature_set_spec.feature_set_type
+            )
+            featureset_version = _get_featureset_version(featureset_type)
+
+            feature_set_info = FeatureSetInfo(
+                entity_spec=feature_set.feature_set_definition.feature_set_spec.entity.entity_spec,
+                feature_set_type=featureset_type,
+                feature_version=featureset_version,
+            )
+
+            feature_set_name = (
+                feature_set.feature_set_definition.feature_set_spec.name.lower()
+            )
+            result[env_name][feature_set_name] = feature_set_info
+    return result
+
+
+def get_entity_type(value_type: ProtoEntitySpec.ValueType):
+    """
+    Normalize entity by the enum
+    """
+
+    if value_type == ProtoEntitySpec.ValueType.STRING:
+        return str
+    elif value_type == ProtoEntitySpec.ValueType.INT:
+        return int
+
+
+def generate_key_unique_name(featureset_name: str):
+    """
+    Generate a name for the feature set key to be registered or treated as a Frogml entity based on its' featuresets'
+    name
+    """
+    return f"{_INTERNAL_KEY_PREFIX}_{featureset_name}_{round(time.time() * 1000)}"
+
+
+def get_batch_source_for_featureset(
+    batch_ds_name: str, feature_registry: FeatureRegistryClient
+):
+    batch_ds = feature_registry.get_data_source_by_name(batch_ds_name)
+    if not batch_ds:
+        raise FrogmlException(
+            f"Trying to register a feature set with a non registered data source -: {batch_ds_name}"
+        )
+
+    ds_spec = batch_ds.data_source.data_source_definition.data_source_spec
+
+    if ds_spec.WhichOneof("type") != "batch_source":
+        raise ValueError(
+            f"Can only register streaming backfill/batch featuresets with batch sources. "
+            f"Source {batch_ds_name} is of type {ds_spec.WhichOneof('type')}"
+        )
+
+    return ds_spec.batch_source

frogml/feature_store/_common/featureset_asterisk_handler.py

@@ -0,0 +1,115 @@
+import dataclasses
+import itertools
+import re
+from typing import Callable, Dict, List
+
+from frogml.feature_store.offline import OfflineClient
+from frogml.model.schema_entities import BaseFeature, Entity, FeatureStoreInput
+
+
+def is_asterisk(feature: BaseFeature) -> bool:
+    """
+    Checks whether the given Feature is a FeatureStoreInput with an '*'
+    Args:
+        feature: input feature
+    Return
+    """
+    p = "^[^.]+.\\*$"
+    # if this is not a FeatureStoreInput, no need to check
+    if isinstance(feature, FeatureStoreInput):
+        return bool(re.match(p, feature.name))
+    return False
+
+
+def inflate_feature(feature: BaseFeature, offline_fs: OfflineClient) -> List:
+    """
+    Inflates a Feature into a list. if the feature contains an '*', a list
+        a feature for each feature in the associate FeatureSet is returned. else -
+        a list containing the original feature is returned.
+    Args:
+        feature: input feature
+        offline_fs: OfflineFeatureStore to query
+    Return:
+        if an '*' exists, returns a List of FeatureStoreInuput for each feature in the feature set
+        else returns a list containing the feature.
+
+    """
+    if is_asterisk(feature):
+        featureset_name = feature.name.lower().split(".")[0]
+
+        p = f"^{featureset_name}+.*$"
+        feature_names = [
+            feature.lower()
+            for feature in offline_fs.get_columns_by_feature_set(featureset_name)
+            if re.match(p, feature.lower())
+        ]
+        return [dataclasses.replace(feature, name=name) for name in feature_names]
+
+    return [feature]
+
+
+def unpack_asterisk_features(
+    features: List[BaseFeature],
+    feature_store_generator: Callable[[], OfflineClient] = lambda: OfflineClient(),
+) -> List[BaseFeature]:
+    """
+    Handles features with an '*'.
+    If a feature of type FeatureStoreInput that matches '<featureset_name>.*' exists,
+    transforms this feature into the list of features actually present in the FeatureSet.
+    Other features remain unchanged.
+    the original order is maintained, and the internal order (inside a FeatureStoreInput
+     that has an '*') is inferred from the order of columns returned from
+     OfflineFeatureStore.get_columns_by_feature_set(...)
+     Args:
+         features: List of features
+         feature_store_generator: A function that generates an instance of OfflineFeatureStore
+                                  If no feature with an '*' is found, the function is not called.
+
+    Returns:
+        the unpacked features
+    """
+    asterisk_found = any([is_asterisk(f) for f in features])
+    if asterisk_found:
+        # at least 1 feature contains an asterisk, map each feature to
+        # a list of features (features with no '*' are mapped to a list of size 1), then flatten it.
+        offline_feature_store = feature_store_generator()
+        inflated_features = [
+            inflate_feature(feature, offline_feature_store) for feature in features
+        ]
+        return list(itertools.chain.from_iterable(inflated_features))
+
+    return features
+
+
+def unpack_asterisk_features_from_key_mapping(
+    key_to_features: Dict[str, List[str]],
+    feature_store_generator: Callable[[], OfflineClient] = lambda: OfflineClient(),
+) -> dict:
+    """
+    Handles features with an '*'.
+    If a feature name that matches '<featureset_name>.*' exists in one of the lists,
+    transforms this feature into the list of feature names actually present in the FeatureSet.
+    Other features remain unchanged.
+    the original order is maintained, and the internal order (inside a feature
+     that has an '*') is inferred from the order of columns returned from
+     OfflineFeatureStore.get_columns_by_feature_set(...)
+     Args:
+         key_to_features: dictionary of entity keys to requested features (same as in the offline store)
+             >>> key_to_features = {'uuid': ['user_purchases.*',
+             >>>                             'user_purchases.avg_purchase_amount']}
+         feature_store_generator: A function that generates an instance of OfflineFeatureStore
+                                  If no feature with an '*' is found, the function is not called.
+    Returns:
+        the unpacked features
+    """
+    inflated_key_to_features = {}
+
+    for entity, features in key_to_features.items():
+        feature_inputs = [
+            FeatureStoreInput(name=feature, entity=Entity(name=entity, type=str))
+            for feature in features
+        ]
+        unpacked = unpack_asterisk_features(feature_inputs, feature_store_generator)
+        feature_list = [feature.name for feature in unpacked]
+        inflated_key_to_features[entity] = feature_list
+    return inflated_key_to_features

frogml/feature_store/_common/functions.py

@@ -0,0 +1,37 @@
+from typing import List
+
+
+def _generate_new_col(col_name: str, duplicate_cols_dict: dict):
+    """
+    Generate column without leading feature set name.
+    Args:
+        col_name: the target column name
+        duplicate_cols_dict: duplicate features dictionary
+    Returns list of the new columns
+    """
+    feature_full_name: List[str] = col_name.split(".")
+    if len(feature_full_name) != 2:
+        return col_name
+    elif len(feature_full_name) == 2 and duplicate_cols_dict[feature_full_name[1]] > 1:
+        return col_name
+    elif len(feature_full_name) == 2 and duplicate_cols_dict[feature_full_name[1]] == 1:
+        return feature_full_name[1]
+
+
+def normalize_cols(cols) -> List[str]:
+    """
+    Normalize cols - try to remove leading feature set name from features
+    Args:
+        cols: list of column
+    Return normalized columns
+    """
+    duplicate_columns_dict = {}
+    for col in cols:
+        feature_full_name: List[str] = col.split(".")
+        if len(feature_full_name) != 2:
+            duplicate_columns_dict[col] = 1
+        else:
+            fs_name = feature_full_name[1]
+            duplicate_columns_dict[fs_name] = duplicate_columns_dict.get(fs_name, 0) + 1
+    cols = [_generate_new_col(col_name, duplicate_columns_dict) for col_name in cols]
+    return cols