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

wedata/feature_store/mlflow_model.py

@@ -0,0 +1,17 @@
+import json
+from typing import Optional, Dict, Any
+import mlflow
+import os
+
+class _FeatureStoreModelWrapper(mlflow.pyfunc.PythonModel):
+    def __init__(self, model):
+        self.model = model
+
+    def predict(self, context, model_input):
+       return self.model.predict(model_input)
+
+# def _load_pyfunc(path):
+#     # Path provided by mlflow is subdirectory of path needed by score_batch
+#     artifact_path = os.path.join(mlflow.pyfunc.DATA, "feature_store")
+#     index = path.find(artifact_path)
+#     return _FeatureStoreModelWrapper(path[:index])

wedata/feature_store/spark_client/spark_client.py

@@ -0,0 +1,289 @@
+from collections import defaultdict
+from typing import Optional, Any, Dict, List
+
+import logging
+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 mlflow.pyfunc import spark_udf
+
+from wedata.feature_store.constants.constants import (
+    _PREBUILT_ENV_URI
+)
+
+from wedata.feature_store.entities.feature import Feature
+from wedata.feature_store.entities.feature_table import FeatureTable
+from wedata.feature_store.entities.function_info import FunctionParameterInfo, FunctionInfo
+from wedata.feature_store.utils.common_utils import unsanitize_identifier, check_spark_table_exists, check_package_version
+
+
+class SparkClient:
+    def __init__(self, spark: SparkSession):
+        self._spark = spark
+
+    def get_current_catalog(self):
+        """
+        获取当前Spark会话的catalog名称（使用spark.catalog.currentCatalog属性）
+
+        返回:
+            str: 当前catalog名称，如果未设置则返回None
+        """
+        _, ok, _ = check_package_version("pyspark", "3.4.0", ">=")
+        print(f"pyspark >= 3.4.0  {ok}")
+        if ok:
+            return unsanitize_identifier(self._spark.catalog.currentCatalog())
+        else:
+            catalog = self._spark.sql("SELECT current_catalog()").first()[0]
+            return unsanitize_identifier(catalog)
+
+
+    def get_current_database(self):
+        """
+        获取Spark上下文中当前设置的数据库名称
+
+        返回:
+            str: 当前数据库名称，如果获取失败则返回None
+        """
+        try:
+            # 使用Spark SQL查询当前数据库
+            df = self._spark.sql("SELECT CURRENT_DATABASE()")
+            # 获取第一行第一列的值并去除特殊字符
+            return unsanitize_identifier(df.first()[0])
+        except Exception:
+            # 捕获所有异常并返回None
+            return None
+
+    def createDataFrame(self, data, schema) -> DataFrame:
+        return self._spark.createDataFrame(data, schema)
+
+
+    def read_table(self, table_name):
+        """读取Spark表数据
+
+        Args:
+            table_name: 表名，支持格式: catalog.schema.table、schema.table
+
+        Returns:
+            DataFrame: 表数据
+
+        Raises:
+            ValueError: 当表不存在或读取失败时抛出
+        """
+        table_name = _translate_spark_table_name(table_name)
+        try:
+            # 验证表是否存在
+            if not  check_spark_table_exists(self._spark, table_name):
+                raise ValueError(f"Table does not exist: {table_name}")
+            return self._spark.table(table_name)
+
+        except Exception as e:
+            raise ValueError(f"Failed to read table {table_name}: {str(e)}")
+
+    def get_features(self, table_name):
+        # 查询列信息
+        table_name = _translate_spark_table_name(table_name)
+        split = table_name.split(".")
+        if len(split) == 2:
+            # db.table_name
+            columns = self._spark.catalog.listColumns(tableName=split[1], dbName=split[0])
+        elif len(split) == 3:
+            # catalog.db.table_name
+            columns = self._spark.catalog.listColumns(tableName=split[2], dbName=split[1])
+        else:
+            columns = self._spark.catalog.listColumns(tableName=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_feature_table(self, table_name):
+        """
+        DLC支持table_name为catalog.schema.table
+        EMR支持table_name为schema.table
+        """
+
+        table_name = _translate_spark_table_name(table_name)
+        # table = self._spark.catalog.getTable(table_name)
+        # 获取表配置信息
+        properties = self._spark.sql(f"SHOW TBLPROPERTIES {table_name}").collect()
+        primary_key_str = next((row.value for row in properties if row.key == "primaryKeys"), None)
+        primary_keys = primary_key_str.split(",") if primary_key_str else []
+        table_id = next((row.value for row in properties if row.key == "table_id"), table_name)
+        description = next((row.value for row in properties if row.key == "comment"), None)
+        timestamp_keys_str = next((row.value for row in properties if row.key == "timestampKeys"), None)
+        timestamp_keys = timestamp_keys_str.split(",") if timestamp_keys_str else []
+        # 获取分区字段信息
+        desc_df = self._spark.sql(f"DESCRIBE EXTENDED {table_name}")
+        partition_info = desc_df.filter("col_name LIKE '_partition%'").collect()
+        partition_columns = []
+        if partition_info:
+            partition_str = partition_info[0]["data_type"]
+            # 从分区字符串中提取分区字段
+            if partition_str.startswith("struct<") and partition_str.endswith(">"):
+                # 去掉struct<>外壳
+                fields_str = partition_str[7:-1]
+                # 分割各个字段定义
+                field_defs = [f.strip() for f in fields_str.split(",") if f.strip()]
+                # 提取字段名
+                partition_columns = [f.split(":")[0].strip() for f in field_defs]
+
+        # 获取特征列信息
+        features = self.get_features(table_name)
+
+        # 构建完整的FeatureTable对象
+        return FeatureTable(
+            name=table_name,
+            table_id=table_id,
+            description=description,
+            primary_keys=primary_keys,
+            partition_columns=partition_columns,
+            features=features,
+            creation_timestamp=None,  # Spark表元数据不包含创建时间戳
+            online_stores=None,
+            notebook_producers=None,
+            job_producers=None,
+            table_data_sources=None,
+            path_data_sources=None,
+            custom_data_sources=None,
+            timestamp_keys=timestamp_keys,
+            tags=None
+        )
+
+    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
+
+    def get_predict_udf(
+        self,
+        model_uri,
+        result_type=None,
+        env_manager=None,
+        params: Optional[dict[str, Any]] = None,
+        prebuilt_env_uri: Optional[str] = None,
+    ):
+        kwargs = {}
+        if result_type:
+            kwargs["result_type"] = result_type
+        if env_manager:
+            kwargs["env_manager"] = env_manager
+        if params:
+            kwargs["params"] = params
+        if prebuilt_env_uri:
+            kwargs[_PREBUILT_ENV_URI] = prebuilt_env_uri
+
+        return spark_udf(self._spark, model_uri, **kwargs)
+
+
+def _translate_spark_table_name(table_name):
+    from wedata.feature_store.constants.engine_types import judge_engine_type, CalculateEngineTypes
+    # 获取表元数据
+    if judge_engine_type() == CalculateEngineTypes.EMR:
+        split_names = table_name.split(".")
+        # print(f"==== EMR TABLE split len({len(split_names)})")
+        if len(split_names) <= 2:
+            return table_name
+        else:
+            table_name = ".".join(table_name.split(".")[1:])
+            return table_name
+    return table_name
+