Repository: dask/dask
Pull request: Avoid shuffle when setting a presorted index with overlapping divisions
- [x] Tests added / passed
- [x] Passes `black dask` / `flake8 dask`

Resolves https://github.com/dask/dask/pull/5076, closes https://github.com/dask/dask/issues/4860

Problem statement:
Repartition randomly dropping some rows
**Bug report**

When I called `set_index(..., sorted=True)`, I got a warning saying partition indices have overlap, which is fine. From here, after calling `repartition`, I lost a row. This must be a bug in the `repartition` code since I don't see why `repartition` would make me lose rows.

```python3
>>> import pandas as pd
>>> import dask.dataframe as dd
>>> A = pd.DataFrame({'key': [1,2,3,4,4,5,6,7], 'value': list('abcd'*2)})
>>> a = dd.from_pandas(A, npartitions=2)
>>> a = a.set_index('key', sorted=True)
/Users/ctj/Documents/dask/dask/dataframe/shuffle.py:600: UserWarning: Partition indices have overlap.
  warnings.warn("Partition indices have overlap.")
>>> a.compute()
    value
key      
1       a
2       b
3       c
4       d
4       a
5       b
6       c
7       d
>>> a = a.repartition(divisions=a.divisions)
>>> a.compute()
    value
key      
1       a
2       b
3       c
4       a
5       b
6       c
7       d
```

Hints from issue discussion:
I found a workaround to this phenomenon. Instead of using `repartition`, I just manually constructed the graph using `methods.boundary_slice`, though it seems like overkill to me.
It's possible that `sorted=True` is assuming that the provided index is
strictly increasing. Are you interested in investigating?

On Thu, May 30, 2019 at 2:46 PM Cody Johnson <notifications@github.com>
wrote:

> I found a workaround to this phenomenon. Instead of using repartition, I
> just manually constructed the graph using methods.boundary_slice, though
> it seems like overkill to me.
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/4860?email_source=notifications&email_token=AAKAOIVQHRK3A64QUZUCAG3PYAVK5A5CNFSM4HRBGJV2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGODWTKLHI#issuecomment-497460637>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AAKAOIQZRQ7ASJEEL7KZ5HTPYAVK5ANCNFSM4HRBGJVQ>
> .
>

I’m currently working on implementing `merge_asof` but am certainly interested after I finish up! Shouldn’t be too much longer…
Does the Dask spec say that all the indices of partition i lie in the interval [ddf.divisions[i], ddf.divisions[i+1]) or is the upper bound allowed to be inclusive? (Except for the last partition)
> all the indices of partition i lie in the interval [ddf.divisions[i], ddf.divisions[i+1])

Yes

> is the upper bound allowed to be inclusive? (Except for the last partition)

No

This seems like a likely culprit in this case to me.

It looks like dask dataframe identified this and raised a not very informative warning

```python-traceback
/Users/ctj/Documents/dask/dask/dataframe/shuffle.py:600: UserWarning: Partition indices have overlap.
```
In that case, how would Dask handle when a ton of the rows have the same index? Is ddf.divisions necessarily a _strictly_ increasing sequence? (Except the last one)
> Is ddf.divisions necessarily a *strictly* increasing sequence

Yes

> how would Dask handle when a ton of the rows have the same index?

Poorly

On Mon, Jul 1, 2019 at 3:27 PM Cody Johnson <notifications@github.com>
wrote:

> In that case, how would Dask handle when a ton of the rows have the same
> index? Is ddf.divisions necessarily a *strictly* increasing sequence?
> (Except the last one)
>
> —
> You are receiving this because you commented.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/4860?email_source=notifications&email_token=AACKZTEZNKN2WU4RYMWGCMLP5IH4NA5CNFSM4HRBGJV2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGODY6JUCA#issuecomment-507288072>,
> or mute the thread
> <https://github.com/notifications/unsubscribe-auth/AACKZTB7STM5I4OAPNYRW6LP5IH4NANCNFSM4HRBGJVQ>
> .
>

Alright, this is very helpful. Thanks Matthew!
What's the next step here? Should that warning be elevated to an error (disclaimer: I haven't really dived into the context to understand the issue).
I’m currently using what Matthew said to fix it. Should be easy.

Warning is a sufficient label because it doesn’t actually mess anything up, only subsequently.

Commit message: flake8

Patch to review:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
index 26048654b0d..c320ad78cb3 100644
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -6294,8 +6294,6 @@ def prefix_reduction(f, ddf, identity, **kwargs):
     ddf : dd.DataFrame
     identity : pd.DataFrame
         an identity element of f, that is f(identity, df) = f(df, identity) = df
-    kwargs : ??
-        keyword arguments of f ??
     """
     dsk = dict()
     name = "prefix_reduction-" + tokenize(f, ddf, identity, **kwargs)
diff --git a/dask/dataframe/io/io.py b/dask/dataframe/io/io.py
index 170c4f84539..d8a0cc00340 100644
--- a/dask/dataframe/io/io.py
+++ b/dask/dataframe/io/io.py
@@ -583,8 +583,7 @@ def from_delayed(
     if divisions == "sorted":
         from ..shuffle import compute_divisions
 
-        divisions = compute_divisions(df)
-        df.divisions = divisions
+        df = compute_divisions(df)
 
     return df
 
diff --git a/dask/dataframe/multi.py b/dask/dataframe/multi.py
index 1270542c62a..e2623dfa742 100644
--- a/dask/dataframe/multi.py
+++ b/dask/dataframe/multi.py
@@ -580,44 +580,6 @@ def merge(
 ###############################################################
 
 
-def fix_overlap(ddf):
-    """ Ensures that ddf.divisions are all distinct and that the upper bound on
-    each partition is exclusive
-    """
-    if not ddf.known_divisions:
-        raise ValueError("Can only fix overlap when divisions are known")
-
-    def body(df, index):
-        return df.drop(index, inplace=True) if index in df else df
-
-    def overlap(df, index):
-        return df.loc[[index]] if index in df else None
-
-    dsk = dict()
-    name = "fix-overlap-" + tokenize(ddf)
-
-    n = len(ddf.divisions) - 1
-    divisions = []
-    for i in range(n):
-        if i > 0 and ddf.divisions[i - 1] == ddf.divisions[i]:
-            frames = dsk[(name, len(divisions) - 1)][1]
-        else:
-            frames = []
-            if i > 0:
-                frames.append((overlap, (ddf._name, i - 1), ddf.divisions[i]))
-            divisions.append(ddf.divisions[i])
-            dsk[(name, len(divisions) - 1)] = (pd.concat, frames)
-
-        if i == n - 1 or ddf.divisions[i + 1] == ddf.divisions[i]:
-            frames.append((ddf._name, i))
-        else:
-            frames.append((body, (ddf._name, i), ddf.divisions[i + 1]))
-    divisions.append(ddf.divisions[-1])
-
-    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[ddf])
-    return new_dd_object(graph, name, ddf._meta, divisions)
-
-
 def most_recent_tail(left, right):
     if len(right.index) == 0:
         return left
@@ -767,9 +729,6 @@ def _concat_compat(frames, left, right):
 
 
 def merge_asof_indexed(left, right, **kwargs):
-    left = fix_overlap(left)
-    right = fix_overlap(right)
-
     dsk = dict()
     name = "asof-join-indexed-" + tokenize(left, right, **kwargs)
     meta = pd.merge_asof(left._meta_nonempty, right._meta_nonempty, **kwargs)
diff --git a/dask/dataframe/shuffle.py b/dask/dataframe/shuffle.py
index bb60e6c7b4b..853791e00f9 100644
--- a/dask/dataframe/shuffle.py
+++ b/dask/dataframe/shuffle.py
@@ -5,7 +5,6 @@
 from operator import getitem
 import uuid
 import tempfile
-import warnings
 
 import tlz as toolz
 import numpy as np
@@ -116,8 +115,6 @@ def set_index(
         ):
             divisions = mins + [maxes[-1]]
             result = set_sorted_index(df, index, drop=drop, divisions=divisions)
-            # There are cases where this still may not be sorted
-            # so sort_index to be sure. https://github.com/dask/dask/issues/2288
             return result.map_partitions(M.sort_index)
 
     return set_partition(
@@ -776,6 +773,49 @@ def set_index_post_series(df, index_name, drop, column_dtype):
     return df2
 
 
+def drop_overlap(df, index):
+    return df.drop(index) if index in df.index else df
+
+
+def get_overlap(df, index):
+    return df.loc[[index]] if index in df.index else None
+
+
+def fix_overlap(ddf, overlap):
+    """ Ensures that the upper bound on each partition of ddf is exclusive """
+    name = "fix-overlap-" + tokenize(ddf, overlap)
+    n = len(ddf.divisions) - 1
+    dsk = {(name, i): (ddf._name, i) for i in range(n)}
+
+    for i in overlap:
+        frames = [(get_overlap, (ddf._name, i - 1), ddf.divisions[i]), (ddf._name, i)]
+        dsk[(name, i)] = (pd.concat, frames)
+        dsk[(name, i - 1)] = (drop_overlap, dsk[(name, i - 1)], ddf.divisions[i])
+
+    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[ddf])
+    return new_dd_object(graph, name, ddf._meta, ddf.divisions)
+
+
+def compute_divisions(df, **kwargs):
+    mins = df.index.map_partitions(M.min, meta=df.index)
+    maxes = df.index.map_partitions(M.max, meta=df.index)
+    mins, maxes = compute(mins, maxes, **kwargs)
+
+    if (
+        sorted(mins) != list(mins)
+        or sorted(maxes) != list(maxes)
+        or any(a > b for a, b in zip(mins, maxes))
+    ):
+        raise ValueError(
+            "Partitions must be sorted ascending with the index", mins, maxes
+        )
+
+    df.divisions = tuple(mins) + (list(maxes)[-1],)
+
+    overlap = [i for i in range(1, len(mins)) if mins[i] >= maxes[i - 1]]
+    return fix_overlap(df, overlap) if overlap else df
+
+
 def set_sorted_index(df, index, drop=True, divisions=None, **kwargs):
     if not isinstance(index, Series):
         meta = df._meta.set_index(index, drop=drop)
@@ -785,7 +825,7 @@ def set_sorted_index(df, index, drop=True, divisions=None, **kwargs):
     result = map_partitions(M.set_index, df, index, drop=drop, meta=meta)
 
     if not divisions:
-        divisions = compute_divisions(result, **kwargs)
+        return compute_divisions(result, **kwargs)
     elif len(divisions) != len(df.divisions):
         msg = (
             "When doing `df.set_index(col, sorted=True, divisions=...)`, "
@@ -800,24 +840,3 @@ def set_sorted_index(df, index, drop=True, divisions=None, **kwargs):
 
     result.divisions = tuple(divisions)
     return result
-
-
-def compute_divisions(df, **kwargs):
-    mins = df.index.map_partitions(M.min, meta=df.index)
-    maxes = df.index.map_partitions(M.max, meta=df.index)
-    mins, maxes = compute(mins, maxes, **kwargs)
-
-    if (
-        sorted(mins) != list(mins)
-        or sorted(maxes) != list(maxes)
-        or any(a > b for a, b in zip(mins, maxes))
-    ):
-        raise ValueError(
-            "Partitions must be sorted ascending with the index", mins, maxes
-        )
-
-    if any(a >= b for a, b in zip(maxes[:1], mins[1:])):
-        warnings.warn("Partition indices have overlap.")
-
-    divisions = tuple(mins) + (list(maxes)[-1],)
-    return divisions
diff --git a/dask/dataframe/tests/test_shuffle.py b/dask/dataframe/tests/test_shuffle.py
index d27adb8c37b..0c32b3f09f8 100644
--- a/dask/dataframe/tests/test_shuffle.py
+++ b/dask/dataframe/tests/test_shuffle.py
@@ -718,9 +718,6 @@ def test_set_index_sorted_true():
     with pytest.raises(ValueError):
         a.set_index(a.z, sorted=True)
 
-    with pytest.warns(UserWarning):
-        a.set_index(a.y, sorted=True)
-
 
 def test_set_index_sorted_single_partition():
     df = pd.DataFrame({"x": [1, 2, 3, 4], "y": [1, 0, 1, 0]})
@@ -805,18 +802,11 @@ def test_compute_divisions():
     a = dd.from_pandas(df, 2, sort=False)
     assert not a.known_divisions
 
-    divisions = compute_divisions(a)
-    b = copy(a)
-    b.divisions = divisions
+    b = compute_divisions(copy(a))
 
     assert_eq(a, b, check_divisions=False)
     assert b.known_divisions
 
-    c = dd.from_pandas(df.set_index("y"), 2, sort=False)
-    # Partitions overlap warning
-    with pytest.warns(UserWarning):
-        compute_divisions(c)
-
 
 def test_empty_partitions():
     # See https://github.com/dask/dask/issues/2408
@@ -973,3 +963,11 @@ def generate_raw_partd_file(compression):
     compressed_data = generate_raw_partd_file(compression="BZ2")
 
     assert len(uncompressed_data) > len(compressed_data)
+
+
+def test_set_index_overlap():
+    A = pd.DataFrame({"key": [1, 2, 3, 4, 4, 5, 6, 7], "value": list("abcd" * 2)})
+    a = dd.from_pandas(A, npartitions=2)
+    a = a.set_index("key", sorted=True)
+    b = a.repartition(divisions=a.divisions)
+    assert_eq(a, b)


Code review comments:
diff_hunk: @@ -776,6 +773,49 @@ def set_index_post_series(df, index_name, drop, column_dtype):
     return df2
 
 
+def drop_overlap(df, index):
+    return df.drop(index) if index in df.index else df
+
+
+def get_overlap(df, index):
+    return df.loc[[index]] if index in df.index else None
+
+
+def fix_overlap(ddf, overlap):
+    """ Ensures that the upper bound on each partition of ddf is exclusive """
+    name = "fix-overlap-" + tokenize(ddf, overlap)
+    n = len(ddf.divisions) - 1
+    dsk = {(name, i): (ddf._name, i) for i in range(n)}
+
+    for i in overlap:
+        frames = [(get_overlap, (ddf._name, i - 1), ddf.divisions[i]), (ddf._name, i)]
+        dsk[(name, i)] = (pd.concat, frames)
+        dsk[(name, i - 1)] = (drop_overlap, dsk[(name, i - 1)], ddf.divisions[i])
+
+    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[ddf])
+    return new_dd_object(graph, name, ddf._meta, ddf.divisions)
+
+
+def compute_divisions(df, **kwargs):
path: dask/dataframe/shuffle.py
text: @user1:
Since this both computes and sets divisions (rather than just computes them), perhaps rename to `with_computed_divisions` or `compute_and_set_divisions`.