Design notes¶
The why behind the choices in Architecture, and the questions still open. If you want to know what the code does, read Architecture. If you want to know why it does it that way, read this.
Decisions taken¶
Column-valued parameters travel in inputs, not kwargs¶
Polars plugins receive inputs: &[Series] (lazy, length-matched) and kwargs (static, JSON-serialised at planning
time). Putting parameters in kwargs would block column-valued parameters, which is the whole differentiator. So every
plugin call passes (value, param_1, ..., param_k) as inputs, and kwargs carries only static config (seed). The
one exception is the constant-parameter sampler fast path below, which passes scalars in kwargs precisely because they
are known not to be column-valued there.
One Rust file per distribution, one plugin function per method that needs Rust¶
Three options were considered:
- one function per
(distribution, method); - one function per method with an internal
DistKindenum dispatch and a row-levelmatch; - Rust only where a Polars expression cannot express the method, Python otherwise.
The result is 1 & 3. Closed-form methods sit happily in Python as pl.Expr; only methods that go through statrs
get a Rust plugin function.
Sampling derives a fresh per-row RNG from (root_seed, row_index)¶
Every sampler needs one property: a deterministic, independent stream per row that depends only on
(root_seed, row_index), never on position within a chunk. That makes output invariant to how Polars chunks or threads
the input, so sample is genuinely elementwise.
The generator is Pcg64Mcg (rand_pcg): cheap to construct (no key schedule), good statistical quality, and output
stable across releases and platforms, so seeded results stay reproducible. The root seed is resolved once per call
(SysRng when seed=None); each row then derives its own generator from (root_seed, row_index).
This replaced an earlier "single ChaCha20Rng advanced once per row in iteration order" design, which coupled rows
across chunks (order-dependent, not streaming-safe). The naive fix, a ChaCha20Rng per row, made sampling markedly
slower (a key schedule plus a keystream block per draw). A one-shot hash-to-uniform would be cheaper still but only
serves distributions needing a single uniform per draw, so it is deliberately not the foundation.
Constant parameters take a sampler fast path¶
sample ships a second plugin, <name>_sample_scalar, used when every parameter is a Python scalar. The general
sampler is built for the differentiator (column-valued parameters), but it makes the common constant-parameter case pay
for machinery it does not use: each scalar is expanded to a full-length pl.repeat column, marshalled across FFI, and
re-validated on every row, and the distribution is rebuilt per row. For a cheap draw (uniform is one multiply-add) that
fixed overhead dominates the draw itself.
The fast path passes the constant parameters in kwargs, validates and builds the distribution once, and sends only the
row index as an input. It keeps the exact (root_seed, row_index) seeding and the same draw, so its output is
byte-identical to the per-row path for any seed; that equality is the contract, pinned by a property test
(test_sample_scalar_fast_path_matches_per_row) rather than left implicit. The result is less per-row work and lower
peak memory (no constant columns), with reproducibility and the chunk- and thread-invariance guarantees untouched.
It is the one deliberate exception to "parameters travel in inputs, not kwargs": admissible precisely because the
path is selected only when the parameters are known scalars, so nothing column-valued is ever forced into kwargs.
Constant parameters validate once, not per row¶
The closed-form moments (mean, variance, std, entropy) and the closed-form methods of Uniform / Bernoulli
/ Exponential do not build a distribution; they compute a Polars expression. But they still route their validation
through a small Rust plugin (normal_sigma, uniform_range, bernoulli_proba, binomial_params, lognormal_sigma,
exponential_rate, beta_params) so an invalid
parameterisation raises the same ComputeError as the sampler and value-keyed methods rather than silently producing a
nonsense moment (see "Invalid parameters raise"). On the general path that plugin runs over the full-length pl.repeat
parameter columns, validating the same constant on every row.
For all-scalar parameters the same plugin is instead called on length-1 pl.lit inputs, so its elementwise closure runs
once. The validated quantity (or, for Binomial.entropy, the support-sum value) is returned behind a pl.when(...)
validity gate; the length-1 condition broadcasts, so the moment stays a length-n column byte-identical to the per-row
path. A length-1 collapse is deliberately not used, because it would break that path equality (column parameters
still yield length-n). This needs no new plugin and no kwargs: it reuses the existing validators, called on fewer rows.
It is the same "constant parameters take a fast path" idea as the sampler, applied to validation: nothing leaves Rust,
and the raise contract is unchanged (pinned by moment_test.py and the *_scalar validation tests). For a constant,
"per row" and "once" are the same check.
samples draws each row's array in one native call¶
sample_iter was rejected for the multi-draw loop body: it advances a single stream in row order, which couples rows
across chunks and breaks the invariance the per-row seeding exists to provide. samples instead uses one stream per
row: row i's size draws are consecutive values from the (root_seed, i) stream, the same stream sample takes its
single draw from. That stays chunk-invariant because the stream is keyed by global row position; what remains rejected
is any stream shared across rows. The per-row stream is what makes samples(size=1) equal sample bit for bit and
samples prefix-stable in size (both pinned by property tests).
It also runs as a single native plugin call that fills the whole Array(inner, size) column in one pass, replacing an
earlier construction of size separate sample calls glued by concat_arr.
Binomial sampling uses rand_distr, not statrs¶
statrs (still true on 0.19) implements Distribution<u64> for Binomial as (0..n).fold(...), one uniform draw per
trial, so a sampled row costs n RNG draws. At n = 10_000 that is 10,000 uniforms per row, turning the sampler
O(n); the constant-factor wins over scipy held only for small n. The binomial sampler therefore draws from
rand_distr::Binomial (inversion for small n*p, BTPE otherwise, both O(1)-amortised), keeping sampling time flat
in n. This is the one place sampling
does not go through statrs; every value-keyed method (pmf, cdf, ppf, ...) still builds the statrs distribution.
Invalid parameters raise, they never silently null¶
An invalid parameter value, scalar or one bad column row (sigma <= 0, max <= min, p outside [0, 1], a
non-finite bound), maps the statrs constructor error through a ComputeError and fails the whole evaluation.
This reverses an earlier "produce null, keep the pipeline running" decision. Silently nulling hides a modelling error: a
user who does not check for nulls gets wrong answers downstream, and an invalid-parameter null is indistinguishable from
a legitimately-null input. Raising is loud, uniform across distributions, and uniform across scalar vs column inputs,
because scalars are coerced to columns and validated per row exactly like columns. Construction rejects only wrong
types. A closed-form distribution cannot raise from a bare pl.Expr, so it routes parameters through one small
validating plugin (see Architecture / Plugin granularity).
Moments that are undefined¶
Every distribution shipped today has finite moments on its valid parameter range, so this policy does not bite yet; it governs distributions on the roadmap. Two cases, handled differently on purpose:
- Permanently undefined (e.g. a Cauchy mean): raise
NotImplementedError. Silently returning null would hide a modelling error from a user who chains.mean().sum(). - Undefined only in part of the parameter range (e.g. a Student-t mean with
df <= 1, defined fordf > 1): returnnull. A user sweeping a parameter across the threshold should not get an exception that breaks the sweep.
Inconsistent on its face, defensible per case. Each class will document which case applies as it lands.