В Pandas DataFrame я хочу создать новый столбец условно на основе значения другого столбца. В моем приложении DataFrame обычно имеет несколько миллионов строк, а число уникальных условных значений невелико, порядка единицы. Производительность чрезвычайно важна: какой самый быстрый способ создать новый столбец?
Я создал примерный пример ниже, и уже пытался и сравнивал различные методы. В примере условное заполнение представлено поиском словаря на основе значения label
столбца (здесь: один из 1, 2, 3
).
lookup_dict = {
1: 100, # arbitrary
2: 200, # arbitrary
3: 300, # arbitrary
}
Затем я ожидаю, что мой DataFrame будет заполнен как:
label output
0 3 300
1 2 200
2 3 300
3 3 300
4 2 200
5 2 200
6 1 100
7 1 100
Ниже приведены 6 различных методов, проверенных на 10- Nlines
линиях (параметр Nlines
в тестовом коде):
- метод 1:
pandas.groupby().apply()
- метод 2:
pandas.groupby().indices.items()
- метод 3:
pandas.Series.map
- метод 4: для цикла на этикетках
- метод 5:
numpy.select
- метод 6: numba
Полный код доступен в конце ответа с моментами выполнения всех методов. Вывод каждого метода считается равным перед сравнением результатов.
метод 1: pandas.groupby().apply()
Я использую pandas.groupby()
на label
, затем заполняет каждый блок тем же значением, используя apply()
.
def fill_output(r):
''' called by groupby().apply(): all r.label values are the same '''
r.loc[:, 'output'] = lookup_dict[r.iloc[0]['label']]
return r
df = df.groupby('label').apply(fill_output)
я получил
>>> method_1_groupby ran in 2.29s (average over 3 iterations)
Обратите внимание, что groupby(). Apply() запускается дважды в первой группе, чтобы определить, какой путь кода использовать (см. Pandas # 2936). Это может замедлить работу небольшого числа групп. Я обманул метод 1, добавив первую фиктивную группу, но я не получил большого улучшения.
метод 2: pandas.groupby().indices.items()
Второй вариант: вместо использования apply
я обращаюсь к указателю directy с помощью groupby().indices.items()
. Это заканчивается в два раза быстрее, чем метод 1, и это метод, который я использовал в течение длительного времени
dgb = df.groupby('label')
for label, idx in dgb.indices.items():
df.loc[idx, 'output'] = lookup_dict[label]
Получил:
method_2_indices ran in 1.21s (average over 3 iterations)
метод 3: pandas.Series.map
Я использовал Pandas.Series.map.
df['output'] = df.label.map(lookup_dict.get)
У меня были очень хорошие результаты в аналогичных случаях, когда количество просмотренных значений было сопоставимо с количеством строк. В данном случае map
заканчивается в два раза медленнее, чем метод 1.
method_3_map работает в 3.07s (в среднем более 3 итераций)
Я отношу это к небольшому числу значений поиска, но может возникнуть проблема с тем, как я его реализовал.
метод 4: для цикла на этикетках
4-й метод довольно наивен: я просто перебираю все метки и выбираю соответствующую часть DataFrame.
for label, value in lookup_dict.items():
df.loc[df.label == label, 'output'] = value
Удивительно, но, в конечном итоге, я получил гораздо более быстрые результаты, чем в предыдущих случаях. Я ожидал, что решения на основе groupby
будут быстрее, чем это, потому что Pandas должен сделать три сравнения с df.label == label
. Результаты доказывают, что я ошибаюсь:
method_4_forloop ran in 0.54s (average over 3 iterations)
метод 5: numpy.select
В пятом методе используется функция select
numpy, основанная на этом qaru.site/info/59317/....
conditions = [df.label == k for k in lookup_dict.keys()]
choices = list(lookup_dict.values())
df['output'] = np.select(conditions, choices)
Это дает наилучшие результаты:
method_5_select ran in 0.29s (average over 3 iterations)
В конце концов, я попробовал numba
в методе 6.
метод 6: numba
Как раз для примера, условные значения заполнения являются жестким кодом в скомпилированной функции. Я не знаю, как дать Numba список как константу времени выполнения:
@jit(int64[:](int64[:]), nopython=True)
def hardcoded_conditional_filling(column):
output = np.zeros_like(column)
i = 0
for c in column:
if c == 1:
output[i] = 100
elif c == 2:
output[i] = 200
elif c == 3:
output[i] = 300
i += 1
return output
df['output'] = hardcoded_conditional_filling(df.label.values)
Я закончил с лучшим временем, быстрее, чем метод 5 на 50%.
method_6_numba ran in 0.19s (average over 3 iterations)
Я не реализовал это по причине, указанной выше: я не знаю, как дать Numba список в качестве постоянной времени выполнения без существенного снижения производительности.
Полный код
import pandas as pd
import numpy as np
from timeit import timeit
from numba import jit, int64
lookup_dict = {
1: 100, # arbitrary
2: 200, # arbitrary
3: 300, # arbitrary
}
Nlines = int(1e7)
# Generate
label = np.round(np.random.rand(Nlines)*2+1).astype(np.int64)
df0 = pd.DataFrame(label, columns=['label'])
# Now the goal is to assign the look_up_dict values to a new column 'output'
# based on the value of label
# Method 1
# using groupby().apply()
def method_1_groupby(df):
def fill_output(r):
''' called by groupby().apply(): all r.label values are the same '''
#print(r.iloc[0]['label']) # activate to reveal the #2936 issue in Pandas
r.loc[:, 'output'] = lookup_dict[r.iloc[0]['label']]
return r
df = df.groupby('label').apply(fill_output)
return df
def method_2_indices(df):
dgb = df.groupby('label')
for label, idx in dgb.indices.items():
df.loc[idx, 'output'] = lookup_dict[label]
return df
def method_3_map(df):
df['output'] = df.label.map(lookup_dict.get)
return df
def method_4_forloop(df):
''' naive '''
for label, value in lookup_dict.items():
df.loc[df.label == label, 'output'] = value
return df
def method_5_select(df):
''' Based on answer from
/info/59317/pandas-conditional-creation-of-a-seriesdataframe-column/411715#411715
'''
conditions = [df.label == k for k in lookup_dict.keys()]
choices = list(lookup_dict.values())
df['output'] = np.select(conditions, choices)
return df
def method_6_numba(df):
''' This works, but it is hardcoded and i don't really know how
to make it compile with list as runtime constants'''
@jit(int64[:](int64[:]), nopython=True)
def hardcoded_conditional_filling(column):
output = np.zeros_like(column)
i = 0
for c in column:
if c == 1:
output[i] = 100
elif c == 2:
output[i] = 200
elif c == 3:
output[i] = 300
i += 1
return output
df['output'] = hardcoded_conditional_filling(df.label.values)
return df
df1 = method_1_groupby(df0)
df2 = method_2_indices(df0.copy())
df3 = method_3_map(df0.copy())
df4 = method_4_forloop(df0.copy())
df5 = method_5_select(df0.copy())
df6 = method_6_numba(df0.copy())
# make sure we havent modified the input (would bias the results)
assert 'output' not in df0.columns
# Test validity
assert (df1 == df2).all().all()
assert (df1 == df3).all().all()
assert (df1 == df4).all().all()
assert (df1 == df5).all().all()
assert (df1 == df6).all().all()
# Compare performances
Nites = 3
print('Compare performances for {0:.1g} lines'.format(Nlines))
print('-'*30)
for method in [
'method_1_groupby', 'method_2_indices',
'method_3_map', 'method_4_forloop',
'method_5_select', 'method_6_numba']:
print('{0} ran in {1:.2f}s (average over {2} iterations)'.format(
method,
timeit("{0}(df)".format(method), setup="from __main__ import df0, {0}; df=df0.copy()".format(method), number=Nites)/Nites,
Nites))
Выход:
Compare performances for 1e+07 lines
------------------------------
method_1_groupby ran in 2.29s (average over 3 iterations)
method_2_indices ran in 1.21s (average over 3 iterations)
method_3_map ran in 3.07s (average over 3 iterations)
method_4_forloop ran in 0.54s (average over 3 iterations)
method_5_select ran in 0.29s (average over 3 iterations)
method_6_numba ran in 0.19s (average over 3 iterations)
Меня интересовало бы любое другое решение, которое могло бы дать лучшие результаты. Первоначально я искал методы на основе Pandas, но я также принимаю решения на основе numba/cython.
редактировать
Добавление методов Chrisb для сравнения:
def method_3b_mapdirect(df):
''' Suggested by /info/16059216/fastest-way-to-create-a-pandas-column-conditionally/24960318#24960318'''
df['output'] = df.label.map(lookup_dict)
return df
def method_7_take(df):
''' Based on answer from
/info/59317/pandas-conditional-creation-of-a-seriesdataframe-column/411715#411715
Exploiting that labels are continuous integers
'''
lookup_arr = np.array(list(lookup_dict.values()))
df['output'] = lookup_arr.take(df['label'] - 1)
return df
С временем выполнения:
method_3_mapdirect ran in 0.23s (average over 3 iterations)
method_7_take ran in 0.11s (average over 3 iterations)
Что делает # 3 быстрее, чем любой другой метод (# 6 в сторону), а также самый элегантный. Используйте # 7, если ваш случай пользователя совместим.