Эффективное применение функции к сгруппированному pandas DataFrame параллельно

Мне часто приходится применять функцию к группам очень большого DataFrame (смешанных типов данных) и хотел бы использовать несколько ядер.

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

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

Ответ 1

Из комментариев выше, кажется, что это запланировано на pandas некоторое время (там также интересный rosetta проект, который я только что заметил).

Однако до тех пор, пока каждая параллельная функциональность не будет включена в pandas, я заметил, что очень просто написать эффективные параллельные расширения без привязки к памяти pandas непосредственно с помощью cython + OpenMP и С++.

Вот краткий пример написания параллельной групповой суммы, использование которой выглядит примерно так:

import pandas as pd
import para_group_demo

df = pd.DataFrame({'a': [1, 2, 1, 2, 1, 1, 0], 'b': range(7)})
print para_group_demo.sum(df.a, df.b)

и вывод:

Примечание Несомненно, эта простая функциональность в конце концов будет частью pandas. Однако некоторые вещи будут более естественными для параллелизма в С++ в течение некоторого времени, и важно знать, как легко объединить это в pandas.

Чтобы сделать это, я написал простейшее расширение с одним источником, код которого следует.

Он начинается с некоторых определений импорта и типов

from libc.stdint cimport int64_t, uint64_t
from libcpp.vector cimport vector
from libcpp.unordered_map cimport unordered_map

cimport cython
from cython.operator cimport dereference as deref, preincrement as inc
from cython.parallel import prange

import pandas as pd

ctypedef unordered_map[int64_t, uint64_t] counts_t
ctypedef unordered_map[int64_t, uint64_t].iterator counts_it_t
ctypedef vector[counts_t] counts_vec_t

Тип С++ unordered_map предназначен для суммирования одним потоком, а vector предназначен для суммирования по всем потокам.

Теперь к функции sum. Он начинается с типизированных просмотров памяти для быстрого доступа:

def sum(crit, vals):
    cdef int64_t[:] crit_view = crit.values
    cdef int64_t[:] vals_view = vals.values

Функция продолжается делением полуподобного на потоки (здесь hardcoded на 4), и каждый поток объединяет записи в своем диапазоне:

    cdef uint64_t num_threads = 4
    cdef uint64_t l = len(crit)
    cdef uint64_t s = l / num_threads + 1
    cdef uint64_t i, j, e
    cdef counts_vec_t counts
    counts = counts_vec_t(num_threads)
    counts.resize(num_threads)
    with cython.boundscheck(False):
        for i in prange(num_threads, nogil=True): 
            j = i * s
            e = j + s
            if e > l:
                e = l
            while j < e:
                counts[i][crit_view[j]] += vals_view[j]
                inc(j)

Когда потоки завершены, функция объединяет все результаты (из разных диапазонов) в один unordered_map:

    cdef counts_t total
    cdef counts_it_t it, e_it
    for i in range(num_threads):
        it = counts[i].begin()
        e_it = counts[i].end()
        while it != e_it:
            total[deref(it).first] += deref(it).second
            inc(it)

Все, что осталось, это создать DataFrame и вернуть результаты:

    key, sum_ = [], []
    it = total.begin()
    e_it = total.end()
    while it != e_it:
        key.append(deref(it).first)
        sum_.append(deref(it).second)
        inc(it)

    df = pd.DataFrame({'key': key, 'sum': sum_})
    df.set_index('key', inplace=True)
    return df