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

wedata/feature_store/entities/feature_spec_constants.py

@@ -0,0 +1,25 @@
+# Field names from feature_spec.proto.
+SOURCE_DATA_COLUMN_INFO = "source_data_column_info"
+FEATURE_COLUMN_INFO = "feature_column_info"
+ON_DEMAND_COLUMN_INFO = "on_demand_column_info"
+INPUT_COLUMNS = "input_columns"
+NAME = "name"
+OUTPUT_NAME = "output_name"
+INPUT_TABLES = "input_tables"
+TABLE_NAME = "table_name"
+TABLE_ID = "table_id"
+SERIALIZATION_VERSION = "serialization_version"
+INPUT_FUNCTIONS = "input_functions"
+INCLUDE = "include"
+DATA_TYPE = "data_type"
+TOPOLOGICAL_ORDERING = "topological_ordering"
+UDF_NAME = "udf_name"
+INPUT_BINDINGS = "input_bindings"
+PARAMETER = "parameter"
+BOUND_TO = "bound_to"
+
+# FeatureSpec YAML source field and allowed values
+SOURCE = "source"
+TRAINING_DATA = "training_data"
+FEATURE_STORE = "feature_store"
+ON_DEMAND_FEATURE = "on_demand_feature"

wedata/feature_store/entities/feature_table.py

@@ -0,0 +1,111 @@
+from typing import Dict
+
+
+class FeatureTable:
+    """
+    特征表实体类
+    """
+
+    def __init__(
+        self,
+        name,
+        table_id,
+        description,
+        primary_keys,
+        partition_columns,
+        features,
+        creation_timestamp=None,
+        online_stores=None,
+        notebook_producers=None,
+        job_producers=None,
+        table_data_sources=None,
+        path_data_sources=None,
+        custom_data_sources=None,
+        timestamp_keys=None,
+        tags=None,
+    ):
+        """Initialize a FeatureTable object."""
+        """初始化特征表对象
+        
+        :param name: 表名
+        :param table_id: 表ID
+        :param description: 描述
+        :param primary_keys: 主键列表
+        :param partition_columns: 分区列列表
+        :param features: 特征列列表
+        :param creation_timestamp: 创建时间戳(可选)
+        :param online_stores: 在线存储配置(可选)
+        :param notebook_producers: Notebook生产者列表(可选)
+        :param job_producers: 作业生产者列表(可选)
+        :param table_data_sources: 表数据源列表(可选)
+        :param path_data_sources: 路径数据源列表(可选)
+        :param custom_data_sources: 自定义数据源列表(可选)
+        :param timestamp_keys: 时间戳键列表(可选)
+        :param tags: 标签字典(可选)
+        """
+        self.name = name
+        self.table_id = table_id
+        self.description = description
+        self.primary_keys = primary_keys
+        self.partition_columns = partition_columns
+        self.features = features
+        self.creation_timestamp = creation_timestamp
+        self.online_stores = online_stores if online_stores is not None else []
+        self.notebook_producers = (
+            notebook_producers if notebook_producers is not None else []
+        )
+        self.job_producers = job_producers if job_producers is not None else []
+        self.table_data_sources = (
+            table_data_sources if table_data_sources is not None else []
+        )
+        self.path_data_sources = (
+            path_data_sources if path_data_sources is not None else []
+        )
+        self.custom_data_sources = (
+            custom_data_sources if custom_data_sources is not None else []
+        )
+        self.timestamp_keys = timestamp_keys if timestamp_keys is not None else []
+        self._tags = tags
+
+    def __str__(self):
+        """
+        返回特征表实例的字符串表示，包含所有关键属性信息
+
+        返回:
+            格式化的字符串，包含表名、ID、描述、主键、分区列、特征数量、
+            时间戳键、创建时间、数据源数量和标签数量等信息
+        """
+        if self.description and len(self.description) > 50:
+            desc = self.description[:50] + "..."
+        else:
+            desc = self.description
+        return (
+            f"FeatureTable(\n"
+            f"  name='{self.name}',\n"
+            f"  table_id='{self.table_id}',\n"
+            f"  description='{desc}',\n"
+            f"  primary_keys={self.primary_keys},\n"
+            f"  partition_columns={self.partition_columns},\n"
+            f"  features={len(self.features)},\n"
+            f"  timestamp_keys={self.timestamp_keys},\n"
+            f"  creation_timestamp={self.creation_timestamp},\n"
+            f"  data_sources=[table:{len(self.table_data_sources)} "
+            f"path:{len(self.path_data_sources)} custom:{len(self.custom_data_sources)}],\n"
+            f"  tags={len(self.tags) if self._tags else 0}\n"
+            f")"
+        )
+
+    @property
+    def tags(self) -> Dict[str, str]:
+        """
+        Get the tags associated with the feature table.
+
+        :return a Dictionary of all tags associated with the feature table as key/value pairs
+        """
+        if self._tags is None:
+            # If no tags are set, self._tags is expected an empty dictionary.
+            raise ValueError(
+                "Internal error: tags have not been fetched for this FeatureTable instance"
+            )
+        return self._tags
+

wedata/feature_store/entities/feature_table_info.py

@@ -0,0 +1,49 @@
+from typing import Optional
+
+
+from wedata.feature_store.common.protos import feature_store_pb2
+
+class FeatureTableInfo:
+    def __init__(
+        self, table_name: str, table_id: str, lookback_window: Optional[float] = None
+    ):
+        if not table_name:
+            raise ValueError("table_name must be non-empty.")
+        if not table_id:
+            raise ValueError("table_id must be non-empty.")
+        self._table_name = table_name
+        self._table_id = table_id
+        self._lookback_window = lookback_window
+
+    @property
+    def table_name(self):
+        return self._table_name
+
+    @property
+    def table_id(self):
+        return self._table_id
+
+    @property
+    def lookback_window(self):
+        return self._lookback_window
+
+    @classmethod
+    def from_proto(cls, feature_table_info_proto):
+        lookback_window = feature_table_info_proto.lookback_window or None
+        # lookback_window = (
+        #     feature_table_info_proto.lookback_window
+        #     if feature_table_info_proto.lookback_window != 0
+        #     else None
+        # )
+        return cls(
+            table_name=feature_table_info_proto.table_name,
+            table_id=feature_table_info_proto.table_id,
+            lookback_window=lookback_window,
+        )
+
+    def to_proto(self):
+        return feature_store_pb2.FeatureTableInfo(
+            table_name=self.table_name,
+            table_id=self.table_id,
+            lookback_window=self.lookback_window,
+        )

wedata/feature_store/entities/function_info.py

@@ -0,0 +1,90 @@
+
+from typing import List, Optional
+
+from wedata.feature_store.common.protos import feature_store_pb2
+
+class FunctionParameterInfo():
+    def __init__(self, name: str, type_text: str):
+        self._name = name
+        self._type_text = type_text
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def type_text(self) -> str:
+        return self._type_text
+
+    @classmethod
+    def from_dict(cls, function_parameter_info_json):
+        return FunctionParameterInfo(
+            function_parameter_info_json["name"],
+            function_parameter_info_json["type_text"],
+        )
+
+
+class FunctionInfo():
+    """
+    Helper entity class that exposes properties in GetFunction's response JSON as attributes.
+    https://docs.databricks.com/api-explorer/workspace/functions/get
+
+    Note: empty fields (e.g. when 0 input parameters) are not included in the response JSON.
+    """
+
+    # Python UDFs have external_language = "Python"
+    PYTHON = "Python"
+
+    def __init__(
+        self,
+        full_name: str,
+        input_params: List[FunctionParameterInfo],
+        routine_definition: Optional[str],
+        external_language: Optional[str],
+    ):
+        self._full_name = full_name
+        self._input_params = input_params
+        self._routine_definition = routine_definition
+        self._external_language = external_language
+
+
+    @property
+    def full_name(self) -> str:
+        return self._full_name
+
+    @property
+    def input_params(self) -> List[FunctionParameterInfo]:
+        return self._input_params
+
+    @property
+    def routine_definition(self) -> Optional[str]:
+        return self._routine_definition
+
+    @property
+    def external_language(self) -> Optional[str]:
+        """
+        Field is None if language is SQL (not an external language).
+        """
+        return self._external_language
+
+    @classmethod
+    def from_dict(cls, function_info_json):
+        input_params = function_info_json.get("input_params", {}).get("parameters", [])
+        return FunctionInfo(
+            full_name=function_info_json["full_name"],
+            input_params=[FunctionParameterInfo.from_dict(p) for p in input_params],
+            routine_definition=function_info_json.get("routine_definition", None),
+            external_language=function_info_json.get("external_language", None),
+        )
+
+    @classmethod
+    def from_proto(cls, function_info_proto):
+        return cls(full_name=function_info_proto.udf_name)
+
+    def to_proto(self):
+        return feature_store_pb2.FunctionInfo(
+            full_name=self.full_name,
+            input_params=self.input_params,
+            routine_definition=self.routine_definition,
+            external_language=self.external_language
+        )

wedata/feature_store/entities/on_demand_column_info.py

@@ -0,0 +1,57 @@
+from typing import Dict
+
+from wedata.feature_store.common.protos import feature_store_pb2
+
+class OnDemandColumnInfo:
+    def __init__(
+        self,
+        udf_name: str,
+        input_bindings: Dict[str, str],
+        output_name: str,
+    ):
+        if not udf_name:
+            raise ValueError("udf_name must be non-empty.")
+        if not output_name:
+            raise ValueError("output_name must be non-empty.")
+
+        self._udf_name = udf_name
+        self._input_bindings = input_bindings
+        self._output_name = output_name
+
+    @property
+    def udf_name(self) -> str:
+        return self._udf_name
+
+    @property
+    def input_bindings(self) -> Dict[str, str]:
+        """
+        input_bindings is serialized as the InputBindings proto message.
+        """
+        return self._input_bindings
+
+    @property
+    def output_name(self) -> str:
+        return self._output_name
+
+    @classmethod
+    def from_proto(cls, on_demand_column_info_proto):
+        input_bindings_dict = {
+            input_binding.parameter: input_binding.bound_to
+            for input_binding in on_demand_column_info_proto.input_bindings
+        }
+        return OnDemandColumnInfo(
+            udf_name=on_demand_column_info_proto.udf_name,
+            input_bindings=input_bindings_dict,
+            output_name=on_demand_column_info_proto.output_name,
+        )
+
+    def to_proto(self):
+        input_bindings_list = [
+            feature_store_pb2.InputBinding(parameter=k, bound_to=v)
+            for k, v in self.input_bindings.items()
+        ]
+        return feature_store_pb2.OnDemandColumnInfo(
+            udf_name=self.udf_name,
+            input_bindings=input_bindings_list,
+            output_name=self.output_name,
+        )

wedata/feature_store/entities/source_data_column_info.py

@@ -0,0 +1,24 @@
+from wedata.feature_store.common.protos import feature_store_pb2
+class SourceDataColumnInfo:
+    def __init__(self, name: str):
+        if not name:
+            raise ValueError("name must be non-empty.")
+        self._name = name
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def output_name(self) -> str:
+        """
+        This field does not exist in the proto, and is provided for convenience.
+        """
+        return self._name
+
+    @classmethod
+    def from_proto(cls, source_data_column_info_proto):
+        return cls(name=source_data_column_info_proto.name)
+
+    def to_proto(self):
+        return feature_store_pb2.SourceDataColumnInfo(name=self._name)

wedata/feature_store/entities/training_set.py

@@ -0,0 +1,135 @@
+from typing import Dict, List, Optional
+
+from pyspark.sql import DataFrame
+
+from wedata.feature_store.entities.feature_table import FeatureTable
+from wedata.feature_store.entities.function_info import FunctionInfo
+from wedata.feature_store.utils.feature_lookup_utils import (
+    join_feature_data_if_not_overridden,
+)
+
+from wedata.feature_store.entities.feature_spec import FeatureSpec
+from wedata.feature_store.utils.feature_spec_utils import (
+    COLUMN_INFO_TYPE_FEATURE,
+    COLUMN_INFO_TYPE_ON_DEMAND,
+    COLUMN_INFO_TYPE_SOURCE,
+    get_feature_execution_groups,
+)
+from wedata.feature_store.utils.on_demand_utils import apply_functions_if_not_overridden
+
+
+class TrainingSet:
+    """
+    .. note::
+
+       Aliases: `!databricks.feature_engineering.training_set.TrainingSet`, `!databricks.feature_store.training_set.TrainingSet`
+
+    Class that defines :obj:`TrainingSet` objects.
+
+    .. note::
+
+       The :class:`TrainingSet` constructor should not be called directly. Instead,
+       call :meth:`create_training_set() <databricks.feature_engineering.client.FeatureEngineeringClient.create_training_set>`.
+    """
+
+    def __init__(
+        self,
+        feature_spec: FeatureSpec,
+        df: DataFrame,
+        labels: List[str],
+        feature_table_metadata_map: Dict[str, FeatureTable],
+        feature_table_data_map: Dict[str, DataFrame],
+        uc_function_infos: Dict[str, FunctionInfo],
+        use_spark_native_join: Optional[bool] = False,
+    ):
+        """Initialize a :obj:`TrainingSet` object."""
+        assert isinstance(
+            labels, list
+        ), f"Expected type `list` for argument `labels`. Got '{labels}' with type '{type(labels)}'."
+
+        self._feature_spec = feature_spec
+        self._df = df
+        self._labels = labels
+        self._feature_table_metadata_map = feature_table_metadata_map
+        self._feature_table_data_map = feature_table_data_map
+        self._uc_function_infos = uc_function_infos
+        self._use_spark_native_join = use_spark_native_join
+        # Perform basic validations and resolve FeatureSpec and label column data types.
+        self._validate_and_inject_dtypes()
+        self._label_data_types = {
+            name: data_type for name, data_type in df.dtypes if name in labels
+        }
+
+    @property
+    def feature_spec(self) -> FeatureSpec:
+        """Define a feature spec."""
+        return self._feature_spec
+
+    def _augment_df(self) -> DataFrame:
+        """
+        Internal helper to augment DataFrame with feature lookups and on-demand features specified in the FeatureSpec.
+        Does not drop excluded columns, and does not overwrite columns that already exist.
+        Return column order is df.columns + feature lookups + on-demand features.
+        """
+        execution_groups = get_feature_execution_groups(
+            self.feature_spec, self._df.columns
+        )
+
+        result_df = self._df
+        # Iterate over all levels and type of DAG nodes in FeatureSpec and execute them.
+        for execution_group in execution_groups:
+            if execution_group.type == COLUMN_INFO_TYPE_SOURCE:
+                continue
+            if execution_group.type == COLUMN_INFO_TYPE_FEATURE:
+                # Apply FeatureLookups
+                result_df = join_feature_data_if_not_overridden(
+                    feature_spec=self.feature_spec,
+                    df=result_df,
+                    features_to_join=execution_group.features,
+                    feature_table_metadata_map=self._feature_table_metadata_map,
+                    feature_table_data_map=self._feature_table_data_map,
+                    use_spark_native_join=self._use_spark_native_join,
+                )
+            elif execution_group.type == COLUMN_INFO_TYPE_ON_DEMAND:
+                # Apply all on-demand UDFs
+                result_df = apply_functions_if_not_overridden(
+                    df=result_df,
+                    functions_to_apply=execution_group.features,
+                    uc_function_infos=self._uc_function_infos,
+                )
+            else:
+                # This should never be reached.
+                raise Exception("Unknown feature execution type:", execution_group.type)
+        return result_df
+
+    def _validate_and_inject_dtypes(self):
+        """
+        Performs validations through _augment_df (e.g. Delta table exists, Delta and feature table dtypes match),
+        then inject the result DataFrame dtypes into the FeatureSpec.
+        """
+        augmented_df = self._augment_df()
+        augmented_df_dtypes = {column: dtype for column, dtype in augmented_df.dtypes}
+
+        # Inject the result DataFrame column types into the respective ColumnInfo
+        for ci in self.feature_spec.column_infos:
+            ci._data_type = augmented_df_dtypes[ci.output_name]
+
+    def load_df(self) -> DataFrame:
+        """
+        Load a :class:`DataFrame <pyspark.sql.DataFrame>`.
+
+        Return a :class:`DataFrame <pyspark.sql.DataFrame>` for training.
+
+        The returned :class:`DataFrame <pyspark.sql.DataFrame>` has columns specified
+        in the ``feature_spec`` and ``labels`` parameters provided
+        in :meth:`create_training_set() <databricks.feature_engineering.client.FeatureEngineeringClient.create_training_set>`.
+
+        :return:
+           A :class:`DataFrame <pyspark.sql.DataFrame>` for training
+        """
+        augmented_df = self._augment_df()
+        # Return only included columns in order defined by FeatureSpec + labels
+        included_columns = [
+            ci.output_name for ci in self.feature_spec.column_infos if ci.include
+        ] + self._labels
+        return augmented_df.select(included_columns)