Skip to main content

molrs/block/
dtype.rs

1//! Data type enumeration and trait for Block columns.
2
3use ndarray::ArrayD;
4
5use super::column::Column;
6use crate::types::{F, I, U};
7
8/// Supported data types for Block columns.
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
10pub enum DType {
11    /// Floating point using the compile-time scalar type [`F`].
12    Float,
13    /// Signed integer using the compile-time scalar type [`I`].
14    Int,
15    /// Boolean
16    Bool,
17    /// Unsigned integer using the compile-time scalar type [`U`].
18    UInt,
19    /// 8-bit unsigned integer
20    U8,
21    /// String
22    String,
23}
24
25impl DType {
26    /// Returns the name of the data type as a string.
27    pub fn name(&self) -> &'static str {
28        match self {
29            DType::Float => "float",
30            DType::Int => "int",
31            DType::Bool => "bool",
32            DType::UInt => "uint",
33            DType::U8 => "u8",
34            DType::String => "string",
35        }
36    }
37}
38
39impl std::fmt::Display for DType {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        write!(f, "{}", self.name())
42    }
43}
44
45/// Trait for types that can be stored in a Block column.
46///
47/// This trait provides the mechanism for generic dispatch when inserting
48/// arrays into a Block. Users don't need to interact with this trait directly.
49pub trait BlockDtype: Sized + 'static {
50    /// Returns the DType for this type.
51    fn dtype() -> DType;
52
53    /// Converts an ArrayD of this type into a Column.
54    fn into_column(arr: ArrayD<Self>) -> Column;
55
56    /// Tries to extract a reference to an ArrayD of this type from a Column.
57    fn from_column(col: &Column) -> Option<&ArrayD<Self>>;
58
59    /// Tries to extract a mutable reference to an ArrayD of this type from a Column.
60    fn from_column_mut(col: &mut Column) -> Option<&mut ArrayD<Self>>;
61}
62
63impl BlockDtype for F {
64    fn dtype() -> DType {
65        DType::Float
66    }
67
68    fn into_column(arr: ArrayD<Self>) -> Column {
69        Column::Float(arr)
70    }
71
72    fn from_column(col: &Column) -> Option<&ArrayD<Self>> {
73        col.as_float()
74    }
75
76    fn from_column_mut(col: &mut Column) -> Option<&mut ArrayD<Self>> {
77        col.as_float_mut()
78    }
79}
80
81impl BlockDtype for I {
82    fn dtype() -> DType {
83        DType::Int
84    }
85
86    fn into_column(arr: ArrayD<Self>) -> Column {
87        Column::Int(arr)
88    }
89
90    fn from_column(col: &Column) -> Option<&ArrayD<Self>> {
91        col.as_int()
92    }
93
94    fn from_column_mut(col: &mut Column) -> Option<&mut ArrayD<Self>> {
95        col.as_int_mut()
96    }
97}
98
99impl BlockDtype for bool {
100    fn dtype() -> DType {
101        DType::Bool
102    }
103
104    fn into_column(arr: ArrayD<Self>) -> Column {
105        Column::Bool(arr)
106    }
107
108    fn from_column(col: &Column) -> Option<&ArrayD<Self>> {
109        col.as_bool()
110    }
111
112    fn from_column_mut(col: &mut Column) -> Option<&mut ArrayD<Self>> {
113        col.as_bool_mut()
114    }
115}
116
117impl BlockDtype for U {
118    fn dtype() -> DType {
119        DType::UInt
120    }
121
122    fn into_column(arr: ArrayD<Self>) -> Column {
123        Column::UInt(arr)
124    }
125
126    fn from_column(col: &Column) -> Option<&ArrayD<Self>> {
127        col.as_uint()
128    }
129
130    fn from_column_mut(col: &mut Column) -> Option<&mut ArrayD<Self>> {
131        col.as_uint_mut()
132    }
133}
134
135impl BlockDtype for u8 {
136    fn dtype() -> DType {
137        DType::U8
138    }
139
140    fn into_column(arr: ArrayD<Self>) -> Column {
141        Column::U8(arr)
142    }
143
144    fn from_column(col: &Column) -> Option<&ArrayD<Self>> {
145        col.as_u8()
146    }
147
148    fn from_column_mut(col: &mut Column) -> Option<&mut ArrayD<Self>> {
149        col.as_u8_mut()
150    }
151}
152
153impl BlockDtype for String {
154    fn dtype() -> DType {
155        DType::String
156    }
157
158    fn into_column(arr: ArrayD<Self>) -> Column {
159        Column::String(arr)
160    }
161
162    fn from_column(col: &Column) -> Option<&ArrayD<Self>> {
163        col.as_string()
164    }
165
166    fn from_column_mut(col: &mut Column) -> Option<&mut ArrayD<Self>> {
167        col.as_string_mut()
168    }
169}