feat: update Cargo.toml and enhance pivot functionality in QDF

This commit is contained in:
Palash Tyagi 2024-11-21 01:29:46 +00:00
parent 99a33237d6
commit 3ffab4c95d
2 changed files with 21 additions and 8 deletions

View File

@ -9,8 +9,8 @@ rustdoc = ["--cfg", "doc"]
[lib] [lib]
name = "msyrs" name = "msyrs"
# path = "src/lib.rs" path = "src/lib.rs"
crate-type = ["cdylib"] crate-type = ["cdylib", "lib"]
[dependencies] [dependencies]

View File

@ -10,6 +10,15 @@ use std::f32::consts::E;
/// The required columns for a Quantamental DataFrame. /// The required columns for a Quantamental DataFrame.
const QDF_INDEX_COLUMNS: [&str; 3] = ["real_date", "cid", "xcat"]; const QDF_INDEX_COLUMNS: [&str; 3] = ["real_date", "cid", "xcat"];
#[allow(dead_code)]
fn get_ticker_from_qdf(df: &DataFrame) -> Result<String, Box<dyn Error>> {
check_quantamental_dataframe(df)?;
let cid = df.column("cid")?.get(0).unwrap().to_string();
let xcat = df.column("xcat")?.get(0).unwrap().to_string();
let ticker = format!("{}_{}", cid, xcat);
Ok(ticker)
}
/// Pivots a dataframe to a format where each ticker a column. /// Pivots a dataframe to a format where each ticker a column.
#[allow(dead_code)] #[allow(dead_code)]
pub fn pivot_dataframe_by_ticker( pub fn pivot_dataframe_by_ticker(
@ -27,7 +36,8 @@ pub fn pivot_dataframe_by_ticker(
// set metric to the first non-index column // set metric to the first non-index column
metric = df.get_column_names()[3].to_string(); metric = df.get_column_names()[3].to_string();
} }
let unique_dates: Column = get_unique_dates(&df)?; let mut unique_dates: Column = get_unique_dates(&df)?;
unique_dates.rename("real_date".into());
// keep only the index columns and the metric column // keep only the index columns and the metric column
let mut keep_cols = QDF_INDEX_COLUMNS let mut keep_cols = QDF_INDEX_COLUMNS
@ -38,18 +48,21 @@ pub fn pivot_dataframe_by_ticker(
keep_cols.push(metric.clone()); keep_cols.push(metric.clone());
let mut out_dfs = split_df_by_tickers(&df, Some(vec![metric.clone()]))?; let mut out_dfs = split_df_by_tickers(&df, Some(vec![metric.clone()]))?;
for (_, odf) in out_dfs.iter_mut() { for (tkr, odf) in out_dfs.iter_mut() {
// select keep_cols odf.rename(&metric, tkr.into())?;
*odf = odf.select(vec!["real_date", &*metric.clone()])?; *odf = odf.select(vec!["real_date", tkr])?;
} }
// create a new dataframe with the unique dates, and iteratively add add the metric columns // create a new dataframe with the unique dates
let mut new_df = DataFrame::new(vec![unique_dates])?; let mut new_df = DataFrame::new(vec![unique_dates])?;
for (_, odf) in out_dfs.iter() { for (_, odf) in out_dfs.iter_mut() {
new_df = new_df.left_join(odf, ["real_date"], ["real_date"])?; new_df = new_df.left_join(odf, ["real_date"], ["real_date"])?;
} }
// left join all dataframes on the date column
Ok(new_df) Ok(new_df)
} }