binius_math/
rows_batch.rs

1// Copyright 2025 Irreducible Inc.
2
3/// This struct represents a batch of rows, each row having the same length equal to `row_len`.
4pub struct RowsBatch<'a, T> {
5	rows: Vec<&'a [T]>,
6	row_len: usize,
7}
8
9impl<'a, T> RowsBatch<'a, T> {
10	/// Create a new `RowsBatch` from a vector of rows and the given row length
11	///
12	/// # Panics
13	/// In case if any of the rows has a length different from `row_len`.
14	#[inline]
15	pub fn new(rows: Vec<&'a [T]>, row_len: usize) -> Self {
16		for row in &rows {
17			assert_eq!(row.len(), row_len);
18		}
19
20		Self { rows, row_len }
21	}
22
23	/// Create a new `RowsBatch` from an iterator of rows and the given row length.
24	///
25	/// # Panics
26	/// In case if any of the rows has a length less than `row_len`.
27	#[inline]
28	pub fn new_from_iter(rows: impl IntoIterator<Item = &'a [T]>, row_len: usize) -> Self {
29		let rows = rows.into_iter().map(|x| &x[..row_len]).collect();
30		Self { rows, row_len }
31	}
32
33	#[inline(always)]
34	pub fn get_ref(&self) -> RowsBatchRef<'_, T> {
35		RowsBatchRef {
36			rows: self.rows.as_slice(),
37			row_len: self.row_len,
38		}
39	}
40}
41
42/// This struct is similar to `RowsBatch`, but it holds a reference to the slice of rows.
43/// Unfortunately due to liftime issues we can't have a single generic struct which is
44/// parameterized by a container type.
45pub struct RowsBatchRef<'a, T> {
46	rows: &'a [&'a [T]],
47	row_len: usize,
48}
49
50impl<'a, T> RowsBatchRef<'a, T> {
51	/// Create a new `RowsBatchRef` from a slice of rows and the given row length.
52	///
53	/// # Panics
54	/// In case if any of the rows has a length different from `n_cols`.
55	#[inline]
56	pub fn new(rows: &'a [&'a [T]], row_len: usize) -> Self {
57		for row in rows {
58			assert_eq!(row.len(), row_len);
59		}
60
61		Self { rows, row_len }
62	}
63
64	#[inline]
65	pub fn iter(&self) -> impl Iterator<Item = &'a [T]> + '_ {
66		self.rows.as_ref().iter().copied()
67	}
68
69	#[inline(always)]
70	pub fn rows(&self) -> &[&'a [T]] {
71		self.rows
72	}
73
74	#[inline(always)]
75	pub fn n_rows(&self) -> usize {
76		self.rows.as_ref().len()
77	}
78
79	#[inline(always)]
80	pub fn row_len(&self) -> usize {
81		self.row_len
82	}
83
84	#[inline(always)]
85	pub fn is_empty(&self) -> bool {
86		self.rows.as_ref().is_empty()
87	}
88
89	pub fn map(&self, indices: impl AsRef<[usize]>) -> RowsBatch<'a, T> {
90		let rows = indices.as_ref().iter().map(|&i| self.rows[i]).collect();
91
92		RowsBatch {
93			rows,
94			row_len: self.row_len,
95		}
96	}
97}
binius_math/rows_batch.rs

binius_math/
rows_batch.rs
