kinda working!

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 }

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"] }

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],
         )

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<Series> {
-    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())
+}

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])