Repository: dask/dask
Pull request: Fix pd.MultiIndex size estimate
Fixes #5002.

Couple of new values with the above test setup:
Actual 3.4e6, old dask 308e6, new dask 3.34e6 (100 x 10000)
Actual 75e6, old dask 1525e6, new dask 57e6 (5 x 1000000)

Still not perfect but definitely an improvement :)


- [x] Tests added / passed
- [x] Passes `black dask` / `flake8 dask`

Problem statement:
dask.sizeof very wrong for MultiIndex
Noticed in `distributed` diagnostics page when scattering large multi-index series:
```
import sys
import pandas as pd
from uuid import uuid4

def object_size(x):  # from dask.sizeof.register_pandas
    if not len(x):
        return 0
    sample = np.random.choice(x, size=20, replace=True)
    sample = list(map(sizeof, sample))
    return sum(sample) / 20 * len(x)

x = pd.MultiIndex.from_product([[uuid4().hex for i in range(5)], range(1_000_000)])
print(x[0])
pd.Series([sys.getsizeof(x), x.memory_usage(deep=True), object_size(x)]) / 1e6
```
Output:
```
0      74.944361
1      74.944337
2    1525.000000
```

Not sure how serious of an issue this is as I don't know what all these estimates are used for.  The scatter itself does seem to be affected for some reason:
```
%time future = client.scatter(<series>, broadcast=False)
Wall time: 1min 10s

%time future = client.scatter(<series>.reset_index(), broadcast=False)
Wall time: 2.06 s
```

Hints from issue discussion:
The `object_size` there looks a little strange for MultiIndex. That forces a conversion to an ndarray of objects, which likely isn't what we want.

@mrocklin do you recall the reason for not just using `index.memory_usage(deep=True)`? 
`memory_usage` is very slow when the object is large, so the subsampling approach seems reasonable; it just doesn't work in this case because repeat values aren't stored multiple times. Is it even feasible to get a good estimate of the size here without traversing the whole array?
Mmm, though the current (unintentional) way of `np.random.choice` converting to an ndarray of tuples isn't exactly fast :)

IMO, I think we're fine with doing deep memory usage here. If we think that number of unique values per level is smallish, then pandas implementation of 

```
        level_nbytes = sum(i.memory_usage(deep=deep) for i in self.levels)
```

With a MI like

```
In [29]: idx = pd.MultiIndex.from_product([list(range(1000)), ['a', 'b']])

In [30]: idx.memory_usage(deep=True)
Out[30]: 55204
```

the `deep` part only affects a `Index(['a', 'b'].memory_usage(deep=True))`. Not too bad.

---

If we really think `deep=True` is problematic in practice, then we'll need to reimplement pandas' way of doing things, but sample the object levels.
If we are concerned with the cost of memory usage on a level with many values, I think we could define a MultiIndex-specific sizeof that applies the current `sizeof_pandas_index` to each level of the MultiIndex, and then adds in the size of the `.codes` (`.labels` for older pandas). Does that sound reasonable @bnaul?
I think that's probably the way to go; changing my example to 10^6 UUIDs, I get
```
[ins] In [105]: %time sys.getsizeof(x)
CPU times: user 1.54 s, sys: 10.3 ms, total: 1.55 s
Wall time: 1.56 s

[ins] In [106]: %time x.memory_usage(deep=True)
CPU times: user 1.45 s, sys: 10.3 ms, total: 1.46 s
Wall time: 1.47 s

[ins] In [108]: %time object_size(x)
CPU times: user 125 µs, sys: 16 µs, total: 141 µs
Wall time: 136 µs
```
Seems like the version you suggested would be about as fast as what we currently have (i.e. way faster than the full computation). I can try taking a stab at it in the next week or so.

Commit message: Fix pd.MultiIndex size estimate

Patch to review:
diff --git a/dask/sizeof.py b/dask/sizeof.py
index 9b72b00bbf8..46a7a4ff171 100644
--- a/dask/sizeof.py
+++ b/dask/sizeof.py
@@ -84,6 +84,13 @@ def sizeof_pandas_index(i):
             p += object_size(i)
         return int(p) + 1000
 
+    @sizeof.register(pd.MultiIndex)
+    def sizeof_pandas_multiindex(i):
+        p = int(sum(object_size(l) for l in i.levels))
+        for c in i.codes if hasattr(i, "codes") else i.labels:
+            p += c.nbytes
+        return int(p) + 1000
+
 
 @sizeof.register_lazy("scipy")
 def register_spmatrix():
diff --git a/dask/tests/test_sizeof.py b/dask/tests/test_sizeof.py
index 06f95c5c674..c18d3ae82d5 100644
--- a/dask/tests/test_sizeof.py
+++ b/dask/tests/test_sizeof.py
@@ -37,10 +37,12 @@ def test_pandas():
     assert sizeof(df.x) >= sizeof(df.index)
     assert sizeof(df.y) >= 100 * 3
     assert sizeof(df.index) >= 20
+    assert sizeof(df.set_index(["x", "y"].index)) < 1800
 
     assert isinstance(sizeof(df), int)
     assert isinstance(sizeof(df.x), int)
     assert isinstance(sizeof(df.index), int)
+    assert isinstance(sizeof(df.set_index(["x", "y"]).index), int)
 
 
 def test_pandas_repeated_column():


Code review comments:
diff_hunk: @@ -37,10 +37,12 @@ def test_pandas():
     assert sizeof(df.x) >= sizeof(df.index)
     assert sizeof(df.y) >= 100 * 3
     assert sizeof(df.index) >= 20
+    assert sizeof(df.set_index(['x', 'y'].index)) < 1800
path: dask/tests/test_sizeof.py
text: @user1:
It looks like this is causing a FutureWarning in Pandas.  I'm not sure exactly why.  (this is reported in the CI)