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

qwak/feature_store/feature_sets/streaming_backfill.py

@@ -0,0 +1,223 @@
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import List, Optional, Set, Union
+
+from _qwak_proto.qwak.feature_store.features.execution_pb2 import (
+    ExecutionSpec as ProtoExecutionSpec,
+)
+from _qwak_proto.qwak.execution.v1.streaming_aggregation_pb2 import (
+    StreamingAggregationBackfillIngestion as ProtoStreamingAggregationBackfillIngestion,
+    BackfillDataSource as ProtoBackfillDataSource,
+    TimeRange as ProtoTimeRange,
+)
+from google.protobuf.timestamp_pb2 import Timestamp as ProtoTimestamp
+from qwak.exceptions import QwakException
+from qwak.feature_store._common.artifact_utils import ArtifactSpec, ArtifactsUploader
+from qwak.feature_store.feature_sets.execution_spec import ClusterTemplate
+from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
+
+_BACKFILL_ = "_qwak_backfill_specification"
+
+
+@dataclass
+class BackfillDataSource:
+    data_source_name: str
+    start_datetime: Optional[datetime] = None
+    end_datetime: Optional[datetime] = None
+
+    def __post_init__(self):
+        self._validate()
+
+    def _to_proto(self) -> ProtoBackfillDataSource:
+        start_timestamp: Optional[ProtoTimestamp] = None
+        end_timestamp: Optional[ProtoTimestamp] = None
+
+        if self.end_datetime:
+            end_timestamp = ProtoTimestamp()
+            end_timestamp.FromDatetime(self.end_datetime.astimezone(timezone.utc))
+
+        if self.start_datetime:
+            start_timestamp = ProtoTimestamp()
+            start_timestamp.FromDatetime(self.start_datetime.astimezone(timezone.utc))
+
+        time_range = ProtoTimeRange(
+            start_timestamp=start_timestamp,
+            end_timestamp=end_timestamp,
+        )
+
+        return ProtoBackfillDataSource(
+            data_source_name=self.data_source_name,
+            time_range=time_range,
+        )
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoBackfillDataSource) -> "BackfillDataSource":
+        start_datetime: Optional[datetime] = None
+        end_datetime: Optional[datetime] = None
+
+        time_range: ProtoTimeRange = proto.time_range
+
+        proto_start_timestamp: Optional[ProtoTimestamp] = (
+            time_range.start_timestamp if time_range.start_timestamp else None
+        )
+        proto_end_timestamp: Optional[ProtoTimestamp] = (
+            time_range.end_timestamp if time_range.end_timestamp else None
+        )
+
+        start_datetime = (
+            datetime.fromtimestamp(
+                proto_start_timestamp.seconds + proto_start_timestamp.nanos / 1e9
+            )
+            if proto_start_timestamp
+            else None
+        )
+
+        end_datetime = (
+            datetime.fromtimestamp(
+                proto_end_timestamp.seconds + proto_end_timestamp.nanos / 1e9
+            )
+            if proto_end_timestamp
+            else None
+        )
+
+        return cls(
+            data_source_name=proto.data_source_name,
+            start_datetime=start_datetime,
+            end_datetime=end_datetime,
+        )
+
+    def _validate(self):
+        if self.start_datetime and self.end_datetime:
+            if self.start_datetime >= self.end_datetime:
+                raise QwakException(
+                    f"Backfill data source {self.data_source_name} has invalid time range: "
+                    f"start_datetime {self.start_datetime} is after or equal end_datetime {self.end_datetime}."
+                )
+
+        if not self.data_source_name:
+            raise QwakException(
+                "Backfill data source must have a valid data source name."
+            )
+
+
+@dataclass
+class StreamingBackfill:
+    featureset_name: str
+    start_datetime: datetime
+    end_datetime: datetime
+    data_sources: List[BackfillDataSource]
+    transform: "SparkSqlTransformation"
+    cluster_template: Optional[ClusterTemplate] = ClusterTemplate.SMALL
+
+    def __post_init__(self):
+        if not self.featureset_name:
+            raise QwakException("featureset_name must be provided for backfill.")
+
+        if not self.start_datetime or not self.end_datetime:
+            raise QwakException(
+                "For Streaming backfill, start_datetime and end_datetime are mandatory fields."
+            )
+
+        if self.start_datetime >= self.end_datetime:
+            raise QwakException(
+                f"Backfill has invalid time range: "
+                f"start_datetime {self.start_datetime} is after or equal end_datetime {self.end_datetime}."
+            )
+
+        if not self.data_sources:
+            raise QwakException(
+                "Trying to create a streaming backfill with no data sources. "
+                "At least one data source has to be provided when trying to create a streaming backfill."
+            )
+
+        if type(self.transform) is not SparkSqlTransformation:
+            raise QwakException(
+                "For backfill, only Spark SQL transformation type is currently supported"
+            )
+
+        self._validate_unique_sources()
+
+    def _validate_unique_sources(self):
+        source_names: List[str] = [
+            data_source.data_source_name for data_source in self.data_sources
+        ]
+        duplicates: Set[str] = {
+            item for item in source_names if source_names.count(item) > 1
+        }
+        if duplicates:
+            raise QwakException(
+                f"A specific data source can only appear once per backfill definition. "
+                f"Found these duplicates: {', '.join(set(duplicates))}"
+            )
+
+    def _to_proto(
+        self,
+        original_instance_module_path: str,
+    ) -> ProtoStreamingAggregationBackfillIngestion:
+        artifact_url: Optional[str] = None
+        artifact_spec: Optional[ArtifactSpec] = ArtifactsUploader.get_artifact_spec(
+            transformation=self.transform,
+            featureset_name=f"{self.featureset_name}-backfill",
+            __instance_module_path__=original_instance_module_path,
+        )
+
+        if artifact_spec:
+            artifact_url = ArtifactsUploader.upload(artifact_spec)
+
+        end_timestamp = ProtoTimestamp()
+        end_timestamp.FromDatetime(self.end_datetime.astimezone(timezone.utc))
+
+        start_timestamp = ProtoTimestamp()
+        start_timestamp.FromDatetime(self.start_datetime.astimezone(timezone.utc))
+
+        return ProtoStreamingAggregationBackfillIngestion(
+            featureset_name=self.featureset_name,
+            start_timestamp=start_timestamp,
+            end_timestamp=end_timestamp,
+            execution_spec=ProtoExecutionSpec(
+                cluster_template=ClusterTemplate.to_proto(self.cluster_template)
+            ),
+            transformation=self.transform._to_proto(artifact_path=artifact_url),
+            data_source_specs=[
+                data_source._to_proto() for data_source in self.data_sources
+            ],
+        )
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoStreamingAggregationBackfillIngestion):
+        backfill_data_sources = [
+            BackfillDataSource._from_proto(ds) for ds in proto.data_source_specs
+        ]
+
+        return cls(
+            featureset_name=proto.featureset_name,
+            start_datetime=datetime.fromtimestamp(
+                proto.start_timestamp.seconds + proto.start_timestamp.nanos / 1e9
+            ),
+            end_datetime=datetime.fromtimestamp(
+                proto.end_timestamp.seconds + proto.end_timestamp.nanos / 1e9
+            ),
+            data_sources=backfill_data_sources,
+            transform=SparkSqlTransformation._from_proto(
+                proto.transformation.sql_transformation
+            ),
+            cluster_template=(
+                ClusterTemplate.from_proto(proto.execution_spec.cluster_template)
+                if proto.execution_spec.cluster_template
+                else None
+            ),
+        )
+
+    @staticmethod
+    def _get_normalized_backfill_sources_spec(
+        data_sources: Union[List[str], List[BackfillDataSource]],
+    ) -> List[BackfillDataSource]:
+        # reformat all data source names to 'BackfillDataSource'
+        return [
+            (
+                BackfillDataSource(data_source_name=data_source)
+                if isinstance(data_source, str)
+                else data_source
+            )
+            for data_source in data_sources
+        ]

qwak/feature_store/feature_sets/transformations/__init__.py

@@ -0,0 +1,33 @@
+from qwak.feature_store.feature_sets.transformations.aggregations.aggregations import (
+    QwakAggregation,
+)
+from qwak.feature_store.feature_sets.transformations.aggregations.windows import Window
+from qwak.feature_store.feature_sets.transformations.functions import (
+    Column,
+    Schema,
+    Type,
+    qwak_pandas_udf,
+)
+from qwak.feature_store.feature_sets.transformations.transformations import (
+    BaseTransformation,
+    KoalasTransformation,
+    PandasOnSparkTransformation,
+    PySparkTransformation,
+    SparkSqlTransformation,
+    UdfTransformation,
+)
+
+__all__ = [
+    "BaseTransformation",
+    "UdfTransformation",
+    "KoalasTransformation",
+    "PySparkTransformation",
+    "SparkSqlTransformation",
+    "PandasOnSparkTransformation",
+    "Window",
+    "QwakAggregation",
+    "qwak_pandas_udf",
+    "Column",
+    "Schema",
+    "Type",
+]

qwak/feature_store/feature_sets/transformations/aggregations/aggregations.py

@@ -0,0 +1,425 @@
+from abc import ABC
+from dataclasses import dataclass
+
+from _qwak_proto.qwak.feature_store.features.aggregation_pb2 import (
+    Avg as ProtoAvg,
+    BooleanAnd as ProtoBooleanAnd,
+    BooleanOr as ProtoBooleanOr,
+    Count as ProtoCount,
+    LastDistinctN as ProtoLastDistinctN,
+    LastN as ProtoLastN,
+    Max as ProtoMax,
+    Min as ProtoMin,
+    Percentile as ProtoPercentile,
+    PopulationSTDEV as ProtoPopulationSTDEV,
+    PopulationVariance as ProtoPopulationVariance,
+    SampleSTDEV as ProtoSampleSTDEV,
+    SampleVariance as ProtoSampleVariance,
+    Sum as ProtoSum,
+)
+from qwak.exceptions import QwakException
+
+
+@dataclass
+class BaseAggregate(ABC):
+    """
+    Base class for Qwak supported aggregations
+    """
+
+    def __init__(self, column):
+        self._column: str = column
+        self._alias_name: str = ""
+
+        self.__post_init__()
+
+    def __post_init__(self):
+        if "." in self._column:
+            raise QwakException(
+                f"Aggregation {self.__repr__()} column {self._column} contains an illegal character '.'"
+            )
+
+    @property
+    def _key(self):
+        raise NotImplementedError
+
+    @property
+    def _proto_class(self):
+        raise NotImplementedError
+
+    def alias(self, name):
+        if "." in name:
+            raise QwakException(
+                f"Aggregation {self.__repr__()} alias {name} contains an illegal character '.'"
+            )
+        self._alias_name = name
+        return self
+
+    def has_alias(self):
+        return bool(self._alias_name)
+
+    def to_proto(self):
+        return self._proto_class(field_name=self._column)
+
+    def get_name(self):
+        return self._alias_name if self.has_alias() else f"{self._key}_{self._column}"
+
+
+class SumAggregate(BaseAggregate):
+    _key = "sum"
+    _proto_class = ProtoSum
+
+
+class CountAggregate(BaseAggregate):
+    _key = "count"
+    _proto_class = ProtoCount
+
+
+class MaxAggregate(BaseAggregate):
+    _key = "max"
+    _proto_class = ProtoMax
+
+
+class MinAggregate(BaseAggregate):
+    _key = "min"
+    _proto_class = ProtoMin
+
+
+class AvgAggregate(BaseAggregate):
+    _key = "avg"
+    _proto_class = ProtoAvg
+
+
+class SampleVarianceAggregate(BaseAggregate):
+    _key = "sample_variance"
+    _proto_class = ProtoSampleVariance
+
+
+class PopulationVarianceAggregate(BaseAggregate):
+    _key = "population_variance"
+    _proto_class = ProtoPopulationVariance
+
+
+class SampleSTDEVAggregate(BaseAggregate):
+    _key = "sample_stdev"
+    _proto_class = ProtoSampleSTDEV
+
+
+class PopulationSTDEVAggregate(BaseAggregate):
+    _key = "population_stdev"
+    _proto_class = ProtoPopulationSTDEV
+
+
+class BooleanAndAggregate(BaseAggregate):
+    _key = "boolean_and"
+    _proto_class = ProtoBooleanAnd
+
+
+class BooleanOrAggregate(BaseAggregate):
+    _key = "boolean_or"
+    _proto_class = ProtoBooleanOr
+
+
+class LastNAggregate(BaseAggregate):
+    _key = "last_n"
+    _proto_class = ProtoLastN
+
+    def __init__(self, column: str, n: int):
+        self.n = n
+        super().__init__(column)
+        self.__post_init__()
+
+    def __post_init__(self):
+        self._validate()
+        super().__post_init__()
+
+    def _validate(self):
+        if not (type(self.n) is int and self.n > 0):
+            raise QwakException("p must be a positive integer")
+
+    def to_proto(self):
+        return self._proto_class(field_name=self._column, n=self.n)
+
+    def get_name(self):
+        return self._alias_name if self.has_alias() else f"last_{self.n}_{self._column}"
+
+
+class LastDistinctNAggregate(BaseAggregate):
+    _key = "last_distinct_n"
+    _proto_class = ProtoLastDistinctN
+
+    def __init__(self, column: str, n: int):
+        self.n = n
+        super().__init__(column)
+        self.__post_init__()
+
+    def __post_init__(self):
+        self._validate()
+        super().__post_init__()
+
+    def _validate(self):
+        if not (type(self.n) is int and self.n > 0):
+            raise QwakException("p must be a positive integer")
+
+    def to_proto(self):
+        return self._proto_class(field_name=self._column, n=self.n)
+
+    def get_name(self):
+        return (
+            self._alias_name
+            if self.has_alias()
+            else f"last_distinct_{self.n}_{self._column}"
+        )
+
+
+class PercentileAggregate(BaseAggregate):
+    _key = "percentile"
+    _proto_class = ProtoPercentile
+
+    def __init__(self, column: str, p: int):
+        self.p = p
+        super().__init__(column)
+        self.__post_init__()
+
+    def __post_init__(self):
+        self._validate()
+        super().__post_init__()
+
+    def _validate(self):
+        if not (type(self.p) is int and 0 < self.p < 100):
+            raise QwakException("p must be an integer between 1 and 99")
+
+    def to_proto(self):
+        return self._proto_class(field_name=self._column, p=self.p)
+
+    def get_name(self):
+        return self._alias_name if self.has_alias() else f"p{self.p}_{self._column}"
+
+
+class QwakAggregation:
+    """
+    Holding all Qwak native supported aggregations. Currently, can be used in Streaming feature sets only.
+    """
+
+    @staticmethod
+    def sum(column):
+        """
+        A sum aggregation.
+
+        Example:
+        >>> from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
+        >>> from qwak.feature_store.feature_sets.transformations import QwakAggregation
+        >>>
+        >>> SparkSqlTransformation(sql="SELECT id, amount, event_timestamp, offset" \
+        >>>                                      "topic, partition FROM transaction_stream")\
+        >>>     .aggregate(QwakAggregation.sum("amount"))\
+        >>>     .by_windows("1 day"),
+        """
+        return SumAggregate(column)
+
+    @staticmethod
+    def count(column):
+        """
+        A count aggregation.
+
+        Example:
+        >>> from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
+        >>> from qwak.feature_store.feature_sets.transformations import QwakAggregation
+        >>>
+        >>> SparkSqlTransformation(sql="SELECT id, amount, event_timestamp, offset" \
+        >>>                                      "topic, partition FROM transaction_stream")\
+        >>>     .aggregate(QwakAggregation.count("amount"))\
+        >>>     .by_windows("1 day"),
+        """
+        return CountAggregate(column)
+
+    @staticmethod
+    def max(column):
+        """
+        A max aggregation.
+
+        Example:
+        >>> from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
+        >>> from qwak.feature_store.feature_sets.transformations import QwakAggregation
+        >>>
+        >>> SparkSqlTransformation(sql="SELECT id, amount, event_timestamp, offset" \
+        >>>                                      "topic, partition FROM transaction_stream")\
+        >>>     .aggregate(QwakAggregation.max("amount"))\
+        >>>     .by_windows("1 day"),
+        """
+        return MaxAggregate(column)
+
+    @staticmethod
+    def min(column):
+        """
+        A min aggregation.
+
+        Example:
+        >>> from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
+        >>> from qwak.feature_store.feature_sets.transformations import QwakAggregation
+        >>>
+        >>> SparkSqlTransformation(sql="SELECT id, amount, event_timestamp, offset" \
+        >>>                                      "topic, partition FROM transaction_stream")\
+        >>>     .aggregate(QwakAggregation.min("amount"))\
+        >>>     .by_windows("1 day"),
+        """
+        return MinAggregate(column)
+
+    @staticmethod
+    def avg(column):
+        """
+        A avarage aggregation.
+
+        Example:
+        >>> from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
+        >>> from qwak.feature_store.feature_sets.transformations import QwakAggregation
+        >>>
+        >>> SparkSqlTransformation(sql="SELECT id, amount, event_timestamp, offset" \
+        >>>                                      "topic, partition FROM transaction_stream")\
+        >>>     .aggregate(QwakAggregation.avg("amount"))\
+        >>>     .by_windows("1 day"),
+        """
+        return AvgAggregate(column)
+
+    @staticmethod
+    def sample_variance(column):
+        """
+        A sample variance aggregation.
+
+        Example:
+        >>> from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
+        >>> from qwak.feature_store.feature_sets.transformations import QwakAggregation
+        >>>
+        >>> SparkSqlTransformation(sql="SELECT id, amount, event_timestamp, offset" \
+        >>>                                      "topic, partition FROM transaction_stream")\
+        >>>     .aggregate(QwakAggregation.sample_variance("amount"))\
+        >>>     .by_windows("1 day"),
+        """
+        return SampleVarianceAggregate(column)
+
+    @staticmethod
+    def population_variance(column):
+        """
+        A population variance aggregation.
+
+        Example:
+        >>> from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
+        >>> from qwak.feature_store.feature_sets.transformations import QwakAggregation
+        >>>
+        >>> SparkSqlTransformation(sql="SELECT id, amount, event_timestamp, offset" \
+        >>>                                      "topic, partition FROM transaction_stream")\
+        >>>     .aggregate(QwakAggregation.population_variance("amount"))\
+        >>>     .by_windows("1 day"),
+        """
+        return PopulationVarianceAggregate(column)
+
+    @staticmethod
+    def sample_stdev(column):
+        """
+        A sample STDEV aggregation.
+
+        Example:
+        >>> from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
+        >>> from qwak.feature_store.feature_sets.transformations import QwakAggregation
+        >>>
+        >>> SparkSqlTransformation(sql="SELECT id, amount, event_timestamp, offset" \
+        >>>                                      "topic, partition FROM transaction_stream")\
+        >>>     .aggregate(QwakAggregation.sample_stdev("amount"))\
+        >>>     .by_windows("1 day"),
+        """
+        return SampleSTDEVAggregate(column)
+
+    @staticmethod
+    def population_stdev(column):
+        """
+        A population STDEV aggregation
+
+        Example:
+        >>> from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
+        >>> from qwak.feature_store.feature_sets.transformations import QwakAggregation
+        >>>
+        >>> SparkSqlTransformation(sql="SELECT id, amount, event_timestamp, offset" \
+        >>>                                      "topic, partition FROM transaction_stream")\
+        >>>     .aggregate(QwakAggregation.population_stdev("amount"))\
+        >>>     .by_windows("1 day"),
+        """
+        return PopulationSTDEVAggregate(column)
+
+    @staticmethod
+    def boolean_and(column):
+        """
+        A BOOLEAN AND aggregation.
+
+        Example:
+        >>> from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
+        >>> from qwak.feature_store.feature_sets.transformations import QwakAggregation
+        >>>
+        >>> SparkSqlTransformation(sql="SELECT id, amount, event_timestamp, offset" \
+        >>>                                      "topic, partition FROM transaction_stream")\
+        >>>     .aggregate(QwakAggregation.boolean_and("is_fraud"))\
+        >>>     .by_windows("1 day"),
+        """
+        return BooleanAndAggregate(column)
+
+    @staticmethod
+    def boolean_or(column):
+        """
+        A BOOLEAN OR aggregation.
+
+        Example:
+        >>> from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
+        >>> from qwak.feature_store.feature_sets.transformations import QwakAggregation
+        >>>
+        >>> SparkSqlTransformation(sql="SELECT id, amount, event_timestamp, offset" \
+        >>>                                      "topic, partition FROM transaction_stream")\
+        >>>     .aggregate(QwakAggregation.boolean_or("is_fraud"))\
+        >>>     .by_windows("1 day"),
+        """
+        return BooleanOrAggregate(column)
+
+    @staticmethod
+    def last_n(column, n):
+        """
+        A last_n aggregation
+
+        Example:
+        >>> from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
+        >>> from qwak.feature_store.feature_sets.transformations import QwakAggregation
+        >>>
+        >>> SparkSqlTransformation(sql="SELECT id, amount, event_timestamp, offset" \
+        >>>                                      "topic, partition FROM transaction_stream")\
+        >>>     .aggregate(QwakAggregation.last_n("amount", n=10))\
+        >>>     .by_windows("1 day"),
+        """
+        return LastNAggregate(column, n=n)
+
+    @staticmethod
+    def last_distinct_n(column, n):
+        """
+        A last_distinct_n aggregation
+
+        Example:
+        >>> from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
+        >>> from qwak.feature_store.feature_sets.transformations import QwakAggregation
+        >>>
+        >>> SparkSqlTransformation(sql="SELECT id, amount, event_timestamp, offset" \
+        >>>                                      "topic, partition FROM transaction_stream")\
+        >>>     .aggregate(QwakAggregation.last_distinct_n("amount", n=10))\
+        >>>     .by_windows("1 day"),
+        """
+        return LastDistinctNAggregate(column, n=n)
+
+    @staticmethod
+    def percentile(column, p):
+        """
+        Calculates a percentile - e.g., p50 (aka median), p90 etc.
+
+        Example:
+        >>> from qwak.feature_store.feature_sets.transformations import SparkSqlTransformation
+        >>> from qwak.feature_store.feature_sets.transformations import QwakAggregation
+        >>>
+        >>> SparkSqlTransformation(sql="SELECT id, amount, event_timestamp, offset" \
+        >>>                                      "topic, partition FROM transaction_stream")\
+        >>>     .aggregate(QwakAggregation.percentile("amount", p=50).alias("median_amount"))\
+        >>>     .by_windows("1 day"),
+        """
+        return PercentileAggregate(column, p=p)