ENH: start enabling proper fancy indexing with multi-level index

Author: wesm

pandas/core/frame.py

@@ -951,7 +951,7 @@ def dropna(self, axis=0, how='any', thresh=None, subset=None):
         how : {'any', 'all'}
             any : if any NA values are present, drop that label
             all : if all values are NA, drop that label
-        thresh : {'any', 'all', int}
+        thresh : int, default None
             int value : require that many non-NA values
         subset : array-like
 

pandas/core/index.py

@@ -375,7 +375,6 @@ class MultiIndex(Index):
     """
     Implements multi-level, a.k.a. hierarchical, index object for pandas objects
 
-
     Parameters
     ----------
     levels : list or tuple of arrays
@@ -559,29 +558,42 @@ def get_tuple_index(self):
     def slice_locs(self, start=None, end=None):
         """
 
-
         Returns
         -------
 
         Notes
         -----
-        This function assumes that the data is sorted, so use at your own peril
+        This function assumes that the data is sorted by the first level
         """
+        level0 = self.levels[0]
+
         if start is None:
-            beg_slice = 0
-        elif start in self:
-            beg_slice = self.indexMap[start]
+            start_slice = 0
+        elif isinstance(start, tuple):
+            pass
         else:
-            beg_slice = self.searchsorted(start, side='left')
+            try:
+                start_label = level0.indexMap[start]
+            except KeyError:
+                start_label = level0.searchsorted(start)
+
+            start_slice = self.labels[0].searchsorted(start_label)
 
         if end is None:
             end_slice = len(self)
-        elif end in self.indexMap:
-            end_slice = self.indexMap[end] + 1
+        elif isinstance(end, tuple):
+            pass
         else:
-            end_slice = self.searchsorted(end, side='right')
+            try:
+                end_label = level0.indexMap[end]
+            except KeyError:
+                end_label = level0.searchsorted(end, side='right')
+                if end_label > 0:
+                    end_label -= 1
 
-        return beg_slice, end_slice
+            end_slice = self.labels[0].searchsorted(end_label, side='right')
+
+        return start_slice, end_slice
 
     def truncate(self, before=None, after=None):
         """

pandas/core/indexing.py

@@ -298,9 +298,4 @@ def crit(x):
 def _need_slice(obj):
     return obj.start is not None or obj.stop is not None
 
-# I don't think this is necessary
-# def _check_step(obj):
-#     if obj.step is not None and obj.step != 1:
-#         raise Exception('steps other than 1 are not supported')
-
 _isboolarr = lambda x: np.asarray(x).dtype == np.bool_

pandas/core/series.py

@@ -41,19 +41,15 @@ def wrapper(self, other):
             if self.index.equals(other.index):
                 return Series(op(self.values, other.values), index=self.index)
 
-            newIndex = self.index + other.index
-
-            try:
-                this_reindexed = self.reindex(newIndex)
-                other_reindexed = other.reindex(newIndex)
-                arr = op(this_reindexed.values, other_reindexed.values)
-            except Exception:
-                arr = Series.combine(self, other, getattr(type(self[0]), name))
-            result = Series(arr, index=newIndex)
-            return result
+            new_index = self.index + other.index
+            this_reindexed = self.reindex(new_index)
+            other_reindexed = other.reindex(new_index)
+            arr = op(this_reindexed.values, other_reindexed.values)
+            return Series(arr, index=new_index)
         elif isinstance(other, DataFrame):
             return NotImplemented
         else:
+            # scalars
             return Series(op(self.values, other), index=self.index)
     return wrapper
 
@@ -84,7 +80,7 @@ def f(self, other, fill_value=None):
 
 class Series(np.ndarray, PandasObject):
     """
-    Generic indexed (labeled) vector (time series or cross-section)
+    Generic indexed (labeled) vector, including time series
 
     Contains values in a numpy-ndarray with an optional bound index
     (also an array of dates, strings, or whatever you want the 'row
@@ -124,7 +120,7 @@ class Series(np.ndarray, PandasObject):
 
     _AXIS_NAMES = dict((v, k) for k, v in _AXIS_NUMBERS.iteritems())
 
-    def __new__(cls, data, index=None, dtype=None, copy=False):
+    def __new__(cls, data, index=None, dtype=None, name=None, copy=False):
         if isinstance(data, Series):
             if index is None:
                 index = data.index
@@ -178,6 +174,7 @@ def __new__(cls, data, index=None, dtype=None, copy=False):
         # Change the class of the array to be the subclass type.
         subarr = subarr.view(cls)
         subarr.index = index
+        subarr.name = name
 
         if subarr.index.is_all_dates():
             subarr = subarr.view(TimeSeries)
@@ -271,12 +268,7 @@ def __getitem__(self, key):
         except TypeError:
             pass
 
-        # boolean indexing, need to check that the data are aligned, otherwise
-        # disallowed
-        if isinstance(key, Series) and key.dtype == np.bool_:
-            if not key.index.equals(self.index):
-                raise Exception('can only boolean index with like-indexed '
-                                'Series or raw ndarrays')
+        self._check_bool_indexer(key)
 
         def _index_with(indexer):
             return Series(self.values[indexer],
@@ -338,7 +330,7 @@ def get(self, key, default=None):
         y : scalar
         """
         if key in self.index:
-            return self._get_val_at(self.index.indexMap[key])
+            return self._get_val_at(self.index.get_loc(key))
         else:
             return default
 
@@ -361,7 +353,7 @@ def __getslice__(self, i, j):
     def __setitem__(self, key, value):
         values = self.values
         try:
-            loc = self.index.indexMap[key]
+            loc = self.index.get_loc(key)
             values[loc] = value
             return
         except KeyError:
@@ -373,12 +365,7 @@ def __setitem__(self, key, value):
             # Could not hash item
             pass
 
-        # boolean indexing, need to check that the data are aligned, otherwise
-        # disallowed
-        if isinstance(key, Series) and key.dtype == np.bool_:
-            if not key.index.equals(self.index):
-                raise Exception('can only boolean index with like-indexed '
-                                'Series or raw ndarrays')
+        self._check_bool_indexer(key)
 
         # special handling of boolean data with NAs stored in object
         # arrays. Sort of an elaborate hack since we can't represent boolean
@@ -396,6 +383,14 @@ def __setitem__(self, key, value):
 
         values[key] = value
 
+    def _check_bool_indexer(self, key):
+        # boolean indexing, need to check that the data are aligned, otherwise
+        # disallowed
+        if isinstance(key, Series) and key.dtype == np.bool_:
+            if not key.index.equals(self.index):
+                raise Exception('can only boolean index with like-indexed '
+                                'Series or raw ndarrays')
+
     def __setslice__(self, i, j, value):
         """Set slice equal to given value(s)"""
         ndarray.__setslice__(self, i, j, value)
@@ -743,12 +738,12 @@ def append(self, other):
         -------
         y : Series
         """
-        newIndex = np.concatenate((self.index, other.index))
-        newIndex = Index(newIndex)
-        newIndex._verify_integrity()
+        new_index = np.concatenate((self.index, other.index))
+        new_index = Index(new_index)
+        new_index._verify_integrity()
 
         new_values = np.concatenate((self, other))
-        return Series(new_values, index=newIndex)
+        return Series(new_values, index=new_index)
 
     def _binop(self, other, func, fill_value=None):
         """
@@ -811,17 +806,17 @@ def combine(self, other, func, fill_value=nan):
         result : Series
         """
         if isinstance(other, Series):
-            newIndex = self.index + other.index
+            new_index = self.index + other.index
 
-            new_values = np.empty(len(newIndex), dtype=self.dtype)
-            for i, idx in enumerate(newIndex):
+            new_values = np.empty(len(new_index), dtype=self.dtype)
+            for i, idx in enumerate(new_index):
                 new_values[i] = func(self.get(idx, fill_value),
                                  other.get(idx, fill_value))
         else:
-            newIndex = self.index
+            new_index = self.index
             new_values = func(self.values, other)
 
-        return Series(new_values, index=newIndex)
+        return Series(new_values, index=new_index)
 
     def combineFirst(self, other):
         """
@@ -837,16 +832,16 @@ def combineFirst(self, other):
             formed as union of two Series
         """
         if self.index.equals(other.index):
-            newIndex = self.index
+            new_index = self.index
             # save ourselves the copying in this case
             this = self
         else:
-            newIndex = self.index + other.index
+            new_index = self.index + other.index
 
-            this = self.reindex(newIndex)
-            other = other.reindex(newIndex)
+            this = self.reindex(new_index)
+            other = other.reindex(new_index)
 
-        result = Series(np.where(isnull(this), other, this), index=newIndex)
+        result = Series(np.where(isnull(this), other, this), index=new_index)
         return result
 
     #----------------------------------------------------------------------
@@ -896,7 +891,9 @@ def _try_mergesort(arr):
                 # stable sort not available for object dtype
                 return arr.argsort()
 
-        if 'missingAtEnd' in kwds:
+        if 'missingAtEnd' in kwds: # pragma: no cover
+            warnings.warn("missingAtEnd is deprecated, use na_last",
+                          FutureWarning)
             na_last = kwds['missingAtEnd']
 
         arr = self.values

pandas/tests/test_series.py

@@ -451,11 +451,11 @@ def test_operators_corner(self):
         result = empty + Series([], index=Index([]))
         self.assert_(len(result) == 0)
 
-        deltas = Series([timedelta(1)] * 5, index=np.arange(5))
-        sub_deltas = deltas[::2]
-
-        deltas5 = deltas * 5
-        deltas = deltas + sub_deltas
+        # TODO: this returned NotImplemented earlier, what to do?
+        # deltas = Series([timedelta(1)] * 5, index=np.arange(5))
+        # sub_deltas = deltas[::2]
+        # deltas5 = deltas * 5
+        # deltas = deltas + sub_deltas
 
         # float + int
         int_ts = self.ts.astype(int)[:-5]
@@ -627,7 +627,6 @@ def test_sort(self):
         self.assert_(np.array_equal(ts.index, self.ts.order().index))
 
     def test_order(self):
-
         ts = self.ts.copy()
         ts[:5] = np.NaN
         vals = ts.values
@@ -640,9 +639,6 @@ def test_order(self):
         self.assert_(np.isnan(result[:5]).all())
         self.assert_(np.array_equal(result[5:], np.sort(vals[5:])))
 
-        # just want to make sure it works
-        result = ts.order(missingAtEnd=False)
-
         # something object-type
         ser = Series(['A', 'B'], [1, 2])
         # no failure