How can I dynamically generate and apply a custom week_mask tuple when adding business days in Polars?

Below is one way to solve the problem. The key idea is that the built‐in function dt.add_business_days expects a constant (Python-level) sequence for its week mask. In other words, you can’t feed it an expression that varies row by row. Instead, you can write a Python function that

1. Builds the week mask tuple from the WKND1 and WKND2 columns.
2. Manually “adds” business days (skipping any day that is a weekend per that mask).

Then use Polars’ row‐wise computation (via pl.struct(...).apply(...)) to generate your NEW_DT.

For example:

PYTHONCopied
import polars as pl
from datetime import timedelta

# Map days to weekday indices (0=Monday ... 6=Sunday)
day_to_idx = {
    "Monday": 0, "Tuesday": 1, "Wednesday": 2, 
    "Thursday": 3, "Friday": 4, "Saturday": 5, "Sunday": 6
}

def add_business_days(dt, n_days, wknd1, wknd2):
    # Construct a set of weekend weekday indexes.
    wknd = set()
    if wknd1 in day_to_idx:
        wknd.add(day_to_idx[wknd1])
    if wknd2 in day_to_idx:
        wknd.add(day_to_idx[wknd2])
    
    # Alternatively, if you want to see the week mask tuple, you can do:
    # days = ["Monday","Tuesday","Wednesday","Thursday","Friday","Saturday","Sunday"]
    # week_mask = tuple(d not in {wknd1, wknd2} for d in days)
    # print("week_mask:", week_mask)   # for debugging
    
    # Iterate until we've added the desired number of business days.
    current_date = dt
    days_remaining = n_days
    while days_remaining > 0:
        current_date += timedelta(days=1)
        # current_date.weekday() returns integer 0=Monday ... 6=Sunday
        if current_date.weekday() not in wknd:
            days_remaining -= 1
    return current_date

# Sample data
data = {
    "DT": ["2025-01-02", "2025-01-09", "2025-01-10", "2025-01-15"],
    "N_DAYS": [2, 2, 2, 1],
    "WKND1": ["Saturday", "Friday", "Saturday", "Saturday"],
    "WKND2": ["Sunday", "Saturday", None, "Sunday"],
}

# Create a Polars DataFrame and cast DT to date
df = pl.DataFrame(data).with_columns(pl.col("DT").str.strptime(pl.Date, fmt="%Y-%m-%d"))

# Use pl.struct() with a row-wise apply to compute NEW_DT dynamically
df = df.with_columns(
    pl.struct(["DT", "N_DAYS", "WKND1", "WKND2"]).apply(
        lambda row: add_business_days(row["DT"], row["N_DAYS"], row["WKND1"], row["WKND2"])
    ).alias("NEW_DT")
)

print(df)

Explanation

1. The function add_business_days builds a set of weekend days based on the weekday indexes of WKND1 and WKND2.
2. It then “steps” day-by-day from the given date (skipping days where the weekday is in the weekend set) until it has added N_DAYS business days.
3. In the Polars pipeline, we use pl.struct to pack the required columns together and then apply the function row-wise.
4. This calculates the NEW_DT per row with the dynamic (record‐dependent) weekend.

This approach sidesteps the limitation that dt.add_business_days can only handle a constant week mask.