ENH: Consolidate and further improve performance of take functions

Author: stephenwlin

pandas/core/common.py

@@ -2644,8 +2644,9 @@ def _reindex_multi(self, new_index, new_columns, copy, fill_value):
         new_columns, col_indexer = self.columns.reindex(new_columns)
 
         if row_indexer is not None and col_indexer is not None:
-            new_values = com.take_2d_multi(self.values, row_indexer,
-                                           col_indexer, fill_value=fill_value)
+            indexer = row_indexer, col_indexer
+            new_values = com.take_2d_multi(self.values, indexer,
+                                           fill_value=fill_value)
             return DataFrame(new_values, index=new_index, columns=new_columns)
         elif row_indexer is not None:
             return self._reindex_with_indexers(new_index, row_indexer,

pandas/core/frame.py

@@ -897,7 +897,7 @@ def _aggregate_series_fast(self, obj, func):
         dummy = obj[:0].copy()
         indexer = _algos.groupsort_indexer(group_index, ngroups)[0]
         obj = obj.take(indexer)
-        group_index = com.ndtake(group_index, indexer)
+        group_index = com.take_nd(group_index, indexer, allow_fill=False)
         grouper = lib.SeriesGrouper(obj, func, group_index, ngroups,
                                     dummy)
         result, counts = grouper.get_result()
@@ -1686,7 +1686,9 @@ def aggregate(self, arg, *args, **kwargs):
                 zipped = zip(result.index.levels, result.index.labels,
                              result.index.names)
                 for i, (lev, lab, name) in enumerate(zipped):
-                    result.insert(i, name, com.ndtake(lev.values, lab))
+                    result.insert(i, name,
+                                  com.take_nd(lev.values, lab,
+                                              allow_fill=False))
                 result = result.consolidate()
             else:
                 values = result.index.values
@@ -2133,7 +2135,7 @@ def __init__(self, data, labels, ngroups, axis=0, keep_internal=False):
     @cache_readonly
     def slabels(self):
         # Sorted labels
-        return com.ndtake(self.labels, self.sort_idx)
+        return com.take_nd(self.labels, self.sort_idx, allow_fill=False)
 
     @cache_readonly
     def sort_idx(self):
@@ -2411,11 +2413,11 @@ def _reorder_by_uniques(uniques, labels):
     mask = labels < 0
 
     # move labels to right locations (ie, unsort ascending labels)
-    labels = com.ndtake(reverse_indexer, labels)
+    labels = com.take_nd(reverse_indexer, labels, allow_fill=False)
     np.putmask(labels, mask, -1)
 
     # sort observed ids
-    uniques = com.ndtake(uniques, sorter)
+    uniques = com.take_nd(uniques, sorter, allow_fill=False)
 
     return uniques, labels
 

pandas/core/groupby.py

@@ -12,7 +12,6 @@
 import pandas.index as _index
 from pandas.lib import Timestamp
 
-from pandas.core.common import ndtake
 from pandas.util.decorators import cache_readonly
 import pandas.core.common as com
 from pandas.util import py3compat
@@ -608,7 +607,8 @@ def union(self, other):
             indexer = (indexer == -1).nonzero()[0]
 
             if len(indexer) > 0:
-                other_diff = ndtake(other.values, indexer)
+                other_diff = com.take_nd(other.values, indexer,
+                                         allow_fill=False)
                 result = com._concat_compat((self.values, other_diff))
                 try:
                     result.sort()
@@ -1037,7 +1037,8 @@ def _join_level(self, other, level, how='left', return_indexers=False):
             rev_indexer = lib.get_reverse_indexer(left_lev_indexer,
                                                   len(old_level))
 
-            new_lev_labels = ndtake(rev_indexer, left.labels[level])
+            new_lev_labels = com.take_nd(rev_indexer, left.labels[level],
+                                         allow_fill=False)
             omit_mask = new_lev_labels != -1
 
             new_labels = list(left.labels)
@@ -1057,8 +1058,9 @@ def _join_level(self, other, level, how='left', return_indexers=False):
             left_indexer = None
 
         if right_lev_indexer is not None:
-            right_indexer = ndtake(right_lev_indexer,
-                                   join_index.labels[level])
+            right_indexer = com.take_nd(right_lev_indexer,
+                                        join_index.labels[level],
+                                        allow_fill=False)
         else:
             right_indexer = join_index.labels[level]
 
@@ -2369,8 +2371,10 @@ def equals(self, other):
             return False
 
         for i in xrange(self.nlevels):
-            svalues = ndtake(self.levels[i].values, self.labels[i])
-            ovalues = ndtake(other.levels[i].values, other.labels[i])
+            svalues = com.take_nd(self.levels[i].values, self.labels[i],
+                                  allow_fill=False)
+            ovalues = com.take_nd(other.levels[i].values, other.labels[i],
+                                  allow_fill=False)
             if not np.array_equal(svalues, ovalues):
                 return False
 

pandas/core/index.py

@@ -120,14 +120,14 @@ def merge(self, other):
         #     union_ref = self.ref_items + other.ref_items
         return _merge_blocks([self, other], self.ref_items)
 
-    def reindex_axis(self, indexer, mask, needs_masking, axis=0,
-                     fill_value=np.nan):
+    def reindex_axis(self, indexer, axis=1, fill_value=np.nan, mask_info=None):
         """
         Reindex using pre-computed indexer information
         """
-        new_values = com.take_fast(self.values, indexer,
-                                   mask, needs_masking, axis=axis,
-                                   fill_value=fill_value)
+        if axis < 1:
+            raise AssertionError('axis must be at least 1, got %d' % axis)
+        new_values = com.take_nd(self.values, indexer, axis,
+                                 fill_value=fill_value, mask_info=mask_info)
         return make_block(new_values, self.items, self.ref_items)
 
     def reindex_items_from(self, new_ref_items, copy=True):
@@ -146,12 +146,9 @@ def reindex_items_from(self, new_ref_items, copy=True):
             new_items = new_ref_items
             new_values = self.values.copy() if copy else self.values
         else:
-            mask = indexer != -1
-            masked_idx = indexer[mask]
-
-            new_values = com.take_fast(self.values, masked_idx,
-                                       mask=None, needs_masking=False,
-                                       axis=0)
+            masked_idx = indexer[indexer != -1]
+            new_values = com.take_nd(self.values, masked_idx, axis=0,
+                                     allow_fill=False)
             new_items = self.items.take(masked_idx)
         return make_block(new_values, new_items, new_ref_items)
 
@@ -221,7 +218,10 @@ def fillna(self, value, inplace=False):
             return make_block(new_values, self.items, self.ref_items)
 
     def astype(self, dtype, copy = True, raise_on_error = True):
-        """ coerce to the new type (if copy=True, return a new copy) raise on an except if raise == True """
+        """
+        Coerce to the new type (if copy=True, return a new copy)
+        raise on an except if raise == True
+        """
         try:
             newb = make_block(com._astype_nansafe(self.values, dtype, copy = copy),
                               self.items, self.ref_items)
@@ -231,12 +231,12 @@ def astype(self, dtype, copy = True, raise_on_error = True):
             newb = self.copy() if copy else self
 
         if newb.is_numeric and self.is_numeric:
-            if newb.shape != self.shape or (not copy and newb.itemsize < self.itemsize):
-                raise TypeError("cannot set astype for copy = [%s] for dtype (%s [%s]) with smaller itemsize that current (%s [%s])" % (copy,
-                                                                                                                                        self.dtype.name,
-                                                                                                                                        self.itemsize,
-                                                                                                                                        newb.dtype.name,
-                                                                                                                                        newb.itemsize))
+            if (newb.shape != self.shape or
+                    (not copy and newb.itemsize < self.itemsize)):
+                raise TypeError("cannot set astype for copy = [%s] for dtype "
+                                "(%s [%s]) with smaller itemsize that current "
+                                "(%s [%s])" % (copy, self.dtype.name,
+                                self.itemsize, newb.dtype.name, newb.itemsize))
         return newb
 
     def convert(self, copy = True, **kwargs):
@@ -356,11 +356,11 @@ def interpolate(self, method='pad', axis=0, inplace=False,
 
         return make_block(values, self.items, self.ref_items)
 
-    def take(self, indexer, axis=1, fill_value=np.nan):
+    def take(self, indexer, axis=1):
         if axis < 1:
             raise AssertionError('axis must be at least 1, got %d' % axis)
-        new_values = com.take_fast(self.values, indexer, None, False,
-                                   axis=axis, fill_value=fill_value)
+        new_values = com.take_nd(self.values, indexer, axis=axis,
+                                 allow_fill=False)
         return make_block(new_values, self.items, self.ref_items)
 
     def get_values(self, dtype):
@@ -1320,15 +1320,9 @@ def reindex_indexer(self, new_axis, indexer, axis=1, fill_value=np.nan):
         if axis == 0:
             return self._reindex_indexer_items(new_axis, indexer, fill_value)
 
-        mask = indexer == -1
-
-        # TODO: deal with length-0 case? or does it fall out?
-        needs_masking = len(new_axis) > 0 and mask.any()
-
         new_blocks = []
         for block in self.blocks:
-            newb = block.reindex_axis(indexer, mask, needs_masking,
-                                      axis=axis, fill_value=fill_value)
+            newb = block.reindex_axis(indexer, axis=axis, fill_value=fill_value)
             new_blocks.append(newb)
 
         new_axes = list(self.axes)
@@ -1354,8 +1348,8 @@ def _reindex_indexer_items(self, new_items, indexer, fill_value):
                 continue
 
             new_block_items = new_items.take(selector.nonzero()[0])
-            new_values = com.take_fast(blk.values, blk_indexer[selector],
-                                       None, False, axis=0)
+            new_values = com.take_nd(blk.values, blk_indexer[selector], axis=0,
+                                     allow_fill=False)
             new_blocks.append(make_block(new_values, new_block_items,
                                          new_items))
 
@@ -1419,8 +1413,8 @@ def _make_na_block(self, items, ref_items, fill_value=np.nan):
         return na_block
 
     def take(self, indexer, axis=1):
-        if axis == 0:
-            raise NotImplementedError
+        if axis < 1:
+            raise AssertionError('axis must be at least 1, got %d' % axis)
 
         indexer = com._ensure_platform_int(indexer)
 
@@ -1433,8 +1427,8 @@ def take(self, indexer, axis=1):
         new_axes[axis] = self.axes[axis].take(indexer)
         new_blocks = []
         for blk in self.blocks:
-            new_values = com.take_fast(blk.values, indexer, None, False,
-                                       axis=axis)
+            new_values = com.take_nd(blk.values, indexer, axis=axis,
+                                     allow_fill=False)
             newb = make_block(new_values, blk.items, self.items)
             new_blocks.append(newb)
 

pandas/core/internals.py

@@ -838,7 +838,7 @@ def _reindex_multi(self, items, major, minor):
             indexer2 = range(len(new_minor))
 
         for i, ind in enumerate(indexer0):
-            com.take_2d_multi(values[ind], indexer1, indexer2,
+            com.take_2d_multi(values[ind], (indexer1, indexer2),
                               out=new_values[i])
 
         return Panel(new_values, items=new_items, major_axis=new_major,