astrors_fork/io/hdus/bintable/
buffer.rs

1use polars::{
2    prelude::{NamedFrom, PlSmallStr},
3    series::Series,
4};
5
6use crate::io::hdus::bintable::*;
7
8#[derive(Debug, PartialEq)]
9pub enum ColumnDataBuffer {
10    L(Vec<bool>),     // Logical
11    X(Vec<u8>),       // Bit
12    B(Vec<i8>),       // Byte
13    I(Vec<i16>),      // Short
14    J(Vec<i32>),      // Int
15    K(Vec<i64>),      // Long
16    A(Vec<String>),   // Char
17    E(Vec<f32>),      // Float
18    D(Vec<f64>),      // Double
19    C(Vec<String>),   // Complex
20    M(Vec<String>),   // Double complex
21    P(Vec<Vec<i32>>), // Array descriptor
22    Q(Vec<Vec<i64>>), // Array descriptor
23}
24
25#[derive(Debug, PartialEq)]
26pub enum ColumnArrayBuffer {
27    L(Vec<Vec<bool>>), // Logical
28    X(Vec<Vec<u8>>),   // Bit
29    B(Vec<Vec<i8>>),   // Byte
30    I(Vec<Vec<i16>>),  // Short
31    J(Vec<Vec<i32>>),  // Int
32    K(Vec<Vec<i64>>),  // Long
33    E(Vec<Vec<f32>>),  // Float
34    D(Vec<Vec<f64>>),  // Double
35}
36
37enum BufferTypes {
38    Scalar(ColumnDataBuffer),
39    Vector(ColumnArrayBuffer),
40}
41
42pub struct Buffer {
43    tform: String,
44    size: i32,
45    buffer: BufferTypes,
46    sub_size: i32,
47    data_letter: String,
48}
49
50impl Buffer {
51    pub fn new(tform: &str, size: i32) -> Self {
52        let tform = tform.to_string();
53        let data_letter = get_first_letter(&tform).to_string();
54
55        let mut sub_size = 1;
56        let mut vec_column = false;
57        if (get_data_bytes_size(&tform) != byte_value_from_str(&tform))
58            & (data_letter != "A")
59            & (data_letter != "C")
60            & (data_letter != "M")
61            & (data_letter != "P")
62            & (data_letter != "Q")
63        {
64            vec_column = true;
65            sub_size = (get_data_bytes_size(&tform) / byte_value_from_str(&tform)) as i32;
66        }
67
68        let buffer: BufferTypes;
69        if vec_column {
70            buffer = BufferTypes::Vector(ColumnArrayBuffer::new(&tform, size, sub_size));
71        } else {
72            buffer = BufferTypes::Scalar(ColumnDataBuffer::new(&tform, size));
73        }
74
75        Buffer {
76            tform,
77            size,
78            buffer,
79            sub_size,
80            data_letter,
81        }
82    }
83
84    pub fn to_series(&self, col_name: &str) -> Series {
85        match &self.buffer {
86            BufferTypes::Scalar(data) => data.to_series(col_name),
87            BufferTypes::Vector(data) => data.to_series(col_name),
88        }
89    }
90
91    pub fn clear(&mut self) {
92        match &mut self.buffer {
93            BufferTypes::Scalar(data) => data.clear(),
94            BufferTypes::Vector(data) => data.clear(),
95        }
96    }
97    pub fn write_on_idx(&mut self, bytes: &[u8], idx: i64) {
98        match &mut self.buffer {
99            BufferTypes::Scalar(data) => data.write_on_idx(bytes, &self.data_letter, idx),
100            BufferTypes::Vector(data) => {
101                data.write_on_idx(bytes, &self.data_letter, idx, self.sub_size)
102            }
103        }
104    }
105
106    pub fn read_var_len_cols(&mut self) {
107        //TODO
108
109        // match &mut self.buffer {
110        //     BufferTypes::Scalar(data) => {
111        //         println!("buffer: {:?}", data);
112        //     }
113        //     BufferTypes::Vector(data) => {}
114        // }
115    }
116}
117
118impl ColumnArrayBuffer {
119    pub fn new(tform: &str, size: i32, sub_size: i32) -> Self {
120        let tform = tform.trim();
121        let tform_type = get_first_letter(tform);
122
123        match tform_type {
124            "L" => ColumnArrayBuffer::L(vec![vec![false; sub_size as usize]; size as usize]),
125            "X" => ColumnArrayBuffer::X(vec![vec![0; sub_size as usize]; size as usize]),
126            "B" => ColumnArrayBuffer::B(vec![vec![0; sub_size as usize]; size as usize]),
127            "I" => ColumnArrayBuffer::I(vec![vec![0; sub_size as usize]; size as usize]),
128            "J" => ColumnArrayBuffer::J(vec![vec![0; sub_size as usize]; size as usize]),
129            "K" => ColumnArrayBuffer::K(vec![vec![0; sub_size as usize]; size as usize]),
130            "E" => ColumnArrayBuffer::E(vec![vec![0.0; sub_size as usize]; size as usize]),
131            "D" => ColumnArrayBuffer::D(vec![vec![0.0; sub_size as usize]; size as usize]),
132            _ => panic!("Unsupported data type for array col"),
133        }
134    }
135
136    pub fn empty(tform: &str) -> Self {
137        let tform = tform.trim();
138        let tform_type = get_first_letter(tform);
139
140        match tform_type {
141            "L" => ColumnArrayBuffer::L(vec![vec![]]),
142            "X" => ColumnArrayBuffer::X(vec![vec![]]),
143            "B" => ColumnArrayBuffer::B(vec![vec![]]),
144            "I" => ColumnArrayBuffer::I(vec![vec![]]),
145            "J" => ColumnArrayBuffer::J(vec![vec![]]),
146            "K" => ColumnArrayBuffer::K(vec![vec![]]),
147            "E" => ColumnArrayBuffer::E(vec![vec![]]),
148            "D" => ColumnArrayBuffer::D(vec![vec![]]),
149            _ => panic!("Unsupported data type for array col"),
150        }
151    }
152
153    pub fn to_series(&self, col_name: &str) -> Series {
154        let col: PlSmallStr = col_name.into();
155        let blank: PlSmallStr = "".into();
156        let series = match self {
157            ColumnArrayBuffer::L(data) => Series::new(
158                col,
159                data.iter()
160                    .map(|vec| Series::new(blank.clone(), vec))
161                    .collect::<Vec<Series>>(),
162            ),
163            ColumnArrayBuffer::X(data) => Series::new(
164                col,
165                data.iter()
166                    .map(|vec| Series::new(blank.clone(), vec))
167                    .collect::<Vec<Series>>(),
168            ),
169            ColumnArrayBuffer::B(data) => Series::new(
170                col,
171                data.iter()
172                    .map(|vec| Series::new(blank.clone(), vec))
173                    .collect::<Vec<Series>>(),
174            ),
175            ColumnArrayBuffer::I(data) => Series::new(
176                col,
177                data.iter()
178                    .map(|vec| Series::new(blank.clone(), vec))
179                    .collect::<Vec<Series>>(),
180            ),
181            ColumnArrayBuffer::J(data) => Series::new(
182                col,
183                data.iter()
184                    .map(|vec| Series::new(blank.clone(), vec))
185                    .collect::<Vec<Series>>(),
186            ),
187            ColumnArrayBuffer::K(data) => Series::new(
188                col,
189                data.iter()
190                    .map(|vec| Series::new(blank.clone(), vec))
191                    .collect::<Vec<Series>>(),
192            ),
193            ColumnArrayBuffer::E(data) => Series::new(
194                col,
195                data.iter()
196                    .map(|vec| Series::new(blank.clone(), vec))
197                    .collect::<Vec<Series>>(),
198            ),
199            ColumnArrayBuffer::D(data) => Series::new(
200                col,
201                data.iter()
202                    .map(|vec| Series::new(blank.clone(), vec))
203                    .collect::<Vec<Series>>(),
204            ),
205        };
206        series
207    }
208
209    pub fn clear(&mut self) {
210        match self {
211            ColumnArrayBuffer::L(data) => data.clear(),
212            ColumnArrayBuffer::X(data) => data.clear(),
213            ColumnArrayBuffer::B(data) => data.clear(),
214            ColumnArrayBuffer::I(data) => data.clear(),
215            ColumnArrayBuffer::J(data) => data.clear(),
216            ColumnArrayBuffer::K(data) => data.clear(),
217            ColumnArrayBuffer::E(data) => data.clear(),
218            ColumnArrayBuffer::D(data) => data.clear(),
219        }
220    }
221
222    pub fn write_on_idx(&mut self, bytes: &[u8], data_letter: &str, idx: i64, sub_size: i32) {
223        match data_letter {
224            "L" => {
225                // parse bytes to bool
226                match self {
227                    ColumnArrayBuffer::L(data) => {
228                        (0..sub_size).for_each(|i| {
229                            data[idx as usize][i as usize] = bytes[i as usize] != 0;
230                        });
231                    }
232                    _ => panic!("Wrong data type"),
233                }
234            }
235            "K" => {
236                // parse bytes to i64
237                match self {
238                    ColumnArrayBuffer::K(data) => {
239                        let mut stbyte = 0;
240                        (0..sub_size).for_each(|i| {
241                            data[idx as usize][i as usize] = i64::from_be_bytes([
242                                bytes[stbyte],
243                                bytes[stbyte + 1],
244                                bytes[stbyte + 2],
245                                bytes[stbyte + 3],
246                                bytes[stbyte + 4],
247                                bytes[stbyte + 5],
248                                bytes[stbyte + 6],
249                                bytes[stbyte + 7],
250                            ]);
251                            stbyte += 8;
252                        });
253                    }
254                    _ => panic!("Wrong data type"),
255                }
256            }
257            _ => panic!("Wrong data type"),
258        }
259    }
260}
261
262impl ColumnDataBuffer {
263    pub fn new(tform: &str, size: i32) -> Self {
264        let tform = tform.trim();
265        let tform_type = get_first_letter(tform);
266
267        match tform_type {
268            "L" => ColumnDataBuffer::L(vec![false; size as usize]),
269            "X" => ColumnDataBuffer::X(vec![0; size as usize]),
270            "B" => ColumnDataBuffer::B(vec![0; size as usize]),
271            "I" => ColumnDataBuffer::I(vec![0; size as usize]),
272            "J" => ColumnDataBuffer::J(vec![0; size as usize]),
273            "K" => ColumnDataBuffer::K(vec![0; size as usize]),
274            "A" => ColumnDataBuffer::A(vec![String::new(); size as usize]),
275            "E" => ColumnDataBuffer::E(vec![0.0; size as usize]),
276            "D" => ColumnDataBuffer::D(vec![0.0; size as usize]),
277            "C" => ColumnDataBuffer::C(vec![String::new(); size as usize]),
278            "M" => ColumnDataBuffer::M(vec![String::new(); size as usize]),
279            "P" => ColumnDataBuffer::P(vec![vec![0; 2]; size as usize]),
280            "Q" => ColumnDataBuffer::Q(vec![vec![0; 2]; size as usize]),
281            _ => ColumnDataBuffer::A(vec![String::new(); size as usize]),
282        }
283    }
284
285    pub fn to_series(&self, col_name: &str) -> Series {
286        let col: PlSmallStr = col_name.into();
287        let blank: PlSmallStr = "".into();
288        let series = match self {
289            ColumnDataBuffer::L(data) => Series::new(col, data),
290            ColumnDataBuffer::X(data) => Series::new(col, data),
291            ColumnDataBuffer::B(data) => Series::new(col, data),
292            ColumnDataBuffer::I(data) => Series::new(col, data),
293            ColumnDataBuffer::J(data) => Series::new(col, data),
294            ColumnDataBuffer::K(data) => Series::new(col, data),
295            ColumnDataBuffer::A(data) => Series::new(col, data),
296            ColumnDataBuffer::E(data) => Series::new(col, data),
297            ColumnDataBuffer::D(data) => Series::new(col, data),
298            ColumnDataBuffer::C(data) => Series::new(col, data),
299            ColumnDataBuffer::M(data) => Series::new(col, data),
300            ColumnDataBuffer::P(data) => {
301                let series_vec: Vec<Series> = data
302                    .into_iter()
303                    .map(|vec| Series::new(blank.clone(), &vec))
304                    .collect();
305                Series::new(col, series_vec)
306            }
307            ColumnDataBuffer::Q(data) => {
308                let series_vec: Vec<Series> = data
309                    .into_iter()
310                    .map(|vec| Series::new(blank.clone(), &vec))
311                    .collect();
312                Series::new(col, series_vec)
313            }
314        };
315        series
316    }
317
318    pub fn clear(&mut self) {
319        match self {
320            ColumnDataBuffer::L(data) => data.clear(),
321            ColumnDataBuffer::X(data) => data.clear(),
322            ColumnDataBuffer::B(data) => data.clear(),
323            ColumnDataBuffer::I(data) => data.clear(),
324            ColumnDataBuffer::J(data) => data.clear(),
325            ColumnDataBuffer::K(data) => data.clear(),
326            ColumnDataBuffer::A(data) => data.clear(),
327            ColumnDataBuffer::E(data) => data.clear(),
328            ColumnDataBuffer::D(data) => data.clear(),
329            ColumnDataBuffer::C(data) => data.clear(),
330            ColumnDataBuffer::M(data) => data.clear(),
331            ColumnDataBuffer::P(data) => data.clear(),
332            ColumnDataBuffer::Q(data) => data.clear(),
333        }
334    }
335
336    pub fn write_on_idx(&mut self, bytes: &[u8], data_letter: &str, idx: i64) {
337        match data_letter {
338            "L" => {
339                // parse bytes to bool
340                match self {
341                    ColumnDataBuffer::L(data) => data[idx as usize] = bytes[0] != 0,
342                    _ => panic!("Wrong data type"),
343                }
344            }
345            "X" => {
346                // parse bytes to u8
347                match self {
348                    ColumnDataBuffer::X(data) => data[idx as usize] = bytes[0],
349                    _ => panic!("Wrong data type"),
350                }
351            }
352            "B" => {
353                // parse bytes to i8
354                match self {
355                    ColumnDataBuffer::B(data) => data[idx as usize] = bytes[0] as i8,
356                    _ => panic!("Wrong data type"),
357                }
358            }
359            "I" => {
360                // parse bytes to i16
361                match self {
362                    ColumnDataBuffer::I(data) => {
363                        data[idx as usize] = i16::from_be_bytes([bytes[0], bytes[1]])
364                    }
365                    _ => panic!("Wrong data type"),
366                }
367            }
368            "J" => {
369                // parse bytes to i32
370                match self {
371                    ColumnDataBuffer::J(data) => {
372                        data[idx as usize] =
373                            i32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
374                    }
375                    _ => panic!("Wrong data type"),
376                }
377            }
378            "K" => {
379                // parse bytes to i64
380                match self {
381                    ColumnDataBuffer::K(data) => {
382                        data[idx as usize] = i64::from_be_bytes([
383                            bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6],
384                            bytes[7],
385                        ])
386                    }
387                    _ => panic!("Wrong data type"),
388                }
389            }
390            "A" => {
391                // parse bytes to String
392                match self {
393                    ColumnDataBuffer::A(data) => {
394                        let string = unsafe { String::from_utf8_unchecked(bytes.to_vec()) }
395                            .trim_end()
396                            .to_string();
397                        data[idx as usize] = string;
398                    }
399                    _ => panic!("Wrong data type"),
400                }
401            }
402            "E" => {
403                // parse bytes to f32
404                match self {
405                    ColumnDataBuffer::E(data) => {
406                        data[idx as usize] =
407                            f32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])
408                    }
409                    _ => panic!("Wrong data type"),
410                }
411            }
412            "D" => {
413                // parse bytes to f64
414                match self {
415                    ColumnDataBuffer::D(data) => {
416                        data[idx as usize] = f64::from_be_bytes([
417                            bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6],
418                            bytes[7],
419                        ])
420                    }
421                    _ => panic!("Wrong data type"),
422                }
423            }
424            "C" => {
425                // parse bytes to String
426                match self {
427                    ColumnDataBuffer::C(data) => {
428                        let string = unsafe { String::from_utf8_unchecked(bytes.to_vec()) }
429                            .trim_end()
430                            .to_string();
431                        data[idx as usize] = string;
432                    }
433                    _ => panic!("Wrong data type"),
434                }
435            }
436            "M" => {
437                // parse bytes to String
438                match self {
439                    ColumnDataBuffer::M(data) => {
440                        let string = unsafe { String::from_utf8_unchecked(bytes.to_vec()) }
441                            .trim_end()
442                            .to_string();
443                        data[idx as usize] = string;
444                    }
445                    _ => panic!("Wrong data type"),
446                }
447            }
448            "P" => {
449                // parse bytes to String
450                match self {
451                    ColumnDataBuffer::P(data) => {
452                        data[idx as usize][0] =
453                            i32::from_be_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]);
454                        data[idx as usize][1] =
455                            i32::from_be_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]);
456                    }
457                    _ => panic!("Wrong data type"),
458                }
459            }
460            "Q" => {
461                // parse bytes to String
462                match self {
463                    ColumnDataBuffer::Q(data) => {
464                        data[idx as usize][0] = i64::from_be_bytes([
465                            bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6],
466                            bytes[7],
467                        ]);
468                        data[idx as usize][0] = i64::from_be_bytes([
469                            bytes[8], bytes[9], bytes[10], bytes[11], bytes[12], bytes[13],
470                            bytes[14], bytes[15],
471                        ]);
472                    }
473                    _ => panic!("Wrong data type"),
474                }
475            }
476            _ => panic!("Wrong data type"),
477        }
478    }
479    //no need for max_len on bintable
480}
astrors_fork/io/hdus/bintable/buffer.rs

astrors_fork/io/hdus/bintable/
buffer.rs
