Concat partitions within worker

coalesce.py

def concat_dfs(df_list):
    """
    Concat a list of cudf dataframes
    """
    return cudf.concat(df_list)

def get_delayed_dict(ddf):
    """
    Returns a dicitionary with the dataframe tasks as keys and 
    the dataframe delayed objects as values
    """
    df_delayed = {}
    for delayed_obj in ddf.to_delayed():
        df_delayed[str(delayed_obj.key)] = delayed_obj
    return df_delayed

def concat_within_workers(client, ddf):
    df_delayed = get_delayed_dict(ddf)
    
    result = []
    for worker, tasks in client.has_what().items():
        worker_task_list = []

        for task in list(tasks):
            if task in df_delayed:
                worker_task_list.append(df_delayed[task])
        concat_tasks = delayed(concat_dfs)(worker_task_list)
        result.append(client.persist(collections=concat_tasks,workers=worker))

    return dask_cudf.from_delayed(result)