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

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

@@ -0,0 +1,70 @@
+import re
+from dataclasses import dataclass
+from enum import Enum
+from typing import List, Type
+
+from _qwak_proto.qwak.feature_store.features.aggregation_pb2 import (
+    TimeUnit as ProtoTimeUnit,
+)
+from qwak.exceptions import QwakException
+
+
+class _WindowTimeUnitMapping(Enum):
+    WEEKS = (ProtoTimeUnit.TIME_UNIT_WEEKS, 7 * 24 * 60 * 60, "w")
+    DAYS = (ProtoTimeUnit.TIME_UNIT_DAYS, 24 * 60 * 60, "d")
+    HOURS = (ProtoTimeUnit.TIME_UNIT_HOURS, 60 * 60, "h")
+    MINUTES = (ProtoTimeUnit.TIME_UNIT_MINUTES, 60, "m")
+    SECONDS = (ProtoTimeUnit.TIME_UNIT_SECONDS, 1, "s")
+
+    # Aliases
+    WEEK = WEEKS
+    DAY = DAYS
+    HOUR = HOURS
+    MINUTE = MINUTES
+    SECOND = SECONDS
+
+
+@dataclass
+class Window:
+    length: int
+    time_unit_proto: ProtoTimeUnit
+    seconds_in_time_unit: int
+    suffix: str
+
+    def get_feature_suffix(self) -> str:
+        return f"{self.length}{self.suffix}"
+
+    @classmethod
+    def _from_string(cls, *windows: str) -> List[Type["Window"]]:
+        time_window_parser = re.compile(
+            rf'(?P<length>\d+)\s+(?P<timeunit>{"|".join([name.lower() for name, member in _WindowTimeUnitMapping.__members__.items()])})$'
+        )
+
+        parsed_windows: List[Window] = list()
+        for window in windows:
+            match = time_window_parser.match(window.lower())
+            if not match:
+                raise QwakException(
+                    f"""
+                        Could not parse given time window: {window}.
+                        Valid time windows should be of the form "X weeks/days/hours/minutes/seconds", where X is a number.
+                        For example - "1 hour", "2 days", "1 week", etc.
+                    """
+                )
+
+            m_dict = match.groupdict()
+            proto, minutes_in_unit, suffix = getattr(
+                _WindowTimeUnitMapping, m_dict["timeunit"].upper()
+            ).value
+
+            additional_window = cls(
+                length=int(m_dict["length"]),
+                seconds_in_time_unit=minutes_in_unit,
+                time_unit_proto=proto,
+                suffix=suffix,
+            )
+
+            if additional_window not in parsed_windows:
+                parsed_windows.append(additional_window)
+
+        return parsed_windows

qwak/feature_store/feature_sets/transformations/functions/__init__.py

@@ -0,0 +1,10 @@
+from qwak.feature_store.feature_sets.transformations.functions.qwak_pandas import (
+    qwak_pandas_udf,
+)
+from qwak.feature_store.feature_sets.transformations.functions.schema import (
+    Column,
+    Schema,
+    Type,
+)
+
+__all__ = ["qwak_pandas_udf", "Column", "Schema", "Type"]

qwak/feature_store/feature_sets/transformations/functions/qwak_pandas.py

@@ -0,0 +1,64 @@
+import inspect as inspect
+from typing import Callable
+
+from qwak.exceptions import QwakException
+from qwak.feature_store.feature_sets.transformations.functions.schema import Schema
+from qwak.tools.logger import get_qwak_logger
+
+logger = get_qwak_logger()
+
+
+def _get_schema_string(user_function: Callable, schema: Schema) -> str:
+    import pandas as pd
+
+    input_signature = inspect.signature(user_function)
+    return_annotation = input_signature.return_annotation
+    if return_annotation is inspect._empty:
+        raise QwakException(
+            f"return type annotation missing for {user_function.__name__}"
+        )
+
+    if return_annotation is pd.DataFrame:
+        # return type is a Pandas DataFrame - all cols must be named
+        if any([c for c in schema.columns if not c.name]):
+            raise QwakException(
+                f"pandas UDFs that return pd.DataFrame must have named columns."
+                f" Please specify the name in the schema constructor"
+                f" the Schema constructor? {user_function.__name__}"
+            )
+        return schema.to_ddl()
+
+    if return_annotation is pd.Series:
+        # return type is a series (meaning the UDF returns a single colum)
+        if len(schema.columns) != 1:
+            raise QwakException(
+                f"Exactly 1 return type must be set for"
+                f" a pandas udf returning pd.Series,"
+                f" Function {user_function.__name__} returning pd.Series must return exactly 1 return type."
+            )
+        return_col = schema.columns[0]
+        if return_col.name:
+            logger.warning(
+                f"return column name specified for {user_function.__name__},"
+                f" however it will be disregarded as it returns pd.Series"
+            )
+        return return_col.type.to_ddl()
+
+    raise QwakException(
+        f"missing or unrecognized type annotation for {user_function.__name__}"
+    )
+
+
+def qwak_pandas_udf(output_schema: Schema) -> Callable:
+    if not isinstance(output_schema, Schema):
+        raise QwakException("Wrong schema type set for qwak_pandas_udf")
+
+    def _qwak_pandas_udf(user_function: Callable) -> Callable:
+        setattr(
+            user_function,
+            "_schema",
+            _get_schema_string(user_function=user_function, schema=output_schema),
+        )
+        return user_function
+
+    return _qwak_pandas_udf

qwak/feature_store/feature_sets/transformations/functions/schema.py

@@ -0,0 +1,57 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import List, Optional
+
+
+class DataType(ABC):
+    @abstractmethod
+    def to_ddl(self) -> str:
+        pass
+
+
+@dataclass
+class SimpleType(DataType):
+    type_name: str
+
+    def to_ddl(self) -> str:
+        return self.type_name
+
+
+@dataclass
+class ArrayType(DataType):
+    type: DataType
+
+    def to_ddl(self) -> str:
+        return f"array<{self.type.to_ddl()}>"
+
+
+class Type:
+    array = ArrayType
+    int = SimpleType("integer")
+    long = SimpleType("long")
+    string = SimpleType("string")
+    double = SimpleType("double")
+    decimal = SimpleType("decimal")
+    float = SimpleType("float")
+    boolean = SimpleType("boolean")
+    timestamp = SimpleType("timestamp")
+    date = SimpleType("date")
+
+
+@dataclass
+class Column:
+    type: DataType
+    name: Optional[str] = None
+
+    def to_ddl(self) -> str:
+        if self.name:
+            return f"{self.name} {self.type.to_ddl()}"
+        return self.type.to_ddl()
+
+
+@dataclass
+class Schema:
+    columns: List[Column]
+
+    def to_ddl(self) -> str:
+        return ", ".join([column.to_ddl() for column in self.columns])

qwak/feature_store/feature_sets/transformations/transformations.py

@@ -0,0 +1,451 @@
+import sys
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Any
+from warnings import warn
+
+if TYPE_CHECKING:
+    import pyspark.pandas as psd
+    import pyspark.sql as psql
+
+import importlib
+
+from _qwak_proto.qwak.feature_store.features.aggregation_pb2 import (
+    AggregationField as ProtoAggregationField,
+    TimeWindow as ProtoTimeWindow,
+)
+from _qwak_proto.qwak.feature_store.features.feature_set_types_pb2 import (
+    KoalasTransformation as ProtoKoalasTransformation,
+    PandasOnSparkTransformation as ProtoPandasOnSparkTransformation,
+    PySparkTransformation as ProtoPySparkTransformation,
+    SqlTransformation as ProtoSqlTransformation,
+    TransformArguments as ProtoTransformArguments,
+    Transformation as ProtoTransformation,
+    UdfTransformation as ProtoUdfTransformation,
+)
+from qwak.exceptions import QwakException
+from qwak.feature_store._common.value import (
+    UPDATE_QWAK_SDK_WITH_FEATURE_STORE_EXTRA_MSG,
+    missing_pyspark_exception_message,
+)
+from qwak.feature_store.feature_sets.transformations.aggregations.aggregations import (
+    BaseAggregate,
+    QwakAggregation,
+)
+from qwak.feature_store.feature_sets.transformations.aggregations.windows import Window
+from qwak.feature_store.feature_sets.transformations.validations.validations_util import (
+    validate_function,
+    validate_qwargs,
+)
+from qwak.tools.logger import logger as qwak_logger
+
+logger = qwak_logger.get_qwak_logger()
+
+
+@dataclass
+class BaseTransformation(ABC):
+    """
+    Base class for Qwak transformations.
+
+    Windows and aggregates are supported only for streaming aggregation featuresets.
+    """
+
+    windows: List[Window] = field(init=False, default_factory=list)
+    aggregations: Tuple[BaseAggregate] = field(init=False, default_factory=tuple)
+    _artifact_path: Optional[str] = field(init=False, default=None)
+
+    def aggregate(self, *aggregations: QwakAggregation):
+        self.aggregations += aggregations
+        return self
+
+    def by_windows(self, *time_windows: str):
+        self.windows += Window._from_string(*set(time_windows))
+        return self
+
+    def get_names(self) -> List[str]:
+        return [a.get_name() for a in self.aggregations]
+
+    def get_features_names(self) -> List[str]:
+        if self.windows and self.aggregations:
+            return [
+                f"{name}_{tu.get_feature_suffix()}"
+                for tu in self.windows
+                for name in self.get_names()
+            ]
+        else:
+            raise QwakException(
+                "Feature names are retrieved for streaming aggregation featuresets only"
+            )
+
+    def _get_aggregations_proto(self) -> List[ProtoAggregationField]:
+        """
+        Converts a user-defined streaming aggregation transformation to a list of proto aggregation fields
+        Time windows are applied to all defined aggregations (cartesian product)
+        """
+
+        time_windows: List[ProtoTimeWindow] = [
+            ProtoTimeWindow(length=window.length, time_unit=window.time_unit_proto)
+            for window in self.windows
+        ]
+
+        proto_aggregations: List[ProtoAggregationField] = []
+
+        for aggregation in self.aggregations:
+            aggregation_params = {
+                aggregation._key: aggregation.to_proto(),
+                "time_windows": time_windows,
+            }
+            if aggregation.has_alias():
+                aggregation_params["field_name_prefix"] = aggregation._alias_name
+
+            proto_aggregations.append(ProtoAggregationField(**aggregation_params))
+
+        return proto_aggregations
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoTransformation):
+        function_mapping = {
+            "sql_transformation": SparkSqlTransformation,
+            "koalas_transformation": KoalasTransformation,
+            "pyspark_transformation": PySparkTransformation,
+            "udf_transformation": UdfTransformation,
+            "pandas_on_spark_transformation": PandasOnSparkTransformation,
+        }
+
+        function_type: str = proto.WhichOneof("type")
+
+        if function_type in function_mapping:
+            function_class = function_mapping.get(function_type)
+            return function_class._from_proto(proto)
+
+        raise QwakException(f"Got unsupported transformation type: {function_type}")
+
+    @abstractmethod
+    def _to_proto(self, artifact_path: Optional[str] = None) -> ProtoTransformation:
+        pass
+
+    def get_functions(self) -> Optional[List[Callable]]:
+        return None
+
+    def _validate_udfs(self):
+        udfs: Optional[List[Callable]] = self.get_functions()
+
+        if udfs and len(udfs) > 0:
+            python_major = sys.version_info[0]
+            python_minor = sys.version_info[1]
+            if f"{python_major}.{python_minor}" != "3.8":
+                raise QwakException(
+                    f"Feature store UDFs are only supported for python 3.8, instead got: "
+                    f"{python_major}.{python_minor}"
+                )
+
+
+@dataclass
+class KoalasTransformation(BaseTransformation):
+    """
+    Koalas transformation, providing the user with the ability to define a Koalas based UDF for the transformation
+    of the FeatureSet. This option will be deprecated in future versions.
+    @param function: The Koalas function defined for the transformation
+    @type: Callable
+    @deprecated
+    """
+
+    function: Optional[Callable] = None
+    qwargs: Optional[Dict[str, str]] = None
+
+    def __post_init__(self):
+        warn(
+            "Koalas transformation is about to be deprecated. "
+            "Please use PySparkTransformation, SparkSqlTransformation or UDFTransformation instead",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+        self.qwargs = self.qwargs if self.qwargs else {}
+        self.qwargs = {str(k): str(v) for k, v in self.qwargs.items()}
+        self._artifact_path = None
+
+        if not self.function:
+            raise QwakException(
+                "Please provide a valid function for the koalas transformation"
+            )
+
+    def _to_proto(self, artifact_path: Optional[str] = None) -> ProtoTransformation:
+        self._validate_udfs()
+
+        if artifact_path:
+            self._artifact_path = artifact_path
+
+        return ProtoTransformation(
+            koalas_transformation=ProtoKoalasTransformation(
+                function_name=self.function.__name__,
+                qwargs=ProtoTransformArguments(qwargs=self.qwargs),
+            ),
+            artifact_path=self._artifact_path,
+        )
+
+    @classmethod
+    def _from_proto(cls, proto: "ProtoTransformation"):
+        koalas_transformation = proto.koalas_transformation
+        qwargs = {}
+        if koalas_transformation.WhichOneof("args_option") == "qwargs":
+            qwargs = koalas_transformation.qwargs.qwargs
+
+        def f():
+            print(
+                f"Loading Koalas UDFs is not supported. Can not load {koalas_transformation.function_name}"
+            )
+
+        f.__name__ = koalas_transformation.function_name
+
+        return cls(function=f, qwargs=qwargs)
+
+    def get_functions(self) -> Optional[List[Callable]]:
+        return [self.function]
+
+
+@dataclass
+class PySparkTransformation(BaseTransformation):
+    """
+    PySpark transformation, providing the user with the ability to define a PySpark based UDF for the transformation
+    of the FeatureSet.
+    @param function: The PySpark function defined for the transformation
+    @type: Callable
+    @deprecated
+    """
+
+    function: Callable[
+        [Dict[str, "psql.DataFrame"], Dict[str, Any]],
+        "psql.DataFrame",
+    ]
+    qwargs: Optional[Dict[str, str]] = None
+
+    def __post_init__(self):
+        self.qwargs = self.qwargs if self.qwargs else {}
+        self._artifact_path = None
+
+        if importlib.util.find_spec("pyspark"):
+            import pyspark
+
+            validate_function(self.function, pyspark.sql.DataFrame, self.__class__)
+        else:
+            raise QwakException(
+                missing_pyspark_exception_message(self.__class__.__name__)
+            )
+
+        validate_qwargs(self.qwargs)
+
+    def _to_proto(self, artifact_path: Optional[str] = None) -> ProtoTransformation:
+        self._validate_udfs()
+
+        if artifact_path:
+            self._artifact_path = artifact_path
+
+        return ProtoTransformation(
+            pyspark_transformation=ProtoPySparkTransformation(
+                function_name=self.function.__name__,
+                qwargs=ProtoTransformArguments(qwargs=self.qwargs),
+            ),
+            artifact_path=self._artifact_path,
+        )
+
+    @classmethod
+    def _from_proto(cls, proto: "ProtoTransformation"):
+        try:
+            import pyspark.sql as ps
+        except Exception:
+            raise QwakException(missing_pyspark_exception_message(cls.__name__))
+
+        pyspark_transformation = proto.pyspark_transformation
+        qwargs = {}
+        if pyspark_transformation.WhichOneof("args_option") == "qwargs":
+            qwargs = pyspark_transformation.qwargs.qwargs
+
+        def f(dfs: Dict[str, ps.DataFrame]) -> ps.DataFrame:
+            print(
+                f"Loading Pyspark UDFs is not supported. Can not load {pyspark_transformation.function_name}"
+            )
+
+        f.__name__ = pyspark_transformation.function_name
+
+        return cls(function=f, qwargs=qwargs)
+
+    def get_functions(self) -> Optional[List[Callable]]:
+        return [self.function]
+
+
+@dataclass
+class SparkSqlTransformation(BaseTransformation):
+    """
+    A Spark SQL transformation
+    :param sql: A valid Spark SQL transformation
+    :param functions: PySpark Pandas UDFs
+    Example transformation:
+    ... code-block:: python
+        SparkSqlTransformation("SELECT user_id, age FROM data_source")
+    Example transformation with additional UDFs:
+    ... code-block:: python
+        @qwak_pandas_udf(output_schema=Schema([
+            Column(type=Type.long)]))
+        def plus_one(column_a: pd.Series) -> pd.Series:
+            return column_a + 1
+
+
+        @qwak_pandas_udf(output_schema=Schema([
+            Column(type=Type.long)]))
+        def mul_by_two(column_a: pd.Series) -> pd.Series:
+            return column_a * 2
+
+        SparkSqlTransformation("SELECT user_id, age FROM data_source", functions=[plus_one, mul_by_two])
+    """
+
+    sql: str = str()
+    functions: Optional[List[Callable]] = None
+
+    def __post_init__(self):
+        if not self.sql:
+            raise QwakException(
+                "SQL statement can not be empty when using a SparkSqlTransformation"
+            )
+
+    @classmethod
+    def _from_proto(cls, proto: "ProtoSqlTransformation"):
+        return cls(sql=proto.sql_transformation.sql)
+
+    def _to_proto(self, artifact_path: Optional[str] = None) -> ProtoTransformation:
+        self._validate_udfs()
+
+        return ProtoTransformation(
+            sql_transformation=ProtoSqlTransformation(
+                sql=self.sql,
+                function_names=(
+                    [function.__name__ for function in self.functions]
+                    if self.functions
+                    else None
+                ),
+            ),
+            artifact_path=artifact_path,
+        )
+
+    def get_functions(self) -> Optional[List[Callable]]:
+        return self.functions
+
+
+@dataclass
+class UdfTransformation(BaseTransformation):
+    """
+    A UDF Transformation
+    :param function: A valid user defined pandas function decorated with the @qwak_pandas_udf decorator
+    Example transformation:
+    ... code-block:: python
+
+    @qwak_pandas_udf(output_schema=Schema([Column(type=Type.long)]))
+    def my_pandas_udf() -> pd.Series:
+        data = np.random.randn(3)
+        series = pd.Series(data)
+        return series
+
+    UdfTransformation(function=my_pandas_udf)
+    """
+
+    function: Callable = ()
+
+    def __post_init__(self):
+        if not self.function:
+            raise QwakException("Please provide a function to the UdfTransformation")
+
+    def _to_proto(self, artifact_path: Optional[str] = None) -> ProtoTransformation:
+        self._validate_udfs()
+
+        if artifact_path:
+            self._artifact_path = artifact_path
+
+        return ProtoTransformation(
+            artifact_path=artifact_path,
+            udf_transformation=ProtoUdfTransformation(
+                function_name=self.function.__name__
+            ),
+        )
+
+    @classmethod
+    def _from_proto(cls, proto: "ProtoTransformation"):
+        return cls(
+            function=lambda x: print("Loading UDFs is not yet supported"),
+        )
+
+    def get_functions(self) -> Optional[List[Callable]]:
+        return [self.function]
+
+
+@dataclass
+class PandasOnSparkTransformation(BaseTransformation):
+    """
+    Pandas On Spark transformation.
+    This class provides the user with the ability to define a Pandas On Spark based UDF as the transformation of the Feature Set.
+
+    Attributes
+    ----------
+    function : Callable
+        The Pandas On Spark function defined for the transformation.
+    qwargs : dict, optional
+        Runtime parameters (e.g., qwak_ingestion_start_timestamp, qwak_ingestion_end_timestamp).
+    """
+
+    function: Callable[
+        [Dict[str, "psd.DataFrame"], Dict[str, Any]],
+        "psd.DataFrame",
+    ]
+    qwargs: Optional[Dict[str, str]] = None
+
+    def __post_init__(self):
+        self.qwargs = self.qwargs if self.qwargs else {}
+        self._artifact_path = None
+        if importlib.util.find_spec("pyspark"):
+            import pyspark.pandas as psd
+
+            validate_function(self.function, psd.DataFrame, self.__class__)
+        else:
+            raise QwakException(
+                missing_pyspark_exception_message(self.__class__.__name__)
+            )
+        validate_qwargs(self.qwargs)
+
+    def _to_proto(self, artifact_path: Optional[str] = None) -> ProtoTransformation:
+        self._validate_udfs()
+
+        if artifact_path:
+            self._artifact_path = artifact_path
+
+        return ProtoTransformation(
+            pandas_on_spark_transformation=ProtoPandasOnSparkTransformation(
+                function_name=self.function.__name__,
+                qwargs=ProtoTransformArguments(qwargs=self.qwargs),
+            ),
+            artifact_path=self._artifact_path,
+        )
+
+    @classmethod
+    def _from_proto(cls, proto: "ProtoTransformation"):
+        try:
+            import pyspark.pandas as ps
+        except Exception:
+            raise QwakException(
+                "Missing 'pyspark' dependency required for PandasOnSpark transformation. "
+                f"{UPDATE_QWAK_SDK_WITH_FEATURE_STORE_EXTRA_MSG}"
+            )
+        pandas_on_spark_transformation = proto.pandas_on_spark_transformation
+        qwargs = {}
+        if pandas_on_spark_transformation.WhichOneof("args_option") == "qwargs":
+            qwargs = pandas_on_spark_transformation.qwargs.qwargs
+
+        def f(dfs: Dict[str, ps.DataFrame]) -> ps.DataFrame:
+            print(
+                f"Loading Pandas On Spark UDFs is not supported. Can not load {pandas_on_spark_transformation.function_name}"
+            )
+
+        f.__name__ = pandas_on_spark_transformation.function_name
+
+        return cls(function=f, qwargs=qwargs)
+
+    def get_functions(self) -> Optional[List[Callable]]:
+        return [self.function]