Repository: dask/dask

Issue:
Per-partition information in dask.dataframe.DataFrame.memory_usage
We currently provide a per-column information in a `pandas.Series` when calling `dask.dataframe.DataFrame.memory_usage` Series. It would also be nice to to have that information on a per-partition-per-column basis. This would allow easier analysis in the case of out-of-memory errors.

I could open a PR that returns this information via a parameter but this would then also change the return-type to a `pandas.DataFrame`. Thus I'm unsure of the right API to implement:

1. `dd.DataFrame.memory_usage(per_partition=True) -> pandas.DataFrame`
2. `dd.DataFrame.memory_usage_detailed() -> pandas.DataFrame`.

Any opinions before I start a PR?


Discussion:
My initial preference is for a separate method (`memory_usage_detailed` or `partitionwise_memory_usage` sound fine), but I think a keyword is also reasonable.
I wonder if it is possible to just use `df.map_partitions(pd.DataFrame.memory_usage)` here.  If this doesn't work, why doesn't it work, and is there some way that we can improve it?
`df.map_partitions(pd.DataFrame.memory_usage)` seems to work for determining per-partition memory usage:

```python
In [1]: import pandas as pd

In [2]: import dask.dataframe as dd

In [3]: df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]})

In [4]: ddf = dd.from_pandas(df, npartitions=3)

In [5]: ddf.map_partitions(pd.DataFrame.memory_usage)
Out[5]:
Dask Series Structure:
npartitions=2
0    int64
2      ...
3      ...
dtype: int64
Dask Name: memory_usage, 4 tasks

In [6]: ddf.map_partitions(pd.DataFrame.memory_usage).compute()
Out[6]:
Index    128
a         16
b         16
Index    132
a         16
b         16
dtype: int64
```

@xhochy does that contain the information you're after?
I want to have information in there that identifies the partition, so we can also group by partition to get the total partition sizes.
Gotcha, I think an additional `map_partitions(pd.Series.sum)` will result in the total memory for each partition

```python
import pandas as pd
import dask.dataframe as dd

df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [5, 6, 7, 8]})
ddf = dd.from_pandas(df, npartitions=2)
mem_usage = ddf.map_partitions(pd.DataFrame.memory_usage)
mem_per_partition = mem_usage.map_partitions(pd.Series.sum)
print("mem_per_partition.compute():")
print(mem_per_partition.compute())
```

outputs

```
mem_per_partition.compute():
0    160
1    164
dtype: int64
```
Yes, that's the expected result. I would really like to have this as a single API call as a I often get problems reports from colleagues where their large dask graph fails with `KilledWorker` due to uneven partition sizes. This would be a good routine for check this. 
fwiw, this comes up fairly frequently in debugging for problems I come across
Echoing Tom's previous comment, I think a separate `memory_usage_per_partition` or `memory_usage_partitionwise` method would be a nice addition 

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -744,6 +744,29 @@ def map_overlap(self, func, before, after, *args, **kwargs):
 
         return map_overlap(func, self, before, after, *args, **kwargs)
 
+    def memory_usage_per_partition(self, index=True, deep=False):
+        """ Return the memory usage of each partition
+
+        Parameters
+        ----------
+        index : bool, default True
+            Specifies whether to include the memory usage of the index in
+            returned Series.
+        deep : bool, default False
+            If True, introspect the data deeply by interrogating
+            ``object`` dtypes for system-level memory consumption, and include
+            it in the returned values.
+
+        Returns
+        -------
+        Series
+            A Series whose index is the parition number and whose values
+            are the memory usage of each partition in bytes.
+        """
+        return self.map_partitions(
+            total_mem_usage, index=index, deep=deep
+        ).clear_divisions()
+
     @insert_meta_param_description(pad=12)
     def reduction(
         self,
@@ -5654,7 +5677,7 @@ def repartition_size(df, size):
         size = parse_bytes(size)
     size = int(size)
 
-    mem_usages = df.map_partitions(total_mem_usage).compute()
+    mem_usages = df.map_partitions(total_mem_usage, deep=True).compute()
 
     # 1. split each partition that is larger than partition_size
     nsplits = 1 + mem_usages // size
@@ -5675,8 +5698,8 @@ def repartition_size(df, size):
     return _repartition_from_boundaries(df, new_partitions_boundaries, new_name)
 
 
-def total_mem_usage(df):
-    mem_usage = df.memory_usage(deep=True)
+def total_mem_usage(df, index=True, deep=False):
+    mem_usage = df.memory_usage(index=index, deep=deep)
     if is_series_like(mem_usage):
         mem_usage = mem_usage.sum()
     return mem_usage


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
@@ -1802,7 +1802,7 @@ def test_repartition_partition_size(use_index, n, partition_size, transform):
     a = dd.from_pandas(df, npartitions=n, sort=use_index)
     b = a.repartition(partition_size=partition_size)
     assert_eq(a, b, check_divisions=False)
-    assert np.alltrue(b.map_partitions(total_mem_usage).compute() <= 1024)
+    assert np.alltrue(b.map_partitions(total_mem_usage, deep=True).compute() <= 1024)
     parts = dask.get(b.dask, b.__dask_keys__())
     assert all(map(len, parts))
 
@@ -3614,6 +3614,34 @@ def test_memory_usage(index, deep):
     )
 
 
+@pytest.mark.parametrize("index", [True, False])
+@pytest.mark.parametrize("deep", [True, False])
+def test_memory_usage_per_partition(index, deep):
+    df = pd.DataFrame(
+        {
+            "x": [1, 2, 3, 4, 5],
+            "y": [1.0, 2.0, 3.0, 4.0, 5.0],
+            "z": ["a", "b", "c", "d", "e"],
+        }
+    )
+    ddf = dd.from_pandas(df, npartitions=2)
+
+    # DataFrame.memory_usage_per_partition
+    expected = pd.Series(
+        part.compute().memory_usage(index=index, deep=deep).sum()
+        for part in ddf.partitions
+    )
+    result = ddf.memory_usage_per_partition(index=index, deep=deep)
+    assert_eq(expected, result)
+
+    # Series.memory_usage_per_partition
+    expected = pd.Series(
+        part.x.compute().memory_usage(index=index, deep=deep) for part in ddf.partitions
+    )
+    result = ddf.x.memory_usage_per_partition(index=index, deep=deep)
+    assert_eq(expected, result)
+
+
 @pytest.mark.parametrize(
     "reduction",
     [