diff --git a/Cargo.toml b/Cargo.toml index 9ae24f2..15abc53 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,3 +13,4 @@ polars-ffi = { version = "0.33.2", default-features = false } polars-plan = { version = "0.33.2", default-feautres = false } polars-lazy = { version = "0.33.2", default-features = false } polars-time = { version = "0.33.2", default-features = false } +chrono = { version = "0.4", default-features = false } \ No newline at end of file diff --git a/src/expression_lib/Cargo.toml b/src/expression_lib/Cargo.toml index f180a75..b27ca69 100644 --- a/src/expression_lib/Cargo.toml +++ b/src/expression_lib/Cargo.toml @@ -12,5 +12,6 @@ crate-type = ["cdylib"] bdays = "0.1.3" polars = { workspace = true, features = ["fmt"], default-features = false } polars-time = { workspace = true, default-features = false } +chrono = { workspace = true, default-features = false } pyo3 = { version = "0.19.0", features = ["extension-module"]} pyo3-polars = { version = "*", path = "../../pyo3-polars", features = ["derive"] } diff --git a/src/expression_lib/expression_lib/__init__.py b/src/expression_lib/expression_lib/__init__.py index 61e0511..6bf738e 100644 --- a/src/expression_lib/expression_lib/__init__.py +++ b/src/expression_lib/expression_lib/__init__.py @@ -5,15 +5,16 @@ lib = _get_shared_lib_location(__file__) -@pl.api.register_expr_namespace("language") -class Language: +@pl.api.register_expr_namespace("bdt") +class BusinessDayTools: def __init__(self, expr: pl.Expr): - self._expr = expr + self._expr = expr.cast(pl.Int32) - def add_bday(self) -> pl.Expr: + def add_bday(self, n) -> pl.Expr: return self._expr._register_plugin( lib=lib, symbol="add_bday", is_elementwise=True, + args = [n], ) diff --git a/src/expression_lib/src/expressions.rs b/src/expression_lib/src/expressions.rs index 01a4869..add644c 100644 --- a/src/expression_lib/src/expressions.rs +++ b/src/expression_lib/src/expressions.rs @@ -1,9 +1,28 @@ use polars::prelude::*; use pyo3_polars::derive::polars_expr; +use bdays::{calendars::WeekendsOnly, HolidayCalendar}; +use chrono::{NaiveDate, NaiveDateTime, Datelike}; +use std::cmp::min; + +// cool, seems to work! let's try it with Datetime? + #[polars_expr(output_type=Date)] fn add_bday(inputs: &[Series]) -> PolarsResult { - let ca = inputs[0].date()?; - Ok(ca.clone().into_series()) -} + let ca = inputs[0].i32()?; + let n = inputs[1].i32()?.get(0).unwrap(); + // wow, this is super-slow, keeps repeatedly converting + // between NaiveDate and timestamp! + // multiple times for each element! max once should be enough! + // or twice if necessary. but not so many + + let out = ca.apply_values( + |x|{ + let w = (x / (3600 * 24 - 3)) % 7 + 1; + let n_days = n + (n + w) / 5 * 2; + x + n_days + } + ); + Ok(out.cast(&DataType::Date).unwrap().into_series()) +} diff --git a/src/run.py b/src/run.py index 17ea385..132afb5 100644 --- a/src/run.py +++ b/src/run.py @@ -1,12 +1,16 @@ import polars as pl -from expression_lib import Language +from expression_lib import BusinessDayTools from datetime import date, datetime +import numpy as np df = pl.DataFrame({ - "names": ["Richard", "Alice", "Bob"], - "dates": [1]*3, - "dates2": [date(2020, 1, 1)]*3, + "dates": pl.datetime_range(date(1, 1, 1), date(9999, 1, 1), eager=True), }) -print(df) +df = df.filter(pl.col('dates').dt.weekday() <6) -print(df.with_columns(dates_plus_1=pl.col('dates2').language.add_bday())) +print(df.with_columns(dates_shifted=pl.col('dates').bdt.add_bday(n=15))[20:28]) +print(df.with_columns(dates_shifted=pl.Series(np.busday_offset(df['dates'], 15)))[20:28]) + +import pandas as pd +dfpd = df.to_pandas() +print((dfpd + pd.tseries.offsets.BusinessDay(15)).iloc[20:28])