daft.functions.add_months#

def add_months(expr: Expression, months: Expression) -> Expression:
    """Adds a number of months to a date or timestamp.

    Mirrors Spark's ``add_months``: when the start day exceeds the number of days in
    the resulting month, the result is clamped to the last day of that month. The
    return type is always Date, even when the input is a Timestamp (the time-of-day
    component is dropped before the shift).

    Args:
        expr: A Date or Timestamp expression.
        months: An integer expression for the number of months to add (may be negative).

    Returns:
        Expression: a Date expression shifted by the given number of months.

    Examples:
        >>> import daft
        >>> from daft.functions import add_months
        >>> df = daft.from_pydict({"d": ["2023-01-31", "2024-01-31"], "n": [1, 1]})
        >>> df = df.with_column("d", df["d"].cast(daft.DataType.date()))
        >>> df = df.with_column("result", add_months(df["d"], df["n"]))
        >>> df.schema()["result"].dtype == daft.DataType.date()
        True
    """
    return Expression._call_builtin_scalar_fn("add_months", expr, months)