Нарезка мультииндексного pandas фрейма данных с большим списком - PullRequest
Новая
       4

Нарезка мультииндексного pandas фрейма данных с большим списком

2 голосов
/ 06 марта 2020

У меня большой массив данных с мультииндексом. Я хотел нарезать этот фрейм данных, используя довольно большой список. Ниже приведен пример кода. Эта операция занимает почти 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')

Есть ли более быстрый способ нарезки в памяти?

1 Ответ

2 голосов
/ 07 марта 2020

Если вы собираетесь разрезать мультииндекс по произвольному списку, использование 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
...