wip, testing is_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));
+    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;
 }