Refactor .extend() and .remove() to use FExpr (#3393)

WIP for #2562

src/core/column/sumprod.h

@@ -27,8 +27,8 @@
 namespace dt {
 
 
-  template <typename T, bool SUM>
-  class SumProd_ColumnImpl : public Virtual_ColumnImpl {
+template <typename T, bool SUM>
+class SumProd_ColumnImpl : public Virtual_ColumnImpl {
   private:
     Column col_;
     Groupby gby_;
@@ -48,7 +48,7 @@ namespace dt {
 
     bool get_element(size_t i, T* out) const override {
       T result = !SUM; // 0 for `sum()` and 1 for `prod()`
-      T value;     
+      T value;
       size_t i0, i1;
       gby_.get_group(i, &i0, &i1);
 
@@ -89,7 +89,7 @@ namespace dt {
       (void)i;
       return col_;
     }
-  };
+};
 
 
 }  // namespace dt

src/core/expr/fexpr.cc

@@ -26,13 +26,13 @@
 #include "expr/fexpr_alias.h"
 #include "expr/fexpr_column.h"
 #include "expr/fexpr_dict.h"
+#include "expr/fexpr_extend_remove.h"
 #include "expr/fexpr_frame.h"
 #include "expr/fexpr_list.h"
 #include "expr/fexpr_literal.h"
 #include "expr/fexpr_slice.h"
 #include "expr/re/fexpr_match.h"
 #include "expr/str/fexpr_len.h"
-#include "documentation.h"
 #include "python/obj.h"
 #include "python/xargs.h"
 #include "utils/exceptions.h"
@@ -234,8 +234,9 @@ oobj PyFExpr::nb__pos__() {
 //----- Other methods ----------------------------------------------------------
 
 oobj PyFExpr::extend(const XArgs& args) {
-  auto arg = args[0].to<oobj>(py::None());
-  return make_binexpr(dt::expr::Op::SETPLUS, robj(this), arg);
+  auto arg = args[0].to_oobj();   
+  return PyFExpr::make(new FExpr_Extend_Remove<true>(ptrExpr(expr_), as_fexpr(arg)));
+
 }
 
 DECLARE_METHOD(&PyFExpr::extend)
@@ -248,8 +249,9 @@ DECLARE_METHOD(&PyFExpr::extend)
 
 
 oobj PyFExpr::remove(const XArgs& args) {
-  auto arg = args[0].to_oobj();
-  return make_binexpr(dt::expr::Op::SETMINUS, robj(this), arg);
+  auto arg = args[0].to_oobj();   
+  return PyFExpr::make(new FExpr_Extend_Remove<false>(ptrExpr(expr_), as_fexpr(arg)));
+
 }
 
 DECLARE_METHOD(&PyFExpr::remove)

src/core/expr/fexpr_extend_remove.cc

@@ -0,0 +1,65 @@
+//------------------------------------------------------------------------------
+// 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 "expr/eval_context.h"
+#include "expr/fexpr_extend_remove.h"
+namespace dt {
+namespace expr {
+
+
+template <bool EXTEND>
+FExpr_Extend_Remove<EXTEND>::FExpr_Extend_Remove(ptrExpr&& arg, ptrExpr&& other) :
+  arg_(std::move(arg)),
+  other_(std::move(other))
+  {}
+
+
+template <bool EXTEND>
+std::string FExpr_Extend_Remove<EXTEND>::repr() const {
+  std::string out = EXTEND? "extend" : "remove";
+  out += '(';
+  out += arg_->repr();
+  out += ", ";
+  out += other_->repr();
+  out += ')';
+  return out;
+}
+
+
+template <bool EXTEND>
+Workframe FExpr_Extend_Remove<EXTEND>::evaluate_n(EvalContext& ctx) const {
+  Workframe wf = arg_->evaluate_n(ctx);
+  Workframe wf_other = other_->evaluate_n(ctx);
+  if (EXTEND){
+    wf.cbind(std::move(wf_other));
+  } else {
+    wf.remove(std::move(wf_other));
+  }
+
+  return wf;
+}
+
+
+template class FExpr_Extend_Remove<true>;
+template class FExpr_Extend_Remove<false>;
+
+
+}}  // dt::expr

src/core/expr/fexpr_extend_remove.h

@@ -0,0 +1,46 @@
+//------------------------------------------------------------------------------
+// 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_EXPR_FEXPR_EXTEND_REMOVE_h
+#define dt_EXPR_FEXPR_EXTEND_REMOVE_h
+#include "expr/fexpr_func.h"
+namespace dt {
+namespace expr {
+
+
+template <bool EXTEND>
+class FExpr_Extend_Remove : public FExpr_Func {
+  private:
+    ptrExpr arg_;
+    ptrExpr other_;
+
+  public:
+    FExpr_Extend_Remove(ptrExpr&& arg, ptrExpr&& other);
+    std::string repr() const override;
+    Workframe evaluate_n(EvalContext& ctx) const override;
+};
+
+extern template class FExpr_Extend_Remove<true>;
+extern template class FExpr_Extend_Remove<false>;
+
+
+}}  // dt::expr
+#endif

tests/test-f.py

@@ -126,19 +126,15 @@ def test_f_columnset_str():
 
 
 def test_f_columnset_extend():
-    assert str(f[:].extend(f.A)) == \
-        "Expr:setplus(FExpr<f[:]>, FExpr<f.A>; )"
-    assert str(f[int].extend(f[str])) == \
-        "Expr:setplus(FExpr<f[int]>, FExpr<f[str]>; )"
-    assert str(f.A.extend(f['B','C'])) == \
-        "Expr:setplus(FExpr<f.A>, FExpr<f[['B', 'C']]>; )"
+    assert str(f[:].extend(f.A)) == "FExpr<extend(f[:], f.A)>"
+    assert str(f[int].extend(f[str])) == "FExpr<extend(f[int], f[str])>"
+    assert str(f.A.extend(f['B','C'])) == "FExpr<extend(f.A, f[['B', 'C']])>"
 
 
 def test_f_columnset_remove():
-    assert str(f[:].remove(f.A)) == "Expr:setminus(FExpr<f[:]>, FExpr<f.A>; )"
-    assert str(f[int].remove(f[0])) == "Expr:setminus(FExpr<f[int]>, FExpr<f[0]>; )"
-    assert str(f.A.remove(f['B','C'])) == \
-        "Expr:setminus(FExpr<f.A>, FExpr<f[['B', 'C']]>; )"
+    assert str(f[:].remove(f.A)) == "FExpr<remove(f[:], f.A)>"
+    assert str(f[int].remove(f[0])) == "FExpr<remove(f[int], f[0])>"
+    assert str(f.A.remove(f['B','C'])) == "FExpr<remove(f.A, f[['B', 'C']])>"
 
 
 
@@ -209,7 +205,7 @@ def test_f_columnset_ltypes(DT):
 def test_columnset_sum(DT):
     assert_equals(DT[:, f[int].extend(f[float])], DT[:, [int, float]])
     assert_equals(DT[:, f[:3].extend(f[-3:])], DT[:, [0, 1, 2, -3, -2, -1]])
-    assert_equals( DT[:, f['A','B','C'].extend(f['E','F', 'G'])], DT[:, [0, 1, 2, -3, -2, -1]])
+    assert_equals(DT[:, f['A','B','C'].extend(f['E','F', 'G'])], DT[:, [0, 1, 2, -3, -2, -1]])
     assert_equals(DT[:, f.A.extend(f.B)], DT[:, ['A', 'B']])
     assert_equals(DT[:, f[:].extend({"extra": f.A + f.C})],
                   dt.cbind(DT, DT[:, {"extra": f.A + f.C}]))