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 intopl.Expr, register plugin calls. Implement methods directly aspl.Exprwhen the closed form is trivial. -
FFI layer (
pyo3-polarsplugin 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.