4. 设计特征存储与特征工程流程

任务目标

构建一个轻量级特征存储系统,统一管理特征的定义、计算逻辑与版本,并保证离线训练与在线推理获取的特征值一致。本案例聚焦于:特征注册、批量计算(离线)、单条计算(在线)以及一致性校验。

核心设计

  • 特征定义:使用Python dataclass声明特征名、来源、计算函数及版本。
  • 存储抽象:离线侧写入Parquet文件,在线侧写入Redis(模拟),通过统一API屏蔽差异。
  • 一致性保障:同一特征仅维护一份计算逻辑,在线与离线共用同一函数,从源头避免偏差。

完整实现

import hashlib
import pandas as pd
from dataclasses import dataclass, field
from typing import Callable, Dict, Any
import redis
import json

# ---------- 1. 特征注册中心 ----------
@dataclass
class FeatureSpec:
    name: str
    version: str
    compute_fn: Callable[[Dict[str, Any]], Any]
    description: str = ""

class FeatureStore:
    def __init__(self, redis_host="localhost", redis_port=6379, db=0):
        self._registry: Dict[str, FeatureSpec] = {}
        self._redis = redis.Redis(host=redis_host, port=redis_port, db=db, decode_responses=True)

    def register(self, spec: FeatureSpec):
        key = f"{spec.name}:{spec.version}"
        self._registry[key] = spec

    def _get_spec(self, name: str, version: str) -> FeatureSpec:
        key = f"{name}:{version}"
        if key not in self._registry:
            raise KeyError(f"特征 {name} 版本 {version} 未注册")
        return self._registry[key]

    # ---------- 2. 离线批量计算 ----------
    def compute_offline(self, df: pd.DataFrame, features: list[tuple[str, str]]) -> pd.DataFrame:
        """输入原始DataFrame,输出包含特征列的DataFrame"""
        result = df.copy()
        for name, version in features:
            spec = self._get_spec(name, version)
            # 逐行应用计算函数(生产环境可用向量化或Spark,此处为演示)
            result[name] = df.apply(lambda row: spec.compute_fn(row.to_dict()), axis=1)
        return result

    # ---------- 3. 在线单条计算与存储 ----------
    def compute_online(self, raw_data: Dict[str, Any], name: str, version: str) -> Any:
        """计算单条特征值,并写入Redis"""
        spec = self._get_spec(name, version)
        value = spec.compute_fn(raw_data)
        # 使用特征名+版本+原始数据哈希作为key,便于追溯
        raw_hash = hashlib.md5(json.dumps(raw_data, sort_keys=True).encode()).hexdigest()[:8]
        redis_key = f"feat:{name}:{version}:{raw_hash}"
        self._redis.set(redis_key, json.dumps({"value": value, "raw": raw_data}))
        return value

    # ---------- 4. 一致性校验 ----------
    def verify_consistency(self, df: pd.DataFrame, features: list[tuple[str, str]]) -> bool:
        """对比离线批量结果与在线逐条计算+存储的结果是否一致"""
        offline_df = self.compute_offline(df, features)
        all_consistent = True
        for idx, row in df.iterrows():
            raw = row.to_dict()
            for name, version in features:
                online_val = self.compute_online(raw, name, version)
                offline_val = offline_df.loc[idx, name]
                if online_val != offline_val:
                    print(f"不一致: 行{idx} 特征{name} 离线={offline_val} 在线={online_val}")
                    all_consistent = False
        return all_consistent

# ---------- 5. 特征定义示例 ----------
def _user_age_feature(raw: Dict[str, Any]) -> int:
    """根据出生年份计算年龄(示例特征)"""
    birth_year = int(raw["birth_year"])
    return 2025 - birth_year

def _user_level_feature(raw: Dict[str, Any]) -> str:
    """根据消费金额划分用户等级(示例特征)"""
    amount = float(raw["total_spend"])
    if amount > 10000:
        return "high"
    elif amount > 1000:
        return "mid"
    return "low"

# ---------- 6. 最小运行闭环 ----------
if __name__ == "__main__":
    # 初始化存储并注册特征
    store = FeatureStore()
    store.register(FeatureSpec(name="user_age", version="v1", compute_fn=_user_age_feature))
    store.register(FeatureSpec(name="user_level", version="v1", compute_fn=_user_level_feature))

    # 模拟原始数据(离线场景)
    raw_df = pd.DataFrame([
        {"user_id": "u1", "birth_year": 1990, "total_spend": 15000},
        {"user_id": "u2", "birth_year": 1985, "total_spend": 500},
        {"user_id": "u3", "birth_year": 2000, "total_spend": 3000},
    ])

    # 执行一致性校验(内部会分别跑离线和在线)
    feature_list = [("user_age", "v1"), ("user_level", "v1")]
    is_consistent = store.verify_consistency(raw_df, feature_list)

    # 查看离线计算结果
    offline_result = store.compute_offline(raw_df, feature_list)
    print("离线特征结果:")
    print(offline_result[["user_id", "user_age", "user_level"]])

    # 查看Redis中存储的在线特征
    print("\nRedis中的在线特征key:")
    keys = store._redis.keys("feat:*")
    for k in keys:
        print(f"  {k}")

    print(f"\n一致性校验通过: {is_consistent}")

运行与验证

  1. 依赖安装:pip install pandas redis,并确保本地Redis服务已启动(默认端口6379)。
  2. 执行:直接运行上述脚本。
  3. 通过条件:
    • 控制台输出“一致性校验通过: True”。
    • 离线结果中user_age依次为35、40、25;user_level依次为high、low、mid。
    • Redis中能看到6个以feat:开头的key(3条数据×2个特征)。

关键说明

  • 本案例使用同一计算函数作为离线和在线的唯一逻辑来源,这是保证一致性的核心手段。生产环境中,离线可替换为Spark批处理,在线可替换为Flink或模型服务,但函数注册与调用模式不变。
  • Redis key中加入原始数据哈希,用于追踪某条特征值由哪条原始记录产生,便于排查数据漂移或回溯。
  • 特征版本号(v1)纳入注册key,支持同一特征多版本并存与灰度切换。