scouter_client/data_utils/
numpy.rs

1use crate::data_utils::{ConvertedData, DataConverter, DataTypes};
2use crate::error::DataError;
3use pyo3::prelude::*;
4
5pub struct NumpyDataConverter;
6
7impl DataConverter for NumpyDataConverter {
8    fn categorize_features<'py>(
9        py: Python<'py>,
10        data: &Bound<'py, PyAny>,
11    ) -> Result<DataTypes, DataError> {
12        let numpy = py.import("numpy")?.getattr("ndarray")?;
13
14        if !data.is_instance(&numpy)? {
15            return Err(DataError::NotNumpyArrayError);
16        }
17
18        let mut string_features = Vec::new();
19        let mut float_features = Vec::new();
20
21        let shape = data.getattr("shape")?.extract::<Vec<usize>>()?;
22        let dtypes = data.getattr("dtype")?;
23
24        if dtypes.getattr("kind")?.extract::<String>()? == "u" {
25            // create vec from shape[1]
26            string_features = (0..shape[1])
27                .map(|i| format!("feature_{i}"))
28                .collect::<Vec<String>>();
29        } else {
30            float_features = (0..shape[1])
31                .map(|i| format!("feature_{i}"))
32                .collect::<Vec<String>>();
33        }
34
35        Ok(DataTypes::new(
36            Vec::new(), // No integer features in numpy
37            float_features,
38            string_features,
39        ))
40    }
41
42    fn process_numeric_features<'py>(
43        data: &Bound<'py, PyAny>,
44        data_types: &DataTypes,
45    ) -> Result<(Option<Bound<'py, PyAny>>, Option<String>), DataError> {
46        if data_types.numeric_features.is_empty() {
47            return Ok((None, None));
48        }
49        let dtype = Some(data.getattr("dtype")?.str()?.to_string());
50
51        Ok((Some(data.clone()), dtype))
52    }
53
54    #[allow(clippy::needless_lifetimes)]
55    fn process_string_features<'py>(
56        data: &Bound<'py, PyAny>,
57        features: &[String],
58    ) -> Result<Option<Vec<Vec<String>>>, DataError> {
59        if features.is_empty() {
60            return Ok(None);
61        }
62
63        Ok(Some(
64            data.call_method1("astype", ("str",))?
65                .call_method0("to_list")?
66                .extract::<Vec<Vec<String>>>()?,
67        ))
68    }
69
70    fn prepare_data<'py>(
71        py: Python<'py>,
72        data: &Bound<'py, PyAny>,
73    ) -> Result<ConvertedData<'py>, DataError> {
74        let data_types = NumpyDataConverter::categorize_features(py, data)?;
75
76        let (numeric_array, dtype) =
77            NumpyDataConverter::process_numeric_features(data, &data_types)?;
78        let string_array =
79            NumpyDataConverter::process_string_features(data, &data_types.string_features)?;
80
81        Ok((
82            data_types.numeric_features,
83            numeric_array,
84            dtype,
85            data_types.string_features,
86            string_array,
87        ))
88    }
89}