PERF: correctly report memory used by Index's

Author: jreback

doc/source/whatsnew/v0.20.0.txt

@@ -314,6 +314,43 @@ New Behavior:
   In [5]: df['a']['2011-12-31 23:59:59']
   Out[5]: 1
 
+.. _whatsnew_0200.api_breaking.memory_usage:
+
+Memory Usage for Index now is now more accurate
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+In previous versions, showing ``.memory_usage()`` on a pandas structure that has an index, would only include actual index values and not include structures that facilitated fast indexing. This will generally be different for ``Index`` and ``MultiIndex`` and less-so for other index types.
+
+Previous Behavior:
+
+.. code-block:: ipython
+
+   In [8]: index = Index(['foo', 'bar', 'baz'])
+
+   In [9]: index.memory_usage(deep=True)
+   Out[9]: 180
+
+   In [10]: index.get_loc('foo')
+   Out[10]: 0
+
+   In [11]: index.memory_usage(deep=True)
+   Out[11]: 180
+
+New Behavior:
+
+.. code-block:: ipython
+
+   In [8]: index = Index(['foo', 'bar', 'baz'])
+
+   In [9]: index.memory_usage(deep=True)
+   Out[9]: 180
+
+   In [10]: index.get_loc('foo')
+   Out[10]: 0
+
+   In [11]: index.memory_usage(deep=True)
+   Out[11]: 276
+
 .. _whatsnew_0200.api:
 
 Other API Changes

pandas/core/base.py

@@ -1067,6 +1067,7 @@ def memory_usage(self, deep=False):
         v = self.values.nbytes
         if deep and is_object_dtype(self):
             v += lib.memory_usage_of_objects(self.values)
+
         return v
 
     def factorize(self, sort=False, na_sentinel=-1):

pandas/index.pyx

@@ -203,6 +203,12 @@ cdef class IndexEngine:
 
         return result
 
+    def sizeof(self, deep=False):
+        """ return the sizeof our mapping """
+        if not self.is_mapping_populated:
+            return 0
+        return self.mapping.sizeof(deep=deep)
+
     property is_unique:
 
         def __get__(self):

pandas/indexes/base.py

@@ -536,6 +536,14 @@ def get_values(self):
         """ return the underlying data as an ndarray """
         return self.values
 
+    @Appender(IndexOpsMixin.memory_usage.__doc__)
+    def memory_usage(self, deep=False):
+        result = super(Index, self).memory_usage(deep=deep)
+
+        # include our engine hashtable
+        result += self._engine.sizeof(deep=deep)
+        return result
+
     # ops compat
     def tolist(self):
         """

pandas/indexes/multi.py

@@ -446,13 +446,19 @@ def _nbytes(self, deep=False):
         return the number of bytes in the underlying data
         deeply introspect the level data if deep=True
 
+        include the engine hashtable
+
         *this is in internal routine*
 
         """
         level_nbytes = sum((i.memory_usage(deep=deep) 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
+        result = level_nbytes + label_nbytes + names_nbytes
+
+        # include our engine hashtable
+        result += self._engine.sizeof(deep=deep)
+        return result
 
     def _format_attrs(self):
         """

pandas/src/hashtable_class_helper.pxi.in

@@ -203,6 +203,7 @@ cdef class ObjectVector:
 
 
 cdef class HashTable:
+
     pass
 
 {{py:
@@ -237,6 +238,12 @@ cdef class {{name}}HashTable(HashTable):
         k = kh_get_{{dtype}}(self.table, key)
         return k != self.table.n_buckets
 
+    def sizeof(self, deep=False):
+        """ return the size of my table in bytes """
+        return self.table.n_buckets * (sizeof(self.table.keys) +
+                                       sizeof(self.table.vals) +
+                                       sizeof(self.table.flags))
+
     cpdef get_item(self, {{dtype}}_t val):
         cdef khiter_t k
         k = kh_get_{{dtype}}(self.table, val)
@@ -464,6 +471,12 @@ cdef class StringHashTable(HashTable):
             kh_destroy_str(self.table)
             self.table = NULL
 
+    def sizeof(self, deep=False):
+        """ return the size of my table in bytes """
+        return self.table.n_buckets * (sizeof(self.table.keys) +
+                                       sizeof(self.table.vals) +
+                                       sizeof(self.table.flags))
+
     cpdef get_item(self, object val):
         cdef:
             khiter_t k
@@ -714,6 +727,12 @@ cdef class PyObjectHashTable(HashTable):
         k = kh_get_pymap(self.table, <PyObject*>key)
         return k != self.table.n_buckets
 
+    def sizeof(self, deep=False):
+        """ return the size of my table in bytes """
+        return self.table.n_buckets * (sizeof(self.table.keys) +
+                                       sizeof(self.table.vals) +
+                                       sizeof(self.table.flags))
+
     cpdef get_item(self, object val):
         cdef khiter_t k
         if val != val or val is None:

pandas/tests/indexes/common.py

@@ -366,6 +366,23 @@ def test_compat(self):
         for ind in self.indices.values():
             self.assertEqual(ind.tolist(), list(ind))
 
+    def test_memory_usage(self):
+        for name, index in compat.iteritems(self.indices):
+            result = index.memory_usage()
+            if len(index):
+                index.get_loc(index[0])
+                result2 = index.memory_usage()
+                result3 = index.memory_usage(deep=True)
+                self.assertTrue(result2 > result)
+
+                if index.inferred_type == 'object':
+                    self.assertTrue(result3 > result2)
+
+            else:
+
+                # we report 0 for no-length
+                self.assertEqual(result, 0)
+
     def test_argsort(self):
         for k, ind in self.indices.items():