postbox.hpp

 
#pragma once
 
#include <functional>
#include <mutex>
#include <string>
#include <unordered_map>
#include <vector>
 
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/shuffler/chunk.hpp>
 
namespace rapidsmpf::shuffler::detail {
 
template <typename KeyType>
class PostBox {
  public:
    using key_type = KeyType;  
 
    template <typename Fn>
    PostBox(Fn&& key_map_fn, size_t num_keys_hint = 0)
        : key_map_fn_(std::move(key_map_fn)) {
        if (num_keys_hint > 0) {
            pigeonhole_.reserve(num_keys_hint);
        }
    }
 
    void insert(Chunk&& chunk);
 
    bool is_empty(PartID pid) const;
 
    [[nodiscard]] Chunk extract(PartID pid, ChunkID cid);
 
    std::unordered_map<ChunkID, Chunk> extract(PartID pid);
 
    std::unordered_map<ChunkID, Chunk> extract_by_key(KeyType key);
 
    std::vector<Chunk> extract_all_ready();
 
    [[nodiscard]] bool empty() const;
 
    [[nodiscard]] std::vector<std::tuple<key_type, ChunkID, std::size_t>> search(
        MemoryType mem_type
    ) const;
 
    [[nodiscard]] std::string str() const;
 
  private:
    // TODO: more fine-grained locking e.g. by locking each partition individually.
    mutable std::mutex mutex_;
    std::function<key_type(PartID)>
        key_map_fn_;  
    std::unordered_map<key_type, std::unordered_map<ChunkID, Chunk>>
        pigeonhole_;  
};
 
template <typename KeyType>
inline std::ostream& operator<<(std::ostream& os, PostBox<KeyType> const& obj) {
    os << obj.str();
    return os;
}
 
}  // namespace rapidsmpf::shuffler::detail