wedata-feature-engineering 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

feature_store/feature_table_client/feature_table_client.py

@@ -1,313 +0,0 @@
-"""
-特征表操作相关工具方法
-"""
-
-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

feature_store/spark_client/spark_client.py

@@ -1,286 +0,0 @@
-from collections import defaultdict
-from typing import List
-
-from pyspark.sql import SparkSession, DataFrame
-from pyspark.sql.catalog import Column
-from pyspark.sql.functions import when, isnull
-from pyspark.sql.types import StructType, StringType, StructField
-
-from feature_store.entities.feature import Feature
-from feature_store.entities.feature_table import FeatureTable
-from feature_store.entities.function_info import FunctionParameterInfo, FunctionInfo
-from feature_store.utils.common_utils import unsanitize_identifier
-from feature_store.utils.utils import sanitize_multi_level_name
-
-
-class SparkClient:
-    def __init__(self, spark: SparkSession):
-        self._spark = spark
-
-    def createDataFrame(self, data, schema) -> DataFrame:
-        return self._spark.createDataFrame(data, schema)
-
-    def read_table(
-            self, qualified_table_name, as_of_delta_timestamp=None, streaming=False
-    ):
-        """
-        Reads a Delta table, optionally as of some timestamp.
-        """
-        if streaming and as_of_delta_timestamp:
-            raise ValueError(
-                "Internal error: as_of_delta_timestamp cannot be specified when"
-                " streaming=True."
-            )
-
-        base_reader = (
-            # By default, Structured Streaming only handles append operations. Because
-            # we have a notion of primary keys, most offline feature store operations
-            # are not appends. For example, FeatureStoreClient.write_table(mode=MERGE)
-            # will issue a MERGE operation.
-            # In order to propagate the non-append operations to the
-            # readStream, we set ignoreChanges to "true".
-            # For more information,
-            # see https://docs.databricks.com/delta/delta-streaming.html#ignore-updates-and-deletes
-            self._spark.readStream.format("delta").option("ignoreChanges", "true")
-            if streaming
-            else self._spark.read.format("delta")
-        )
-
-        if as_of_delta_timestamp:
-            return base_reader.option("timestampAsOf", as_of_delta_timestamp).table(
-                sanitize_multi_level_name(qualified_table_name)
-            )
-        else:
-            return base_reader.table(sanitize_multi_level_name(qualified_table_name))
-
-    def get_current_catalog(self):
-        """
-        Get current set catalog in the spark context.
-        """
-        try:
-            df = self._spark.sql("SELECT CURRENT_CATALOG()").collect()
-            return unsanitize_identifier(df[0][0])
-        except Exception as e:
-            return None
-
-    def get_current_database(self):
-        """
-        Get current set database in the spark context.
-        """
-        try:
-            df = self._spark.sql("SELECT CURRENT_DATABASE()").collect()
-            return unsanitize_identifier(df[0][0])
-        except Exception as e:
-            return None
-
-    def read_table(self, table_name):
-        """读取Spark表数据
-
-        Args:
-            table_name: 表名，支持格式: catalog.schema.table、schema.table 或 table
-
-        Returns:
-            DataFrame: 表数据
-
-        Raises:
-            ValueError: 当表不存在或读取失败时抛出
-        """
-        try:
-            # 解析表名
-            parts = table_name.split('.')
-            if len(parts) == 3:
-                catalog, schema, table = parts
-            elif len(parts) == 2:
-                schema, table = parts
-            else:
-                table = table_name
-
-            # 验证表是否存在
-            if not self._spark.catalog.tableExists(table):
-                raise ValueError(f"表不存在: {table_name}")
-
-            return self._spark.table(table)
-
-        except Exception as e:
-            raise ValueError(f"读取表 {table_name} 失败: {str(e)}")
-
-
-    def get_features(self, table_name):
-        # 解析表名
-        parts = table_name.split('.')
-        if len(parts) == 3:
-            # 对于三部分名称(catalog.schema.table)，使用schema.table格式
-            _, schema, table = parts
-            full_table_name = f"{schema}.{table}"
-        elif len(parts) == 2:
-            # 对于两部分名称(schema.table)，直接使用
-            full_table_name = table_name
-        else:
-            # 单表名，使用当前数据库
-            current_db = self.get_current_database()
-            if not current_db:
-                raise ValueError("无法确定当前数据库")
-            full_table_name = f"{current_db}.{table_name}"
-
-        # 使用dbName.tableName格式查询列信息
-        columns = self._spark.catalog.listColumns(tableName=full_table_name)
-        return [
-            Feature(
-                feature_table=table_name,
-                feature_id=f"{table_name}_{row.name}",
-                name=row.name,
-                data_type=row.dataType,
-                description=row.description or ""
-            ) for row in columns
-        ]
-
-    def get_online_stores(self, table_name):
-        return None
-
-
-
-    def get_feature_table(self, table_name):
-
-        # 获取表元数据
-        table = self._spark.catalog.getTable(table_name)
-
-        parts = table_name.split('.')
-        if len(parts) == 3:
-            # 对于三部分名称(catalog.schema.table)，只使用表名部分
-            table_to_describe = parts[2]
-        elif len(parts) == 2:
-            # 对于两部分名称(schema.table)，只使用表名部分
-            table_to_describe = parts[1]
-        else:
-            # 单表名，直接使用
-            table_to_describe = table_name
-        # 获取表详细信息
-        table_details = self._spark.sql(f"DESCRIBE TABLE EXTENDED {table_to_describe}").collect()
-
-        table_properties = {}
-        for row in table_details:
-            if row.col_name == "Table Properties":
-                props = row.data_type[1:-1].split(", ")
-                table_properties = dict(p.split("=") for p in props if "=" in p)
-
-        # 获取特征列信息
-        features = self.get_features(table_name)
-
-        # 构建完整的FeatureTable对象
-        return FeatureTable(
-            name=table_name,
-            table_id=table_properties.get("table_id", table_name),
-            description=table.description or table_properties.get("description", table_name),
-            primary_keys=table_properties.get("primaryKeys", "").split(",") if table_properties.get("primaryKeys") else [],
-            partition_columns=table.partitionColumnNames if hasattr(table, 'partitionColumnNames') else [],
-            features=features,
-            creation_timestamp=None,  # Spark表元数据不包含创建时间戳
-            online_stores=self.get_online_stores(table_name),
-            notebook_producers=None,
-            job_producers=None,
-            table_data_sources=None,
-            path_data_sources=None,
-            custom_data_sources=None,
-            timestamp_keys=table_properties.get("timestamp_keys"),
-            tags=table_properties.get("tags")
-        )
-
-    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.createDataFrame(
-            [full_routine_name.split(".") for full_routine_name in full_routine_names],
-            routine_name_schema,
-        )
-        routines_table = self.read_table(
-            "system.information_schema.routines"
-        )
-        parameters_table = self.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
-
-