Skip to main content

trs_dataframe/dataframe/
key.rs

1//use flexstr::SharedStr as SharedString;
2#[cfg(feature = "python")]
3use pyo3::prelude::*;
4#[cfg(feature = "utoipa")]
5use utoipa::ToSchema;
6
7use smartstring::alias::String as SString;
8
9use crate::DataType;
10
11/// [`Key`] holds information about key. It's used to identify the feature
12/// Mostly is represented as a string but inside is stored as a hash and the [`DataType`]
13/// of the feature/value which is stored in [`crate::DataFrame`]
14#[derive(Clone, Eq, serde::Deserialize, serde::Serialize, Default)]
15#[cfg_attr(feature = "python", pyo3::pyclass)]
16#[cfg_attr(feature = "utoipa", derive(ToSchema))]
17pub struct Key {
18    pub key: u32,
19    #[cfg_attr(feature = "utoipa", schema(schema_with = smart_string_schema))]
20    pub name: SString,
21    pub ctype: DataType,
22}
23
24#[cfg(feature = "utoipa")]
25fn smart_string_schema() -> utoipa::openapi::Object {
26    utoipa::openapi::ObjectBuilder::new()
27        .schema_type(utoipa::openapi::schema::Type::String)
28        .build()
29}
30
31impl Key {
32    pub fn name(&self) -> &str {
33        self.name.as_str()
34    }
35
36    pub fn id(&self) -> u32 {
37        self.key
38    }
39    pub fn key(&self) -> crate::Key {
40        self.clone()
41    }
42}
43impl PartialEq for Key {
44    fn eq(&self, other: &Self) -> bool {
45        self.key == other.key //&& self.name() == other.name()
46    }
47}
48impl PartialOrd for Key {
49    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
50        Some(self.cmp(other))
51    }
52}
53impl Ord for Key {
54    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
55        self.key.cmp(&other.key)
56    }
57}
58
59#[cfg(feature = "python")]
60#[pymethods]
61impl Key {
62    #[new]
63    #[pyo3(signature = (name, ctype=None))]
64    pub fn init(name: String, ctype: Option<DataType>) -> Self {
65        Self::new(name.as_str(), ctype.unwrap_or(DataType::Unknown))
66    }
67
68    #[pyo3(name = "name")]
69    pub fn py_name(&self) -> &str {
70        self.name()
71    }
72
73    #[pyo3(name = "id")]
74    pub fn py_id(&self) -> u32 {
75        self.key
76    }
77    #[pyo3(name = "dtype")]
78    pub fn py_type(&self) -> DataType {
79        self.ctype
80    }
81}
82
83impl std::fmt::Display for Key {
84    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
85        write!(f, "{}", self.name)
86    }
87}
88impl std::fmt::Debug for Key {
89    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
90        write!(f, "{self}")
91    }
92}
93
94impl From<&str> for Key {
95    fn from(name: &str) -> Self {
96        Self {
97            key: xxhash_rust::const_xxh32::xxh32(name.as_bytes(), 0),
98            name: name.into(),
99            ctype: DataType::Unknown,
100        }
101    }
102}
103impl From<SString> for Key {
104    fn from(name: SString) -> Self {
105        Self::from(name.as_str())
106    }
107}
108impl From<String> for Key {
109    fn from(name: String) -> Self {
110        Self::from(name.as_str())
111    }
112}
113
114impl From<&String> for Key {
115    fn from(name: &String) -> Self {
116        Self::from(name.as_str())
117    }
118}
119
120impl Key {
121    pub fn new(name: &str, ctype: DataType) -> Self {
122        Self {
123            key: xxhash_rust::const_xxh32::xxh32(name.as_bytes(), 0),
124            name: name.into(),
125            ctype,
126        }
127    }
128}
129
130impl std::hash::Hash for Key {
131    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
132        state.write_u32(self.key);
133    }
134}
135
136#[cfg(test)]
137mod test {
138    use super::*;
139    #[test]
140    fn dummy_test() {
141        let key = Key::from("test");
142        assert_eq!(key.key, 1042293711);
143        assert_eq!(key.id(), 1042293711);
144        assert_eq!(key.key(), key);
145        assert_eq!(key.name, "test");
146        let s = String::from("test");
147        let key1 = Key::from(s);
148        assert_eq!(key1.name, "test");
149        assert_eq!(key1, key);
150        assert!(key1.cmp(&key) == std::cmp::Ordering::Equal);
151        let s = SString::from("test");
152        let key = Key::from(s);
153        assert_eq!(key.key, 1042293711);
154        assert_eq!(key.name, "test");
155        let key = Key::new("test", DataType::Unknown);
156        assert_eq!(key.key, 1042293711);
157        assert_eq!(key.name, "test");
158        assert_eq!(key.ctype, DataType::Unknown);
159        assert_eq!(format!("{}", key), "test");
160        let s = format!("{:?}", key);
161        println!("{}", s);
162        let key = Key::from(&("test".to_string()));
163        assert_eq!(key.key, 1042293711);
164        assert_eq!(key.name, "test");
165        let key_serialized = serde_json::to_string(&key).expect("BUG: Cannot serialize");
166        let key_deserialized: Key =
167            serde_json::from_str(&key_serialized).expect("BUG: Cannot deserialize");
168        assert_eq!(key, key_deserialized);
169    }
170
171    #[cfg(feature = "python")]
172    #[test]
173    fn py_test() {
174        let key = Key::init("test".into(), Some(DataType::String));
175        assert_eq!(key.key, 1042293711);
176        assert_eq!(key.name, "test");
177        assert_eq!(key.ctype, DataType::String);
178        assert_eq!(key.py_name(), "test");
179        assert_eq!(key.py_id(), 1042293711);
180    }
181}