tencent-wedata-feature-engineering-dev 0.1.0__py3-none-any.whl

wedata/feature_store/common/protos/feature_store_pb2.py

@@ -0,0 +1,49 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: feature_store.proto
+"""Generated protocol buffer code."""
+from google.protobuf.internal import builder as _builder
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf import reflection as _reflection
+from google.protobuf import message as _message
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x13\x66\x65\x61ture_store.proto\x12\x1b\x66\x65\x61ture_store.common.protos\"$\n\x14SourceDataColumnInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\"\x84\x01\n\x11\x46\x65\x61tureColumnInfo\x12\x12\n\ntable_name\x18\x01 \x01(\t\x12\x14\n\x0c\x66\x65\x61ture_name\x18\x02 \x01(\t\x12\x12\n\nlookup_key\x18\x03 \x03(\t\x12\x13\n\x0boutput_name\x18\x04 \x01(\t\x12\x1c\n\x14timestamp_lookup_key\x18\x05 \x03(\t\"~\n\x12OnDemandColumnInfo\x12\x10\n\x08udf_name\x18\x01 \x01(\t\x12\x41\n\x0einput_bindings\x18\x02 \x03(\x0b\x32).feature_store.common.protos.InputBinding\x12\x13\n\x0boutput_name\x18\x03 \x01(\t\"3\n\x0cInputBinding\x12\x11\n\tparameter\x18\x01 \x01(\t\x12\x10\n\x08\x62ound_to\x18\x02 \x01(\t\"\xfc\x02\n\nColumnInfo\x12T\n\x17source_data_column_info\x18\x01 \x01(\x0b\x32\x31.feature_store.common.protos.SourceDataColumnInfoH\x00\x12M\n\x13\x66\x65\x61ture_column_info\x18\x02 \x01(\x0b\x32..feature_store.common.protos.FeatureColumnInfoH\x00\x12P\n\x15on_demand_column_info\x18\x03 \x01(\x0b\x32/.feature_store.common.protos.OnDemandColumnInfoH\x00\x12\x14\n\x07include\x18\x04 \x01(\x08H\x01\x88\x01\x01\x12\x11\n\tdata_type\x18\x05 \x01(\t\x12!\n\x14topological_ordering\x18\x06 \x01(\x05H\x02\x88\x01\x01\x42\x06\n\x04infoB\n\n\x08_includeB\x17\n\x15_topological_ordering\"Q\n\x10\x46\x65\x61tureTableInfo\x12\x12\n\ntable_name\x18\x01 \x01(\t\x12\x10\n\x08table_id\x18\x02 \x01(\t\x12\x17\n\x0flookback_window\x18\x03 \x01(\x01\"3\n\x15\x46unctionParameterInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\"\xa2\x01\n\x0c\x46unctionInfo\x12\x11\n\tfull_name\x18\x01 \x01(\t\x12H\n\x0cinput_params\x18\x02 \x03(\x0b\x32\x32.feature_store.common.protos.FunctionParameterInfo\x12\x1a\n\x12routine_definition\x18\x03 \x01(\t\x12\x19\n\x11\x65xternal_language\x18\x04 \x01(\t\"\xb1\x02\n\x0b\x46\x65\x61tureSpec\x12>\n\rinput_columns\x18\x01 \x03(\x0b\x32\'.feature_store.common.protos.ColumnInfo\x12\x43\n\x0cinput_tables\x18\x02 \x03(\x0b\x32-.feature_store.common.protos.FeatureTableInfo\x12\x42\n\x0finput_functions\x18\x03 \x03(\x0b\x32).feature_store.common.protos.FunctionInfo\x12\x14\n\x0cworkspace_id\x18\x04 \x01(\x03\x12$\n\x1c\x66\x65\x61ture_store_client_version\x18\x05 \x01(\t\x12\x1d\n\x15serialization_version\x18\x06 \x01(\x05\x62\x06proto3')
+
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'feature_store_pb2', globals())
+_FEATURESPEC = DESCRIPTOR.message_types_by_name['FeatureSpec']
+FeatureSpec = _reflection.GeneratedProtocolMessageType('FeatureSpec', (_message.Message,), {
+  'DESCRIPTOR' : _FEATURESPEC,
+  '__module__' : 'feature_spec_pb2'
+  # @@protoc_insertion_point(class_scope:featurestorecommon.FeatureSpec)
+})
+if _descriptor._USE_C_DESCRIPTORS == False:
+
+  DESCRIPTOR._options = None
+  _SOURCEDATACOLUMNINFO._serialized_start=52
+  _SOURCEDATACOLUMNINFO._serialized_end=88
+  _FEATURECOLUMNINFO._serialized_start=91
+  _FEATURECOLUMNINFO._serialized_end=223
+  _ONDEMANDCOLUMNINFO._serialized_start=225
+  _ONDEMANDCOLUMNINFO._serialized_end=351
+  _INPUTBINDING._serialized_start=353
+  _INPUTBINDING._serialized_end=404
+  _COLUMNINFO._serialized_start=407
+  _COLUMNINFO._serialized_end=787
+  _FEATURETABLEINFO._serialized_start=789
+  _FEATURETABLEINFO._serialized_end=870
+  _FUNCTIONPARAMETERINFO._serialized_start=872
+  _FUNCTIONPARAMETERINFO._serialized_end=923
+  _FUNCTIONINFO._serialized_start=926
+  _FUNCTIONINFO._serialized_end=1088
+  _FEATURESPEC._serialized_start=1091
+  _FEATURESPEC._serialized_end=1396
+# @@protoc_insertion_point(module_scope)

wedata/feature_store/common/store_config/redis.py

@@ -0,0 +1,48 @@
+# -*- coding: utf-8 -*-
+
+__doc__ = """
+Feature Redis存储配置
+"""
+
+
+class RedisStoreConfig:
+    def __init__(self, host='localhost', port=6379, db=0, password=None, instance_id=None):
+        self._host = host
+        self._port = port
+        self._db = db
+        self._password = password
+        self._instance_id = instance_id
+
+    @property
+    def host(self):
+        return self._host
+
+    @property
+    def port(self):
+        return self._port
+
+    @property
+    def db(self):
+        return self._db
+
+    @property
+    def password(self):
+        return self._password
+
+    @property
+    def instance_id(self):
+        return self._instance_id
+
+    @property
+    def connection_string(self):
+        if self.password:
+            connection_string = f"{self.host}:{self.port},db={self.db},password={self._password}"
+        else:
+            connection_string = f"{self.host}:{self.port},db={self.db}"
+        return connection_string
+
+    def __repr__(self):
+        return f"RedisStoreConfig(host={self.host}, port={self.port}, db={self.db}, instance_id={self.instance_id})"
+
+    def __str__(self):
+        return self.__repr__()

wedata/feature_store/constants/constants.py

@@ -0,0 +1,59 @@
+
+OVERWRITE = "overwrite"
+APPEND = "append"
+PATH = "path"
+TABLE = "table"
+CUSTOM = "custom"
+PREDICTION_COLUMN_NAME = "prediction"
+MODEL_DATA_PATH_ROOT = "feature_store"
+RAW_MODEL_FOLDER = "raw_model"
+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"}
+FEATURE_STORE_CLIENT = "FeatureStoreClient"
+
+
+_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"
+
+# MLflow模型相关常量(原mlflow_model_constants.py)
+# Module name of the original mlflow_model
+MLFLOW_MODEL_NAME = "wedata.feature_store.mlflow_model"
+
+# FeatureStoreClient.log_model将记录包含'raw_model'文件夹的模型
+# 该文件夹存储原始模型的MLmodel文件，用于推理
+RAW_MODEL_FOLDER = "raw_model"
+
+# ML模型文件名常量
+ML_MODEL = "MLmodel"
+
+# 特征查找客户端的PyPI包名
+FEATURE_LOOKUP_CLIENT_PIP_PACKAGE = "tencent-wedata-feature-engineering"
+
+# 特征查找版本号
+FEATURE_LOOKUP_CLIENT_MAJOR_VERSION = "0.1.0"
+
+# 特征存储内部数据目录
+FEATURE_STORE_INTERNAL_DATA_DIR = "_wedata_internal/"
+WEDATA_DEFAULT_FEATURE_STORE_DATABASE = "WEDATA_DEFAULT_FEATURE_STORE_DATABASE"
+
+# 特征表属性
+FEATURE_TABLE_KEY = "wedata.feature_table"
+FEATURE_TABLE_VALUE = "true"
+
+FEATURE_TABLE_PROJECT = "wedata.feature_project_id"
+FEATURE_TABLE_TIMESTAMP = "timestampKeys"
+FEATURE_TABLE_BACKUP_PRIMARY_KEY = "primaryKeys"    # 备用标识，主键
+FEATURE_DLC_TABLE_PRIMARY_KEY = "dlc.ao.data.govern.sorted.keys"

wedata/feature_store/constants/engine_types.py

@@ -0,0 +1,34 @@
+from enum import Enum
+import os
+
+
+class EngineTypes(Enum):
+    HIVE_ENGINE = "hive"
+    ICEBERG_ENGINE = "iceberg"
+
+    @classmethod
+    def get_engine(cls, engine_name: str) -> 'EngineTypes':
+        try:
+            return cls(engine_name.lower())
+        except ValueError:
+            raise ValueError(f"Invalid engine type: {engine_name}. Supported engine types: {list(cls)}")
+
+
+class CalculateEngineTypes(Enum):
+    DLC = "dlc"
+    EMR = "emr"
+
+    @classmethod
+    def get_calculate_engine(cls, engine_name: str) -> 'CalculateEngineTypes':
+        try:
+            return cls(engine_name.lower())
+        except ValueError:
+            raise ValueError(f"Invalid engine type: {engine_name}. Supported engine types: {list(cls)}")
+
+
+def judge_engine_type() -> 'CalculateEngineTypes':
+    if os.environ.get("DLC_REGION", ""):
+        return CalculateEngineTypes.DLC
+    else:
+        return CalculateEngineTypes.EMR
+

wedata/feature_store/entities/column_info.py

@@ -0,0 +1,138 @@
+import copy
+from typing import Optional, Union
+
+from wedata.feature_store.entities.feature_column_info import FeatureColumnInfo
+from wedata.feature_store.entities.feature_spec_constants import SOURCE_DATA_COLUMN_INFO, FEATURE_COLUMN_INFO, \
+    ON_DEMAND_COLUMN_INFO
+from wedata.feature_store.entities.on_demand_column_info import OnDemandColumnInfo
+from wedata.feature_store.entities.source_data_column_info import SourceDataColumnInfo
+
+from wedata.feature_store.common.protos import feature_store_pb2
+
+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 or None
+        # data_type = (
+        #     column_info_proto.data_type
+        #     if column_info_proto.HasField("data_type")
+        #     else None
+        # )
+        topological_ordering = column_info_proto.topological_ordering or 0
+        # 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
+    def to_proto(self):
+        column_info = feature_store_pb2.ColumnInfo(
+            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

wedata/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
+)

wedata/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,
+        )

wedata/feature_store/entities/feature_column_info.py

@@ -0,0 +1,72 @@
+from typing import List, Optional
+
+from wedata.feature_store.common.protos import feature_store_pb2
+
+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),
+        )
+
+    def to_proto(self):
+        return feature_store_pb2.FeatureColumnInfo(
+            table_name=self.table_name,
+            feature_name=self.feature_name,
+            lookup_key=self.lookup_key,
+            output_name=self.output_name,
+            timestamp_lookup_key=self.timestamp_lookup_key,
+        )

wedata/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