Koalas: df.apply(func_without_typehint, axis=1) is not running in parallel

[patryk-oleniuk]

I was hoping that each operation, per row is gonna be executed in parallel.
However this code (and many other experiments I did, including logging to MLflow) shows that it just executes it sequentially for every row.

import time
import databricks.koalas as ks

kdf = ks.DataFrame({"col":[i for i in range(12)]})

def do_job( value ) -> float:
  time.sleep(5)
  return value

kdf.apply(lambda x: do_job(x["col"]) , axis=1)

This is taking 60 sec with > 12 Spark workers.
I was expecting every row to be executed in parallel, that would result in around 5 seconds.

Am I doing something wrong?

My workaround:

kdf.to_spark().rdd.map(lambda x: do_job(x["col"])).collect()

[HyukjinKwon]

There are two possibilities I can think of at the moment:

• When the records is lower then compute.shortcut_limit (default 1000), it executes via just normal pandas API as a shortcut because pandas will be faster in small datasets

• Can you check your index type of the dataframe (see https://koalas.readthedocs.io/en/latest/user_guide/options.html#default-index-type)? It might be the default index type is sequence, and makes all data into single node with single partition. Might be worth trying with distributed type.

[patryk-oleniuk]

Thanks for extremely fast reply.
This helped!

ks.set_option('compute.shortcut_limit', 1)

Distributing the index didn't.
Any idea how to improve the function to execute apply in parallel automatically even for smaller datasets?

[HyukjinKwon]

Can you try this?

import time
import databricks.koalas as ks

kdf = ks.DataFrame({"col":[i for i in range(12)]})

def do_job( df ) -> float:
  value = df["col"]
  time.sleep(5)
  return value

start = time.time()
kdf.apply(do_job , axis=1)
time.time() - start

When kdf.apply takes a function with a type hint, it wouldn't try the shortcut path as it can know the type right away from the given type.

The reason it tried the shortcut was, lambda x: do_job(x["col"]) lambda didn't have the type hint. This seems taking 7~8 secs in my local.

[patryk-oleniuk]

Yes, that works for me too, and by default (without any custom settings, also for more amount of workers and more rows).

Interesting how the type hint can change execution tree. I guess this is the "execute first 1000 rows in pandas and see the output type" kind of behavior..

Anyway, works like a charm now.

[HyukjinKwon]

Yes .. we should definitely make such logic more visible and easy to use. Thanks for your feedback. I will think about how to make it easier and visible more.

[ueshin]

One more thing I want to mention in general is:
The parallelization is based on Spark partitions.
Even if the setting is proper and the function has type hint, if all the data are, let's say, in only one partition, the function will be executed sequentially.

[HyukjinKwon]

I think this was partially documented at https://koalas.readthedocs.io/en/latest/user_guide/transform_apply.html