wedata-feature-engineering 0.1.8__tar.gz → 0.1.13__tar.gz

{wedata-feature-engineering-0.1.8 → wedata-feature-engineering-0.1.13}/PKG-INFO

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

wedata-feature-engineering-0.1.13/tests/test_feature_store.py

@@ -0,0 +1,298 @@
+# This is a test script for FeatureStoreClient
+from datetime import date
+
+import pandas as pd
+from pyspark.sql import SparkSession, DataFrame
+from sklearn.ensemble import RandomForestClassifier
+
+import mlflow.sklearn
+
+from wedata.feature_store.client import FeatureStoreClient
+from pyspark.sql.types import StructType, StructField, StringType, TimestampType, IntegerType, DoubleType, DateType
+
+from wedata.feature_store.entities.feature_lookup import FeatureLookup
+from wedata.feature_store.entities.training_set import TrainingSet
+
+
+# 创建FeatureStoreClient实例
+def create_client() -> FeatureStoreClient:
+    spark = SparkSession.builder \
+        .appName("FeatureStoreDemo") \
+        .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension") \
+        .config("spark.sql.catalog.spark_catalog", "org.apache.spark.sql.delta.catalog.DeltaCatalog") \
+        .config("spark.jars.packages", "io.delta:delta-core_2.12:2.4.0") \
+        .enableHiveSupport() \
+        .getOrCreate()
+
+    # 创建FeatureStoreClient实例
+    client = FeatureStoreClient(spark)
+    return client
+
+# 创建特征表
+def create_table(client: FeatureStoreClient):
+    user_data = [
+        (1001, 25, "F", 120.5, date(2020, 5, 15)),  # user_id, age, gender, avg_purchase, member_since
+        (1002, 30, "M", 200.0, date(2019, 3, 10)),
+        (1003, 35, "F", 180.3, date(2021, 1, 20))
+    ]
+
+    # 定义schema
+    user_schema = StructType([
+        StructField("user_id", IntegerType(), False, metadata={"comment": "用户唯一标识ID"}),
+        StructField("age", IntegerType(), True, metadata={"comment": "用户年龄"}),
+        StructField("gender", StringType(), True, metadata={"comment": "用户性别(F-女性,M-男性)"}),
+        StructField("avg_purchase", DoubleType(), True, metadata={"comment": "用户平均消费金额"}),
+        StructField("member_since", DateType(), True, metadata={"comment": "用户注册日期"})
+    ])
+
+    # 创建DataFrame
+    user_df = client.spark.createDataFrame(user_data, user_schema)
+
+    display(user_df)
+
+    client.create_table(
+        name="user_features",  # 表名
+        primary_keys=["user_id"],      # 主键
+        df=user_df,                   # 数据
+        partition_columns=["member_since"],  # 按注册日期分区
+        description="用户基本特征和消费行为特征",  # 描述
+        tags={  # 业务标签
+            "create_by": "tencent",
+            "sensitivity": "internal"
+        }
+    )
+
+    # 商品数据样例
+    product_data = [
+        (5001, "电子", 599.0, 0.85, date(2024, 1, 1)),
+        (5002, "服装", 199.0, 0.92, date(2023, 11, 15)),
+        (5003, "家居", 299.0, 0.78, date(2024, 2, 20))
+    ]
+
+    # 定义schema
+    product_schema = StructType([
+        StructField("product_id", IntegerType(), False),
+        StructField("category", StringType(), True),
+        StructField("price", DoubleType(), True),
+        StructField("popularity", DoubleType(), True),
+        StructField("release_date", DateType(), True)
+    ])
+
+    # 创建DataFrame
+    product_df = client.spark.createDataFrame(product_data, product_schema)
+
+    display(product_df)
+
+    # 创建商品特征表
+    client.create_table(
+        name="product_features",
+        primary_keys=["product_id"],
+        df=product_df,
+        description="商品基本属性和受欢迎程度",
+        tags={  # 业务标签
+            "feature_table": "true",
+            "sensitivity": "internal"
+        }
+    )
+
+
+# 追加写入数据
+def append_data(client: FeatureStoreClient):
+    user_data = [
+        (1004, 45, "F", 120.5, date(2020, 5, 15)),
+        (1005, 55, "M", 200.0, date(2019, 3, 10)),
+        (1006, 65, "F", 180.3, date(2021, 1, 20))
+    ]
+
+    user_schema = StructType([
+        StructField("user_id", IntegerType(), False, metadata={"comment": "用户唯一标识ID"}),
+        StructField("age", IntegerType(), True, metadata={"comment": "用户年龄"}),
+        StructField("gender", StringType(), True, metadata={"comment": "用户性别(F-女性,M-男性)"}),
+        StructField("avg_purchase", DoubleType(), True, metadata={"comment": "用户平均消费金额"}),
+        StructField("member_since", DateType(), True, metadata={"comment": "用户注册日期"})
+    ])
+
+    user_df = client.spark.createDataFrame(user_data, user_schema)
+
+    display(user_df)
+
+    client.write_table(
+        name="user_features",
+        df=user_df,
+        mode="append"
+    )
+
+    product_data = [
+        (5007, "食品", 599.0, 0.85, date(2024, 1, 1)),
+        (5008, "玩具", 199.0, 0.92, date(2023, 11, 15)),
+        (5009, "电脑", 299.0, 0.78, date(2024, 2, 20))
+    ]
+
+    product_schema = StructType([
+        StructField("product_id", IntegerType(), False, metadata={"comment": "商品唯一标识ID"}),
+        StructField("category", StringType(), True, metadata={"comment": "商品类别"}),
+        StructField("price", DoubleType(), True, metadata={"comment": "商品价格(元)"}),
+        StructField("popularity", DoubleType(), True, metadata={"comment": "商品受欢迎程度(0-1)"}),
+        StructField("release_date", DateType(), True, metadata={"comment": "商品发布日期"})
+    ])
+
+    product_df = client.spark.createDataFrame(product_data, product_schema)
+
+    display(product_df)
+
+    client.write_table(
+        name="product_features",
+        df=product_df,
+        mode="append"
+    )
+
+# 读取特征表数据
+def read_table(client: FeatureStoreClient):
+
+    # 读取用户特征表
+    user_df = client.read_table("user_features")
+    display(user_df)
+
+    # 读取商品特征表
+    product_df = client.read_table("product_features")
+    display(product_df)
+
+# 获取特征表元数据
+def get_table(client: FeatureStoreClient):
+    feature_table_user = client.get_table(name="user_features")
+    print(feature_table_user)
+
+
+# 创建训练集
+def create_training_set(client: FeatureStoreClient)  -> TrainingSet:
+
+    # 订单数据样例
+    order_data = [
+        (9001, 1001, 5001, date(2025, 3, 1), 1, 0),
+        (9002, 1002, 5002, date(2025, 3, 2), 2, 1),
+        (9003, 1003, 5003, date(2025, 3, 3), 1, 0)
+    ]
+
+    # 定义schema
+    order_schema = StructType([
+        StructField("order_id", IntegerType(), False, metadata={"comment": "订单唯一标识ID"}),
+        StructField("user_id", IntegerType(), True, metadata={"comment": "用户ID"}),
+        StructField("product_id", IntegerType(), True, metadata={"comment": "商品ID"}),
+        StructField("order_date", DateType(), True, metadata={"comment": "订单日期"}),
+        StructField("quantity", IntegerType(), True, metadata={"comment": "购买数量"}),
+        StructField("is_returned", IntegerType(), True, metadata={"comment": "是否退货(0-未退货,1-已退货)"})
+    ])
+
+    # 创建DataFrame
+    order_df = client.spark.createDataFrame(order_data, order_schema)
+
+    # 查看订单数据
+    display(order_df)
+
+    # 定义用户特征查找
+    user_feature_lookup = FeatureLookup(
+        table_name="user_features",
+        feature_names=["age", "gender", "avg_purchase"],  # 选择需要的特征列
+        lookup_key="user_id"  # 关联键
+    )
+
+    # 定义商品特征查找
+    product_feature_lookup = FeatureLookup(
+        table_name="product_features",
+        feature_names=["category", "price", "popularity"],  # 选择需要的特征列
+        lookup_key="product_id"  # 关联键
+    )
+
+    # 创建训练集
+    training_set = client.create_training_set(
+        df=order_df,  # 基础数据
+        feature_lookups=[user_feature_lookup, product_feature_lookup],  # 特征查找配置
+        label="is_returned",  # 标签列
+        exclude_columns=["order_id"]  # 排除不需要的列
+    )
+
+    # 获取最终的训练DataFrame
+    training_df = training_set.load_df()
+
+    # 查看训练数据
+    display(training_df)
+
+    return training_set
+
+
+# 查看df中数据
+def display(df):
+
+    """
+    打印DataFrame的结构和数据
+
+    参数:
+        df (DataFrame): 要打印的Spark DataFrame
+        num_rows (int): 要显示的行数，默认为20
+        truncate (bool): 是否截断过长的列，默认为True
+    """
+    # 打印表结构
+    print("=== 表结构 ===")
+    df.printSchema()
+
+    # 打印数据
+    print("\n=== 数据示例 ===")
+    df.show(20, True)
+
+    # 打印行数统计
+    print(f"\n总行数: {df.count()}")
+
+
+def log_model(client: FeatureStoreClient,
+              training_set: TrainingSet
+              ):
+
+    # 初始化模型
+    model = RandomForestClassifier(
+        n_estimators=100,  # 增加树的数量提高模型稳定性
+        random_state=42    # 固定随机种子保证可复现性
+    )
+
+    # 获取训练数据并转换为Pandas格式
+    train_pd = training_set.load_df().toPandas()
+
+    # 特征工程处理
+    # 1. 处理分类特征
+    train_pd['gender'] = train_pd['gender'].map({'F': 0, 'M': 1})
+    train_pd = pd.get_dummies(train_pd, columns=['category'])
+
+    # 2. 处理日期特征（转换为距今天数）
+    current_date = pd.to_datetime('2025-04-19')  # 使用参考信息中的当前时间
+    train_pd['order_days'] = (current_date - pd.to_datetime(train_pd['order_date'])).dt.days
+    train_pd = train_pd.drop('order_date', axis=1)
+
+    # 3. 创建交互特征（价格*数量）
+    train_pd['total_amount'] = train_pd['price'] * train_pd['quantity']
+
+    # 分离特征和标签
+    X = train_pd.drop("is_returned", axis=1)
+    y = train_pd["is_returned"]
+
+    # 训练模型
+    model.fit(X, y)
+    # 记录模型到MLflow
+    with mlflow.start_run():
+        client.log_model(
+            model=model,
+            artifact_path="return_prediction_model",  # 更符合业务场景的路径名
+            flavor=mlflow.sklearn,
+            training_set=training_set,
+            registered_model_name="product_return_prediction_model"  # 更准确的模型名称
+        )
+
+
+
+# Press the green button in the gutter to run the script.
+if __name__ == '__main__':
+    client = create_client()
+    #create_table(client)
+    #append_data(client)
+    #read_table(client)
+    #get_table(client)
+    training_set = create_training_set(client)
+    #log_model(client, training_set)

{wedata-feature-engineering-0.1.8 → wedata-feature-engineering-0.1.13}/wedata/__init__.py

@@ -3,4 +3,4 @@ WeData Feature Engineering
 A toolkit for automated feature engineering
 """
 
-__version__ = "0.1.8"
+__version__ = "0.1.13"

{wedata-feature-engineering-0.1.8 → wedata-feature-engineering-0.1.13}/wedata/feature_store/client.py

@@ -178,7 +178,7 @@ class FeatureStoreClient:
                     raise ValueError("FeatureLookup must specify a table_name")
                 # 先校验表名格式是否合法
                 common_utils.validate_table_name(feature.table_name)
-                # 再构建完整表名
+                # 再构建完整表名，并赋值给FeatureLookup对象
                 feature.table_name = common_utils.build_full_table_name(feature.table_name)
 
         features = feature_lookups
@@ -202,6 +202,7 @@ class FeatureStoreClient:
             flavor: ModuleType,
             training_set: Optional[TrainingSet] = None,
             registered_model_name: Optional[str] = None,
+            model_registry_uri: Optional[str] = None,
             await_registration_for: int = mlflow.tracking._model_registry.DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
             infer_input_example: bool = False,
             **kwargs,
@@ -218,6 +219,7 @@ class FeatureStoreClient:
              flavor: MLflow模型类型模块(如mlflow.sklearn)
              training_set: 训练模型使用的TrainingSet对象(可选)
              registered_model_name: 要注册的模型名称(可选)
+             model_registry_uri: 模型注册中心地址(可选)
              await_registration_for: 等待模型注册完成的秒数(默认300秒)
              infer_input_example: 是否自动记录输入示例(默认False)
 
@@ -231,6 +233,7 @@ class FeatureStoreClient:
             flavor=flavor,
             training_set=training_set,
             registered_model_name=registered_model_name,
+            model_registry_uri=model_registry_uri,
             await_registration_for=await_registration_for,
             infer_input_example=infer_input_example,
             **kwargs

{wedata-feature-engineering-0.1.8 → wedata-feature-engineering-0.1.13}/wedata/feature_store/feature_table_client/feature_table_client.py

@@ -113,7 +113,7 @@ class FeatureTableClient:
         try:
             if self._spark.catalog.tableExists(table_name):
                 raise ValueError(
-                    f"Table '{table_name}' already exists\n"
+                    f"Table '{name}' already exists\n"
                     "Solutions:\n"
                     "1. Use a different table name\n"
                     "2. Drop the existing table: spark.sql(f'DROP TABLE {name}')\n"
@@ -125,11 +125,6 @@ class FeatureTableClient:
         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_key}' must be a primary key")
-            timestamp_keys_ddl.append(f"`{timestamp_key}` TIMESTAMP")
 
         #从环境变量获取额外标签
         env_tags = {
@@ -142,6 +137,7 @@ class FeatureTableClient:
         tbl_properties = {
             "feature_table": "TRUE",
             "primaryKeys": ",".join(primary_keys),
+            "timestampKeys": ",".join(timestamp_keys) if timestamp_keys else "",
             "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()}
@@ -171,7 +167,7 @@ class FeatureTableClient:
         CREATE TABLE {table_name} (
             {', '.join(columns_ddl)}
         )
-        USING PARQUET
+        USING iceberg
         {partition_expr}
         TBLPROPERTIES (
             {', '.join(f"'{k}'='{self._escape_sql_value(v)}'" for k, v in tbl_properties.items())}
@@ -293,13 +289,13 @@ class FeatureTableClient:
         try:
             # 检查表是否存在
             if not self._spark.catalog.tableExists(table_name):
-                raise ValueError(f"表 '{table_name}' 不存在")
+                raise ValueError(f"Table '{name}' does not exist")
 
             # 读取表数据
             return self._spark.read.table(table_name)
 
         except Exception as e:
-            raise ValueError(f"读取表 '{table_name}' 失败: {str(e)}") from e
+            raise ValueError(f"Failed to read table '{name}': {str(e)}") from e
 
     def drop_table(self, name: str):
 
@@ -327,15 +323,17 @@ class FeatureTableClient:
         try:
             # 检查表是否存在
             if not self._spark.catalog.tableExists(table_name):
-                raise ValueError(f"表 '{table_name}' 不存在")
+                print(f"Table '{name}' does not exist")
+                return
 
             # 执行删除
             self._spark.sql(f"DROP TABLE {table_name}")
+            print(f"Table '{name}' dropped")
 
         except ValueError as e:
             raise  # 直接抛出已知的ValueError
         except Exception as e:
-            raise RuntimeError(f"删除表 '{table_name}' 失败: {str(e)}") from e
+            raise RuntimeError(f"Failed to delete table '{name}': {str(e)}") from e
 
     def get_table(
             self,
@@ -365,4 +363,4 @@ class FeatureTableClient:
         try:
             return spark_client.get_feature_table(table_name)
         except Exception as e:
-            raise ValueError(f"获取表'{name}'元数据失败: {str(e)}") from e
+            raise ValueError(f"Failed to get metadata for table '{name}': {str(e)}") from e

{wedata-feature-engineering-0.1.8 → wedata-feature-engineering-0.1.13}/wedata/feature_store/spark_client/spark_client.py

@@ -9,13 +9,39 @@ from pyspark.sql.types import StructType, StringType, StructField
 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, sanitize_multi_level_name
+from wedata.feature_store.utils.common_utils import unsanitize_identifier
 
 
 class SparkClient:
     def __init__(self, spark: SparkSession):
         self._spark = spark
 
+    def _parse_table_name(self, table_name):
+        """解析表名并返回表名部分
+
+        参数:
+            table_name: 完整表名，支持格式: catalog.schema.table、schema.table 或 table
+
+        返回:
+            str: 解析后的表名部分
+        """
+        if not isinstance(table_name, str):
+            raise ValueError("Table name must be string type")
+
+        table_name = table_name.strip()
+        if not table_name:
+            raise ValueError("Table name cannot be empty")
+
+        parts = table_name.split('.')
+        if len(parts) == 3:
+            # 对于三部分名称(catalog.schema.table)，只使用表名部分
+            return parts[2]
+        elif len(parts) == 2:
+            # 对于两部分名称(schema.table)，只使用表名部分
+            return parts[1]
+        else:
+            # 单表名，直接使用
+            return table_name
 
     def get_current_catalog(self):
         """
@@ -66,19 +92,13 @@ class SparkClient:
         """
         try:
             # 解析表名
-            parts = table_name.split('.')
-            if len(parts) == 3:
-                catalog, schema, table = parts
-            elif len(parts) == 2:
-                schema, table = parts
-            else:
-                table = table_name
+            schema_table_name = self._parse_table_name(table_name)
 
             # 验证表是否存在
-            if not self._spark.catalog.tableExists(table):
+            if not self._spark.catalog.tableExists(schema_table_name):
                 raise ValueError(f"表不存在: {table_name}")
 
-            return self._spark.table(table)
+            return self._spark.table(schema_table_name)
 
         except Exception as e:
             raise ValueError(f"读取表 {table_name} 失败: {str(e)}")
@@ -86,23 +106,10 @@ class SparkClient:
 
     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}"
+        schema_table_name = self._parse_table_name(table_name)
 
         # 使用dbName.tableName格式查询列信息
-        columns = self._spark.catalog.listColumns(tableName=full_table_name)
+        columns = self._spark.catalog.listColumns(tableName=schema_table_name)
         return [
             Feature(
                 feature_table=table_name,
@@ -114,28 +121,26 @@ class SparkClient:
         ]
 
     def get_feature_table(self, table_name):
+        # 解析表名
+        schema_table_name = self._parse_table_name(table_name)
 
         # 获取表元数据
-        table = self._spark.catalog.getTable(table_name)
+        table = self._spark.catalog.getTable(schema_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_details = self._spark.sql(f"DESCRIBE TABLE EXTENDED {schema_table_name}").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)
+                table_properties = {}
+                for p in props:
+                    if "=" in p:
+                        parts = p.split("=", 1)
+                        key = parts[0].strip()
+                        value = parts[1].strip() if len(parts) > 1 else ""
+                        table_properties[key] = value
 
         # 获取特征列信息
         features = self.get_features(table_name)

{wedata-feature-engineering-0.1.8 → wedata-feature-engineering-0.1.13}/wedata/feature_store/training_set_client/training_set_client.py

@@ -186,6 +186,7 @@ class TrainingSetClient:
             flavor: ModuleType,
             training_set: Optional[TrainingSet],
             registered_model_name: Optional[str],
+            model_registry_uri: Optional[str],
             await_registration_for: int,
             infer_input_example: bool,
             **kwargs,
@@ -334,8 +335,7 @@ class TrainingSetClient:
             except Exception:
                 input_example = None
 
-            # todo:
-            #feature_spec.save(data_path)
+            feature_spec.save(data_path)
 
             # Log the packaged model. If no run is active, this call will create an active run.
             mlflow.pyfunc.log_model(
@@ -355,13 +355,12 @@ class TrainingSetClient:
             # If the user provided an explicit model_registry_uri when constructing the FeatureStoreClient,
             # we respect this by setting the registry URI prior to reading the model from Model
             # Registry.
-            # todo:
-            # if self._model_registry_uri:
-            #     # This command will override any previously set registry_uri.
-            #     mlflow.set_registry_uri(self._model_registry_uri)
+        if model_registry_uri is not None:
+             # This command will override any previously set registry_uri.
+            mlflow.set_registry_uri(model_registry_uri)
 
             mlflow.register_model(
                 "runs:/%s/%s" % (run_id, artifact_path),
                 registered_model_name,
                 await_registration_for=await_registration_for,
-                )
+            )

{wedata-feature-engineering-0.1.8 → wedata-feature-engineering-0.1.13}/wedata_feature_engineering.egg-info/PKG-INFO

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

{wedata-feature-engineering-0.1.8 → wedata-feature-engineering-0.1.13}/wedata_feature_engineering.egg-info/SOURCES.txt

@@ -1,5 +1,6 @@
 README.md
 setup.py
+tests/test_feature_store.py
 wedata/__init__.py
 wedata/feature_store/__init__.py
 wedata/feature_store/client.py