[Python] Dask dataframe running parallel and partitioned by columns

I have a dataframe with multiple companies and countries' data that I am trying to transform in parallel using a function. The data takes a format like this but is much larger and with many more clients:



I have made a dask dataframe that combines the companygroupname and country to make a blended column that I am trying to use as an index for partitions.

I don't want to use npartitions > 1 when making the dask dataframe as that can randomly put half of my client's data in one partition and then the other half in another partition (as far as I can tell)

Here is what I have tried, my question is how can I properly split up my dataframe into specified partitions in order to process each one of them in parallel through my transformation function?

Many thanks in advance for any help

def calculate_predictions(df):

df['CompanyCountryCombined'] = df['CompanyGroupName'].astype(str) + '-' + df['Country']

# Convert pandas DataFrame to Dask DataFrame with one partition
ddf = dd.from_pandas(df, npartitions=1)
ddf = ddf.set_index('CompanyCountryCombined', sorted=True).repartition(partition_size="10MB") # Here I am trying to repartition the data across the combined index
list_of_output_columns = [] # defined elsewhere

# Define meta DataFrame to specify the output format
meta = pd.DataFrame(columns=[list_of_output_columns])

ddf = ddf.sort_values(by=['CompanyCountryCombined','Date'])

models = {} # defined elsewhere too

# Processing the partitions through the process_forecast_group function which is also defined elsewhere
final_results = ddf.groupby('CompanyCountryCombined').apply(
lambda partition: process_forecast_group(partition, models),
meta=meta
).compute()

return final_results

Continue reading...