frogml-core 0.0.115__py3-none-any.whl → 0.0.117__py3-none-any.whl

feature_store/data_sources/batch/snowflake.py

@@ -0,0 +1,79 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
+    BatchSource as ProtoBatchSource,
+)
+from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
+    SnowflakeSource as ProtoSnowflakeSource,
+)
+from frogml._proto.qwak.feature_store.sources.data_source_pb2 import (
+    DataSourceSpec as ProtoDataSourceSpec,
+)
+from frogml.core.exceptions import FrogmlException
+from frogml.feature_store.data_sources.batch._batch import BaseBatchSource
+
+
+@dataclass
+class SnowflakeSource(BaseBatchSource):
+    host: str
+    username_secret_name: str
+    password_secret_name: str
+    database: str
+    schema: str
+    warehouse: str
+    table: Optional[str] = None
+    query: Optional[str] = None
+    repository: Optional[str] = None
+
+    def __post_init__(self):
+        self._validate()
+
+    def _validate(self):
+        if self.table and self.query:
+            raise FrogmlException("Only one of query and table may be set")
+        if not self.table and not self.query:
+            raise FrogmlException("One of table or query must be set")
+
+        if not self.username_secret_name:
+            raise FrogmlException("username_secret_name must be set!")
+
+        if not self.password_secret_name:
+            raise FrogmlException("password_secret_name must be set!")
+
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
+        return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
+            batch_source=ProtoBatchSource(
+                name=self.name,
+                description=self.description,
+                date_created_column=self.date_created_column,
+                snowflakeSource=ProtoSnowflakeSource(
+                    host=self.host,
+                    username_secret_name=self.username_secret_name,
+                    password_secret_name=self.password_secret_name,
+                    database=self.database,
+                    schema=self.schema,
+                    warehouse=self.warehouse,
+                    table=self.table,
+                    query=self.query,
+                ),
+            ),
+        )
+
+    @classmethod
+    def _from_proto(cls, proto):
+        snowflake = proto.snowflakeSource
+        return cls(
+            name=proto.name,
+            date_created_column=proto.date_created_column,
+            description=proto.description,
+            host=snowflake.host,
+            username_secret_name=snowflake.username_secret_name,
+            password_secret_name=snowflake.password_secret_name,
+            database=snowflake.database,
+            schema=snowflake.schema,
+            warehouse=snowflake.warehouse,
+            table=snowflake.table,
+            query=snowflake.query,
+        )

feature_store/data_sources/batch/vertica.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
+    BatchSource as ProtoBatchSource,
+)
+from frogml._proto.qwak.feature_store.sources.batch_pb2 import (
+    VerticaSource as ProtoVerticaSource,
+)
+from frogml._proto.qwak.feature_store.sources.data_source_pb2 import (
+    DataSourceSpec as ProtoDataSourceSpec,
+)
+from frogml.feature_store.data_sources.batch._batch import BaseBatchSource
+
+
+@dataclass
+class VerticaSource(BaseBatchSource):
+    host: str
+    port: int
+    database: str
+    schema: str
+    table: str
+    username_secret_name: str
+    password_secret_name: str
+    repository: Optional[str] = None
+
+    @classmethod
+    def _from_proto(cls, proto):
+        vertica = proto.verticaSource
+        return cls(
+            name=proto.name,
+            date_created_column=proto.date_created_column,
+            description=proto.description,
+            host=vertica.host,
+            username_secret_name=vertica.username_secret_name,
+            password_secret_name=vertica.password_secret_name,
+            database=vertica.database,
+            schema=vertica.schema,
+            port=vertica.port,
+            table=vertica.table,
+        )
+
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
+        return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
+            batch_source=ProtoBatchSource(
+                name=self.name,
+                description=self.description,
+                date_created_column=self.date_created_column,
+                verticaSource=ProtoVerticaSource(
+                    host=self.host,
+                    username_secret_name=self.username_secret_name,
+                    password_secret_name=self.password_secret_name,
+                    database=self.database,
+                    schema=self.schema,
+                    port=self.port,
+                    table=self.table,
+                ),
+            ),
+        )

feature_store/data_sources/streaming/_streaming.py

@@ -0,0 +1,9 @@
+from abc import ABC
+from dataclasses import dataclass
+
+from frogml.feature_store.data_sources.base import BaseSource
+
+
+@dataclass
+class BaseStreamingSource(BaseSource, ABC):
+    pass

feature_store/data_sources/streaming/kafka/__init__.py

@@ -0,0 +1,25 @@
+from frogml.feature_store.data_sources.streaming.kafka.authentication import (
+    PlainAuthentication,
+    SaslAuthentication,
+    SaslMechanism,
+    SecurityProtocol,
+    SslAuthentication,
+)
+from frogml.feature_store.data_sources.streaming.kafka.deserialization import (
+    CustomDeserializer,
+    GenericDeserializer,
+    MessageFormat,
+)
+from frogml.feature_store.data_sources.streaming.kafka.kafka import KafkaSource
+
+__all__ = [
+    "KafkaSource",
+    "CustomDeserializer",
+    "GenericDeserializer",
+    "SaslAuthentication",
+    "SslAuthentication",
+    "PlainAuthentication",
+    "MessageFormat",
+    "SaslMechanism",
+    "SecurityProtocol",
+]

feature_store/data_sources/streaming/kafka/authentication.py

@@ -0,0 +1,124 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+from typing import Type
+
+from frogml._proto.qwak.feature_store.sources.streaming_pb2 import (
+    Authentication as ProtoAuthentication,
+)
+from frogml._proto.qwak.feature_store.sources.streaming_pb2 import Plain as ProtoPlain
+from frogml._proto.qwak.feature_store.sources.streaming_pb2 import Sasl as ProtoSasl
+from frogml._proto.qwak.feature_store.sources.streaming_pb2 import (
+    SaslMechanism as ProtoSaslMechanism,
+)
+from frogml._proto.qwak.feature_store.sources.streaming_pb2 import (
+    SecurityProtocol as ProtoSecurityProtocol,
+)
+from frogml._proto.qwak.feature_store.sources.streaming_pb2 import Ssl as ProtoSsl
+from frogml.core.clients.secret_service import SecretServiceClient
+from frogml.core.exceptions import FrogmlException
+
+
+class SaslMechanism(Enum):
+    SCRAMSHA256 = ProtoSaslMechanism.SCRAMSHA256
+    SCRAMSHA512 = ProtoSaslMechanism.SCRAMSHA512
+    PLAIN = ProtoSaslMechanism.PLAIN
+
+
+class SecurityProtocol(Enum):
+    SASL_SSL = ProtoSecurityProtocol.SASL_SSL
+
+
+@dataclass
+class BaseAuthentication(ABC):
+    @abstractmethod
+    def _to_proto(self) -> ProtoAuthentication:
+        pass
+
+    @classmethod
+    def _from_proto(
+        cls,
+        proto_authentication_method: ProtoAuthentication,
+    ) -> Type["BaseAuthentication"]:
+        proto_authentication_method = getattr(
+            proto_authentication_method,
+            proto_authentication_method.WhichOneof("type"),
+        )
+        if isinstance(proto_authentication_method, ProtoPlain):
+            return PlainAuthentication._from_proto(proto_authentication_method)
+        elif isinstance(proto_authentication_method, ProtoSsl):
+            return SslAuthentication._from_proto(proto_authentication_method)
+        elif isinstance(proto_authentication_method, ProtoSasl):
+            return SaslAuthentication._from_proto(proto_authentication_method)
+        else:
+            raise FrogmlException(
+                f"Got unsupported authentication method {proto_authentication_method}"
+            )
+
+
+@dataclass
+class PlainAuthentication(BaseAuthentication):
+    def _to_proto(self) -> ProtoAuthentication:
+        return ProtoAuthentication(plain_configuration=ProtoPlain())
+
+    @classmethod
+    def _from_proto(
+        cls, proto_authentication_method: ProtoPlain
+    ) -> Type["PlainAuthentication"]:
+        return cls()
+
+
+@dataclass
+class SslAuthentication(BaseAuthentication):
+    def _to_proto(self) -> ProtoAuthentication:
+        return ProtoAuthentication(ssl_configuration=ProtoSsl())
+
+    @classmethod
+    def _from_proto(
+        cls, proto_authentication_method: ProtoSsl
+    ) -> Type["SslAuthentication"]:
+        return cls()
+
+
+@dataclass
+class SaslAuthentication(BaseAuthentication):
+    username_secret: str
+    password_secret: str
+    sasl_mechanism: SaslMechanism
+    security_protocol: SecurityProtocol
+
+    def _to_proto(self) -> ProtoAuthentication:
+        self._validate()
+        return ProtoAuthentication(
+            sasl_configuration=ProtoSasl(
+                username_secret=self.username_secret,
+                password_secret=self.password_secret,
+                sasl_mechanism=self.sasl_mechanism.value,
+                security_protocol=self.security_protocol.value,
+            )
+        )
+
+    @classmethod
+    def _from_proto(
+        cls, proto_authentication_method: ProtoSasl
+    ) -> Type["SaslAuthentication"]:
+        return cls(
+            username_secret=proto_authentication_method.username_secret,
+            password_secret=proto_authentication_method.password_secret,
+            sasl_mechanism=SaslMechanism(proto_authentication_method.sasl_mechanism),
+            security_protocol=SecurityProtocol(
+                proto_authentication_method.security_protocol
+            ),
+        )
+
+    def _validate(self):
+        secret_service_client = SecretServiceClient()
+
+        if not secret_service_client.get_secret(self.username_secret):
+            raise FrogmlException(
+                f"Secret for username {self.username_secret} does not exist"
+            )
+        if not secret_service_client.get_secret(self.password_secret):
+            raise FrogmlException(
+                f"Secret for password {self.password_secret} does not exist"
+            )

feature_store/data_sources/streaming/kafka/deserialization.py

@@ -0,0 +1,115 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Callable, Optional, Type
+
+from frogml._proto.qwak.feature_store.sources.streaming_pb2 import (
+    CustomDeserializer as ProtoCustomDeserializer,
+)
+from frogml._proto.qwak.feature_store.sources.streaming_pb2 import (
+    Deserialization as ProtoDeserialization,
+)
+from frogml._proto.qwak.feature_store.sources.streaming_pb2 import (
+    GenericDeserializer as ProtoGenericDeserializer,
+)
+from frogml._proto.qwak.feature_store.sources.streaming_pb2 import (
+    MessageFormat as ProtoMessageFormat,
+)
+from frogml.core.exceptions import FrogmlException
+
+
+class MessageFormat(Enum):
+    JSON = ProtoMessageFormat.JSON
+    AVRO = ProtoMessageFormat.AVRO
+
+
+class Deserializer(ABC):
+    @abstractmethod
+    def _to_proto(self, artifact_path: Optional[str] = None) -> ProtoDeserialization:
+        pass
+
+    @abstractmethod
+    def _get_function(self) -> Optional[Callable]:
+        pass
+
+    @classmethod
+    def _from_proto(
+        cls,
+        proto_deserializer: ProtoDeserialization,
+    ) -> Type["Deserializer"]:
+        deserializer = getattr(
+            proto_deserializer, proto_deserializer.WhichOneof("type")
+        )
+
+        if isinstance(deserializer, ProtoGenericDeserializer):
+            return GenericDeserializer._from_proto(proto_deserializer=deserializer)
+        elif isinstance(deserializer, ProtoCustomDeserializer):
+            return CustomDeserializer._from_proto(proto_deserializer=deserializer)
+        else:
+            raise FrogmlException(f"Got unsupported deserializer type {deserializer}")
+
+
+@dataclass
+class GenericDeserializer(Deserializer):
+    message_format: MessageFormat
+    schema: str
+
+    def _to_proto(self, artifact_path: Optional[str] = None) -> ProtoDeserialization:
+        # TODO: add backend schema validation
+        return ProtoDeserialization(
+            generic_deserializer=ProtoGenericDeserializer(
+                deserializer_format=self.message_format.value, schema=self.schema
+            )
+        )
+
+    def _get_function(self) -> Optional[Callable]:
+        return None
+
+    @classmethod
+    def _from_proto(
+        cls, proto_deserializer: ProtoGenericDeserializer
+    ) -> Type["GenericDeserializer"]:
+        return cls(
+            message_format=MessageFormat(proto_deserializer.deserializer_format),
+            schema=proto_deserializer.schema,
+        )
+
+
+@dataclass
+class CustomDeserializer(Deserializer):
+    function: Callable
+    _artifact_path: Optional[str] = field(init=False, default=None)
+
+    def _to_proto(self, artifact_path: Optional[str] = None) -> ProtoDeserialization:
+        self._validate()
+
+        if artifact_path:
+            self._artifact_path = artifact_path
+
+        # TODO: add backend schema validation
+        return ProtoDeserialization(
+            custom_deserializer=ProtoCustomDeserializer(
+                function_name=self.function.__name__, artifact_path=self._artifact_path
+            )
+        )
+
+    @classmethod
+    def _from_proto(
+        cls, proto_deserializer: ProtoCustomDeserializer
+    ) -> Type["CustomDeserializer"]:
+        def dummy_deserializer(df):
+            return df
+
+        custom_function: Callable = dummy_deserializer
+        custom_function.__name__ = proto_deserializer.function_name
+
+        return cls(function=custom_function)
+
+    def _validate(self):
+        if self.function is None:
+            raise FrogmlException("deserialization function must be set!")
+        if self.function.__name__ == "<lambda>":
+            raise FrogmlException("Custom Deserializer can not be set with a lambda")
+
+    def _get_function(self) -> Optional[Callable]:
+        return self.function

feature_store/data_sources/streaming/kafka/kafka.py

@@ -0,0 +1,126 @@
+import inspect
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Callable, Dict, Optional, Type
+
+from frogml._proto.qwak.feature_store.sources.data_source_pb2 import (
+    DataSourceSpec as ProtoDataSourceSpec,
+)
+from frogml._proto.qwak.feature_store.sources.streaming_pb2 import (
+    KafkaSourceV1 as ProtoKafkaSourceV1,
+)
+from frogml._proto.qwak.feature_store.sources.streaming_pb2 import (
+    StreamingSource as ProtoStreamingSource,
+)
+from frogml.core.exceptions import FrogmlException
+from frogml.feature_store._common.artifact_utils import ArtifactSpec
+from frogml.feature_store.data_sources.streaming._streaming import (
+    BaseStreamingSource,
+)
+from frogml.feature_store.data_sources.streaming.kafka.authentication import (
+    BaseAuthentication,
+    SslAuthentication,
+)
+from frogml.feature_store.data_sources.streaming.kafka.deserialization import (
+    Deserializer,
+)
+
+
+@dataclass
+class KafkaSource(BaseStreamingSource):
+    bootstrap_servers: str
+
+    # Deserialization
+    deserialization: Deserializer
+
+    # secret configs, the value is resolved to the secret,
+    # s.t. (key, value) -> (key, get_secret(value))
+    # not all configs will be respected, this is a best-effort
+    secret_configs: Dict[str, str] = field(default_factory=lambda: {})
+
+    # passthrough configs - not all configs will be respected,
+    # this is a best-effort
+    passthrough_configs: Dict[str, str] = field(default_factory=lambda: {})
+
+    # the following 3 are pairwise mutually exclusive
+    assign: Optional[str] = None
+    subscribe: Optional[str] = None
+    subscribe_pattern: Optional[str] = None
+
+    authentication_method: Optional[BaseAuthentication] = field(
+        default_factory=lambda: SslAuthentication()
+    )
+
+    repository: Optional[str] = None
+
+    def __post_init__(self):
+        self._validate()
+
+    def _validate(self):
+        num_defined = len(
+            [
+                _
+                for _ in [self.assign, self.subscribe, self.subscribe_pattern]
+                if _ is not None
+            ]
+        )
+        if num_defined != 1:
+            raise FrogmlException(
+                "Exactly one of (assign, subscribe, subscribe_pattern) must be defined!"
+            )
+
+    def _get_artifacts(self) -> Optional["ArtifactSpec"]:
+        deserializer_function: Optional[Callable] = self.deserialization._get_function()
+
+        if deserializer_function:
+            return ArtifactSpec(
+                artifact_name=self.name,
+                root_module_path=Path(
+                    os.path.abspath(inspect.getfile(deserializer_function))
+                ).parent,
+                artifact_object=self.deserialization,
+                callables=[deserializer_function],
+                suffix="resource.zip",
+            )
+
+    def _to_proto(self, artifact_url: Optional[str] = None) -> ProtoDataSourceSpec:
+        return ProtoDataSourceSpec(
+            data_source_repository_name=self.repository,
+            stream_source=ProtoStreamingSource(
+                name=self.name,
+                description=self.description,
+                kafkaSourceV1=ProtoKafkaSourceV1(
+                    bootstrap_servers=self.bootstrap_servers,
+                    assign=self.assign,
+                    subscribe=self.subscribe,
+                    subscribe_pattern=self.subscribe_pattern,
+                    secret_configs=self.secret_configs,
+                    passthrough_configs=self.passthrough_configs,
+                    authentication_method=self.authentication_method._to_proto(),
+                    deserialization=self.deserialization._to_proto(
+                        artifact_path=artifact_url
+                    ),
+                ),
+            ),
+        )
+
+    @classmethod
+    def _from_proto(cls, proto) -> Type["KafkaSource"]:
+        kafka = proto.kafkaSourceV1
+        topic_configuration_key = kafka.WhichOneof("topic_configuration")
+        oneof_args = {topic_configuration_key: getattr(kafka, topic_configuration_key)}
+        return cls(
+            name=proto.name,
+            description=proto.description,
+            bootstrap_servers=kafka.bootstrap_servers,
+            secret_configs=kafka.secret_configs,
+            passthrough_configs=kafka.passthrough_configs,
+            deserialization=Deserializer._from_proto(
+                proto_deserializer=kafka.deserialization
+            ),
+            authentication_method=BaseAuthentication._from_proto(
+                proto_authentication_method=kafka.authentication_method
+            ),
+            **oneof_args,
+        )

feature_store/feature_sets/base_feature_set.py

@@ -0,0 +1,151 @@
+import re
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List, Optional, Tuple
+
+from frogml._proto.qwak.feature_store.entities.entity_pb2 import (
+    EntityDefinition,
+    EntitySpec,
+)
+from frogml._proto.qwak.feature_store.features.feature_set_pb2 import FeatureSetSpec
+from frogml.feature_store._common.feature_set_utils import generate_key_unique_name
+from frogml.core.feature_store.entities.entity import Entity
+from frogml.core.feature_store.validations.validation_response import (
+    SuccessValidationResponse,
+)
+
+if TYPE_CHECKING:
+    try:
+        import pandas as pd
+    except ImportError:
+        pass
+
+from frogml.core.clients.feature_store import FeatureRegistryClient
+from frogml.core.exceptions import FrogmlException
+from frogml.core.feature_store.validations.validation_options import (
+    FeatureSetValidationOptions,
+)
+
+
+@dataclass
+class BaseFeatureSet(ABC):
+    name: str
+    data_sources: List[str]
+    entity: str = str()
+    key: str = str()
+    repository: Optional[str] = None
+    __instance_module_path__: Optional[str] = None
+
+    def _validate(self):
+        self._validate_feature_set_name()
+
+    def _validate_feature_set_name(self):
+        """
+        Verifies that the name follows these rules (rfc 1123):
+        * A name can start or end with a letter or a number
+        * A name MUST NOT start or end with a '-' (dash)
+        * A name MUST NOT consist of all numeric values
+        * A name can be up to 63 characters
+        """
+        regex = re.compile(r"^(?![_-])\D[a-zA-Z0-9_-]{0,61}[a-zA-Z0-9](?![_-])$")
+        if not regex.match(self.name):
+            raise FrogmlException(
+                f"feature set name: {self.name} is invalid, name must follow these conditions:"
+                f"can start or end with a letter or a number"
+                f"name MUST NOT start or end with a '-' (dash)"
+                f"name MUST NOT consist of all numeric values"
+            )
+
+    def _get_entity_definition(
+        self, feature_registry: FeatureRegistryClient
+    ) -> EntityDefinition:
+        # fs has a key def but not an entity def, return an entity that is valid but not registered
+        # This is done mainly for the validation flow/get sample, to avoid registering temporary entities for keys
+        if self.key and not self.entity:
+            new_key_name: str = generate_key_unique_name(self.name)
+
+            new_key_spec: EntitySpec = Entity(
+                name=new_key_name,
+                description=f"Key of feature set {self.name}",
+                key=self.key,
+            )._to_proto()
+
+            # return new entity def with no id
+            entity_definition: EntityDefinition = EntityDefinition(
+                entity_spec=new_key_spec
+            )
+        else:
+            feature_set_entity = feature_registry.get_entity_by_name(self.entity)
+
+            if not feature_set_entity:
+                raise FrogmlException(
+                    f"Trying to register a feature set with a non existing entity or bad key configuration-: {self.entity}"
+                )
+
+            entity_definition: EntityDefinition = (
+                feature_set_entity.entity.entity_definition
+            )
+
+        return entity_definition
+
+    @staticmethod
+    @abstractmethod
+    def _from_proto(cls, proto: FeatureSetSpec):
+        pass
+
+    @abstractmethod
+    def _to_proto(
+        self,
+        git_commit,
+        features,
+        feature_registry,
+        artifact_url: Optional[str] = None,
+        **kwargs,
+    ) -> Tuple[FeatureSetSpec, Optional[str]]:
+        pass
+
+    def get_sample(
+        self,
+        number_of_rows: int = 10,
+        validation_options: Optional[FeatureSetValidationOptions] = None,
+    ) -> "pd.DataFrame":
+        """
+        Fetches a sample of the Feature set transformation by loading requested sample of data from the data source
+        and executing the transformation on that data.
+
+        :param number_of_rows: number of rows requests
+        :param validation_options: validation options
+        :returns Sample Pandas Dataframe
+
+        Example:
+
+        ... code-block:: python
+            @batch.feature_set(entity="users", data_sources=["snowflake_users_table"])
+            @batch.backfill(start_date=datetime(2022, 1, 1))
+            def user_features():
+                return SparkSqlTransformation("SELECT user_id, age FROM data_source")
+
+            sample_features = user_features.get_sample()
+            print(sample_feature)
+            #	    user_id	         timestamp	        user_features.age
+            # 0	      1	        2021-01-02 17:00:00	              23
+            # 1	      1	        2021-01-01 12:00:00	              51
+            # 2	      2	        2021-01-02 12:00:00	              66
+            # 3	      2	        2021-01-01 18:00:00	              34
+        """
+        from frogml.feature_store.validations.validator import (
+            FeaturesOperatorValidator,
+        )
+
+        v = FeaturesOperatorValidator()
+
+        response, _ = v.validate_featureset(
+            featureset=self,
+            sample_size=number_of_rows,
+            validation_options=validation_options,
+        )
+
+        if isinstance(response, SuccessValidationResponse):
+            return response.sample
+        else:
+            raise FrogmlException(f"Sampling failed: \n{response}")  # noqa: E231