cudf 0.3.1

Safe Rust bindings for NVIDIA libcudf -- GPU-accelerated DataFrame operations
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
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316

//! GPU-accelerated groupby aggregation.
//!
//! Provides a builder-pattern API for performing groupby operations on
//! GPU-resident tables.
//!
//! # Examples
//!
//! ```rust,no_run
//! use cudf::{Column, Table};
//! use cudf::groupby::GroupBy;
//! use cudf::aggregation::AggregationKind;
//!
//! // Keys: group labels
//! let keys_col = Column::from_slice(&[1i32, 1, 2, 2, 3]).unwrap();
//! let keys = Table::new(vec![keys_col]).unwrap();
//!
//! // Values: columns to aggregate
//! let val_a = Column::from_slice(&[10.0f64, 20.0, 30.0, 40.0, 50.0]).unwrap();
//! let val_b = Column::from_slice(&[1i32, 2, 3, 4, 5]).unwrap();
//! let values = Table::new(vec![val_a, val_b]).unwrap();
//!
//! let result = GroupBy::new(&keys)
//!     .agg(0, AggregationKind::Sum)
//!     .agg(1, AggregationKind::Mean)
//!     .execute(&values)
//!     .unwrap();
//! ```

use crate::aggregation::{Aggregation, AggregationKind};
use crate::column::Column;
use crate::error::{CudfError, Result};
use crate::table::Table;
use crate::types::checked_i32;

/// A groupby operation builder.
///
/// Accumulates aggregation requests and executes them against a values table.
/// The result is a [`Table`] with key columns first, followed by aggregation
/// result columns in the order they were added.
pub struct GroupBy<'a> {
    keys: &'a Table,
    requests: Vec<(usize, Aggregation)>,
}

impl<'a> GroupBy<'a> {
    /// Create a new groupby builder with the given key columns.
    pub fn new(keys: &'a Table) -> Self {
        Self {
            keys,
            requests: Vec::new(),
        }
    }

    /// Add an aggregation request for a value column.
    ///
    /// `column` is the 0-based index into the values table.
    /// `kind` specifies the aggregation to compute.
    ///
    /// Returns `self` for method chaining.
    pub fn agg(mut self, column: usize, kind: AggregationKind) -> Self {
        self.requests.push((column, Aggregation::new(kind)));
        self
    }

    /// Add a pre-built [`Aggregation`] for a value column.
    ///
    /// Use this when you need to configure null handling or other options
    /// before creating the aggregation.
    pub fn agg_with(mut self, column: usize, aggregation: Aggregation) -> Self {
        self.requests.push((column, aggregation));
        self
    }

    /// Execute the groupby, returning a table with key columns followed by
    /// aggregation result columns.
    ///
    /// # Errors
    ///
    /// Returns an error if column indices are out of bounds, aggregation
    /// types are incompatible with column types, or a GPU error occurs.
    pub fn execute(self, values: &Table) -> Result<Table> {
        if self.requests.is_empty() {
            return Err(CudfError::InvalidArgument(
                "groupby requires at least one aggregation request".to_string(),
            ));
        }

        let mut builder = cudf_cxx::groupby::ffi::groupby_new(&self.keys.inner);

        for (col_idx, agg) in self.requests {
            cudf_cxx::groupby::ffi::groupby_add_request(
                builder.pin_mut(),
                checked_i32(col_idx)?,
                agg.inner,
            );
        }

        let raw = cudf_cxx::groupby::ffi::groupby_execute(builder.pin_mut(), &values.inner)
            .map_err(CudfError::from_cxx)?;

        Ok(Table { inner: raw })
    }

    /// Execute the groupby and return only the key columns result.
    ///
    /// Useful when you need just the unique group keys.
    pub fn execute_keys(self, values: &Table) -> Result<Table> {
        if self.requests.is_empty() {
            return Err(CudfError::InvalidArgument(
                "groupby requires at least one aggregation request".to_string(),
            ));
        }

        let mut builder = cudf_cxx::groupby::ffi::groupby_new(&self.keys.inner);

        for (col_idx, agg) in self.requests {
            cudf_cxx::groupby::ffi::groupby_add_request(
                builder.pin_mut(),
                checked_i32(col_idx)?,
                agg.inner,
            );
        }

        let raw = cudf_cxx::groupby::ffi::groupby_execute_keys(builder.pin_mut(), &values.inner)
            .map_err(CudfError::from_cxx)?;

        Ok(Table { inner: raw })
    }

    /// Execute the groupby and return only the aggregation result columns.
    pub fn execute_values(self, values: &Table) -> Result<Table> {
        if self.requests.is_empty() {
            return Err(CudfError::InvalidArgument(
                "groupby requires at least one aggregation request".to_string(),
            ));
        }

        let mut builder = cudf_cxx::groupby::ffi::groupby_new(&self.keys.inner);

        for (col_idx, agg) in self.requests {
            cudf_cxx::groupby::ffi::groupby_add_request(
                builder.pin_mut(),
                checked_i32(col_idx)?,
                agg.inner,
            );
        }

        let raw = cudf_cxx::groupby::ffi::groupby_execute_values(builder.pin_mut(), &values.inner)
            .map_err(CudfError::from_cxx)?;

        Ok(Table { inner: raw })
    }
}

// ── GroupBy Scan ──────────────────────────────────────────────

/// Aggregation type for groupby scan operations.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GroupByScanOp {
    /// Cumulative sum within group.
    Sum = 0,
    /// Cumulative min within group.
    Min = 2,
    /// Cumulative max within group.
    Max = 3,
    /// Cumulative count within group.
    Count = 11,
    /// Rank within group.
    Rank = 12,
}

/// A builder for groupby scan operations.
///
/// Groupby scan computes cumulative (prefix) aggregations within each group.
///
/// # Examples
///
/// ```rust,no_run
/// use cudf::{Column, Table};
/// use cudf::groupby::{GroupByScan, GroupByScanOp};
///
/// let keys = Table::new(vec![
///     Column::from_slice(&[1i32, 1, 2, 2]).unwrap(),
/// ]).unwrap();
/// let values = Table::new(vec![
///     Column::from_slice(&[10.0f64, 20.0, 30.0, 40.0]).unwrap(),
/// ]).unwrap();
///
/// let result = GroupByScan::new(&keys)
///     .scan(0, GroupByScanOp::Sum)
///     .execute(&values)
///     .unwrap();
/// ```
pub struct GroupByScan<'a> {
    keys: &'a Table,
    requests: Vec<(usize, GroupByScanOp)>,
}

impl<'a> GroupByScan<'a> {
    /// Create a new groupby scan builder with the given key columns.
    pub fn new(keys: &'a Table) -> Self {
        Self {
            keys,
            requests: Vec::new(),
        }
    }

    /// Add a scan request for a value column.
    pub fn scan(mut self, column: usize, op: GroupByScanOp) -> Self {
        self.requests.push((column, op));
        self
    }

    /// Execute the grouped scan.
    ///
    /// Returns a table with key columns followed by scan result columns.
    pub fn execute(self, values: &Table) -> Result<Table> {
        if self.requests.is_empty() {
            return Err(CudfError::InvalidArgument(
                "groupby scan requires at least one request".to_string(),
            ));
        }

        let mut builder = cudf_cxx::groupby::ffi::groupby_scan_new(&self.keys.inner);

        for (col_idx, op) in self.requests {
            cudf_cxx::groupby::ffi::groupby_scan_add_request(
                builder.pin_mut(),
                checked_i32(col_idx)?,
                op as i32,
            );
        }

        let raw = cudf_cxx::groupby::ffi::groupby_scan_execute(builder.pin_mut(), &values.inner)
            .map_err(CudfError::from_cxx)?;

        Ok(Table { inner: raw })
    }
}

// ── GroupBy Get Groups ────────────────────────────────────────

/// Result of a groupby get_groups operation.
pub struct GroupByGroups {
    /// The grouped key rows.
    pub keys: Table,
    /// Group boundary offsets (length = num_groups + 1).
    pub offsets: Column,
    /// The grouped values (if values were provided).
    pub values: Option<Table>,
}

/// Policy for replacing nulls within groups.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GroupByReplacePolicy {
    /// Fill with the preceding non-null value.
    Forward = 0,
    /// Fill with the following non-null value.
    Backward = 1,
}

impl Table {
    /// Get the group structure for this table used as keys.
    ///
    /// Returns the grouped keys and group boundary offsets.
    pub fn groupby_get_groups(&self) -> Result<GroupByGroups> {
        let mut raw =
            cudf_cxx::groupby::ffi::groupby_get_groups(&self.inner).map_err(CudfError::from_cxx)?;
        let keys_raw = cudf_cxx::groupby::ffi::groupby_groups_take_keys(raw.pin_mut())
            .map_err(CudfError::from_cxx)?;
        let offsets_raw = cudf_cxx::groupby::ffi::groupby_groups_take_offsets(raw.pin_mut())
            .map_err(CudfError::from_cxx)?;
        Ok(GroupByGroups {
            keys: Table { inner: keys_raw },
            offsets: Column { inner: offsets_raw },
            values: None,
        })
    }

    /// Get the group structure, also reordering `values` by group.
    ///
    /// Returns grouped keys, offsets, and the values table reordered
    /// so rows in the same group are contiguous.
    pub fn groupby_get_groups_with_values(&self, values: &Table) -> Result<GroupByGroups> {
        let mut raw =
            cudf_cxx::groupby::ffi::groupby_get_groups_with_values(&self.inner, &values.inner)
                .map_err(CudfError::from_cxx)?;
        let keys_raw = cudf_cxx::groupby::ffi::groupby_groups_take_keys(raw.pin_mut())
            .map_err(CudfError::from_cxx)?;
        let offsets_raw = cudf_cxx::groupby::ffi::groupby_groups_take_offsets(raw.pin_mut())
            .map_err(CudfError::from_cxx)?;
        let values_raw = cudf_cxx::groupby::ffi::groupby_groups_take_values(raw.pin_mut())
            .map_err(CudfError::from_cxx)?;
        Ok(GroupByGroups {
            keys: Table { inner: keys_raw },
            offsets: Column { inner: offsets_raw },
            values: Some(Table { inner: values_raw }),
        })
    }

    /// Replace nulls within groups using forward or backward fill.
    ///
    /// `policies` must have one entry per column in `values`.
    ///
    /// Returns a table with key columns followed by null-replaced value columns.
    pub fn groupby_replace_nulls(
        &self,
        values: &Table,
        policies: &[GroupByReplacePolicy],
    ) -> Result<Table> {
        let p: Vec<i32> = policies.iter().map(|&p| p as i32).collect();
        let raw = cudf_cxx::groupby::ffi::groupby_replace_nulls(&self.inner, &values.inner, &p)
            .map_err(CudfError::from_cxx)?;
        Ok(Table { inner: raw })
    }
}