etop_core/datasets/
cryo_dataset.rs1use crate::{DataSpec, DataWarehouse, EtopError, InputDataset};
2use etop_format::ColumnFormatShorthand;
3use polars::prelude::*;
4use std::collections::HashMap;
5
6#[derive(Clone)]
8pub struct CryoDataset {
9 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 .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}