Skip to content

Architecture

The canonical "what polars-stats is and how it is wired". For the rationale behind these choices and the open questions, see Design notes.

Three layers

  • Python API layer (polars_stats/): Distribution classes (Normal, Bernoulli, ...). Coerce params into pl.Expr, register plugin calls. Implement methods directly as pl.Expr when the closed form is trivial.

  • FFI layer (pyo3-polars plugin functions). One #[polars_expr] per (distribution, method) where Rust is required. Per-row null propagation; an invalid parameter on a row raises.

  • Math layer (statrs): Trusted upstream for pdf, cdf, ppf, moments, sampling

Public method surface

The method surface (pdf/pmf, cdf, sf, ppf, isf, the log_* family, mean, variance, std, median, entropy, sample, samples) is defined on the abstract base classes ContinuousDistribution and DiscreteDistribution. The catalogue and the full table live in the API reference.

Template-method split: every value-keyed method is concrete in the base: it coerces the argument with as_expr and delegates everything else to a private hook (_pdf, _cdf, _ppf, ...). The hook receives an already-coerced pl.Expr and owns the whole per-row contract, the null and NaN rows included (null in, null out; NaN in, NaN out, matching scipy). Nothing sits between the hook and the caller; Design notes explain why a when / then / otherwise above it could not be trusted with those rows.

Every hook dispatches to Rust by name. A subclass declares _distribution_name ("normal", "beta", ...), and the base hook for each value-keyed method calls the plugin <name>_<method> (normal_cdf, beta_ln_pdf, ...) through _value_plugin, or its <name>_<method>_scalar twin when every parameter is constant. The valid names are Literal aliases in polars_stats/_typing.py, so a typo is a type error. A subclass overrides a hook only where the Rust body does not exist yet: Beta and Binomial compose log_cdf / log_sf as log(cdf) / log(sf), which underflow in the tails (see Accuracy). Override the private _x hook, never the public method: the public methods only coerce, and median (ppf(0.5)) and std (variance().sqrt()) are composed from the hooks and public moments, so a public override is bypassed by everything built on it.

_isf is solved against q itself, in Rust: _ppf(1 - q) quantises a small quantile before the inverse runs, and polars would form 1 - q ahead of the Rust dtype gate.

The moments follow the same split at the parameter level. A closed-form moment is a polars expression over the raw parameters, gated by _moment on the distribution's _validated_params: one validating plugin call (normal_sigma, uniform_range, ...) that raises on an invalid parameterisation and is null on a null one. A moment with no closed form (Beta.entropy, Binomial.entropy) calls its own parameter-keyed plugin through _param_plugin, which validates inside. Both run once, on length-1 literals, when the parameters are constant.

Column-valued parameters

Every distribution __init__ coerces each parameter with a single shared helper, coerce_param (coerce_n for count parameters like Binomial's n). The accepted inputs and their coercions are tabulated in Reference / Parameters and contracts.

A scalar becomes pl.lit(value), a length-1 scalar column, and Rust owns row-alignment: polars broadcasts nothing into a plugin, so align_inputs broadcasts every length-1 input up to the call's row count before any cast. It aligns by length, not by expression kind, so a user-written .first() or .max() parameter is handled like a literal. The value-keyed drivers take one shortcut ahead of that. When every parameter is length 1 they validate and build it once and never expand it, so a pl.lit or an aggregate parameter costs what a Python scalar does.

An expression whose inputs are all constant is therefore a scalar column, so df.select(Normal(0.0, 1.0).mean()) is one row; any column-valued input sets the length instead. See Reference / Parameters and contracts.

On the Rust side, a plugin function receives inputs: &[Series] of (value, param_1, ..., param_k), passes each through the numeric dtype gate (coerce_f64: a numeric or Null-typed column casts to Float64Chunked, anything else raises), iterates per chunk, propagates nulls, and constructs the statrs distribution per row. kwargs carries only static config that cannot be column-valued: the sampler seed, plus the constant parameters in the sampler fast path below (the one place a parameter rides in kwargs, valid precisely because there it is known to be a scalar).

Plugin granularity

One Rust file per distribution, one #[polars_expr] per method that genuinely needs Rust: the closed-form moments (Bernoulli's mean, variance, entropy; Exponential's, Uniform's, ...) live in Python as pl.Expr and compose with the expression engine. Methods that go through statrs (sampling, transcendental pdf/pmf, cdf, inverse_cdf, native ln_pdf/ln_pmf, native sf) get a Rust plugin function, and so do the elementary value-keyed closed forms: they branch on the evaluation value, which puts the validator inside a when arm where polars can mask it away (see Design notes).

Parameter validation needs Rust. A bare pl.Expr cannot raise per row, so any method computed in Python routes its parameters through one small validating plugin that raises on a bad row and returns a reused quantity. That covers the closed-form moments, since the value-keyed methods already validate inside their own plugin: normal_sigma validates sigma > 0, so Normal.mean() raises exactly as Normal.pdf() does, and uniform_range returns max - min and raises on max <= min for Uniform's five moments, as bernoulli_proba, exponential_rate and geometric_p do for Bernoulli's, Exponential's and Geometric's. Either way a pure-math mean makes one FFI round-trip and reports an invalid parameterisation consistently, trading a little throughput for a uniform error surface.

Sampling

sample(seed) returns one draw per row from the row-specific distribution; samples(size, seed) returns an Array(inner=..., shape=size), drawn as one native multi-draw plugin call per frame: row i's size draws are consecutive values from the one per-row stream keyed (seed, i), so samples(size=1) matches sample bit for bit and growing size extends each row's array without changing existing draws.

Per-row seeding: the root seed is resolved once per plugin call (SysRng when seed=None), then each row derives its own Pcg64Mcg generator from (root_seed, row_index) via two splitmix64 mixing draws. The row index arrives as a regular input column, so it tracks the partition under over / group_by.

The index is an identity, not a parameter, which is why it is the one input that must never be broadcast. A parameter can be longer than the frame, so the per-row samplers size the index by the call's row count (row_index_expr(params)) rather than by pl.len(). A pl.len()-sized index would be length 1 in that case, broadcast to position 0 on every row, and return one draw repeated at full height with no error. The plugin cannot repair it either, because the streaming engine splits such a call into one-row morsels and the flattened index is all it ever sees.

Identical (root_seed, row_index) always yields an identical stream, which is exactly what makes sample elementwise and invariant to chunking and thread count. Pcg64Mcg is cheap to construct (a handful of integer ops, no key schedule), passes TestU01 BigCrush, and is stable across rand_pcg releases and platforms, so seeded results are reproducible across OS and architecture.

Constant-parameter fast path: when every distribution parameter is a Python scalar (the common case), the sampler takes a dedicated plugin, <name>_sample_scalar. The parameters travel in kwargs and are validated once; only the row index crosses FFI, instead of one broadcast column per parameter that the general plugin would marshal and re-validate on every row. The shared sample_by_index helper in rng.rs resolves the seed once and maps the dense, non-null index straight into the typed output. Each distribution writes the two fast-path shells by hand over sample_by_index / samples_by_index, carrying its parameters in SampleScalarKwargs<P> / SamplesScalarKwargs<P>, the generic seed / size wrappers in rng.rs around the distribution's own <Name>Params. The shells reuse the same (root_seed, row_index) seeding and call the same named draw as the per-row path, so output is byte-identical for the same seed (a property test pins that equality); column-valued parameters still take the general per-row plugin, itself a shell over the sample_per_row_* / samples_per_row_* drivers in the same file.

Earlier ChaCha20 design (removed)

A previous design advanced a single ChaCha20Rng once per row in iteration order, which coupled rows across chunks and was not streaming-safe. The naive fix, constructing a ChaCha20Rng per row, made sampling markedly slower (a key schedule plus a keystream block per draw). Per-row Pcg64Mcg is both correct and cheap and replaced it; statrs's sampling traits consume it directly.

Null and error contract

The full table is in Reference / Parameters and contracts. An invalid parameter value raises a ComputeError and fails the evaluation; null is reserved for null inputs. Construction rejects only wrong types. There is no early Python validation of parameter values, so a bad scalar and a bad column row surface identically.

Stack

The math runs on statrs 0.19; the plugin glue is pyo3-polars over pyo3 (abi3); per-row seeded RNG is rand_pcg (Pcg64Mcg), with rand for SysRng; serde deserialises the static seed kwarg. The full dependency rationale and the deliberately-excluded crates are in Contributing / Stack, and the repository layout is in Contributing / Repository layout. Supported Python, Polars, and OS versions are in Reference / Compatibility.