packed_data.hpp

 
#pragma once
 
#include <rapidsmpf/buffer/content_description.hpp>
#include <rapidsmpf/buffer/packed_data.hpp>
#include <rapidsmpf/streaming/core/message.hpp>
 
namespace rapidsmpf::streaming {
 
struct PackedDataChunk {
    PackedData data;
};
 
inline ContentDescription get_content_description(PackedDataChunk const& obj) {
    return ContentDescription{
        {{obj.data.data->mem_type(), obj.data.data->size}},
        ContentDescription::Spillable::YES
    };
}
 
Message to_message(
    std::uint64_t sequence_number, std::unique_ptr<PackedDataChunk> chunk
) {
    auto cd = get_content_description(*chunk);
    return Message{
        sequence_number,
        std::move(chunk),
        cd,
        [](Message const& msg, MemoryReservation& reservation) -> Message {
            auto const& self = msg.get<PackedDataChunk>();
            auto chunk = std::make_unique<PackedDataChunk>(self.data.copy(reservation));
            auto cd = get_content_description(*chunk);
            return Message{msg.sequence_number(), std::move(chunk), cd, msg.copy_cb()};
        }
    };
}
 
}  // namespace rapidsmpf::streaming