DOC/REF: add docstrings and DRY it up

Author: cpcloud

doc/source/v0.13.1.txt

@@ -128,7 +128,6 @@ API changes
       import pandas.core.common as com
       com.array_equivalent(np.array([0, np.nan]), np.array([0, np.nan]))
       np.array_equal(np.array([0, np.nan]), np.array([0, np.nan]))
-- Add nsmallest and nlargest Series methods (:issue:`3960`)
 
 - ``DataFrame.apply`` will use the ``reduce`` argument to determine whether a
   ``Series`` or a ``DataFrame`` should be returned when the ``DataFrame`` is

doc/source/v0.14.0.txt

@@ -643,6 +643,7 @@ Enhancements
                 values='Quantity', aggfunc=np.sum)
 
 - str.wrap implemented (:issue:`6999`)
+- Add nsmallest and nlargest Series methods (:issue:`3960`)
 
 - `PeriodIndex` fully supports partial string indexing like `DatetimeIndex` (:issue:`7043`)
 

pandas/algos.pyx

@@ -739,7 +739,8 @@ def _check_minp(win, minp, N):
 # Physical description: 366 p.
 #               Series: Prentice-Hall Series in Automatic Computation
 
-ctypedef fused kth_type:
+
+ctypedef fused numeric:
     int8_t
     int16_t
     int32_t
@@ -754,17 +755,25 @@ ctypedef fused kth_type:
     float64_t
 
 
-cdef void swap_kth(kth_type *a, kth_type *b):
-    cdef kth_type t
+cdef inline Py_ssize_t swap(numeric *a, numeric *b) except -1:
+    cdef numeric t
+
+    # cython doesn't allow pointer dereference so use array syntax
     t = a[0]
     a[0] = b[0]
     b[0] = t
+    return 0
 
 
-cpdef kth_type kth_smallest(kth_type[:] a, Py_ssize_t k):
+@cython.boundscheck(False)
+@cython.wraparound(False)
+cpdef numeric kth_smallest(numeric[:] a, Py_ssize_t k):
     cdef:
-        Py_ssize_t i, j, l = 0, n = a.size, m = n - 1
-        kth_type x
+        Py_ssize_t i, j, l, m, n = a.size
+        numeric x
+
+    l = 0
+    m = n - 1
 
     while l < m:
         x = a[k]
@@ -775,7 +784,7 @@ cpdef kth_type kth_smallest(kth_type[:] a, Py_ssize_t k):
             while a[i] < x: i += 1
             while x < a[j]: j -= 1
             if i <= j:
-                swap_kth(&a[i], &a[j])
+                swap(&a[i], &a[j])
                 i += 1; j -= 1
 
             if i > j: break
@@ -801,7 +810,7 @@ cdef inline kth_smallest_c(float64_t* a, Py_ssize_t k, Py_ssize_t n):
             while a[i] < x: i += 1
             while x < a[j]: j -= 1
             if i <= j:
-                swap_kth(&a[i], &a[j])
+                swap(&a[i], &a[j])
                 i += 1; j -= 1
 
             if i > j: break
@@ -811,7 +820,7 @@ cdef inline kth_smallest_c(float64_t* a, Py_ssize_t k, Py_ssize_t n):
     return a[k]
 
 
-cpdef kth_type median(kth_type[:] arr):
+cpdef numeric median(numeric[:] arr):
     '''
     A faster median
     '''

pandas/core/algorithms.py

@@ -9,9 +9,7 @@
 import pandas.core.common as com
 import pandas.algos as algos
 import pandas.hashtable as htable
-import pandas.compat as compat
-from pandas.compat import filter, string_types
-from pandas.util.decorators import deprecate_kwarg
+from pandas.compat import string_types
 
 def match(to_match, values, na_sentinel=-1):
     """
@@ -413,6 +411,90 @@ def group_position(*args):
     return result
 
 
+_dtype_map = {'datetime64[ns]': 'int64', 'timedelta64[ns]': 'int64'}
+
+
+def _finalize_nsmallest(arr, kth_val, n, take_last, narr):
+    ns, = np.nonzero(arr <= kth_val)
+    inds = ns[arr[ns].argsort(kind='mergesort')][:n]
+
+    if take_last:
+        # reverse indices
+        return narr - 1 - inds
+    return inds
+
+
+def nsmallest(arr, n, take_last=False):
+    '''
+    Find the indices of the n smallest values of a numpy array.
+
+    Note: Fails silently with NaN.
+
+    '''
+    if take_last:
+        arr = arr[::-1]
+
+    narr = len(arr)
+    n = min(n, narr)
+
+    sdtype = str(arr.dtype)
+    arr = arr.view(_dtype_map.get(sdtype, sdtype))
+
+    kth_val = algos.kth_smallest(arr.copy(), n - 1)
+    return _finalize_nsmallest(arr, kth_val, n, take_last, narr)
+
+
+def nlargest(arr, n, take_last=False):
+    """
+    Find the indices of the n largest values of a numpy array.
+
+    Note: Fails silently with NaN.
+    """
+    sdtype = str(arr.dtype)
+    arr = arr.view(_dtype_map.get(sdtype, sdtype))
+    return nsmallest(-arr, n, take_last=take_last)
+
+
+def select_n_slow(dropped, n, take_last, method):
+    reverse_it = take_last or method == 'nlargest'
+    ascending = method == 'nsmallest'
+    slc = np.s_[::-1] if reverse_it else np.s_[:]
+    return dropped[slc].order(ascending=ascending).head(n)
+
+
+_select_methods = {'nsmallest': nsmallest, 'nlargest': nlargest}
+
+
+def select_n(series, n, take_last, method):
+    """Implement n largest/smallest.
+
+    Parameters
+    ----------
+    n : int
+    take_last : bool
+    method : str, {'nlargest', 'nsmallest'}
+
+    Returns
+    -------
+    nordered : Series
+    """
+    dtype = series.dtype
+    if not issubclass(dtype.type, (np.integer, np.floating, np.datetime64,
+                                   np.timedelta64)):
+        raise TypeError("Cannot use method %r with dtype %s" % (method, dtype))
+
+    if n <= 0:
+        return series[[]]
+
+    dropped = series.dropna()
+
+    if n >= len(series):
+        return select_n_slow(dropped, n, take_last, method)
+
+    inds = _select_methods[method](dropped.values, n, take_last)
+    return dropped.iloc[inds]
+
+
 _rank1d_functions = {
     'float64': algos.rank_1d_float64,
     'int64': algos.rank_1d_int64,

pandas/core/series.py

@@ -35,6 +35,7 @@
 
 import pandas.core.array as pa
 import pandas.core.ops as ops
+from pandas.core.algorithms import select_n
 
 import pandas.core.common as com
 import pandas.core.datetools as datetools
@@ -1699,17 +1700,7 @@ def _try_kind_sort(arr):
         good = ~bad
         idx = pa.arange(len(self))
 
-        def _try_kind_sort(arr, kind='mergesort'):
-            # easier to ask forgiveness than permission
-            try:
-                # if kind==mergesort, it can fail for object dtype
-                return arr.argsort(kind=kind)
-            except TypeError:
-                # stable sort not available for object dtype
-                # uses the argsort default quicksort
-                return arr.argsort(kind='quicksort')
-
-        argsorted = _try_kind_sort(arr[good], kind=kind)
+        argsorted = _try_kind_sort(arr[good])
 
         if not ascending:
             argsorted = argsorted[::-1]
@@ -1733,49 +1724,70 @@ def _try_kind_sort(arr, kind='mergesort'):
             return result.__finalize__(self)
 
     def nlargest(self, n=5, take_last=False):
-        '''
-        Returns the largest n rows:
+        """Return the largest `n` elements.
 
-        May be faster than .order(ascending=False).head(n).
+        Parameters
+        ----------
+        n : int
+            Return this many descending sorted values
+        take_last : bool
+            Where there are duplicate values, take the last duplicate
 
-        '''
-        # TODO remove need for dropna ?
-        dropped = self.dropna()
+        Returns
+        -------
+        top_n : Series
+            The n largest values in the Series, in sorted order
 
-        from pandas.tools.util import nlargest
+        Notes
+        -----
+        Faster than ``.order(ascending=False).head(n)`` for small `n` relative
+        to the size of the ``Series`` object.
 
-        if dropped.dtype == object:
-            try:
-                dropped = dropped.astype(float)
-            except (NotImplementedError, TypeError):
-                return dropped.order(ascending=False).head(n)
+        See Also
+        --------
+        Series.nsmallest
 
-        inds = nlargest(dropped.values, n, take_last)
-        if len(inds) == 0:
-            # TODO remove this special case
-            return dropped[[]]
-        return dropped.iloc[inds]
+        Examples
+        --------
+        >>> import pandas as pd
+        >>> import numpy as np
+        >>> s = pd.Series(np.random.randn(1e6))
+        >>> s.nlargest(10)  # only sorts up to the N requested
+        """
+        return select_n(self, n=n, take_last=take_last, method='nlargest')
 
     def nsmallest(self, n=5, take_last=False):
-        '''
-        Returns the smallest n rows.
+        """Return the smallest `n` elements.
 
-        May be faster than .order().head(n).
+        Parameters
+        ----------
+        n : int
+            Return this many ascending sorted values
+        take_last : bool
+            Where there are duplicate values, take the last duplicate
 
-        '''
-        # TODO remove need for dropna ?
-        dropped = self.dropna()
+        Returns
+        -------
+        bottom_n : Series
+            The n smallest values in the Series, in sorted order
 
-        from pandas.tools.util import nsmallest
-        try:
-            inds = nsmallest(dropped.values, n, take_last)
-        except NotImplementedError:
-            return dropped.order().head(n)
-
-        if len(inds) == 0:
-            # TODO remove this special case
-            return dropped[[]]
-        return dropped.iloc[inds]
+        Notes
+        -----
+        Faster than ``.order().head(n)`` for small `n` relative to
+        the size of the ``Series`` object.
+
+        See Also
+        --------
+        Series.nlargest
+
+        Examples
+        --------
+        >>> import pandas as pd
+        >>> import numpy as np
+        >>> s = pd.Series(np.random.randn(1e6))
+        >>> s.nsmallest(10)  # only sorts up to the N requested
+        """
+        return select_n(self, n=n, take_last=take_last, method='nsmallest')
 
     def sortlevel(self, level=0, ascending=True, sort_remaining=True):
         """

pandas/tests/test_series.py

@@ -4002,19 +4002,47 @@ def test_nsmallest_nlargest(self):
         # float, int, datetime64 (use i8), timedelts64 (same),
         # object that are numbers, object that are strings
 
-        s_list = [Series([3, 2, 1, 2, 5]),
-                  Series([3., 2., 1., 2., 5.]),
-                  Series([3., 2, 1, 2, 5], dtype='object'),
-                  Series([3., 2, 1, 2, '5'], dtype='object'),
-                  Series(pd.to_datetime(['2003', '2002', '2001', '2002', '2005']))]
+        base = [3, 2, 1, 2, 5]
+
+        s_list = [
+            Series(base, dtype='int8'),
+            Series(base, dtype='int16'),
+            Series(base, dtype='int32'),
+            Series(base, dtype='int64'),
+            Series(base, dtype='float32'),
+            Series(base, dtype='float64'),
+            Series(base, dtype='uint8'),
+            Series(base, dtype='uint16'),
+            Series(base, dtype='uint32'),
+            Series(base, dtype='uint64'),
+            Series(base).astype('timedelta64[ns]'),
+            Series(pd.to_datetime(['2003', '2002', '2001', '2002', '2005'])),
+        ]
+
+        raising = [
+            Series([3., 2, 1, 2, '5'], dtype='object'),
+            Series([3., 2, 1, 2, 5], dtype='object'),
+            Series([3., 2, 1, 2, 5], dtype='complex256'),
+            Series([3., 2, 1, 2, 5], dtype='complex128'),
+        ]
+
+        for r in raising:
+            dt = r.dtype
+            msg = "Cannot use method 'n(larg|small)est' with dtype %s" % dt
+            args = 2, len(r), 0, -1
+            methods = r.nlargest, r.nsmallest
+            for method, arg in product(methods, args):
+                with tm.assertRaisesRegexp(TypeError, msg):
+                    method(arg)
 
         for s in s_list:
 
             assert_series_equal(s.nsmallest(2), s.iloc[[2, 1]])
             assert_series_equal(s.nsmallest(2, take_last=True), s.iloc[[2, 3]])
 
             assert_series_equal(s.nlargest(3), s.iloc[[4, 0, 1]])
-            assert_series_equal(s.nlargest(3, take_last=True), s.iloc[[4, 0, 3]])
+            assert_series_equal(s.nlargest(3, take_last=True),
+                                s.iloc[[4, 0, 3]])
 
             empty = s.iloc[0:0]
             assert_series_equal(s.nsmallest(0), empty)
@@ -4025,7 +4053,8 @@ def test_nsmallest_nlargest(self):
             assert_series_equal(s.nsmallest(len(s)), s.order())
             assert_series_equal(s.nsmallest(len(s) + 1), s.order())
             assert_series_equal(s.nlargest(len(s)), s.iloc[[4, 0, 1, 3, 2]])
-            assert_series_equal(s.nlargest(len(s) + 1), s.iloc[[4, 0, 1, 3, 2]])
+            assert_series_equal(s.nlargest(len(s) + 1),
+                                s.iloc[[4, 0, 1, 3, 2]])
 
         s = Series([3., np.nan, 1, 2, 5])
         assert_series_equal(s.nlargest(), s.iloc[[4, 0, 3, 2]])