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

#include "rolling_shim.h"
#include <cudf/rolling.hpp>
#include <cudf/aggregation.hpp>
#include <cudf/utilities/default_stream.hpp>
#include <stdexcept>

namespace cudf_shims {

namespace {

/// Create a rolling aggregation from an integer kind.
std::unique_ptr<cudf::rolling_aggregation> make_rolling_agg(int32_t agg_kind) {
    switch (agg_kind) {
        case 0: return cudf::make_sum_aggregation<cudf::rolling_aggregation>();
        case 1: return cudf::make_min_aggregation<cudf::rolling_aggregation>();
        case 2: return cudf::make_max_aggregation<cudf::rolling_aggregation>();
        case 3: return cudf::make_count_aggregation<cudf::rolling_aggregation>();
        case 4: return cudf::make_mean_aggregation<cudf::rolling_aggregation>();
        case 5: return cudf::make_collect_list_aggregation<cudf::rolling_aggregation>();
        case 6: return cudf::make_row_number_aggregation<cudf::rolling_aggregation>();
        case 7: return cudf::make_lead_aggregation<cudf::rolling_aggregation>(1);
        case 8: return cudf::make_lag_aggregation<cudf::rolling_aggregation>(1);
        default:
            throw std::runtime_error("unknown rolling aggregation kind: " + std::to_string(agg_kind));
    }
}

} // anonymous namespace

std::unique_ptr<OwnedColumn> rolling_window(
    const OwnedColumn& col,
    int32_t preceding,
    int32_t following,
    int32_t min_periods,
    int32_t agg_kind)
{
    auto agg = make_rolling_agg(agg_kind);

    auto result = cudf::rolling_window(
        col.view(),
        preceding,
        following,
        min_periods,
        *agg);

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

std::unique_ptr<OwnedColumn> grouped_rolling_window(
    const OwnedTable& group_keys,
    const OwnedColumn& input,
    int32_t preceding,
    int32_t following,
    int32_t min_periods,
    int32_t agg_kind)
{
    auto agg = make_rolling_agg(agg_kind);

    auto result = cudf::grouped_rolling_window(
        group_keys.view(),
        input.view(),
        preceding,
        following,
        min_periods,
        *agg);

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

std::unique_ptr<OwnedColumn> rolling_window_variable(
    const OwnedColumn& col,
    const OwnedColumn& preceding_col,
    const OwnedColumn& following_col,
    int32_t min_periods,
    int32_t agg_kind)
{
    auto agg = make_rolling_agg(agg_kind);

    auto result = cudf::rolling_window(
        col.view(),
        preceding_col.view(),
        following_col.view(),
        min_periods,
        *agg);

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

bool is_valid_rolling_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_rolling_agg to get the correct aggregation (same mapping as rolling_window()),
    // then extract its Kind — avoids assuming RollingAgg numbering == aggregation::Kind.
    auto agg = make_rolling_agg(agg_kind);
    return cudf::is_valid_rolling_aggregation(src_type, agg->kind);
}

} // namespace cudf_shims