tea_dtype/
datatype.rs

1#[cfg(feature = "time")]
2use tea_time::TimeUnitTrait;
3#[cfg(feature = "time")]
4use tea_time::{DateTime, TimeDelta, TimeUnit};
5
6/// Represents the various data types supported by the system.
7///
8/// This enum encapsulates both primitive and complex data types, including
9/// numeric types, strings, optional types, and time-related types (when the "time" feature is enabled).
10#[derive(PartialEq, Eq, Debug)]
11pub enum DataType {
12    /// Boolean type
13    Bool,
14    /// 32-bit floating point type
15    F32,
16    /// 64-bit floating point type
17    F64,
18    /// 32-bit signed integer type
19    I32,
20    /// 64-bit signed integer type
21    I64,
22    /// 8-bit unsigned integer type
23    U8,
24    /// 64-bit unsigned integer type
25    U64,
26    /// Platform-specific size type
27    Usize,
28    /// String slice type
29    Str,
30    /// Owned String type
31    String,
32    /// Generic Object type
33    Object,
34    /// Optional Boolean type
35    OptBool,
36    /// Optional 32-bit floating point type
37    OptF32,
38    /// Optional 64-bit floating point type
39    OptF64,
40    /// Optional 32-bit signed integer type
41    OptI32,
42    /// Optional 64-bit signed integer type
43    OptI64,
44    /// Optional platform-specific size type
45    OptUsize,
46    /// Vector of platform-specific size type
47    VecUsize,
48    /// DateTime type with specified TimeUnit (only available with "time" feature)
49    #[cfg(feature = "time")]
50    DateTime(TimeUnit),
51    /// TimeDelta type (only available with "time" feature)
52    #[cfg(feature = "time")]
53    TimeDelta,
54    /// Unknown type
55    Unknown,
56}
57
58/// A trait for types that can provide their corresponding DataType.
59pub trait GetDataType {
60    /// Returns the DataType corresponding to the implementing type.
61    ///
62    /// # Returns
63    ///
64    /// A `DataType` enum variant representing the type.
65    fn dtype() -> DataType
66    where
67        Self: Sized;
68}
69
70macro_rules! impl_datatype {
71    ($tyname:ident, $physical:ty) => {
72        impl GetDataType for $physical {
73            #[inline(always)]
74            fn dtype() -> DataType {
75                DataType::$tyname
76            }
77        }
78    };
79}
80
81impl DataType {
82    /// Checks if the DataType is a floating-point type.
83    ///
84    /// # Returns
85    ///
86    /// `true` if the DataType is F32, F64, OptF32, or OptF64, `false` otherwise.
87    #[inline(always)]
88    pub fn is_float(&self) -> bool {
89        use DataType::*;
90        matches!(self, F32 | F64 | OptF32 | OptF64)
91    }
92
93    /// Checks if the DataType is an integer type.
94    ///
95    /// # Returns
96    ///
97    /// `true` if the DataType is I32, I64, Usize, U64, OptUsize, OptI32, or OptI64, `false` otherwise.
98    #[inline(always)]
99    pub fn is_int(&self) -> bool {
100        use DataType::*;
101        matches!(self, I32 | I64 | Usize | U64 | OptUsize | OptI32 | OptI64)
102    }
103
104    /// Checks if the DataType is a time-related type.
105    ///
106    /// # Returns
107    ///
108    /// `true` if the DataType is DateTime (when the "time" feature is enabled), `false` otherwise.
109    #[inline]
110    pub fn is_time(&self) -> bool {
111        #[cfg(feature = "time")]
112        {
113            use DataType::*;
114            matches!(self, DateTime(_))
115        }
116        #[cfg(not(feature = "time"))]
117        {
118            false
119        }
120    }
121
122    /// Converts the DataType to its floating-point equivalent.
123    ///
124    /// # Returns
125    ///
126    /// A new DataType that represents the floating-point equivalent of the current type.
127    pub fn float(self) -> Self {
128        use DataType::*;
129        match self {
130            F32 => F32,
131            I32 => F32,
132            I64 => F64,
133            Usize => F64,
134            U64 => F64,
135            OptUsize => F64,
136            OptI32 => OptF32,
137            OptI64 => OptF64,
138            OptF32 => OptF32,
139            OptF64 => OptF64,
140            _ => F64,
141        }
142    }
143
144    /// Converts the DataType to its integer equivalent.
145    ///
146    /// # Returns
147    ///
148    /// A new DataType that represents the integer equivalent of the current type.
149    pub fn int(self) -> Self {
150        use DataType::*;
151        match self {
152            I32 => I32,
153            F32 => I32,
154            F64 => I64,
155            Usize => Usize,
156            U64 => I64,
157            OptUsize => OptUsize,
158            OptI32 => OptI32,
159            OptI64 => OptI64,
160            OptF32 => OptI32,
161            OptF64 => OptI64,
162            _ => I64,
163        }
164    }
165}
166
167impl_datatype!(Bool, bool);
168impl_datatype!(U8, u8);
169impl_datatype!(F32, f32);
170impl_datatype!(F64, f64);
171impl_datatype!(I32, i32);
172impl_datatype!(I64, i64);
173impl_datatype!(U64, u64);
174impl_datatype!(Usize, usize);
175impl_datatype!(String, String);
176impl_datatype!(OptBool, Option<bool>);
177impl_datatype!(OptF32, Option<f32>);
178impl_datatype!(OptF64, Option<f64>);
179impl_datatype!(OptI32, Option<i32>);
180impl_datatype!(OptI64, Option<i64>);
181impl_datatype!(OptUsize, Option<usize>);
182impl_datatype!(VecUsize, Vec<usize>);
183
184#[cfg(feature = "time")]
185impl<U: TimeUnitTrait> GetDataType for DateTime<U> {
186    #[inline(always)]
187    fn dtype() -> DataType {
188        DataType::DateTime(U::unit())
189    }
190}
191
192#[cfg(feature = "time")]
193impl_datatype!(TimeDelta, TimeDelta);
194
195impl GetDataType for &str {
196    #[inline(always)]
197    fn dtype() -> DataType {
198        DataType::Str
199    }
200}
201
202#[cfg(feature = "polars")]
203const fn into_pl_dtype(dt: &DataType) -> tea_deps::polars::prelude::DataType {
204    use tea_deps::polars::prelude::{DataType as PlDataType, UnknownKind};
205    match dt {
206        DataType::Bool => PlDataType::Boolean,
207        DataType::F32 => PlDataType::Float32,
208        DataType::F64 => PlDataType::Float64,
209        DataType::I32 => PlDataType::Int32,
210        DataType::I64 => PlDataType::Int64,
211        DataType::U8 => PlDataType::UInt8,
212        DataType::U64 => PlDataType::UInt64,
213        DataType::Usize => PlDataType::UInt64,
214        DataType::String => PlDataType::String,
215        _ => PlDataType::Unknown(UnknownKind::Any),
216    }
217}
218
219#[cfg(feature = "polars")]
220const fn from_pl_dtype(dt: &tea_deps::polars::prelude::DataType) -> DataType {
221    use tea_deps::polars::prelude::DataType as PlDataType;
222    match dt {
223        PlDataType::Boolean => DataType::Bool,
224        PlDataType::Float32 => DataType::F32,
225        PlDataType::Float64 => DataType::F64,
226        PlDataType::Int32 => DataType::I32,
227        PlDataType::Int64 => DataType::I64,
228        PlDataType::UInt8 => DataType::U8,
229        PlDataType::UInt64 => DataType::U64,
230        PlDataType::String => DataType::String,
231        _ => DataType::Unknown,
232    }
233}
234
235#[cfg(feature = "polars")]
236impl From<DataType> for tea_deps::polars::prelude::DataType {
237    #[inline(always)]
238    fn from(dt: DataType) -> Self {
239        into_pl_dtype(&dt)
240    }
241}
242
243#[cfg(feature = "polars")]
244impl From<&DataType> for tea_deps::polars::prelude::DataType {
245    #[inline(always)]
246    fn from(dt: &DataType) -> Self {
247        into_pl_dtype(dt)
248    }
249}
250
251#[cfg(feature = "polars")]
252impl From<tea_deps::polars::prelude::DataType> for DataType {
253    #[inline(always)]
254    fn from(dt: tea_deps::polars::prelude::DataType) -> Self {
255        from_pl_dtype(&dt)
256    }
257}
258
259#[cfg(feature = "polars")]
260impl From<&tea_deps::polars::prelude::DataType> for DataType {
261    #[inline(always)]
262    fn from(dt: &tea_deps::polars::prelude::DataType) -> Self {
263        from_pl_dtype(dt)
264    }
265}
266
267#[cfg(test)]
268mod tests {
269    use super::*;
270
271    #[test]
272    #[cfg(feature = "polars")]
273    fn test_from_datatype() {
274        use tea_deps::polars::prelude::DataType as PlDataType;
275        let dt: PlDataType = DataType::Bool.into();
276        assert_eq!(dt, PlDataType::Boolean);
277        let dt: PlDataType = DataType::F32.into();
278        assert_eq!(dt, PlDataType::Float32);
279        let dt: PlDataType = DataType::F64.into();
280        assert_eq!(dt, PlDataType::Float64);
281        let dt: PlDataType = DataType::I32.into();
282        assert_eq!(dt, PlDataType::Int32);
283        let dt: PlDataType = DataType::I64.into();
284        assert_eq!(dt, PlDataType::Int64);
285        let dt: PlDataType = DataType::U8.into();
286        assert_eq!(dt, PlDataType::UInt8);
287        let dt: PlDataType = DataType::U64.into();
288        assert_eq!(dt, PlDataType::UInt64);
289        let dt: PlDataType = DataType::Usize.into();
290        assert_eq!(dt, PlDataType::UInt64);
291        let dt: PlDataType = DataType::String.into();
292        assert_eq!(dt, PlDataType::String);
293        let dt: PlDataType = DataType::Object.into();
294        assert!(matches!(dt, PlDataType::Unknown(_)));
295    }
296    #[test]
297    #[cfg(feature = "polars")]
298    fn test_from_pldatatype() {
299        use tea_deps::polars::prelude::{DataType as PlDataType, UnknownKind};
300
301        assert_eq!(DataType::from(PlDataType::Boolean), DataType::Bool);
302        assert_eq!(DataType::from(PlDataType::Float32), DataType::F32);
303        assert_eq!(DataType::from(PlDataType::Float64), DataType::F64);
304        assert_eq!(DataType::from(PlDataType::Int32), DataType::I32);
305        assert_eq!(DataType::from(PlDataType::Int64), DataType::I64);
306        assert_eq!(DataType::from(PlDataType::UInt8), DataType::U8);
307        assert_eq!(DataType::from(PlDataType::UInt64), DataType::U64);
308        assert_eq!(DataType::from(PlDataType::String), DataType::String);
309
310        // For Unknown type, we expect it to be mapped to Unknown
311        if let DataType::Unknown = DataType::from(PlDataType::Unknown(UnknownKind::Any)) {
312            // Test passes
313        } else {
314            panic!("Expected Unknown DataType for PlDataType::Object");
315        }
316
317        // Test with reference
318        assert_eq!(DataType::from(&PlDataType::Boolean), DataType::Bool);
319    }
320}
tea_dtype/datatype.rs

tea_dtype/
datatype.rs
