Skip to content

Commit

Permalink
Doc and move tests
Browse files Browse the repository at this point in the history
  • Loading branch information
TomAugspurger committed Apr 26, 2018
1 parent 74b2c09 commit 5db6624
Show file tree
Hide file tree
Showing 3 changed files with 85 additions and 46 deletions.
49 changes: 40 additions & 9 deletions pandas/core/algorithms.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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:
Expand Down
37 changes: 0 additions & 37 deletions pandas/tests/test_algos.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
45 changes: 45 additions & 0 deletions pandas/tests/test_take.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

0 comments on commit 5db6624

Please sign in to comment.