diff --git a/polars_business/perf.py b/polars_business/perf.py index dd72919..cdf2f78 100644 --- a/polars_business/perf.py +++ b/polars_business/perf.py @@ -1,8 +1,8 @@ import timeit import numpy as np -# BENCHMARKS = [1, 2, 3, 4] -BENCHMARKS = [1] +BENCHMARKS = [1, 2, 3, 4] +# BENCHMARKS = [1] # BENCHMARK 1: NO HOLIDAYS INVOLVED @@ -50,10 +50,10 @@ def time_it(statement): ) print("NumPy: ", time_it("result_np = np.busday_offset(input_dates, 17)")) -# uncomment, too slow... -print( - "pandas: ", time_it("result_pd = df_pd['ts'] + pd.tseries.offsets.BusinessDay(17)") -) + # uncomment, too slow... + # print( + # "pandas: ", time_it("result_pd = df_pd['ts'] + pd.tseries.offsets.BusinessDay(17)") + # ) # BENCHMARK 2: WITH HOLIDAYS diff --git a/polars_business/polars_business/Cargo.toml b/polars_business/polars_business/Cargo.toml index 9b7af44..ceb4685 100644 --- a/polars_business/polars_business/Cargo.toml +++ b/polars_business/polars_business/Cargo.toml @@ -17,6 +17,7 @@ polars-ops = { workspace = true, default-features = false } pyo3 = { version = "0.20.0", features = ["extension-module"] } pyo3-polars = { git = "https://github.com/pola-rs/pyo3-polars", rev = "0a3df66ce626504e418491cb4b1d31bd3ee7ed6c", features = ["derive"] } serde = { version = "1", features = ["derive"] } +chrono = "0.4.31" [target .'cfg(target_os = "linux")'] jemallocator = { version = "0.5", features = ["disable_initial_exec_tls"] } diff --git a/polars_business/polars_business/src/business_days.rs b/polars_business/polars_business/src/business_days.rs index 9e299ee..db7ffed 100644 --- a/polars_business/polars_business/src/business_days.rs +++ b/polars_business/polars_business/src/business_days.rs @@ -1,6 +1,7 @@ use ahash::AHashMap; use polars::prelude::arity::try_binary_elementwise; use polars::prelude::*; +use chrono::NaiveDateTime; pub(crate) fn weekday(x: i32) -> i32 { // the first modulo might return a negative number, so we add 7 and take @@ -67,6 +68,11 @@ fn roll(n_days: i32, x_weekday: i32, weekend: &[i32]) -> i32 { n_days } +fn its_a_business_date_error_message(x: i32) -> PolarsResult{ + let date = NaiveDateTime::from_timestamp_opt(x as i64 * 24 * 60 * 60, 0).unwrap().format("%Y-%m-%d"); + polars_bail!(ComputeError: format!("date {} is not a business date, cannot advance. `roll` argument coming soon.", date)) +} + pub(crate) fn calculate_n_days_with_holidays( x: i32, n: i32, @@ -76,13 +82,13 @@ pub(crate) fn calculate_n_days_with_holidays( let x_weekday = weekday(x_mod_7); if x_weekday >= 5 { - polars_bail!(ComputeError: format!("date {} is not a business date, cannot advance. `roll` argument coming soon.", x)) + return its_a_business_date_error_message(x); }; let mut n_days = calculate_n_days_without_holidays_fast(n, x_weekday); if holidays.binary_search(&x).is_ok() { - polars_bail!(ComputeError: format!("date {} is not a business date, cannot advance. `roll` argument coming soon.", x)) + return its_a_business_date_error_message(x); } let mut count_hols = count_holidays(x, x + n_days, holidays); while count_hols > 0 { @@ -111,14 +117,14 @@ pub(crate) fn calculate_n_days_with_weekend_and_holidays( let x_weekday = weekday(x_mod_7); if weekend.contains(&x_weekday) { - polars_bail!(ComputeError: format!("date {} is not a business date, cannot advance. `roll` argument coming soon.", x)) + return its_a_business_date_error_message(x); }; let mut n_days = calculate_n_days_without_holidays_slow(x_weekday, n, weekend.len() as i32, cache); if holidays.binary_search(&x).is_ok() { - polars_bail!(ComputeError: format!("date {} is not a business date, cannot advance. `roll` argument coming soon.", x)) + return its_a_business_date_error_message(x); } let mut count_hols = count_holidays(x, x + n_days, holidays); while count_hols > 0 { @@ -157,7 +163,7 @@ pub(crate) fn calculate_n_days_with_weekend( let n_weekend = weekend.len() as i32; if weekend.contains(&x_weekday) { - polars_bail!(ComputeError: format!("date {} is not a business date, cannot advance. `roll` argument coming soon.", x)) + return its_a_business_date_error_message(x); }; Ok(calculate_n_days_without_holidays_slow( @@ -235,7 +241,7 @@ pub(crate) fn impl_advance_n_days( ca.try_apply(|x_date| { let x_weekday = weekday(x_date); if x_weekday >= 5 { - polars_bail!(ComputeError: format!("date {} is not a business date, cannot advance. `roll` argument coming soon.", x_date)) + return its_a_business_date_error_message(x_date); } Ok(x_date+(calculate_n_days_without_holidays_fast(n, x_weekday))) }) @@ -267,7 +273,7 @@ pub(crate) fn impl_advance_n_days( if holidays.is_empty() && weekend == [5, 6] { let x_weekday = weekday(x_date); if x_weekday >= 5 { - polars_bail!(ComputeError: format!("date {} is not a business date, cannot advance. `roll` argument coming soon.", x_date)) + return its_a_business_date_error_message(x_date).map(|x| Some(x)); } Ok(Some( x_date + (calculate_n_days_without_holidays_fast(n, x_weekday)), @@ -315,7 +321,7 @@ pub(crate) fn impl_advance_n_days( let x_date = (x / multiplier) as i32; let x_weekday = weekday(x_date); if x_weekday >= 5 { - polars_bail!(ComputeError: format!("date {} is not a business date, cannot advance. `roll` argument coming soon.", x_date)) + return its_a_business_date_error_message(x_date).map(|x| x as i64); } Ok(x+(calculate_n_days_without_holidays_fast(n, x_weekday) as i64 *multiplier)) }) @@ -358,7 +364,7 @@ pub(crate) fn impl_advance_n_days( if holidays.is_empty() && weekend == [5, 6] { let x_weekday = weekday(x_date); if x_weekday >= 5 { - polars_bail!(ComputeError: format!("date {} is not a business date, cannot advance. `roll` argument coming soon.", x_date)) + return its_a_business_date_error_message(x_date).map(|x| Some(x as i64)); } Ok(Some( x + (calculate_n_days_without_holidays_fast(n, x_weekday) as i64 diff --git a/polars_business/run.py b/polars_business/run.py index 865e70d..9be26d2 100644 --- a/polars_business/run.py +++ b/polars_business/run.py @@ -8,7 +8,7 @@ start = datetime(2000, 1, 18) n = 11 weekend = ["Sat", "Sun"] -holidays = [] # [date(2000, 1, 29)] +holidays = [date(2000, 1, 18)] weekmask = [0 if reverse_mapping[i] in weekend else 1 for i in range(7)] df = pl.DataFrame({"dates": [start]})