Skip to content

Commit

Permalink
BUG: enable multiple column selection in SparseDataFrame via __getite…
Browse files Browse the repository at this point in the history
…m__ close #1585
  • Loading branch information
wesm committed Jul 11, 2012
1 parent 77b49b9 commit 060ae8f
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 12 deletions.
1 change: 1 addition & 0 deletions RELEASE.rst
Expand Up @@ -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
============
Expand Down
25 changes: 14 additions & 11 deletions pandas/sparse/frame.py
Expand Up @@ -19,6 +19,7 @@

from pandas.sparse.series import SparseSeries
from pandas.util.decorators import Appender
import pandas.lib as lib


class _SparseMockBlockManager(object):
Expand Down Expand Up @@ -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

Expand Down
9 changes: 8 additions & 1 deletion pandas/sparse/tests/test_sparse.py
Expand Up @@ -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)
Expand Down

0 comments on commit 060ae8f

Please sign in to comment.