BUG: enable multiple column selection in SparseDataFrame via __getitem__ close #1585

Author: wesm

RELEASE.rst

@@ -56,6 +56,7 @@ pandas 0.8.1
     entry (#1586)
   - Fix resampling of tz-aware time series with "anchored" freq (#1591)
   - Fix DataFrame.rank error on integer data (#1589)
+  - Selection of multiple SparseDataFrame columns by list in __getitem__ (#1585)
 
 pandas 0.8.0
 ============

pandas/sparse/frame.py

@@ -19,6 +19,7 @@
 
 from pandas.sparse.series import SparseSeries
 from pandas.util.decorators import Appender
+import pandas.lib as lib
 
 
 class _SparseMockBlockManager(object):
@@ -314,26 +315,28 @@ def _set_columns(self, cols):
     index = property(fget=_get_index, fset=_set_index)
     columns = property(fget=_get_columns, fset=_set_columns)
 
-    def __getitem__(self, item):
+    def __getitem__(self, key):
         """
         Retrieve column or slice from DataFrame
         """
         try:
             # unsure about how kludgy this is
-            s = self._series[item]
-            s.name = item
+            s = self._series[key]
+            s.name = key
             return s
         except (TypeError, KeyError):
-            if isinstance(item, slice):
-                date_rng = self.index[item]
+            if isinstance(key, slice):
+                date_rng = self.index[key]
                 return self.reindex(date_rng)
 
-            elif isinstance(item, np.ndarray):
-                if len(item) != len(self.index):
-                    raise Exception('Item wrong length %d instead of %d!' %
-                                    (len(item), len(self.index)))
-                newIndex = self.index[item]
-                return self.reindex(newIndex)
+            elif isinstance(key, (np.ndarray, list)):
+                if isinstance(key, list):
+                    key = lib.list_to_object_array(key)
+
+                # also raises Exception if object array with NA values
+                if com._is_bool_indexer(key):
+                    key = np.asarray(key, dtype=bool)
+                return self._getitem_array(key)
             else: # pragma: no cover
                 raise
 

pandas/sparse/tests/test_sparse.py

@@ -892,7 +892,14 @@ def test_scalar_ops(self):
         pass
 
     def test_getitem(self):
-        pass
+        # #1585 select multiple columns
+        sdf = SparseDataFrame(index=[0, 1, 2], columns=['a', 'b','c'])
+
+        result = sdf[['a', 'b']]
+        exp = sdf.reindex(columns=['a', 'b'])
+        assert_sp_frame_equal(result, exp)
+
+        self.assertRaises(Exception, sdf.__getitem__, ['a', 'd'])
 
     def test_set_value(self):
         res = self.frame.set_value('foobar', 'B', 1.5)