Skip to content

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.

Aryanto, M.Si

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:

  1. Algorithm optimization: Use a better algorithm. O(n log n) beats O(n²) regardless of language.
  2. Vectorization: Replace loops with NumPy/Pandas vectorized operations.
  3. Library upgrade: Switch from Pandas to Polars or DuckDB for analytical workloads.
  4. JIT compilation: Use Numba for numerical Python code.
  5. Cython: Write C extensions for hot paths that can’t be vectorized.
  6. 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:

  1. Single-threaded: Pandas operations run on a single CPU core.
  2. Eager evaluation: Every operation is executed immediately, even if it could be optimized by deferring execution.
  3. 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):

OperationPandasPolarsDuckDB
Group-by aggregation12.3s1.8s1.5s
Filter + aggregate8.7s1.2s1.0s
Join (100M × 1M)25.1s3.4s2.8s
Window functions18.5s2.1s1.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

CriterionNumbaCython
Ease of use★★★★★★★★
First-call latency1-2s JITNone (pre-compiled)
String supportLimitedFull
C interopLimitedExcellent
DebuggingDifficultStandard C tools
DeploymentSimple (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:

  1. The component is a bottleneck: It accounts for > 30% of total runtime.
  2. Python optimization is exhausted: Vectorization, Numba, and Cython have been tried and are insufficient.
  3. The component is stable: The logic won’t change frequently. Rewriting in C every time the requirements change is expensive.
  4. 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

  1. ctypes/cffi: Call C functions from Python directly. Simplest approach, minimal overhead.
  2. pybind11: Create Python bindings for C++ code. More convenient than ctypes for complex APIs.
  3. Rust + PyO3: Write Rust code with Python bindings. Growing ecosystem, excellent performance.
  4. 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.