1use 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
16pub trait ColumnAccess {
18 fn as_float_view(&self) -> Option<ArrayViewD<'_, F>>;
20 fn as_int_view(&self) -> Option<ArrayViewD<'_, I>>;
22 fn as_bool_view(&self) -> Option<ArrayViewD<'_, bool>>;
24 fn as_uint_view(&self) -> Option<ArrayViewD<'_, U>>;
26 fn as_u8_view(&self) -> Option<ArrayViewD<'_, u8>>;
28 fn as_string_view(&self) -> Option<ArrayViewD<'_, String>>;
30 fn nrows(&self) -> Option<usize>;
32 fn dtype(&self) -> DType;
34 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
114pub trait BlockAccess {
116 fn get_float_view(&self, key: &str) -> Option<ArrayViewD<'_, F>>;
118 fn get_int_view(&self, key: &str) -> Option<ArrayViewD<'_, I>>;
120 fn get_bool_view(&self, key: &str) -> Option<ArrayViewD<'_, bool>>;
122 fn get_uint_view(&self, key: &str) -> Option<ArrayViewD<'_, U>>;
124 fn get_u8_view(&self, key: &str) -> Option<ArrayViewD<'_, u8>>;
126 fn get_string_view(&self, key: &str) -> Option<ArrayViewD<'_, String>>;
128 fn nrows(&self) -> Option<usize>;
130 fn len(&self) -> usize;
132 fn is_empty(&self) -> bool;
134 fn contains_key(&self, key: &str) -> bool;
136 fn column_keys(&self) -> Vec<&str>;
138 fn column_dtype(&self, key: &str) -> Option<DType>;
140 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}