ml-data-pipeline-architecture

ML 数据管道架构

使用 Polars、Arrow 和 ClickHouse 构建高效 ML 数据管道的模式。

ADR：2026-01-22-polars-preference-hook（效率偏好框架）

注意：PreToolUse 钩子会强制优先使用 Polars。若要使用 Pandas，请在文件顶部添加 # polars-exception: <reason>。

自我演进技能：此技能会在使用过程中不断改进。如果说明有误、参数已变更或需要采用变通方案，请立即修复此文件，不要拖延。仅针对真实且可复现的问题进行更新。

何时使用此技能

在以下情况下使用此技能：

• 为数据管道选择 Polars 还是 Pandas
• 使用零拷贝 Arrow 模式优化内存用量
• 将 ClickHouse 中的数据加载到 PyTorch DataLoader
• 为大型数据集实现惰性求值
• 将现有 Pandas 代码迁移到 Polars

---

1. 决策树：Polars 还是 Pandas

Dataset size?
├─ < 1M rows → Pandas OK (simpler API, richer ecosystem)
├─ 1M-10M rows → Consider Polars (2-5x faster, less memory)
└─ > 10M rows → Use Polars (required for memory efficiency)

Operations?
├─ Simple transforms → Either works
├─ Group-by aggregations → Polars 5-10x faster
├─ Complex joins → Polars with lazy evaluation
└─ Streaming/chunked → Polars scan_* functions

Integration?
├─ scikit-learn heavy → Pandas (better interop)
├─ PyTorch/custom → Polars + Arrow (zero-copy to tensor)
└─ ClickHouse source → Arrow stream → Polars (optimal)

---

2. 零拷贝管道架构

Pandas 的问题

# BAD: 3 memory copies
df = pd.read_sql(query, conn)     # Copy 1: DB → pandas
X = df[features].values           # Copy 2: pandas → numpy
tensor = torch.from_numpy(X)      # Copy 3: numpy → tensor
# Peak memory: 3x data size

使用 Arrow 的解决方案

# GOOD: 0-1 memory copies
import clickhouse_connect
import polars as pl
import torch

client = clickhouse_connect.get_client(...)
arrow_table = client.query_arrow("SELECT * FROM bars")  # Arrow in DB memory
df = pl.from_arrow(arrow_table)                          # Zero-copy view
X = df.select(features).to_numpy()                       # Single allocation
tensor = torch.from_numpy(X)                             # View (no copy)
# Peak memory: 1.2x data size

---

3. ClickHouse 集成模式

模式 A：Arrow 流（推荐）

def query_arrow(client, query: str) -> pl.DataFrame:
    """ClickHouse → Arrow → Polars (zero-copy chain)."""
    arrow_table = client.query_arrow(f"{query} FORMAT ArrowStream")
    return pl.from_arrow(arrow_table)

# Usage
df = query_arrow(client, "SELECT * FROM bars WHERE ts >= '2024-01-01'")

模式 B：原生 Polars（更简单）

# Polars has native ClickHouse support (see pola.rs for version requirements)
df = pl.read_database_uri(
    query="SELECT * FROM bars",
    uri="clickhouse://user:pass@host/db"
)

模式 C：Parquet 导出（批处理作业）

# For reproducible batch processing
client.query("SELECT * FROM bars INTO OUTFILE 'data.parquet' FORMAT Parquet")
df = pl.scan_parquet("data.parquet")  # Lazy, memory-mapped

---

4. PyTorch DataLoader 集成

最小改动模式

from torch.utils.data import TensorDataset, DataLoader

# Accept both pandas and polars
def prepare_data(df) -> tuple[torch.Tensor, torch.Tensor]:
    if isinstance(df, pd.DataFrame):
        df = pl.from_pandas(df)

    X = df.select(features).to_numpy()
    y = df.select(target).to_numpy()

    return (
        torch.from_numpy(X).float(),
        torch.from_numpy(y).float()
    )

X, y = prepare_data(df)
dataset = TensorDataset(X, y)
loader = DataLoader(dataset, batch_size=32, pin_memory=True)

自定义 PolarsDataset（大型数据）

class PolarsDataset(torch.utils.data.Dataset):
    """Memory-efficient dataset from Polars DataFrame."""

    def __init__(self, df: pl.DataFrame, features: list[str], target: str):
        self.arrow = df.to_arrow()  # Arrow backing for zero-copy slicing
        self.features = features
        self.target = target

    def __len__(self) -> int:
        return self.arrow.num_rows

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        row = self.arrow.slice(idx, 1)
        x = torch.tensor([row[f][0].as_py() for f in self.features], dtype=torch.float32)
        y = torch.tensor(row[self.target][0].as_py(), dtype=torch.float32)
        return x, y

---

5. 惰性求值模式

管道组合

# Define transformations lazily (no computation yet)
pipeline = (
    pl.scan_parquet("raw_data.parquet")
    .filter(pl.col("timestamp") >= start_date)
    .with_columns([
        (pl.col("close").pct_change()).alias("returns"),
        (pl.col("volume").log()).alias("log_volume"),
    ])
    .select(features + [target])
)

# Execute only when needed
train_df = pipeline.filter(pl.col("timestamp") < split_date).collect()
test_df = pipeline.filter(pl.col("timestamp") >= split_date).collect()

大型文件的流式处理

# Process file in chunks (never loads full file)
def process_large_file(path: str, chunk_size: int = 100_000):
    reader = pl.scan_parquet(path)

    for batch in reader.iter_batches(n_rows=chunk_size):
        # Process each chunk
        features = compute_features(batch)
        yield features.to_numpy()

---

6. 模式验证

使用 Pydantic 验证配置

from pydantic import BaseModel, field_validator

class FeatureConfig(BaseModel):
    features: list[str]
    target: str
    seq_len: int = 15

    @field_validator("features")
    @classmethod
    def validate_features(cls, v):
        required = {"returns_vs", "momentum_z", "atr_pct"}
        missing = required - set(v)
        if missing:
            raise ValueError(f"Missing required features: {missing}")
        return v

DataFrame 模式验证

def validate_schema(df: pl.DataFrame, required: list[str], stage: str) -> None:
    """Fail-fast schema validation."""
    missing = [c for c in required if c not in df.columns]
    if missing:
        raise ValueError(
            f"[{stage}] Missing columns: {missing}\n"
            f"Available: {sorted(df.columns)}"
        )

---

7. 性能基准测试

操作	Pandas	Polars	加速比
读取 CSV（1GB）	45s	4s	11x
筛选行	2.1s	0.4s	5x
分组聚合	3.8s	0.3s	13x
排序	5.2s	0.4s	13x
内存峰值	10GB	2.5GB	4x

基准测试：50M 行、20 列、MacBook M2

---

8. 迁移检查清单

阶段 1：添加 Arrow 支持

• 在依赖项中添加 polars = "<version>"（参见 PyPI）
• 在数据客户端中实现 query_arrow()
• 使用内存分析器验证零拷贝

阶段 2：在入口点使用 Polars

• 在训练器入口添加 pl.from_pandas() 包装器
• 更新 prepare_sequences()，使其同时接受两种类型
• 转换后添加模式验证

阶段 3：全面启用惰性求值

• 将文件读取转换为 pl.scan_*
• 以惰性方式组合转换操作
• 仅在调用 .to_numpy() 之前调用 .collect()

---

9. 应避免的反模式

不要：不必要地混用 API

# BAD: Convert back and forth
df_polars = pl.from_pandas(df_pandas)
df_pandas_again = df_polars.to_pandas()  # Why?

不要：过早收集数据

# BAD: Defeats lazy evaluation
df = pl.scan_parquet("data.parquet").collect()  # Full load
filtered = df.filter(...)  # After the fact

# GOOD: Filter before collect
df = pl.scan_parquet("data.parquet").filter(...).collect()

不要：忽视内存压力

# BAD: Loads entire file
df = pl.read_parquet("huge_file.parquet")

# GOOD: Stream in chunks
for batch in pl.scan_parquet("huge_file.parquet").iter_batches():
    process(batch)

---

参考资料

• Polars 用户指南
• Polars 迁移指南
• Apache Arrow Python
• ClickHouse Python 客户端
• PyTorch 数据加载
• Polars 偏好钩子 ADR

---

故障排查

问题	原因	解决方案
加载期间内存激增	过早收集	使用惰性求值，仅在需要时调用 collect()
Arrow 转换失败	数据类型不受支持	检查是否存在 object 列，并转换为原生类型
ClickHouse 连接错误	端口或凭据错误	验证 host:8123（HTTP）或 host:9000（原生）
零拷贝不起作用	中间进行了 pandas 转换	移除 to_pandas() 调用，保持使用 Arrow/Polars
Polars 钩子阻止代码运行	未使用例外说明的 pandas 代码	在文件顶部添加 # polars-exception: reason 注释
group-by 操作缓慢	对大型数据集使用 pandas	迁移到 Polars，可提速 5-10 倍
架构验证失败	列名区分大小写	验证源数据中的准确列名
PyTorch DataLoader 内存溢出	将整个数据集加载到内存	使用基于 Arrow 的 PolarsDataset 实现惰性访问
Parquet 扫描性能不佳	未使用谓词下推	在 collect() 前添加筛选条件以进行惰性求值
张量类型不匹配	Float64 与 Float32 不匹配	转换为 numpy 前，使用 .cast(pl.Float32) 显式转换

执行后反思

此技能完成后，请在结束任务前进行反思：

0. 确认所在位置。 — 编辑前，找到此 SKILL.md 的规范路径。
1. 哪里失败了？ — 修复导致问题的指令。
2. 有什么效果超出预期？ — 将其提升为推荐做法。
3. 哪里出现了偏差？ — 修复与实际情况不再匹配的脚本、参考资料或依赖项。
4. 记录下来。 — 添加一条演进日志，记录触发原因、修复内容和证据。

不要拖延。下一次调用会继承你留下的一切。