Как усреднить списки в разных столбцах Polars LazyFrame ⇐ Python

[Anonymous]

У меня есть Polars LazyFrame, который имеет 3 столбца типа списка с нулевым значением[f64], что-то вроде этого.

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

import polars as pl

lf = pl.DataFrame({
"1": [
[0.0, 1.1, 2.2],
[0.0, 1.1, 2.2],
[0.0, 1.1, 2.2],
None,
],
"2": [
[0.3, 1.3, 2.3],
[0.4, 1.4, 2.4],
[0.5, 1.5, 2.5],
None,
],
"3": [
[0.7, 1.7, 2.7],
None,
[0.9, 1.9, 2.9],
None,
],
}).lazy()

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

┌─────────────────┬─────────────────┬─────────────────┐
│ 1               ┆ 2               ┆ 3               │
│ ---             ┆ ---             ┆ ---             │
│ list[f64]       ┆ list[f64]       ┆ list[f64]       │
╞═════════════════╪═════════════════╪═════════════════╡
│ [0.0, 1.1, 2.2] ┆ [0.3, 1.3, 2.3] ┆ [0.7, 1.7, 2.7] │
│ [0.0, 1.1, 2.2] ┆ [0.4, 1.4, 2.4] ┆ null            │
│ [0.0, 1.1, 2.2] ┆ [0.5, 1.5, 2.5] ┆ [0.9, 1.9, 2.9] │
│ null            ┆ null            ┆ null            │
└─────────────────┴─────────────────┴─────────────────┘

Мне нужно добавить столбец со средним значением трех списков столбцов, кроме того.

• когда в строке есть только ноль, тогда avg будет заполненным нулями списком фиксированной длины 3
• когда один элемент имеет значение NULL среднее значение будет вычисляться для ненулевых списков

Под «средним значением списков» я подразумеваю поэлементную сумму, деленную на количество списков, участвующих в сумме.
Итак, в первой строке я хочу:

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

[
0.0 + 0.3 + 0.7,
1.1 + 1.3 + 1.7,
2.2 + 2.3 + 2.7
] / 3
=
[
1.0, 4.1, 7.2
] / 3
=
[0.33, 1.36, 2.40]

Во второй строке: [0,0 + 0,4, 1,1 + 1,4, 2,2 + 2,4] / 2 = [0,4, 2,8, 4,6] / 2 = [0,2, 1,4, 2,3].
В последней строке: [0,0, 0,0, 0.0].
Я нашел способ суммировать столбцы

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

lf.select(
pl.sum_horizontal(
pl.all().list.explode()
).reshape((1, -1)).alias("sum"),
).collect()

Но это работает только в том случае, если все элементы в строке не равны нулю.


Подробнее здесь: https://stackoverflow.com/questions/769 ... -lazyframe