Repository: dask/dask

Issue:
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
```


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.

Resolving patch:
diff --git a/dask/sizeof.py b/dask/sizeof.py
--- 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():


Test patch:
diff --git a/dask/tests/test_sizeof.py b/dask/tests/test_sizeof.py
--- a/dask/tests/test_sizeof.py
+++ b/dask/tests/test_sizeof.py
@@ -43,6 +43,15 @@ def test_pandas():
     assert isinstance(sizeof(df.index), int)
 
 
+def test_pandas_multiindex():
+    pd = pytest.importorskip("pandas")
+    index = pd.MultiIndex.from_product([range(5), ["a", "b", "c", "d", "e"]])
+    actual_size = sys.getsizeof(index) + 1000  # adjust for serialization overhead
+
+    assert 0.5 * actual_size < sizeof(index) < 2 * actual_size
+    assert isinstance(sizeof(index), int)
+
+
 def test_pandas_repeated_column():
     pd = pytest.importorskip("pandas")
     df = pd.DataFrame({"x": [1, 2, 3]})