API: Added axis to take (pandas-dev#20999)

Author: TomAugspurger

pandas/core/algorithms.py

@@ -1448,7 +1448,7 @@ def func(arr, indexer, out, fill_value=np.nan):
     return func
 
 
-def take(arr, indices, allow_fill=False, fill_value=None):
+def take(arr, indices, axis=0, allow_fill=False, fill_value=None):
     """
     Take elements from an array.
 
@@ -1461,6 +1461,8 @@ def take(arr, indices, allow_fill=False, fill_value=None):
         to an ndarray.
     indices : sequence of integers
         Indices to be taken.
+    axis : int, default 0
+        The axis over which to select values.
     allow_fill : bool, default False
         How to handle negative values in `indices`.
 
@@ -1476,6 +1478,9 @@ def take(arr, indices, allow_fill=False, fill_value=None):
         This may be ``None``, in which case the default NA value for
         the type (``self.dtype.na_value``) is used.
 
+        For multi-dimensional `arr`, each *element* is filled with
+        `fill_value`.
+
     Returns
     -------
     ndarray or ExtensionArray
@@ -1529,10 +1534,11 @@ def take(arr, indices, allow_fill=False, fill_value=None):
     if allow_fill:
         # Pandas style, -1 means NA
         validate_indices(indices, len(arr))
-        result = take_1d(arr, indices, allow_fill=True, fill_value=fill_value)
+        result = take_1d(arr, indices, axis=axis, allow_fill=True,
+                         fill_value=fill_value)
     else:
         # NumPy style
-        result = arr.take(indices)
+        result = arr.take(indices, axis=axis)
     return result
 
 

pandas/tests/test_take.py

@@ -447,6 +447,29 @@ def test_2d_datetime64(self):
         expected[:, [2, 4]] = datetime(2007, 1, 1)
         tm.assert_almost_equal(result, expected)
 
+    def test_take_axis_0(self):
+        arr = np.arange(12).reshape(4, 3)
+        result = algos.take(arr, [0, -1])
+        expected = np.array([[0, 1, 2], [9, 10, 11]])
+        tm.assert_numpy_array_equal(result, expected)
+
+        # allow_fill=True
+        result = algos.take(arr, [0, -1], allow_fill=True, fill_value=0)
+        expected = np.array([[0, 1, 2], [0, 0, 0]])
+        tm.assert_numpy_array_equal(result, expected)
+
+    def test_take_axis_1(self):
+        arr = np.arange(12).reshape(4, 3)
+        result = algos.take(arr, [0, -1], axis=1)
+        expected = np.array([[0, 2], [3, 5], [6, 8], [9, 11]])
+        tm.assert_numpy_array_equal(result, expected)
+
+        # allow_fill=True
+        result = algos.take(arr, [0, -1], axis=1, allow_fill=True,
+                            fill_value=0)
+        expected = np.array([[0, 0], [3, 0], [6, 0], [9, 0]])
+        tm.assert_numpy_array_equal(result, expected)
+
 
 class TestExtensionTake(object):
     # The take method found in pd.api.extensions