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    /// Returns the human-readable column name.
33    pub fn name(&self) -> &str {
34        self.name.as_str()
35    }
36
37    /// Returns the xxhash32 id derived from the name.
38    pub fn id(&self) -> u32 {
39        self.key
40    }
41
42    /// Returns a clone of this key.
43    pub fn key(&self) -> crate::Key {
44        self.clone()
45    }
46}
47impl PartialEq for Key {
48    fn eq(&self, other: &Self) -> bool {
49        self.key == other.key //&& self.name() == other.name()
50    }
51}
52impl PartialOrd for Key {
53    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
54        Some(self.cmp(other))
55    }
56}
57impl Ord for Key {
58    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
59        self.key.cmp(&other.key)
60    }
61}
62
63#[cfg(feature = "python")]
64#[pymethods]
65impl Key {
66    #[new]
67    #[pyo3(signature = (name, ctype=None))]
68    pub fn init(name: String, ctype: Option<DataType>) -> Self {
69        Self::new(name.as_str(), ctype.unwrap_or(DataType::Unknown))
70    }
71
72    #[pyo3(name = "name")]
73    pub fn py_name(&self) -> &str {
74        self.name()
75    }
76
77    #[pyo3(name = "id")]
78    pub fn py_id(&self) -> u32 {
79        self.key
80    }
81    #[pyo3(name = "dtype")]
82    pub fn py_type(&self) -> DataType {
83        self.ctype
84    }
85}
86
87impl std::fmt::Display for Key {
88    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89        write!(f, "{}", self.name)
90    }
91}
92impl std::fmt::Debug for Key {
93    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
94        write!(f, "{self}")
95    }
96}
97
98impl From<&str> for Key {
99    fn from(name: &str) -> Self {
100        Self {
101            key: xxhash_rust::const_xxh32::xxh32(name.as_bytes(), 0),
102            name: name.into(),
103            ctype: DataType::Unknown,
104        }
105    }
106}
107impl From<SString> for Key {
108    fn from(name: SString) -> Self {
109        Self::from(name.as_str())
110    }
111}
112impl From<String> for Key {
113    fn from(name: String) -> Self {
114        Self::from(name.as_str())
115    }
116}
117
118impl From<&String> for Key {
119    fn from(name: &String) -> Self {
120        Self::from(name.as_str())
121    }
122}
123
124impl Key {
125    /// Creates a key with the given name and data type.
126    pub fn new(name: &str, ctype: DataType) -> Self {
127        Self {
128            key: xxhash_rust::const_xxh32::xxh32(name.as_bytes(), 0),
129            name: name.into(),
130            ctype,
131        }
132    }
133}
134
135impl std::hash::Hash for Key {
136    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
137        state.write_u32(self.key);
138    }
139}
140
141#[cfg(test)]
142mod test {
143    use super::*;
144    #[test]
145    fn dummy_test() {
146        let key = Key::from("test");
147        assert_eq!(key.key, 1042293711);
148        assert_eq!(key.id(), 1042293711);
149        assert_eq!(key.key(), key);
150        assert_eq!(key.name, "test");
151        let s = String::from("test");
152        let key1 = Key::from(s);
153        assert_eq!(key1.name, "test");
154        assert_eq!(key1, key);
155        assert!(key1.cmp(&key) == std::cmp::Ordering::Equal);
156        let s = SString::from("test");
157        let key = Key::from(s);
158        assert_eq!(key.key, 1042293711);
159        assert_eq!(key.name, "test");
160        let key = Key::new("test", DataType::Unknown);
161        assert_eq!(key.key, 1042293711);
162        assert_eq!(key.name, "test");
163        assert_eq!(key.ctype, DataType::Unknown);
164        assert_eq!(format!("{}", key), "test");
165        let s = format!("{:?}", key);
166        println!("{}", s);
167        let key = Key::from(&("test".to_string()));
168        assert_eq!(key.key, 1042293711);
169        assert_eq!(key.name, "test");
170        let key_serialized = serde_json::to_string(&key).expect("BUG: Cannot serialize");
171        let key_deserialized: Key =
172            serde_json::from_str(&key_serialized).expect("BUG: Cannot deserialize");
173        assert_eq!(key, key_deserialized);
174    }
175
176    #[cfg(feature = "python")]
177    #[test]
178    fn py_test() {
179        let key = Key::init("test".into(), Some(DataType::String));
180        assert_eq!(key.key, 1042293711);
181        assert_eq!(key.name, "test");
182        assert_eq!(key.ctype, DataType::String);
183        assert_eq!(key.py_name(), "test");
184        assert_eq!(key.py_id(), 1042293711);
185    }
186}
trs_dataframe/dataframe/key.rs

trs_dataframe/dataframe/
key.rs
