Dask map_overlap passes partitions in non-chronological order

Dask's df.map_overlap() function passes the partitions in non chronological order. This is extremely inefficient for computations on large datasets with a lot of overlap. Of course, we are running parallelized anyway, so the order of operations is not guaranteed, but this seems like a design flaw that contributes to a higher than necessary memory load. Does anyone have an idea of how this could be solved? Here is a MVE that prints the start and end dates. You can see that the ordering seems to be random. Any ideas how to solve this?

import pandas as pd
import dask.dataframe as dd
import numpy as np

date_range = pd.date_range(start='2023-06-01', end='2023-07-01', freq='T')
df = pd.DataFrame({
    'timestamp': date_range,
    'quantity': np.random.rand(len(date_range))
})

df = dd.from_pandas(df, npartitions=5)
df['timestamp'] = df['timestamp'].dt.ceil('1h')
df = df.set_index('timestamp', sorted=True)

df = df.repartition(freq='1h')

def custom_func(ser):
    print(f'from: {ser.index.min()}, to: {ser.index.max()}')
    return pd.Series(np.mean(ser), index=ser.index)

results = df['quantity'].map_overlap(custom_func, before=pd.Timedelta(days=7), after=0, meta=pd.Series(dtype='float64'), align_dataframes=False, enforce_metadata=False).drop_duplicates()

results.compute()

This is not a complete answer, but hopefully gives you some promising directions to explore.

It's possible to construct a custom DAG that satisfies your desired ordering. Assuming your data is indexed/order, you can use the explicit partition ranges to run the computations. For example:

from dask import compute
from dask.datasets import timeseries
df = timeseries(end='2000-01-10')

n_partitions = df.npartitions
days_offset = 7
partition_combinations = [(i-days_offset, i) for i in range(n_partitions) if i-days_offset>=0]

all_results = []
for start, end in partition_combinations:
   print(start, end)
   subset_result = df.partitions[start:end].mean()
   all_results.append(subset_result)

computed = compute(*all_results)

Note that there is no .map or .map_overlap used since we are explicitly controlling the range of partitions used in a calculation. After all the lazy/delayed computations are formed, the last line will trigger the actual computations.

I found a potential solution to this issue. By manually sorting the partitions based on their timestamps before using the map_overlap() function, we can ensure that the computation proceeds in a more chronological order, thus potentially reducing memory pressure and improving overall efficiency.

Here's the updated code using the partition sorting:

import pandas as pd
import dask.dataframe as dd
import numpy as np

date_range = pd.date_range(start='2023-06-01', end='2023-07-01', freq='T')
df = pd.DataFrame({
    'timestamp': date_range,
    'quantity': np.random.rand(len(date_range))
})

df = dd.from_pandas(df, npartitions=5)
df['timestamp'] = df['timestamp'].dt.ceil('1h')
df = df.set_index('timestamp', sorted=True)

# Repartitioning is not necessary for the solution, so we'll remove it.

def custom_func(ser):
    print(f'from: {ser.index.min()}, to: {ser.index.max()}')
    return pd.Series(np.mean(ser), index=ser.index)

# Manually sort the partitions based on the timestamp before applying map_overlap
df = df.map_partitions(lambda part: part.sort_index(), meta=df._meta)

results = df['quantity'].map_overlap(custom_func, before=pd.Timedelta(days=7), after=0, meta=pd.Series(dtype='float64'), align_dataframes=False, enforce_metadata=False).drop_duplicates()

results.compute()

By using the map_partitions function with a lambda function to sort the partitions before calling map_overlap(), we can mitigate the non-chronological order issue. This should help optimize the memory load and enhance the overall performance of our computations, especially when dealing with extensive overlaps in the data.