diff --git a/src/main.rs b/src/main.rs index 833979f..85017af 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,5 @@ use msyrs::download::jpmaqsdownload::{JPMaQSDownload, JPMaQSDownloadGetIndicatorArgs}; - +use msyrs::utils::dftools::is_quantamental_dataframe; fn main() { println!("Authentication to DataQuery API"); let mut jpamqs_download = JPMaQSDownload::default(); @@ -8,7 +8,7 @@ fn main() { jpamqs_download.check_connection().unwrap(); println!("Retrieving catalogue of tickers"); - let mut start = std::time::Instant::now(); + let start = std::time::Instant::now(); let tickers = jpamqs_download.get_catalogue().unwrap(); println!( "Retrieved catalogue with {} tickers in {:?}", @@ -25,40 +25,46 @@ fn main() { // let mut df_deets = Vec::new(); println!("Retrieving indicators for {} tickers", sel_tickers.len()); - start = std::time::Instant::now(); - let all_metrics: Vec = ["value", "grading", "eop_lag", "mop_lag"] - .iter() - .map(|x| x.to_string()) - .collect(); - let res = jpamqs_download.save_indicators_as_csv( - JPMaQSDownloadGetIndicatorArgs { + // start = std::time::Instant::now(); + // let all_metrics: Vec = ["value", "grading", "eop_lag", "mop_lag"] + // .iter() + // .map(|x| x.to_string()) + // .collect(); + // let res = jpamqs_download.save_indicators_as_csv( + // JPMaQSDownloadGetIndicatorArgs { + // tickers: sel_tickers.clone(), + // metrics: all_metrics, + // ..Default::default() + // }, + // "./data/", + // ); + + // match res { + // Ok(_) => println!( + // "Saved indicators for {} tickers in {:?}", + // sel_tickers.len(), + // start.elapsed() + // ), + // Err(e) => println!("Error saving indicators: {:?}", e), + // } + + let res_df = jpamqs_download + .get_indicators_qdf(JPMaQSDownloadGetIndicatorArgs { tickers: sel_tickers.clone(), - metrics: all_metrics, + start_date: "2024-11-05".to_string(), ..Default::default() - }, - "./data/", + }) + .unwrap(); + + println!( + "Retrieved indicators for {} tickers in {:?}", + sel_tickers.len(), + start.elapsed() ); - match res { - Ok(_) => println!( - "Saved indicators for {} tickers in {:?}", - sel_tickers.len(), - start.elapsed() - ), - Err(e) => println!("Error saving indicators: {:?}", e), + if !is_quantamental_dataframe(&res_df) { + println!("DataFrame is not a quantamental DataFrame"); + } else { + println!("DataFrame is a quantamental DataFrame"); } - - // let mut res_df: DataFrame = jpamqs_download - // .get_indicators_qdf(JPMaQSDownloadGetIndicatorArgs { - // tickers: sel_tickers.clone(), - // start_date: "2024-11-05".to_string(), - // ..Default::default() - // }) - // .unwrap(); - - // println!( - // "Retrieved indicators for {} tickers in {:?}", - // sel_tickers.len(), - // start.elapsed() - // ); } diff --git a/src/utils/dftools.rs b/src/utils/dftools.rs index 95355f4..6d38743 100644 --- a/src/utils/dftools.rs +++ b/src/utils/dftools.rs @@ -3,16 +3,49 @@ use polars::prelude::*; const QDF_IDX_COLUMNS: [&str; 3] = ["real_date", "cid", "xcat"]; /// Check if a DataFrame is a quantamental DataFrame. -fn is_quantamental_dataframe(df: &DataFrame) -> bool { +pub fn is_quantamental_dataframe(df: &DataFrame) -> bool { let columns = df .get_column_names() .iter() .map(|s| s.as_str()) .collect::>(); let has_idx_columns = QDF_IDX_COLUMNS.iter().all(|col| columns.contains(col)); + if !has_idx_columns { + return false; + } + let real_date_col = df.select(["real_date"]); + match real_date_col { + Ok(_) => {} + Err(_) => return false, + }; - // check if real_date is date-type - let real_date = df.column("real_date").unwrap(); + let is_date_dtype = real_date_col + .unwrap() + .dtypes() + .iter() + .all(|dtype| dtype == &DataType::Date); - has_idx_columns + if !is_date_dtype { + return false; + } + + let cid_col = df.select(["cid"]); + match cid_col { + Ok(_) => {} + Err(_) => return false, + }; + + let xcat_col = df.select(["xcat"]); + match xcat_col { + Ok(_) => {} + Err(_) => return false, + }; + + // has atleast 1 more column + let has_other_columns = columns.len() > 3; + if !has_other_columns { + return false; + } + + return true; }