diff --git a/polars_xdt/functions.py b/polars_xdt/functions.py index 3084b41..6caee50 100644 --- a/polars_xdt/functions.py +++ b/polars_xdt/functions.py @@ -23,7 +23,9 @@ if parse_version(pl.__version__) < parse_version("0.20.16"): - from polars.utils.udfs import _get_shared_lib_location + from polars.utils.udfs import ( # type: ignore[import-not-found] + _get_shared_lib_location, + ) lib: str | Path = _get_shared_lib_location(__file__) else: diff --git a/polars_xdt/ranges.py b/polars_xdt/ranges.py index 1a3d852..52168ad 100644 --- a/polars_xdt/ranges.py +++ b/polars_xdt/ranges.py @@ -10,7 +10,7 @@ if TYPE_CHECKING: from datetime import date, datetime, timedelta - from polars.type_aliases import ClosedInterval, IntoExprColumn, TimeUnit + from polars.type_aliases import ClosedInterval, IntoExprColumn @overload @@ -20,8 +20,6 @@ def date_range( interval: str | timedelta = "1d", *, closed: ClosedInterval = ..., - time_unit: TimeUnit | None = ..., - time_zone: str | None = ..., eager: Literal[False] = ..., weekend: Sequence[str] = ..., holidays: Sequence[date] | None = ..., @@ -35,8 +33,6 @@ def date_range( interval: str | timedelta = "1d", *, closed: ClosedInterval = ..., - time_unit: TimeUnit | None = ..., - time_zone: str | None = ..., eager: Literal[True], weekend: Sequence[str] = ..., holidays: Sequence[date] | None = ..., @@ -50,8 +46,6 @@ def date_range( interval: str | timedelta = "1d", *, closed: ClosedInterval = ..., - time_unit: TimeUnit | None = ..., - time_zone: str | None = ..., eager: bool = ..., weekend: Sequence[str] = ..., holidays: Sequence[date] | None = ..., @@ -59,13 +53,11 @@ def date_range( def date_range( # noqa: PLR0913 - start: date | datetime | IntoExprColumn, - end: date | datetime | IntoExprColumn, + start: date | IntoExprColumn, + end: date | IntoExprColumn, interval: str | timedelta = "1bd", *, closed: ClosedInterval = "both", - time_unit: TimeUnit | None = None, - time_zone: str | None = None, eager: bool = False, weekend: Sequence[str] = ("Sat", "Sun"), holidays: Sequence[date] | None = None, @@ -87,12 +79,6 @@ def date_range( # noqa: PLR0913 "Examples" section below). closed : {'both', 'left', 'right', 'none'} Define which sides of the range are closed (inclusive). - time_unit : {None, 'ns', 'us', 'ms'} - Time unit of the resulting ``Datetime`` data type. - Only takes effect if the output column is of type ``Datetime``. - time_zone - Time zone of the resulting ``Datetime`` data type. - Only takes effect if the output column is of type ``Datetime``. eager Evaluate immediately and return a ``Series``. If set to ``False`` (default), return an expression instead. @@ -152,8 +138,6 @@ def date_range( # noqa: PLR0913 end, interval, closed=closed, - time_unit=time_unit, - time_zone=time_zone, eager=False, ) expr = expr.filter(~expr.dt.date().is_in(holidays)) diff --git a/src/expressions.rs b/src/expressions.rs index ab71e33..29d35f1 100644 --- a/src/expressions.rs +++ b/src/expressions.rs @@ -203,3 +203,25 @@ fn ewma_by_time(inputs: &[Series], kwargs: EwmTimeKwargs) -> PolarsResult polars_bail!(InvalidOperation: "First argument should be a date or datetime type."), } } + +#[polars_expr(output_type_func=same_output)] +fn truncate(inputs: &[Series]) -> PolarsResult { + let ser = &inputs[0]; + let ca = ser.datetime()?; + let tu = ca.time_unit(); + let tz = ca.time_zone(); + let ca = &ca.0; + let duration = 7200000000; + let out = ca + .apply_values(|t| { + let mut remainder = t % duration; + if remainder < 0 { + remainder += duration + } + t - remainder + }) + .into_datetime(tu, tz.clone()) + .into_series(); + + Ok(out) +}