Skip to main content

molrs_core/block/
column_view.rs

1//! Zero-copy borrowed view of a [`Column`].
2//!
3//! `ColumnView<'a>` borrows the underlying ndarray data from a [`Column`] without
4//! copying, providing read-only access with the same API surface as `Column`.
5
6use ndarray::ArrayViewD;
7
8use super::column::Column;
9use super::dtype::DType;
10use crate::types::{F, I, U};
11
12/// A borrowed, read-only view of a [`Column`].
13///
14/// Each variant holds an `ArrayViewD` that borrows from the corresponding
15/// `ArrayD` inside an owned `Column`. No data is copied.
16pub enum ColumnView<'a> {
17    /// Borrowed float column.
18    Float(ArrayViewD<'a, F>),
19    /// Borrowed signed integer column.
20    Int(ArrayViewD<'a, I>),
21    /// Borrowed boolean column.
22    Bool(ArrayViewD<'a, bool>),
23    /// Borrowed unsigned integer column.
24    UInt(ArrayViewD<'a, U>),
25    /// Borrowed u8 column.
26    U8(ArrayViewD<'a, u8>),
27    /// Borrowed string column.
28    String(ArrayViewD<'a, String>),
29}
30
31impl<'a> ColumnView<'a> {
32    /// Returns the number of rows (axis-0 length) of this column view.
33    ///
34    /// Returns `None` if the array has rank 0.
35    pub fn nrows(&self) -> Option<usize> {
36        match self {
37            ColumnView::Float(a) => a.shape().first().copied(),
38            ColumnView::Int(a) => a.shape().first().copied(),
39            ColumnView::Bool(a) => a.shape().first().copied(),
40            ColumnView::UInt(a) => a.shape().first().copied(),
41            ColumnView::U8(a) => a.shape().first().copied(),
42            ColumnView::String(a) => a.shape().first().copied(),
43        }
44    }
45
46    /// Returns the data type of this column view.
47    pub fn dtype(&self) -> DType {
48        match self {
49            ColumnView::Float(_) => DType::Float,
50            ColumnView::Int(_) => DType::Int,
51            ColumnView::Bool(_) => DType::Bool,
52            ColumnView::UInt(_) => DType::UInt,
53            ColumnView::U8(_) => DType::U8,
54            ColumnView::String(_) => DType::String,
55        }
56    }
57
58    /// Returns the shape of the underlying array view.
59    pub fn shape(&self) -> &[usize] {
60        match self {
61            ColumnView::Float(a) => a.shape(),
62            ColumnView::Int(a) => a.shape(),
63            ColumnView::Bool(a) => a.shape(),
64            ColumnView::UInt(a) => a.shape(),
65            ColumnView::U8(a) => a.shape(),
66            ColumnView::String(a) => a.shape(),
67        }
68    }
69
70    /// Returns a view of the float data, or `None` if this column view is not `Float`.
71    pub fn as_float(&self) -> Option<ArrayViewD<'a, F>> {
72        match self {
73            ColumnView::Float(a) => Some(a.clone()),
74            _ => None,
75        }
76    }
77
78    /// Returns a view of the integer data, or `None` if not `Int`.
79    pub fn as_int(&self) -> Option<ArrayViewD<'a, I>> {
80        match self {
81            ColumnView::Int(a) => Some(a.clone()),
82            _ => None,
83        }
84    }
85
86    /// Returns a view of the boolean data, or `None` if not `Bool`.
87    pub fn as_bool(&self) -> Option<ArrayViewD<'a, bool>> {
88        match self {
89            ColumnView::Bool(a) => Some(a.clone()),
90            _ => None,
91        }
92    }
93
94    /// Returns a view of the unsigned integer data, or `None` if not `UInt`.
95    pub fn as_uint(&self) -> Option<ArrayViewD<'a, U>> {
96        match self {
97            ColumnView::UInt(a) => Some(a.clone()),
98            _ => None,
99        }
100    }
101
102    /// Returns a view of the u8 data, or `None` if not `U8`.
103    pub fn as_u8(&self) -> Option<ArrayViewD<'a, u8>> {
104        match self {
105            ColumnView::U8(a) => Some(a.clone()),
106            _ => None,
107        }
108    }
109
110    /// Returns a view of the string data, or `None` if not `String`.
111    pub fn as_string(&self) -> Option<ArrayViewD<'a, String>> {
112        match self {
113            ColumnView::String(a) => Some(a.clone()),
114            _ => None,
115        }
116    }
117
118    /// Creates an owned [`Column`] by cloning the viewed data.
119    pub fn to_owned(&self) -> Column {
120        match self {
121            ColumnView::Float(a) => Column::from_float(a.to_owned()),
122            ColumnView::Int(a) => Column::from_int(a.to_owned()),
123            ColumnView::Bool(a) => Column::from_bool(a.to_owned()),
124            ColumnView::UInt(a) => Column::from_uint(a.to_owned()),
125            ColumnView::U8(a) => Column::from_u8(a.to_owned()),
126            ColumnView::String(a) => Column::from_string(a.to_owned()),
127        }
128    }
129}
130
131impl<'a> From<&'a Column> for ColumnView<'a> {
132    fn from(col: &'a Column) -> Self {
133        match col {
134            Column::Float(a) => ColumnView::Float(a.view()),
135            Column::Int(a) => ColumnView::Int(a.view()),
136            Column::Bool(a) => ColumnView::Bool(a.view()),
137            Column::UInt(a) => ColumnView::UInt(a.view()),
138            Column::U8(a) => ColumnView::U8(a.view()),
139            Column::String(a) => ColumnView::String(a.view()),
140        }
141    }
142}
143
144impl std::fmt::Debug for ColumnView<'_> {
145    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
146        match self {
147            ColumnView::Float(a) => write!(f, "ColumnView::Float(shape={:?})", a.shape()),
148            ColumnView::Int(a) => write!(f, "ColumnView::Int(shape={:?})", a.shape()),
149            ColumnView::Bool(a) => write!(f, "ColumnView::Bool(shape={:?})", a.shape()),
150            ColumnView::UInt(a) => write!(f, "ColumnView::UInt(shape={:?})", a.shape()),
151            ColumnView::U8(a) => write!(f, "ColumnView::U8(shape={:?})", a.shape()),
152            ColumnView::String(a) => write!(f, "ColumnView::String(shape={:?})", a.shape()),
153        }
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160    use ndarray::Array1;
161
162    #[test]
163    fn test_from_column_float() {
164        let col = Column::from_float(Array1::from_vec(vec![1.0 as F, 2.0, 3.0]).into_dyn());
165        let view = ColumnView::from(&col);
166        assert_eq!(view.dtype(), DType::Float);
167        assert_eq!(view.nrows(), Some(3));
168        assert_eq!(view.shape(), &[3]);
169        assert!(view.as_float().is_some());
170        assert!(view.as_int().is_none());
171    }
172
173    #[test]
174    fn test_from_column_int() {
175        let col = Column::from_int(Array1::from_vec(vec![1 as I, 2, 3]).into_dyn());
176        let view = ColumnView::from(&col);
177        assert_eq!(view.dtype(), DType::Int);
178        assert!(view.as_int().is_some());
179        assert!(view.as_float().is_none());
180    }
181
182    #[test]
183    fn test_from_column_bool() {
184        let col = Column::from_bool(Array1::from_vec(vec![true, false]).into_dyn());
185        let view = ColumnView::from(&col);
186        assert_eq!(view.dtype(), DType::Bool);
187        assert!(view.as_bool().is_some());
188    }
189
190    #[test]
191    fn test_from_column_uint() {
192        let col = Column::from_uint(Array1::from_vec(vec![1 as U, 2]).into_dyn());
193        let view = ColumnView::from(&col);
194        assert_eq!(view.dtype(), DType::UInt);
195        assert!(view.as_uint().is_some());
196    }
197
198    #[test]
199    fn test_from_column_u8() {
200        let col = Column::from_u8(Array1::from_vec(vec![1u8, 2]).into_dyn());
201        let view = ColumnView::from(&col);
202        assert_eq!(view.dtype(), DType::U8);
203        assert!(view.as_u8().is_some());
204    }
205
206    #[test]
207    fn test_from_column_string() {
208        let col = Column::from_string(
209            Array1::from_vec(vec!["a".to_string(), "b".to_string()]).into_dyn(),
210        );
211        let view = ColumnView::from(&col);
212        assert_eq!(view.dtype(), DType::String);
213        assert!(view.as_string().is_some());
214    }
215
216    #[test]
217    fn test_to_owned_roundtrip() {
218        let col = Column::from_float(Array1::from_vec(vec![1.0 as F, 2.0, 3.0]).into_dyn());
219        let view = ColumnView::from(&col);
220        let owned = view.to_owned();
221        assert_eq!(owned.dtype(), DType::Float);
222        assert_eq!(owned.nrows(), Some(3));
223        assert_eq!(
224            owned.as_float().unwrap().as_slice_memory_order().unwrap(),
225            &[1.0, 2.0, 3.0]
226        );
227    }
228
229    #[test]
230    fn test_zero_copy() {
231        let col = Column::from_float(Array1::from_vec(vec![1.0 as F, 2.0, 3.0]).into_dyn());
232        let view = ColumnView::from(&col);
233        // The view's data pointer should match the original column's data pointer
234        let orig_ptr = col.as_float().unwrap().as_ptr();
235        let view_ptr = view.as_float().unwrap().as_ptr();
236        assert_eq!(orig_ptr, view_ptr);
237    }
238}