Repository: dask/dask

Issue:
Implement DataFrame.coalesce 
I would love to use repartition of DataFrames without having to set index (#1380), yet even simple `coalesce` (merging some partitions together without a need to shuffle/rebalance partitions, like it is done in PySpark) would work in my use-case.

Can this be easily done in Dask? Please, advise where to look for the inspiration (I'm currently trying to get a sense of how `repartition` and `set_index` are done, though it is said that they are quite complex and still not complete for distributed use while `coalesce` should be much easier to implement).



Discussion:
Just to make sure I understand, you're looking to do this on dataframes for which divisions are not known, correct?  If you are interested in contributing a solution for dask.dataframe then you might take a look at `dask.bag.Bag.repartition` which solves a very similar problem.

Correct. Thank you, I will take a look there.


Resolving patch:
diff --git a/dask/bag/core.py b/dask/bag/core.py
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -1068,15 +1068,19 @@ def repartition(self, npartitions):
             raise NotImplementedError(
               "Repartition only supports going to fewer partitions\n"
               " old: %d  new: %d" % (self.npartitions, npartitions))
-        size = self.npartitions / npartitions
-        L = [int(i * size) for i in range(npartitions + 1)]
-        name = 'repartition-%d-%s' % (npartitions, self.name)
-        dsk = dict(((name, i), (list,
-                                (toolz.concat,
-                                 [(self.name, j) for j in range(L[i], L[i + 1])]
-                                 )))
-                   for i in range(npartitions))
-        return Bag(merge(self.dask, dsk), name, npartitions)
+        npartitions_ratio = self.npartitions / npartitions
+        new_partitions_boundaries = [int(old_partition_index * npartitions_ratio)
+                                        for old_partition_index in range(npartitions + 1)]
+        new_name = 'repartition-%d-%s' % (npartitions, tokenize(self))
+        dsk = {(new_name, new_partition_index):
+                (list,
+                 (toolz.concat,
+                  [(self.name, old_partition_index)
+                    for old_partition_index in range(
+                        new_partitions_boundaries[new_partition_index],
+                        new_partitions_boundaries[new_partition_index + 1])]))
+                for new_partition_index in range(npartitions)}
+        return Bag(dsk=merge(self.dask, dsk), name=new_name, npartitions=npartitions)
 
     def accumulate(self, binop, initial=no_default):
         """Repeatedly apply binary function to a sequence, accumulating results.
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -1370,9 +1370,9 @@ class DataFrame(_Frame):
     _partition_type = pd.DataFrame
     _token_prefix = 'dataframe-'
 
-    def __new__(cls, dask, name, columns, divisions):
+    def __new__(cls, dsk, name, columns, divisions):
         result = object.__new__(cls)
-        result.dask = dask
+        result.dask = dsk
         result._name = name
 
         result._pd, result._known_dtype = cls._build_pd(columns)
@@ -2630,12 +2630,23 @@ def _is_single_last_div(x):
 
 def repartition_npartitions(df, npartitions):
     """ Repartition dataframe to a smaller number of partitions """
-    npartitions = min(npartitions, df.npartitions)
-    k = int(math.ceil(df.npartitions / npartitions))
-    divisions = df.divisions[::k]
-    if len(divisions) <= npartitions:
-        divisions = divisions + (df.divisions[-1],)
-    return df.repartition(divisions=divisions)
+    npartitions_ratio = df.npartitions / npartitions
+    new_partitions_boundaries = [int(new_partition_index * npartitions_ratio)
+                                    for new_partition_index in range(npartitions + 1)]
+    new_name = 'repartition-%d-%s' % (npartitions, tokenize(df))
+    dsk = {(new_name, new_partition_index):
+            (pd.concat,
+             [(df._name, old_partition_index)
+                for old_partition_index in range(
+                    new_partitions_boundaries[new_partition_index],
+                    new_partitions_boundaries[new_partition_index + 1])])
+            for new_partition_index in range(npartitions)}
+    return DataFrame(
+        dsk=merge(df.dask, dsk),
+        name=new_name,
+        columns=df.columns,
+        divisions=[df.divisions[new_partition_index]
+                    for new_partition_index in new_partitions_boundaries])
 
 
 def repartition(df, divisions=None, force=False):


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -1155,22 +1155,23 @@ def test_repartition_on_pandas_dataframe():
 
 
 def test_repartition_npartitions():
-    df = pd.DataFrame({'x': [1, 2, 3, 4, 5, 6], 'y': list('abdabd')},
-                      index=[10, 20, 30, 40, 50, 60])
-    for n in [1, 2, 4, 5]:
-        for k in [1, 2, 4, 5]:
-            if k > n:
-                continue
-            a = dd.from_pandas(df, npartitions=n)
-            k = min(a.npartitions, k)
-
-            b = a.repartition(npartitions=k)
-            eq(a, b)
-            assert b.npartitions == k
-
-    a = dd.from_pandas(df, npartitions=1)
-    with pytest.raises(ValueError):
-        a.repartition(npartitions=5)
+    for use_index in (True, False):
+        df = pd.DataFrame({'x': [1, 2, 3, 4, 5, 6], 'y': list('abdabd')},
+                          index=[10, 20, 30, 40, 50, 60])
+        for n in [1, 2, 4, 5]:
+            for k in [1, 2, 4, 5]:
+                if k > n:
+                    continue
+                a = dd.from_pandas(df, npartitions=n, sort=use_index)
+                k = min(a.npartitions, k)
+
+                b = a.repartition(npartitions=k)
+                eq(a, b)
+                assert b.npartitions == k
+
+        a = dd.from_pandas(df, npartitions=1)
+        with pytest.raises(ValueError):
+            a.repartition(npartitions=5)
 
 
 def test_embarrassingly_parallel_operations():