algeon-cugraph 0.1.2

Safe libcugraph execution component for Algeon
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

use cudf::{TypeId as CudfTypeId, interop::contract::TypeId as InteropTypeId};

use crate::types::DataType;

macro_rules! map_type_id_to_cugraph {
    ($type_id:expr, $type:path) => {
        match $type_id {
            <$type>::Int8 => Some(DataType::Int8),
            <$type>::Int16 => Some(DataType::Int16),
            <$type>::Int32 => Some(DataType::Int32),
            <$type>::Int64 => Some(DataType::Int64),
            <$type>::UInt8 => Some(DataType::Uint8),
            <$type>::UInt16 => Some(DataType::Uint16),
            <$type>::UInt32 => Some(DataType::Uint32),
            <$type>::UInt64 => Some(DataType::Uint64),
            <$type>::Float32 => Some(DataType::Float32),
            <$type>::Float64 => Some(DataType::Float64),
            <$type>::Bool8 => Some(DataType::Bool),
            _ => None,
        }
    };
}

/// Convert a cudf [`cudf::TypeId`] to a cugraph [`DataType`].
///
/// Returns `None` for cudf types that have no direct cugraph equivalent
/// (e.g. timestamps, strings, structs).
pub fn cudf_type_to_cugraph(type_id: CudfTypeId) -> Option<DataType> {
    map_type_id_to_cugraph!(type_id, CudfTypeId)
}

pub(super) fn borrowed_type_to_cugraph(type_id: InteropTypeId) -> Option<DataType> {
    map_type_id_to_cugraph!(type_id, InteropTypeId)
}

/// Convert a cugraph [`DataType`] to a cudf [`cudf::TypeId`].
///
/// Returns `None` for cuGraph types without a direct cuDF equivalent, including [`DataType::SizeT`].
pub fn cugraph_type_to_cudf(dtype: DataType) -> Option<CudfTypeId> {
    match dtype {
        DataType::Int8 => Some(CudfTypeId::Int8),
        DataType::Int16 => Some(CudfTypeId::Int16),
        DataType::Int32 => Some(CudfTypeId::Int32),
        DataType::Int64 => Some(CudfTypeId::Int64),
        DataType::Uint8 => Some(CudfTypeId::UInt8),
        DataType::Uint16 => Some(CudfTypeId::UInt16),
        DataType::Uint32 => Some(CudfTypeId::UInt32),
        DataType::Uint64 => Some(CudfTypeId::UInt64),
        DataType::Float32 => Some(CudfTypeId::Float32),
        DataType::Float64 => Some(CudfTypeId::Float64),
        DataType::Bool => Some(CudfTypeId::Bool8),
        DataType::SizeT => None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    const MAPPED_DATA_TYPES: &[DataType] = &[
        DataType::Int8,
        DataType::Int16,
        DataType::Int32,
        DataType::Int64,
        DataType::Uint8,
        DataType::Uint16,
        DataType::Uint32,
        DataType::Uint64,
        DataType::Float32,
        DataType::Float64,
        DataType::Bool,
    ];

    const MAPPED_TYPE_IDS: &[CudfTypeId] = &[
        CudfTypeId::Int8,
        CudfTypeId::Int16,
        CudfTypeId::Int32,
        CudfTypeId::Int64,
        CudfTypeId::UInt8,
        CudfTypeId::UInt16,
        CudfTypeId::UInt32,
        CudfTypeId::UInt64,
        CudfTypeId::Float32,
        CudfTypeId::Float64,
        CudfTypeId::Bool8,
    ];

    #[test]
    fn cugraph_to_cudf_roundtrip() {
        for &dt in MAPPED_DATA_TYPES {
            let cudf_tid =
                cugraph_type_to_cudf(dt).unwrap_or_else(|| panic!("cugraph_type_to_cudf returned None for {:?}", dt));
            let back = cudf_type_to_cugraph(cudf_tid)
                .unwrap_or_else(|| panic!("cudf_type_to_cugraph returned None for {:?}", cudf_tid));
            assert_eq!(dt, back, "round-trip failed: {:?} -> {:?} -> {:?}", dt, cudf_tid, back);
        }
    }

    #[test]
    fn cudf_to_cugraph_roundtrip() {
        for &tid in MAPPED_TYPE_IDS {
            let cugraph_dt =
                cudf_type_to_cugraph(tid).unwrap_or_else(|| panic!("cudf_type_to_cugraph returned None for {:?}", tid));
            let back = cugraph_type_to_cudf(cugraph_dt)
                .unwrap_or_else(|| panic!("cugraph_type_to_cudf returned None for {:?}", cugraph_dt));
            assert_eq!(tid, back, "round-trip failed: {:?} -> {:?} -> {:?}", tid, cugraph_dt, back);
        }
    }

    #[test]
    fn borrowed_type_to_cugraph_matches_cudf_mapping() {
        let borrowed = [
            (InteropTypeId::Int8, CudfTypeId::Int8),
            (InteropTypeId::Int16, CudfTypeId::Int16),
            (InteropTypeId::Int32, CudfTypeId::Int32),
            (InteropTypeId::Int64, CudfTypeId::Int64),
            (InteropTypeId::UInt8, CudfTypeId::UInt8),
            (InteropTypeId::UInt16, CudfTypeId::UInt16),
            (InteropTypeId::UInt32, CudfTypeId::UInt32),
            (InteropTypeId::UInt64, CudfTypeId::UInt64),
            (InteropTypeId::Float32, CudfTypeId::Float32),
            (InteropTypeId::Float64, CudfTypeId::Float64),
            (InteropTypeId::Bool8, CudfTypeId::Bool8),
        ];

        for (borrowed_tid, cudf_tid) in borrowed {
            assert_eq!(
                borrowed_type_to_cugraph(borrowed_tid),
                cudf_type_to_cugraph(cudf_tid),
                "borrowed and cudf mappings diverged for {:?} / {:?}",
                borrowed_tid,
                cudf_tid
            );
        }
    }
}