etop_core/datasets/
cryo_dataset.rs

1use crate::{DataSpec, DataWarehouse, EtopError, InputDataset};
2use etop_format::ColumnFormatShorthand;
3use polars::prelude::*;
4use std::collections::HashMap;
5
6/// cryo dataset
7#[derive(Clone)]
8pub struct CryoDataset {
9    /// name
10    pub name: String,
11}
12
13impl DataSpec for CryoDataset {
14    fn name(&self) -> String {
15        self.name.to_string()
16    }
17
18    fn row_noun(&self) -> String {
19        self.name.to_string()
20    }
21
22    fn inputs(&self) -> Vec<InputDataset> {
23        vec![InputDataset::Raw(self.name.to_string())]
24    }
25
26    fn transform(
27        &self,
28        warehouse: &DataWarehouse,
29        start_block: Option<u32>,
30        end_block: Option<u32>,
31    ) -> Result<DataFrame, EtopError> {
32        let mut df = warehouse.get_dataset(self.name.as_str())?;
33        df = crate::filter_by_block_number(df, start_block, end_block)?;
34
35        for (name, dtype) in df.schema().iter() {
36            if !dtype.is_integer() & !dtype.is_float() {
37                match dtype {
38                    DataType::Utf8 => {}
39                    _ => {
40                        df = df
41                            .lazy()
42                            .drop_columns(vec![name])
43                            // .with_column(col(name).cast(DataType::Utf8))
44                            .collect()
45                            .map_err(EtopError::PolarsError)?;
46                    }
47                }
48            }
49        }
50
51        Ok(df)
52    }
53
54    fn default_columns(&self) -> Option<Vec<String>> {
55        None
56    }
57
58    fn default_column_formats(&self) -> Option<HashMap<String, ColumnFormatShorthand>> {
59        None
60    }
61}