abstractqqq/polars_ds_extension

Polars extension for general data science use cases

PDS (polars_ds) — A Data Science Toolkit for Polars

What it is: PDS is a Python package that adds data-science tools on top of Polars, a fast dataframe library. It's designed for people who already work with Polars and want statistics, machine-learning models, and feature engineering without switching to other tools like Pandas or scikit-learn.

Install: pip install polars-ds

Main Capabilities

1. Machine-Learning Metrics (Parallel)

You can compute metrics like ROC AUC and log loss on different slices of your data at once, using Polars' lazy evaluation. Example from the README:

df.lazy().group_by("segments").agg(
    pds.query_roc_auc("actual", "predicted").alias("roc_auc"),
    pds.query_log_loss("actual", "predicted").alias("log_loss"),
).collect()

This returns a table with one row per segment, each with its own ROC AUC and log-loss value.

2. Quick Modeling (Non-Persistent)

You can run simple models directly inside a dataframe expression — no separate training step. Available models:

  • Linear regression (with optional L1/L2 regularization → Lasso, Ridge, Elastic Net)
  • Non-negative linear regression
  • Rolling and recursive linear regression
  • Binary logistic regression (with L1/L2)
  • A statsmodels-style regression summary table (coefficients, standard errors, t-stats, p-values, R²)

Example: pds.lin_reg(pl.col("x1"), pl.col("x2"), target=pl.col("y"), add_bias=False, return_pred=True) returns predictions as a struct column.

3. Distance Functions

  • String distances: Levenshtein, Optimal String Alignment, Jaro-Winkler (powered by RapidFuzz, so very fast)
  • Array/list distances: squared L2 (Euclidean) distance between two columns of arrays

4. Statistical Tests

Built-in tests you can run inside a group_by aggregation:

  • Independent t-test (with optional unequal variance)
  • Chi-squared test
  • F-test (ANOVA-style, comparing group means)

Each returns a struct with the statistic and p-value.

5. Convenience Helpers

pds.E(['a', 'b'], ["min", "max", "n_unique", "len"]) generates multiple aggregation expressions at once, so you don't have to write them out manually.

6. Data Transformation Pipeline

A Blueprint/Pipeline system for building reusable preprocessing steps:

  • Filter rows
  • Impute missing values (linear or median)
  • Add custom expressions (log transforms, clipping, etc.)
  • Scale numeric columns (standard scaling)
  • One-hot encode, weight-of-evidence encode, target encode
  • Materialize into a Pipeline object that can transform new data

It supports batched processing via collect_batches() for streaming machine-learning workflows.

7. Nearest-Neighbor Queries

Find all points within a given radius of each point (using squared L2 distance), returning the neighbor IDs. The README notes this is not fast for large k or high dimensions.

8. Compatibility with Other Dataframes

Via the polars_ds.compat module, PDS functions can work with other eager dataframes (like Pandas) as long as columns implement to_numpy(). Example: pds2.lin_reg(df_pd["x1"], df_pd["x2"], target=df_pd["y"], return_pred=True).

Other Notes

  • Performance: The README claims it's generally faster than Pandas + SciPy/Sklearn/NumPy, especially when evaluating many expressions at once and on multi-core machines.
  • Dependencies: Requires Polars ≥ 1.4.0. Python 3.9 uses Polars 1.x; Polars 2 requires Python ≥ 3.10.
  • Build from source: Requires Rust and maturin. You can set RUSTFLAGS for LTS CPU support or native CPU.
  • Credits: Uses Statrs (statistics), faer (linear algebra), and RapidFuzz (string similarity).
  • Maintenance: Mostly maintained by one person; AI-assisted code is reviewed by the maintainer, and AI-generated PRs are scrutinized.

Documentation: polars-ds-extension.readthedocs.io

Related

  • Project
  • Project
  • Project
  • Project