Skip to main content

molrs_core/block/
access.rs

1//! Unified access traits for owned and borrowed column/block types.
2//!
3//! [`ColumnAccess`] and [`BlockAccess`] provide a common read-only interface
4//! that is implemented by both the owned types ([`Column`], [`Block`]) and
5//! their zero-copy view counterparts ([`ColumnView`], [`BlockView`]).
6
7use ndarray::ArrayViewD;
8
9use super::Block;
10use super::block_view::BlockView;
11use super::column::Column;
12use super::column_view::ColumnView;
13use super::dtype::DType;
14use crate::types::{F, I, U};
15
16/// Unified read-only access for [`Column`] and [`ColumnView`].
17pub trait ColumnAccess {
18    /// Returns a float array view, or `None` if the column is not `Float`.
19    fn as_float_view(&self) -> Option<ArrayViewD<'_, F>>;
20    /// Returns an int array view, or `None` if the column is not `Int`.
21    fn as_int_view(&self) -> Option<ArrayViewD<'_, I>>;
22    /// Returns a bool array view, or `None` if the column is not `Bool`.
23    fn as_bool_view(&self) -> Option<ArrayViewD<'_, bool>>;
24    /// Returns a uint array view, or `None` if the column is not `UInt`.
25    fn as_uint_view(&self) -> Option<ArrayViewD<'_, U>>;
26    /// Returns a u8 array view, or `None` if the column is not `U8`.
27    fn as_u8_view(&self) -> Option<ArrayViewD<'_, u8>>;
28    /// Returns a string array view, or `None` if the column is not `String`.
29    fn as_string_view(&self) -> Option<ArrayViewD<'_, String>>;
30    /// Returns the number of rows (axis-0 length), or `None` if rank 0.
31    fn nrows(&self) -> Option<usize>;
32    /// Returns the data type of this column.
33    fn dtype(&self) -> DType;
34    /// Returns the shape of the underlying array as an owned `Vec`.
35    fn shape(&self) -> Vec<usize>;
36}
37
38impl ColumnAccess for Column {
39    fn as_float_view(&self) -> Option<ArrayViewD<'_, F>> {
40        self.as_float().map(|a| a.view())
41    }
42
43    fn as_int_view(&self) -> Option<ArrayViewD<'_, I>> {
44        self.as_int().map(|a| a.view())
45    }
46
47    fn as_bool_view(&self) -> Option<ArrayViewD<'_, bool>> {
48        self.as_bool().map(|a| a.view())
49    }
50
51    fn as_uint_view(&self) -> Option<ArrayViewD<'_, U>> {
52        self.as_uint().map(|a| a.view())
53    }
54
55    fn as_u8_view(&self) -> Option<ArrayViewD<'_, u8>> {
56        self.as_u8().map(|a| a.view())
57    }
58
59    fn as_string_view(&self) -> Option<ArrayViewD<'_, String>> {
60        self.as_string().map(|a| a.view())
61    }
62
63    fn nrows(&self) -> Option<usize> {
64        Column::nrows(self)
65    }
66
67    fn dtype(&self) -> DType {
68        Column::dtype(self)
69    }
70
71    fn shape(&self) -> Vec<usize> {
72        Column::shape(self).to_vec()
73    }
74}
75
76impl ColumnAccess for ColumnView<'_> {
77    fn as_float_view(&self) -> Option<ArrayViewD<'_, F>> {
78        self.as_float()
79    }
80
81    fn as_int_view(&self) -> Option<ArrayViewD<'_, I>> {
82        self.as_int()
83    }
84
85    fn as_bool_view(&self) -> Option<ArrayViewD<'_, bool>> {
86        self.as_bool()
87    }
88
89    fn as_uint_view(&self) -> Option<ArrayViewD<'_, U>> {
90        self.as_uint()
91    }
92
93    fn as_u8_view(&self) -> Option<ArrayViewD<'_, u8>> {
94        self.as_u8()
95    }
96
97    fn as_string_view(&self) -> Option<ArrayViewD<'_, String>> {
98        self.as_string()
99    }
100
101    fn nrows(&self) -> Option<usize> {
102        ColumnView::nrows(self)
103    }
104
105    fn dtype(&self) -> DType {
106        ColumnView::dtype(self)
107    }
108
109    fn shape(&self) -> Vec<usize> {
110        ColumnView::shape(self).to_vec()
111    }
112}
113
114/// Unified read-only access for [`Block`] and [`BlockView`].
115pub trait BlockAccess {
116    /// Gets a float array view for `key` if present and of correct type.
117    fn get_float_view(&self, key: &str) -> Option<ArrayViewD<'_, F>>;
118    /// Gets an int array view for `key` if present and of correct type.
119    fn get_int_view(&self, key: &str) -> Option<ArrayViewD<'_, I>>;
120    /// Gets a bool array view for `key` if present and of correct type.
121    fn get_bool_view(&self, key: &str) -> Option<ArrayViewD<'_, bool>>;
122    /// Gets a uint array view for `key` if present and of correct type.
123    fn get_uint_view(&self, key: &str) -> Option<ArrayViewD<'_, U>>;
124    /// Gets a u8 array view for `key` if present and of correct type.
125    fn get_u8_view(&self, key: &str) -> Option<ArrayViewD<'_, u8>>;
126    /// Gets a string array view for `key` if present and of correct type.
127    fn get_string_view(&self, key: &str) -> Option<ArrayViewD<'_, String>>;
128    /// Returns the common axis-0 length, or `None` if empty.
129    fn nrows(&self) -> Option<usize>;
130    /// Number of columns.
131    fn len(&self) -> usize;
132    /// Returns `true` if there are no columns.
133    fn is_empty(&self) -> bool;
134    /// Returns `true` if the block contains the specified key.
135    fn contains_key(&self, key: &str) -> bool;
136    /// Returns column keys as a `Vec`.
137    fn column_keys(&self) -> Vec<&str>;
138    /// Returns the data type of the column with the given key, if it exists.
139    fn column_dtype(&self, key: &str) -> Option<DType>;
140    /// Returns the shape of the column with the given key, if it exists.
141    fn column_shape(&self, key: &str) -> Option<Vec<usize>>;
142}
143
144impl BlockAccess for Block {
145    fn get_float_view(&self, key: &str) -> Option<ArrayViewD<'_, F>> {
146        self.get_float(key).map(|a| a.view())
147    }
148
149    fn get_int_view(&self, key: &str) -> Option<ArrayViewD<'_, I>> {
150        self.get_int(key).map(|a| a.view())
151    }
152
153    fn get_bool_view(&self, key: &str) -> Option<ArrayViewD<'_, bool>> {
154        self.get_bool(key).map(|a| a.view())
155    }
156
157    fn get_uint_view(&self, key: &str) -> Option<ArrayViewD<'_, U>> {
158        self.get_uint(key).map(|a| a.view())
159    }
160
161    fn get_u8_view(&self, key: &str) -> Option<ArrayViewD<'_, u8>> {
162        self.get_u8(key).map(|a| a.view())
163    }
164
165    fn get_string_view(&self, key: &str) -> Option<ArrayViewD<'_, String>> {
166        self.get_string(key).map(|a| a.view())
167    }
168
169    fn nrows(&self) -> Option<usize> {
170        Block::nrows(self)
171    }
172
173    fn len(&self) -> usize {
174        Block::len(self)
175    }
176
177    fn is_empty(&self) -> bool {
178        Block::is_empty(self)
179    }
180
181    fn contains_key(&self, key: &str) -> bool {
182        Block::contains_key(self, key)
183    }
184
185    fn column_keys(&self) -> Vec<&str> {
186        self.keys().collect()
187    }
188
189    fn column_dtype(&self, key: &str) -> Option<DType> {
190        self.get(key).map(|col| col.dtype())
191    }
192
193    fn column_shape(&self, key: &str) -> Option<Vec<usize>> {
194        self.get(key).map(|col| col.shape().to_vec())
195    }
196}
197
198impl BlockAccess for BlockView<'_> {
199    fn get_float_view(&self, key: &str) -> Option<ArrayViewD<'_, F>> {
200        self.get_float(key)
201    }
202
203    fn get_int_view(&self, key: &str) -> Option<ArrayViewD<'_, I>> {
204        self.get_int(key)
205    }
206
207    fn get_bool_view(&self, key: &str) -> Option<ArrayViewD<'_, bool>> {
208        self.get_bool(key)
209    }
210
211    fn get_uint_view(&self, key: &str) -> Option<ArrayViewD<'_, U>> {
212        self.get_uint(key)
213    }
214
215    fn get_u8_view(&self, key: &str) -> Option<ArrayViewD<'_, u8>> {
216        self.get_u8(key)
217    }
218
219    fn get_string_view(&self, key: &str) -> Option<ArrayViewD<'_, String>> {
220        self.get_string(key)
221    }
222
223    fn nrows(&self) -> Option<usize> {
224        BlockView::nrows(self)
225    }
226
227    fn len(&self) -> usize {
228        BlockView::len(self)
229    }
230
231    fn is_empty(&self) -> bool {
232        BlockView::is_empty(self)
233    }
234
235    fn contains_key(&self, key: &str) -> bool {
236        BlockView::contains_key(self, key)
237    }
238
239    fn column_keys(&self) -> Vec<&str> {
240        self.keys().copied().collect()
241    }
242
243    fn column_dtype(&self, key: &str) -> Option<DType> {
244        BlockView::dtype(self, key)
245    }
246
247    fn column_shape(&self, key: &str) -> Option<Vec<usize>> {
248        self.get(key).map(|col_view| col_view.shape().to_vec())
249    }
250}
251
252#[cfg(test)]
253mod tests {
254    use super::*;
255    use ndarray::Array1;
256
257    fn make_block() -> Block {
258        let mut block = Block::new();
259        block
260            .insert("x", Array1::from_vec(vec![1.0 as F, 2.0, 3.0]).into_dyn())
261            .unwrap();
262        block
263            .insert("id", Array1::from_vec(vec![10 as I, 20, 30]).into_dyn())
264            .unwrap();
265        block
266    }
267
268    #[test]
269    fn test_column_access_on_column() {
270        let col = Column::from_float(Array1::from_vec(vec![1.0 as F, 2.0]).into_dyn());
271        assert!(ColumnAccess::as_float_view(&col).is_some());
272        assert!(ColumnAccess::as_int_view(&col).is_none());
273        assert_eq!(ColumnAccess::nrows(&col), Some(2));
274        assert_eq!(ColumnAccess::dtype(&col), DType::Float);
275        assert_eq!(ColumnAccess::shape(&col), vec![2]);
276    }
277
278    #[test]
279    fn test_column_access_on_column_view() {
280        let col = Column::from_int(Array1::from_vec(vec![1 as I, 2, 3]).into_dyn());
281        let view = ColumnView::from(&col);
282        assert!(ColumnAccess::as_int_view(&view).is_some());
283        assert!(ColumnAccess::as_float_view(&view).is_none());
284        assert_eq!(ColumnAccess::nrows(&view), Some(3));
285        assert_eq!(ColumnAccess::dtype(&view), DType::Int);
286    }
287
288    #[test]
289    fn test_block_access_on_block() {
290        let block = make_block();
291        assert!(BlockAccess::get_float_view(&block, "x").is_some());
292        assert!(BlockAccess::get_int_view(&block, "id").is_some());
293        assert_eq!(BlockAccess::nrows(&block), Some(3));
294        assert_eq!(BlockAccess::len(&block), 2);
295        assert!(!BlockAccess::is_empty(&block));
296        assert!(BlockAccess::contains_key(&block, "x"));
297        assert!(!BlockAccess::contains_key(&block, "missing"));
298    }
299
300    #[test]
301    fn test_block_access_on_block_view() {
302        let block = make_block();
303        let view = BlockView::from(&block);
304        assert!(BlockAccess::get_float_view(&view, "x").is_some());
305        assert!(BlockAccess::get_int_view(&view, "id").is_some());
306        assert_eq!(BlockAccess::nrows(&view), Some(3));
307        assert_eq!(BlockAccess::len(&view), 2);
308        assert!(!BlockAccess::is_empty(&view));
309        assert!(BlockAccess::contains_key(&view, "x"));
310    }
311
312    #[test]
313    fn test_generic_function_with_block_access() {
314        fn count_float_columns(b: &impl BlockAccess) -> usize {
315            b.column_keys()
316                .iter()
317                .filter(|k| b.get_float_view(k).is_some())
318                .count()
319        }
320
321        let block = make_block();
322        assert_eq!(count_float_columns(&block), 1);
323
324        let view = BlockView::from(&block);
325        assert_eq!(count_float_columns(&view), 1);
326    }
327}