group_by_key_rdd.h

#ifndef CPARK_GROUP_BY_KEY_RDD_H
#define CPARK_GROUP_BY_KEY_RDD_H
 
#include <concepts>
#include <memory>
#include <ranges>
#include <shared_mutex>
#include <mutex>
 
#include "base_rdd.h"
#include "utils.h"
 
namespace cpark {
 
template <concepts::KeyValueRdd R>
  requires concepts::KeyValueRdd<R>
class GroupByKeyRdd : public BaseRdd<GroupByKeyRdd<R>> {
public:
  using Base = BaseRdd<GroupByKeyRdd<R>>;
  friend Base;
 
public:
  using KeyType = utils::RddKeyType<R>;
  using ValueType = utils::RddValueType<R>;
 
  class Iterator : public std::forward_iterator_tag {
  public:
    using value_type =
        std::pair<KeyType, std::ranges::subrange<std::ranges::iterator_t<std::vector<ValueType>>>>;
    using difference_type = std::ptrdiff_t;
 
  public:
    explicit Iterator(
        std::ranges::subrange<std::ranges::iterator_t<std::ranges::range_value_t<R>>> original_data,
        bool is_end)
        : original_data_{std::move(original_data)},
          inner_data_mutex_{std::make_shared<std::shared_mutex>()},
          is_end_{is_end} {}
 
    explicit Iterator(bool is_end) : is_end_{is_end} {}
 
    Iterator() = default;
 
    value_type operator*() const {
      waitOrEvaluate();
      // Note that all operations to the shared data after any call of waitOrEvaluate() are read operations,
      // so mutex is no longer needed.
      checkOrInitInnerIterator();
      return {inner_iterator_->first, std::ranges::subrange(inner_iterator_->second)};
    }
 
    Iterator& operator++() {
      waitOrEvaluate();
      // Note that all operations to the shared data after any call of waitOrEvaluate() are read operations,
      // so mutex is no longer needed.
      checkOrInitInnerIterator();
      ++inner_iterator_;
      return *this;
    }
 
    Iterator operator++(int) {
      auto old = *this;
      ++(*this);
      return old;
    }
 
    bool operator==(const Iterator& other) const {
      if (isEnd() && other.isEnd()) {
        return true;
      }
      if (isEnd() || other.isEnd()) {
        return false;
      }
      return inner_data_ == other.inner_data_ && inner_iterator_ == other.inner_iterator_;
    }
 
    bool operator!=(const Iterator& other) const { return !(*this == other); }
 
  private:
    void evaluate() const {
      inner_data_ = std::make_shared<std::unordered_map<KeyType, std::vector<ValueType>>>();
      for (const auto& [key, value] : original_data_) {
        (*inner_data_)[key].push_back(value);
      }
    }
 
    void waitOrEvaluate() const {
      {
        std::shared_lock the_shared_lock(*inner_data_mutex_);
        if (inner_data_) {
          return;
        }
      }
      {
        std::unique_lock the_unique_lock(*inner_data_mutex_);
        if (inner_data_) {
          return;
        }
        evaluate();
      }
    }
 
    void checkOrInitInnerIterator() const {
      if (!inner_iterator_init_) {
        inner_iterator_ = std::ranges::begin(*inner_data_);
        inner_iterator_init_ = true;
      }
    }
 
    bool isEnd() const {
      return is_end_ || (inner_iterator_init_ && inner_iterator_ == std::ranges::end(*inner_data_));
    }
 
  private:
    // The shared state that stores the results of the group operation.
    mutable std::shared_ptr<std::unordered_map<KeyType, std::vector<ValueType>>> inner_data_{
        nullptr};
    // An iterator pointing tho the current element in the shared state.
    mutable std::ranges::iterator_t<std::unordered_map<KeyType, std::vector<ValueType>>>
        inner_iterator_{};
    // Whether `inner_iterator_` is initialized.
    mutable bool inner_iterator_init_{false};
    // The mutex to protect the shared state.
    mutable std::shared_ptr<std::shared_mutex> inner_data_mutex_{nullptr};
    std::ranges::subrange<std::ranges::iterator_t<std::ranges::range_value_t<R>>> original_data_;
    bool is_end_{false};
  };
 
public:
  explicit GroupByKeyRdd(const R& prev) : Base{prev} {
    for (const cpark::concepts::Split auto& split : prev) {
      splits_.emplace_back(std::ranges::subrange(Iterator{split, false}, Iterator{true}),
                           Base::context_);
      splits_.back().addDependency(split);
    }
  }
 
  GroupByKeyRdd(const GroupByKeyRdd&) = default;
  GroupByKeyRdd& operator=(const GroupByKeyRdd&) = default;
 
private:
  constexpr auto beginImpl() const { return std::ranges::begin(splits_); }
 
  constexpr auto endImpl() const { return std::ranges::end(splits_); }
 
private:
  std::vector<ViewSplit<std::ranges::subrange<Iterator>>> splits_;
};
 
// TODO: Helper classes to be added later.
 
 
}  // namespace cpark
 
#endif // CPARK_GROUP_BY_KEY_RDD_H