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

feature_store/spark_client/spark_client.py

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

feature_store/training_set_client/training_set_client.py

@@ -0,0 +1,196 @@
+import json
+import logging
+import os
+from collections import defaultdict
+from types import ModuleType
+from typing import Any, Dict, List, Optional, Set, Union
+
+import mlflow
+import yaml
+from mlflow.models import Model, ModelSignature
+from mlflow.utils.file_utils import TempDir, YamlSafeDumper, read_yaml
+from pyspark.sql import DataFrame
+from pyspark.sql.functions import struct
+
+from feature_store.entities.feature_column_info import FeatureColumnInfo
+from feature_store.entities.feature_function import FeatureFunction
+from feature_store.entities.feature_lookup import FeatureLookup
+from feature_store.entities.feature_spec import FeatureSpec
+from feature_store.entities.training_set import TrainingSet
+from feature_store.spark_client.spark_client import SparkClient
+
+from feature_store.constants.constants import (
+    _NO_RESULT_TYPE_PASSED,
+    _PREBUILT_ENV_URI,
+    _USE_SPARK_NATIVE_JOIN,
+    _WARN,
+    MODEL_DATA_PATH_ROOT,
+    PREDICTION_COLUMN_NAME,
+)
+
+from feature_store.utils import common_utils, training_set_utils
+from feature_store.utils.feature_spec_utils import convert_to_yaml_string
+
+_logger = logging.getLogger(__name__)
+
+FEATURE_SPEC_GRAPH_MAX_COLUMN_INFO = 1000
+
+
+class TrainingSetClient:
+    def __init__(
+        self,
+        spark_client: SparkClient
+    ):
+        self._spark_client = spark_client
+
+    def create_training_set(
+        self,
+        feature_spec: FeatureSpec,
+        feature_column_infos: List[FeatureColumnInfo],
+        label_names: List[str],
+        df: DataFrame,
+        ft_metadata: training_set_utils._FeatureTableMetadata,
+        kwargs,
+    ):
+        uc_function_infos = training_set_utils.get_uc_function_infos(
+            self._spark_client,
+            {odci.udf_name for odci in feature_spec.on_demand_column_infos},
+        )
+
+        # TODO(divyagupta-db): Move validation from _validate_join_feature_data in feature_lookup_utils.py
+        #  to a helper function called here and in score_batch.
+
+        # Add consumer of each feature and instrument as final step
+        consumer_feature_table_map = defaultdict(list)
+        for feature in feature_column_infos:
+            consumer_feature_table_map[feature.table_name].append(feature.feature_name)
+        consumed_udf_names = [f.udf_name for f in feature_spec.function_infos]
+
+        # Spark query planning is known to cause spark driver to crash if there are many feature tables to PiT join.
+        # See https://docs.google.com/document/d/1EyA4vvlWikTJMeinsLkxmRAVNlXoF1eqoZElOdqlWyY/edit
+        # So we disable native join by default.
+        training_set_utils.warn_if_non_photon_for_native_spark(
+            kwargs.get(_USE_SPARK_NATIVE_JOIN, False), self._spark_client
+        )
+        return TrainingSet(
+            feature_spec,
+            df,
+            label_names,
+            ft_metadata.feature_table_metadata_map,
+            ft_metadata.feature_table_data_map,
+            uc_function_infos,
+            kwargs.get(_USE_SPARK_NATIVE_JOIN, False),
+        )
+
+    def create_training_set_from_feature_lookups(
+        self,
+        df: DataFrame,
+        feature_lookups: List[Union[FeatureLookup, FeatureFunction]],
+        label: Union[str, List[str], None],
+        exclude_columns: List[str],
+        **kwargs,
+    ) -> TrainingSet:
+
+        # 获取特征查找列表和特征函数列表
+        features = feature_lookups
+        feature_lookups = [f for f in features if isinstance(f, FeatureLookup)]
+        feature_functions = [f for f in features if isinstance(f, FeatureFunction)]
+
+        # 如果未提供标签，则用空列表初始化label_names
+        label_names = common_utils.as_list(label, [])
+        del label
+
+        # 校验数据集和标签
+        training_set_utils.verify_df_and_labels(df, label_names, exclude_columns)
+
+        # 获取特征表元数据
+        ft_metadata = training_set_utils.get_table_metadata(
+            self._spark_client,
+            {fl.table_name for fl in feature_lookups}
+        )
+
+        column_infos = training_set_utils.get_column_infos(
+            feature_lookups,
+            feature_functions,
+            ft_metadata,
+            df_columns=df.columns,
+            label_names=label_names,
+        )
+
+        training_set_utils.validate_column_infos(
+            self._spark_client,
+            ft_metadata,
+            column_infos.source_data_column_infos,
+            column_infos.feature_column_infos,
+            column_infos.on_demand_column_infos,
+            label_names,
+        )
+
+        # Build feature_spec locally for comparison with the feature spec yaml generated by the
+        # FeatureStore backend. This will be removed once the migration is validated.
+        feature_spec = training_set_utils.build_feature_spec(
+            feature_lookups,
+            ft_metadata,
+            column_infos,
+            exclude_columns
+        )
+
+        return self.create_training_set(
+            feature_spec,
+            column_infos.feature_column_infos,
+            label_names,
+            df,
+            ft_metadata,
+            kwargs=kwargs,
+        )
+
+
+
+
+
+    def create_feature_spec(
+        self,
+        name: str,
+        features: List[Union[FeatureLookup, FeatureFunction]],
+        sparkClient: SparkClient,
+        exclude_columns: List[str] = [],
+    ) -> FeatureSpec:
+
+        feature_lookups = [f for f in features if isinstance(f, FeatureLookup)]
+        feature_functions = [f for f in features if isinstance(f, FeatureFunction)]
+
+        # Maximum of 100 FeatureFunctions is supported
+        if len(feature_functions) > training_set_utils.MAX_FEATURE_FUNCTIONS:
+            raise ValueError(
+                f"A maximum of {training_set_utils.MAX_FEATURE_FUNCTIONS} FeatureFunctions are supported."
+            )
+
+        # Get feature table metadata and column infos
+        ft_metadata = training_set_utils.get_table_metadata(
+            self._spark_client,
+            {fl.table_name for fl in feature_lookups}
+        )
+        column_infos = training_set_utils.get_column_infos(
+            feature_lookups,
+            feature_functions,
+            ft_metadata,
+        )
+
+        column_infos = training_set_utils.add_inferred_source_columns(column_infos)
+
+        training_set_utils.validate_column_infos(
+            self._spark_client,
+            ft_metadata,
+            column_infos.source_data_column_infos,
+            column_infos.feature_column_infos,
+            column_infos.on_demand_column_infos,
+        )
+
+        feature_spec = training_set_utils.build_feature_spec(
+            feature_lookups,
+            ft_metadata,
+            column_infos,
+            exclude_columns
+        )
+
+        return feature_spec

{wedata_feature_engineering-0.1.0.dist-info → wedata_feature_engineering-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: wedata-feature-engineering
-Version: 0.1.0
+Version: 0.1.2
 Summary: Wedata Feature Engineering Library
 Home-page: 
 Author: meahqian

wedata_feature_engineering-0.1.2.dist-info/RECORD

@@ -0,0 +1,30 @@
+feature_store/__init__.py,sha256=CP3YAMoy3pSTWRYzTza_CYBnGbTv_KzycVEBMQCeiD8,101
+feature_store/client.py,sha256=FG1xK460rD859iSY4VA75XeYhqStJD8Wlr0sRxk25LI,5267
+feature_store/constants/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+feature_store/constants/constants.py,sha256=exW3kiFLDyCmU9cYHFjcvIQhPWEpFtkogLXeB9Arfd8,827
+feature_store/entities/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+feature_store/entities/column_info.py,sha256=WezowI46YHDym5ZlbhCJDqhKbVcjXjnjt7dQdy3XqYM,4164
+feature_store/entities/data_type.py,sha256=VpHS6Fr3TphQQ8NbAcEnDJ-8eOZV6ivYuWxv3pAM2RM,3394
+feature_store/entities/environment_variables.py,sha256=ZEFml5H9MQuzBKM074mUrFYu-Sga4Knmxqiwpke2WGc,1679
+feature_store/entities/feature.py,sha256=wX8fTBlJq3GYdj9rrBDCY3kFgcVBBAiOOZdxEhnQkNQ,1241
+feature_store/entities/feature_column_info.py,sha256=-TGxRafYUaNKe0YzHus2XbfRaVrMv7pcffMdbtTT4nA,2031
+feature_store/entities/feature_function.py,sha256=R17INrCE-U_Uj9KLbFz69aYlOkTETTwQHMMo470F4lQ,1865
+feature_store/entities/feature_lookup.py,sha256=zUDMdDIboitOffYRZlurf_O_4UeBPmE5YS0PyCS2Fqg,7912
+feature_store/entities/feature_spec.py,sha256=F4MiKEyvKZSBh6Uv7V4vVLbamZ9fRClaC3HCrUeynDE,20079
+feature_store/entities/feature_spec_constants.py,sha256=YWDBfRiNDe6fUJFUBo3V4WYg2xsljoPAE-ZejfFZCgM,785
+feature_store/entities/feature_table.py,sha256=4ghopIvJcoIlyFiSEuTkOcDWn88c1Kt6q5LWM4BYEHI,6073
+feature_store/entities/feature_table_info.py,sha256=2vUaVdW_jw1dRAlmJWvBRueuMeuqWu_NYB9SlxLI7Uw,1126
+feature_store/entities/function_info.py,sha256=l0kmiq2R_QNfSMJ7y0xZohlMiemgYSr1dN5vzV8ijIs,7314
+feature_store/entities/on_demand_column_info.py,sha256=Eh5ieaj1TxC7DG6ipBZzH2ZyY0bwkLrDOkuZjgYr4gY,1297
+feature_store/entities/source_data_column_info.py,sha256=a9jQOJvehwDIrKPwsP6W9YRBSPNK2nZYypE6-p80CwA,542
+feature_store/entities/training_set.py,sha256=9H2uGnUxTAsk93Om50QxRELbeFCocwGMze2VexPVJWI,5569
+feature_store/feature_table_client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+feature_store/feature_table_client/feature_table_client.py,sha256=uir33K7oigrSnjTT6VbNOp0Nb22-X3JHd1_92kWjrow,10754
+feature_store/spark_client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+feature_store/spark_client/spark_client.py,sha256=vd-NCE9IGC0Ygqr-QSVY0teuWsQSkq_BFV4Mn6xMMNU,11578
+feature_store/training_set_client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+feature_store/training_set_client/training_set_client.py,sha256=Aa80xVXVE1KBdgplL9qqR8ftD5A5r2pfBttAhmySrB0,6696
+wedata_feature_engineering-0.1.2.dist-info/METADATA,sha256=IALf_mmflM-eRTOIOqVDJg5OoVVfLXBUHofIdC1T_wI,493
+wedata_feature_engineering-0.1.2.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+wedata_feature_engineering-0.1.2.dist-info/top_level.txt,sha256=15761LgVdJ7tJWbdlYk0EZ560G9k6C4TE42dfLx8d0I,14
+wedata_feature_engineering-0.1.2.dist-info/RECORD,,

wedata_feature_engineering-0.1.0.dist-info/RECORD

@@ -1,6 +0,0 @@
-feature_store/__init__.py,sha256=CP3YAMoy3pSTWRYzTza_CYBnGbTv_KzycVEBMQCeiD8,101
-feature_store/client.py,sha256=FG1xK460rD859iSY4VA75XeYhqStJD8Wlr0sRxk25LI,5267
-wedata_feature_engineering-0.1.0.dist-info/METADATA,sha256=gTLzZnQR7SI9-sp0NWq8JYYwF1edzt2iodypP9mkvNk,493
-wedata_feature_engineering-0.1.0.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-wedata_feature_engineering-0.1.0.dist-info/top_level.txt,sha256=15761LgVdJ7tJWbdlYk0EZ560G9k6C4TE42dfLx8d0I,14
-wedata_feature_engineering-0.1.0.dist-info/RECORD,,