partition_by_rdd.h

#ifndef CPARK_PARTITION_BY_RDD_H
#define CPARK_PARTITION_BY_RDD_H
 
#include <concepts>
#include <ranges>
 
#include "base_rdd.h"
#include "filter_rdd.h"
#include "merged_view.h"
#include "utils.h"
 
namespace cpark {
 
template <concepts::KeyValueRdd R, typename Partitioner = std::hash<utils::RddKeyType<R>>>
class PartitionByRdd : public BaseRdd<PartitionByRdd<R, Partitioner>> {
public:
  using Base = BaseRdd<PartitionByRdd<R, Partitioner>>;
  friend Base;
 
public:
  PartitionByRdd(const R& prev, Partitioner partitioner)
      : Base{prev}, partitioner_{std::move(partitioner)} {
    for (size_t i : std::views::iota(0u, Base::splits_num_)) {
      auto split_view =
          FilterView(MergedSameView(prev), PartitionerHelper{i, Base::splits_num_, &partitioner_});
      splits_.emplace_back(std::move(split_view), Base::context_);
      for (const concepts::Split auto split : prev) {
        splits_.back().addDependency(split);
      }
    }
  }
 
  explicit PartitionByRdd(const R& prev) : PartitionByRdd{prev, Partitioner{}} {}
 
  PartitionByRdd(const PartitionByRdd&) = default;
  PartitionByRdd& operator=(const PartitionByRdd&) = default;
 
private:
  constexpr auto beginImpl() const { return std::ranges::begin(splits_); }
 
  constexpr auto endImpl() const { return std::ranges::end(splits_); }
 
private:
  class PartitionerHelper {
  public:
    PartitionerHelper(size_t split_index, size_t splits_num, Partitioner* partitioner)
        : split_index_{split_index}, splits_num_{splits_num}, partitioner_{partitioner} {}
 
    bool operator()(const utils::RddElementType<R>& x) const {
      return (*partitioner_)(x.first) % splits_num_ == split_index_;
    }
 
  private:
    size_t split_index_, splits_num_;
    Partitioner* partitioner_;
  };
 
private:
  std::vector<ViewSplit<FilterView<MergedSameView<R>, PartitionerHelper>>> splits_;
  Partitioner partitioner_;
};
 
template <typename... Args>
class PartitionBy;
 
template <typename Partitioner>
class PartitionBy<Partitioner> {
public:
  explicit PartitionBy(Partitioner partitioner) : partitioner_{std::move(partitioner)} {}
 
  template <concepts::Rdd R, typename K = utils::RddKeyType<R>>
    requires std::invocable<Partitioner, K> &&
             std::convertible_to<std::invoke_result_t<Partitioner, K>, size_t>
  auto operator()(const R& r) const {
    return PartitionByRdd(r, partitioner_);
  }
 
private:
  Partitioner partitioner_;
};
 
template <>
class PartitionBy<> {
public:
  template <concepts::Rdd R>
  auto operator()(const R& r) const {
    return PartitionByRdd(r);
  }
};
 
template <cpark::concepts::Rdd R, typename P>
auto operator|(const R& r, const P& p) {
  return p(r);
}
 
 
}  // namespace cpark
 
#endif  // CPARK_PARTITION_BY_RDD_H