Polars Expression Plugins: Rust PII Masking from Scratch

NovaOwl Intermediate 9h ago 139 views 12 likes 2 min read

Stop using map_elements for heavy string manipulation if you have millions of rows. I recently hit a wall with payment note columns containing PII—things like credit card numbers and RUT IDs—where the overhead of calling a Python function for every single row was killing my performance. When you use map_elements, Polars effectively pauses its vectorized engine and hands every single value back to the Python interpreter, creating a massive bottleneck of str objects and regex allocations.

The fix isn't just "finding a faster regex"; it's moving the logic into a native Rust plugin. By writing the per-value work in Rust and registering it as an expression, Python is called exactly once for the entire column, not once per row.

Building the Bridge: Rust to Python

A Polars plugin consists of a compiled dynamic library and a small Python wrapper. On the Rust side, you use the pyo3-polars macro to annotate your function.

use pyo3_polars::derive::polars_expr;
use pyo3_polars::export::polars_core::prelude::*;

#[polars_expr(output_type=String)]
fn mask_digits(inputs: &[Series]) -> PolarsResult<Series> {
 let ca = inputs[0].str()?;
 let out: StringChunked =
 ca.apply(|opt| opt.map(|s| std::borrow::Cow::Owned(mask_cell(s))));
 Ok(out.into_series())
}

Then, you register this in Python so Polars knows where to find the compiled binary:

from pathlib import Path
import polars as pl
from polars.plugins import register_plugin_function

PLUGIN = Path(__file__).parent

def mask_digits(expr):
 col = pl.col(expr) if isinstance(expr, str) else expr
 return register_plugin_function(
 plugin_path=PLUGIN,
 function_name="mask_digits",
 args=[col],
 is_elementwise=True,
 )

Once registered, it integrates seamlessly into your AI workflow or data pipeline: df.with_columns(mask_digits("notes")).

Configuration and Deployment

To get this to compile, your Cargo.toml needs to be specific. The cdylib crate type is mandatory because it allows CPython to dlopen the library.

[lib]
crate-type = ["cdylib"]

[dependencies]
pyo3 = { version = "0.25", features = ["extension-module"] }
pyo3-polars = { version = "0.23", features = ["derive"] }
polars-core = "0.50"
regex = "1"
once_cell = "1"

One major "gotcha" for anyone attempting this deployment: pyo3 and pyo3-polars versioning is extremely strict. If you bump one without the other, you'll get cryptic link errors that look like compiler bugs. Treat them as a coupled pair.

Using maturin develop is the easiest way to compile and install the plugin directly into your environment. The result is a massive productivity gain—you get the expressive API of Polars with the raw execution speed of Rust.

AI ProgrammingAI Codingpythonrustpolars

All Replies (3)

C
ChrisCat Intermediate 9h ago
did u run into any memory issues with huge strings or did rust handle it fine?
0 Reply
P
PatFounder Advanced 9h ago
Switched to plugins for my logs last month; the speed difference is actually night and day.
0 Reply
T
Taylor27 Intermediate 9h ago
Does this actually beat a vectorized regex in Python, or is the speedup negligible?
0 Reply

Write a Reply

Markdown supported