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

• Polars User Guide

• Polars Migration Guide

• Apache Arrow Python

• ClickHouse Python Client

• PyTorch Data Loading

• Polars Preference Hook ADR

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