table_chunk.hpp

 
#pragma once
 
#include <array>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>
 
#include <cudf/packed_types.hpp>
#include <cudf/table/table.hpp>
#include <cudf/table/table_view.hpp>
#include <rmm/cuda_stream_view.hpp>
 
#include <coro/task.hpp>
 
#include <rapidsmpf/memory/content_description.hpp>
#include <rapidsmpf/memory/packed_data.hpp>
#include <rapidsmpf/owning_wrapper.hpp>
#include <rapidsmpf/streaming/core/context.hpp>
#include <rapidsmpf/streaming/core/memory_reserve_or_wait.hpp>
#include <rapidsmpf/streaming/core/message.hpp>
 
namespace rapidsmpf::streaming {
 
class TableChunk {
  public:
    enum class ExclusiveView : bool {
        NO,
        YES,
    };
 
    TableChunk(std::unique_ptr<cudf::table> table, rmm::cuda_stream_view stream);
 
    TableChunk(
        cudf::table_view table_view,
        rmm::cuda_stream_view stream,
        OwningWrapper&& owner,
        ExclusiveView exclusive_view
    );
 
    TableChunk(std::unique_ptr<PackedData> packed_data);
 
    ~TableChunk() = default;
 
    TableChunk(TableChunk&& other) noexcept;
 
    TableChunk& operator=(TableChunk&& other) noexcept;
    TableChunk(TableChunk const&) = delete;
    TableChunk& operator=(TableChunk const&) = delete;
 
    [[nodiscard]] rmm::cuda_stream_view stream() const noexcept;
 
    [[nodiscard]] std::size_t data_alloc_size(MemoryType mem_type) const;
 
    [[nodiscard]] bool is_available() const noexcept;
 
    [[nodiscard]] std::size_t make_available_cost() const noexcept;
 
    [[nodiscard]] TableChunk make_available(MemoryReservation& reservation);
 
    [[nodiscard]] TableChunk make_available(MemoryReservation&& reservation);
 
    [[nodiscard]] coro::task<TableChunk> make_available(
        std::shared_ptr<Context> ctx,
        std::int64_t net_memory_delta = MemoryReserveOrWait::missing_net_memory_delta
    );
 
    [[nodiscard]] cudf::table_view table_view() const;
 
    [[nodiscard]] bool is_spillable() const;
 
    [[nodiscard]] TableChunk copy(MemoryReservation& reservation) const;
 
    [[nodiscard]] std::pair<cudf::size_type, cudf::size_type> shape() const noexcept;
 
  private:
    OwningWrapper owner_{};
 
    // At most, one of the following unique pointers is non-null. If all of them are null,
    // the TableChunk is a non-owning view.
    // TODO: use a variant and drop the unique pointers?
    std::unique_ptr<cudf::table> table_;
    std::unique_ptr<PackedData> packed_data_;
 
    // Has value iff this TableChunk is available.
    std::optional<cudf::table_view> table_view_;
 
    // Zero initialized data allocation size (one for each memory type).
    std::array<std::size_t, MEMORY_TYPES.size()> data_alloc_size_ = {};
    std::size_t make_available_cost_;  // For now, only device memory cost is tracked.
 
    rmm::cuda_stream_view stream_;
    bool is_spillable_;
};
 
ContentDescription get_content_description(TableChunk const& obj);
 
Message to_message(std::uint64_t sequence_number, std::unique_ptr<TableChunk> chunk);
 
}  // namespace rapidsmpf::streaming