qwak-core 0.1.37__py3-none-any.whl → 0.5.19__py3-none-any.whl

qwak/feature_store/feature_sets/batch.py

@@ -1,16 +1,17 @@
-import ast
 import functools
 import inspect
 from dataclasses import dataclass
 from datetime import datetime
 from inspect import Signature
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
-
-from _qwak_proto.qwak.features_operator.v3.features_operator_async_service_pb2 import (
-    FeatureSetValidationOptions as ProtoFeatureSetValidationOptions,
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
+
+from qwak.feature_store._common.artifact_utils import ArtifactSpec, ArtifactsUploader
+from qwak.feature_store._common.feature_set_utils import get_batch_source_for_featureset
+from qwak.feature_store.feature_sets._utils._featureset_utils import FeaturesetUtils
+from qwak.feature_store.feature_sets.base_feature_set import BaseFeatureSet
+from qwak.feature_store.validations.validation_options import (
+    FeatureSetValidationOptions,
 )
-from qwak.feature_store._common.functions import print_validation_outputs
-from qwak.feature_store.validation_options import FeatureSetValidationOptions
 
 if TYPE_CHECKING:
     try:
@@ -27,26 +28,23 @@ from _qwak_proto.qwak.feature_store.features.feature_set_types_pb2 import (
     FeatureSetBatchSource as ProtoFeatureSetBatchSource,
     FeatureSetType as ProtoFeatureSetType,
 )
-from _qwak_proto.qwak.features_operator.v3.features_operator_pb2 import (
-    ValidationSuccessResponse,
-)
-from qwak.clients.feature_store import FeatureRegistryClient
-from qwak.clients.feature_store.operator_client import FeaturesOperatorClient
-from qwak.exceptions import QwakException
 from qwak.feature_store.entities.entity import Entity
 from qwak.feature_store.feature_sets.backfill import Backfill
 from qwak.feature_store.feature_sets.context import Context
 from qwak.feature_store.feature_sets.execution_spec import ClusterTemplate
-from qwak.feature_store.feature_sets.metadata import Metadata
+from qwak.feature_store.feature_sets.metadata import (
+    Metadata,
+    get_metadata_from_function,
+    set_metadata_on_function,
+)
 from qwak.feature_store.feature_sets.read_policies import ReadPolicy, ReadPolicyABC
-from qwak.feature_store.feature_sets.transformations import BatchTransformation
+from qwak.feature_store.feature_sets.transformations import BaseTransformation
 from typeguard import typechecked
 
 # decorator attributes
 _BACKFILL_POLICY_ATTRIBUTE = "_qwak_backfill_policy"
 _EXECUTION_SPECIFICATION_ATTRIBUTE = "_qwak_execution_specification"
 _SCHEDULING_POLICY_ATTRIBUTE = "_qwak_scheduling_policy"
-_METADATA_ATTRIBUTE = "_qwak_metadata"
 
 # Default timestamp column name for TimeFrame
 _DEFAULT_TIMEFRAME_TS_COL_NAME = "qwak_window_end_ts"
@@ -54,16 +52,20 @@ _DEFAULT_TIMEFRAME_TS_COL_NAME = "qwak_window_end_ts"
 
 def feature_set(
     *,
-    entity: str,
     data_sources: Union[Dict[str, ReadPolicyABC], List[str]],
     timestamp_column_name: str = None,
     name: str = None,
+    entity: Optional[str] = None,
+    key: Optional[str] = None,
+    repository: Optional[str] = None,
 ):
     """
     Define a batch scheduled feature set. Default scheduling policy is every 4 hours.
 
     :param name: The name of the feature set. If not defined, taken from the function name
-    :type entity: a string reference to a Qwak entity, or the entity definition itself
+    :param entity: a string reference to a Qwak entity, or the entity definition itself.
+                  Only one of, entity or key, can be provided.
+    :param key: a column name in the feature set which is the key. Only one of, entity or key, can be provided.
     :param data_sources: a string reference to a Qwak data source, or the data source definition itself
     :param timestamp_column_name: Timestamp column the feature set should consider when ingesting records.
     If more than one data source is specified - the field is mandatory. If only one data source is specified - the
@@ -71,11 +73,17 @@ def feature_set(
 
     Example:
 
-    .. code-block:: python
+    ... code-block:: python
         @batch.feature_set(entity="users", data_sources=["snowflake_users_table"])
         def user_features():
             return SparkSqlTransformation("SELECT user_id, age FROM data_source")
 
+    Example with key:
+
+    ... code-block:: python
+    @batch.feature_set(key="user_id", data_sources=["snowflake_users_table"])
+    def user_features():
+        return SparkSqlTransformation("SELECT user_id, age FROM data_source")
     """
 
     def decorator(function):
@@ -84,28 +92,25 @@ def feature_set(
             user_transformation = function(context=Context())
         else:
             user_transformation = function()
-        if not isinstance(user_transformation, BatchTransformation):
-            raise ValueError(
-                "Function must return a valid batch transformation function"
-            )
+
+        FeaturesetUtils.validate_base_featureset_decorator(
+            user_transformation=user_transformation, entity=entity, key=key
+        )
 
         fs_name = name or function.__name__
-        batch_feature_set = BatchFeatureSetV1(
+        batch_feature_set = BatchFeatureSet(
             name=fs_name,
-            entity=entity,
+            entity=entity if entity else None,
+            key=key if key else None,
+            repository=repository if repository else None,
             data_sources=data_sources,
             timestamp_column_name=timestamp_column_name,
             transformation=user_transformation,
             scheduling_policy=getattr(
                 function, _SCHEDULING_POLICY_ATTRIBUTE, "0 */4 * * *"
             ),
-            metadata=getattr(
-                function,
-                _METADATA_ATTRIBUTE,
-                Metadata(
-                    description=fs_name,
-                    display_name=fs_name,
-                ),
+            metadata=get_metadata_from_function(
+                function, description=fs_name, display_name=fs_name
             ),
             cluster_template=getattr(
                 function,
@@ -147,7 +152,7 @@ def scheduling(*, cron_expression: Optional[str]):
 
     Example:
 
-    .. code-block:: python
+    ... code-block:: python
 
         @batch.feature_set(entity="users", data_sources=["snowflake_users_table"])
         @batch.scheduling(cron_expression="0 8 * * *")
@@ -167,17 +172,23 @@ def scheduling(*, cron_expression: Optional[str]):
 
 
 @typechecked
-def metadata(*, owner: str, description: str, display_name: str):
+def metadata(
+    *,
+    owner: Optional[str] = None,
+    description: Optional[str] = None,
+    display_name: Optional[str] = None,
+    version_comment: Optional[str] = None,
+):
     """
     Sets additional user provided metadata
 
     :param owner: feature set owner
     :param description: General description of the feature set
     :param display_name: Human readable name of the feature set
-
+    :param version_comment: Comment which describes the version
     Example:
 
-    .. code-block:: python
+    ... code-block:: python
 
         @batch.feature_set(entity="users", data_sources=["snowflake_users_table"])
         @batch.metadata(
@@ -192,10 +203,12 @@ def metadata(*, owner: str, description: str, display_name: str):
 
     def decorator(user_transformation):
         _validate_decorator_ordering(user_transformation)
-        setattr(
-            user_transformation,
-            _METADATA_ATTRIBUTE,
-            Metadata(owner=owner, description=description, display_name=display_name),
+        set_metadata_on_function(
+            function=user_transformation,
+            owner=owner,
+            description=description,
+            display_name=display_name,
+            version_comment=version_comment,
         )
 
         return user_transformation
@@ -216,7 +229,7 @@ def backfill(
 
     Example:
 
-    .. code-block:: python
+    ... code-block:: python
         @batch.feature_set(entity="users", data_sources=["snowflake_users_table"])
         @batch.backfill(start_date=datetime(2022, 1, 1))
         def user_features():
@@ -240,7 +253,7 @@ def backfill(
 @typechecked
 def execution_specification(
     *,
-    cluster_template: ClusterTemplate = None,
+    cluster_template: ClusterTemplate,
 ):
     """
     Set the execution specification of the cluster running the feature set
@@ -249,7 +262,7 @@ def execution_specification(
 
     Cluster template example:
 
-    .. code-block:: python
+    ... code-block:: python
         @batch.feature_set(entity="users", data_sources=["snowflake_users_table"])
         @batch.execution_specification(cluster_template=ClusterTemplate.MEDIUM)
         def user_features():
@@ -269,50 +282,26 @@ def execution_specification(
 
 
 def _validate_decorator_ordering(user_transformation):
-    if isinstance(user_transformation, BatchFeatureSetV1):
+    if isinstance(user_transformation, BatchFeatureSet):
         raise ValueError(
-            "Wrong decorator ordering - @batch.feature_set should be to top most decorator"
+            "Wrong decorator ordering - @batch.feature_set should be the top most decorator"
         )
 
 
 @dataclass
-class BatchFeatureSetV1:
-    name: str
-    entity: str
-    data_sources: List[str]
-    timestamp_column_name: str = None
-    scheduling_policy: str = None
-    transformation: BatchTransformation = None
-    metadata: Metadata = None
+class BatchFeatureSet(BaseFeatureSet):
+    timestamp_column_name: str = str()
+    scheduling_policy: str = str()
+    transformation: Optional[BaseTransformation] = None
+    metadata: Optional[Metadata] = None
     cluster_template: Optional[ClusterTemplate] = None
     backfill: Optional[Backfill] = None
-    __instance_module_path__: Optional[str] = None
-
-    def _get_entity_definition(self, feature_registry):
-        feature_set_entity = feature_registry.get_entity_by_name(self.entity)
-        if not feature_set_entity:
-            raise QwakException(
-                f"Trying to register a feature set with a non existing entity -: {self.entity}"
-            )
-        return feature_set_entity.entity.entity_definition
 
-    def _get_batch_source(self, batch_ds_name, feature_registry):
-        batch_ds = feature_registry.get_data_source_by_name(batch_ds_name)
-        if not batch_ds:
-            raise QwakException(
-                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 batch feature sets with batch sources. Source {batch_ds_name} is of type {ds_spec.WhichOneof('type')}"
-            )
-
-        return ds_spec.batch_source
+    def __post_init__(self):
+        self._validate()
 
     @classmethod
-    def from_proto(cls, proto: FeatureSetSpec):
+    def _from_proto(cls, proto: FeatureSetSpec):
         batch_v1_def = proto.feature_set_type.batch_feature_set_v1
 
         return cls(
@@ -321,16 +310,18 @@ class BatchFeatureSetV1:
             data_sources=[
                 ds.data_source.name for ds in batch_v1_def.feature_set_batch_sources
             ],
+            repository=proto.featureset_repository_name,
             timestamp_column_name=batch_v1_def.timestamp_column_name,
             scheduling_policy=batch_v1_def.scheduling_policy,
-            transformation=BatchTransformation._from_proto(batch_v1_def.transformation),
+            transformation=BaseTransformation._from_proto(batch_v1_def.transformation),
             metadata=Metadata.from_proto(proto.metadata),
-            cluster_template=ClusterTemplate.from_proto(batch_v1_def.execution_spec),
+            cluster_template=ClusterTemplate.from_cluster_template_number(
+                batch_v1_def.execution_spec.cluster_template
+            ),
             backfill=Backfill.from_proto(batch_v1_def.backfill),
         )
 
     def _get_data_sources(self, feature_registry) -> List[ProtoFeatureSetBatchSource]:
-        actual_data_sources = {}
         if isinstance(self.data_sources, list):
             actual_data_sources = {
                 ds_name: ReadPolicy.NewOnly for ds_name in self.data_sources
@@ -340,10 +331,12 @@ class BatchFeatureSetV1:
 
         return [
             ProtoFeatureSetBatchSource(
-                data_source=self._get_batch_source(ds, feature_registry),
-                read_policy=read_policy().to_proto()
-                if inspect.isclass(read_policy)
-                else read_policy.to_proto(),
+                data_source=get_batch_source_for_featureset(ds, feature_registry),
+                read_policy=(
+                    read_policy().to_proto()
+                    if inspect.isclass(read_policy)
+                    else read_policy.to_proto()
+                ),
             )
             for ds, read_policy in actual_data_sources.items()
         ]
@@ -370,38 +363,60 @@ class BatchFeatureSetV1:
         # if we got here - only one data source is defined
         return data_sources[0].data_source.date_created_column
 
+    def _validate(self):
+        super()._validate()
+
     def _to_proto(
         self,
         git_commit,
         features,
         feature_registry,
-    ) -> FeatureSetSpec:
+        artifact_url: Optional[str] = str(),
+        **kwargs,
+    ) -> Tuple[FeatureSetSpec, Optional[str]]:
+        self._validate()
+
         data_sources = self._get_data_sources(feature_registry)
 
-        return FeatureSetSpec(
-            name=self.name,
-            metadata=self.metadata.to_proto(),
-            git_commit=git_commit,
-            features=features,
-            entity=self._get_entity_definition(feature_registry),
-            feature_set_type=ProtoFeatureSetType(
-                batch_feature_set_v1=ProtoBatchFeatureSetV1(
-                    online_sink=True,
-                    offline_sink=True,
-                    timestamp_column_name=self._get_timestamp_column(data_sources),
-                    scheduling_policy=self.scheduling_policy,
-                    feature_set_batch_sources=data_sources,
-                    execution_spec=ProtoExecutionSpec(
-                        cluster_template=ClusterTemplate.to_proto(self.cluster_template)
-                        if self.cluster_template
-                        else None
-                    ),
-                    transformation=self.transformation._to_proto()
-                    if self.transformation
-                    else None,
-                    backfill=self.backfill.to_proto() if self.backfill else None,
-                )
+        if not artifact_url:
+            artifact: Optional[ArtifactSpec] = ArtifactsUploader.get_artifact_spec(
+                self.transformation, self.name, self.__instance_module_path__
+            )
+            if artifact:
+                artifact_url = ArtifactsUploader.upload(artifact)
+
+        return (
+            FeatureSetSpec(
+                name=self.name,
+                metadata=self.metadata.to_proto(),
+                git_commit=git_commit,
+                features=features,
+                entity=self._get_entity_definition(feature_registry),
+                featureset_repository_name=self.repository,
+                feature_set_type=ProtoFeatureSetType(
+                    batch_feature_set_v1=ProtoBatchFeatureSetV1(
+                        online_sink=True,
+                        offline_sink=True,
+                        timestamp_column_name=self._get_timestamp_column(data_sources),
+                        scheduling_policy=self.scheduling_policy,
+                        feature_set_batch_sources=data_sources,
+                        execution_spec=ProtoExecutionSpec(
+                            cluster_template=(
+                                ClusterTemplate.to_proto(self.cluster_template)
+                                if self.cluster_template
+                                else None
+                            )
+                        ),
+                        transformation=(
+                            self.transformation._to_proto(artifact_path=artifact_url)
+                            if self.transformation
+                            else None
+                        ),
+                        backfill=self.backfill.to_proto() if self.backfill else None,
+                    )
+                ),
             ),
+            artifact_url,
         )
 
     def get_sample(
@@ -414,12 +429,12 @@ class BatchFeatureSetV1:
         and executing the transformation on that data.
 
         :param number_of_rows: number of rows requests
-        :param validation_options: vlidation options
+        :param validation_options: validation options
         :returns Sample Pandas Dataframe
 
-        Example:
+        Example with entity:
 
-        .. code-block:: python
+        ... code-block:: python
             @batch.feature_set(entity="users", data_sources=["snowflake_users_table"])
             @batch.backfill(start_date=datetime(2022, 1, 1))
             def user_features():
@@ -432,47 +447,23 @@ class BatchFeatureSetV1:
             # 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
-        """
-
-        """
-        Fetches a sample of the Featureset transformation by loading requested sample of data from the data source
-        and executing the transformation on that data.
-        """
-
-        try:
-            import pandas as pd
-        except ImportError:
-            raise QwakException("Missing Pandas dependency required for getting sample")
 
-        if 0 >= number_of_rows > 1000:
-            raise ValueError(
-                f"`number_rows` must be under 1000 and positive, got: {number_of_rows}"
-            )
+        Example with key:
 
-        operator_client = FeaturesOperatorClient()
-        registry_client = FeatureRegistryClient()
-
-        featureset_spec = self._to_proto(
-            git_commit=None, features=None, feature_registry=registry_client
-        )
-        validation_options_proto: Optional[ProtoFeatureSetValidationOptions] = (
-            validation_options.to_proto() if validation_options else None
-        )
+        ... code-block:: python
+            @batch.feature_set(key="user_id", 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")
 
-        result = operator_client.validate_featureset_blocking(
-            featureset_spec=featureset_spec,
-            resource_path=None,
-            num_samples=number_of_rows,
-            validation_options=validation_options_proto,
+            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
+        """
+        return super().get_sample(
+            number_of_rows=number_of_rows, validation_options=validation_options
         )
-
-        response = getattr(result, result.WhichOneof("type"))
-        if isinstance(response, ValidationSuccessResponse):
-            print_validation_outputs(response)
-
-            return pd.read_json(
-                path_or_buf=ast.literal_eval(response.sample),
-                dtype=response.spark_column_description,
-            )
-        else:
-            raise QwakException(f"Sampling failed:\n{response}")

qwak/feature_store/feature_sets/execution_spec.py

@@ -44,6 +44,12 @@ class ClusterTemplate(AutoName):
             ClusterTemplate._proto_to_cluster_template.value[proto.cluster_template]
         )
 
+    @classmethod
+    def from_cluster_template_number(cls, cluster_template_number: int):
+        return cls(
+            ClusterTemplate._proto_to_cluster_template.value[cluster_template_number]
+        )
+
     @staticmethod
     def to_proto(template):
         if not template:

qwak/feature_store/feature_sets/metadata.py

@@ -2,22 +2,72 @@ from dataclasses import dataclass
 
 from _qwak_proto.qwak.feature_store.features.feature_set_pb2 import (
     FeatureSetUserMetadata as ProtoMetadata,
+    FeaturesetVersionUserMetadata as ProtoFeaturesetVersionUserMetadata,
 )
 
+_METADATA_ATTRIBUTE = "_qwak_metadata"
+
+
+def set_metadata_on_function(
+    function,
+    owner: str = "",
+    description: str = "",
+    display_name: str = "",
+    version_comment="",
+):
+    setattr(
+        function,
+        _METADATA_ATTRIBUTE,
+        Metadata(
+            owner=owner,
+            description=description,
+            display_name=display_name,
+            version_comment=version_comment,
+        ),
+    )
+
+
+def get_metadata_from_function(
+    function,
+    owner: str = "",
+    description: str = "",
+    display_name: str = "",
+    version_comment="",
+):
+    return getattr(
+        function,
+        _METADATA_ATTRIBUTE,
+        Metadata(
+            owner=owner,
+            description=description,
+            display_name=display_name,
+            version_comment=version_comment,
+        ),
+    )
+
 
 @dataclass
 class Metadata:
     owner: str = ""
     description: str = ""
     display_name: str = ""
+    version_comment: str = ""
 
     @staticmethod
     def from_proto(metadata: ProtoMetadata):
-        return Metadata(metadata.owner, metadata.description, metadata.display_name)
+        return Metadata(
+            owner=metadata.owner,
+            description=metadata.description,
+            display_name=metadata.display_name,
+            version_comment=metadata.featureset_version_user_metadata.comment,
+        )
 
     def to_proto(self):
         return ProtoMetadata(
             owner=self.owner,
             description=self.description,
             display_name=self.display_name,
+            featureset_version_user_metadata=ProtoFeaturesetVersionUserMetadata(
+                comment=self.version_comment
+            ),
         )