BUG/ENH: reindex boolean Series passed to getitem/setitem, GH #429

Author: wesm

RELEASE.rst

@@ -31,6 +31,8 @@ pandas 0.6.1
 
   - Rename `names` argument in DataFrame.from_records to `columns`. Add
     deprecation warning
+  - Boolean get/set operations on Series with boolean Series will reindex
+    instead of requiring that the indexes be exactly equal (GH #429)
 
 **New features / modules**
 

pandas/core/series.py

@@ -275,7 +275,7 @@ def _index_with(indexer):
         # special handling of boolean data with NAs stored in object
         # arrays. Since we can't represent NA with dtype=bool
         if _is_bool_indexer(key):
-            self._check_bool_indexer(key)
+            key = self._check_bool_indexer(key)
             key = np.asarray(key, dtype=bool)
             return _index_with(key)
 
@@ -392,17 +392,12 @@ def __setitem__(self, key, value):
             # Could not hash item
             pass
 
-        self._check_bool_indexer(key)
+        key = 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
         # NA. Hmm
         if isinstance(key, np.ndarray) and key.dtype == np.object_:
-            mask = isnull(key)
-            if mask.any():
-                raise ValueError('cannot index with vector containing '
-                                 'NA / NaN values')
-
             if set([True, False]).issubset(set(key)):
                 key = np.asarray(key, dtype=bool)
                 values[key] = value
@@ -413,10 +408,18 @@ def __setitem__(self, key, value):
     def _check_bool_indexer(self, key):
         # boolean indexing, need to check that the data are aligned, otherwise
         # disallowed
+        result = key
         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')
+                result = key.reindex(self.index)
+
+        if isinstance(result, np.ndarray) and result.dtype == np.object_:
+            mask = isnull(result)
+            if mask.any():
+                raise ValueError('cannot index with vector containing '
+                                 'NA / NaN values')
+
+        return result
 
     def __setslice__(self, i, j, value):
         """Set slice equal to given value(s)"""

pandas/tests/test_series.py

@@ -477,6 +477,23 @@ def test_ix_setitem_corner(self):
         self.assertRaises(Exception, self.series.ix.__setitem__,
                           inds + ['foo'], 5)
 
+    def test_get_set_boolean_different_order(self):
+        ordered = self.series.order()
+
+        # setting
+        copy = self.series.copy()
+        copy[ordered > 0] = 0
+
+        expected = self.series.copy()
+        expected[expected > 0] = 0
+
+        assert_series_equal(copy, expected)
+
+        # getting
+        sel = self.series[ordered > 0]
+        exp = self.series[self.series > 0]
+        assert_series_equal(sel, exp)
+
     def test_repr(self):
         str(self.ts)
         str(self.series)