Пользовательская функция ленивого фрейма данных Polars для строк ⇐ Python

[Anonymous]

Я пытаюсь запустить пользовательскую функцию в ленивом кадре данных построчно.
Сама по себе функция не имеет значения, поэтому я использую softmax в качестве дублера. Все, что здесь имеет значение, это то, что его невозможно вычислить с помощью выражений pl.
Я дошел до этого:

Код: Выделить всё

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()

Однако, если я хочу получить результирующий отложенный df, содержащий столбцы, отличные от столбцов (например, id), я застряну, потому что

Код: Выделить всё

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

больше не работает, потому что pl.col(cols).map_batches выполняется по столбцу...
Не похоже, что это необычный вариант использования, поэтому мне интересно, не упускаю ли я что-то.>

Подробнее здесь: https://stackoverflow.com/questions/794 ... -over-rows