trs_dataframe/
lib.rs

1pub mod candidate;
2pub mod dataframe;
3pub mod error;
4pub use candidate::CandidateData;
5pub use data_value;
6pub use data_value::DataValue;
7pub use dataframe::DataFrame;
8pub mod filter;
9pub mod utils;
10pub use dataframe::join::{JoinBy, JoinById, JoinRelation};
11pub use dataframe::{
12    column_store::{ColumnFrame, KeyIndex},
13    key::Key,
14};
15pub type MLChefMap = halfbrown::HashMap<smartstring::alias::String, Vec<DataValue>>;
16pub use ndarray;
17
18#[cfg(feature = "jmalloc")]
19#[global_allocator]
20static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
21
22#[cfg(feature = "polars-df")]
23pub use polars;
24/// Data type for the values in the dataframe
25#[derive(
26    Debug, Clone, Copy, serde::Serialize, serde::Deserialize, PartialEq, Eq, Hash, Default,
27)]
28#[cfg_attr(feature = "python", pyo3::pyclass(eq, eq_int))]
29#[cfg_attr(feature = "utoipa", derive(utoipa::ToSchema))]
30pub enum DataType {
31    Bool,
32    U32,
33    I32,
34    U8,
35    U64,
36    I64,
37    F32,
38    F64,
39    String,
40    Bytes,
41    #[default]
42    Unknown,
43    Vec,
44    Map,
45}
46
47/// Autodetector for the data type from [`DataValue`]
48pub fn detect_dtype(value: &DataValue) -> DataType {
49    use DataValue::*;
50    match value {
51        Bool(_) => DataType::Bool,
52        I32(_) => DataType::I32,
53        U32(_) => DataType::U32,
54        I64(_) => DataType::I64,
55        U64(_) => DataType::U64,
56        F32(_) => DataType::F32,
57        F64(_) => DataType::F64,
58        String(_) => DataType::String,
59        Bytes(_) => DataType::Bytes,
60        Vec(_) => DataType::Vec,
61        Map(_) => DataType::Map,
62        _ => DataType::Unknown,
63    }
64}
65
66#[cfg(feature = "python")]
67use pyo3::prelude::*;
68
69#[cfg(feature = "python")]
70///
71/// ```
72/// use pyo3::prelude::*;
73///
74///  fn main() {
75///     let result = pyo3::Python::with_gil(|py| -> PyResult<()> {
76///         let module = PyModule::new(py, "trs_dataframe")?;
77///         let _m = trs_dataframe::trs_dataframe(py, module)?;
78///         Ok(())
79///         });
80///     assert!(result.is_ok(), "{:?}", result);
81///  }
82///
83/// ```
84#[pymodule]
85pub fn trs_dataframe(_py: pyo3::Python<'_>, m: pyo3::Bound<'_, PyModule>) -> pyo3::PyResult<()> {
86    m.add_class::<DataFrame>()?;
87    m.add_class::<JoinRelation>()?;
88    m.add_class::<Key>()?;
89    Ok(())
90}
91#[cfg(test)]
92mod test {
93    use super::*;
94    use rstest::*;
95
96    #[rstest]
97    #[case(DataType::Bool, DataValue::Bool(true))]
98    #[case(DataType::I32, DataValue::I32(1))]
99    #[case(DataType::U32, DataValue::U32(1))]
100    #[case(DataType::I64, DataValue::I64(1))]
101    #[case(DataType::U64, DataValue::U64(1))]
102    #[case(DataType::F32, DataValue::F32(1.0))]
103    #[case(DataType::F64, DataValue::F64(1.0))]
104    #[case(DataType::String, DataValue::String("1".into()))]
105    #[case(DataType::Bytes, DataValue::Bytes(b"1".to_vec()))]
106    #[case(DataType::Vec, DataValue::Vec(vec![DataValue::I32(1)]))]
107    #[case(DataType::Map, DataValue::Map(std::collections::HashMap::new()))]
108    #[case(DataType::Unknown, DataValue::Null)]
109    fn detection_test(#[case] dtype: DataType, #[case] value: DataValue) {
110        assert_eq!(detect_dtype(&value), dtype);
111        let serde_dtype: DataType =
112            serde_json::from_str(&serde_json::to_string(&dtype).expect("BUG: cannot serialize"))
113                .expect("BUG: cannot deserialize");
114        assert_eq!(serde_dtype, dtype);
115    }
116}
trs_dataframe/lib.rs

trs_dataframe/
lib.rs
