ENH: implemented BlockManager.take for axis >=1 and seriously sped up DataFrame.take

Author: wesm

pandas/core/frame.py

@@ -1090,19 +1090,22 @@ def take(self, indices, axis=0):
         -------
         taken : DataFrame
         """
-        if axis == 0:
-            new_index = self.index.take(indices)
-            new_columns = self.columns
-        else:
-            new_index = self.index
-            new_columns = self.columns.take(indices)
-
-        # TODO: implement take on BlockManager
         if self._data.is_mixed_dtype():
-            return self.reindex(index=new_index, columns=new_columns)
-
-        new_values = self.values.take(indices, axis=axis)
-        return DataFrame(new_values, index=new_index, columns=new_columns)
+            if axis == 0:
+                new_data = self._data.take(indices, axis=1)
+                return DataFrame(new_data)
+            else:
+                return self.reindex(columns=new_columns)
+        else:
+            new_values = self.values.take(indices, axis=axis)
+            if axis == 0:
+                new_columns = self.columns
+                new_index = self.index.take(indices)
+            else:
+                new_columns = self.columns.take(indices)
+                new_index = self.index
+            return DataFrame(new_values, index=new_index,
+                             columns=new_columns)
 
     #----------------------------------------------------------------------
     # Reindex-based selection methods

pandas/core/internals.py

@@ -598,6 +598,20 @@ def reindex_items(self, new_items):
 
         return BlockManager(new_blocks, new_axes)
 
+    def take(self, indices, axis=1):
+        if axis == 0:
+            raise NotImplementedError
+
+        new_axes = list(self.axes)
+        new_axes[axis] = self.axes[axis].take(indices)
+        new_blocks = []
+        for blk in self.blocks:
+            newb = make_block(blk.values.take(indices, axis=axis), blk.items,
+                              self.items)
+            new_blocks.append(newb)
+
+        return BlockManager(new_blocks, new_axes)
+
     def merge(self, other, lsuffix=None, rsuffix=None):
         assert(self._is_indexed_like(other))
 

pandas/tests/test_tseries.py

@@ -76,7 +76,7 @@ def test_inner_join_indexer():
     a = np.array([1, 2, 3, 4, 5], dtype=np.int64)
     b = np.array([0, 3, 5, 7, 9], dtype=np.int64)
 
-    index, ares, bres = lib.inner_join_indexer(a, b)
+    index, ares, bres = lib.inner_join_indexer_int64(a, b)
 
     index_exp = np.array([3, 5], dtype=np.int64)
     assert_almost_equal(index, index_exp)
@@ -90,7 +90,7 @@ def test_outer_join_indexer():
     a = np.array([1, 2, 3, 4, 5], dtype=np.int64)
     b = np.array([0, 3, 5, 7, 9], dtype=np.int64)
 
-    index, ares, bres = lib.outer_join_indexer(a, b)
+    index, ares, bres = lib.outer_join_indexer_int64(a, b)
 
     index_exp = np.array([0, 1, 2, 3, 4, 5, 7, 9], dtype=np.int64)
     assert_almost_equal(index, index_exp)