wedata-feature-engineering 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

feature_store/constants/constants.py

@@ -0,0 +1,28 @@
+
+OVERWRITE = "overwrite"
+APPEND = "append"
+PATH = "path"
+TABLE = "table"
+CUSTOM = "custom"
+PREDICTION_COLUMN_NAME = "prediction"
+MODEL_DATA_PATH_ROOT = "feature_store"
+UTF8_BYTES_PER_CHAR = 4
+MAX_PRIMARY_KEY_STRING_LENGTH_CHARS = 100
+MAX_PRIMARY_KEY_STRING_LENGTH_BYTES = (
+    MAX_PRIMARY_KEY_STRING_LENGTH_CHARS * UTF8_BYTES_PER_CHAR
+)
+STREAMING_TRIGGER_CONTINUOUS = "continuous"
+STREAMING_TRIGGER_ONCE = "once"
+STREAMING_TRIGGER_PROCESSING_TIME = "processingTime"
+DEFAULT_WRITE_STREAM_TRIGGER = {STREAMING_TRIGGER_PROCESSING_TIME: "5 seconds"}
+_DEFAULT_PUBLISH_STREAM_TRIGGER = {STREAMING_TRIGGER_PROCESSING_TIME: "5 minutes"}
+
+
+_WARN = "WARN"
+_ERROR = "ERROR"
+_SOURCE_FORMAT_DELTA = "delta"
+
+_NO_RESULT_TYPE_PASSED = "NO_RESULT_TYPE"
+_USE_SPARK_NATIVE_JOIN = "use_spark_native_join"
+_PREBUILT_ENV_URI = "prebuilt_env_uri"
+

feature_store/entities/column_info.py

@@ -0,0 +1,117 @@
+import copy
+from typing import Optional, Union
+
+from feature_store.entities.feature_column_info import FeatureColumnInfo
+from feature_store.entities.feature_spec_constants import SOURCE_DATA_COLUMN_INFO, FEATURE_COLUMN_INFO, \
+    ON_DEMAND_COLUMN_INFO
+from feature_store.entities.on_demand_column_info import OnDemandColumnInfo
+from feature_store.entities.source_data_column_info import SourceDataColumnInfo
+
+
+class ColumnInfo:
+    """
+    ColumnInfo's structure and properties are mapped 1:1 to the ColumnInfo proto message, unless specified otherwise.
+    """
+
+    def __init__(
+        self,
+        info: Union[SourceDataColumnInfo, FeatureColumnInfo, OnDemandColumnInfo],
+        include: bool,
+        data_type: Optional[str] = None,
+        topological_ordering: Optional[int] = None,
+    ):
+        if not isinstance(
+            info, (SourceDataColumnInfo, FeatureColumnInfo, OnDemandColumnInfo)
+        ):
+            raise ValueError(
+                "info must be one of SourceDataColumnInfo, FeatureColumnInfo, OnDemandColumnInfo."
+            )
+        self._info = info
+        self._include = include
+        self._data_type = data_type
+        self._topological_ordering = topological_ordering
+
+    @property
+    def info(
+        self,
+    ) -> Union[SourceDataColumnInfo, FeatureColumnInfo, OnDemandColumnInfo]:
+        return self._info
+
+    @property
+    def include(self) -> bool:
+        return self._include
+
+    @property
+    def data_type(self) -> Optional[str]:
+        """
+        FeatureSpecs before v7 are not required to have data types.
+        """
+        return self._data_type
+
+    @property
+    def topological_ordering(self) -> Optional[int]:
+        """
+        FeatureSpecs before v8 are not required to have topological ordering.
+        """
+        return self._topological_ordering
+
+    @property
+    def output_name(self) -> str:
+        """
+        This field does not exist in the proto, and is provided for convenience.
+        """
+        return self.info.output_name
+
+    def with_topological_ordering(self, ordering: int):
+        new_column_info = copy.copy(self)
+        new_column_info._topological_ordering = ordering
+        return new_column_info
+
+    @classmethod
+    def from_proto(cls, column_info_proto):
+        if column_info_proto.HasField(SOURCE_DATA_COLUMN_INFO):
+            info = SourceDataColumnInfo.from_proto(
+                column_info_proto.source_data_column_info
+            )
+        elif column_info_proto.HasField(FEATURE_COLUMN_INFO):
+            info = FeatureColumnInfo.from_proto(column_info_proto.feature_column_info)
+        elif column_info_proto.HasField(ON_DEMAND_COLUMN_INFO):
+            info = OnDemandColumnInfo.from_proto(
+                column_info_proto.on_demand_column_info
+            )
+        else:
+            raise ValueError("Unsupported info type: " + str(column_info_proto))
+
+        data_type = (
+            column_info_proto.data_type
+            if column_info_proto.HasField("data_type")
+            else None
+        )
+        topological_ordering = (
+            column_info_proto.topological_ordering
+            if column_info_proto.HasField("topological_ordering")
+            else None
+        )
+        return ColumnInfo(
+            info=info,
+            include=column_info_proto.include,
+            data_type=data_type,
+            topological_ordering=topological_ordering,
+        )
+
+    # def to_proto(self):
+    #     column_info = ProtoColumnInfo(
+    #         include=self.include,
+    #         data_type=self.data_type,
+    #         topological_ordering=self.topological_ordering,
+    #     )
+    #     if isinstance(self.info, SourceDataColumnInfo):
+    #         column_info.source_data_column_info.CopyFrom(self.info.to_proto())
+    #     elif isinstance(self.info, FeatureColumnInfo):
+    #         column_info.feature_column_info.CopyFrom(self.info.to_proto())
+    #     elif isinstance(self.info, OnDemandColumnInfo):
+    #         column_info.on_demand_column_info.CopyFrom(self.info.to_proto())
+    #     else:
+    #         raise ValueError("Unsupported info type: " + str(self.info))
+    #
+    #     return column_info

feature_store/entities/data_type.py

@@ -0,0 +1,92 @@
+import json
+import re
+from typing import Any
+
+from pyspark.sql.types import ArrayType, DataType, DecimalType, MapType, StructType
+
+
+
+class DataType(_ProtoEnumEntity):
+    """Online store types."""
+
+    INTEGER = ProtoDataType.Value("INTEGER")
+    FLOAT = ProtoDataType.Value("FLOAT")
+    BOOLEAN = ProtoDataType.Value("BOOLEAN")
+    STRING = ProtoDataType.Value("STRING")
+    DOUBLE = ProtoDataType.Value("DOUBLE")
+    LONG = ProtoDataType.Value("LONG")
+    TIMESTAMP = ProtoDataType.Value("TIMESTAMP")
+    DATE = ProtoDataType.Value("DATE")
+    SHORT = ProtoDataType.Value("SHORT")
+    ARRAY = ProtoDataType.Value("ARRAY")
+    MAP = ProtoDataType.Value("MAP")
+    BINARY = ProtoDataType.Value("BINARY")
+    DECIMAL = ProtoDataType.Value("DECIMAL")
+    STRUCT = ProtoDataType.Value("STRUCT")
+
+    _FIXED_DECIMAL = re.compile("decimal\\(\\s*(\\d+)\\s*,\\s*(\\d+)\\s*\\)")
+
+    @classmethod
+    def _enum_type(cls) -> Any:
+        return ProtoDataType
+
+    @classmethod
+    def from_spark_type(cls, spark_type):
+        return cls.from_string(spark_type.typeName())
+
+    @classmethod
+    def spark_type_to_string(cls, spark_type):
+        return DataType.to_string(DataType.from_spark_type(spark_type))
+
+    @classmethod
+    def top_level_type_supported(cls, spark_type: DataType) -> bool:
+        """
+        Checks whether the provided Spark data type is supported by Feature Store, only considering
+        the top-level type for nested data types.
+
+        Details on nested types:
+          ArrayType: The elementType is not checked. Will return True.
+          MapType: The keyType and valueType are not checked. Will return True.
+          StructType: The struct fieds are not checked. Will return True.
+        """
+        cls.init()
+        return spark_type.typeName().upper() in cls._STRING_TO_ENUM
+
+    @classmethod
+    def to_complex_spark_type(cls, json_value):
+        """
+        Constructs a complex Spark DataType from its compact JSON representation.
+
+        Examples:
+            - Input: '"decimal(1,2)"'
+              Output: DecimalType(1,2)
+            - Input: '{"containsNull":false,"elementType":"integer","type":"array"}'
+              Output: ArrayType(IntegerType,false)
+            - Input: '{"keyType":"integer","type":"map","valueContainsNull":True,"valueType":"integer"}'
+              Output: MapType(IntegerType,IntegerType,true)
+        """
+        if not json_value:
+            raise ValueError("Empty JSON value cannot be converted to Spark DataType")
+
+        json_data = json.loads(json_value)
+        if not isinstance(json_data, dict):
+            # DecimalType does not have fromJson() method
+            if json_value == "decimal":
+                return DecimalType()
+            if cls._FIXED_DECIMAL.match(json_data):
+                m = cls._FIXED_DECIMAL.match(json_data)
+                return DecimalType(int(m.group(1)), int(m.group(2)))
+
+        if json_data["type"].upper() == cls.to_string(cls.ARRAY):
+            return ArrayType.fromJson(json_data)
+
+        if json_data["type"].upper() == cls.to_string(cls.MAP):
+            return MapType.fromJson(json_data)
+
+        if json_data["type"].upper() == cls.to_string(cls.STRUCT):
+            return StructType.fromJson(json_data)
+
+        else:
+            raise ValueError(
+                f"Spark type {json_data['type']} cannot be converted to a complex Spark DataType"
+            )

feature_store/entities/environment_variables.py

@@ -0,0 +1,55 @@
+import os
+
+
+class _EnvironmentVariable:
+    """
+    Represents an environment variable for the feature store client for custom configurations as needed.
+    """
+
+    def __init__(self, name, type_, default):
+        self.name = name
+        self.type = type_
+        self.default = default
+
+    @property
+    def defined(self):
+        return self.name in os.environ
+
+    def get_raw(self):
+        return os.getenv(self.name)
+
+    def set(self, value):
+        os.environ[self.name] = str(value)
+
+    def unset(self):
+        os.environ.pop(self.name, None)
+
+    def get(self):
+        """
+        Reads the value of the environment variable if it exists and converts it to the desired
+        type. Otherwise, returns the default value.
+        """
+        if (val := self.get_raw()) is not None:
+            try:
+                return self.type(val)
+            except Exception as e:
+                raise ValueError(
+                    f"Failed to convert {val!r} to {self.type} for {self.name}: {e}"
+                )
+        return self.default
+
+    def __str__(self):
+        return f"{self.name} (default: {self.default}, type: {self.type.__name__})"
+
+    def __repr__(self):
+        return repr(self.name)
+
+    def __format__(self, format_spec: str) -> str:
+        return self.name.__format__(format_spec)
+
+
+# The threshold (in MB) where a broadcast join will be performed for the asof join for point in time feature join
+# Default is 20MB as benchmarks show diminishing returns with broadcast past this value.The default spark broadcast join threshold is 10MB
+BROADCAST_JOIN_THRESHOLD = _EnvironmentVariable(
+    "BROADCAST_JOIN_THRESHOLD", int, 20 * 1024 * 1024
+)

feature_store/entities/feature.py

@@ -0,0 +1,53 @@
+
+
+class Feature:
+    def __init__(
+        self,
+        feature_table,
+        feature_id,
+        name,
+        data_type,
+        description,
+        data_type_details=None,
+    ):
+        self._feature_table = feature_table
+        self._name = name
+        self._data_type = data_type
+        self._description = description
+        self._data_type_details = data_type_details
+        self._feature_id = feature_id
+
+    @property
+    def feature_table(self):
+        return self._feature_table
+
+    @property
+    def feature_id(self):
+        return self._feature_id
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def data_type(self):
+        return self._data_type
+
+    @property
+    def data_type_details(self):
+        return self._data_type_details
+
+    @property
+    def description(self):
+        return self._description
+
+    @classmethod
+    def from_proto(cls, feature_proto):
+        return cls(
+            feature_table=feature_proto.table,
+            feature_id=feature_proto.id,
+            name=feature_proto.name,
+            data_type=feature_proto.data_type,
+            data_type_details=feature_proto.data_type_details,
+            description=feature_proto.description,
+        )

feature_store/entities/feature_column_info.py

@@ -0,0 +1,64 @@
+from typing import List, Optional
+
+
+
+class FeatureColumnInfo:
+    def __init__(
+        self,
+        table_name: str,
+        feature_name: str,
+        lookup_key: List[str],
+        output_name: str,
+        timestamp_lookup_key: Optional[List[str]] = None,
+    ):
+        if timestamp_lookup_key is None:
+            timestamp_lookup_key = []
+        if not table_name:
+            raise ValueError("table_name must be non-empty.")
+        if not feature_name:
+            raise ValueError("feature_name must be non-empty.")
+        if not isinstance(lookup_key, list):
+            raise ValueError("lookup_key must be a list.")
+        if not lookup_key or "" in lookup_key or None in lookup_key:
+            raise ValueError("lookup_key must be non-empty.")
+        if not output_name:
+            raise ValueError("output_name must be non-empty.")
+        if not isinstance(timestamp_lookup_key, list):
+            raise ValueError("timestamp_lookup_key must be a list.")
+
+        self._table_name = table_name
+        self._feature_name = feature_name
+        self._lookup_key = lookup_key
+        self._output_name = output_name
+        self._timestamp_lookup_key = timestamp_lookup_key
+
+    @property
+    def table_name(self):
+        return self._table_name
+
+    @property
+    def lookup_key(self):
+        return self._lookup_key
+
+    @property
+    def feature_name(self):
+        return self._feature_name
+
+    @property
+    def output_name(self):
+        return self._output_name
+
+    @property
+    def timestamp_lookup_key(self):
+        return self._timestamp_lookup_key
+
+    @classmethod
+    def from_proto(cls, feature_column_info_proto):
+        return cls(
+            table_name=feature_column_info_proto.table_name,
+            feature_name=feature_column_info_proto.feature_name,
+            lookup_key=list(feature_column_info_proto.lookup_key),
+            output_name=feature_column_info_proto.output_name,
+            timestamp_lookup_key=list(feature_column_info_proto.timestamp_lookup_key),
+        )
+

feature_store/entities/feature_function.py

@@ -0,0 +1,55 @@
+from typing import Dict, Optional
+
+
+class FeatureFunction:
+
+    """
+     特征方法类
+
+     特征方法是用户定义的函数，用于将特征表中的特征组合成新特征，特征方法可以是任何用户定义的函数，例如Python UDF。
+
+    特征方法类有以下属性：
+    - udf_name：要调用的Python UDF的名称。
+    - input_bindings：用于将Python UDF的输入映射到训练集中的特征的字典。
+    - output_name：如果提供，则会将此特征重命名为 :meth:`create_training_set() <databricks.feature_engineering.client.FeatureEngineeringClient.create_training_set>` 返回的 :class:`TrainingSet <databricks.ml_features.training_set.TrainingSet>` 中的特征。
+
+    """
+
+    def __init__(
+        self,
+        *,
+        udf_name: str,
+        input_bindings: Optional[Dict[str, str]] = None,
+        output_name: Optional[str] = None,
+    ):
+        """Initialize a FeatureFunction object. See class documentation."""
+        # UC function names are always lowercase.
+        self._udf_name = udf_name.lower()
+        self._input_bindings = input_bindings if input_bindings else {}
+        self._output_name = output_name
+
+    @property
+    def udf_name(self) -> str:
+        """
+        The name of the Python UDF called by this FeatureFunction.
+        """
+        return self._udf_name
+
+    @property
+    def input_bindings(self) -> Dict[str, str]:
+        """
+        The input to use for each argument of the Python UDF.
+
+        For example:
+
+        `{"x": "feature1", "y": "input1"}`
+        """
+        return self._input_bindings
+
+    @property
+    def output_name(self) -> Optional[str]:
+        """
+        The output name to use for the results of this FeatureFunction.
+        If empty, defaults to the fully qualified `udf_name` when evaluated.
+        """
+        return self._output_name

feature_store/entities/feature_lookup.py

@@ -0,0 +1,179 @@
+import copy
+import datetime
+import logging
+from typing import Dict, List, Optional, Union
+
+from feature_store.utils import common_utils
+
+_logger = logging.getLogger(__name__)
+
+
+class FeatureLookup:
+
+    """
+    特征查找类
+
+    特征查找类用于指定特征表中的特征，并将其与训练集中的特征进行关联。
+
+    特征查找类有以下属性：
+
+    - table_name：特征表的名称。
+    - lookup_key：用于在特征表和训练集之间进行联接的键。lookup_key必须是训练集中的列。lookup_key的类型和顺序必须与特征表的主键匹配。
+    - feature_names：要从特征表中查找的特征的名称。如果您的模型需要主键作为特征，则可以将它们声明为独立的FeatureLookups。
+    - rename_outputs：如果提供，则会将特征重命名为 :meth:`create_training_set() <databricks.feature_engineering.client.FeatureEngineeringClient.create_training_set>`返回的 :class:`TrainingSet <databricks.ml_features.training_set.TrainingSet>` 中的特征。
+    - timestamp_lookup_key：用于在特征表和训练集之间进行联接的时间戳键。timestamp_lookup_key必须是训练集中的列。timestamp_lookup_key的类型必须与特征表的时间戳键的类型匹配。
+    - lookback_window: 当对特征表执行时间点查找时使用的回溯窗口，该查找针对传递给 :meth:`create_training_set() <databricks.feature_engineering.client.FeatureEngineeringClient.create_training_set>` 方法的数据帧。特征存储将检索在数据帧的``timestamp_lookup_key``指定时间戳之前且在``lookback_window``时间范围内的最新特征值，如果不存在这样的特征值则返回null。当设置为0时，仅返回特征表中的精确匹配项。
+    - feature_name：特征名称。**已弃用**。使用 `feature_names`。
+    - output_name：如果提供，则会将此特征重命名为 :meth:`create_training_set() <databricks.feature_engineering.client.FeatureEngineeringClient.create_training_set>` 返回的 :class:`TrainingSet <databricks.ml_features.training_set.TrainingSet>` 中的特征。**已弃用**。使用 `rename_outputs`。
+
+    示例：
+
+    from databricks.feature_store import FeatureLookup
+
+    lookup = FeatureLookup(
+        table_name="my_feature_table",
+        lookup_key="my_lookup_key",
+        feature_names=["my_feature_1", "my_feature_2"],
+        rename_outputs={"my_feature_1": "my_feature_1_renamed"},
+        timestamp_lookup_key="my_timestamp_lookup_key",
+        lookback_window=datetime.timedelta(days=1)
+    )
+
+    """
+
+    def __init__(
+        self,
+        table_name: str,
+        lookup_key: Union[str, List[str]],
+        *,
+        feature_names: Union[str, List[str], None] = None,
+        rename_outputs: Optional[Dict[str, str]] = None,
+        timestamp_lookup_key: Optional[str] = None,
+        lookback_window: Optional[datetime.timedelta] = None,
+        **kwargs,
+    ):
+        """Initialize a FeatureLookup object. See class documentation."""
+
+        self._feature_name_deprecated = kwargs.pop("feature_name", None)
+        self._output_name_deprecated = kwargs.pop("output_name", None)
+
+        if kwargs:
+            raise TypeError(
+                f"FeatureLookup got unexpected keyword argument(s): {list(kwargs.keys())}"
+            )
+
+        self._table_name = table_name
+
+        if type(timestamp_lookup_key) is list:
+            if len(timestamp_lookup_key) == 0:
+                timestamp_lookup_key = None
+            elif len(timestamp_lookup_key) == 1:
+                timestamp_lookup_key = timestamp_lookup_key[0]
+            else:
+                raise ValueError(
+                    f"Setting multiple timestamp lookup keys is not supported."
+                )
+
+        if rename_outputs is not None and not isinstance(rename_outputs, dict):
+            raise ValueError(
+                f"Unexpected type for rename_outputs: {type(rename_outputs)}"
+            )
+
+        self._feature_names = common_utils.as_list(feature_names, default=[])
+
+        # Make sure the user didn't accidentally pass in any nested lists/dicts in feature_names
+        for fn in self._feature_names:
+            if not isinstance(fn, str):
+                raise ValueError(
+                    f"Unexpected type for element in feature_names: {type(self._feature_names)}, only strings allowed in list"
+                )
+
+        if lookback_window is not None:
+            if not timestamp_lookup_key:
+                raise ValueError(
+                    f"Unexpected lookback_window value: {lookback_window}, lookback windows can only be applied on time series "
+                    f"feature tables. Use timestamp_lookup_key to perform point-in-time lookups with lookback window."
+                )
+            if not isinstance(
+                lookback_window, datetime.timedelta
+            ) or lookback_window < datetime.timedelta(0):
+                raise ValueError(
+                    f"Unexpected value for lookback_window: {lookback_window}, only non-negative datetime.timedelta allowed."
+                )
+
+        self._lookup_key = copy.copy(lookup_key)
+        self._timestamp_lookup_key = copy.copy(timestamp_lookup_key)
+        self._lookback_window = copy.copy(lookback_window)
+
+        self._rename_outputs = {}
+        if rename_outputs is not None:
+            self._rename_outputs = rename_outputs.copy()
+
+        self._inject_deprecated_feature_name()
+        self._inject_deprecated_output_name()
+
+    @property
+    def table_name(self):
+        """The table name to use in this FeatureLookup."""
+        return self._table_name
+
+    @property
+    def lookup_key(self):
+        """The lookup key(s) to use in this FeatureLookup."""
+        return self._lookup_key
+
+    @property
+    def feature_name(self):
+        """The feature name to use in this FeatureLookup. **Deprecated**. Use `feature_names`."""
+        return self._feature_name_deprecated
+
+    @property
+    def feature_names(self):
+        """The feature names to use in this FeatureLookup."""
+        return self._feature_names
+
+    @property
+    def output_name(self):
+        """The output name to use in this FeatureLookup. **Deprecated**. Use `feature_names`."""
+        if self._output_name_deprecated:
+            return self._output_name_deprecated
+        else:
+            return self._feature_name_deprecated
+
+    @property
+    def timestamp_lookup_key(self):
+        return self._timestamp_lookup_key
+
+    @property
+    def lookback_window(self):
+        """A lookback window applied only for point-in-time lookups."""
+        return self._lookback_window
+
+    def _get_feature_names(self):
+        return self._feature_names
+
+    def _get_output_name(self, feature_name):
+        """Lookup the renamed output, or fallback to the feature name itself if no mapping is present"""
+        return self._rename_outputs.get(feature_name, feature_name)
+
+    def _inject_deprecated_feature_name(self):
+        if self._feature_name_deprecated:
+            if len(self._feature_names) > 0:
+                raise ValueError(
+                    "Use either feature_names or feature_name parameter, but not both."
+                )
+            _logger.warning(
+                f'The feature_name parameter is deprecated. Use "feature_names".'
+            )
+            self._feature_names = [self._feature_name_deprecated]
+
+    def _inject_deprecated_output_name(self):
+        if len(self._feature_names) == 1 and self._output_name_deprecated:
+            if len(self._rename_outputs) > 0:
+                raise ValueError(
+                    "Use either output_name or rename_outputs parameter, but not both."
+                )
+            _logger.warning(
+                f'The output_name parameter is deprecated.  Use "rename_outputs".'
+            )
+            self._rename_outputs[self._feature_names[0]] = self._output_name_deprecated