channel_metadata.hpp

 
#pragma once
 
#include <cstdint>
#include <memory>
#include <optional>
#include <vector>
 
#include <cuda_runtime_api.h>
 
#include <cudf/types.hpp>
 
#include <rapidsmpf/memory/buffer_resource.hpp>
#include <rapidsmpf/memory/content_description.hpp>
#include <rapidsmpf/streaming/core/message.hpp>
#include <rapidsmpf/streaming/cudf/table_chunk.hpp>
 
namespace rapidsmpf::streaming {
 
struct HashScheme {
    std::vector<cudf::size_type> column_indices;  
    int modulus;  
 
    bool operator==(HashScheme const&) const = default;
};
 
struct OrderKey {
    cudf::size_type column_index;  
    cudf::order order;  
    cudf::null_order null_order;  
 
    bool operator==(OrderKey const&) const = default;
 
    bool operator!=(OrderKey const&) const = default;
};
 
struct OrderScheme {
    std::vector<OrderKey> keys;  
    std::shared_ptr<TableChunk> boundaries;  
    bool strict_boundaries{false};
 
    OrderScheme() = default;
 
    OrderScheme(
        std::vector<OrderKey> keys,
        std::shared_ptr<TableChunk> boundaries,
        bool strict_boundaries = false
    );
 
    [[nodiscard]] OrderScheme with_keys(std::vector<OrderKey> new_keys) const;
 
    [[nodiscard]] bool boundaries_aligned_with(
        OrderScheme const& other, rapidsmpf::BufferResource const& br
    ) const;
};
 
struct PartitioningSpec {
    enum class Type : std::uint8_t {
        NONE,  
        INHERIT,  
        HASH,  
        ORDER,  
    };
 
    Type type = Type::NONE;  
    std::optional<HashScheme> hash;  
    std::optional<OrderScheme> order;  
 
    static PartitioningSpec none() {
        return {};
    }
 
    static PartitioningSpec inherit() {
        return {.type = Type::INHERIT, .hash = std::nullopt, .order = std::nullopt};
    }
 
    static PartitioningSpec from_hash(HashScheme h) {
        return {.type = Type::HASH, .hash = std::move(h), .order = std::nullopt};
    }
 
    static PartitioningSpec from_order(OrderScheme o);
};
 
struct Partitioning {
    PartitioningSpec inter_rank;
    PartitioningSpec local;
};
 
struct ChannelMetadata {
    std::uint64_t local_count{};  
    Partitioning partitioning;  
    bool duplicated{};  
 
    ChannelMetadata() = default;
 
    ChannelMetadata(
        std::uint64_t local_count, Partitioning partitioning = {}, bool duplicated = false
    )
        : local_count(local_count),
          partitioning(std::move(partitioning)),
          duplicated(duplicated) {}
};
 
Message to_message(std::uint64_t sequence_number, std::unique_ptr<ChannelMetadata> m);
 
}  // namespace rapidsmpf::streaming