diff --git a/pandas/core/algorithms.py b/pandas/core/algorithms.py index 796e5c8c6eb91..d456e2542402a 100644 --- a/pandas/core/algorithms.py +++ b/pandas/core/algorithms.py @@ -1455,19 +1455,21 @@ def take(arr, indexer, allow_fill=False, fill_value=None): Parameters ---------- - arr : ndarray or ExtensionArray + arr : sequence + Non array-likes (sequences without a dtype) are coereced + to an ndarray. indexer : sequence of integers - Indices to be taken. See Notes for how negative indicies - are handled. + Indices to be taken. allow_fill : bool, default False How to handle negative values in `indexer`. - For False values (the default), negative values in `indexer` - indiciate slices from the right. + * False: negative values in `indexer` indicate + slices from the right (the default) + + * True: negative values in `indexer` indicate + missing values. These values are set to `fill_value`. Any other + other negative values raise a ``ValueError``. - For True values, indicies where `indexer` is ``-1`` indicate - missing values. These values are set to `fill_value`. Any other - other negative value should raise a ``ValueError``. fill_value : any, optional Fill value to use for NA-indicies when `allow_fill` is True. This may be ``None``, in which case the default NA value for @@ -1486,13 +1488,42 @@ def take(arr, indexer, allow_fill=False, fill_value=None): When the indexer contains negative values other than ``-1`` and `allow_fill` is True. + Notes + ----- + When `allow_fill` is False, `indexer` may be whatever dimensionality + is accepted by NumPy for `arr`. + + When `allow_fill` is True, `indexer` should be 1-D. + See Also -------- numpy.take + + Examples + -------- + >>> from pandas.api.extensions import take + + With the default ``allow_fill=False``, negative numbers indicate + slices from the right. + + >>> take(np.array([10, 20, 30]), [0, 0, -1]) + array([10, 10, 30]) + + Setting ``allow_fill=True`` will place `fill_value` in those positions. + + >>> take(np.array([10, 20, 30]), [0, 0, -1], allow_fill=True) + array([10., 10., nan]) + + >>> take(np.array([10, 20, 30]), [0, 0, -1], allow_fill=True, + ... fill_value=-10) + array([ 10, 10, -10]) """ from pandas.core.indexing import validate_indices - # Do we require int64 here? + if not is_array_like(arr): + arr = np.asarray(arr) + + # Do we require int64 or intp here? indexer = np.asarray(indexer, dtype='int') if allow_fill: diff --git a/pandas/tests/test_algos.py b/pandas/tests/test_algos.py index 3c8c91ab30d3f..8a8a6f7de70d7 100644 --- a/pandas/tests/test_algos.py +++ b/pandas/tests/test_algos.py @@ -1564,40 +1564,3 @@ def test_index(self): idx = Index(['1 day', '1 day', '-1 day', '-1 day 2 min', '2 min', '2 min'], dtype='timedelta64[ns]') tm.assert_series_equal(algos.mode(idx), exp) - - -class TestTake(object): - - def test_bounds_check_large(self): - arr = np.array([1, 2]) - with pytest.raises(IndexError): - algos.take(arr, [2, 3], allow_fill=True) - - with pytest.raises(IndexError): - algos.take(arr, [2, 3], allow_fill=False) - - def test_bounds_check_small(self): - arr = np.array([1, 2, 3], dtype=np.int64) - indexer = [0, -1, -2] - with pytest.raises(ValueError): - algos.take(arr, indexer, allow_fill=True) - - result = algos.take(arr, indexer) - expected = np.array([1, 3, 2], dtype=np.int64) - tm.assert_numpy_array_equal(result, expected) - - @pytest.mark.parametrize('allow_fill', [True, False]) - def test_take_empty(self, allow_fill): - arr = np.array([], dtype=np.int64) - # empty take is ok - result = algos.take(arr, [], allow_fill=allow_fill) - tm.assert_numpy_array_equal(arr, result) - - with pytest.raises(IndexError): - algos.take(arr, [0], allow_fill=allow_fill) - - def test_take_na_empty(self): - result = algos.take(np.array([]), [-1, -1], allow_fill=True, - fill_value=0.0) - expected = np.array([0., 0.]) - tm.assert_numpy_array_equal(result, expected) diff --git a/pandas/tests/test_take.py b/pandas/tests/test_take.py index 7b97b0e975df3..2b78c91f9dac5 100644 --- a/pandas/tests/test_take.py +++ b/pandas/tests/test_take.py @@ -3,6 +3,7 @@ from datetime import datetime import numpy as np +import pytest from pandas.compat import long import pandas.core.algorithms as algos import pandas.util.testing as tm @@ -445,3 +446,47 @@ def test_2d_datetime64(self): expected = arr.take(indexer, axis=1) expected[:, [2, 4]] = datetime(2007, 1, 1) tm.assert_almost_equal(result, expected) + + +class TestExtensionTake(object): + # The take method found in pd.api.extensions + + def test_bounds_check_large(self): + arr = np.array([1, 2]) + with pytest.raises(IndexError): + algos.take(arr, [2, 3], allow_fill=True) + + with pytest.raises(IndexError): + algos.take(arr, [2, 3], allow_fill=False) + + def test_bounds_check_small(self): + arr = np.array([1, 2, 3], dtype=np.int64) + indexer = [0, -1, -2] + with pytest.raises(ValueError): + algos.take(arr, indexer, allow_fill=True) + + result = algos.take(arr, indexer) + expected = np.array([1, 3, 2], dtype=np.int64) + tm.assert_numpy_array_equal(result, expected) + + @pytest.mark.parametrize('allow_fill', [True, False]) + def test_take_empty(self, allow_fill): + arr = np.array([], dtype=np.int64) + # empty take is ok + result = algos.take(arr, [], allow_fill=allow_fill) + tm.assert_numpy_array_equal(arr, result) + + with pytest.raises(IndexError): + algos.take(arr, [0], allow_fill=allow_fill) + + def test_take_na_empty(self): + result = algos.take(np.array([]), [-1, -1], allow_fill=True, + fill_value=0.0) + expected = np.array([0., 0.]) + tm.assert_numpy_array_equal(result, expected) + + def test_take_coerces_list(self): + arr = [1, 2, 3] + result = algos.take(arr, [0, 0]) + expected = np.array([1, 1]) + tm.assert_numpy_array_equal(result, expected)