channel_metadata.hpp

 
#pragma once
 
#include <cstdint>
#include <optional>
#include <vector>
 
#include <cudf/types.hpp>
 
#include <rapidsmpf/streaming/core/message.hpp>
 
namespace rapidsmpf::streaming {
 
struct HashScheme {
    std::vector<cudf::size_type> column_indices;  
    int modulus;  
 
    bool operator==(HashScheme const&) const = default;
};
 
struct PartitioningSpec {
    enum class Type : std::uint8_t {
        NONE,  
        INHERIT,  
        HASH,  
    };
 
    Type type = Type::NONE;  
    std::optional<HashScheme> hash;  
 
    static PartitioningSpec none() {
        return {};
    }
 
    static PartitioningSpec inherit() {
        return {.type = Type::INHERIT, .hash = std::nullopt};
    }
 
    static PartitioningSpec from_hash(HashScheme h) {
        return {.type = Type::HASH, .hash = std::move(h)};
    }
 
    bool operator==(PartitioningSpec const&) const = default;
};
 
struct Partitioning {
    PartitioningSpec inter_rank;
    PartitioningSpec local;
 
    bool operator==(Partitioning const&) const = default;
};
 
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) {}
 
    bool operator==(ChannelMetadata const&) const = default;
};
 
Message to_message(std::uint64_t sequence_number, std::unique_ptr<ChannelMetadata> m);
 
}  // namespace rapidsmpf::streaming