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

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]

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.

[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.