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

#include "join_shim.h"
#include <cudf/join/join.hpp>
#include <cudf/join/filtered_join.hpp>
#include <cudf/table/table.hpp>
#include <cudf/column/column_factories.hpp>

namespace cudf_shims {

namespace {

/// Package a pair of gather map device_uvectors into a 2-column OwnedTable.
std::unique_ptr<OwnedTable> package_gather_maps(
    std::unique_ptr<rmm::device_uvector<cudf::size_type>> left_map,
    std::unique_ptr<rmm::device_uvector<cudf::size_type>> right_map)
{
    // Capture sizes BEFORE release() — C++ function argument evaluation
    // order is unspecified, so release() could be evaluated before size().
    auto left_size = static_cast<cudf::size_type>(left_map->size());
    auto left_buf = left_map->release();
    auto left_col = std::make_unique<cudf::column>(
        cudf::data_type{cudf::type_id::INT32},
        left_size,
        std::move(left_buf),
        rmm::device_buffer{}, 0);

    auto right_size = static_cast<cudf::size_type>(right_map->size());
    auto right_buf = right_map->release();
    auto right_col = std::make_unique<cudf::column>(
        cudf::data_type{cudf::type_id::INT32},
        right_size,
        std::move(right_buf),
        rmm::device_buffer{}, 0);

    std::vector<std::unique_ptr<cudf::column>> cols;
    cols.push_back(std::move(left_col));
    cols.push_back(std::move(right_col));

    auto table = std::make_unique<cudf::table>(std::move(cols));
    return std::make_unique<OwnedTable>(std::move(table));
}

} // anonymous namespace

std::unique_ptr<OwnedTable> inner_join(
    const OwnedTable& left_keys, const OwnedTable& right_keys)
{
    auto [left_map, right_map] = cudf::inner_join(
        left_keys.view(), right_keys.view());
    return package_gather_maps(std::move(left_map), std::move(right_map));
}

std::unique_ptr<OwnedTable> left_join(
    const OwnedTable& left_keys, const OwnedTable& right_keys)
{
    auto [left_map, right_map] = cudf::left_join(
        left_keys.view(), right_keys.view());
    return package_gather_maps(std::move(left_map), std::move(right_map));
}

std::unique_ptr<OwnedTable> full_join(
    const OwnedTable& left_keys, const OwnedTable& right_keys)
{
    auto [left_map, right_map] = cudf::full_join(
        left_keys.view(), right_keys.view());
    return package_gather_maps(std::move(left_map), std::move(right_map));
}

std::unique_ptr<OwnedTable> cross_join(
    const OwnedTable& left, const OwnedTable& right)
{
    auto result = cudf::cross_join(left.view(), right.view());
    return std::make_unique<OwnedTable>(std::move(result));
}

namespace {

/// Package a single gather map device_uvector into a 1-column OwnedTable.
std::unique_ptr<OwnedTable> package_single_map(
    std::unique_ptr<rmm::device_uvector<cudf::size_type>> map)
{
    // Capture size BEFORE release() — eval order is unspecified in C++
    auto map_size = static_cast<cudf::size_type>(map->size());
    auto map_buf = map->release();
    auto col = std::make_unique<cudf::column>(
        cudf::data_type{cudf::type_id::INT32},
        map_size,
        std::move(map_buf),
        rmm::device_buffer{}, 0);

    std::vector<std::unique_ptr<cudf::column>> cols;
    cols.push_back(std::move(col));

    auto table = std::make_unique<cudf::table>(std::move(cols));
    return std::make_unique<OwnedTable>(std::move(table));
}

} // anonymous namespace

std::unique_ptr<OwnedTable> left_semi_join(
    const OwnedTable& left_keys, const OwnedTable& right_keys)
{
    cudf::filtered_join fj(right_keys.view(), cudf::null_equality::EQUAL, cudf::set_as_build_table::RIGHT, 0.5);
    auto map = fj.semi_join(left_keys.view());
    return package_single_map(std::move(map));
}

std::unique_ptr<OwnedTable> left_anti_join(
    const OwnedTable& left_keys, const OwnedTable& right_keys)
{
    cudf::filtered_join fj(right_keys.view(), cudf::null_equality::EQUAL, cudf::set_as_build_table::RIGHT, 0.5);
    auto map = fj.anti_join(left_keys.view());
    return package_single_map(std::move(map));
}

// ── Hash Join ─────────────────────────────────────────────────

std::unique_ptr<OwnedHashJoin> hash_join_create(const OwnedTable& build) {
    auto hj = std::make_unique<cudf::hash_join>(
        build.view(),
        cudf::null_equality::EQUAL);
    return std::make_unique<OwnedHashJoin>(std::move(hj));
}

std::unique_ptr<OwnedTable> hash_join_inner(
    const OwnedHashJoin& hj, const OwnedTable& probe)
{
    auto [left_map, right_map] = hj.inner->inner_join(probe.view());
    return package_gather_maps(std::move(left_map), std::move(right_map));
}

std::unique_ptr<OwnedTable> hash_join_left(
    const OwnedHashJoin& hj, const OwnedTable& probe)
{
    auto [left_map, right_map] = hj.inner->left_join(probe.view());
    return package_gather_maps(std::move(left_map), std::move(right_map));
}

std::unique_ptr<OwnedTable> hash_join_full(
    const OwnedHashJoin& hj, const OwnedTable& probe)
{
    auto [left_map, right_map] = hj.inner->full_join(probe.view());
    return package_gather_maps(std::move(left_map), std::move(right_map));
}

int64_t hash_join_inner_size(const OwnedHashJoin& hj, const OwnedTable& probe) {
    return static_cast<int64_t>(hj.inner->inner_join_size(probe.view()));
}

int64_t hash_join_left_size(const OwnedHashJoin& hj, const OwnedTable& probe) {
    return static_cast<int64_t>(hj.inner->left_join_size(probe.view()));
}

int64_t hash_join_full_size(const OwnedHashJoin& hj, const OwnedTable& probe) {
    return static_cast<int64_t>(hj.inner->full_join_size(probe.view()));
}

} // namespace cudf_shims