wip

src/py/qdf.pyi

@@ -8,17 +8,16 @@ __all__ = [
 ]
 
 # qdf/load_qdf.pyi
-def load_qdf(*args, **kwargs) -> DataFrame: 
-    ...
+def load_qdf(*args, **kwargs) -> DataFrame: ...
 
 # qdf/load_qdf_from_download_bank.pyi
-def load_qdf_from_download_bank(*args, **kwargs) -> DataFrame: 
-    ...
+def load_qdf_from_download_bank(*args, **kwargs) -> DataFrame: ...
 
 # qdf/reduce_dataframe.pyi
-def reduce_dataframe(*args, **kwargs) -> DataFrame: 
-    ...
+def reduce_dataframe(*args, **kwargs) -> DataFrame: ...
 
 # qdf/update_dataframe.pyi
-def update_dataframe(*args, **kwargs) -> DataFrame: 
-    ...
+def update_dataframe(*args, **kwargs) -> DataFrame: ...
+
+# qdf/pivots.pyi
+def pivot_dataframe_by_ticker(*args, **kwargs) -> DataFrame: ...

src/py/qdf.rs

@@ -9,6 +9,7 @@ pub fn qdf(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(load_qdf_from_download_bank, m)?)?;
     m.add_function(wrap_pyfunction!(reduce_dataframe, m)?)?;
     m.add_function(wrap_pyfunction!(update_dataframe, m)?)?;
+    m.add_function(wrap_pyfunction!(pivot_dataframe_by_ticker, m)?)?;
     Ok(())
 }
 
@@ -76,3 +77,12 @@ pub fn update_dataframe(
             .unwrap(),
     ))
 }
+
+/// Python wrapper for pivoting a dataframe by ticker.
+/// See [`crate::utils::qdf::pivots::pivot_dataframe_by_ticker`] for full documentation.
+#[pyfunction]
+pub fn pivot_dataframe_by_ticker(df: PyDataFrame, metric: Option<String>) -> PyResult<PyDataFrame> {
+    Ok(PyDataFrame(
+        crate::utils::qdf::pivots::pivot_dataframe_by_ticker(df.into(), metric).unwrap(),
+    ))
+}

src/utils/qdf/pivots.rs

@@ -37,16 +37,20 @@ pub fn pivot_dataframe_by_ticker(
         .collect::<Vec<String>>();
     keep_cols.push(metric.clone());
 
-    let 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() {
+        // select keep_cols
+        *odf = odf.select(vec!["real_date", &*metric.clone()])?;
+    }
 
     // create a new dataframe with the unique dates, and iteratively add add the metric columns
     let mut new_df = DataFrame::new(vec![unique_dates])?;
 
-    
-
-
-
+    for (_, odf) in out_dfs.iter() {
+        new_df = new_df.left_join(odf, ["real_date"], ["real_date"])?;
+    }
 
+    Ok(new_df)
 }
 
 /// Splits a dataframe by ticker.
@@ -92,6 +96,9 @@ fn split_df_by_tickers(
     Ok(df_outs)
 }
 
+/// Splits a QDF container data for a single ticker into a Vec of DataFrames (one per metric).
+/// The resulting DataFrames will have the "real_date" column and the metric column.
+#[allow(dead_code)]
 fn single_ticker_qdf_to_timeseries(mut df: DataFrame) -> Result<Vec<DataFrame>, Box<dyn Error>> {
     let mut df_vec = Vec::new();