astrors_fork/io/hdus/table/
buffer.rs

1use polars::{
2    prelude::{NamedFrom, PlSmallStr},
3    series::Series,
4};
5use rayon::prelude::*;
6
7#[derive(Debug, PartialEq)]
8pub enum ColumnDataBuffer {
9    I(Vec<i32>),
10    E(Vec<f32>),
11    D(Vec<f64>),
12    A(Vec<String>),
13    F(Vec<f32>),
14}
15
16impl ColumnDataBuffer {
17    pub fn new(tform: &str, size: i32) -> Self {
18        let tform = tform.trim();
19        let tform_type = tform.chars().next().unwrap();
20        match tform_type {
21            'I' => ColumnDataBuffer::I(vec![0; size as usize]),
22            'E' => ColumnDataBuffer::E(vec![0.0; size as usize]),
23            'D' => ColumnDataBuffer::D(vec![0.0; size as usize]),
24            'A' => ColumnDataBuffer::A(vec![String::new(); size as usize]),
25            'F' => ColumnDataBuffer::F(vec![0.0; size as usize]),
26            _ => panic!("Wrong data type"),
27        }
28    }
29
30    pub fn max_len(&self) -> usize {
31        match self {
32            ColumnDataBuffer::I(data) => data
33                .par_iter()
34                .map(|x| x.to_string().len())
35                .max()
36                .unwrap_or(0),
37            ColumnDataBuffer::E(data) => data
38                .par_iter()
39                .map(|x| x.to_string().len())
40                .max()
41                .unwrap_or(0),
42            ColumnDataBuffer::D(data) => data
43                .par_iter()
44                .map(|x| x.to_string().len())
45                .max()
46                .unwrap_or(0),
47            ColumnDataBuffer::A(data) => data.par_iter().map(|x| x.len()).max().unwrap_or(0),
48            ColumnDataBuffer::F(data) => data
49                .par_iter()
50                .map(|x| x.to_string().len())
51                .max()
52                .unwrap_or(0),
53        }
54    }
55
56    pub fn byte_value(&self) -> usize {
57        match self {
58            ColumnDataBuffer::I(_data) => 4,
59            ColumnDataBuffer::E(_data) => 4,
60            ColumnDataBuffer::D(_data) => 8,
61            ColumnDataBuffer::A(_data) => 1,
62            ColumnDataBuffer::F(_data) => 4,
63        }
64    }
65
66    pub fn byte_value_from_str(data_type: &str) -> usize {
67        match data_type {
68            "I" => 4,
69            "E" => 4,
70            "D" => 8,
71            "A" => 1,
72            "F" => 4,
73            _ => panic!("Wrong data type"),
74        }
75    }
76
77    pub fn to_series(&self, col_name: &str) -> Series {
78        let col_name: PlSmallStr = col_name.into();
79        let series = match self {
80            ColumnDataBuffer::I(data) => Series::new(col_name, data),
81            ColumnDataBuffer::E(data) => Series::new(col_name, data),
82            ColumnDataBuffer::D(data) => Series::new(col_name, data),
83            ColumnDataBuffer::A(data) => Series::new(col_name, data),
84            ColumnDataBuffer::F(data) => Series::new(col_name, data),
85        };
86        series
87    }
88
89    pub fn clear(&mut self) {
90        match self {
91            ColumnDataBuffer::I(data) => data.clear(),
92            ColumnDataBuffer::E(data) => data.clear(),
93            ColumnDataBuffer::D(data) => data.clear(),
94            ColumnDataBuffer::A(data) => data.clear(),
95            ColumnDataBuffer::F(data) => data.clear(),
96        }
97    }
98
99    pub fn write_on_idx(&mut self, bytes: &[u8], data_type: char, idx: i64) {
100        let string = String::from_utf8_lossy(&bytes)
101            .trim_end()
102            .trim_start()
103            .to_string();
104        match data_type {
105            'I' => {
106                // parse bytes to i32
107                match self {
108                    ColumnDataBuffer::I(data) => {
109                        data[idx as usize] = string.parse::<i32>().unwrap()
110                    }
111                    _ => panic!("Wrong data type"),
112                }
113            }
114            'E' => {
115                // parse bytes to f32
116                match self {
117                    ColumnDataBuffer::E(data) => {
118                        data[idx as usize] = string.parse::<f32>().unwrap()
119                    }
120                    _ => panic!("Wrong data type"),
121                }
122            }
123            'D' => {
124                // parse bytes to f64
125                match self {
126                    ColumnDataBuffer::D(data) => {
127                        data[idx as usize] = string.parse::<f64>().unwrap()
128                    }
129                    _ => panic!("Wrong data type"),
130                }
131            }
132            'A' => {
133                // parse bytes to String
134                match self {
135                    ColumnDataBuffer::A(data) => data[idx as usize] = string,
136                    _ => panic!("Wrong data type"),
137                }
138            }
139            'F' => {
140                // parse bytes to f32
141                match self {
142                    ColumnDataBuffer::F(data) => {
143                        data[idx as usize] = string.parse::<f32>().unwrap()
144                    }
145                    _ => panic!("Wrong data type"),
146                }
147            }
148            _ => panic!("Wrong data type"),
149        }
150    }
151    //no need for max_len on bintable
152}
astrors_fork/io/hdus/table/buffer.rs

astrors_fork/io/hdus/table/
buffer.rs
