[ENH] nth function

docs/api/dt/nth.rst

@@ -0,0 +1,79 @@
+
+.. xfunction:: datatable.nth
+    :src: src/core/expr/fexpr_nth.cc pyfn_nth
+    :cvar: doc_dt_nth
+    :tests: tests/test-reduce.py
+    :signature: nth(cols, n=0, skipna=None)
+
+    Return the ``nth`` row for an ``Expr``.
+
+    Parameters
+    ----------
+    cols: FExpr | iterable
+        Input columns or an iterable.
+
+    n: int
+        The number of the row to be returned.
+
+    skipna: None | "any" | "all"
+        Drop the nulls before counting which row is the nth row.
+        Needs to be ``None``, ``any``, or ``all``.
+
+    return: Expr | ...
+        One-row f-expression that has the same names, stypes and
+        number of columns as `cols`. 
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> from datatable import dt, f, by
+        >>>
+        >>> df = dt.Frame({'A': [1, 1, 2, 1, 2],
+        ...                'B': [None, 2, 3, 4, 5],
+        ...                'C': [1, 2, 1, 1, 2]})
+        >>> df
+           |     A      B      C
+           | int32  int32  int32
+        -- + -----  -----  -----
+         0 |     1     NA      1
+         1 |     1      2      2
+         2 |     2      3      1
+         3 |     1      4      1
+         4 |     2      5      2
+        [5 rows x 3 columns]
+
+    Get the third row of column A::
+
+        >>> df[:, dt.nth(f.A, n=2)]
+           |     A
+           | int32
+        -- + -----
+         0 |     2
+        [1 row x 1 column]
+
+    Get the third row for multiple columns::
+
+        >>> df[:, dt.nth(f[:], n=2)]
+           |     A      B      C
+           | int32  int32  int32
+        -- + -----  -----  -----
+         0 |     2      3      1
+
+
+In the presence of :func:`by()`, it returns the nth row of the specified columns per group::
+
+        >>> df[:, [dt.nth(f.A, n = 2), dt.nth(f.B, n = -1)], by(f.C)]
+           |     C      A      B
+           | int32  int32  int32
+        -- + -----  -----  -----
+         0 |     1      1      4
+         1 |     2     NA      5
+        [2 rows x 3 columns]
+
+
+
+    See Also
+    --------
+    - :func:`first()` -- function that returns the first row.
+    - :func:`last()` -- function that returns the last row.

docs/api/fexpr.rst

@@ -327,6 +327,7 @@
     .median()       <fexpr/median>
     .min()          <fexpr/min>
     .nunique()      <fexpr/nunique>
+    .nth()          <fexpr/nth>
     .prod()         <fexpr/prod>
     .re_match()     <fexpr/re_match>
     .remove()       <fexpr/remove>

docs/api/fexpr/nth.rst

@@ -0,0 +1,7 @@
+
+.. xmethod:: datatable.FExpr.nth
+    :src: src/core/expr/fexpr.cc PyFExpr::nth
+    :cvar: doc_FExpr_nth
+    :signature: nth(n, skipna)
+
+    Equivalent to :func:`dt.nth(cols, n, skipna)`.

docs/api/index-api.rst

@@ -189,6 +189,8 @@ Functions
       - Find the smallest element per column
     * - :func:`ngroup()`
       - Number each group
+    * - :func:`nth()`
+      - Return the nth row.
     * - :func:`nunique()`
       - Count the number of unique values per column
     * - :func:`prod()`
@@ -273,6 +275,7 @@ Other
     median()          <dt/median>
     min()             <dt/min>
     ngroup()          <dt/ngroup>
+    nth()             <dt/nth>
     nunique()         <dt/nunique>
     prod()            <dt/prod>
     qcut()            <dt/qcut>

docs/releases/v1.1.0.rst

@@ -114,6 +114,8 @@
 
     -[new] Added reducer functions :func:`dt.countna()` and :func:`dt.nunique()`. [#2999]
 
+    -[new] Added function :func:`dt.nth()` to retrieve the n-th row. [#3128]
+
     -[new] Class :class:`dt.FExpr` now has method :meth:`.nunique()`,
       which behaves exactly as the equivalent base level function :func:`dt.nunique()`.
 

src/core/column/nth.h

@@ -0,0 +1,86 @@
+//------------------------------------------------------------------------------
+// Copyright 2022 H2O.ai
+//
+// Permission is hereby granted, free of charge, to any person obtaining a
+// copy of this software and associated documentation files (the "Software"),
+// to deal in the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
+// IN THE SOFTWARE.
+//------------------------------------------------------------------------------
+#ifndef dt_NTH_h
+#define dt_NTH_h
+#include "column/virtual.h"
+#include "stype.h"
+namespace dt {
+
+template<typename T>
+class Nth_ColumnImpl : public Virtual_ColumnImpl {
+  private:
+    Column col_;
+    Groupby gby_;
+    bool is_grouped_;
+    int32_t n_;
+    size_t : 32;
+
+  public:
+    Nth_ColumnImpl(Column&& col, const Groupby& gby, bool is_grouped, int32_t n)
+      : Virtual_ColumnImpl(gby.size(), col.stype()),
+        col_(std::move(col)),
+        gby_(gby),
+        is_grouped_(is_grouped),
+        n_(n)
+    {
+      xassert(col_.can_be_read_as<T>());
+    }
+
+
+    ColumnImpl* clone() const override {
+      return new Nth_ColumnImpl(Column(col_), gby_, is_grouped_, n_);
+    }
+
+
+    size_t n_children() const noexcept override {
+      return 1;
+    }
+
+
+    const Column& child(size_t i) const override {
+      xassert(i == 0);  (void)i;
+      return col_;
+    }
+
+    bool get_element(size_t i, T* out) const override { 
+      xassert(i < gby_.size());
+      size_t i0, i1;
+      gby_.get_group(i, &i0, &i1);
+
+      // Note, when `n_` is negative it is cast to `size_t`, that is
+      // an unsigned type. Then, when adding `i1`, we rely on `size_t`
+      // wrap-around.
+      size_t ni = (n_ >= 0)? static_cast<size_t>(n_) + i0
+                           : static_cast<size_t>(n_) + i1;
+      bool isvalid = false;
+      if (ni >= i0 && ni < i1){
+        ni = is_grouped_?i:ni;
+        isvalid = col_.get_element(ni, out);
+      }
+      return isvalid;
+    }
+}; 
+
+}  // namespace dt
+#endif
+
+

src/core/documentation.h

@@ -52,6 +52,7 @@ extern const char* doc_dt_mean;
 extern const char* doc_dt_median;
 extern const char* doc_dt_min;
 extern const char* doc_dt_ngroup;
+extern const char* doc_dt_nth;
 extern const char* doc_dt_nunique;
 extern const char* doc_dt_qcut;
 extern const char* doc_dt_rbind;
@@ -302,6 +303,7 @@ extern const char* doc_FExpr_max;
 extern const char* doc_FExpr_mean;
 extern const char* doc_FExpr_median;
 extern const char* doc_FExpr_min;
+extern const char* doc_FExpr_nth;
 extern const char* doc_FExpr_nunique;
 extern const char* doc_FExpr_prod;
 extern const char* doc_FExpr_remove;

src/core/expr/fexpr.cc

@@ -531,6 +531,20 @@ DECLARE_METHOD(&PyFExpr::min)
     ->name("min")
     ->docs(dt::doc_FExpr_min);
 
+oobj PyFExpr::nth(const XArgs& args) {
+  auto nthFn = oobj::import("datatable", "nth");
+  oobj n = args[0].to_oobj() ? args[0].to_oobj() 
+                             : py::oint(0);
+  oobj skipna = args[1].to_oobj_or_none();
+  return nthFn.call({this, n, skipna});
+}
+
+DECLARE_METHOD(&PyFExpr::nth)
+    ->name("nth")
+    ->arg_names({"n", "skipna"})
+    ->n_positional_or_keyword_args(2)
+    ->docs(dt::doc_FExpr_nth);
+
 
 oobj PyFExpr::nunique(const XArgs&) {
   auto nuniqueFn = oobj::import("datatable", "nunique");

src/core/expr/fexpr.h

@@ -197,6 +197,7 @@ class PyFExpr : public py::XObject<PyFExpr> {
     py::oobj mean(const py::XArgs&);
     py::oobj median(const py::XArgs&);
     py::oobj min(const py::XArgs&);
+    py::oobj nth(const py::XArgs&);
     py::oobj nunique(const py::XArgs&);
     py::oobj prod(const py::XArgs&);
     py::oobj remove(const py::XArgs&);