ARKB for ARKF

This commit is contained in:
Elijah McMorris 2024-04-24 18:01:39 -07:00
parent 3ad11c2881
commit 8ed0fb751c
Signed by: NexVeridian
SSH key fingerprint: SHA256:bsA1SKZxuEcEVHAy3gY1HUeM5ykRJl0U0kQHQn0hMg8

View file

@ -376,6 +376,16 @@ impl Ark {
.otherwise(col("company")) .otherwise(col("company"))
.alias("company"), .alias("company"),
]) ])
.with_columns(vec![
when(col("company").eq(lit("ARK BITCOIN ETF HOLDCO (ARKF)")))
.then(lit("ARKB"))
.otherwise(col("ticker"))
.alias("ticker"),
when(col("company").eq(lit("ARK BITCOIN ETF HOLDCO (ARKF)")))
.then(lit("ARKB"))
.otherwise(col("company"))
.alias("company"),
])
.collect() .collect()
{ {
df = x; df = x;
@ -723,4 +733,40 @@ mod tests {
Ok(()) Ok(())
} }
#[test]
#[serial]
fn arkf_format_arkb() -> Result<(), Error> {
let test_df = df![
"date" => ["2024-01-01", "2024-01-02"],
"ticker" => [None::<&str>, Some("TSLA")],
"cusip" => ["123abc", "TESLA"],
"company" => ["ARK BITCOIN ETF HOLDCO (ARKF)", "TESLA"],
"market_value" => [100, 400],
"shares" => [10, 20],
"share_price" => [10, 20],
"weight" => [10.00, 20.00]
]?;
Ark::write_df_parquet("data/test/ARKF.parquet".into(), test_df.clone().into())?;
let read = Ark::new(Source::Read, Ticker::ARKF, Some("data/test".to_owned()))?.collect()?;
fs::remove_file("data/test/ARKF.parquet")?;
let df = Ark::df_format(read.into())?.collect()?;
assert_eq!(
df,
df![
"date" => ["2024-01-01", "2024-01-02"],
"ticker" => ["ARKB", "TSLA"],
"cusip" => ["123abc", "TESLA"],
"company" => ["ARKB", "TESLA"],
"market_value" => [100, 400],
"shares" => [10, 20],
"share_price" => [10, 20],
"weight" => [10.00, 20.00]
]?
);
Ok(())
}
} }