メモ:dplyrのrow_number()の実装を追う

このissueを解決してなんかコントリビュート感を出したい。が、C++力なさすぎて実装をまったく追えないのでメモ。

github.com

arrange()の場合

まず、Hadleyが「We did work around this for arrange() though,」って言ってるのでそっちがどうなってるか見てみる。

arrange_implはこのへん。値をぜんぶquosureにしたあと、なにやらOrderVisitorsというやつをつくってそれをapply()するとインデックスが出てくるらしい。

  variables.names() = quosures.names();

  OrderVisitors o(variables, ascending, nargs);
  IntegerVector index = o.apply();

  DataFrameSubsetVisitors visitors(data, data.names());
  List res = visitors.subset(index, get_class(data));

(https://github.com/tidyverse/dplyr/blob/7e90e95258010c8c43d856f3f9af927546a9c833/src/arrange.cpp#L65-L71)

これにDataFrameクラスが渡ったときの実装はここ。

  OrderVisitors(DataFrame data) :
    visitors(data.size()), n(data.size()), nrows(data.nrows())
  {
    for (int i = 0; i < n; i++)
      visitors[i]  = order_visitor(data[i], true, i);
  }

(https://github.com/tidyverse/dplyr/blob/7e90e95258010c8c43d856f3f9af927546a9c833/inst/include/dplyr/Order.h#L23-L28)

order_visitor()はこのあたり。

template <bool ascending>
OrderVisitor* order_visitor_asc_vector(SEXP vec);

inline OrderVisitor* order_visitor(SEXP vec, const bool ascending, const int i) {
  try {
    if (ascending) {
      return order_visitor_asc<true>(vec);
    }
    else {
      return order_visitor_asc<false>(vec);
    }
  }
  catch (const Rcpp::exception& e) {
    bad_pos_arg(i + 1, e.what());
  }
}

(https://github.com/tidyverse/dplyr/blob/7e90e95258010c8c43d856f3f9af927546a9c833/inst/include/dplyr/OrderVisitorImpl.h#L204-L219)

なにやら場合分けがいろいろある。文字列(STRSXP)なのでOrderCharacterVectorVisitorImplが選ばれる。

template <bool ascending>
inline OrderVisitor* order_visitor_asc_vector(SEXP vec) {
  switch (TYPEOF(vec)) {
  case INTSXP:
    return new OrderVectorVisitorImpl<INTSXP, ascending, Vector<INTSXP > >(vec);
  case REALSXP:
    return new OrderVectorVisitorImpl<REALSXP, ascending, Vector<REALSXP> >(vec);
  case LGLSXP:
    return new OrderVectorVisitorImpl<LGLSXP, ascending, Vector<LGLSXP > >(vec);
  case STRSXP:
    return new OrderCharacterVectorVisitorImpl<ascending>(vec);
  case CPLXSXP:
    return new OrderVectorVisitorImpl<CPLXSXP, ascending, Vector<CPLXSXP > >(vec);
  case VECSXP:
  {

(https://github.com/tidyverse/dplyr/blob/7e90e95258010c8c43d856f3f9af927546a9c833/inst/include/dplyr/OrderVisitorImpl.h#L252-L266)

OrderCharacterVectorVisitorImplはこのへん。

template <bool ascending>
class OrderCharacterVectorVisitorImpl : public OrderVisitor {
public:
  OrderCharacterVectorVisitorImpl(const CharacterVector& vec_) :
    vec(vec_),
    orders(CharacterVectorOrderer(vec).get())
  {}

  inline bool equal(int i, int j) const {
    return orders.equal(i, j);
  }

  inline bool before(int i, int j) const {
    return orders.before(i, j);
  }

  SEXP get() {
    return vec;
  }

(https://github.com/tidyverse/dplyr/blob/7e90e95258010c8c43d856f3f9af927546a9c833/inst/include/dplyr/OrderVisitorImpl.h#L75-L93)

CharacterVectorOrdererはこのへんと見せかけつつ、

https://github.com/tidyverse/dplyr/blob/ca16b62e2481e2273e53fd3c291dde6bd32e4ead/inst/include/dplyr/CharacterVectorOrderer.h

実装はこっち。

https://github.com/tidyverse/dplyr/blob/354b0851cbf7ef41ce9314269e6e677c8d4523dc/src/api.cpp#L95

で、このあたりがみそっぽい。

  // retrieve unique strings from the set
  int n_uniques = set.size();
  CharacterVector uniques(set.begin(), set.end());
  CharacterVector s_uniques = Language("sort", uniques).fast_eval();

  // order the uniques with a callback to R
  IntegerVector o = r_match(uniques, s_uniques);

要はRの実装であるsortを使ってるから大丈夫という。

row_number()の場合

このへん。

template <bool ascending>
Result* row_number_asc(const RObject& data) {
  switch (TYPEOF(data)) {
  case INTSXP:
    return new RowNumber<INTSXP, ascending>(data);
  case REALSXP:
    return new RowNumber<REALSXP, ascending>(data);
  case STRSXP:
    return new RowNumber<STRSXP, ascending>(data);
  default:
    return 0;
  }
}

(https://github.com/tidyverse/dplyr/blob/354b0851cbf7ef41ce9314269e6e677c8d4523dc/src/hybrid_window.cpp#L14-L26)

RowNumberはこのへん。

https://github.com/tidyverse/dplyr/blob/4bb35fbf4dbdc5c366f6f6dba4847b3fc6f785ac/inst/include/dplyr/Result/Rank.h#L261-L262

で、このあたりが実際に並べ替えてるところっぽい。

      Visitor visitor(slice);
      std::sort(tmp.begin(), tmp.begin() + m, Comparer(visitor));

Comparerはそのちょっと上でtypedefされている。

  typedef Compare_Single_OrderVisitor<Visitor> Comparer;

Compare_Single_OrderVisitorはこのあたり。

template <typename OrderVisitorClass>
class Compare_Single_OrderVisitor {
public:
  Compare_Single_OrderVisitor(const OrderVisitorClass& obj_) : obj(obj_) {}

  inline bool operator()(int i, int j) const {
    if (i == j) return false;
    if (obj.equal(i, j)) return i < j;
    return obj.before(i, j);
  }

private:
  const OrderVisitorClass& obj;
};

(https://github.com/tidyverse/dplyr/blob/78b3e9f54f36d75b4a4795a94058d04a646fb6d6/inst/include/dplyr/Order.h#L55-L68)

objは何かというと、ちょっと戻ってComperer()に渡されているvisitorの型Visitorはここで定義されている。

  typedef OrderVectorVisitorImpl<RTYPE, ascending, Slice> Visitor;

OrderVectorVisitorImplの実装をみるとこんな感じ。

// version used for ascending = true
template <int RTYPE, bool ascending, typename VECTOR>
class OrderVectorVisitorImpl : public OrderVisitor {
  typedef comparisons<RTYPE> compare;

public:
  /**
   * The type of data : int, double, SEXP, Rcomplex
   */
  typedef typename Rcpp::traits::storage_type<RTYPE>::type STORAGE;

  OrderVectorVisitorImpl(const VECTOR& vec_) : vec(vec_) {}

  inline bool equal(int i, int j) const {
    return compare::equal_or_both_na(vec[i], vec[j]);
  }

  inline bool before(int i, int j) const {
    return compare::is_less(vec[i], vec[j]);
  }

  SEXP get() {
    return vec;
  }

private:
  VECTOR vec;
};

(https://github.com/tidyverse/dplyr/blob/4bb35fbf4dbdc5c366f6f6dba4847b3fc6f785ac/inst/include/dplyr/OrderVisitorImpl.h#L16-L43)

で、is_less()はこんな感じで<で比較してるだけなので、文字列を扱うとだめっぽい。

  static inline bool is_less(STORAGE lhs, STORAGE rhs) {
    if (is_na(lhs)) return false;
    if (is_na(rhs)) return true;

    return lhs < rhs;
  }

(https://github.com/tidyverse/dplyr/blob/ca16b62e2481e2273e53fd3c291dde6bd32e4ead/inst/include/dplyr/comparisons.h#L10-L15)