20
20
// IN THE SOFTWARE.
21
21
// ------------------------------------------------------------------------------
22
22
#include " column/const.h"
23
+ #include " column/func_binary.h"
24
+ #include " column/isna.h"
23
25
#include " column/nth.h"
24
26
#include " documentation.h"
25
27
#include " expr/fexpr_func.h"
26
28
#include " expr/eval_context.h"
27
29
#include " python/xargs.h"
30
+ #include < iostream>
28
31
namespace dt {
29
32
namespace expr {
30
33
31
34
32
- template <bool SKIPNA>
35
+ template <size_t SKIPNA>
33
36
class FExpr_Nth : public FExpr_Func {
34
37
private:
35
38
ptrExpr arg_;
@@ -46,11 +49,57 @@ class FExpr_Nth : public FExpr_Func {
46
49
out += arg_->repr ();
47
50
out += " , n=" ;
48
51
out += std::to_string (n_);
49
- out += " , skipna=" ;
50
- out += SKIPNA? " True" : " False" ;
52
+ if (SKIPNA == 0 ) {
53
+ out += " , skipna=None" ;
54
+ } else if (SKIPNA == 1 ) {
55
+ out += " , skipna=any" ;
56
+ } else if (SKIPNA == 2 ) {
57
+ out += " , skipna=all" ;
58
+ }
51
59
out += ' )' ;
52
60
return out;
53
61
}
62
+
63
+ static Column make_isna_col (Column&& col) {
64
+ switch (col.stype ()) {
65
+ case SType::VOID: return Const_ColumnImpl::make_bool_column (col.nrows (), true );
66
+ case SType::BOOL:
67
+ case SType::INT8: return Column (new Isna_ColumnImpl<int8_t >(std::move (col)));
68
+ case SType::INT16: return Column (new Isna_ColumnImpl<int16_t >(std::move (col)));
69
+ case SType::DATE32:
70
+ case SType::INT32: return Column (new Isna_ColumnImpl<int32_t >(std::move (col)));
71
+ case SType::TIME64:
72
+ case SType::INT64: return Column (new Isna_ColumnImpl<int64_t >(std::move (col)));
73
+ case SType::FLOAT32: return Column (new Isna_ColumnImpl<float >(std::move (col)));
74
+ case SType::FLOAT64: return Column (new Isna_ColumnImpl<double >(std::move (col)));
75
+ case SType::STR32:
76
+ case SType::STR64: return Column (new Isna_ColumnImpl<CString>(std::move (col)));
77
+ default : throw RuntimeError ();
78
+ }
79
+ }
80
+
81
+ template <typename T>
82
+ static Column make_bool_col (Column&& a, Column&& b, SType BOOL) {
83
+ xassert (compatible_type<T>(stype));
84
+ size_t nrows = a.nrows ();
85
+ a.cast_inplace (SType::BOOL);
86
+ b.cast_inplace (SType::BOOL);
87
+ if (SKIPNA == 1 ) {
88
+ return Column (new FuncBinary1_ColumnImpl<T, T, T>(
89
+ std::move (a), std::move (b),
90
+ [](T x, T y){ return x | y; },
91
+ nrows, SType::BOOL
92
+ ));
93
+ }
94
+ if (SKIPNA == 2 ) {
95
+ return Column (new FuncBinary1_ColumnImpl<T, T, T>(
96
+ std::move (a), std::move (b),
97
+ [](T x, T y){ return x & y; },
98
+ nrows, SType::BOOL
99
+ ));
100
+ }
101
+
102
+ }
54
103
55
104
56
105
Workframe evaluate_n (EvalContext &ctx) const override {
@@ -67,16 +116,12 @@ class FExpr_Nth : public FExpr_Func {
67
116
);
68
117
Column coli = evaluate1 (wf.retrieve_column (i), gby, is_grouped, n_);
69
118
outputs.add_column (std::move (coli), wf.retrieve_name (i), Grouping::GtoONE);
70
- // auto coli = inputs.retrieve_column(i);
71
- // outputs.add_column(
72
- // evaluate1(std::move(coli), gby, n_),
73
- // inputs.retrieve_name(i),
74
- // Grouping::GtoONE
75
- // );
76
119
}
77
120
return outputs;
78
121
}
79
122
123
+
124
+
80
125
81
126
Column evaluate1 (Column&& col, const Groupby& gby, bool is_grouped, const int32_t n) const {
82
127
SType stype = col.stype ();
@@ -111,17 +156,37 @@ class FExpr_Nth : public FExpr_Func {
111
156
static py::oobj pyfn_nth (const py::XArgs& args) {
112
157
auto arg = args[0 ].to_oobj ();
113
158
auto n = args[1 ].to <py::oobj>(py::oint (0 ));
114
- auto skipna = args[2 ].to <bool >(false );
159
+ auto skipna = args[2 ].to_oobj_or_none ();
160
+ if (!skipna.is_none ()) {
161
+ if (!skipna.is_string ()) {
162
+ throw TypeError () << " The argument for the `skipna` parameter "
163
+ <<" in function datatable.nth() should either be None, "
164
+ <<" or a string, instead got " <<skipna.typeobj ();
165
+
166
+ }
167
+ std::string skip_na = skipna.to_string ();
168
+ if (skip_na != " any" && skip_na != " all" ) {
169
+ throw TypeError () << " The argument for the `skipna` parameter "
170
+ <<" in function datatable.nth() should either be None, "
171
+ <<" any or all, instead got " <<skipna.repr ();
172
+ }
173
+ }
115
174
if (!n.is_int ()) {
116
175
throw TypeError () << " The argument for the `nth` parameter "
117
176
<<" in function datatable.nth() should be an integer, "
118
177
<<" instead got " <<n.typeobj ();
119
178
}
120
- if (skipna) {
121
- return PyFExpr::make (new FExpr_Nth<true >(as_fexpr (arg), n));
122
- } else {
123
- return PyFExpr::make (new FExpr_Nth<false >(as_fexpr (arg), n));
124
- }
179
+ if (!skipna.is_none ()) {
180
+ std::string skip_na = skipna.to_string ();
181
+ if (skip_na == " any" ) {
182
+ return PyFExpr::make (new FExpr_Nth<1 >(as_fexpr (arg), n));
183
+ }
184
+ if (skip_na == " all" ) {
185
+ return PyFExpr::make (new FExpr_Nth<2 >(as_fexpr (arg), n));
186
+ }
187
+
188
+ }
189
+ return PyFExpr::make (new FExpr_Nth<0 >(as_fexpr (arg), n));
125
190
}
126
191
127
192
@@ -131,7 +196,7 @@ DECLARE_PYFN(&pyfn_nth)
131
196
->arg_names({" cols" , " n" , " skipna" })
132
197
->n_positional_args(1 )
133
198
->n_positional_or_keyword_args(2 )
134
- ->n_required_args(2 );
199
+ ->n_required_args(1 );
135
200
136
201
137
202
}} // dt::expr
0 commit comments