Optimizing Python for Production Data Science: Cython, Polars, Vectorization, and When to Drop to C
A practical guide to optimizing Python for production data science workloads, covering Cython, Polars/DuckDB, vectorization strategies, profiling techniques, and the decision framework for when to drop from Python to C.
Optimizing Python for Production Data Science: Cython, Polars, Vectorization, and When to Drop to C
Python is the lingua franca of data science, but it’s also notoriously slow. A naive Python loop that processes 10 million rows can take hours; the same operation in C takes seconds. For notebook-based exploration, this doesn’t matter — you can wait. But when your feature engineering pipeline must process terabytes of data daily within a cost budget, or your model inference must serve predictions in under 10ms, Python’s performance limitations become a real constraint.
After nine years of optimizing Python for production data science, I’ve developed a pragmatic framework for when and how to optimize. This post covers the full spectrum — from simple vectorization tricks to Cython extensions to knowing when Python is simply the wrong tool.
The Optimization Hierarchy
Before diving into specific techniques, here’s the hierarchy I follow. Each level should be exhausted before moving to the next:
- Algorithm optimization: Use a better algorithm. O(n log n) beats O(n²) regardless of language.
- Vectorization: Replace loops with NumPy/Pandas vectorized operations.
- Library upgrade: Switch from Pandas to Polars or DuckDB for analytical workloads.
- JIT compilation: Use Numba for numerical Python code.
- Cython: Write C extensions for hot paths that can’t be vectorized.
- Drop to C/Rust: For performance-critical components that need maximum speed.
Critical principle: Always profile before optimizing. I’ve seen teams spend weeks optimizing code that accounted for 2% of total runtime. Profile first, optimize the bottleneck, repeat.
Profiling: Finding the Bottleneck
cProfile for Function-Level Profiling
import cProfile
import pstats
def feature_pipeline(data):
# ... pipeline code ...
pass
# Profile the pipeline
cProfile.run('feature_pipeline(data)', 'profile_output')
# Analyze results
stats = pstats.Stats('profile_output')
stats.sort_stats('cumulative')
stats.print_stats(20) # Top 20 functions by cumulative time
cProfile tells you which functions are slow, but not which lines within those functions are slow. For that, you need line-level profiling.
line_profiler for Line-Level Profiling
# Install: pip install line_profiler
# Usage: kernprof -l -v my_script.py
@profile
def compute_features(df):
features = {}
for col in df.columns: # Line 1: 0.1%
features[f'{col}_mean'] = df[col].mean() # Line 2: 45.2%
features[f'{col}_std'] = df[col].std() # Line 3: 42.1%
features[f'{col}_skew'] = df[col].skew() # Line 4: 12.6%
return features
Now you know exactly where the time is spent. In this example, mean(), std(), and skew() are the bottlenecks — and they’re already vectorized, so further optimization requires a different approach (like using Polars or computing all three in a single pass).
memory_profiler for Memory Profiling
# Install: pip install memory_profiler
# Usage: python -m memory_profiler my_script.py
@profile
def load_and_process(filepath):
df = pd.read_parquet(filepath) # Line 1: +2.3 GB
df = df.dropna() # Line 2: +0.1 GB (temporary copy)
features = compute_features(df) # Line 3: +0.5 GB
return features # Total: ~3 GB peak
Memory profiling is critical for production pipelines that process large datasets. A pipeline that works on your 16GB laptop might OOM on a production container with 4GB.
Vectorization: The First Optimization
The 100x Rule
Vectorized NumPy/Pandas operations are typically 10-100x faster than equivalent Python loops. This is because vectorized operations push the loop into C code inside NumPy, avoiding Python’s per-element overhead.
Common Vectorization Patterns
Pattern 1: Replace apply with vectorized operations
# SLOW: apply with lambda
df['is_weekend'] = df['date'].apply(lambda x: x.weekday() >= 5)
# FAST: vectorized dt accessor
df['is_weekend'] = df['date'].dt.weekday >= 5
# SPEEDUP: ~50x
Pattern 2: Replace iterrows with itertuples or vectorized
# SLOW: iterrows
for idx, row in df.iterrows():
df.at[idx, 'result'] = row['a'] * row['b'] + row['c']
# FAST: vectorized
df['result'] = df['a'] * df['b'] + df['c']
# SPEEDUP: ~100x
Pattern 3: Use np.where instead of apply with if/else
# SLOW: apply with conditional
df['category'] = df['value'].apply(lambda x: 'high' if x > 100 else 'low')
# FAST: np.where
df['category'] = np.where(df['value'] > 100, 'high', 'low')
# SPEEDUP: ~30x
Pattern 4: Use np.select for multiple conditions
# SLOW: apply with multiple conditions
def categorize(x):
if x < 10: return 'tiny'
elif x < 100: return 'small'
elif x < 1000: return 'medium'
else: return 'large'
df['size'] = df['value'].apply(categorize)
# FAST: np.select
conditions = [
df['value'] < 10,
df['value'] < 100,
df['value'] < 1000,
]
choices = ['tiny', 'small', 'medium']
df['size'] = np.select(conditions, choices, default='large')
# SPEEDUP: ~40x
When Vectorization Isn’t Possible
Some operations can’t be easily vectorized:
- Complex stateful computations (cumulative operations with custom logic)
- String operations with complex regex patterns
- Iterative algorithms (EM, Gibbs sampling)
- Operations that depend on previous results (sequential dependencies)
For these, you need Numba, Cython, or a different tool entirely.
Numba: JIT Compilation for Numerical Python
When to Use Numba
Numba compiles Python functions to machine code at runtime using LLVM. It’s ideal for:
- Numerical loops that can’t be vectorized
- Custom aggregation functions
- Mathematical computations with tight inner loops
Numba Example
from numba import jit
import numpy as np
# Without Numba: ~5 seconds for 10M elements
def rolling_custom_stat(values, window):
n = len(values)
result = np.empty(n)
for i in range(n):
start = max(0, i - window + 1)
window_vals = values[start:i + 1]
# Custom statistic: trimmed mean
sorted_vals = np.sort(window_vals)
trim = max(1, len(sorted_vals) // 10)
result[i] = np.mean(sorted_vals[trim:-trim])
return result
# With Numba: ~0.05 seconds for 10M elements (100x speedup)
@jit(nopython=True)
def rolling_custom_stat_numba(values, window):
n = len(values)
result = np.empty(n)
for i in range(n):
start = max(0, i - window + 1)
window_vals = values[start:i + 1]
sorted_vals = np.sort(window_vals)
trim = max(1, len(sorted_vals) // 10)
result[i] = np.mean(sorted_vals[trim:-trim])
return result
Numba Limitations
- First-call overhead: Numba compiles on first call, which can take 1-2 seconds. Subsequent calls are fast.
- Limited Python support: Numba doesn’t support all Python features. No dictionaries, limited string support, no Pandas operations.
- Numerical code only: Numba is designed for numerical computations. For string processing or complex data structures, use Cython instead.
Polars and DuckDB: Replacing Pandas for Analytical Workloads
Why Pandas Is Slow
Pandas is slow for three reasons:
- Single-threaded: Pandas operations run on a single CPU core.
- Eager evaluation: Every operation is executed immediately, even if it could be optimized by deferring execution.
- Memory copies: Many Pandas operations create copies of the data, doubling memory usage.
Polars: The Modern Alternative
Polars is a DataFrame library written in Rust that addresses all three Pandas limitations:
- Multi-threaded: Automatically parallelizes operations across CPU cores.
- Lazy evaluation: Optimizes query plans before execution.
- Zero-copy: Minimizes memory allocations and copies.
import polars as pl
# Pandas (single-threaded, eager)
import pandas as pd
df = pd.read_parquet('large_file.parquet')
result = df.groupby('category').agg({
'value': ['mean', 'std', 'count'],
'amount': 'sum',
}).reset_index()
# Polars (multi-threaded, lazy)
lf = pl.scan_parquet('large_file.parquet') # Lazy - doesn't read yet
result = (
lf.group_by('category')
.agg([
pl.col('value').mean().alias('value_mean'),
pl.col('value').std().alias('value_std'),
pl.col('value').count().alias('value_count'),
pl.col('amount').sum().alias('amount_sum'),
])
.collect() # Executes the optimized query plan
)
Performance Comparison
For a typical analytical workload (group-by aggregation on 100M rows):
| Operation | Pandas | Polars | DuckDB |
|---|---|---|---|
| Group-by aggregation | 12.3s | 1.8s | 1.5s |
| Filter + aggregate | 8.7s | 1.2s | 1.0s |
| Join (100M × 1M) | 25.1s | 3.4s | 2.8s |
| Window functions | 18.5s | 2.1s | 1.9s |
Polars and DuckDB are typically 5-10x faster than Pandas for analytical workloads. For my production pipelines, this translates to significant cost savings — the same work that required a 32-core Spark cluster can now run on a single machine.
DuckDB: SQL for Python
DuckDB is an embedded analytical database that runs inside your Python process. It’s ideal for:
- Ad-hoc analytical queries on Parquet/CSV files.
- Complex SQL queries that would be awkward in Pandas/Polars.
- Integration with existing SQL-based workflows.
import duckdb
# Query Parquet files directly
result = duckdb.sql("""
SELECT
category,
AVG(value) as avg_value,
STDDEV(value) as std_value,
COUNT(*) as n
FROM 'data/*.parquet'
WHERE date >= '2025-01-01'
GROUP BY category
HAVING COUNT(*) > 100
ORDER BY avg_value DESC
""").df() # Convert to Pandas DataFrame
When to Use What
- Pandas: Small datasets (< 1M rows), interactive exploration, compatibility with visualization libraries.
- Polars: Large datasets (1M-1B rows), production pipelines, when you need lazy evaluation and multi-threading.
- DuckDB: SQL-comfortable teams, ad-hoc analysis on files, complex analytical queries.
- Spark: Distributed datasets (> 1B rows), when data doesn’t fit on a single machine.
Cython: When You Need C Performance
When Cython Is the Right Tool
Cython is ideal when:
- You have a hot loop that can’t be vectorized and Numba can’t handle it.
- You need to interface with C/C++ libraries.
- You need predictable, consistent performance (Numba’s JIT compilation can vary).
- You’re building a reusable library that others will use.
Cython Example
# feature_engine.pyx
import numpy as np
cimport numpy as np
from libc.math cimport sqrt, fabs
def compute_ema(np.ndarray[double, ndim=1] values, double alpha):
"""
Compute exponential moving average.
Pure Python: ~2.5 seconds for 10M elements.
Cython: ~0.03 seconds for 10M elements (80x speedup).
"""
cdef int n = values.shape[0]
cdef np.ndarray[double, ndim=1] result = np.empty(n, dtype=np.float64)
cdef double ema = values[0]
cdef int i
result[0] = ema
for i in range(1, n):
ema = alpha * values[i] + (1 - alpha) * ema
result[i] = ema
return result
Cython Compilation
# setup.py
from setuptools import setup
from Cython.Build import cythonize
import numpy as np
setup(
ext_modules=cythonize("feature_engine.pyx"),
include_dirs=[np.get_include()],
)
python setup.py build_ext --inplace
Cython vs. Numba Decision Matrix
| Criterion | Numba | Cython |
|---|---|---|
| Ease of use | ★★★★★ | ★★★ |
| First-call latency | 1-2s JIT | None (pre-compiled) |
| String support | Limited | Full |
| C interop | Limited | Excellent |
| Debugging | Difficult | Standard C tools |
| Deployment | Simple (pip install) | Requires compilation |
My recommendation: Start with Numba for numerical code. Switch to Cython if you need string handling, C interop, or consistent performance without JIT warmup.
When to Drop to C/Rust
The Decision Framework
Dropping from Python to C or Rust is a significant investment. I only do it when:
- The component is a bottleneck: It accounts for > 30% of total runtime.
- Python optimization is exhausted: Vectorization, Numba, and Cython have been tried and are insufficient.
- The component is stable: The logic won’t change frequently. Rewriting in C every time the requirements change is expensive.
- The performance requirement is strict: Sub-millisecond latency, real-time processing, or extreme throughput.
What to Write in C/Rust
From my experience, the components most often worth rewriting in C/Rust:
- Custom distance/similarity functions: Cosine similarity, DTW, custom metrics for specific domains.
- Feature extraction from unstructured data: Text tokenization, image preprocessing, audio feature extraction.
- Tree traversal for ensemble models: When you need sub-millisecond inference for tree-based models.
- Data serialization/deserialization: Custom binary formats for high-throughput data pipelines.
Integration Approaches
- ctypes/cffi: Call C functions from Python directly. Simplest approach, minimal overhead.
- pybind11: Create Python bindings for C++ code. More convenient than ctypes for complex APIs.
- Rust + PyO3: Write Rust code with Python bindings. Growing ecosystem, excellent performance.
- Cython as glue: Use Cython to interface between Python and C code.
# Example: Calling a C function via ctypes
import ctypes
import numpy as np
# Load the compiled C library
lib = ctypes.CDLL('./libfeature.so')
# Define the function signature
lib.cosine_similarity.argtypes = [
ctypes.POINTER(ctypes.c_float), # vector a
ctypes.POINTER(ctypes.c_float), # vector b
ctypes.c_int, # dimension
]
lib.cosine_similarity.restype = ctypes.c_float
# Call from Python
a = np.array([1.0, 2.0, 3.0], dtype=np.float32)
b = np.array([4.0, 5.0, 6.0], dtype=np.float32)
similarity = lib.cosine_similarity(
a.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
b.ctypes.data_as(ctypes.POINTER(ctypes.c_float)),
len(a),
)
Real-World Optimization Case Study
The Problem
A feature engineering pipeline that processed 500GB of clickstream data daily. The pipeline was taking 14 hours on a 16-core machine, exceeding the 8-hour budget.
Profiling Results
Total runtime: 14.2 hours
- Data loading (Pandas read_parquet): 2.1 hours (15%)
- Feature computation (groupby + agg): 8.5 hours (60%)
- Data writing (Pandas to_parquet): 1.8 hours (13%)
- Other (preprocessing, validation): 1.8 hours (12%)
Optimization Steps
Step 1: Switch from Pandas to Polars (estimated 5x improvement)
- Data loading: 2.1h → 0.4h
- Feature computation: 8.5h → 1.7h
- Data writing: 1.8h → 0.4h
- Total: 14.2h → 3.2h
Step 2: Use lazy evaluation (estimated 20% improvement)
- Polars lazy evaluation optimizes the query plan, reducing unnecessary intermediate computations.
- Total: 3.2h → 2.6h
Step 3: Parallelize across cores (Polars does this automatically)
- Verify that Polars is using all 16 cores. Set
POLARS_MAX_THREADS=16. - Total: 2.6h → 2.1h
Step 4: Optimize the remaining bottleneck (custom aggregation)
- One custom aggregation function couldn’t be expressed in Polars. Rewrite in Cython.
- Total: 2.1h → 1.8h
Final Result
14.2 hours → 1.8 hours. An 8x improvement, well within the 8-hour budget. The optimizations were:
- Pandas → Polars (biggest impact)
- Lazy evaluation
- Cython for one custom function
No C/Rust was needed. The Python ecosystem is powerful enough for most data science workloads if you use the right tools.
Lessons Learned
Lesson 1: Profile Before Optimizing
I’ve seen teams spend weeks optimizing code that accounted for 2% of runtime. Always profile first. The bottleneck is rarely where you think it is.
Lesson 2: Algorithm First, Implementation Second
A better algorithm beats a faster implementation. If you can reduce O(n²) to O(n log n), that’s worth more than any micro-optimization.
Lesson 3: Polars Is the New Default
For any analytical workload on a single machine, Polars should be your default choice over Pandas. The performance improvement is dramatic, the API is similar, and the lazy evaluation model prevents many common performance mistakes.
Lesson 4: Don’t Optimize What Doesn’t Need Optimizing
If a pipeline runs in 5 minutes and the budget is 8 hours, don’t spend time optimizing it. Focus optimization effort on actual bottlenecks that affect production SLAs.
Lesson 5: Measure End-to-End, Not Just the Fast Part
Optimizing one component might shift the bottleneck to another. Always measure end-to-end runtime after each optimization step. I’ve seen cases where optimizing the compute phase made the I/O phase the new bottleneck, requiring a different optimization strategy.
Conclusion
Python performance optimization for data science is a systematic process, not a bag of tricks. Start with profiling to find the bottleneck. Apply the optimization hierarchy: algorithm → vectorization → library upgrade → JIT → Cython → C. Measure after each step. The Python ecosystem (NumPy, Polars, DuckDB, Numba, Cython) is powerful enough for the vast majority of production data science workloads. You rarely need to drop to C — but when you do, the investment is justified by the performance requirements.