wip, testing is_quantamental_dataframe

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<String> = ["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 mut res_df: DataFrame = jpamqs_download
-    //     .get_indicators_qdf(JPMaQSDownloadGetIndicatorArgs {
+    // start = std::time::Instant::now();
+    // let all_metrics: Vec<String> = ["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(),
-    //         start_date: "2024-11-05".to_string(),
+    //         metrics: all_metrics,
     //         ..Default::default()
-    //     })
-    //     .unwrap();
+    //     },
+    //     "./data/",
+    // );
 
-    // println!(
-    //     "Retrieved indicators for {} tickers in {:?}",
+    // 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(),
+            start_date: "2024-11-05".to_string(),
+            ..Default::default()
+        })
+        .unwrap();
+
+    println!(
+        "Retrieved indicators for {} tickers in {:?}",
+        sel_tickers.len(),
+        start.elapsed()
+    );
+
+    if !is_quantamental_dataframe(&res_df) {
+        println!("DataFrame is not a quantamental DataFrame");
+    } else {
+        println!("DataFrame is a quantamental DataFrame");
+    }
 }

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::<Vec<&str>>();
     let has_idx_columns = QDF_IDX_COLUMNS.iter().all(|col| columns.contains(col));
-
-    // check if real_date is date-type
-    let real_date = df.column("real_date").unwrap();
-
-    has_idx_columns
+    if !has_idx_columns {
+        return false;
+    }
+    let real_date_col = df.select(["real_date"]);
+    match real_date_col {
+        Ok(_) => {}
+        Err(_) => return false,
+    };
+
+    let is_date_dtype = real_date_col
+        .unwrap()
+        .dtypes()
+        .iter()
+        .all(|dtype| dtype == &DataType::Date);
+
+    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;
 }