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,