feat: RKLB

This commit is contained in:
Elijah McMorris 2025-06-01 17:34:35 -07:00
parent 1ad252cd00
commit e297be713e
Signed by: NexVeridian
SSH key fingerprint: SHA256:bsA1SKZxuEcEVHAy3gY1HUeM5ykRJl0U0kQHQn0hMg8

View file

@ -17,6 +17,7 @@ pub enum Ticker {
XYZ,
CASH_USD,
TSM,
RKLB,
}
impl Ticker {
@ -38,6 +39,7 @@ impl Ticker {
Self::XYZ => Self::xyz(df),
Self::CASH_USD => Self::cash_usd(df),
Self::TSM => Self::tsm(df),
Self::RKLB => Self::rklb(df),
}
}
@ -250,6 +252,30 @@ impl Ticker {
Ok(df.into())
}
fn rklb(df: DF) -> Result<DF, Error> {
let mut df = df.collect()?;
if let Ok(x) = df
.clone()
.lazy()
.with_columns(vec![
when(col("company").eq(lit("ROCKET LAB")))
.then(lit("RKLB"))
.otherwise(col("ticker"))
.alias("ticker"),
when(col("company").eq(lit("ROCKET LAB USA")))
.then(lit("ROCKET LAB"))
.otherwise(col("company"))
.alias("company")
])
.collect()
{
df = x;
}
Ok(df.into())
}
}
#[cfg(test)]
@ -364,6 +390,17 @@ mod tests {
&[Some("TMSC")]
)?,
)]
#[case::rklb(
Ticker::RKLB,
defualt_df(
&[Some("RKLB"), Some("RKLB"), None::<&str>],
&[Some("ROCKET LAB"), Some("ROCKET LAB USA"), Some("ROCKET LAB")],
)?,
defualt_df(
&[Some("RKLB"), Some("RKLB"), Some("RKLB")],
&[Some("ROCKET LAB"), Some("ROCKET LAB"), Some("ROCKET LAB")]
)?,
)]
fn matrix(
#[case] ticker: Ticker,
#[case] input: DataFrame,