Function to add weekdays from date

Question

I need to create a generic function using SQL on Databricks that adds a number of weekdays to a date. It should receive two parameters, the initial_date and the number of weekdays to add. I was able to do that in a query, but when I try to create the function I receive the error

UNSUPPORTED_SUBQUERY_EXPRESSION_CATEGORY.CORRELATED_REFERENCE] Unsupported subquery expression: Expressions referencing the outer query are not supported outside of WHERE/HAVING clauses: "date_add(dt_initial, day_offset) AS business_date"

Follow my function, any idea is welcome!

CREATE OR REPLACE FUNCTION add_weekdays(dt_initial DATE, num_days INT)
RETURNS DATE
RETURN (
WITH days AS (
    SELECT EXPLODE(SEQUENCE(0, num_days * 2)) AS day_offset
),
valid_days AS (
    SELECT DATE_ADD(dt_initial, day_offset) AS business_date,
           ROW_NUMBER() OVER (ORDER BY day_offset) AS rn
    FROM days
    WHERE WEEKDAY(DATE_ADD(dt_initial, day_offset)) BETWEEN 0 AND 4
)
SELECT business_date FROM valid_days WHERE rn = num_days + 1 limit 1
);

Answer 1

Here is a python scalar function implementation:

CREATE OR REPLACE FUNCTION ADD_WEEKDAYS(
  X DATE COMMENT 'DATE TO ADD N WEEKDAYS', 
  N INT COMMENT 'INTEGER WEEKDAYS TO ADD TO X'
)
  RETURNS DATE
  LANGUAGE PYTHON
  COMMENT 'ADDS N WEEKDAYS TO A DATE (SKIPS WEEKENDS)'
  AS $$
    from datetime import timedelta

    def add_weekdays(x, n):
        """Adds n weekdays to a date (skips weekends)"""
        current_date = x
        added_days = 0
        while added_days < n:
            current_date += timedelta(days = 1)
            if current_date.weekday() < 5:
                added_days += 1
        return current_date

    return add_weekdays(X, N) if X and N else None
  $$

SELECT ADD_WEEKDAYS("2025-03-01", 1) AS X;