cudf-cxx 0.2.0

cxx-based FFI bridge between Rust and NVIDIA libcudf C++ API
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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189

#include "reduction_shim.h"
#include <cudf/reduction.hpp>
#include <cudf/utilities/default_stream.hpp>
#include <stdexcept>

namespace cudf_shims {

namespace {

/// Create a reduce aggregation from an integer kind.
std::unique_ptr<cudf::reduce_aggregation> make_reduce_agg(int32_t agg_kind) {
    switch (agg_kind) {
        case 0: return cudf::make_sum_aggregation<cudf::reduce_aggregation>();
        case 1: return cudf::make_product_aggregation<cudf::reduce_aggregation>();
        case 2: return cudf::make_min_aggregation<cudf::reduce_aggregation>();
        case 3: return cudf::make_max_aggregation<cudf::reduce_aggregation>();
        case 4: return cudf::make_sum_of_squares_aggregation<cudf::reduce_aggregation>();
        case 5: return cudf::make_mean_aggregation<cudf::reduce_aggregation>();
        case 6: return cudf::make_variance_aggregation<cudf::reduce_aggregation>();
        case 7: return cudf::make_std_aggregation<cudf::reduce_aggregation>();
        case 8: return cudf::make_any_aggregation<cudf::reduce_aggregation>();
        case 9: return cudf::make_all_aggregation<cudf::reduce_aggregation>();
        case 10: return cudf::make_median_aggregation<cudf::reduce_aggregation>();
        default:
            throw std::runtime_error("unknown reduce aggregation kind: " + std::to_string(agg_kind));
    }
}

/// Create a scan aggregation from an integer kind.
std::unique_ptr<cudf::scan_aggregation> make_scan_agg(int32_t agg_kind) {
    switch (agg_kind) {
        case 0: return cudf::make_sum_aggregation<cudf::scan_aggregation>();
        case 1: return cudf::make_product_aggregation<cudf::scan_aggregation>();
        case 2: return cudf::make_min_aggregation<cudf::scan_aggregation>();
        case 3: return cudf::make_max_aggregation<cudf::scan_aggregation>();
        default:
            throw std::runtime_error("unknown scan aggregation kind: " + std::to_string(agg_kind));
    }
}

/// Create a segmented reduce aggregation from an integer kind.
std::unique_ptr<cudf::segmented_reduce_aggregation> make_segmented_reduce_agg(int32_t agg_kind) {
    switch (agg_kind) {
        case 0: return cudf::make_sum_aggregation<cudf::segmented_reduce_aggregation>();
        case 1: return cudf::make_product_aggregation<cudf::segmented_reduce_aggregation>();
        case 2: return cudf::make_min_aggregation<cudf::segmented_reduce_aggregation>();
        case 3: return cudf::make_max_aggregation<cudf::segmented_reduce_aggregation>();
        case 4: return cudf::make_sum_of_squares_aggregation<cudf::segmented_reduce_aggregation>();
        case 5: return cudf::make_mean_aggregation<cudf::segmented_reduce_aggregation>();
        case 6: return cudf::make_variance_aggregation<cudf::segmented_reduce_aggregation>();
        case 7: return cudf::make_std_aggregation<cudf::segmented_reduce_aggregation>();
        case 8: return cudf::make_any_aggregation<cudf::segmented_reduce_aggregation>();
        case 9: return cudf::make_all_aggregation<cudf::segmented_reduce_aggregation>();
        case 10: return cudf::make_median_aggregation<cudf::segmented_reduce_aggregation>();
        default:
            throw std::runtime_error("unknown segmented reduce aggregation kind: " + std::to_string(agg_kind));
    }
}

} // anonymous namespace

std::unique_ptr<OwnedScalar> reduce(
    const OwnedColumn& col,
    int32_t agg_kind,
    int32_t output_type_id)
{
    auto agg = make_reduce_agg(agg_kind);
    auto out_type = cudf::data_type{static_cast<cudf::type_id>(output_type_id)};

    auto result = cudf::reduce(
        col.view(),
        *agg,
        out_type);

    return std::make_unique<OwnedScalar>(std::move(result));
}

std::unique_ptr<OwnedColumn> scan(
    const OwnedColumn& col,
    int32_t agg_kind,
    bool inclusive)
{
    auto agg = make_scan_agg(agg_kind);
    auto scan_type = inclusive
        ? cudf::scan_type::INCLUSIVE
        : cudf::scan_type::EXCLUSIVE;

    auto result = cudf::scan(
        col.view(),
        *agg,
        scan_type);

    return std::make_unique<OwnedColumn>(std::move(result));
}

std::unique_ptr<OwnedColumn> segmented_reduce(
    const OwnedColumn& col,
    const OwnedColumn& offsets,
    int32_t agg_kind,
    int32_t output_type_id,
    bool include_nulls)
{
    auto agg = make_segmented_reduce_agg(agg_kind);
    auto out_type = cudf::data_type{static_cast<cudf::type_id>(output_type_id)};
    auto null_policy = include_nulls
        ? cudf::null_policy::INCLUDE
        : cudf::null_policy::EXCLUDE;

    auto result = cudf::segmented_reduce(
        col.view(),
        offsets.view(),
        *agg,
        out_type,
        null_policy);

    return std::make_unique<OwnedColumn>(std::move(result));
}

std::unique_ptr<MinMaxResult> minmax(const OwnedColumn& col) {
    auto [min_scalar, max_scalar] = cudf::minmax(col.view());
    auto result = std::make_unique<MinMaxResult>();
    result->min_val = std::make_unique<OwnedScalar>(std::move(min_scalar));
    result->max_val = std::make_unique<OwnedScalar>(std::move(max_scalar));
    return result;
}

std::unique_ptr<OwnedScalar> minmax_take_min(MinMaxResult& result) {
    return std::move(result.min_val);
}

std::unique_ptr<OwnedScalar> minmax_take_max(MinMaxResult& result) {
    return std::move(result.max_val);
}

std::unique_ptr<OwnedScalar> reduce_with_init(
    const OwnedColumn& col,
    int32_t agg_kind,
    int32_t output_type_id,
    const OwnedScalar& init)
{
    auto agg = make_reduce_agg(agg_kind);
    auto out_type = cudf::data_type{static_cast<cudf::type_id>(output_type_id)};

    auto result = cudf::reduce(
        col.view(),
        *agg,
        out_type,
        std::optional<std::reference_wrapper<cudf::scalar const>>{*init.inner});

    return std::make_unique<OwnedScalar>(std::move(result));
}

bool is_valid_reduction_aggregation(int32_t source_type_id, int32_t agg_kind)
{
    auto src_type = cudf::data_type{static_cast<cudf::type_id>(source_type_id)};
    // Use make_reduce_agg to get the correct aggregation (same mapping as reduce()),
    // then extract its Kind — avoids assuming ReduceOp numbering == aggregation::Kind.
    auto agg = make_reduce_agg(agg_kind);
    return cudf::reduction::is_valid_aggregation(src_type, agg->kind);
}


std::unique_ptr<OwnedScalar> reduce_var_with_ddof(
    const OwnedColumn& col,
    int32_t ddof,
    int32_t output_type_id)
{
    auto agg = cudf::make_variance_aggregation<cudf::reduce_aggregation>(
        static_cast<cudf::size_type>(ddof));
    auto out_type = cudf::data_type{static_cast<cudf::type_id>(output_type_id)};

    auto result = cudf::reduce(col.view(), *agg, out_type);
    return std::make_unique<OwnedScalar>(std::move(result));
}

std::unique_ptr<OwnedScalar> reduce_std_with_ddof(
    const OwnedColumn& col,
    int32_t ddof,
    int32_t output_type_id)
{
    auto agg = cudf::make_std_aggregation<cudf::reduce_aggregation>(
        static_cast<cudf::size_type>(ddof));
    auto out_type = cudf::data_type{static_cast<cudf::type_id>(output_type_id)};

    auto result = cudf::reduce(col.view(), *agg, out_type);
    return std::make_unique<OwnedScalar>(std::move(result));
}

} // namespace cudf_shims