Polars Lazy Dataframe benutzerdefinierte Funktion über Zeilen ⇐ Python

[Anonymous]

Ich versuche, eine benutzerdefinierte Funktion zeilenweise auf einem Lazy-Datenrahmen auszuführen.
Die Funktion selbst spielt keine Rolle, daher verwende ich Softmax als Ersatz. Wichtig dabei ist nur, dass es nicht über pl-Ausdrücke berechenbar ist.
Ich komme so weit:

Code: Select all

import polars as pl
import numpy as np

def softmax(t):
a = np.exp(np.array(t))
return tuple(t/np.sum(t))

ldf = pl.DataFrame({ 'id': [1,2,3], 'a': [0.2,0.1,0.3], 'b': [0.4,0.1,0.3], 'c': [0.4,0.8,0.4]}).lazy()

cols = ['a','b','c']
redict = { f'column_{i}':c for i,c in enumerate(cols) }

ldf.select(cols).map_batches(lambda bdf: bdf.map_rows(softmax).rename(redict)).collect()

Wenn ich jedoch ein resultierendes Lazy-DF erhalten möchte, das andere Spalten als Spalten (z. B. id) enthält, bleibe ich hängen, weil

Code: Select all

ldf.with_columns(pl.col(cols).map_batches(lambda bdf: bdf.map_rows(softmax).rename(redict))).collect()

funktioniert nicht mehr, da pl.col(cols).map_batches spaltenweise ausgeführt wird...
Dies scheint kein ungewöhnlicher Anwendungsfall zu sein, daher frage ich mich, ob ich etwas übersehe.