PERF: set multiindex labels with coerced dtype (GH8456)

Author: jreback

doc/source/whatsnew/v0.15.1.txt

@@ -20,6 +20,30 @@ users upgrade to this version.
 API changes
 ~~~~~~~~~~~
 
+- Represent ``MultiIndex`` labels with a dtype that utilizes memory based on the level size. In prior versions, the memory usage was a constant 8 bytes per element in each level. In addition, in prior versions, the *reported* memory usage was incorrect as it didn't show the usage for the memory occupied by the underling data array. (:issue:`8456`)
+
+  .. ipython:: python
+
+     dfi = DataFrame(1,index=pd.MultiIndex.from_product([['a'],range(1000)]),columns=['A'])
+
+  previous behavior:
+
+  .. code-block:: python
+
+     # this was underreported and actually took (in < 0.15.1) about 24008 bytes
+     In [1]: dfi.memory_usage(index=True)
+     Out[1]:
+     Index    8000
+     A        8000
+     dtype: int64
+
+
+  current behavior:
+
+  .. ipython:: python
+
+     dfi.memory_usage(index=True)
+
 - ``groupby`` with ``as_index=False`` will not add erroneous extra columns to
   result (:issue:`8582`):
 

pandas/core/base.py

@@ -232,7 +232,6 @@ def __repr__(self):
     __setitem__ = __setslice__ = __delitem__ = __delslice__ = _disabled
     pop = append = extend = remove = sort = insert = _disabled
 
-
 class FrozenNDArray(PandasObject, np.ndarray):
 
     # no __array_finalize__ for now because no metadata
@@ -540,4 +539,3 @@ def duplicated(self, take_last=False):
 
     def _update_inplace(self, result, **kwargs):
         raise NotImplementedError
-

pandas/core/common.py

@@ -725,6 +725,7 @@ def _get_take_nd_function(ndim, arr_dtype, out_dtype, axis=0, mask_info=None):
             return func
 
     def func(arr, indexer, out, fill_value=np.nan):
+        indexer = _ensure_int64(indexer)
         _take_nd_generic(arr, indexer, out, axis=axis,
                          fill_value=fill_value, mask_info=mask_info)
     return func
@@ -817,6 +818,7 @@ def take_nd(arr, indexer, axis=0, out=None, fill_value=np.nan,
     func = _get_take_nd_function(arr.ndim, arr.dtype, out.dtype,
                                  axis=axis, mask_info=mask_info)
 
+    indexer = _ensure_int64(indexer)
     func(arr, indexer, out, fill_value)
 
     if flip_order:
@@ -965,15 +967,14 @@ def diff(arr, n, axis=0):
 
 def _coerce_indexer_dtype(indexer, categories):
     """ coerce the indexer input array to the smallest dtype possible """
-    indexer = np.array(indexer,copy=False)
     l = len(categories)
     if l < _int8_max:
-        return indexer.astype('int8')
+        return _ensure_int8(indexer)
     elif l < _int16_max:
-        return indexer.astype('int16')
+        return _ensure_int16(indexer)
     elif l < _int32_max:
-        return indexer.astype('int32')
-    return indexer.astype('int64')
+        return _ensure_int32(indexer)
+    return _ensure_int64(indexer)
 
 def _coerce_to_dtypes(result, dtypes):
     """ given a dtypes and a result set, coerce the result elements to the

pandas/core/frame.py

@@ -1564,7 +1564,7 @@ def memory_usage(self, index=False):
         result = Series([ c.values.nbytes for col, c in self.iteritems() ],
                         index=self.columns)
         if index:
-             result = Series(self.index.values.nbytes,
+             result = Series(self.index.nbytes,
                         index=['Index']).append(result)
         return result
 

pandas/core/index.py

@@ -7,6 +7,7 @@
 from pandas import compat
 import numpy as np
 
+from sys import getsizeof
 import pandas.tslib as tslib
 import pandas.lib as lib
 import pandas.algos as _algos
@@ -17,7 +18,7 @@
 from pandas.core.common import isnull, array_equivalent
 import pandas.core.common as com
 from pandas.core.common import (_values_from_object, is_float, is_integer,
-                                ABCSeries, _ensure_object)
+                                ABCSeries, _ensure_object, _ensure_int64)
 from pandas.core.config import get_option
 
 # simplify
@@ -2680,13 +2681,13 @@ def _set_labels(self, labels, level=None, copy=False, validate=True,
             raise ValueError('Length of labels must match length of levels.')
 
         if level is None:
-            new_labels = FrozenList(_ensure_frozen(v, copy=copy)._shallow_copy()
-                                    for v in labels)
+            new_labels = FrozenList(_ensure_frozen(lab, lev, copy=copy)._shallow_copy()
+                                    for lev, lab in zip(self.levels, labels))
         else:
             level = [self._get_level_number(l) for l in level]
             new_labels = list(self._labels)
-            for l, v in zip(level, labels):
-                new_labels[l] = _ensure_frozen(v, copy=copy)._shallow_copy()
+            for l, lev, lab in zip(level, self.levels, labels):
+                new_labels[l] = _ensure_frozen(lab, lev, copy=copy)._shallow_copy()
             new_labels = FrozenList(new_labels)
 
         self._labels = new_labels
@@ -2824,6 +2825,14 @@ def _array_values(self):
     def dtype(self):
         return np.dtype('O')
 
+    @cache_readonly
+    def nbytes(self):
+        """ return the number of bytes in the underlying data """
+        level_nbytes = sum(( i.nbytes for i in self.levels ))
+        label_nbytes = sum(( i.nbytes for i in self.labels ))
+        names_nbytes = sum(( getsizeof(i) for i in self.names ))
+        return level_nbytes + label_nbytes + names_nbytes
+
     def __repr__(self):
         encoding = get_option('display.encoding')
         attrs = [('levels', default_pprint(self.levels)),
@@ -4361,7 +4370,7 @@ def insert(self, loc, item):
                 lev_loc = level.get_loc(k)
 
             new_levels.append(level)
-            new_labels.append(np.insert(labels, loc, lev_loc))
+            new_labels.append(np.insert(_ensure_int64(labels), loc, lev_loc))
 
         return MultiIndex(levels=new_levels, labels=new_labels,
                           names=self.names, verify_integrity=False)
@@ -4474,8 +4483,8 @@ def _ensure_index(index_like, copy=False):
     return Index(index_like)
 
 
-def _ensure_frozen(array_like, copy=False):
-    array_like = np.asanyarray(array_like, dtype=np.int_)
+def _ensure_frozen(array_like, categories, copy=False):
+    array_like = com._coerce_indexer_dtype(array_like, categories)
     array_like = array_like.view(FrozenNDArray)
     if copy:
         array_like = array_like.copy()

pandas/tests/test_frame.py

@@ -6766,6 +6766,15 @@ def test_info_memory_usage(self):
         size_df = np.size(df.columns.values)  # index=False; default
         self.assertEqual(size_df, np.size(df.memory_usage()))
 
+        # test for validity
+        DataFrame(1,index=['a'],columns=['A']).memory_usage(index=True)
+        DataFrame(1,index=['a'],columns=['A']).index.nbytes
+        DataFrame(1,index=pd.MultiIndex.from_product([['a'],range(1000)]),columns=['A']).index.nbytes
+        DataFrame(1,index=pd.MultiIndex.from_product([['a'],range(1000)]),columns=['A']).index.values.nbytes
+        DataFrame(1,index=pd.MultiIndex.from_product([['a'],range(1000)]),columns=['A']).memory_usage(index=True)
+        DataFrame(1,index=pd.MultiIndex.from_product([['a'],range(1000)]),columns=['A']).index.nbytes
+        DataFrame(1,index=pd.MultiIndex.from_product([['a'],range(1000)]),columns=['A']).index.values.nbytes
+
     def test_dtypes(self):
         self.mixed_frame['bool'] = self.mixed_frame['A'] > 0
         result = self.mixed_frame.dtypes

pandas/tests/test_index.py

@@ -37,6 +37,7 @@
 class Base(object):
     """ base class for index sub-class tests """
     _holder = None
+    _compat_props = ['shape', 'ndim', 'size', 'itemsize', 'nbytes']
 
     def verify_pickle(self,index):
         unpickled = self.round_trip_pickle(index)
@@ -90,9 +91,12 @@ def test_ndarray_compat_properties(self):
         self.assertTrue(idx.transpose().equals(idx))
 
         values = idx.values
-        for prop in ['shape', 'ndim', 'size', 'itemsize', 'nbytes']:
+        for prop in self._compat_props:
             self.assertEqual(getattr(idx, prop), getattr(values, prop))
 
+        # test for validity
+        idx.nbytes
+        idx.values.nbytes
 
 class TestIndex(Base, tm.TestCase):
     _holder = Index
@@ -1837,6 +1841,7 @@ def test_pickle_compat_construction(self):
 class TestMultiIndex(Base, tm.TestCase):
     _holder = MultiIndex
     _multiprocess_can_split_ = True
+    _compat_props = ['shape', 'ndim', 'size', 'itemsize']
 
     def setUp(self):
         major_axis = Index(['foo', 'bar', 'baz', 'qux'])
@@ -1865,6 +1870,24 @@ def f():
                 pass
         tm.assertRaisesRegexp(ValueError,'The truth value of a',f)
 
+    def test_labels_dtypes(self):
+
+        # GH 8456
+        i = MultiIndex.from_tuples([('A', 1), ('A', 2)])
+        self.assertTrue(i.labels[0].dtype == 'int8')
+        self.assertTrue(i.labels[1].dtype == 'int8')
+
+        i = MultiIndex.from_product([['a'],range(40)])
+        self.assertTrue(i.labels[1].dtype == 'int8')
+        i = MultiIndex.from_product([['a'],range(400)])
+        self.assertTrue(i.labels[1].dtype == 'int16')
+        i = MultiIndex.from_product([['a'],range(40000)])
+        self.assertTrue(i.labels[1].dtype == 'int32')
+
+        i = pd.MultiIndex.from_product([['a'],range(1000)])
+        self.assertTrue((i.labels[0]>=0).all())
+        self.assertTrue((i.labels[1]>=0).all())
+
     def test_hash_error(self):
         with tm.assertRaisesRegexp(TypeError,
                                    "unhashable type: %r" %