Add linear regression aggregation functions (regr_slope, regr_intercept, regr_r2, regr_count)
Description
I'd like to propose adding the core SQL-standard linear regression aggregation functions to Polars:
pl.regr_slope(y, x)— slope of the least-squares linear fit ofyonxpl.regr_intercept(y, x)— intercept of the fitpl.regr_r2(y, x)— coefficient of determinationpl.regr_count(y, x)— number of rows where bothyandxare non-null
These are standard aggregates in PostgreSQL, Spark SQL, DuckDB and DataFusion, with SQL-standard argument order (dependent variable y first) and semantics: null pairs are excluded pairwise; zero variance in x yields null for slope/intercept/r2; zero variance in y alone yields r2 = 1.
df.group_by("g").agg(
pl.regr_slope("y", "x"),
pl.regr_intercept("y", "x"),
pl.regr_r2("y", "x"),
)
# and via the SQL interface:
df.sql("SELECT g, REGR_SLOPE(y, x) FROM self GROUP BY g")Why it fits well into Polars:
- Polars already has
pl.corr/pl.cov, but nothing regression-shaped; today users have to composecov(x, y) / var(x)manually (and the intercept/r2 variants are even more verbose), which is error-prone around null handling. - The implementation can reuse the numerically stable
PearsonStatemoment state that already exists inpolars-compute(weight,mean_x,mean_y,dp_xx,dp_xy,dp_yy) — that state is exactly the sufficient statistic for all of these aggregates, so no new kernels are needed and the streaming engine gets full support (including streaming group-by) via aGroupedReductionanalogous to the existing Pearson one. - The SQL interface gets
REGR_SLOPE/REGR_INTERCEPT/REGR_R2/REGR_COUNTfor free, improving PostgreSQL compatibility.
The full family in PostgreSQL also includes regr_avgx, regr_avgy, regr_sxx, regr_syy, regr_sxy; I'd suggest starting with the four above and following up with the rest if there is interest, since they share the same machinery.
I have a working implementation (both engines, SQL, tests validated against DuckDB) and would be happy to open a PR.
Source: pola-rs/polars