Using Dask Dataframes instead of Pandas Dataframe in prefect tasks

[Sinha-Ujjawal]

Hi,
Can we use dask dataframe instead of pandas dataframe in prefect tasks (running on same dask executor)?

Thx

[jcrist]

Yes, but with a few caveats:

• You'll need to use a worker_client to submit work: https://distributed.dask.org/en/latest/task-launch.html#connection-with-context-manager

• You shouldn't persist dask objects to prefect Results, since those objects refer to resources stored on a dask cluster, not actual final values. You could either

  • set checkpoint=False on functions that return dask objects (like a dask.dataframe)
  • or always ensure you call df.compute() before returning from a prefect task, so that the result is a pandas object not a dask object.

from dask.distributed import worker_client


@task
def calling_compute_in_a_task(filepath):
    with worker_client():
        df = dd.read_csv(filepath)
        return df.describe().compute()


@task(checkpoint=False)
def using_checkpoint_false(filepath):
    with worker_client():
        return dd.read_csv(filepath)

@task
def compute_describe(df):
    with worker_client():
        return df.describe().compute()

[emkademy]

I am getting ValueError: No workers found. Isn't it possible to run this with LocalDaskExecutor ?

[mlessig]

Does anyone happen to know how to make this work with Prefect 2.0 and Dask .persist() rather than Dask .compute()? Dask .compute() just returns as Pandas dataframe, so you loose all parallelism.