0.3.4

src/util.rs

@@ -21,7 +21,6 @@ pub enum Ticker {
     ARKW,
     ARKX,
 }
-
 impl Ticker {
     pub fn value(&self) -> &str {
         match *self {
@@ -36,83 +35,111 @@ impl Ticker {
     }
 }
 
-pub fn merge_csv_to_parquet(ticker: Ticker) -> Result<(), Box<dyn Error>> {
+pub enum DF {
-    let mut dfs = vec![];
+    LazyFrame(LazyFrame),
-    for x in glob(&format!("data/csv/{}/*", ticker))?.filter_map(Result::ok) {
+    DataFrame(DataFrame),
-        dfs.push(LazyCsvReader::new(x).finish()?);
+}
+impl From<LazyFrame> for DF {
+    fn from(lf: LazyFrame) -> Self {
+        DF::LazyFrame(lf)
+    }
+}
+impl From<DataFrame> for DF {
+    fn from(df: DataFrame) -> Self {
+        DF::DataFrame(df)
+    }
+}
+impl DF {
+    pub fn collect(self) -> Result<DataFrame, Box<dyn Error>> {
+        match self {
+            DF::LazyFrame(x) => Ok(x.collect()?),
+            DF::DataFrame(x) => Ok(x),
+        }
+    }
+    pub fn lazy(self) -> LazyFrame {
+        match self {
+            DF::LazyFrame(x) => x,
+            DF::DataFrame(x) => x.lazy(),
         }
-    let mut df = concat(dfs, false, true)?;
-
-    if read_parquet(ticker).is_ok() {
-        let df_old = read_parquet(ticker)?;
-        df = concat_df(vec![df_old, df])?;
-        write_parquet(ticker, df_sort(df.collect()?)?)?;
-    } else {
-        write_parquet(ticker, df_format(df)?)?;
     }
-
-    Ok(())
 }
 
 pub enum Source {
+    Read,
     Ark,
     ApiIncremental,
     ApiFull,
 }
-
+struct Ark {
-pub fn update_parquet(ticker: Ticker, source: Source) -> Result<(), Box<dyn Error>> {
+    df: DF,
-    let mut df = read_parquet(ticker)?;
+    ticker: Ticker,
+}
+impl Ark {
+    pub fn new(source: Source, ticker: Ticker) -> Result<Self, Box<dyn Error>> {
+        let mut ark = Self {
+            df: Self::read_parquet(ticker)?.into(),
+            ticker,
+        };
 
         let update = match source {
-        Source::Ark => get_csv_ark(ticker)?,
+            Source::Read => None,
+            Source::Ark => Some(ark.get_csv_ark()?),
             Source::ApiIncremental => {
-            let last_day = df
+                let last_day = ark
-                .clone()
+                    .df
                     .collect()?
                     .column("date")
                     .unwrap()
                     .max()
                     .and_then(NaiveDate::from_num_days_from_ce_opt);
-            get_api(ticker, last_day)?
+                Some(ark.get_api(last_day)?)
             }
-        Source::ApiFull => get_api(ticker, None)?,
+            Source::ApiFull => Some(ark.get_api(None)?),
         };
 
-    df = concat_df(vec![df, update])?;
+        if let Some(update) = update {
-    write_parquet(ticker, df_sort(df.collect()?)?)?;
+            ark.df = Self::concat_df(vec![ark.df, update.into()])?.into();
-    Ok(())
+        }
-}
+        Ok(ark)
-
-fn concat_df(mut dfs: Vec<LazyFrame>) -> Result<LazyFrame, Box<dyn Error>> {
-    // with dedupe and format
-    for x in &mut dfs {
-        *x = df_format(x.to_owned())?.lazy();
     }
-    let mut df = concat(dfs, false, true)?;
-    df = df.unique_stable(None, UniqueKeepStrategy::First);
-    Ok(df)
-}
 
-pub fn read_parquet(ticker: Ticker) -> Result<LazyFrame, Box<dyn Error>> {
+    pub fn write_parquet(mut self) -> Result<Self, Box<dyn Error>> {
+        ParquetWriter::new(File::create(format!(
+            "data/parquet/{}.parquet",
+            self.ticker
+        ))?)
+        .finish(&mut self.df.collect()?)?;
+        Ok(self)
+    }
+
+    pub fn df_sort(&self) -> Result<&Self, Box<dyn Error>> {
+        let df = self
+            .df
+            .collect()?
+            .sort(["date", "weight"], vec![false, true])?;
+        self.df = df.into();
+        Ok(self)
+    }
+
+    fn read_parquet(ticker: Ticker) -> Result<LazyFrame, Box<dyn Error>> {
         let df = LazyFrame::scan_parquet(
             format!("data/parquet/{}.parquet", ticker),
             ScanArgsParquet::default(),
         )?;
         Ok(df)
-}
+    }
 
-pub fn write_parquet(ticker: Ticker, mut df: DataFrame) -> Result<(), Box<dyn Error>> {
+    fn concat_df(mut dfs: Vec<DF>) -> Result<DF, Box<dyn Error>> {
-    ParquetWriter::new(File::create(format!("data/parquet/{}.parquet", ticker))?)
+        // with dedupe and format
-        .finish(&mut df)?;
+        // for x in &mut dfs {
-    Ok(())
+        //     *x = Self::df_format(x.to_owned())?.lazy();
-}
+        // }
+        let mut df = concat(dfs, false, true)?;
+        df = df.unique_stable(None, UniqueKeepStrategy::First);
+        Ok(df.into())
+    }
 
-pub fn df_sort(df: DataFrame) -> Result<DataFrame, Box<dyn Error>> {
+    fn df_format(df: LazyFrame) -> Result<DataFrame, Box<dyn Error>> {
-    let df = df.sort(["date", "weight"], vec![false, true])?;
-    Ok(df)
-}
-
-pub fn df_format(df: LazyFrame) -> Result<DataFrame, Box<dyn Error>> {
         let mut df = df.collect()?;
 
         if df.get_column_names().contains(&"market_value_($)") {
@@ -233,37 +260,64 @@ pub fn df_format(df: LazyFrame) -> Result<DataFrame, Box<dyn Error>> {
                 "share_price",
                 "weight",
             ])?;
-    } else {
+        } else if !df
+            .get_column_names()
+            .eq(&["date", "ticker", "cusip", "company", "weight"])
+        {
             df = df.select(["date", "ticker", "cusip", "company", "weight"])?;
         }
 
         Ok(df)
-}
+    }
 
-pub fn get_api(ticker: Ticker, last_day: Option<NaiveDate>) -> Result<LazyFrame, Box<dyn Error>> {
+    pub fn get_api(self, last_day: Option<NaiveDate>) -> Result<LazyFrame, Box<dyn Error>> {
-    let url = match (ticker, last_day) {
+        let tic = self.ticker;
-        (Ticker::ARKVC, Some(last_day)) => format!(
+        let url = match (tic, last_day) {
+            (self::Ticker::ARKVC, Some(last_day)) => format!(
                 "https://api.nexveridian.com/arkvc_holdings?end={}",
                 last_day
             ),
-        (ticker, Some(last_day)) => format!(
+            (tic, Some(last_day)) => format!(
                 "https://api.nexveridian.com/ark_holdings?ticker={}&end={}",
-            ticker, last_day
+                tic.value(),
+                last_day
             ),
-        (Ticker::ARKVC, None) => "https://api.nexveridian.com/arkvc_holdings".to_owned(),
+            (self::Ticker::ARKVC, None) => "https://api.nexveridian.com/arkvc_holdings".to_owned(),
-        (ticker, None) => {
+            (tic, None) => {
-            format!("https://api.nexveridian.com/ark_holdings?ticker={}", ticker)
+                format!(
+                    "https://api.nexveridian.com/ark_holdings?ticker={}",
+                    tic.value()
+                )
             }
         };
         Reader::Json.get_data_url(url)
-}
+    }
 
-pub fn get_csv_ark(ticker: Ticker) -> Result<LazyFrame, Box<dyn Error>> {
+    pub fn get_csv_ark(self) -> Result<LazyFrame, Box<dyn Error>> {
-    let url = match ticker {
+        let url = match self.ticker {
-        Ticker::ARKVC => "https://ark-ventures.com/wp-content/uploads/funds-etf-csv/ARK_VENTURE_FUND_HOLDINGS.csv".to_owned(),
+            self::Ticker::ARKVC => "https://ark-ventures.com/wp-content/uploads/funds-etf-csv/ARK_VENTURE_FUND_HOLDINGS.csv".to_owned(),
-        _ => format!("https://ark-funds.com/wp-content/uploads/funds-etf-csv/ARK_{}_ETF_{}_HOLDINGS.csv", ticker.value(), ticker),
+            _ => format!("https://ark-funds.com/wp-content/uploads/funds-etf-csv/ARK_{}_ETF_{}_HOLDINGS.csv", self.ticker.value(), self.ticker),
         };
         Reader::Csv.get_data_url(url)
+    }
+
+    pub fn merge_old_csv_to_parquet(ticker: Ticker) -> Result<(), Box<dyn Error>> {
+        let mut dfs = vec![];
+        for x in glob(&format!("data/csv/{}/*", ticker))?.filter_map(Result::ok) {
+            dfs.push(LazyCsvReader::new(x).finish()?);
+        }
+        let mut df = concat(dfs, false, true)?;
+
+        if Self::read_parquet(ticker).is_ok() {
+            let df_old = Self::read_parquet(ticker)?;
+            df = Self::concat_df(vec![df_old, df])?;
+            write_parquet(ticker, df_sort(df.collect()?)?)?;
+        } else {
+            write_parquet(ticker, df_format(df)?)?;
+        }
+
+        Ok(())
+    }
 }
 
 pub enum Reader {