batch_t.hpp

/****************************************************************************
 *  hipipe library
 *  Copyright (c) 2018, Iterait a.s.
 *  Author(s) Filip Matzner, Adam Blazek
 *
 *  This file is distributed under the MIT License.
 *  See the accompanying file LICENSE.txt for the complete license agreement.
 ****************************************************************************/
 
#pragma once
 
#include <hipipe/build_config.hpp>
#include <hipipe/core/stream/column_t.hpp>
 
#include <memory>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <typeindex>
#include <typeinfo>
#include <unordered_map>
 
namespace hipipe::stream {
 
class batch {
private:
 
    std::unordered_map<std::type_index, std::unique_ptr<abstract_column>> columns_;
 
    template<typename Column>
    void throw_check_contains() const
    {
        if (!columns_.count(std::type_index{typeid(Column)})) {
            throw std::runtime_error{
              std::string{"Trying to retrieve column `"} + Column{}.name()
              + "`, but the batch contains no such column."};
        }
    }
 
public:
 
    // constructors //
 
    batch() = default;
    batch(const batch&) = delete;
    batch(batch&&) = default;
 
    // direct access //
 
    template<typename Column>
    std::unique_ptr<abstract_column>& at()
    {
        throw_check_contains<Column>();
        return columns_.at(std::type_index{typeid(Column)});
    }
 
    // value extraction //
 
    template<typename Column>
    typename Column::data_type& extract()
    {
        throw_check_contains<Column>();
        return columns_.at(std::type_index{typeid(Column)})->extract<Column>();
    }
 
    template<typename Column>
    const typename Column::data_type& extract() const
    {
        throw_check_contains<Column>();
        return columns_.at(std::type_index{typeid(Column)})->extract<Column>();
    }
 
    // column insertion/rewrite //
 
    template<typename Column, typename... Args>
    void insert_or_assign(Args&&... args)
    {
        static_assert(std::is_constructible_v<Column, Args&&...>,
          "Cannot construct the given column from the provided arguments.");
        columns_.insert_or_assign(
          std::type_index{typeid(Column)},
          std::make_unique<Column>(std::forward<Args>(args)...));
    }
 
    template<typename Column>
    void raw_insert_or_assign(std::unique_ptr<abstract_column> column_ptr)
    {
        columns_.insert_or_assign(std::type_index{typeid(Column)}, std::move(column_ptr));
    }
 
    // column check //
 
    std::size_t size() const
    {
        return columns_.size();
    }
 
    template<typename Column>
    bool contains() const
    {
        return columns_.count(std::type_index{typeid(Column)});
    }
 
    // column removal //
 
    template<typename Column>
    void erase()
    {
        throw_check_contains<Column>();
        columns_.erase(std::type_index{typeid(Column)});
    }
 
    // batching utilities //
 
    std::size_t batch_size() const
    {
        if (columns_.empty()) return 0;
        std::size_t batch_size = columns_.begin()->second->size();
        for (auto it = ++columns_.begin(); it != columns_.end(); ++it) {
            if (it->second->size() != batch_size) {
                throw std::runtime_error{"hipipe: Canot deduce a batch size from a batch "
                  "with columns of different size (`" + it->second->name() + "`)."};
            }
            batch_size = it->second->size();
        }
        return batch_size;
    }
 
    batch take(std::size_t n)
    {
        batch new_batch;
        for (const auto& [key, col] : columns_) {
            new_batch.columns_.insert_or_assign(key, col->take(n));
        }
        return new_batch;
    }
 
    void push_back(batch rhs)
    {
        for (auto& [key, col] : rhs.columns_) {
            if (!columns_.count(key)) {
                columns_[key] = std::move(col);
            } else {
                columns_.at(key)->push_back(std::move(col));
            }
        }
    }
 
    #ifdef HIPIPE_BUILD_PYTHON
    boost::python::dict to_python()
    {
        boost::python::dict res;
        for (auto it = columns_.begin(); it != columns_.end(); ++it) {
            res[it->second->name()] = it->second->to_python();
        }
        columns_.clear();
        return res;
    }
    #endif
};
 
 
using batch_t = batch;
 
}  // namespace hipipe::stream