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

astrors_fork/io/hdus/bintable/
buffer.rs
