Repository: dask/dask

Issue:
Feature request: Bag Repartition "partition_size" keyword argument
Hi there,

Love this project, use it extensively in my day to day.

This is a feature request for the [```Bag.repartition```](https://docs.dask.org/en/latest/bag-api.html#dask.bag.Bag.repartition) function. The [```Dataframe.repartition```](https://docs.dask.org/en/latest/dataframe-api.html#dask.dataframe.DataFrame.repartition) function has a nice ```partition_size``` keyword which repartitions the dataframe by a specified size. It would be really nice to have the same argument for ```Bag.repartition```. Right now I have to estimate the number of partitions I'll need to get my data down to 100mb per partition prior to calling repartition, so this would be a quality of life improvement for me.

Cheers and thank you for this great project!


Discussion:
Hi @ankur-gos thank you for the suggestion! This seems like a reasonable suggestion and one that could be supported. Would you be interested in working on this?

Resolving patch:
diff --git a/dask/bag/core.py b/dask/bag/core.py
--- a/dask/bag/core.py
+++ b/dask/bag/core.py
@@ -41,6 +41,7 @@
 from ..bytes import open_files
 from ..context import globalmethod
 from ..core import quote, istask, get_dependencies, reverse_dict, flatten
+from ..sizeof import sizeof
 from ..delayed import Delayed, unpack_collections
 from ..highlevelgraph import HighLevelGraph
 from ..multiprocessing import get as mpget
@@ -56,6 +57,8 @@
     ensure_bytes,
     ensure_unicode,
     key_split,
+    parse_bytes,
+    iter_chunks,
 )
 from . import chunk
 
@@ -1566,62 +1569,39 @@ def to_delayed(self, optimize_graph=True):
             dsk = self.__dask_optimize__(dsk, keys)
         return [Delayed(k, dsk) for k in keys]
 
-    def repartition(self, npartitions):
-        """ Changes the number of partitions of the bag.
+    def repartition(self, npartitions=None, partition_size=None):
+        """ Repartition Bag across new divisions.
 
-        This can be used to reduce or increase the number of partitions
-        of the bag.
+        Parameters
+        ----------
+        npartitions : int, optional
+            Number of partitions of output.
+        partition_size : int or string, optional
+            Max number of bytes of memory for each partition. Use numbers or
+            strings like 5MB.
+
+            .. warning::
+
+               This keyword argument triggers computation to determine
+               the memory size of each partition, which may be expensive.
+
+        Notes
+        -----
+        Exactly one of ``npartitions`` or ``partition_size`` should be specified.
+        A ``ValueError`` will be raised when that is not the case.
 
         Examples
         --------
         >>> b.repartition(5)  # set to have 5 partitions  # doctest: +SKIP
         """
-        new_name = "repartition-%d-%s" % (npartitions, tokenize(self, npartitions))
-        if npartitions == self.npartitions:
-            return self
-        elif npartitions < self.npartitions:
-            ratio = self.npartitions / npartitions
-            new_partitions_boundaries = [
-                int(old_partition_index * ratio)
-                for old_partition_index in range(npartitions + 1)
-            ]
-
-            dsk = {}
-            for new_partition_index in range(npartitions):
-                value = (
-                    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],
-                            )
-                        ],
-                    ),
-                )
-                dsk[new_name, new_partition_index] = value
-        else:  # npartitions > self.npartitions
-            ratio = npartitions / self.npartitions
-            split_name = "split-%s" % tokenize(self, npartitions)
-            dsk = {}
-            last = 0
-            j = 0
-            for i in range(self.npartitions):
-                new = last + ratio
-                if i == self.npartitions - 1:
-                    k = npartitions - j
-                else:
-                    k = int(new - last)
-                dsk[(split_name, i)] = (split, (self.name, i), k)
-                for jj in range(k):
-                    dsk[(new_name, j)] = (operator.getitem, (split_name, i), jj)
-                    j += 1
-                last = new
-
-        graph = HighLevelGraph.from_collections(new_name, dsk, dependencies=[self])
-        return Bag(graph, name=new_name, npartitions=npartitions)
+        if sum([partition_size is not None, npartitions is not None]) != 1:
+            raise ValueError(
+                "Please provide exactly one ``npartitions`` or ``partition_size`` keyword arguments"
+            )
+        if npartitions is not None:
+            return repartition_npartitions(self, npartitions)
+        elif partition_size is not None:
+            return repartition_size(self, partition_size)
 
     def accumulate(self, binop, initial=no_default):
         """ Repeatedly apply binary function to a sequence, accumulating results.
@@ -2498,3 +2478,129 @@ def to_dataframe(seq, columns, dtypes):
         seq = list(seq)
     res = pd.DataFrame(seq, columns=list(columns))
     return res.astype(dtypes, copy=False)
+
+
+def repartition_npartitions(bag, npartitions):
+    """ Changes the number of partitions of the bag.
+
+    This can be used to reduce or increase the number of partitions
+    of the bag.
+    """
+    if npartitions == bag.npartitions:
+        return bag
+
+    new_name = "repartition-%d-%s" % (npartitions, tokenize(bag, npartitions))
+    if bag.npartitions > npartitions:
+        ratio = bag.npartitions / npartitions
+        new_partitions_boundaries = [
+            int(old_partition_index * ratio)
+            for old_partition_index in range(npartitions + 1)
+        ]
+        return _repartition_from_boundaries(bag, new_partitions_boundaries, new_name)
+    else:  # npartitions > bag.npartitions
+        div, mod = divmod(npartitions, bag.npartitions)
+        nsplits = [div] * bag.npartitions
+        nsplits[-1] += mod
+        return _split_partitions(bag, nsplits, new_name)
+
+
+def total_mem_usage(partition):
+    from copy import deepcopy
+    from itertools import chain
+
+    # if repartition is called multiple times prior to calling compute(), the partitions
+    # will be itertools.chain objects. Copy the object to avoid consuming the iterable.
+    if isinstance(partition, chain):
+        partition = reify(deepcopy(partition))
+    return sizeof(partition)
+
+
+def repartition_size(bag, size):
+    """
+    Repartition bag so that new partitions have approximately `size` memory usage each
+    """
+    if isinstance(size, str):
+        size = parse_bytes(size)
+    size = int(size)
+    mem_usages = bag.map_partitions(total_mem_usage).compute()
+
+    # 1. split each partition that is larger than partition size
+    nsplits = [1 + mem_usage // size for mem_usage in mem_usages]
+    if any((nsplit > 1 for nsplit in nsplits)):
+        split_name = "repartition-split-{}".format(tokenize(bag, size))
+        bag = _split_partitions(bag, nsplits, split_name)
+        # update mem_usages to account for the split partitions
+        split_mem_usages = []
+        for n, usage in zip(nsplits, mem_usages):
+            split_mem_usages.extend([usage / n] * n)
+        mem_usages = split_mem_usages
+
+    # 2. now that all partitions are less than size, concat them up to size
+    assert all((mem_usage <= size for mem_usage in mem_usages))
+    new_npartitions = list(map(len, iter_chunks(mem_usages, size)))
+    new_partitions_boundaries = accumulate(operator.add, new_npartitions)
+    new_name = "repartition-{}".format(tokenize(bag, size))
+    return _repartition_from_boundaries(bag, new_partitions_boundaries, new_name)
+
+
+def _split_partitions(bag, nsplits, new_name):
+    """Split a Dask bag into new partitions
+
+    Parameters
+    ----------
+    bag: Dask bag
+    nsplits: List[int]
+        Number of target bags for each partition
+        The length of nsplits should be the same as bag.npartitions
+    new_name: str
+
+    See Also
+    --------
+    repartition_npartitions
+    repartition_size
+    """
+    if len(nsplits) != bag.npartitions:
+        raise ValueError("nsplits should have len={}".format(bag.npartitions))
+    dsk = {}
+    split_name = "split-{}".format(tokenize(bag, nsplits))
+    j = 0
+    for i, k in enumerate(nsplits):
+        if k == 1:
+            dsk[new_name, j] = (bag.name, i)
+            j += 1
+        else:
+            dsk[split_name, i] = (split, (bag.name, i), k)
+            for jj in range(k):
+                dsk[new_name, j] = (operator.getitem, (split_name, i), jj)
+                j += 1
+
+    graph = HighLevelGraph.from_collections(new_name, dsk, dependencies=[bag])
+    return Bag(graph, name=new_name, npartitions=sum(nsplits))
+
+
+def _repartition_from_boundaries(bag, new_partitions_boundaries, new_name):
+    if not isinstance(new_partitions_boundaries, list):
+        new_partitions_boundaries = list(new_partitions_boundaries)
+    if new_partitions_boundaries[0] > 0:
+        new_partitions_boundaries.insert(0, 0)
+    if new_partitions_boundaries[-1] < bag.npartitions:
+        new_partitions_boundaries.append(bag.npartitions)
+    num_new_partitions = len(new_partitions_boundaries) - 1
+    dsk = {}
+    for new_partition_index in range(num_new_partitions):
+        value = (
+            list,
+            (
+                toolz.concat,
+                [
+                    (bag.name, old_partition_index)
+                    for old_partition_index in range(
+                        new_partitions_boundaries[new_partition_index],
+                        new_partitions_boundaries[new_partition_index + 1],
+                    )
+                ],
+            ),
+        )
+        dsk[new_name, new_partition_index] = value
+    graph = HighLevelGraph.from_collections(new_name, dsk, dependencies=[bag])
+    return Bag(graph, name=new_name, npartitions=num_new_partitions)
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -40,6 +40,7 @@
     OperatorMethodMixin,
     is_arraylike,
     typename,
+    iter_chunks,
 )
 from ..array.core import Array, normalize_arg
 from ..array.utils import zeros_like_safe
@@ -5915,34 +5916,6 @@ def total_mem_usage(df, index=True, deep=False):
     return mem_usage
 
 
-def iter_chunks(sizes, max_size):
-    """Split sizes into chunks of total max_size each
-
-    Parameters
-    ----------
-    sizes : iterable of numbers
-        The sizes to be chunked
-    max_size : number
-        Maximum total size per chunk.
-        It must be greater or equal than each size in sizes
-    """
-    chunk, chunk_sum = [], 0
-    iter_sizes = iter(sizes)
-    size = next(iter_sizes, None)
-    while size is not None:
-        assert size <= max_size
-        if chunk_sum + size <= max_size:
-            chunk.append(size)
-            chunk_sum += size
-            size = next(iter_sizes, None)
-        else:
-            assert chunk
-            yield chunk
-            chunk, chunk_sum = [], 0
-    if chunk:
-        yield chunk
-
-
 def repartition_npartitions(df, npartitions):
     """ Repartition dataframe to a smaller number of partitions """
     new_name = "repartition-%d-%s" % (npartitions, tokenize(df))
diff --git a/dask/utils.py b/dask/utils.py
--- a/dask/utils.py
+++ b/dask/utils.py
@@ -1384,3 +1384,31 @@ def ndimlist(seq):
         return 1
     else:
         return 1 + ndimlist(seq[0])
+
+
+def iter_chunks(sizes, max_size):
+    """Split sizes into chunks of total max_size each
+
+    Parameters
+    ----------
+    sizes : iterable of numbers
+        The sizes to be chunked
+    max_size : number
+        Maximum total size per chunk.
+        It must be greater or equal than each size in sizes
+    """
+    chunk, chunk_sum = [], 0
+    iter_sizes = iter(sizes)
+    size = next(iter_sizes, None)
+    while size is not None:
+        assert size <= max_size
+        if chunk_sum + size <= max_size:
+            chunk.append(size)
+            chunk_sum += size
+            size = next(iter_sizes, None)
+        else:
+            assert chunk
+            yield chunk
+            chunk, chunk_sum = [], 0
+    if chunk:
+        yield chunk


Test patch:
diff --git a/dask/bag/tests/test_bag.py b/dask/bag/tests/test_bag.py
--- a/dask/bag/tests/test_bag.py
+++ b/dask/bag/tests/test_bag.py
@@ -26,6 +26,7 @@
     inline_singleton_lists,
     optimize,
     from_delayed,
+    total_mem_usage,
 )
 from dask.bag.utils import assert_eq
 from dask.delayed import Delayed
@@ -647,7 +648,7 @@ def test_read_text_large():
 def test_read_text_encoding():
     with tmpfile() as fn:
         with open(fn, "wb") as f:
-            f.write((u"你好！" + os.linesep).encode("gb18030") * 100)
+            f.write(("你好！" + os.linesep).encode("gb18030") * 100)
         b = db.read_text(fn, blocksize=100, encoding="gb18030")
         c = db.read_text(fn, encoding="gb18030")
         assert len(b.dask) > 5
@@ -683,11 +684,11 @@ def test_from_s3():
     pytest.importorskip("s3fs")
 
     five_tips = (
-        u"total_bill,tip,sex,smoker,day,time,size\n",
-        u"16.99,1.01,Female,No,Sun,Dinner,2\n",
-        u"10.34,1.66,Male,No,Sun,Dinner,3\n",
-        u"21.01,3.5,Male,No,Sun,Dinner,3\n",
-        u"23.68,3.31,Male,No,Sun,Dinner,2\n",
+        "total_bill,tip,sex,smoker,day,time,size\n",
+        "16.99,1.01,Female,No,Sun,Dinner,2\n",
+        "10.34,1.66,Male,No,Sun,Dinner,3\n",
+        "21.01,3.5,Male,No,Sun,Dinner,3\n",
+        "23.68,3.31,Male,No,Sun,Dinner,2\n",
     )
 
     # test compressed data
@@ -947,7 +948,7 @@ def test_to_textfiles_name_function_warn():
 
 
 def test_to_textfiles_encoding():
-    b = db.from_sequence([u"汽车", u"苹果", u"天气"], npartitions=2)
+    b = db.from_sequence(["汽车", "苹果", "天气"], npartitions=2)
     for ext, myopen in ext_open:
         with tmpdir() as dir:
             c = b.to_textfiles(
@@ -960,7 +961,7 @@ def test_to_textfiles_encoding():
             text = f.read()
             if hasattr(text, "decode"):
                 text = text.decode("gb18030")
-            assert u"天气" in text
+            assert "天气" in text
             f.close()
 
 
@@ -1011,12 +1012,12 @@ def test_string_namespace():
 
 
 def test_string_namespace_with_unicode():
-    b = db.from_sequence([u"Alice Smith", u"Bob Jones", "Charlie Smith"], npartitions=2)
+    b = db.from_sequence(["Alice Smith", "Bob Jones", "Charlie Smith"], npartitions=2)
     assert list(b.str.lower()) == ["alice smith", "bob jones", "charlie smith"]
 
 
 def test_str_empty_split():
-    b = db.from_sequence([u"Alice Smith", u"Bob Jones", "Charlie Smith"], npartitions=2)
+    b = db.from_sequence(["Alice Smith", "Bob Jones", "Charlie Smith"], npartitions=2)
     assert list(b.str.split()) == [
         ["Alice", "Smith"],
         ["Bob", "Jones"],
@@ -1065,7 +1066,7 @@ def test_bag_class_extend():
 
 
 def test_gh715():
-    bin_data = u"\u20ac".encode("utf-8")
+    bin_data = "\u20ac".encode("utf-8")
     with tmpfile() as fn:
         with open(fn, "wb") as f:
             f.write(bin_data)
@@ -1164,16 +1165,52 @@ def test_zip(npartitions, hi=1000):
 
 @pytest.mark.parametrize("nin", [1, 2, 7, 11, 23])
 @pytest.mark.parametrize("nout", [1, 2, 5, 12, 23])
-def test_repartition(nin, nout):
+def test_repartition_npartitions(nin, nout):
     b = db.from_sequence(range(100), npartitions=nin)
     c = b.repartition(npartitions=nout)
-
     assert c.npartitions == nout
     assert_eq(b, c)
     results = dask.get(c.dask, c.__dask_keys__())
     assert all(results)
 
 
+@pytest.mark.parametrize(
+    "nin, nout",
+    [(1, 1), (2, 1), (5, 1), (1, 2), (2, 2), (5, 2), (1, 5), (2, 5), (5, 5),],
+)
+def test_repartition_partition_size(nin, nout):
+    b = db.from_sequence(range(1, 100), npartitions=nin)
+    total_mem = sum(b.map_partitions(total_mem_usage).compute())
+    c = b.repartition(partition_size=(total_mem // nout))
+    assert c.npartitions >= nout
+    assert_eq(b, c)
+
+
+def test_multiple_repartition_partition_size():
+    b = db.from_sequence(range(1, 100), npartitions=1)
+    total_mem = sum(b.map_partitions(total_mem_usage).compute())
+
+    c = b.repartition(partition_size=(total_mem // 2))
+    assert c.npartitions >= 2
+    assert_eq(b, c)
+
+    d = c.repartition(partition_size=(total_mem // 5))
+    assert d.npartitions >= 5
+    assert_eq(c, d)
+
+
+def test_repartition_partition_size_complex_dtypes():
+    np = pytest.importorskip("numpy")
+
+    b = db.from_sequence([np.array(range(100)) for _ in range(4)], npartitions=1)
+    total_mem = sum(b.map_partitions(total_mem_usage).compute())
+
+    new_partition_size = total_mem // 4
+    c = b.repartition(partition_size=new_partition_size)
+    assert c.npartitions >= 4
+    assert_eq(b, c)
+
+
 def test_repartition_names():
     b = db.from_sequence(range(100), npartitions=5)
     c = b.repartition(2)
@@ -1187,6 +1224,12 @@ def test_repartition_names():
     assert b is c
 
 
+def test_repartition_input_errors():
+    with pytest.raises(ValueError):
+        bag = db.from_sequence(range(10))
+        bag.repartition(npartitions=5, partition_size="5MiB")
+
+
 def test_accumulate():
     parts = [[1, 2, 3], [4, 5], [], [6, 7]]
     dsk = dict((("test", i), p) for (i, p) in enumerate(parts))
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
@@ -22,7 +22,6 @@
     _concat,
     Scalar,
     has_parallel_type,
-    iter_chunks,
     total_mem_usage,
     is_broadcastable,
 )
@@ -1918,21 +1917,6 @@ def test_repartition_partition_size_arg():
     assert b.npartitions == 1
 
 
-def test_iter_chunks():
-    sizes = [14, 8, 5, 9, 7, 9, 1, 19, 8, 19]
-    assert list(iter_chunks(sizes, 19)) == [
-        [14],
-        [8, 5],
-        [9, 7],
-        [9, 1],
-        [19],
-        [8],
-        [19],
-    ]
-    assert list(iter_chunks(sizes, 28)) == [[14, 8, 5], [9, 7, 9, 1], [19, 8], [19]]
-    assert list(iter_chunks(sizes, 67)) == [[14, 8, 5, 9, 7, 9, 1], [19, 8, 19]]
-
-
 def test_repartition_npartitions_same_limits():
     df = pd.DataFrame(
         {"x": [1, 2, 3]},
diff --git a/dask/tests/test_utils.py b/dask/tests/test_utils.py
--- a/dask/tests/test_utils.py
+++ b/dask/tests/test_utils.py
@@ -28,6 +28,7 @@
     parse_timedelta,
     parse_bytes,
     is_arraylike,
+    iter_chunks,
 )
 from dask.utils_test import inc
 from dask.highlevelgraph import HighLevelGraph
@@ -566,3 +567,18 @@ def test_is_arraylike():
     assert is_arraylike(np.empty(())) is True
     assert is_arraylike(np.empty((0,))) is True
     assert is_arraylike(np.empty((0, 0))) is True
+
+
+def test_iter_chunks():
+    sizes = [14, 8, 5, 9, 7, 9, 1, 19, 8, 19]
+    assert list(iter_chunks(sizes, 19)) == [
+        [14],
+        [8, 5],
+        [9, 7],
+        [9, 1],
+        [19],
+        [8],
+        [19],
+    ]
+    assert list(iter_chunks(sizes, 28)) == [[14, 8, 5], [9, 7, 9, 1], [19, 8], [19]]
+    assert list(iter_chunks(sizes, 67)) == [[14, 8, 5, 9, 7, 9, 1], [19, 8, 19]]