Class Forest

cpp/src/arrow/dataset/forest_internal.h:33–122

A Forest is a view of a sorted range which carries an ancestry relation in addition to an ordering relation: each element's descendants appear directly after it. This can be used to efficiently skip subtrees when iterating through the range.

/// to an ordering relation: each element's descendants appear directly after it.
/// This can be used to efficiently skip subtrees when iterating through the range.
class Forest {
 public:
  Forest() = default;

  /// \brief Construct a Forest viewing the range [0, size).
  Forest(int size, std::function<bool(int, int)> is_ancestor) : size_(size) {
    std::vector<int> descendant_counts(size, 0);

    std::vector<int> parent_stack;

    for (int i = 0; i < size; ++i) {
      while (parent_stack.size() != 0) {
        if (is_ancestor(parent_stack.back(), i)) break;

        // parent_stack.back() has no more descendants; finalize count and pop
        descendant_counts[parent_stack.back()] = i - 1 - parent_stack.back();
        parent_stack.pop_back();
      }

      parent_stack.push_back(i);
    }

    // finalize descendant_counts for anything left in the stack
    while (parent_stack.size() != 0) {
      descendant_counts[parent_stack.back()] = size - 1 - parent_stack.back();
      parent_stack.pop_back();
    }

    descendant_counts_ = std::make_shared<std::vector<int>>(std::move(descendant_counts));
  }

  /// \brief Returns the number of nodes in this forest.
  int size() const { return size_; }

  bool Equals(const Forest& other) const {
    auto it = descendant_counts_->begin();
    return size_ == other.size_ &&
           std::equal(it, it + size_, other.descendant_counts_->begin());
  }

  struct Ref {
    int num_descendants() const { return forest->descendant_counts_->at(i); }

    bool IsAncestorOf(const Ref& ref) const {
      return i < ref.i && i + 1 + num_descendants() > ref.i;
    }

    explicit operator bool() const { return forest != NULLPTR; }

    const Forest* forest;
    int i;
  };

  /// \brief Visit with eager pruning. Visitors must return Result<bool>, using
  /// true to indicate a subtree should be visited and false to indicate that the
  /// subtree should be skipped.
  template <typename PreVisitor, typename PostVisitor>
  Status Visit(PreVisitor&& pre, PostVisitor&& post) const {