Более быстрая альтернатива функции Pandas `isin`

У меня очень большой фрейм данных df, который выглядит так:

ID       Value1    Value2
1345      3.2      332
1355      2.2      32
2346      1.0      11
3456      8.9      322

И у меня есть список, содержащий подмножество идентификаторов ID_list. Мне нужно иметь подмножество df для ID, содержащегося в ID_list.

В настоящее время я использую df_sub=df[df.ID.isin(ID_list)] для этого. Но это занимает много времени. ID, содержащийся в ID_list, не имеет никакого шаблона, поэтому он не находится в определенном диапазоне. (И мне нужно применить одну и ту же операцию ко многим подобным фреймам данных. Мне было интересно, есть ли какой-нибудь более быстрый способ сделать это. Будет ли это помогать, если make ID в качестве индекса?

Спасибо!

Ответ 1

ОБНОВЛЕНИЕ 2: Здесь приведена ссылка на более свежий взгляд на выполнение различных операций pandas, хотя в него, похоже, не включены слияние и объединение.

https://github.com/mm-mansour/Fast-Pandas

ОБНОВЛЕНИЕ 1: Эти тесты были для довольно старой версии панд и, вероятно, не все еще актуальны. См. комментарий Майка ниже на merge.

Это зависит от размера ваших данных, но для больших наборов данных DataFrame.join, кажется, будет правильным решением. Для этого необходимо, чтобы ваш индекс DataFrame был вашим "идентификатором", а Series или DataFrame, к которому вы присоединяетесь, чтобы иметь индекс, который является вашим "ID_list". Ряд также должен иметь name для использования с join, который вводится как новое поле под названием name. Вам также нужно указать inner join, чтобы получить что-то вроде isin, потому что join по умолчанию использует left join. Синтаксис запроса in, похоже, имеет те же скоростные характеристики, что и isin для больших наборов данных.

Если вы работаете с небольшими наборами данных, вы получаете другое поведение, и на самом деле становится проще использовать понимание списка или применять его к словарю, чем с помощью isin.

В противном случае вы можете попытаться увеличить скорость с помощью Cython.

# I'm ignoring that the index is defaulting to a sequential number. You
# would need to explicitly assign your IDs to the index here, e.g.:
# >>> l_series.index = ID_list
mil = range(1000000)
l = mil
l_series = pd.Series(l)

df = pd.DataFrame(l_series, columns=['ID'])


In [247]: %timeit df[df.index.isin(l)]
1 loops, best of 3: 1.12 s per loop

In [248]: %timeit df[df.index.isin(l_series)]
1 loops, best of 3: 549 ms per loop

# index vs column does not make a difference here
In [304]: %timeit df[df.ID.isin(l_series)]
1 loops, best of 3: 541 ms per loop

In [305]: %timeit df[df.index.isin(l_series)]
1 loops, best of 3: 529 ms per loop

# query 'in' syntax has the same performance as 'isin'
In [249]: %timeit df.query('index in @l')
1 loops, best of 3: 1.14 s per loop

In [250]: %timeit df.query('index in @l_series')
1 loops, best of 3: 564 ms per loop

# ID must be the index for DataFrame.join and l_series must have a name.
# join defaults to a left join so we need to specify inner for existence.
In [251]: %timeit df.join(l_series, how='inner')
10 loops, best of 3: 93.3 ms per loop

# Smaller datasets.
df = pd.DataFrame([1,2,3,4], columns=['ID'])
l = range(10000)
l_dict = dict(zip(l, l))
l_series = pd.Series(l)
l_series.name = 'ID_list'


In [363]: %timeit df.join(l_series, how='inner')
1000 loops, best of 3: 733 µs per loop

In [291]: %timeit df[df.ID.isin(l_dict)]
1000 loops, best of 3: 742 µs per loop

In [292]: %timeit df[df.ID.isin(l)]
1000 loops, best of 3: 771 µs per loop

In [294]: %timeit df[df.ID.isin(l_series)]
100 loops, best of 3: 2 ms per loop

# It actually faster to use apply or a list comprehension for these small cases.
In [296]: %timeit df[[x in l_dict for x in df.ID]]
1000 loops, best of 3: 203 µs per loop

In [299]: %timeit df[df.ID.apply(lambda x: x in l_dict)]
1000 loops, best of 3: 297 µs per loop