Как я могу sort_index, groupby и применить функцию в Dask?

Question

У меня есть функция, которая применяется к фрейму данных pandas, я рассматриваю возможность использования dask для повышения производительности

Это мой существующий код:

df.reset_index(
        level=0,
        inplace=True,
    )

df = df.sort_index().groupby(
            ['col1', 'col2', 'col3'],
            as_index=False).apply(
        myfunction
    )

Я пытаюсь преобразовать этос синтаксисом dask и удалось добраться до:

from dask import dataframe as dd
from multiprocessing import cpu_count

nCores = cpu_count()

df = dd.from_pandas(
        df,
        npartitions=nCores
    ).reset_index().set_index().groupby(
        ['col1', 'col2', 'col3']
        ).apply(
            myfunction
        ).compute()

Похоже, вы можете передать только столбец к set_index, и в dask нет эквивалента sort_index().Как я могу написать этот код панды в синтаксисе dask?

Answer 1

Для тех, кто заинтересован в подобном решении, вот версия кода в синтаксисе dask, которая работает, пожалуйста, обратите внимание, что я сортирую индекс перед передачей его в dask.

from dask import dataframe as dd
from multiprocessing import cpu_count

nCores = cpu_count()

df.sort_index(inplace=True)

df = dd.from_pandas(
        df,
        npartitions=nCores
    ).map_partitions(
        lambda df : df.groupby(
            ['col1', 'col2', 'col3']
            ).apply(
                my_function
            )
        ).compute()

Важно, чтоDask версия была не быстрее, чем подход панды;но очень близкоmy_function в этом случае векторизовано и применяется к каждому объекту groupby, в основном, с использованием массивов numpy.