Repository: pandas-dev/pandas

Issue:
API: expose `axis` keyword in pandas.core.algorithms.take
I haven't quite replaced cyberpandas' take with pandas' yet. I think one reason is that I need to be able to pass axis through to `pandas.core.algorithms.take`

We'll have to be careful. The default for `take_nd` is 0, while for `ndarray.take` it's None. I think we want the default to be 0.

I still have to figure out one other thing, but then hopefully things will work.



Discussion:

Resolving patch:
diff --git a/pandas/core/algorithms.py b/pandas/core/algorithms.py
--- a/pandas/core/algorithms.py
+++ b/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
 
 


Test patch:
diff --git a/pandas/tests/test_take.py b/pandas/tests/test_take.py
--- a/pandas/tests/test_take.py
+++ b/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