ML Data Pipeline Architecture Skill

ML Data Pipeline Architecture

Patterns for efficient ML data pipelines using Polars, Arrow, and ClickHouse.

ADR: 2026-01-22-polars-preference-hook (efficiency preferences framework)

Note: A PreToolUse hook enforces Polars preference. To use Pandas, add # polars-exception: <reason> at file top.

When to Use This Skill

Use this skill when:

Deciding between Polars and Pandas for a data pipeline
Optimizing memory usage with zero-copy Arrow patterns
Loading data from ClickHouse into PyTorch DataLoaders
Implementing lazy evaluation for large datasets
Migrating existing Pandas code to Polars

1. Decision Tree: Polars vs 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. Zero-Copy Pipeline Architecture

The Problem with 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

The Solution with 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 Integration Patterns

Pattern A: Arrow Stream (Recommended)

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'")

Pattern B: Polars Native (Simpler)

# 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"
)

Pattern C: Parquet Export (Batch Jobs)

# 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 Integration

Minimal Change Pattern

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)

Custom PolarsDataset (Large Data)

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. Lazy Evaluation Patterns

Pipeline Composition

# 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()

Streaming for Large Files

# 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. Schema Validation

Pydantic for Config

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 Schema Validation

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. Performance Benchmarks

| Operation | Pandas | Polars | Speedup | | -------------- | ------ | ------ | ------- | | Read CSV (1GB) | 45s | 4s | 11x | | Filter rows | 2.1s | 0.4s | 5x | | Group-by agg | 3.8s | 0.3s | 13x | | Sort | 5.2s | 0.4s | 13x | | Memory peak | 10GB | 2.5GB | 4x |

Benchmark: 50M rows, 20 columns, MacBook M2

8. Migration Checklist

Phase 1: Add Arrow Support

[ ] Add polars = "<version>" to dependencies (see PyPI)
[ ] Implement query_arrow() in data client
[ ] Verify zero-copy with memory profiler

Phase 2: Polars at Entry Points

[ ] Add pl.from_pandas() wrapper at trainer entry
[ ] Update prepare_sequences() to accept both types
[ ] Add schema validation after conversion

Phase 3: Full Lazy Evaluation

[ ] Convert file reads to pl.scan_*
[ ] Compose transformations lazily
[ ] Call .collect() only before .to_numpy()

9. Anti-Patterns to Avoid

DON'T: Mix APIs Unnecessarily

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

DON'T: Collect Too Early

# 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()

DON'T: Ignore Memory Pressure

# 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)

References

Troubleshooting

| Issue | Cause | Solution | | --------------------------- | -------------------------------- | ---------------------------------------------------- | | Memory spike during load | Collecting too early | Use lazy evaluation, call collect() only when needed | | Arrow conversion fails | Unsupported data type | Check for object columns, convert to native types | | ClickHouse connection error | Wrong port or credentials | Verify host:8123 (HTTP) or host:9000 (native) | | Zero-copy not working | Intermediate pandas conversion | Remove to_pandas() calls, stay in Arrow/Polars | | Polars hook blocking code | Pandas used without exception | Add # polars-exception: reason comment at file top | | Slow group-by operations | Using pandas for large datasets | Migrate to Polars for 5-10x speedup | | Schema validation failure | Column names case-sensitive | Verify exact column names from source | | PyTorch DataLoader OOM | Loading full dataset into memory | Use PolarsDataset with Arrow backing for lazy access | | Parquet scan performance | Not using predicate pushdown | Add filters before collect() for lazy evaluation | | Type mismatch in tensor | Float64 vs Float32 mismatch | Explicitly cast with .cast(pl.Float32) before numpy |

Agent Skills: ML Data Pipeline Architecture

Install this agent skill to your local

Skill Files