数据工程与 AI 流水线完全指南：从 Spark 到 Kafka

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, count, avg, when
from pyspark.ml.feature import VectorAssembler
from pyspark.ml.classification import RandomForestClassifier

spark = SparkSession.builder \
    .appName("AI Feature Engineering") \
    .config("spark.sql.adaptive.enabled", "true") \
    .getOrCreate()

# 加载并转换数据
df = spark.read.parquet("s3://data-lake/features/")
df_clean = df.dropna() \
    .withColumn("age_group", when(col("age") < 30, "young")
                              .when(col("age") < 50, "middle")
                              .otherwise("senior"))

# 聚合统计
stats_df = df_clean.groupBy("category").agg(
    count("*").alias("cnt"),
    avg("score").alias("avg_score")
)

# 构建特征向量
assembler = VectorAssembler(
    inputCols=["age", "income", "score"],
    outputCol="features"
)
ml_df = assembler.transform(df_clean)

# 训练模型
rf = RandomForestClassifier(numTrees=100, featuresCol="features")
model = rf.fit(ml_df)

from kafka import KafkaProducer, KafkaConsumer
import json
import numpy as np

# 生产者：发送实时特征数据
producer = KafkaProducer(
    bootstrap_servers=['localhost:9092'],
    value_serializer=lambda v: json.dumps(v).encode('utf-8')
)

def send_feature(data):
    producer.send('ml-features', value={
        'timestamp': data['ts'],
        'features': data['features'].tolist(),
        'user_id': data['user_id']
    })

# 消费者：实时 AI 推理
consumer = KafkaConsumer(
    'ml-features',
    bootstrap_servers=['localhost:9092'],
    value_deserializer=lambda m: json.loads(m.decode('utf-8'))
)

import onnxruntime as ort
session = ort.InferenceSession('model.onnx')

for msg in consumer:
    features = np.array(msg.value['features']).reshape(1, -1)
    prediction = session.run(None, {'input': features.astype(np.float32)})[0]
    print(f"User {msg.value['user_id']}: {prediction}")

from airflow.decorators import dag, task
from datetime import datetime

@dag(schedule='@daily', start_date=datetime(2026, 1, 1))
def ml_training_pipeline():

    @task
    def extract_data():
        # 提取数据
        return {'rows': 10000, 'path': '/tmp/data.parquet'}

    @task
    def preprocess(info: dict):
        # 预处理
        import pandas as pd
        df = pd.read_parquet(info['path'])
        df_clean = df.dropna()
        return df_clean.to_dict()

    @task
    def train_model(data: dict):
        # 模型训练
        from sklearn.ensemble import GradientBoostingClassifier
        import mlflow

        with mlflow.start_run():
            model = GradientBoostingClassifier()
            # training logic here
            mlflow.sklearn.log_model(model, "model")
        return "model_trained"

    @task
    def validate_model(status: str):
        # 模型验证
        if status == "model_trained":
            return {"accuracy": 0.92, "passed": True}
        return {"passed": False}

    raw = extract_data()
    processed = preprocess(raw)
    result = train_model(processed)
    validate_model(result)

dag_instance = ml_training_pipeline()

models/
├── staging/          # 清洗原始源数据
│   ├── stg_orders.sql
│   └── stg_customers.sql
├── intermediate/     # 业务逻辑的中间转换
│   └── int_order_items.sql
└── marts/           # 面向最终业务的表
    ├── fct_orders.sql
    └── dim_customers.sql

-- models/marts/fct_orders.sql
SELECT
    o.order_id,
    o.customer_id,
    c.customer_name,
    o.total_amount,
    o.created_at
FROM {{ ref('stg_orders') }} o
LEFT JOIN {{ ref('dim_customers') }} c
    ON o.customer_id = c.customer_id
WHERE o.status = 'completed'

# schema.yml
models:
  - name: fct_orders
    description: '已完成订单事实表'
    columns:
      - name: order_id
        description: '订单唯一 ID'
        tests:
          - unique
          - not_null
      - name: total_amount
        tests:
          - not_null
          - dbt_utils.accepted_range:
              min_value: 0

from delta.tables import DeltaTable
from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension") \
    .config("spark.sql.catalog.spark_catalog",
            "org.apache.spark.sql.delta.catalog.DeltaCatalog") \
    .getOrCreate()

# Upsert (Merge) 操作
delta_table = DeltaTable.forPath(spark, "/delta/users")
updates_df = spark.read.parquet("/tmp/updates/")

delta_table.alias("target").merge(
    updates_df.alias("source"),
    "target.user_id = source.user_id"
).whenMatchedUpdateAll() \
 .whenNotMatchedInsertAll() \
 .execute()

# 时间旅行：查询昨天的数据
df_yesterday = spark.read.format("delta") \
    .option("timestampAsOf", "2026-03-16") \
    .load("/delta/users")

# 按版本查询
df_v1 = spark.read.format("delta") \
    .option("versionAsOf", 1) \
    .load("/delta/users")

# 查看变更历史
delta_table.history().show()

spark.sql("""
    OPTIMIZE delta.`/delta/events`
    ZORDER BY (user_id, event_date)
""")

from feast import FeatureStore, Entity, FeatureView, Field
from feast.types import Float32, Int64

# 定义特征
user_stats = FeatureView(
    name="user_stats",
    entities=["user_id"],
    schema=[
        Field(name="avg_purchase_amount", dtype=Float32),
        Field(name="purchase_count_30d", dtype=Int64),
    ],
    ttl=timedelta(days=30),
)

# 在线特征查询(实时推理)
store = FeatureStore(repo_path=".")
features = store.get_online_features(
    features=["user_stats:avg_purchase_amount", "user_stats:purchase_count_30d"],
    entity_rows=[{"user_id": 1001}],
).to_dict()

import great_expectations as gx

context = gx.get_context()
datasource = context.sources.add_pandas("my_datasource")

# 定义期望(Expectation)
suite = context.add_expectation_suite("orders_suite")
validator = context.get_validator(
    batch_request=batch_request,
    expectation_suite_name="orders_suite"
)

validator.expect_column_values_to_not_be_null("order_id")
validator.expect_column_values_to_be_between("total_amount", min_value=0, max_value=100000)
validator.expect_column_values_to_be_unique("order_id")
validator.save_expectation_suite()

# 执行校验
checkpoint = context.add_or_update_checkpoint(
    name="orders_checkpoint",
    validations=[{"batch_request": batch_request, "expectation_suite_name": "orders_suite"}],
)
result = checkpoint.run()

from evidently.report import Report
from evidently.metric_preset import DataDriftPreset

report = Report(metrics=[DataDriftPreset()])
report.run(reference_data=train_df, current_data=prod_df)
report.show(mode='inline')