切片具有大列表的多索引熊猫数据框

Posted 2023-02-23

【中文标题】切片具有大列表的多索引熊猫数据框

【英文标题】：Slicing a multi-index pandas dataframe with a large list

【发布时间】：2020-06-19 14:32:00

【问题描述】：

我有一个带有多索引的大型数据框。我想使用一个相当大的列表来切片这个数据框。下面是一个示例代码。此操作大约需要 10 秒。

import pandas as pd
import numpy as np

df = pd.DataFrame(
    
        "x": np.repeat(np.arange(10000), 50),
        "y": np.repeat(np.arange(50), 10000),
        "val": np.random.rand(50*10000)
    
).set_index(["x", "y"])

large_list = range(5000,10000)

slice = df.loc[(large_list, slice(None)),:] # Takes 10 seconds on my machine

作为比较，如果我将此数据帧写入 hdf 文件并使用与我的切片操作相同的where 条件读取它，则只需 1.5 秒！

df.to_hdf("sample.hdf", key="df", append=True)
df1 = pd.read_hdf("sample.hdf", "df", where='x in large_list')

有没有更快的方法来切入内存？

【问题讨论】：

df.loc[5000:10000] 也可以正常工作，并显着加快速度。它符合您的用例吗？

不，我基本上想传入一个iter 对象而不是slice 本身。它可能是一个包含随机项目的列表。

【参考方案1】：

如果您的意图是按任意列表对多索引进行切片，则使用 query 会快得多

创建从5000 到10000 的任意列表

np.random.seed(0)
large_list =  np.random.choice(list(range(5000, 10000)), 5000, replace=False)

In [2245]: large_list
Out[2245]: array([5398, 8833, 9836, ..., 6653, 7607, 7732])

x = df.query('x in @large_list')

比较结果

In [2246]: y = df.loc[(large_list, slice(None)),:]
In [2249]: np.allclose(x, y)
Out[2249]: True

【讨论】：

这正是我想要的。非常感谢您向我介绍query！！

@najeem：很高兴我能帮上忙。快乐编码:)