Ускорение операций с большими массивами и наборами данных (Pandas медленнее, Numpy лучше, дальнейшие улучшения?)

Question

У меня большой набор данных, состоящий из миллионов строк и около 6 столбцов. Данные в настоящее время находятся в фрейме данных Pandas, и я ищу самый быстрый способ работы с ним. Например, допустим, я хочу отбросить все строки, в которых значение в одном столбце равно «1».

Вот мой минимальный рабочий пример:

# Create dummy data arrays and pandas dataframe
array_size = int(5e6)
array1 = np.random.rand(array_size)
array2 = np.random.rand(array_size)
array3 = np.random.rand(array_size)
array_condition = np.random.randint(0, 3, size=array_size)

df = pd.DataFrame({'array_condition': array_condition, 'array1': array1, 'array2': array2, 'array3': array3})

def method1():
    df_new = df.drop(df[df.array_condition == 1].index)

РЕДАКТИРОВАТЬ: Как указал Генри Йикв комментариях более быстрый подход Pandas заключается в следующем:

def method1b():
    df_new = df[df.array_condition != 1]

Я считаю, что Pandas может быть довольно медленным в такого рода вещах, поэтому я также реализовал метод, использующий numpy, обрабатывая каждый столбец как отдельныймассив:

def method2():
    masking = array_condition != 1
    array1_new = array1[masking]
    array2_new = array2[masking]
    array3_new = array3[masking]
    array_condition_new = array_condition[masking]    

И результаты:

%timeit method1()
625 ms ± 7.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit methodb()
158 ms ± 7.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit method2()
138 ms ± 3.8 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Таким образом, мы видим небольшое значительное повышение производительности при использовании numpy. Однако это происходит за счет гораздо менее читаемого кода (т. Е. Необходимости создавать маску и применять ее к каждому массиву). Этот метод также не выглядит настолько масштабируемым, как если бы у меня было, скажем, 30 столбцов данных, мне понадобится много строк кода, которые применяют маску к каждому массиву! Кроме того, было бы полезно разрешить необязательные столбцы, чтобы этот метод мог не работать при работе с пустыми массивами.

Поэтому у меня есть 2 вопроса:

1) Есть ли очиститель? / более гибкий способ реализовать это в numpy?

2) Или лучше, есть ли способ более высокой производительности, который я мог бы использовать здесь? например, JIT (numba?), Cython или что-то еще?

PS, на практике можно использовать операции на месте, заменяя старый массив новым после удаления данных

Answer 1

Часть 1: Панды и (возможно) Numpy

Сравните ваши method1b и method2 :

• method1b генерирует DataFrame , что, вероятно, то, что вы хотите,
• method2 генерирует массив Numpy , так что, чтобы получить полностью сопоставимый результат,впоследствии вы должны сгенерировать DataFrame из него.

Поэтому я изменил ваш method2 на:

def method2():
    masking = array_condition != 1
    array1_new = array1[masking]
    array2_new = array2[masking]
    array3_new = array3[masking]
    array_condition_new = array_condition[masking]
    df_new = pd.DataFrame({ 'array_condition': array_condition[masking],
        'array1': array1_new, 'array2': array2_new, 'array3': array3_new})

и затем сравнил выполнениераз (используя % timeit ).

В результате моя (расширенная) версия method2 выполнена примерно на 5% длиннее , чем method1b (проверьте самостоятельно).

Так что я считаю, что если речь идет о single операции, то, вероятно, лучше остаться с Pandas .

Но если вы хотите выполнить с вашим источником DataFrame пару последовательных операций и / или вы удовлетворены результатом как массив Numpy , стоит:

• вызов arr = df.values для получения базового массива Numpy .
• Выполнить все необходимые операции сон использует Numpy методы.
• (Опционально) создает DataFrame из окончательного результата.

Я пытался Numpy версия method1b :

def method3():
    a = df.values
    arr = a[a[:,0] != 1]

но время выполнения было примерно на 40% больше .

Возможно, причина в том, что массив Numpy имеет все элементы одного типа, поэтому столбец array_condition приводится к float , а затем создается весь массив Numpy , что занимает некоторое время.

Часть 2: Numpy и Numba

Альтернативой для рассмотрения является использование пакета Numba - компилятора Python Just-In-Time.

Я сделал такой тест:

Создан массив Numpy (в качестве предварительного шага):

a = df.values

Причина в том, что JIT-скомпилированные методы могут использовать Numpy методы и типы, но не те, что Панды .

Для выполнения теста я использовал почти тот же метод, что и выше, но с @njit аннотация (требуется из numba import njit ):

@njit
def method4():
    arr = a[a[:,0] != 1]

На этот раз:

• Время выполнения составляло около 45% временидля method1b .
• Но поскольку a = df.values было выполнено до цикла тестирования, есть сомнения, сопоставим ли этот результат с более ранними тестами.

В любом случае, попробуйте Numba самостоятельно, возможно, это будет интересный вариант для вас.

Answer 2

Вы можете использовать здесь numpy.where . Он преобразует логическую маску в массив индексов, делая жизнь намного дешевле. Комбинируя это с numpy.vstack, можно выполнять некоторые операции с дешевой памятью:

def method3():
    wh = np.where(array_condition == 1)
    return np.vstack(tuple(col[wh] for col in (array1, array2, array3)))

Это дает следующие временные интервалы:

>>> %timeit method2()
180 ms ± 6.66 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
>>> %timeit method3()
96.9 ms ± 2.5 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Распаковка кортежей позволяет операции довольно легко выполнять операции с памятью,когда объект снова объединен, он меньше. Если вам нужно извлечь столбцы непосредственно из DataFrame, может пригодиться следующий фрагмент кода:

def method3b():
    wh = np.where(array_condition == 1)
    col_names = ['array1','array2','array3']
    return np.vstack(tuple(col[wh] for col in tuple(df[col_name].to_numpy() 
        for col_name in col_names)))

Это позволяет получать столбцы по имени из DataFrame, которые затем распаковываются в кортеж на лету. ,Скорость примерно одинакова:

>>> %timeit method3b()
96.6 ms ± 3.09 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Наслаждайтесь!