ps.with_columns is a drop-in replacement for with_columns from polars that can handle some additional use cases like functions that need to peek at the full data for evaluation. It works efficiently on both DataFrame and LazyFrame.
Details
The key idea is FrameExpr — an expression that needs a peek at the data (schema or a small aggregation) before it resolves into a regular Polars expression. This unlocks operations like deriving Enum categories from the data, lumping rare levels, or reordering factor levels by a summary statistic, while keeping the rest of your pipeline lazy.
How FrameExpr stays efficient
ps.with_columns resolves each FrameExpr in two phases. First it runs a small aggregation (e.g. unique().sort() to discover categories) against the current lazy plan — so any preceding .filter() or .select() is already embedded and Polars’ predicate/projection pushdown keeps the peek cheap. Then it uses the result to build a concrete pl.Expr (e.g. .cast(pl.Enum(["a", "b", "c"]))) that goes back into the lazy plan and executes normally.
# Only the filtered rows are scanned for category discovery;# the cast itself remains lazy.lf = pl.scan_parquet("events.parquet")result = ( lf.filter(pl.col("country") =="DE") .ps.with_columns(pl.col("status").ps_enum.make()) .filter(pl.col("status") =="active") .collect())
See the FrameExpr docstring for the full explanation, including when the peek is larger and notes on parallel evaluation.
Calling arbitrary functions
Sometimes there’s no Polars expression for what you need. ps.F, ps.B, and ps.E wrap arbitrary functions (numpy, scipy, plain Python, …) so they can be called directly on expressions, in place of hand-rolled pl.struct(...).map_batches(...) /map_elements(...).
ps.F is the right default whenever a function needs to see the complete input, not a sample or a batch — clustering is the clearest example. scipy’s fclusterdata takes every point at once and assigns cluster labels; there’s no Polars equivalent, and critically, it cannot be computed correctly on a slice of the data. Only the pl.Expr argument (pl.concat_arr("x", "y")) is resolved against the data — t and criterion are forwarded to fclusterdata unchanged:
shape: (8, 4)
┌────────┬───────┬──────┬─────────┐
│ region ┆ x ┆ y ┆ cluster │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ f64 ┆ f64 ┆ i32 │
╞════════╪═══════╪══════╪═════════╡
│ A ┆ 0.0 ┆ 0.0 ┆ 1 │
│ A ┆ 0.0 ┆ 0.2 ┆ 1 │
│ A ┆ 0.0 ┆ 9.8 ┆ 1 │
│ A ┆ 0.0 ┆ 10.0 ┆ 1 │
│ B ┆ 100.0 ┆ 0.0 ┆ 2 │
│ B ┆ 100.0 ┆ 0.2 ┆ 2 │
│ B ┆ 100.0 ┆ 9.8 ┆ 2 │
│ B ┆ 100.0 ┆ 10.0 ┆ 2 │
└────────┴───────┴──────┴─────────┘
ps.B is the right choice when fn genuinely doesn’t care about batching (e.g., np.logaddexp (the numerically-stable way to compute log(exp(a) + exp(b)) for which there is no equivalent in polars).
ps.E is for functions that only accept scalars, not arrays at all — like a hand-rolled edit distance, useful for catching typos against a reference list:
def levenshtein(a, b):iflen(a) <len(b): a, b = b, a prev =list(range(len(b) +1))for i, ca inenumerate(a, 1): curr = [i] + [0] *len(b)for j, cb inenumerate(b, 1): curr[j] =min(prev[j] +1, curr[j -1] +1, prev[j -1] + (ca != cb)) prev = currreturn prev[-1]df3 = pl.DataFrame({"typed": ["aplpe", "bananna", "orange"], "correct": ["apple", "banana", "orange"]})df3.with_columns(dist=ps.E(levenshtein)(pl.col("typed"), pl.col("correct")))