packed_data.hpp

 
#pragma once
 
#include <cstdint>
#include <memory>
#include <vector>
 
#include <rmm/cuda_stream_view.hpp>
#include <rmm/device_buffer.hpp>
 
#include <rapidsmpf/buffer/buffer.hpp>
#include <rapidsmpf/buffer/resource.hpp>
#include <rapidsmpf/error.hpp>
 
namespace rapidsmpf {
 
struct PackedData {
    std::unique_ptr<std::vector<std::uint8_t>> metadata;  
    std::unique_ptr<Buffer> data;  
 
    PackedData(
        std::unique_ptr<std::vector<std::uint8_t>> metadata, std::unique_ptr<Buffer> data
    )
        : metadata{std::move(metadata)}, data{std::move(data)} {
        RAPIDSMPF_EXPECTS(
            this->metadata != nullptr, "the metadata pointer cannot be null"
        );
        RAPIDSMPF_EXPECTS(this->data != nullptr, "the GPU data pointer cannot be null");
        RAPIDSMPF_EXPECTS(
            (this->metadata->size() > 0 || this->data->size == 0),
            "Empty Metadata and non-empty GPU data is not allowed"
        );
    }
 
    ~PackedData() = default;
 
    PackedData(PackedData&&) = default;
 
    PackedData& operator=(PackedData&&) = default;
    PackedData(PackedData const&) = delete;
    PackedData& operator=(PackedData const&) = delete;
 
    [[nodiscard]] bool empty() const {
        return metadata->empty() && data->size == 0;
    }
 
    [[nodiscard]] rmm::cuda_stream_view stream() const {
        return data->stream();
    }
 
    PackedData copy(MemoryReservation& reservation) const {
        auto dst = reservation.br()->allocate(data->size, data->stream(), reservation);
        buffer_copy(*dst, *data, data->size);
        return PackedData{
            std::make_unique<std::vector<std::uint8_t>>(*metadata), std::move(dst)
        };
    }
};
 
}  // namespace rapidsmpf