Polars 从 n 列中减去 numpy 1xn 数组

Question

我正在与极地斗争。我有一个数据框和一个 numpy 数组。我想减去它们。

import polars as pl
import pandas as pd

df = pl.DataFrame(np.random.randn(6, 4), schema=['#', 'x', 'y', 'z'])

arr = np.array([-10, -20, -30])


df.select(
    pl.col(r'^(x|y|z)$') # ^[xyz]$
).map_rows(
    lambda x: np.array(x) - arr
)

# ComputeError: expected tuple, got ndarray

但是如果我尝试计算范数，那么它就有效：

df.select(
    pl.col(r'^(x|y|z)$')
).map_rows(
    lambda x: np.sum((np.array(x) - arr)**2)**0.5
)
shape: (6, 1)
┌───────────┐
│ map       │
│ ---       │
│ f64       │
╞═══════════╡
│ 38.242255 │
│ 37.239545 │
│ 38.07624  │
│ 36.688312 │
│ 38.419194 │
│ 36.262196 │
└───────────┘

# check if it is correct:
np.sum((df.to_pandas()[['x', 'y', 'z']].to_numpy() - arr)**2, axis=1) ** 0.5
>>> array([38.24225488, 37.23954478, 38.07623986, 36.68831161, 38.41919409,
       36.2621962 ])

在 pandas 中可以这样做：

df.to_pandas()[['x', 'y', 'z']] - arr

x   y   z
0   10.143819   21.875335   29.682364
1   10.360651   21.116404   28.871060
2   9.777666    20.846593   30.325185
3   9.394726    19.357053   29.716592
4   9.223525    21.618511   30.390805
5   9.751234    21.667080   27.393393

一种有效的方法是对每一列分别进行操作。但这意味着很多相同的代码，特别是当列数增加时：

df.select(
    pl.col('x') - arr[0], pl.col('y') - arr[1], pl.col('z') - arr[2]
)

Answer 1

这个问题涉及到一些事情。

首先，你真的真的不想使用

apply

，除非你正在做一些自定义的Python函数

apply 表达式将列的元素传递给 python 函数。 请注意，您现在正在运行 python，这会很慢。

没有真正的极地方法可以做你想做的事。 当 Polars 看到

pl.col(r'^[x|y|z]$').expr

时，它将识别适合正则表达式的每一列，然后将有一个线程执行表达式其余部分的工作。 该表达式不知道它的顺序在哪里。 它只知道它的数据是什么以及它应该做什么。 因此，您无法在

expr

中放入任何内容，让它知道要访问数组中的哪个元素。

要得到你想要的东西，你必须做类似 @ignoring_gravity 的事情，但你可以使用

re

模块。

import re
df.select(pl.col(col)-arr[i] 
          for i, col in enumerate(filter(re.compile(r'^[x|y|z]$').match, df.columns)))

Answer 2

避免

re

导入的另一个选项是：

res = df.select(
    pl.col(col) - c
    for col, c in zip(df.select(pl.col(r'^[x|y|z]$')).columns, arr)
)

对于非常小的数据帧来说，这稍微慢一些（我猜是因为它随后由正则表达式速度决定），但对于较大的数据帧来说同样快。

Answer 3

可以匹配pandas输出

In [15]: df.to_pandas()[['x', 'y', 'z']] - arr
Out[15]:
           x          y          z
0  10.342991  21.258934  29.083287
1  10.136803  21.543558  28.168207
2  11.900141  19.557348  29.490541
3   9.192346  19.498689  28.195094
4   9.219745  20.330358  29.005278
5  11.853378  19.458095  30.357041

与

In [17]: df.select([pl.col(col)-arr[i] for i, col in enumerate(['x', 'y', 'z'])])
Out[17]:
shape: (6, 3)
┌───────────┬───────────┬───────────┐
│ x         ┆ y         ┆ z         │
│ ---       ┆ ---       ┆ ---       │
│ f64       ┆ f64       ┆ f64       │
╞═══════════╪═══════════╪═══════════╡
│ 10.342991 ┆ 21.258934 ┆ 29.083287 │
│ 10.136803 ┆ 21.543558 ┆ 28.168207 │
│ 11.900141 ┆ 19.557348 ┆ 29.490541 │
│ 9.192346  ┆ 19.498689 ┆ 28.195094 │
│ 9.219745  ┆ 20.330358 ┆ 29.005278 │
│ 11.853378 ┆ 19.458095 ┆ 30.357041 │
└───────────┴───────────┴───────────┘

Answer 4

我短暂地看到了我正在寻找的答案，但评论已被删除。

解决方案是返回一个元组：

df.select(
    pl.col(r'^(x|y|z)$')
).map_rows(
    # lambda x: np.array(x) - arr  # old code
    lambda x: tuple(np.array(x) - arr)  # new code
)