将每一行映射到一个列表中并收集所有

Question

我有一个数据框：

df = pl.DataFrame(
    {
        "t_left": [0.0, 1.0, 2.0, 3.0],
        "t_right": [1.0, 2.0, 3.0, 4.0],
        "counts": [1, 2, 3, 4],
    }
)

我想将每一行映射到一个列表中，然后收集所有值（要传递到 例如matplotlib.hist）

我可以手动完成，如下所示：

times = []
for t_left, t_right, counts in df.rows():
    times.extend(np.linspace(t_left, t_right, counts + 1)[1:])

但这对于我的数据集来说太慢了。

我对 python 和 Polars 都是全新的，所以我想知道是否有更好的方法来实现这一点。

编辑

完整的复制粘贴示例以重现。

import polars as pl
import numpy as np

size = 1000000
df = pl.DataFrame(
    {
        "t_left": np.random.rand(size),
        "t_right": np.random.rand(size) + 1,
        "counts": [1] * size,
    }
)

times = []
for t_left, t_right, counts in df.rows():
    times.extend(np.linspace(t_left, t_right, counts + 1)[1:])

Answer 1

由于您需要将最终结果作为一个大列表，因此您实际上不需要使用列表，您可以将explode()数据帧分成多行，这样您就可以对其运行矢量化操作：

\n

因此，以您的示例数据为例：

\n

df = pl.DataFrame(\n    {\n        "t_left": [0.0, 1.0, 2.0, 3.0],\n        "t_right": [1.0, 2.0, 3.0, 4.0],\n        "counts": [1, 2, 3, 4],\n    }\n)\n\xe2\x94\x8c\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x90\n\xe2\x94\x82 t_left \xe2\x94\x86 t_right \xe2\x94\x86 counts \xe2\x94\x82\n\xe2\x94\x82 ---    \xe2\x94\x86 ---     \xe2\x94\x86 ---    \xe2\x94\x82\n\xe2\x94\x82 f64    \xe2\x94\x86 f64     \xe2\x94\x86 i64    \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82 0.0    \xe2\x94\x86 1.0     \xe2\x94\x86 1      \xe2\x94\x82\n\xe2\x94\x82 1.0    \xe2\x94\x86 2.0     \xe2\x94\x86 2      \xe2\x94\x82\n\xe2\x94\x82 2.0    \xe2\x94\x86 3.0     \xe2\x94\x86 3      \xe2\x94\x82\n\xe2\x94\x82 3.0    \xe2\x94\x86 4.0     \xe2\x94\x86 4      \xe2\x94\x82\n\xe2\x94\x94\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x98\n

\n

现在，让我们创建一些可用于计算的列。要计算均匀分布的数字，您需要一个起点 -t_left步骤 -(t_right - t_left) / counts然后您需要重复counts多次。

\n

# Thanks @Hericks for pointing out that code could be simplified even more\ntimes = df.select(\n    start = pl.col(\'t_left\'),\n    step = ((pl.col("t_right") - pl.col("t_left")) / pl.col("counts")),\n    i = pl.int_ranges(1, pl.col(\'counts\') + 1)\n).explode(\'i\')\n\n# my old code\n# times = df.select(\n#    start = pl.col(\'t_left\').repeat_by(\'counts\').explode(),\n#    step = ((pl.col("t_right") - pl.col("t_left")) / pl.col("counts")).repeat_by(\'counts\').explode(),\n#    i = pl.int_ranges(1, pl.col(\'counts\') + 1).explode()\n#)\n\n\xe2\x94\x8c\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xac\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x90\n\xe2\x94\x82 start \xe2\x94\x86 step     \xe2\x94\x86 i   \xe2\x94\x82\n\xe2\x94\x82 ---   \xe2\x94\x86 ---      \xe2\x94\x86 --- \xe2\x94\x82\n\xe2\x94\x82 f64   \xe2\x94\x86 f64      \xe2\x94\x86 i64 \xe2\x94\x82\n\xe2\x95\x9e\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xaa\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\x90\xe2\x95\xa1\n\xe2\x94\x82 0.0   \xe2\x94\x86 1.0      \xe2\x94\x86 1   \xe2\x94\x82\n\xe2\x94\x82 1.0   \xe2\x94\x86 0.5      \xe2\x94\x86 1   \xe2\x94\x82\n\xe2\x94\x82 1.0   \xe2\x94\x86 0.5      \xe2\x94\x86 2   \xe2\x94\x82\n\xe2\x94\x82 2.0   \xe2\x94\x86 0.333333 \xe2\x94\x86 1   \xe2\x94\x82\n\xe2\x94\x82 \xe2\x80\xa6     \xe2\x94\x86 \xe2\x80\xa6        \xe2\x94\x86 \xe2\x80\xa6   \xe2\x94\x82\n\xe2\x94\x82 3.0   \xe2\x94\x86 0.25     \xe2\x94\x86 1   \xe2\x94\x82\n\xe2\x94\x82 3.0   \xe2\x94\x86 0.25     \xe2\x94\x86 2   \xe2\x94\x82\n\xe2\x94\x82 3.0   \xe2\x94\x86 0.25     \xe2\x94\x86 3   \xe2\x94\x82\n\xe2\x94\x82 3.0   \xe2\x94\x86 0.25     \xe2\x94\x86 4   \xe2\x94\x82\n\xe2\x94\x94\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\xb4\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x80\xe2\x94\x98\n

\n

现在，只需应用公式并转换结果即可to_list()：

\n

times.select(\n    res = pl.col(\'start\') + pl.col(\'step\') * pl.col(\'i\')\n)[\'res\'].to_list()\n\n[1.0, 1.5, 2.0, 2.3333333333333335, 2.6666666666666665, 3.0, 3.25, 3.5, 3.75, 4.0]\n

\n