[ENH] nth function

docs/api/dt/nth.rst

@@ -0,0 +1,89 @@
+
+.. 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)
+
+    Return the ``nth`` row for an ``Expr``.
+
+    Parameters
+    ----------
+    cols: FExpr | iterable
+        Input columns or an iterable.
+
+    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
+
+    Replicate :func:`first()`::
+
+        >>> df[:, dt.nth(f.A, n=0)]
+           |     A
+           | int32
+        -- + -----
+         0 |     1
+        [1 row x 1 column]
+
+    Replicate :func:`last()`::
+
+        >>> df[:, dt.nth(f.A, n=-1)]
+           |     A
+           | int32
+        -- + -----
+         0 |     2
+        [1 row x 1 column]
+
+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/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)
+
+    Equivalent to :func:`dt.nth(cols, n)`.

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

@@ -100,6 +100,8 @@
 
     -[new] Added reducer functions :func:`dt.countna()` and :func:`dt.nunique()`. [#2999]
 
+    -[new] Added function :func:`dt.nth()` to retrieve specified 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,85 @@
+//------------------------------------------------------------------------------
+// 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_;
+    int32_t n_;
+    size_t : 32;
+
+  public:
+    Nth_ColumnImpl(Column&& col, const Groupby& gby, int32_t n)
+      : Virtual_ColumnImpl(gby.size(), col.stype()),
+        col_(std::move(col)),
+        gby_(gby),
+        n_(n)
+    {
+      xassert(col_.can_be_read_as<T>());
+    }
+
+
+    ColumnImpl* clone() const override {
+      return new Nth_ColumnImpl(Column(col_), gby_, 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){
+        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

@@ -532,6 +532,21 @@ 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)
+    ->n_required_args(1)
+    ->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&);

src/core/expr/fexpr_nth.cc

@@ -0,0 +1,131 @@
+//------------------------------------------------------------------------------
+// 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.
+//------------------------------------------------------------------------------
+#include "column/const.h"
+#include "column/nth.h"
+#include "documentation.h"
+#include "expr/fexpr_func.h"
+#include "expr/eval_context.h"
+#include "python/xargs.h"
+namespace dt {
+namespace expr {
+
+
+template<bool SKIPNA>
+class FExpr_Nth : public FExpr_Func {
+  private:
+    ptrExpr arg_;
+    int32_t n_;
+
+  public:
+    FExpr_Nth(ptrExpr&& arg, py::oobj n)
+      : arg_(std::move(arg))
+       {n_ = n.to_int32_strict();}
+
+    std::string repr() const override {
+      std::string out = "nth";
+      out += '(';
+      out += arg_->repr();
+      out += ", nth=";
+      out += std::to_string(n_);
+      out += ", skipna=";
+      out += SKIPNA? "True" : "False";
+      out += ')';
+      return out;
+    }
+
+
+    Workframe evaluate_n(EvalContext &ctx) const override {
+      Workframe inputs = arg_->evaluate_n(ctx);
+      Workframe outputs(ctx);
+      Groupby gby = ctx.get_groupby();
+
+      if (!gby) gby = Groupby::single_group(ctx.nrows()); 
+
+      for (size_t i = 0; i < inputs.ncols(); ++i) {
+        auto coli = inputs.retrieve_column(i);
+        outputs.add_column(
+          evaluate1(std::move(coli), gby, n_),
+          inputs.retrieve_name(i),
+          Grouping::GtoONE
+        );
+      }
+      return outputs;
+    }
+
+
+    Column evaluate1(Column&& col, const Groupby& gby, const int32_t n) const {
+      SType stype = col.stype();
+      switch (stype) {
+        case SType::VOID:    return Column(new ConstNa_ColumnImpl(gby.size()));
+        case SType::BOOL:
+        case SType::INT8:    return make<int8_t>(std::move(col), gby, n);
+        case SType::INT16:   return make<int16_t>(std::move(col), gby, n);
+        case SType::DATE32:
+        case SType::INT32:   return make<int32_t>(std::move(col), gby, n);
+        case SType::TIME64:
+        case SType::INT64:   return make<int64_t>(std::move(col), gby, n);
+        case SType::FLOAT32: return make<float>(std::move(col), gby, n);
+        case SType::FLOAT64: return make<double>(std::move(col), gby, n);
+        case SType::STR32:   return make<CString>(std::move(col), gby, n);
+        case SType::STR64:   return make<CString>(std::move(col), gby, n);
+        default:
+          throw TypeError()
+            << "Invalid column of type `" << stype << "` in " << repr();
+      }
+    }
+
+
+    template <typename T>
+    Column make(Column&& col, const Groupby& gby, int32_t n) const {
+      return Column(new Nth_ColumnImpl<T>(std::move(col), gby, n));
+    }
+
+};
+
+
+static py::oobj pyfn_nth(const py::XArgs& args) {
+  auto arg = args[0].to_oobj();
+  auto n = args[1].to_oobj();
+  auto skipna = args[2].to<bool>(false);
+  if (!n.is_int()) {
+    throw TypeError() << "The argument for the `nth` parameter "
+                <<"in function datatable.nth() should be an integer, "
+                <<"instead got "<<n.typeobj();
+  }
+  if (skipna) {
+    return PyFExpr::make(new FExpr_Nth<true>(as_fexpr(arg), n));
+  } else {
+    return PyFExpr::make(new FExpr_Nth<false>(as_fexpr(arg), n));
+  }
+}
+
+
+DECLARE_PYFN(&pyfn_nth)
+    ->name("nth")
+    ->docs(doc_dt_nth)
+    ->arg_names({"cols", "nth", "skipna"})
+    ->n_positional_args(1)
+    ->n_positional_or_keyword_args(2)
+    ->n_required_args(2);
+
+
+}}  // dt::expr

src/core/expr/head_reduce_unary.cc

@@ -184,6 +184,7 @@ static Column compute_gfirstlast(Column&& arg, const Groupby&) {
 
 
 
+
 //------------------------------------------------------------------------------
 // mean(A)
 //------------------------------------------------------------------------------