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

feature_store/entities/function_info.py

@@ -0,0 +1,184 @@
+from collections import defaultdict
+from typing import List, Optional
+
+from pyspark.sql import Column, DataFrame
+from pyspark.sql.functions import isnull, when
+from pyspark.sql.types import StringType, StructField, StructType
+
+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),
+        )
+
+
+class InformationSchemaSparkClient:
+    """
+    Internal client to retrieve Unity Catalog metadata from system.information_schema.
+    https://docs.databricks.com/sql/language-manual/sql-ref-information-schema.html
+    """
+
+    def _get_routines_with_parameters(self, full_routine_names: List[str]) -> DataFrame:
+        """
+        Retrieve the routines with their parameters from information_schema.routines, information_schema.parameters.
+        Return DataFrame only contains routines that 1. exist and 2. the caller has GetFunction permission on.
+
+        Note: The returned DataFrame contains the cartesian product of routines and parameters.
+        For efficiency, routines table columns are only present in the first row for each routine.
+        """
+        routine_name_schema = StructType(
+            [
+                StructField("specific_catalog", StringType(), False),
+                StructField("specific_schema", StringType(), False),
+                StructField("specific_name", StringType(), False),
+            ]
+        )
+        routine_names_df = self._spark_client.createDataFrame(
+            [full_routine_name.split(".") for full_routine_name in full_routine_names],
+            routine_name_schema,
+        )
+        routines_table = self._spark_client.read_table(
+            "system.information_schema.routines"
+        )
+        parameters_table = self._spark_client.read_table(
+            "system.information_schema.parameters"
+        )
+
+        # Inner join routines table to filter out non-existent routines.
+        # Left join parameters as routines may have no parameters.
+        full_routines_with_parameters_df = routine_names_df.join(
+            routines_table, on=routine_names_df.columns, how="inner"
+        ).join(parameters_table, on=routine_names_df.columns, how="left")
+
+        # Return only relevant metadata from information_schema, sorted by routine name + parameter order.
+        # For efficiency, only preserve routine column values in the first of each routine's result rows.
+        # The first row will have parameter.ordinal_value is None (no parameters) or equals 0 (first parameter).
+        def select_if_first_row(col: Column) -> Column:
+            return when(
+                isnull(parameters_table.ordinal_position)
+                | (parameters_table.ordinal_position == 0),
+                col,
+            ).otherwise(None)
+
+        return full_routines_with_parameters_df.select(
+            routine_names_df.columns
+            + [
+                select_if_first_row(routines_table.routine_definition).alias(
+                    "routine_definition"
+                ),
+                select_if_first_row(routines_table.external_language).alias(
+                    "external_language"
+                ),
+                parameters_table.ordinal_position,
+                parameters_table.parameter_name,
+                parameters_table.full_data_type,
+            ]
+        ).sort(routine_names_df.columns + [parameters_table.ordinal_position])
+
+    def get_functions(self, full_function_names: List[str]) -> List[FunctionInfo]:
+        """
+        Retrieves and maps Unity Catalog functions' metadata as FunctionInfos.
+        """
+        # Avoid unnecessary Spark calls and return if empty.
+        if not full_function_names:
+            return []
+
+        # Collect dict of routine name -> DataFrame rows describing the routine.
+        routines_with_parameters_df = self._get_routines_with_parameters(
+            full_routine_names=full_function_names
+        )
+        routine_infos = defaultdict(list)
+        for r in routines_with_parameters_df.collect():
+            routine_name = f"{r.specific_catalog}.{r.specific_schema}.{r.specific_name}"
+            routine_infos[routine_name].append(r)
+
+        # Mock GetFunction DNE error, since information_schema does not throw.
+        for function_name in full_function_names:
+            if not function_name in routine_infos:
+                raise ValueError(f"Function '{function_name}' does not exist.")
+
+        # Map routine_infos into FunctionInfos.
+        function_infos = []
+        for function_name in full_function_names:
+            routine_info = routine_infos[function_name][0]
+            input_params = [
+                FunctionParameterInfo(name=p.parameter_name, type_text=p.full_data_type)
+                for p in routine_infos[function_name]
+                if p.ordinal_position is not None
+            ]
+            function_infos.append(
+                FunctionInfo(
+                    full_name=function_name,
+                    input_params=input_params,
+                    routine_definition=routine_info.routine_definition,
+                    external_language=routine_info.external_language,
+                )
+            )
+        return function_infos

feature_store/entities/on_demand_column_info.py

@@ -0,0 +1,44 @@
+from typing import Dict
+
+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,
+        )

feature_store/entities/source_data_column_info.py

@@ -0,0 +1,21 @@
+
+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)

feature_store/entities/training_set.py

@@ -0,0 +1,134 @@
+from typing import Dict, List, Optional
+
+from pyspark.sql import DataFrame
+
+from feature_store.entities.feature_table import FeatureTable
+from feature_store.entities.function_info import FunctionInfo
+from feature_store.utils.feature_lookup_utils import (
+    join_feature_data_if_not_overridden,
+)
+
+from feature_store.entities.feature_spec import FeatureSpec
+from feature_store.utils.feature_spec_utils import (
+    COLUMN_INFO_TYPE_FEATURE,
+    COLUMN_INFO_TYPE_ON_DEMAND,
+    COLUMN_INFO_TYPE_SOURCE,
+    get_feature_execution_groups,
+)
+
+
+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)

feature_store/feature_table_client/feature_table_client.py

@@ -0,0 +1,313 @@
+"""
+特征表操作相关工具方法
+"""
+
+from typing import Union, List, Dict, Optional, Sequence, Any
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql.streaming import StreamingQuery
+from pyspark.sql.types import StructType
+import os
+
+from feature_store.constants.constants import APPEND, DEFAULT_WRITE_STREAM_TRIGGER
+
+
+class FeatureTableClient:
+    """特征表操作类"""
+
+    def __init__(
+            self,
+            spark: SparkSession
+    ):
+        self._spark = spark
+
+    @staticmethod
+    def _normalize_params(
+            param: Optional[Union[str, Sequence[str]]],
+            default_type: type = list
+    ) -> list:
+        """统一处理参数标准化"""
+        if param is None:
+            return default_type()
+        return list(param) if isinstance(param, Sequence) else [param]
+
+    @staticmethod
+    def _validate_schema(df: DataFrame, schema: StructType):
+        """校验DataFrame和schema的有效性和一致性"""
+        # 检查是否同时为空
+        if df is None and schema is None:
+            raise ValueError("必须提供DataFrame或schema其中之一")
+
+        # 检查schema匹配
+        if df is not None and schema is not None:
+            df_schema = df.schema
+            if df_schema != schema:
+                diff_fields = set(df_schema.fieldNames()).symmetric_difference(set(schema.fieldNames()))
+                raise ValueError(
+                    f"DataFrame与schema不匹配。差异字段: {diff_fields if diff_fields else '字段类型不一致'}"
+                )
+
+    @staticmethod
+    def _validate_table_name(name: str):
+        """验证特征表命名规范"""
+        if name.count('.') < 2:
+            raise ValueError("特征表名称需符合<catalog>.<schema>.<table>格式")
+
+    @staticmethod
+    def _validate_key_conflicts(primary_keys: List[str], timestamp_keys: List[str]):
+        """校验主键与时间戳键是否冲突"""
+        conflict_keys = set(timestamp_keys) & set(primary_keys)
+        if conflict_keys:
+            raise ValueError(f"时间戳键与主键冲突: {conflict_keys}")
+
+    @staticmethod
+    def _escape_sql_value(value: str) -> str:
+        """转义SQL值中的特殊字符"""
+        return value.replace("'", "''")
+
+    def create_table(
+            self,
+            name: str,
+            primary_keys: Union[str, List[str]],
+            df: Optional[DataFrame] = None,
+            *,
+            timestamp_keys: Union[str, List[str], None] = None,
+            partition_columns: Union[str, List[str], None] = None,
+            schema: Optional[StructType] = None,
+            description: Optional[str] = None,
+            tags: Optional[Dict[str, str]] = None
+    ):
+        """
+        创建特征表（支持批流数据写入）
+
+        Args:
+            name: 特征表全称（格式：<table>）
+            primary_keys: 主键列名（支持复合主键）
+            df: 初始数据（可选，用于推断schema）
+            timestamp_keys: 时间戳键（用于时态特征）
+            partition_columns: 分区列（优化存储查询）
+            description: 业务描述
+            tags: 业务标签
+
+        Returns:
+            FeatureTable实例
+
+        Raises:
+            ValueError: 当schema与数据不匹配时
+        """
+        # 参数标准化
+        primary_keys = self._normalize_params(primary_keys)
+        timestamp_keys = self._normalize_params(timestamp_keys)
+        partition_columns = self._normalize_params(partition_columns)
+
+        # 元数据校验
+        self._validate_schema(df, schema)
+        #self._validate_table_name(name)
+        self._validate_key_conflicts(primary_keys, timestamp_keys)
+
+        # 表名 格式：<catalog>.<schema>.<table>  catalog默认值：DataLakeCatalog，schema默认值：feature_store
+        table_name = f'DataLakeCatalog.feature_store.{name}'
+
+        # 检查表是否存在
+        try:
+            if self._spark.catalog.tableExists(table_name):
+                raise ValueError(
+                    f"表 '{table_name}' 已存在\n"
+                    "解决方案：\n"
+                    "1. 使用不同的表名\n"
+                    "2. 删除现有表: spark.sql(f'DROP TABLE {name}')\n"
+                )
+        except Exception as e:
+            raise ValueError(f"检查表存在性时出错: {str(e)}") from e
+
+        # 推断表schema
+        table_schema = schema or df.schema
+
+        # 构建时间戳键属性
+        timestamp_keys_ddl = []
+        for timestamp_key in timestamp_keys:
+            if timestamp_key not in primary_keys:
+                raise ValueError(f"时间戳键 '{timestamp_key}' 必须是主键")
+            timestamp_keys_ddl.append(f"`{timestamp_key}` TIMESTAMP")
+
+        #从环境变量获取额外标签
+        env_tags = {
+            "project_id": os.getenv("WEDATA_PROJECT_ID", ""),  # wedata项目ID
+            "engine_name": os.getenv("WEDATA_NOTEBOOK_ENGINE", ""),  # wedata引擎名称
+            "user_uin": os.getenv("WEDATA_USER_UIN", "")  # wedata用户UIN
+        }
+
+        # 构建表属性（通过TBLPROPERTIES）
+        tbl_properties = {
+            "feature_table": "TRUE",
+            "primaryKeys": ",".join(primary_keys),
+            "comment": description or "",
+            **{f"{k}": v for k, v in (tags or {}).items()},
+            **{f"feature_{k}": v for k, v in (env_tags or {}).items()}
+        }
+
+        # 构建列定义
+        columns_ddl = []
+        for field in table_schema.fields:
+            data_type = field.dataType.simpleString().upper()
+            col_def = f"`{field.name}` {data_type}"
+            if not field.nullable:
+                col_def += " NOT NULL"
+            # 添加字段注释(如果metadata中有comment)
+            if field.metadata and "comment" in field.metadata:
+                comment = self._escape_sql_value(field.metadata["comment"])
+                col_def += f" COMMENT '{comment}'"
+            columns_ddl.append(col_def)
+
+        # 构建分区表达式
+        partition_expr = (
+            f"PARTITIONED BY ({', '.join([f'`{c}`' for c in partition_columns])})"
+            if partition_columns else ""
+        )
+
+        # 核心建表语句
+        ddl = f"""
+        CREATE TABLE {table_name} (
+            {', '.join(columns_ddl)}
+        )
+        USING PARQUET
+        {partition_expr}
+        TBLPROPERTIES (
+            {', '.join(f"'{k}'='{self._escape_sql_value(v)}'" for k, v in tbl_properties.items())}
+        )
+        """
+
+        # 打印sql
+        print(f"create table ddl: {ddl}")
+
+        # 执行DDL
+        try:
+            self._spark.sql(ddl)
+            if df is not None:
+                df.write.insertInto(table_name)
+        except Exception as e:
+            raise ValueError(f"建表失败: {str(e)}") from e
+
+    def write_table(
+            self,
+            name: str,
+            df: DataFrame,
+            mode: str = APPEND,
+            checkpoint_location: Optional[str] = None,
+            trigger: Optional[Dict[str, Any]] = DEFAULT_WRITE_STREAM_TRIGGER
+    ) -> Optional[StreamingQuery]:
+        """
+        写入特征表数据（支持批处理和流式写入）
+
+        Args:
+            name: 特征表名称（格式：<table>）
+            df: 要写入的数据（DataFrame）
+            mode: 写入模式（append/overwrite）
+            checkpoint_location: 流式写入的检查点位置（仅流式写入需要）
+            trigger: 流式写入触发条件（仅流式写入需要）
+
+        Returns:
+            如果是流式写入返回StreamingQuery对象，否则返回None
+
+        Raises:
+            ValueError: 当参数不合法时抛出
+        """
+
+        # 验证写入模式
+        valid_modes = ["append", "overwrite"]
+        if mode not in valid_modes:
+            raise ValueError(f"无效的写入模式 '{mode}'，可选值: {valid_modes}")
+
+        # 完整表名格式：<catalog>.<schema>.<table>
+        table_name = f'DataLakeCatalog.feature_store.{name}'
+
+        # 判断是否是流式DataFrame
+        is_streaming = df.isStreaming
+
+        try:
+            if is_streaming:
+                # 流式写入
+                if not checkpoint_location:
+                    raise ValueError("流式写入必须提供checkpoint_location参数")
+
+                writer = df.writeStream \
+                    .format("parquet") \
+                    .outputMode(mode) \
+                    .option("checkpointLocation", checkpoint_location)
+
+                if trigger:
+                    writer = writer.trigger(**trigger)
+
+                return writer.toTable(table_name)
+            else:
+                # 批处理写入
+                df.write \
+                    .mode(mode) \
+                    .insertInto(table_name)
+                return None
+
+        except Exception as e:
+            raise ValueError(f"写入表'{table_name}'失败: {str(e)}") from e
+
+    def read_table(
+            self,
+            name: str
+    ) -> DataFrame:
+        """
+        从特征表中读取数据
+
+        Args:
+            name: 特征表名称（格式：<table>）
+
+        Returns:
+            包含表数据的DataFrame
+
+        Raises:
+            ValueError: 当表不存在或读取失败时抛出
+        """
+        # 构建完整表名
+        table_name = f'DataLakeCatalog.feature_store.{name}'
+
+        try:
+            # 检查表是否存在
+            if not self._spark.catalog.tableExists(table_name):
+                raise ValueError(f"表 '{table_name}' 不存在")
+
+            # 读取表数据
+            return self._spark.read.table(table_name)
+
+        except Exception as e:
+            raise ValueError(f"读取表 '{table_name}' 失败: {str(e)}") from e
+
+    def drop_table(
+            self,
+            name: str
+    ) -> None:
+        """
+        删除特征表（表不存在时抛出异常）
+
+        Args:
+            name: 特征表名称（格式：<table>）
+
+        Raises:
+            ValueError: 当表不存在时抛出
+            RuntimeError: 当删除操作失败时抛出
+
+        示例:
+            # 基本删除
+            drop_table("user_features")
+        """
+        # 构建完整表名
+        table_name = f'DataLakeCatalog.feature_store.{name}'
+
+        try:
+            # 检查表是否存在
+            if not self._spark.catalog.tableExists(table_name):
+                raise ValueError(f"表 '{table_name}' 不存在")
+
+            # 执行删除
+            self._spark.sql(f"DROP TABLE {table_name}")
+
+        except ValueError as e:
+            raise  # 直接抛出已知的ValueError
+        except Exception as e:
+            raise RuntimeError(f"删除表 '{table_name}' 失败: {str(e)}") from e