Vectorized statistical distributions, built on the Polars engine

If you do statistical work in Polars today, you leave Polars to do it: .to_numpy(), a trip through scipy.stats, and back. That cuts the lazy query in half and materializes everything. The other way out, a map_elements UDF, is slower still and holds the GIL.

scipy broadcasts parameter arrays fine, so a different distribution per row, where the mean and standard deviation are themselves columns, is already vectorized there. It hands the result back as a NumPy array, and the lazy plan ends there.

This talk walks through a Polars expression plugin that exposes scipy.stats-style distributions natively inside Polars expressions, with column-valued parameters as a first-class feature. The math is Rust (statrs); the surface is Polars. We will cover the plugin architecture (pyo3-polars), why column-valued parameters reshape the API, how seeded sampling stays reproducible across platforms, chunk layouts and both engines, and the null and error contract.

The thread running through it is what we got wrong before we got it right: the abstraction we built and then changed, the parameterization we flipped, the contract we had to rewrite. Expect concrete code, real benchmarks against scipy + polars.


The problem: Polars has become a widely adopted dataframe engine for data engineering and ML feature pipelines, but it has no native statistical distribution layer. Users fall back to three bad options:

  1. round-trip through scipy.stats (breaks laziness, materializes columns, costs memory),
  2. Python UDFs via map_elements (slow, GIL-bound, not vectorized),
  3. hand-rolled per-distribution expressions (duplicated and error-prone)

The first is what people actually do, and scipy broadcasts parameter arrays, so the per-row case (the probability of a reading under a Normal whose mean and standard deviation are stored as columns) is already vectorized there. The result comes back as a NumPy array. The collect() cuts the query in half, pushdown stops at that boundary, keeping the result row-aligned through joins, filters and over is the caller's job, and NumPy has no null, so an invalid parameter returns NaN instead of raising.

What we built is a Polars expression plugin that puts scipy.stats-style distributions directly into Polars expressions. Two properties set it apart from the sampling-focused plugins that already exist:

  • Column-valued parameters across the whole surface: any parameter can be a scalar or a Polars expression, for pdf / cdf / sf / ppf, their log variants, the moments and sampling alike. Normal(mu=pl.col("mu"), sigma=pl.col("sigma")).cdf(pl.col("x")) evaluates one distribution per row, fully vectorized.
  • Lazy-native: every method returns a pl.Expr, so nothing materializes and both the in-memory and streaming engines are supported.

What the audience will take away: this is a practical, cross-ecosystem talk: a Python-facing library whose performance comes from Rust. Concretely, attendees will see:

  1. Plugin anatomy: how a pyo3-polars expression plugin is wired across three layers (Python API, Rust FFI, the statrs math layer), why parameters must travel as length-matched inputs rather than static kwargs to stay column-valued, and the one fast path that is allowed to break that rule.
  2. API design under a real constraint. How the column-valued requirement forces choices that a scipy clone never faces: parameter coercion that keeps expressions elementwise under group_by and over, row alignment that has to happen in Rust because Polars broadcasts nothing into a plugin, and why parameters carry each distribution's own convention (mu / sigma, rate, min / max) instead of scipy's loc / scale.
  3. Correctness and reproducibility: a null and error contract where a null input gives a null output but an invalid parameter raises instead of quietly returning NaN, and sampling keyed on (seed, row index), so a seeded column repeats across runs, chunk layouts, thread counts, operating systems and both engines.
  4. Benchmarks and honest trade-offs: real numbers against scipy + polars where a like-for-like comparison exists, the costs we accept (an FFI round trip per call, a catalogue far smaller than scipy's hundred-plus), plus the lessons that only show up in hindsight.

What we got wrong first: the most useful part of the talk is the design history, not the final state. We will walk through the macro scaffolding we wrote to generate the per-distribution plugin shells and then replaced with traits and generic drivers, the parameter names we shipped as statrs's (mean, std_dev) and renamed to (mu, sigma), the seeding design that advanced one ChaCha20 stream in row order until we noticed it made results depend on how Polars chunked the frame, and the null-versus-raise contract we shipped as "null the bad row, keep the pipeline running" and then reversed, because a silently nulled parameter error is indistinguishable from a null input. Each is a small, transferable lesson about building numerical libraries on top of someone else's math and someone else's engine.

Francesco Bruzzesi

Data science tech lead at intella.tech · Mathematician at heart · Open source enthusiast