Скалярное произведение по строкам в Polars ⇐ Python

[Anonymous]

У меня есть кадр данных с двумя значениями столбцов и весами типа list[i64], и я хотел бы выполнить построчное скалярное произведение этих двух.

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

df = pl.DataFrame({
'values': [[0], [0, 2], [0, 2, 4], [2, 4, 0], [4, 0, 8]],
'weights': [[3], [2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]
})

Есть один способ: сначала поместить значения и веса в структуру, а затем выполнить .map_elements для каждой строки:

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

df.with_columns(
pl.struct(['values', 'weights'])
.map_elements(
lambda x: np.dot(x['values'], x['weights']), return_dtype=pl.Float64
).alias('dot')
)

Но, как указано в документации, map_elements в целом работает намного медленнее, чем собственные полярные выражения, поэтому я пытался реализовать их в собственных выражениях.
Я пробовал следующее:

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

df.with_columns(
pl.concat_list('values', 'weights').alias('combined'),
pl.concat_list('values', 'weights').list.eval(pl.element().slice(0, pl.len() // 2)).alias('values1'),
pl.concat_list('values', 'weights').list.eval(pl.element().slice(pl.len() // 2, pl.len() // 2)).alias('values2'),
pl.concat_list('values', 'weights').list.eval(
pl.element().slice(0, pl.len() // 2).dot(pl.element().slice(pl.len() // 2, pl.len() // 2))
).list.first().alias('dot'),
pl.concat_list('values', 'weights').list.eval(
pl.element().slice(0, pl.len() // 2) + pl.element().slice(pl.len() // 2, pl.len() // 2)
).alias('sum'),
)

Я ожидал, что столбец с точкой будет [0, 6, 16, 10, 28], но оказалось следующее.

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

shape: (5, 7)
┌───────────┬───────────┬─────────────┬───────────┬───────────┬─────┬────────────┐
│ values    ┆ weights   ┆ combined    ┆ values1   ┆ values2   ┆ dot ┆ sum        │
│ ---       ┆ ---       ┆ ---         ┆ ---       ┆ ---       ┆ --- ┆ ---        │
│ list[i64] ┆ list[i64] ┆ list[i64]   ┆ list[i64] ┆ list[i64] ┆ i64 ┆ list[i64]  │
╞═══════════╪═══════════╪═════════════╪═══════════╪═══════════╪═════╪════════════╡
│ [0]       ┆ [3]       ┆ [0, 3]      ┆ [0]       ┆ [3]       ┆ 0   ┆ [0]        │
│ [0, 2]    ┆ [2, 3]    ┆ [0, 2, … 3] ┆ [0, 2]    ┆ [2, 3]    ┆ 4   ┆ [0, 4]     │
│ [0, 2, 4] ┆ [1, 2, 3] ┆ [0, 2, … 3] ┆ [0, 2, 4] ┆ [1, 2, 3] ┆ 20  ┆ [0, 4, 8]  │
│ [2, 4, 0] ┆ [1, 2, 3] ┆ [2, 4, … 3] ┆ [2, 4, 0] ┆ [1, 2, 3] ┆ 20  ┆ [4, 8, 0]  │
│ [4, 0, 8] ┆ [1, 2, 3] ┆ [4, 0, … 3] ┆ [4, 0, 8] ┆ [1, 2, 3] ┆ 80  ┆ [8, 0, 16] │
└───────────┴───────────┴─────────────┴───────────┴───────────┴─────┴────────────┘

Обратите внимание, что даже сумма не такая, как я ожидал. Кажется, что первый фрагмент добавляется вместо второго
Я делаю что-то не так? Как лучше всего выполнить скалярное произведение по строкам в Polars?

Подробнее здесь: https://stackoverflow.com/questions/786 ... -in-polars