table_chunk.hpp

 
#pragma once
 
#include <array>
#include <memory>
#include <optional>
 
#include <cudf/contiguous_split.hpp>
#include <cudf/table/table.hpp>
#include <cudf/table/table_view.hpp>
 
#include <rapidsmpf/buffer/content_description.hpp>
#include <rapidsmpf/buffer/packed_data.hpp>
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/owning_wrapper.hpp>
#include <rapidsmpf/streaming/core/channel.hpp>
#include <rapidsmpf/streaming/core/context.hpp>
#include <rapidsmpf/streaming/core/node.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,
        std::size_t device_alloc_size,
        rmm::cuda_stream_view stream,
        OwningWrapper&& owner,
        ExclusiveView exclusive_view
    );
 
    TableChunk(
        std::unique_ptr<cudf::packed_columns> packed_columns, rmm::cuda_stream_view stream
    );
 
    TableChunk(std::unique_ptr<PackedData> packed_data);
 
    ~TableChunk() = default;
 
    TableChunk(TableChunk&&) = default;
 
    TableChunk& operator=(TableChunk&&) = default;
    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]] cudf::table_view table_view() const;
 
    [[nodiscard]] bool is_spillable() const;
 
    [[nodiscard]] TableChunk copy(MemoryReservation& reservation) const;
 
  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<cudf::packed_columns> packed_columns_;
    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