ENH: added a cache of DataFrame columns to improve perf with many calls to __getitem__

Author: wesm

pandas/core/frame.py

@@ -163,6 +163,7 @@ def __init__(self, data=None, index=None, columns=None, dtype=None,
             raise PandasError('DataFrame constructor not properly called!')
 
         self._data = mgr
+        self._series_cache = {}
 
     def _init_dict(self, data, index, columns, dtype=None):
         """
@@ -613,8 +614,15 @@ def _get_dtype_counts(self):
     #----------------------------------------------------------------------
     # properties for index and columns
 
+    def _get_columns(self):
+        return self._data.axes[0]
+
+    def _set_columns(self, value):
+        self._data.set_axis(0, value)
+        self._series_cache.clear()
+    columns = property(fset=_set_columns, fget=_get_columns)
+
     # reference underlying BlockManager
-    columns = AxisProperty(0)
     index = AxisProperty(1)
 
     def as_matrix(self, columns=None):
@@ -653,6 +661,7 @@ def __setstate__(self, state):
         else: # pragma: no cover
             # old pickling format, for compatibility
             self._unpickle_matrix_compat(state)
+        self._series_cache = {}
 
     def _unpickle_frame_compat(self, state): # pragma: no cover
         from pandas.core.common import _unpickle_array
@@ -770,8 +779,14 @@ def _getitem_multilevel(self, key):
             return self._getitem_single(key)
 
     def _getitem_single(self, key):
+        res = self._series_cache.get(key)
+        if res is not None:
+            return res
+
         values = self._data.get(key)
-        return Series(values, index=self.index)
+        res = Series(values, index=self.index)
+        self._series_cache[key] = res
+        return res
 
     def __setitem__(self, key, value):
         """
@@ -834,6 +849,11 @@ def _set_item(self, key, value):
         value = np.atleast_2d(value) # is this a hack?
         self._data.set(key, value)
 
+        try:
+            del self._series_cache[key]
+        except KeyError:
+            pass
+
     def _sanitize_column(self, value):
         # Need to make sure new columns (which go into the BlockManager as new
         # blocks) are always copied
@@ -864,6 +884,11 @@ def __delitem__(self, key):
         """
         self._data.delete(key)
 
+        try:
+            del self._series_cache[key]
+        except KeyError:
+            pass
+
     def pop(self, item):
         """
         Return column and drop from frame. Raise KeyError if not
@@ -1219,9 +1244,11 @@ def rename(self, index=None, columns=None):
 
     def _rename_index_inplace(self, mapper):
         self._data = self._data.rename_axis(mapper, axis=1)
+        self._series_cache.clear()
 
     def _rename_columns_inplace(self, mapper):
         self._data = self._data.rename_items(mapper)
+        self._series_cache.clear()
 
     #----------------------------------------------------------------------
     # Arithmetic / combination related

pandas/core/indexing.py

@@ -90,18 +90,20 @@ def __init__(self, frame):
         self.frame = frame
 
     def __getitem__(self, key):
-        frame = self.frame
         if isinstance(key, slice):
             return self._fancy_getitem_axis(key, axis=0)
         elif isinstance(key, tuple):
-            if len(key) != 2:
-                raise Exception('only length 2 tuple supported')
-            return self._fancy_getitem_tuple(*key)
+            return self._getitem_tuple(key)
         elif _is_list_like(key):
             return self._fancy_getitem(key, axis=0)
         else:
             return self._fancy_getitem_axis(key, axis=0)
 
+    def _getitem_tuple(self, key):
+        if len(key) != 2:
+            raise Exception('only length 2 tuple supported')
+        return self._fancy_getitem_tuple(*key)
+
     def __setitem__(self, key, value):
         # also has the side effect of consolidating in-place
         if self.frame._is_mixed_type:

pandas/core/series.py

@@ -18,6 +18,7 @@
 from pandas.core.daterange import DateRange
 from pandas.core.generic import PandasObject
 from pandas.core.index import Index, MultiIndex
+from pandas.core.indexing import _SeriesIndexer, _maybe_droplevels
 import pandas.core.datetools as datetools
 import pandas._tseries as _tseries
 
@@ -1477,8 +1478,6 @@ def _lastTimeWithValue(self): # pragma: no cover
     _ix = None
     @property
     def ix(self):
-        from pandas.core.indexing import _SeriesIndexer
-
         if self._ix is None:
             self._ix = _SeriesIndexer(self)
 

pandas/tests/test_internals.py

@@ -1,3 +1,5 @@
+# pylint: disable=W0102
+
 import unittest
 
 import numpy as np