Pandas apply in parallel when axis=0

Question

I want to apply some function on all pandas columns in parallel. For example, I want to do this in parallel:

def my_sum(x, a):
    return x + a


df = pd.DataFrame({'num_legs': [2, 4, 8, 0],
                   'num_wings': [2, 0, 0, 0]})
df.apply(lambda x: my_sum(x, 2), axis=0)

I know there is a swifter package, but it doesn't support axis=0 in apply:

NotImplementedError: Swifter cannot perform axis=0 applies on large datasets. Dask currently does not have an axis=0 apply implemented. More details at https://github.com/jmcarpenter2/swifter/issues/10

Dask also doesn't support this for axis=0 (according to documentation in swifter).

I have googled several sources but couldn't find an easy solution.

Can't believe this is so complicated in pandas.

Answer 1

Koalas provides a way to perform computation on a dataframe in parallel. It accepts the same commands as pandas but performs them on a Apache Spark engine in the background.

Note that you do need the parallel infrastructure available in order to use it properly.

On their blog post they compare the following chunks of code:

pandas:

import pandas as pd
df = pd.DataFrame({'x': [1, 2], 'y': [3, 4], 'z': [5, 6]})
# Rename columns
df.columns = [‘x’, ‘y’, ‘z1’]
# Do some operations in place
df[‘x2’] = df.x * df.x

Koalas:

import databricks.koalas as ks
df = ks.DataFrame({'x': [1, 2], 'y': [3, 4], 'z': [5, 6]})
# Rename columns
df.columns = [‘x’, ‘y’, ‘z1’]
# Do some operations in place
df[‘x2’] = df.x * df.x

Answer 2

The real answer is hidden in comments, so I will add this answer: use mapply.

import pandas as pd
import mapply

mapply.init(n_workers=-1, chunk_size=1)

def my_sum(x, a):
    return x + a

df = pd.DataFrame({'num_legs': [2, 4, 8, 0],
                   'num_wings': [2, 0, 0, 0]})
df.mapply(lambda x: my_sum(x, 2), axis=0)

I tried swifter and pandarallel, but swifter simply didn't work with columns, and pandarallel seemed to duplicate the work on all workers. Only mapply worked.

Answer 3

You can use the dask delayed interface to set up a custom workflow:

import pandas as pd
import dask
import distributed

# start local cluster, by default one worker per core
client = distributed.Client() 

@dask.delayed
def my_sum(x, a):
    return x + a

df = pd.DataFrame({'num_legs': [2, 4, 8, 0],
                   'num_wings': [2, 0, 0, 0]})    

# Here, we mimic the apply command. However, we do not
# actually run any computation. Instead, that line of code 
# results in a list of delayed objects, which contain the 
# information what computation should be performed eventually
delayeds = [my_sum(df[column], 2) for column in df.columns]

# send the list of delayed objects to the cluster, which will 
# start computing the result in parallel. 
# It returns future objects, pointing to the computation while
# it is still running
futures = client.compute(delayeds)

# get all the results, as soon as they are ready. This returns 
# a list of pandas Series objects, each is one column of the 
# output dataframe
computed_columns = client.gather(futures)

# create dataframe out of individual columns
computed_df = pd.concat(computed_columns, axis = 1)

Alternatively, you could also use the multiprocessing backend of dask:

import pandas as pd
import dask

@dask.delayed
def my_sum(x, a):
    return x + a

df = pd.DataFrame({'num_legs': [2, 4, 8, 0],
                   'num_wings': [2, 0, 0, 0]})    

# same as above
delayeds = [my_sum(df[column], 2) for column in df.columns]

# run the computation using the dask's multiprocessing backend
computed_columns = dask.compute(delayeds, scheduler = 'processes')

# create dataframe out of individual columns
computed_df = pd.concat(computed_columns, axis = 1)

Answer 4

In my opinion, this case should be tackled focusing on how the data is split over the available resources. Dask offers map_partitions which applies a Python function on each DataFrame partition. Of course, the number of rows per partition that your workstation can deal with depends on the available hardware resources. Here is an example based on the information you provided in your question:

# imports
import dask
from dask import dataframe as dd
import multiprocessing as mp
import numpy as np
import pandas as pd

# range for values to be randomly generated
range_ = {
    "min": 0,
    "max": 100
}

# rows and columns for the fake dataframe
df_shape = (
                int(1e8), # rows
                2 # columns
            )

# some fake data
data_in = pd.DataFrame(np.random.randint(range_["min"], range_["max"], size = df_shape), columns = ["legs", "wings"])

# function to apply adding some value a to the partition
def my_sum(x, a):
    return x + a
"""
applies my_sum on the partitions rowwise (axis = 0)

number of partitions = cpu_count

the scheduler can be:
"threads": Uses a ThreadPool in the local process
"processes": Uses a ProcessPool to spread work between processes
"single-threaded": Uses a for-loop in the current thread
"""
data_out = dd.from_pandas(data_in, npartitions = mp.cpu_count()).map_partitions(
        lambda df: df.apply(
            my_sum, axis = 0, a = 2
        )
).compute(scheduler = "threads")

# inspection
print(data_in.head(5))
print(data_out.head(5))

This implementation was tested on a random generated dataframe with 100,000,000 rows and 2 columns.

Workstation Specs
CPU: Intel(R) Core(TM) i7-8750H CPU @ 2.20GHz
Memory Total: 16698340 kB
OS: Ubuntu 18.04.4 LTS