Skip to main content

High-performance machine learning metrics for Polars DataFrames

Project description

Polarbearings ๐Ÿปโ€โ„๏ธ๐Ÿงญ

High-performance machine learning metrics implemented as native Polars expressions.

Features

  • Fast where it counts: Native Polars expressions โ€” large wins on grouped and probabilistic metrics (see Performance)
  • Weighted: Optional sample weights on nearly every metric
  • Flexible labels: Any positive class โ€” 1, 100, "cancer", True โ€” via pos_label
  • Correct: scikit-learn-faithful, verified with property-based testing
  • Composable: Plain Polars expressions โ€” drop into group_by, over, and lazy pipelines
  • Type-safe: Full type hints

Installation

uv add polarbearings
# or: pip install polarbearings

The only runtime dependency is Polars.

Quick Start

import polars as pl
from polarbearings import roc_auc

df = pl.DataFrame({
    "label": [0, 0, 1, 1, 1],
    "score": [0.1, 0.4, 0.35, 0.8, 0.9]
})

# Every metric is a Polars expression โ€” use it anywhere an expression is allowed.
df.select(roc_auc("label", "score"))
# shape: (1, 1)
# โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”
# โ”‚ roc_auc_label_score โ”‚
# โ”‚ f64                 โ”‚
# โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก
# โ”‚ 0.833333            โ”‚
# โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜

Why Polarbearings

Every metric is a pure Polars expression. That single design choice is where the strengths come from โ€” things a scikit-learn wrapper or a compiled plugin can't easily match:

  • Group-wise metrics at scale โ€” metrics drop straight into group_by().agg() and Polars parallelizes across groups (15โ€“65x faster than a Python loop calling scikit-learn per segment).
  • A whole metric suite in one pass โ€” bundle every metric into a single df.select([...]); Polars shares the column scans and parallelizes across the independent outputs. 13 metrics on 10M rows run in ~1.05 s vs scikit-learn's ~6.1 s (5.8x) โ€” see Performance.
  • Sample weights almost everywhere โ€” nearly every metric accepts an optional weight column, including ROC AUC, log loss, MCC, and Cohen's kappa (a few exceptions are noted per metric).
  • Any positive class โ€” pos_label accepts 1, 100, "cancer", True, โ€ฆ no remapping your labels to 0/1.
  • scikit-learn-faithful โ€” names and edge-case semantics mirror scikit-learn (e.g. null for undefined cases), so migration is ~1:1.
  • Composable & lazy-friendly โ€” pure expressions fold into lazy query plans, over() windows, and the rest of your Polars pipeline.
  • Zero build, one dependency โ€” pure Python emitting Polars expressions; no compiled extension, installs anywhere Polars runs, supports Polars 1.0.0+.

A whole evaluation report in one select

Because every metric is just an expression, a full report is one df.select(...) โ€” Polars reads each column once and fans the work across the independent outputs:

from polarbearings import (
    precision, recall, f1_score, roc_auc, average_precision,
    log_loss, brier_score, confusion_matrix,
)

df.select(
    precision("label", "prob"),
    recall("label", "prob"),
    f1_score("label", "prob"),
    roc_auc("label", "prob"),
    average_precision("label", "prob"),
    log_loss("label", "prob"),
    brier_score("label", "prob"),
    confusion_matrix("label", "prob"),   # struct {threshold, tp, fp, fn, tn}
)
# One tidy row, one column per metric (8 columns; abbreviated):
# shape: (1, 8)
# โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”
# โ”‚ precision_label_prob_0.5 โ”† recall_label_prob_0.5 โ”†  โ€ฆ  โ”† brier_score_label_prob โ”‚
# โ”‚ f64                      โ”† f64                   โ”†     โ”† f64                    โ”‚
# โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก
# โ”‚ 0.714286                 โ”† 0.833333              โ”†  โ€ฆ  โ”† 0.161944               โ”‚
# โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜

benchmarks/bench_multi_metric.py reproduces this timing across sizes from 100 to 1M rows.

Metrics

~40 metrics across ranking, probabilistic, classification, calibration, and regression families, plus curve generators, bootstrap CIs, deterministic splitting, and threshold utilities. The full metrics reference โ†’ docs/guides/METRICS.md documents every metric's semantics, edge cases, and scikit-learn correspondence.

Family What's in it
Ranking roc_auc, average_precision, gini_coefficient, dcg_score, ndcg_score
Probabilistic log_loss, brier_score
Classification precision, recall, f1_score, fbeta_score, specificity, accuracy, balanced_accuracy, matthews_corrcoef, cohens_kappa, jaccard_score, confusion_matrix
Thresholds threshold_sweep, quantiles / equal_width / linspace specs
Curves roc_curve, pr_curve, det_curve, expected_cost, confusion_curve
Regression mae, mse, rmse, r2_score, mape, smape, MSLE/RMSLE, huber_loss, log_cosh_loss, pinball, Tweedie/Poisson/gamma deviance, Dยฒ scores, max_error, median_absolute_error
Calibration calibration_curve, expected_calibration_error (ECE), maximum_calibration_error (MCE)
Weights & CIs balanced_sample_weight, balanced_class_weights, bootstrap_ci, bootstrap_weight
Splitting hash_split, hash_splits, hash_fold, hash_uniform
from polarbearings import roc_auc, f1_score, mae, calibration_curve

# Classification + probabilistic metrics are plain expressions:
df.select(roc_auc("label", "prob"), f1_score("label", "prob", threshold=0.7))

# Regression metrics too:
reg.select(mae("y", "pred", weight="w"))

# Curve helpers take a (Lazy)Frame and return a plot-ready LazyFrame:
calibration_curve(df, "label", "prob", n_bins=10, strategy="quantile").collect()

Cross-cutting behaviour

Four behaviours are shared by the metrics rather than specific to one โ€” full details in the metrics reference:

  • Output names โ€” each expression is pre-aliased <metric>_<target>_<score> (plus suffixes for a weight, non-default pos_label, or threshold), so roc_auc("label", "score") yields a roc_auc_label_score column. Chain .alias("auc") to rename it.
  • Sample weights โ€” pass weight="col" to nearly any metric. Six omit it (dcg_score, ndcg_score, max_error, median_absolute_error, d2_absolute_error_score, d2_pinball_score) where a weighted form is undefined or not cleanly expressible.
  • Custom positive class โ€” pos_label accepts ints, strings, or booleans; no remapping to 0/1.
  • Missing values โ€” any null/NaN in target, score, or weight makes a metric return null (loud, not silent), scoped to the evaluation context. Drop rows yourself for complete-case behaviour. Curve helpers are the exception โ€” they drop incomplete rows.

Diagnostic plots

notebooks/diagnostics.ipynb gives examples of ROC/PR/DET/calibration curves and bootstrap confidence bands with these helpers and visualizing them with Plotly.

Use Cases

Polarbearings is a good fit for:

  1. Large-scale model evaluation โ€” millions of predictions, efficiently
  2. Group-wise metrics โ€” per-segment analysis via Polars' group_by
  3. Streaming / out-of-core โ€” lazy expressions over LazyFrames
  4. Composed evaluation reports โ€” a whole metric suite in one pass
from polarbearings import roc_auc

# ROC AUC per customer segment, in one parallelized pass:
df.group_by("segment").agg(roc_auc("label", "score"))
# shape: (2, 2)
# โ”Œโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ฌโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”
# โ”‚ segment โ”‚ roc_auc_label_score โ”‚
# โ•žโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ชโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•ก
# โ”‚ A       โ”‚ 1.0                 โ”‚
# โ”‚ B       โ”‚ 0.5                 โ”‚
# โ””โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”ดโ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”˜

Performance

Polarbearings runs every metric as a native Polars expression. The advantage is large and real where there's work to parallelize, and honest where there isn't. Numbers below are speedup vs scikit-learn (scikit-learn time รท polarbearings time; higher = faster), median of clean benchmark runs.

Where polarbearings wins big โ€” grouped, probabilistic, and ranking metrics. Ranges span polars 1.0.0 and 1.41.2; see PERFORMANCE.md for the per-version, per-size matrix.

Metric 100k rows 10M rows
Grouped metrics (per segment) ~15โ€“60x โ€”
Precision / F1 ~5โ€“22x ~14โ€“23x
Brier Score ~9โ€“10x ~7โ€“8x
Log Loss ~5โ€“6x ~4โ€“5x
ROC AUC ~4.5โ€“6x ~5x
13 metrics in one select โ€” 5.8x

Where it's at parity or slower โ€” trivial reductions (MAE, MSE, MAPE, Rยฒ) are roughly even with scikit-learn at small-to-mid sizes and slower on a single very large array, where NumPy's tight single-threaded loop beats one Polars expression. Reach for polarbearings on grouped/composed pipelines and the probabilistic metrics; for a one-off MAE over a giant array, NumPy is fine.

See docs/technical/PERFORMANCE.md for the full per-metric breakdown, the size-scaling curve, and the ceteris-paribus Polars version comparison.

Development

This project uses uv for dependency management and just for task running. Run just to list all recipes.

uv sync --all-groups   # install dependencies
just test              # run tests        (or: uv run pytest)
just quality           # lint + type-check
just check             # fast local check: lint + type-check + tests (not full CI)
just test-compat       # test against min / mid / latest Polars
just bench             # benchmarks vs scikit-learn

Testing combines unit tests, property-based tests (Hypothesis), and scikit-learn compatibility tests across multiple Polars versions โ€” see docs/guides/TESTING.md.

Requirements

  • Python: 3.11+
  • Polars: 1.0.0+

Roadmap

  • Publish to PyPI (first tagged release)
  • KS statistic & lift/gain curves (churn / credit workflows)
  • Ranking metrics: precision@k, recall@k, MRR, MAP
  • Weighted (linear / quadratic) Cohen's kappa for ordinal targets
  • .metrics Polars expression namespace (pl.col(...).metrics.roc_auc(...))

More candidate metrics and their implementation notes live in docs/FUTURE_IDEAS.md.

Contributing

Contributions are welcome! Please:

  1. Run just check to verify your changes locally (lint + type-check + tests)
  2. Submit a Pull Request with a clear description

License

MIT License - see LICENSE file for details.

Acknowledgments

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

polarbearings-0.2.2.tar.gz (111.0 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

polarbearings-0.2.2-py3-none-any.whl (65.0 kB view details)

Uploaded Python 3

File details

Details for the file polarbearings-0.2.2.tar.gz.

File metadata

  • Download URL: polarbearings-0.2.2.tar.gz
  • Upload date:
  • Size: 111.0 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: uv/0.11.29 {"installer":{"name":"uv","version":"0.11.29","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}

File hashes

Hashes for polarbearings-0.2.2.tar.gz
Algorithm Hash digest
SHA256 bd2f1360730c527d76c9cb62a0b153c69440d05b92baca47f5517ca645fde905
MD5 7ad15810c17f2624a0f2c758deedc192
BLAKE2b-256 5d61711af8060d90b3cf96044d49025f5dc58d4ceda8f3993154f9a606e07d5a

See more details on using hashes here.

File details

Details for the file polarbearings-0.2.2-py3-none-any.whl.

File metadata

  • Download URL: polarbearings-0.2.2-py3-none-any.whl
  • Upload date:
  • Size: 65.0 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: uv/0.11.29 {"installer":{"name":"uv","version":"0.11.29","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}

File hashes

Hashes for polarbearings-0.2.2-py3-none-any.whl
Algorithm Hash digest
SHA256 191b07508f7c9988fc8de3d1964b6e8e80acf445793d5837859a09a621ed53eb
MD5 d95145608ae622cdad78f0565e9a2033
BLAKE2b-256 58f64d59a20911822890b55d0c9573c744c492116568faad9c835e02614bd2d9

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page