Repository: dask/dask

Issue:
map_overlap with multiple dask dataframes
Hi,

I tried to use dask.dataframe.rolling.map_overlap function with multiple dask dataframes. I was expecting the same behavior as map_partitions but unlike map_partitions, only the first dask dataframe is materialized at execution time (https://docs.dask.org/en/latest/generated/dask.dataframe.rolling.map_overlap.html).

I would like to know if we could have the same behavior on map_overlap ?

As an example, I would like to write:

```
ddf1 = dask.dataframe.from_pandas(...)
ddf2 = dask.dataframe.from_pandas(...)

def diff(dd1, dd2):
	return dd1 - dd2
	
dask.dataframe.rolling.map_overlap(diff, ddf1, ddf2, before=5, after=5)
```

What do you think about it ?

Any ideas to address this kind of issues are welcomed :)

Thanks for your help,

Fabien



Discussion:
@faulaire Welcome to Dask, and thanks for reporting this! It does look like using multiple DataFrames isn't implemented in `map_overlap`, but it might be a reasonable case to support. 

One non-ideal workaround for simple operations could be to [convert the DataFrames to Dask Arrays](https://docs.dask.org/en/stable/generated/dask.dataframe.DataFrame.to_dask_array.html), and use [`da.map_overlap`](https://docs.dask.org/en/stable/generated/dask.array.map_overlap.html) which does allow for multiple Arrays. I'll keep looking into this!
Hi @pavithraes , ok, thanks for your answer.

I already started to implement the support for multiple dataframes on my side. I need some times to tests all cases but I'll raise a pull request before the end of this month.
> I already started to implement the support for multiple dataframes on my side. I need some times to tests all cases but I'll raise a pull request before the end of this month.

Thank you! Looking forward to it!
I'd like to change the API for map_overlap function. Because it can now accept multiple dataframes, it should have the same signature as map_partitions:

My proposal is;

```python
def map_overlap(
    func,
    before,
    after,
    *args,
    meta=no_default,
    enforce_metadata=True,
    transform_divisions=True,
    align_dataframes=True,
    **kwargs,
):
```

What's your opinion about this change ? 

Moreover, because I don't want to copy/paste code from map_partitions, is it ok if I create some helper functions from map_partitions code in order to share code to handle meta, partition alignement... ?

WIP PR : #9065 
Hi,

This pull request (draft) : #9065 implement multi dataframe handling.

Can you please have a look ?

I was not able to materialize arguments passed in kwargs, if someone can help me on this point it would be great.

Thanks


@faulaire Thanks for opening a PR, apologies for not getting to this sooner!

I see your PR was closed, shall we re-open it?


@pavithraes, my fault. I reorganized my forked of dask and remove the existing one for a new one with multiple branches (I'm not a git expert, I'm sure it's possible to do it in a smoother way). I'll raise a new PR soon with some changes. The first one was buggy in distributed mode.

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -861,15 +861,39 @@ def map_overlap(self, func, before, after, *args, **kwargs):
         ----------
         func : function
             Function applied to each partition.
-        before : int
-            The number of rows to prepend to partition ``i`` from the end of
+        before : int or timedelta
+            The rows to prepend to partition ``i`` from the end of
             partition ``i - 1``.
-        after : int
-            The number of rows to append to partition ``i`` from the beginning
+        after : int or timedelta
+            The rows to append to partition ``i`` from the beginning
             of partition ``i + 1``.
         args, kwargs :
-            Arguments and keywords to pass to the function. The partition will
-            be the first argument, and these will be passed *after*.
+            Positional and keyword arguments to pass to the function.
+            Positional arguments are computed on a per-partition basis, while
+            keyword arguments are shared across all partitions. The partition
+            itself will be the first positional argument, with all other
+            arguments passed *after*. Arguments can be ``Scalar``, ``Delayed``,
+            or regular Python objects. DataFrame-like args (both dask and
+            pandas) will be repartitioned to align (if necessary) before
+            applying the function; see ``align_dataframes`` to control this
+            behavior.
+        enforce_metadata : bool, default True
+            Whether to enforce at runtime that the structure of the DataFrame
+            produced by ``func`` actually matches the structure of ``meta``.
+            This will rename and reorder columns for each partition,
+            and will raise an error if this doesn't work or types don't match.
+        transform_divisions : bool, default True
+            Whether to apply the function onto the divisions and apply those
+            transformed divisions to the output.
+        align_dataframes : bool, default True
+            Whether to repartition DataFrame- or Series-like args
+            (both dask and pandas) so their divisions align before applying
+            the function. This requires all inputs to have known divisions.
+            Single-partition inputs will be split into multiple partitions.
+
+            If False, all inputs must have either the same number of partitions
+            or a single partition. Single-partition inputs will be broadcast to
+            every partition of multi-partition inputs.
         $META
 
         Notes
@@ -891,8 +915,6 @@ def map_overlap(self, func, before, after, *args, **kwargs):
 
         5. Trim ``after`` rows from the end of all but the last partition.
 
-        Note that the index and divisions are assumed to remain unchanged.
-
         Examples
         --------
         Given a DataFrame, Series, or Index, such as:
@@ -6408,18 +6430,8 @@ def map_partitions(
             ) from e
 
     dfs = [df for df in args if isinstance(df, _Frame)]
-    meta_index = getattr(make_meta(dfs[0]), "index", None) if dfs else None
-    if parent_meta is None and dfs:
-        parent_meta = dfs[0]._meta
 
-    if meta is no_default:
-        # Use non-normalized kwargs here, as we want the real values (not
-        # delayed values)
-        meta = _emulate(func, *args, udf=True, **kwargs)
-        meta_is_emulated = True
-    else:
-        meta = make_meta(meta, index=meta_index, parent_meta=parent_meta)
-        meta_is_emulated = False
+    meta = _get_meta_map_partitions(args, dfs, func, kwargs, meta, parent_meta)
 
     if all(isinstance(arg, Scalar) for arg in args):
         layer = {
@@ -6432,20 +6444,6 @@ def map_partitions(
         }
         graph = HighLevelGraph.from_collections(name, layer, dependencies=args)
         return Scalar(graph, name, meta)
-    elif not (has_parallel_type(meta) or is_arraylike(meta) and meta.shape):
-        if not meta_is_emulated:
-            warnings.warn(
-                "Meta is not valid, `map_partitions` expects output to be a pandas object. "
-                "Try passing a pandas object as meta or a dict or tuple representing the "
-                "(name, dtype) of the columns. In the future the meta you passed will not work.",
-                FutureWarning,
-            )
-        # If `meta` is not a pandas object, the concatenated results will be a
-        # different type
-        meta = make_meta(_concat([meta]), index=meta_index)
-
-    # Ensure meta is empty series
-    meta = make_meta(meta, parent_meta=parent_meta)
 
     args2 = []
     dependencies = []
@@ -6472,25 +6470,9 @@ def map_partitions(
         if collections:
             simple = False
 
-    if align_dataframes:
-        divisions = dfs[0].divisions
-    else:
-        # Unaligned, dfs is a mix of 1 partition and 1+ partition dataframes,
-        # use longest divisions found
-        divisions = max((d.divisions for d in dfs), key=len)
-
-    if transform_divisions and isinstance(dfs[0], Index) and len(dfs) == 1:
-        try:
-            divisions = func(
-                *[pd.Index(a.divisions) if a is dfs[0] else a for a in args], **kwargs
-            )
-            if isinstance(divisions, pd.Index):
-                divisions = methods.tolist(divisions)
-        except Exception:
-            pass
-        else:
-            if not valid_divisions(divisions):
-                divisions = [None] * (dfs[0].npartitions + 1)
+    divisions = _get_divisions_map_partitions(
+        align_dataframes, transform_divisions, dfs, func, args, kwargs
+    )
 
     if has_keyword(func, "partition_info"):
         partition_info = {
@@ -6524,6 +6506,69 @@ def func(partition_info, *args, **kwargs):
     return new_dd_object(graph, name, meta, divisions)
 
 
+def _get_divisions_map_partitions(
+    align_dataframes, transform_divisions, dfs, func, args, kwargs
+):
+    """
+    Helper to get divisions for map_partitions and map_overlap output.
+    """
+    if align_dataframes:
+        divisions = dfs[0].divisions
+    else:
+        # Unaligned, dfs is a mix of 1 partition and 1+ partition dataframes,
+        # use longest divisions found
+        divisions = max((d.divisions for d in dfs), key=len)
+    if transform_divisions and isinstance(dfs[0], Index) and len(dfs) == 1:
+        try:
+            divisions = func(
+                *[pd.Index(a.divisions) if a is dfs[0] else a for a in args], **kwargs
+            )
+            if isinstance(divisions, pd.Index):
+                divisions = methods.tolist(divisions)
+        except Exception:
+            pass
+        else:
+            if not valid_divisions(divisions):
+                divisions = [None] * (dfs[0].npartitions + 1)
+    return divisions
+
+
+def _get_meta_map_partitions(args, dfs, func, kwargs, meta, parent_meta):
+    """
+    Helper to generate metadata for map_partitions and map_overlap output.
+    """
+    meta_index = getattr(make_meta(dfs[0]), "index", None) if dfs else None
+    if parent_meta is None and dfs:
+        parent_meta = dfs[0]._meta
+    if meta is no_default:
+        # Use non-normalized kwargs here, as we want the real values (not
+        # delayed values)
+        meta = _emulate(func, *args, udf=True, **kwargs)
+        meta_is_emulated = True
+    else:
+        meta = make_meta(meta, index=meta_index, parent_meta=parent_meta)
+        meta_is_emulated = False
+
+    if not (has_parallel_type(meta) or is_arraylike(meta) and meta.shape) and not all(
+        isinstance(arg, Scalar) for arg in args
+    ):
+        if not meta_is_emulated:
+            warnings.warn(
+                "Meta is not valid, `map_partitions` and `map_overlap` expects output to be a pandas object. "
+                "Try passing a pandas object as meta or a dict or tuple representing the "
+                "(name, dtype) of the columns. In the future the meta you passed will not work.",
+                FutureWarning,
+            )
+        # If `meta` is not a pandas object, the concatenated results will be a
+        # different type
+        meta = make_meta(_concat([meta]), index=meta_index)
+
+    # Ensure meta is empty series
+    meta = make_meta(meta, parent_meta=parent_meta)
+
+    return meta
+
+
 def apply_and_enforce(*args, **kwargs):
     """Apply a function, and enforce the output to match meta
 
diff --git a/dask/dataframe/rolling.py b/dask/dataframe/rolling.py
--- a/dask/dataframe/rolling.py
+++ b/dask/dataframe/rolling.py
@@ -6,18 +6,30 @@
 from pandas.api.types import is_datetime64_any_dtype
 from pandas.core.window import Rolling as pd_Rolling
 
+from dask.array.core import normalize_arg
 from dask.base import tokenize
+from dask.blockwise import BlockwiseDepDict
 from dask.dataframe import methods
-from dask.dataframe.core import _emulate
-from dask.dataframe.utils import make_meta
+from dask.dataframe.core import (
+    Scalar,
+    _Frame,
+    _get_divisions_map_partitions,
+    _get_meta_map_partitions,
+    _maybe_from_pandas,
+    apply_and_enforce,
+    new_dd_object,
+    no_default,
+    partitionwise_graph,
+)
+from dask.dataframe.multi import _maybe_align_partitions
+from dask.delayed import unpack_collections
 from dask.highlevelgraph import HighLevelGraph
-from dask.utils import M, derived_from, funcname, has_keyword
+from dask.utils import M, apply, derived_from, funcname, has_keyword
 
+CombinedOutput = type("CombinedOutput", (tuple,), {})
 
-def overlap_chunk(
-    func, prev_part, current_part, next_part, before, after, args, kwargs
-):
 
+def _combined_parts(prev_part, current_part, next_part, before, after):
     msg = (
         "Partition size is less than overlapping "
         "window size. Try using ``df.repartition`` "
@@ -34,50 +46,110 @@ def overlap_chunk(
 
     parts = [p for p in (prev_part, current_part, next_part) if p is not None]
     combined = methods.concat(parts)
-    out = func(combined, *args, **kwargs)
-    if prev_part is None:
+
+    return CombinedOutput(
+        (
+            combined,
+            len(prev_part) if prev_part is not None else None,
+            len(next_part) if next_part is not None else None,
+        )
+    )
+
+
+def overlap_chunk(func, before, after, *args, **kwargs):
+    dfs = [df for df in args if isinstance(df, CombinedOutput)]
+    combined, prev_part_length, next_part_length = dfs[0]
+
+    args = [arg[0] if isinstance(arg, CombinedOutput) else arg for arg in args]
+
+    out = func(*args, **kwargs)
+
+    if prev_part_length is None:
         before = None
     if isinstance(before, datetime.timedelta):
-        before = len(prev_part)
+        before = prev_part_length
 
     expansion = None
     if combined.shape[0] != 0:
         expansion = out.shape[0] // combined.shape[0]
     if before and expansion:
         before *= expansion
-    if next_part is None:
+    if next_part_length is None:
         return out.iloc[before:]
     if isinstance(after, datetime.timedelta):
-        after = len(next_part)
+        after = next_part_length
     if after and expansion:
         after *= expansion
     return out.iloc[before:-after]
 
 
-def map_overlap(func, df, before, after, *args, **kwargs):
+def map_overlap(
+    func,
+    df,
+    before,
+    after,
+    *args,
+    meta=no_default,
+    enforce_metadata=True,
+    transform_divisions=True,
+    align_dataframes=True,
+    **kwargs,
+):
     """Apply a function to each partition, sharing rows with adjacent partitions.
 
     Parameters
     ----------
     func : function
-        Function applied to each partition.
-    df : dd.DataFrame, dd.Series
+        The function applied to each partition. If this function accepts
+        the special ``partition_info`` keyword argument, it will recieve
+        information on the partition's relative location within the
+        dataframe.
+    df: dd.DataFrame, dd.Series
+    args, kwargs :
+        Positional and keyword arguments to pass to the function.
+        Positional arguments are computed on a per-partition basis, while
+        keyword arguments are shared across all partitions. The partition
+        itself will be the first positional argument, with all other
+        arguments passed *after*. Arguments can be ``Scalar``, ``Delayed``,
+        or regular Python objects. DataFrame-like args (both dask and
+        pandas) will be repartitioned to align (if necessary) before
+        applying the function; see ``align_dataframes`` to control this
+        behavior.
+    enforce_metadata : bool, default True
+        Whether to enforce at runtime that the structure of the DataFrame
+        produced by ``func`` actually matches the structure of ``meta``.
+        This will rename and reorder columns for each partition,
+        and will raise an error if this doesn't work or types don't match.
     before : int or timedelta
         The rows to prepend to partition ``i`` from the end of
         partition ``i - 1``.
     after : int or timedelta
         The rows to append to partition ``i`` from the beginning
         of partition ``i + 1``.
-    args, kwargs :
-        Arguments and keywords to pass to the function. The partition will
-        be the first argument, and these will be passed *after*.
+    transform_divisions : bool, default True
+        Whether to apply the function onto the divisions and apply those
+        transformed divisions to the output.
+    align_dataframes : bool, default True
+        Whether to repartition DataFrame- or Series-like args
+        (both dask and pandas) so their divisions align before applying
+        the function. This requires all inputs to have known divisions.
+        Single-partition inputs will be split into multiple partitions.
+
+        If False, all inputs must have either the same number of partitions
+        or a single partition. Single-partition inputs will be broadcast to
+        every partition of multi-partition inputs.
+    $META
 
     See Also
     --------
     dd.DataFrame.map_overlap
     """
+    args = (df,) + args
+
+    dfs = [df for df in args if isinstance(df, _Frame)]
+
     if isinstance(before, datetime.timedelta) or isinstance(after, datetime.timedelta):
-        if not is_datetime64_any_dtype(df.index._meta_nonempty.inferred_type):
+        if not is_datetime64_any_dtype(dfs[0].index._meta_nonempty.inferred_type):
             raise TypeError(
                 "Must have a `DatetimeIndex` when using string offset "
                 "for `before` and `after`"
@@ -91,31 +163,178 @@ def map_overlap(func, df, before, after, *args, **kwargs):
         ):
             raise ValueError("before and after must be positive integers")
 
-    if "token" in kwargs:
-        func_name = kwargs.pop("token")
-        token = tokenize(df, before, after, *args, **kwargs)
+    name = kwargs.pop("token", None)
+    parent_meta = kwargs.pop("parent_meta", None)
+
+    assert callable(func)
+    if name is not None:
+        token = tokenize(meta, before, after, *args, **kwargs)
     else:
-        func_name = "overlap-" + funcname(func)
-        token = tokenize(func, df, before, after, *args, **kwargs)
+        name = "overlap-" + funcname(func)
+        token = tokenize(func, meta, before, after, *args, **kwargs)
+    name = f"{name}-{token}"
+
+    if align_dataframes:
+        args = _maybe_from_pandas(args)
+        try:
+            args = _maybe_align_partitions(args)
+        except ValueError as e:
+            raise ValueError(
+                f"{e}. If you don't want the partitions to be aligned, and are "
+                "calling `map_overlap` directly, pass `align_dataframes=False`."
+            ) from e
+
+    meta = _get_meta_map_partitions(args, dfs, func, kwargs, meta, parent_meta)
+
+    if all(isinstance(arg, Scalar) for arg in args):
+        layer = {
+            (name, 0): (
+                apply,
+                func,
+                (tuple, [(arg._name, 0) for arg in args]),
+                kwargs,
+            )
+        }
+        graph = HighLevelGraph.from_collections(name, layer, dependencies=args)
+        return Scalar(graph, name, meta)
+
+    args2 = []
+    dependencies = []
+
+    divisions = _get_divisions_map_partitions(
+        align_dataframes, transform_divisions, dfs, func, args, kwargs
+    )
+
+    def _handle_frame_argument(arg):
+        dsk = {}
 
-    if "meta" in kwargs:
-        meta = kwargs.pop("meta")
+        prevs_parts_dsk, prevs = _get_previous_partitions(arg, before)
+        dsk.update(prevs_parts_dsk)
+
+        nexts_parts_dsk, nexts = _get_nexts_partitions(arg, after)
+        dsk.update(nexts_parts_dsk)
+
+        name_a = "overlap-concat-" + tokenize(arg)
+        for i, (prev, current, next) in enumerate(
+            zip(prevs, arg.__dask_keys__(), nexts)
+        ):
+            key = (name_a, i)
+            dsk[key] = (_combined_parts, prev, current, next, before, after)
+
+        graph = HighLevelGraph.from_collections(name_a, dsk, dependencies=[arg])
+        return new_dd_object(graph, name_a, meta, divisions)
+
+    for arg in args:
+        if isinstance(arg, _Frame):
+            arg = _handle_frame_argument(arg)
+            args2.append(arg)
+            dependencies.append(arg)
+            continue
+        arg = normalize_arg(arg)
+        arg2, collections = unpack_collections(arg)
+        if collections:
+            args2.append(arg2)
+            dependencies.extend(collections)
+        else:
+            args2.append(arg)
+
+    kwargs3 = {}
+    simple = True
+    for k, v in kwargs.items():
+        v = normalize_arg(v)
+        v, collections = unpack_collections(v)
+        dependencies.extend(collections)
+        kwargs3[k] = v
+        if collections:
+            simple = False
+
+    if has_keyword(func, "partition_info"):
+        partition_info = {
+            (i,): {"number": i, "division": division}
+            for i, division in enumerate(divisions[:-1])
+        }
+
+        args2.insert(0, BlockwiseDepDict(partition_info))
+        orig_func = func
+
+        def func(partition_info, *args, **kwargs):
+            return orig_func(*args, **kwargs, partition_info=partition_info)
+
+    if enforce_metadata:
+        dsk = partitionwise_graph(
+            apply_and_enforce,
+            name,
+            func,
+            before,
+            after,
+            *args2,
+            dependencies=dependencies,
+            _func=overlap_chunk,
+            _meta=meta,
+            **kwargs3,
+        )
     else:
-        meta = _emulate(func, df, *args, **kwargs)
-    meta = make_meta(meta, index=df._meta.index, parent_meta=df._meta)
+        kwargs4 = kwargs if simple else kwargs3
+        dsk = partitionwise_graph(
+            overlap_chunk,
+            name,
+            func,
+            before,
+            after,
+            *args2,
+            **kwargs4,
+            dependencies=dependencies,
+        )
 
-    name = f"{func_name}-{token}"
-    name_a = "overlap-prepend-" + tokenize(df, before)
-    name_b = "overlap-append-" + tokenize(df, after)
-    df_name = df._name
+    graph = HighLevelGraph.from_collections(name, dsk, dependencies=dependencies)
+    return new_dd_object(graph, name, meta, divisions)
 
+
+def _get_nexts_partitions(df, after):
+    """
+    Helper to get the nexts partitions required for the overlap
+    """
     dsk = {}
+    df_name = df._name
 
     timedelta_partition_message = (
         "Partition size is less than specified window. "
         "Try using ``df.repartition`` to increase the partition size"
     )
 
+    name_b = "overlap-append-" + tokenize(df, after)
+    if after and isinstance(after, Integral):
+        nexts = []
+        for i in range(1, df.npartitions):
+            key = (name_b, i)
+            dsk[key] = (M.head, (df_name, i), after)
+            nexts.append(key)
+        nexts.append(None)
+    elif isinstance(after, datetime.timedelta):
+        # TODO: Do we have a use-case for this? Pandas doesn't allow negative rolling windows
+        deltas = pd.Series(df.divisions).diff().iloc[1:-1]
+        if (after > deltas).any():
+            raise ValueError(timedelta_partition_message)
+
+        nexts = []
+        for i in range(1, df.npartitions):
+            key = (name_b, i)
+            dsk[key] = (_head_timedelta, (df_name, i - 0), (df_name, i), after)
+            nexts.append(key)
+        nexts.append(None)
+    else:
+        nexts = [None] * df.npartitions
+    return dsk, nexts
+
+
+def _get_previous_partitions(df, before):
+    """
+    Helper to get the previous partitions required for the overlap
+    """
+    dsk = {}
+    df_name = df._name
+
+    name_a = "overlap-prepend-" + tokenize(df, before)
     if before and isinstance(before, Integral):
 
         prevs = [None]
@@ -171,44 +390,7 @@ def map_overlap(func, df, before, after, *args, **kwargs):
                 prevs.append(key)
     else:
         prevs = [None] * df.npartitions
-
-    if after and isinstance(after, Integral):
-        nexts = []
-        for i in range(1, df.npartitions):
-            key = (name_b, i)
-            dsk[key] = (M.head, (df_name, i), after)
-            nexts.append(key)
-        nexts.append(None)
-    elif isinstance(after, datetime.timedelta):
-        # TODO: Do we have a use-case for this? Pandas doesn't allow negative rolling windows
-        deltas = pd.Series(df.divisions).diff().iloc[1:-1]
-        if (after > deltas).any():
-            raise ValueError(timedelta_partition_message)
-
-        nexts = []
-        for i in range(1, df.npartitions):
-            key = (name_b, i)
-            dsk[key] = (_head_timedelta, (df_name, i - 0), (df_name, i), after)
-            nexts.append(key)
-        nexts.append(None)
-    else:
-        nexts = [None] * df.npartitions
-
-    for i, (prev, current, next) in enumerate(zip(prevs, df.__dask_keys__(), nexts)):
-        dsk[(name, i)] = (
-            overlap_chunk,
-            func,
-            prev,
-            current,
-            next,
-            before,
-            after,
-            args,
-            kwargs,
-        )
-
-    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[df])
-    return df._constructor(graph, name, meta, df.divisions)
+    return dsk, prevs
 
 
 def _head_timedelta(current, next_, after):


Test patch:
diff --git a/dask/dataframe/tests/test_rolling.py b/dask/dataframe/tests/test_rolling.py
--- a/dask/dataframe/tests/test_rolling.py
+++ b/dask/dataframe/tests/test_rolling.py
@@ -1,3 +1,5 @@
+import datetime
+
 import numpy as np
 import pandas as pd
 import pytest
@@ -24,16 +26,23 @@
     )
 )[:N]
 
+idx_constant_freq = (pd.date_range("2016-01-01", freq="1s", periods=100))[:N]
+
+ts_data = {
+    "a": np.random.randn(N).cumsum(),
+    "b": np.random.randint(100, size=(N,)),
+    "c": np.random.randint(100, size=(N,)),
+    "d": np.random.randint(100, size=(N,)),
+    "e": np.random.randint(100, size=(N,)),
+}
+
 ts = pd.DataFrame(
-    {
-        "a": np.random.randn(N).cumsum(),
-        "b": np.random.randint(100, size=(N,)),
-        "c": np.random.randint(100, size=(N,)),
-        "d": np.random.randint(100, size=(N,)),
-        "e": np.random.randint(100, size=(N,)),
-    },
+    ts_data,
     index=idx,
 )
+
+ts_constant_freq = pd.DataFrame(ts_data, index=idx_constant_freq)
+
 dts = dd.from_pandas(ts, 3)
 
 
@@ -58,15 +67,119 @@ def test_map_overlap(npartitions):
         assert_eq(res, sol)
 
 
-def test_map_overlap_names():
-    npartitions = 3
+@pytest.mark.parametrize("npartitions", [1, 4])
+@pytest.mark.parametrize("enforce_metadata", [True, False])
+@pytest.mark.parametrize("transform_divisions", [True, False])
+@pytest.mark.parametrize("align_dataframes", [True, False])
+@pytest.mark.parametrize(
+    "overlap_setup",
+    [
+        (df, 0, 3),
+        (df, 3, 0),
+        (df, 3, 3),
+        (df, 0, 0),
+        (
+            ts_constant_freq,
+            datetime.timedelta(seconds=3),
+            datetime.timedelta(seconds=3),
+        ),
+        (ts_constant_freq, datetime.timedelta(seconds=3), 0),
+    ],
+)
+def test_map_overlap_multiple_dataframes(
+    npartitions, enforce_metadata, transform_divisions, align_dataframes, overlap_setup
+):
+    dataframe, before, after = overlap_setup
+
+    ddf = dd.from_pandas(dataframe, npartitions)
+    ddf2 = dd.from_pandas(dataframe * 2, npartitions)
+
+    def get_shifted_sum_arg(overlap):
+        return (
+            overlap.seconds - 1 if isinstance(overlap, datetime.timedelta) else overlap
+        )
+
+    before_shifted_sum, after_shifted_sum = get_shifted_sum_arg(
+        before
+    ), get_shifted_sum_arg(after)
+
+    # DataFrame
+    res = ddf.map_overlap(
+        shifted_sum,
+        before,
+        after,
+        before_shifted_sum,
+        after_shifted_sum,
+        ddf2,
+        align_dataframes=align_dataframes,
+        transform_divisions=transform_divisions,
+        enforce_metadata=enforce_metadata,
+    )
+    sol = shifted_sum(dataframe, before_shifted_sum, after_shifted_sum, dataframe * 2)
+    assert_eq(res, sol)
+
+    # Series
+    res = ddf.b.map_overlap(
+        shifted_sum,
+        before,
+        after,
+        before_shifted_sum,
+        after_shifted_sum,
+        ddf2.b,
+        align_dataframes=align_dataframes,
+        transform_divisions=transform_divisions,
+        enforce_metadata=enforce_metadata,
+    )
+    sol = shifted_sum(
+        dataframe.b, before_shifted_sum, after_shifted_sum, dataframe.b * 2
+    )
+    assert_eq(res, sol)
+
+
+@pytest.mark.parametrize("npartitions", [1, 4])
+@pytest.mark.parametrize("enforce_metadata", [True, False])
+@pytest.mark.parametrize("transform_divisions", [True, False])
+@pytest.mark.parametrize("align_dataframes", [True, False])
+def test_map_overlap_names(
+    npartitions, enforce_metadata, transform_divisions, align_dataframes
+):
     ddf = dd.from_pandas(df, npartitions)
 
-    res = ddf.map_overlap(shifted_sum, 0, 3, 0, 3, c=2)
-    res2 = ddf.map_overlap(shifted_sum, 0, 3, 0, 3, c=2)
+    res = ddf.map_overlap(
+        shifted_sum,
+        0,
+        3,
+        0,
+        3,
+        c=2,
+        align_dataframes=align_dataframes,
+        transform_divisions=transform_divisions,
+        enforce_metadata=enforce_metadata,
+    )
+    res2 = ddf.map_overlap(
+        shifted_sum,
+        0,
+        3,
+        0,
+        3,
+        c=2,
+        align_dataframes=align_dataframes,
+        transform_divisions=transform_divisions,
+        enforce_metadata=enforce_metadata,
+    )
     assert set(res.dask) == set(res2.dask)
 
-    res3 = ddf.map_overlap(shifted_sum, 0, 3, 0, 3, c=3)
+    res3 = ddf.map_overlap(
+        shifted_sum,
+        0,
+        3,
+        0,
+        3,
+        c=3,
+        align_dataframes=align_dataframes,
+        transform_divisions=transform_divisions,
+        enforce_metadata=enforce_metadata,
+    )
     assert res3._name != res._name
     # Difference is just the final map
     diff = res3.dask.keys() - res.dask.keys()