From 3c19920695ead380c6a0e85acee5847d0f2ab4c4 Mon Sep 17 00:00:00 2001 From: sinhrks Date: Thu, 17 Mar 2016 01:40:07 +0900 Subject: [PATCH] API: Index.take inconsistently handle fill_value --- doc/source/whatsnew/v0.18.1.txt | 9 ++- pandas/indexes/base.py | 50 +++++++++++-- pandas/indexes/category.py | 22 ++---- pandas/indexes/multi.py | 37 ++++++++-- pandas/tests/indexes/test_base.py | 28 +++++++ pandas/tests/indexes/test_category.py | 97 +++++++++++++++++++++++++ pandas/tests/indexes/test_multi.py | 40 ++++++++++ pandas/tests/indexes/test_numeric.py | 55 ++++++++++++++ pandas/tests/indexes/test_range.py | 27 +++++++ pandas/tseries/base.py | 23 +++--- pandas/tseries/period.py | 8 -- pandas/tseries/tests/test_period.py | 32 ++++++++ pandas/tseries/tests/test_timedeltas.py | 32 ++++++++ pandas/tseries/tests/test_timeseries.py | 63 ++++++++++++++++ 14 files changed, 476 insertions(+), 47 deletions(-) diff --git a/doc/source/whatsnew/v0.18.1.txt b/doc/source/whatsnew/v0.18.1.txt index 1179a347e4c46..486d7f3a067bc 100644 --- a/doc/source/whatsnew/v0.18.1.txt +++ b/doc/source/whatsnew/v0.18.1.txt @@ -29,7 +29,6 @@ New features Enhancements ~~~~~~~~~~~~ - .. _whatsnew_0181.partial_string_indexing: Partial string indexing on ``DateTimeIndex`` when part of a ``MultiIndex`` @@ -59,6 +58,14 @@ Other Enhancements - ``pd.read_csv()`` now supports opening ZIP files that contains a single CSV, via extension inference or explict ``compression='zip'`` (:issue:`12175`) - ``pd.read_csv()`` now supports opening files using xz compression, via extension inference or explicit ``compression='xz'`` is specified; ``xz`` compressions is also supported by ``DataFrame.to_csv`` in the same way (:issue:`11852`) - ``pd.read_msgpack()`` now always gives writeable ndarrays even when compression is used (:issue:`12359`). +- ``Index.take`` now handles ``allow_fill`` and ``fill_value`` consistently (:issue:`12631`) + +.. ipython:: python + + idx = pd.Index([1., 2., 3., 4.], dtype='float') + idx.take([2, -1]) # default, allow_fill=True, fill_value=None + idx.take([2, -1], fill_value=True) + .. _whatsnew_0181.api: diff --git a/pandas/indexes/base.py b/pandas/indexes/base.py index 588521e6810b3..d5e0c71087fdf 100644 --- a/pandas/indexes/base.py +++ b/pandas/indexes/base.py @@ -1329,24 +1329,60 @@ def _ensure_compat_concat(indexes): return indexes - def take(self, indices, axis=0, allow_fill=True, fill_value=None): - """ - return a new Index of the values selected by the indexer + _index_shared_docs['take'] = """ + return a new Index of the values selected by the indices For internal compatibility with numpy arrays. - # filling must always be None/nan here - # but is passed thru internally + Parameters + ---------- + indices : list + Indices to be taken + axis : int, optional + The axis over which to select values, always 0. + allow_fill : bool, default True + fill_value : bool, default None + If allow_fill=True and fill_value is not None, indices specified by + -1 is regarded as NA. If Index doesn't hold NA, raise ValueError See also -------- numpy.ndarray.take """ - + @Appender(_index_shared_docs['take']) + def take(self, indices, axis=0, allow_fill=True, fill_value=None): indices = com._ensure_platform_int(indices) - taken = self.values.take(indices) + if self._can_hold_na: + taken = self._assert_take_fillable(self.values, indices, + allow_fill=allow_fill, + fill_value=fill_value, + na_value=self._na_value) + else: + if allow_fill and fill_value is not None: + msg = 'Unable to fill values because {0} cannot contain NA' + raise ValueError(msg.format(self.__class__.__name__)) + taken = self.values.take(indices) return self._shallow_copy(taken) + def _assert_take_fillable(self, values, indices, allow_fill=True, + fill_value=None, na_value=np.nan): + """ Internal method to handle NA filling of take """ + indices = com._ensure_platform_int(indices) + + # only fill if we are passing a non-None fill_value + if allow_fill and fill_value is not None: + if (indices < -1).any(): + msg = ('When allow_fill=True and fill_value is not None, ' + 'all indices must be >= -1') + raise ValueError(msg) + taken = values.take(indices) + mask = indices == -1 + if mask.any(): + taken[mask] = na_value + else: + taken = values.take(indices) + return taken + @cache_readonly def _isnan(self): """ return if each value is nan""" diff --git a/pandas/indexes/category.py b/pandas/indexes/category.py index a9390c76f26a6..5844c69c57f0e 100644 --- a/pandas/indexes/category.py +++ b/pandas/indexes/category.py @@ -459,21 +459,13 @@ def _convert_list_indexer(self, keyarr, kind=None): return None - def take(self, indexer, axis=0, allow_fill=True, fill_value=None): - """ - For internal compatibility with numpy arrays. - - # filling must always be None/nan here - # but is passed thru internally - assert isnull(fill_value) - - See also - -------- - numpy.ndarray.take - """ - - indexer = com._ensure_platform_int(indexer) - taken = self.codes.take(indexer) + @Appender(_index_shared_docs['take']) + def take(self, indices, axis=0, allow_fill=True, fill_value=None): + indices = com._ensure_platform_int(indices) + taken = self._assert_take_fillable(self.codes, indices, + allow_fill=allow_fill, + fill_value=fill_value, + na_value=-1) return self._create_from_codes(taken) def delete(self, loc): diff --git a/pandas/indexes/multi.py b/pandas/indexes/multi.py index 4a77282b3877a..49a9c3d622c1b 100644 --- a/pandas/indexes/multi.py +++ b/pandas/indexes/multi.py @@ -1,3 +1,4 @@ + # pylint: disable=E1101,E1103,W0232 import datetime import warnings @@ -11,7 +12,7 @@ from pandas.compat import range, zip, lrange, lzip, map from pandas import compat -from pandas.core.base import FrozenList +from pandas.core.base import FrozenList, FrozenNDArray import pandas.core.base as base from pandas.util.decorators import (Appender, cache_readonly, deprecate, deprecate_kwarg) @@ -1003,12 +1004,38 @@ def __getitem__(self, key): names=self.names, sortorder=sortorder, verify_integrity=False) - def take(self, indexer, axis=None): - indexer = com._ensure_platform_int(indexer) - new_labels = [lab.take(indexer) for lab in self.labels] - return MultiIndex(levels=self.levels, labels=new_labels, + @Appender(_index_shared_docs['take']) + def take(self, indices, axis=0, allow_fill=True, fill_value=None): + indices = com._ensure_platform_int(indices) + taken = self._assert_take_fillable(self.labels, indices, + allow_fill=allow_fill, + fill_value=fill_value, + na_value=-1) + return MultiIndex(levels=self.levels, labels=taken, names=self.names, verify_integrity=False) + def _assert_take_fillable(self, values, indices, allow_fill=True, + fill_value=None, na_value=None): + """ Internal method to handle NA filling of take """ + # only fill if we are passing a non-None fill_value + if allow_fill and fill_value is not None: + if (indices < -1).any(): + msg = ('When allow_fill=True and fill_value is not None, ' + 'all indices must be >= -1') + raise ValueError(msg) + taken = [lab.take(indices) for lab in self.labels] + mask = indices == -1 + if mask.any(): + masked = [] + for new_label in taken: + label_values = new_label.values() + label_values[mask] = na_value + masked.append(base.FrozenNDArray(label_values)) + taken = masked + else: + taken = [lab.take(indices) for lab in self.labels] + return taken + def append(self, other): """ Append a collection of Index options together diff --git a/pandas/tests/indexes/test_base.py b/pandas/tests/indexes/test_base.py index 24a20f7adf624..41933247961e9 100644 --- a/pandas/tests/indexes/test_base.py +++ b/pandas/tests/indexes/test_base.py @@ -1240,6 +1240,34 @@ def test_nan_first_take_datetime(self): exp = Index([idx[-1], idx[0], idx[1]]) tm.assert_index_equal(res, exp) + def test_take_fill_value(self): + # GH 12631 + idx = pd.Index(list('ABC'), name='xxx') + result = idx.take(np.array([1, 0, -1])) + expected = pd.Index(list('BAC'), name='xxx') + tm.assert_index_equal(result, expected) + + # fill_value + result = idx.take(np.array([1, 0, -1]), fill_value=True) + expected = pd.Index(['B', 'A', np.nan], name='xxx') + tm.assert_index_equal(result, expected) + + # allow_fill=False + result = idx.take(np.array([1, 0, -1]), allow_fill=False, + fill_value=True) + expected = pd.Index(['B', 'A', 'C'], name='xxx') + tm.assert_index_equal(result, expected) + + msg = ('When allow_fill=True and fill_value is not None, ' + 'all indices must be >= -1') + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -2]), fill_value=True) + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -5]), fill_value=True) + + with tm.assertRaises(IndexError): + idx.take(np.array([1, -5])) + def test_reindex_preserves_name_if_target_is_list_or_ndarray(self): # GH6552 idx = pd.Index([0, 1, 2]) diff --git a/pandas/tests/indexes/test_category.py b/pandas/tests/indexes/test_category.py index d0aae9ba281f3..a8534309c115c 100644 --- a/pandas/tests/indexes/test_category.py +++ b/pandas/tests/indexes/test_category.py @@ -708,3 +708,100 @@ def test_fillna_categorical(self): with tm.assertRaisesRegexp(ValueError, 'fill value must be in categories'): idx.fillna(2.0) + + def test_take_fill_value(self): + # GH 12631 + + # numeric category + idx = pd.CategoricalIndex([1, 2, 3], name='xxx') + result = idx.take(np.array([1, 0, -1])) + expected = pd.CategoricalIndex([2, 1, 3], name='xxx') + tm.assert_index_equal(result, expected) + tm.assert_categorical_equal(result.values, expected.values) + + # fill_value + result = idx.take(np.array([1, 0, -1]), fill_value=True) + expected = pd.CategoricalIndex([2, 1, np.nan], categories=[1, 2, 3], + name='xxx') + tm.assert_index_equal(result, expected) + tm.assert_categorical_equal(result.values, expected.values) + + # allow_fill=False + result = idx.take(np.array([1, 0, -1]), allow_fill=False, + fill_value=True) + expected = pd.CategoricalIndex([2, 1, 3], name='xxx') + tm.assert_index_equal(result, expected) + tm.assert_categorical_equal(result.values, expected.values) + + # object category + idx = pd.CategoricalIndex(list('CBA'), categories=list('ABC'), + ordered=True, name='xxx') + result = idx.take(np.array([1, 0, -1])) + expected = pd.CategoricalIndex(list('BCA'), categories=list('ABC'), + ordered=True, name='xxx') + tm.assert_index_equal(result, expected) + tm.assert_categorical_equal(result.values, expected.values) + + # fill_value + result = idx.take(np.array([1, 0, -1]), fill_value=True) + expected = pd.CategoricalIndex(['B', 'C', np.nan], + categories=list('ABC'), ordered=True, + name='xxx') + tm.assert_index_equal(result, expected) + tm.assert_categorical_equal(result.values, expected.values) + + # allow_fill=False + result = idx.take(np.array([1, 0, -1]), allow_fill=False, + fill_value=True) + expected = pd.CategoricalIndex(list('BCA'), categories=list('ABC'), + ordered=True, name='xxx') + tm.assert_index_equal(result, expected) + tm.assert_categorical_equal(result.values, expected.values) + + msg = ('When allow_fill=True and fill_value is not None, ' + 'all indices must be >= -1') + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -2]), fill_value=True) + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -5]), fill_value=True) + + with tm.assertRaises(IndexError): + idx.take(np.array([1, -5])) + + def test_take_fill_value_datetime(self): + + # datetime category + idx = pd.DatetimeIndex(['2011-01-01', '2011-02-01', '2011-03-01'], + name='xxx') + idx = pd.CategoricalIndex(idx) + result = idx.take(np.array([1, 0, -1])) + expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', '2011-03-01'], + name='xxx') + expected = pd.CategoricalIndex(expected) + tm.assert_index_equal(result, expected) + + # fill_value + result = idx.take(np.array([1, 0, -1]), fill_value=True) + expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', 'NaT'], + name='xxx') + exp_cats = pd.DatetimeIndex(['2011-01-01', '2011-02-01', '2011-03-01']) + expected = pd.CategoricalIndex(expected, categories=exp_cats) + tm.assert_index_equal(result, expected) + + # allow_fill=False + result = idx.take(np.array([1, 0, -1]), allow_fill=False, + fill_value=True) + expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', '2011-03-01'], + name='xxx') + expected = pd.CategoricalIndex(expected) + tm.assert_index_equal(result, expected) + + msg = ('When allow_fill=True and fill_value is not None, ' + 'all indices must be >= -1') + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -2]), fill_value=True) + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -5]), fill_value=True) + + with tm.assertRaises(IndexError): + idx.take(np.array([1, -5])) diff --git a/pandas/tests/indexes/test_multi.py b/pandas/tests/indexes/test_multi.py index f70ea49bd4c29..f37e62e52b618 100644 --- a/pandas/tests/indexes/test_multi.py +++ b/pandas/tests/indexes/test_multi.py @@ -1514,6 +1514,46 @@ def test_take_preserve_name(self): taken = self.index.take([3, 0, 1]) self.assertEqual(taken.names, self.index.names) + def test_take_fill_value(self): + # GH 12631 + vals = [['A', 'B'], + [pd.Timestamp('2011-01-01'), pd.Timestamp('2011-01-02')]] + idx = pd.MultiIndex.from_product(vals, names=['str', 'dt']) + + result = idx.take(np.array([1, 0, -1])) + exp_vals = [('A', pd.Timestamp('2011-01-02')), + ('A', pd.Timestamp('2011-01-01')), + ('B', pd.Timestamp('2011-01-02'))] + expected = pd.MultiIndex.from_tuples(exp_vals, names=['str', 'dt']) + tm.assert_index_equal(result, expected) + + # fill_value + result = idx.take(np.array([1, 0, -1]), fill_value=True) + exp_vals = [('A', pd.Timestamp('2011-01-02')), + ('A', pd.Timestamp('2011-01-01')), + (np.nan, pd.NaT)] + expected = pd.MultiIndex.from_tuples(exp_vals, names=['str', 'dt']) + tm.assert_index_equal(result, expected) + + # allow_fill=False + result = idx.take(np.array([1, 0, -1]), allow_fill=False, + fill_value=True) + exp_vals = [('A', pd.Timestamp('2011-01-02')), + ('A', pd.Timestamp('2011-01-01')), + ('B', pd.Timestamp('2011-01-02'))] + expected = pd.MultiIndex.from_tuples(exp_vals, names=['str', 'dt']) + tm.assert_index_equal(result, expected) + + msg = ('When allow_fill=True and fill_value is not None, ' + 'all indices must be >= -1') + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -2]), fill_value=True) + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -5]), fill_value=True) + + with tm.assertRaises(IndexError): + idx.take(np.array([1, -5])) + def test_join_level(self): def _check_how(other, how): join_index, lidx, ridx = other.join(self.index, how=how, diff --git a/pandas/tests/indexes/test_numeric.py b/pandas/tests/indexes/test_numeric.py index 325e0df14a07e..56e19c8b1d7d9 100644 --- a/pandas/tests/indexes/test_numeric.py +++ b/pandas/tests/indexes/test_numeric.py @@ -343,6 +343,34 @@ def test_fillna_float64(self): exp = Index([1.0, 'obj', 3.0], name='x') self.assert_index_equal(idx.fillna('obj'), exp) + def test_take_fill_value(self): + # GH 12631 + idx = pd.Float64Index([1., 2., 3.], name='xxx') + result = idx.take(np.array([1, 0, -1])) + expected = pd.Float64Index([2., 1., 3.], name='xxx') + tm.assert_index_equal(result, expected) + + # fill_value + result = idx.take(np.array([1, 0, -1]), fill_value=True) + expected = pd.Float64Index([2., 1., np.nan], name='xxx') + tm.assert_index_equal(result, expected) + + # allow_fill=False + result = idx.take(np.array([1, 0, -1]), allow_fill=False, + fill_value=True) + expected = pd.Float64Index([2., 1., 3.], name='xxx') + tm.assert_index_equal(result, expected) + + msg = ('When allow_fill=True and fill_value is not None, ' + 'all indices must be >= -1') + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -2]), fill_value=True) + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -5]), fill_value=True) + + with tm.assertRaises(IndexError): + idx.take(np.array([1, -5])) + class TestInt64Index(Numeric, tm.TestCase): _holder = Int64Index @@ -757,6 +785,33 @@ def test_take_preserve_name(self): taken = index.take([3, 0, 1]) self.assertEqual(index.name, taken.name) + def test_take_fill_value(self): + # GH 12631 + idx = pd.Int64Index([1, 2, 3], name='xxx') + result = idx.take(np.array([1, 0, -1])) + expected = pd.Int64Index([2, 1, 3], name='xxx') + tm.assert_index_equal(result, expected) + + # fill_value + msg = "Unable to fill values because Int64Index cannot contain NA" + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -1]), fill_value=True) + + # allow_fill=False + result = idx.take(np.array([1, 0, -1]), allow_fill=False, + fill_value=True) + expected = pd.Int64Index([2, 1, 3], name='xxx') + tm.assert_index_equal(result, expected) + + msg = "Unable to fill values because Int64Index cannot contain NA" + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -2]), fill_value=True) + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -5]), fill_value=True) + + with tm.assertRaises(IndexError): + idx.take(np.array([1, -5])) + def test_int_name_format(self): index = Index(['a', 'b', 'c'], name=0) s = Series(lrange(3), index) diff --git a/pandas/tests/indexes/test_range.py b/pandas/tests/indexes/test_range.py index 84ed1049936d3..f41c252f44d39 100644 --- a/pandas/tests/indexes/test_range.py +++ b/pandas/tests/indexes/test_range.py @@ -647,6 +647,33 @@ def test_take_preserve_name(self): taken = index.take([3, 0, 1]) self.assertEqual(index.name, taken.name) + def test_take_fill_value(self): + # GH 12631 + idx = pd.RangeIndex(1, 4, name='xxx') + result = idx.take(np.array([1, 0, -1])) + expected = pd.Int64Index([2, 1, 3], name='xxx') + tm.assert_index_equal(result, expected) + + # fill_value + msg = "Unable to fill values because RangeIndex cannot contain NA" + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -1]), fill_value=True) + + # allow_fill=False + result = idx.take(np.array([1, 0, -1]), allow_fill=False, + fill_value=True) + expected = pd.Int64Index([2, 1, 3], name='xxx') + tm.assert_index_equal(result, expected) + + msg = "Unable to fill values because RangeIndex cannot contain NA" + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -2]), fill_value=True) + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -5]), fill_value=True) + + with tm.assertRaises(IndexError): + idx.take(np.array([1, -5])) + def test_print_unicode_columns(self): df = pd.DataFrame({u("\u05d0"): [1, 2, 3], "\u05d1": [4, 5, 6], diff --git a/pandas/tseries/base.py b/pandas/tseries/base.py index 7584b99dbdb97..48e17fd84a3b2 100644 --- a/pandas/tseries/base.py +++ b/pandas/tseries/base.py @@ -12,6 +12,7 @@ import pandas.tslib as tslib import pandas.lib as lib from pandas.core.index import Index +from pandas.indexes.base import _index_shared_docs from pandas.util.decorators import Appender, cache_readonly import pandas.tseries.frequencies as frequencies import pandas.algos as _algos @@ -258,22 +259,22 @@ def sort_values(self, return_indexer=False, ascending=True): return self._simple_new(sorted_values, **attribs) + @Appender(_index_shared_docs['take']) def take(self, indices, axis=0, allow_fill=True, fill_value=None): - """ - Analogous to ndarray.take - """ indices = com._ensure_int64(indices) + maybe_slice = lib.maybe_indices_to_slice(indices, len(self)) if isinstance(maybe_slice, slice): return self[maybe_slice] - taken = self.asi8.take(com._ensure_platform_int(indices)) - - # only fill if we are passing a non-None fill_value - if allow_fill and fill_value is not None: - mask = indices == -1 - if mask.any(): - taken[mask] = tslib.iNaT - return self._shallow_copy(taken, freq=None) + + taken = self._assert_take_fillable(self.asi8, indices, + allow_fill=allow_fill, + fill_value=fill_value, + na_value=tslib.iNaT) + + # keep freq in PeriodIndex, reset otherwise + freq = self.freq if isinstance(self, com.ABCPeriodIndex) else None + return self._shallow_copy(taken, freq=freq) def get_duplicates(self): values = Index.get_duplicates(self) diff --git a/pandas/tseries/period.py b/pandas/tseries/period.py index b34af4e62845b..e9a9796f9c48d 100644 --- a/pandas/tseries/period.py +++ b/pandas/tseries/period.py @@ -851,14 +851,6 @@ def _format_native_types(self, na_rep=u('NaT'), date_format=None, values[imask] = np.array([formatter(dt) for dt in values[imask]]) return values - def take(self, indices, axis=0, allow_fill=True, fill_value=None): - """ - Analogous to ndarray.take - """ - indices = com._ensure_platform_int(indices) - taken = self.asi8.take(indices, axis=axis) - return self._simple_new(taken, self.name, freq=self.freq) - def append(self, other): """ Append a collection of Index options together diff --git a/pandas/tseries/tests/test_period.py b/pandas/tseries/tests/test_period.py index e0dad2995f91c..947a200e3b63d 100644 --- a/pandas/tseries/tests/test_period.py +++ b/pandas/tseries/tests/test_period.py @@ -2815,6 +2815,38 @@ def test_take(self): self.assertEqual(taken.freq, index.freq) self.assertEqual(taken.name, expected.name) + def test_take_fill_value(self): + # GH 12631 + idx = pd.PeriodIndex(['2011-01-01', '2011-02-01', '2011-03-01'], + name='xxx', freq='D') + result = idx.take(np.array([1, 0, -1])) + expected = pd.PeriodIndex(['2011-02-01', '2011-01-01', '2011-03-01'], + name='xxx', freq='D') + tm.assert_index_equal(result, expected) + + # fill_value + result = idx.take(np.array([1, 0, -1]), fill_value=True) + expected = pd.PeriodIndex(['2011-02-01', '2011-01-01', 'NaT'], + name='xxx', freq='D') + tm.assert_index_equal(result, expected) + + # allow_fill=False + result = idx.take(np.array([1, 0, -1]), allow_fill=False, + fill_value=True) + expected = pd.PeriodIndex(['2011-02-01', '2011-01-01', '2011-03-01'], + name='xxx', freq='D') + tm.assert_index_equal(result, expected) + + msg = ('When allow_fill=True and fill_value is not None, ' + 'all indices must be >= -1') + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -2]), fill_value=True) + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -5]), fill_value=True) + + with tm.assertRaises(IndexError): + idx.take(np.array([1, -5])) + def test_joins(self): index = period_range('1/1/2000', '1/20/2000', freq='D') diff --git a/pandas/tseries/tests/test_timedeltas.py b/pandas/tseries/tests/test_timedeltas.py index 4bdd0ed462852..e285856eca3f7 100644 --- a/pandas/tseries/tests/test_timedeltas.py +++ b/pandas/tseries/tests/test_timedeltas.py @@ -1613,6 +1613,38 @@ def test_take(self): self.assertIsNone(taken.freq) self.assertEqual(taken.name, expected.name) + def test_take_fill_value(self): + # GH 12631 + idx = pd.TimedeltaIndex(['1 days', '2 days', '3 days'], + name='xxx') + result = idx.take(np.array([1, 0, -1])) + expected = pd.TimedeltaIndex(['2 days', '1 days', '3 days'], + name='xxx') + tm.assert_index_equal(result, expected) + + # fill_value + result = idx.take(np.array([1, 0, -1]), fill_value=True) + expected = pd.TimedeltaIndex(['2 days', '1 days', 'NaT'], + name='xxx') + tm.assert_index_equal(result, expected) + + # allow_fill=False + result = idx.take(np.array([1, 0, -1]), allow_fill=False, + fill_value=True) + expected = pd.TimedeltaIndex(['2 days', '1 days', '3 days'], + name='xxx') + tm.assert_index_equal(result, expected) + + msg = ('When allow_fill=True and fill_value is not None, ' + 'all indices must be >= -1') + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -2]), fill_value=True) + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -5]), fill_value=True) + + with tm.assertRaises(IndexError): + idx.take(np.array([1, -5])) + def test_isin(self): index = tm.makeTimedeltaIndex(4) diff --git a/pandas/tseries/tests/test_timeseries.py b/pandas/tseries/tests/test_timeseries.py index 615390b5209b6..8ed3ef93cbdb4 100644 --- a/pandas/tseries/tests/test_timeseries.py +++ b/pandas/tseries/tests/test_timeseries.py @@ -3297,6 +3297,69 @@ def test_take(self): self.assertEqual(taken.tz, expected.tz) self.assertEqual(taken.name, expected.name) + def test_take_fill_value(self): + # GH 12631 + idx = pd.DatetimeIndex(['2011-01-01', '2011-02-01', '2011-03-01'], + name='xxx') + result = idx.take(np.array([1, 0, -1])) + expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', '2011-03-01'], + name='xxx') + tm.assert_index_equal(result, expected) + + # fill_value + result = idx.take(np.array([1, 0, -1]), fill_value=True) + expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', 'NaT'], + name='xxx') + tm.assert_index_equal(result, expected) + + # allow_fill=False + result = idx.take(np.array([1, 0, -1]), allow_fill=False, + fill_value=True) + expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', '2011-03-01'], + name='xxx') + tm.assert_index_equal(result, expected) + + msg = ('When allow_fill=True and fill_value is not None, ' + 'all indices must be >= -1') + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -2]), fill_value=True) + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -5]), fill_value=True) + + with tm.assertRaises(IndexError): + idx.take(np.array([1, -5])) + + def test_take_fill_value_with_timezone(self): + idx = pd.DatetimeIndex(['2011-01-01', '2011-02-01', '2011-03-01'], + name='xxx', tz='US/Eastern') + result = idx.take(np.array([1, 0, -1])) + expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', '2011-03-01'], + name='xxx', tz='US/Eastern') + tm.assert_index_equal(result, expected) + + # fill_value + result = idx.take(np.array([1, 0, -1]), fill_value=True) + expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', 'NaT'], + name='xxx', tz='US/Eastern') + tm.assert_index_equal(result, expected) + + # allow_fill=False + result = idx.take(np.array([1, 0, -1]), allow_fill=False, + fill_value=True) + expected = pd.DatetimeIndex(['2011-02-01', '2011-01-01', '2011-03-01'], + name='xxx', tz='US/Eastern') + tm.assert_index_equal(result, expected) + + msg = ('When allow_fill=True and fill_value is not None, ' + 'all indices must be >= -1') + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -2]), fill_value=True) + with tm.assertRaisesRegexp(ValueError, msg): + idx.take(np.array([1, 0, -5]), fill_value=True) + + with tm.assertRaises(IndexError): + idx.take(np.array([1, -5])) + def test_map_bug_1677(self): index = DatetimeIndex(['2012-04-25 09:30:00.393000']) f = index.asof