polars: create date_range natively

polars/polars-core/src/lib.rs

@@ -19,6 +19,8 @@ pub mod series;
 pub mod testing;
 #[cfg(test)]
 mod tests;
+#[cfg(feature = "temporal")]
+pub mod time;
 pub(crate) mod vector_hasher;
 
 #[cfg(any(feature = "dtype-categorical", feature = "object"))]

polars/polars-core/src/time.rs

@@ -0,0 +1,13 @@
+use crate::datatypes::Int64Chunked;
+use crate::prelude::DatetimeChunked;
+pub use polars_time::*;
+
+pub fn date_range(
+    start: TimeNanoseconds,
+    stop: TimeNanoseconds,
+    every: Duration,
+    closed: ClosedWindow,
+    name: &str,
+) -> DatetimeChunked {
+    Int64Chunked::new_vec(name, date_range_vec(start, stop, every, closed)).into_date()
+}

polars/polars-time/src/calendar.rs

@@ -67,7 +67,21 @@ pub fn date_range(
                 t += every.duration()
             }
         }
-        _ => unimplemented!(),
+        ClosedWindow::Right => {
+            t += every.duration();
+            while t <= stop {
+                ts.push(t);
+                t += every.duration()
+            }
+        }
+        ClosedWindow::None => {
+            t += every.duration();
+            while t < stop {
+                ts.push(t);
+                t += every.duration()
+            }
+        }
     }
+    debug_assert!(size >= ts.len());
     ts
 }

polars/polars-time/src/duration.rs

@@ -174,6 +174,7 @@ impl Duration {
     }
 
     /// Estimated duration of the window duration. Not a very good one if months != 0.
+    #[inline]
     pub const fn duration(&self) -> TimeNanoseconds {
         self.months * 30 * 24 * 3600 * NS_SECOND + self.nsecs
     }

polars/polars-time/src/lib.rs

@@ -14,4 +14,7 @@ mod test;
 mod unit;
 mod window;
 
-pub use {calendar::date_range, duration::Duration, window::Window};
+pub use {
+    calendar::date_range as date_range_vec, duration::Duration, groupby::ClosedWindow,
+    unit::TimeNanoseconds, window::Window,
+};

py-polars/polars/internals/functions.py

@@ -1,15 +1,15 @@
 from datetime import datetime, timedelta
 from typing import Optional, Sequence, Union, overload
 
-import numpy as np
-
 from polars import internals as pli
-from polars.datatypes import Datetime, py_type_to_dtype
+from polars.datatypes import py_type_to_dtype
+from polars.utils import _datetime_to_pl_timestamp, _timedelta_to_pl_duration
 
 try:
     from polars.polars import concat_df as _concat_df
     from polars.polars import concat_lf as _concat_lf
     from polars.polars import concat_series as _concat_series
+    from polars.polars import py_date_range as _py_date_range
     from polars.polars import py_diag_concat_df as _diag_concat_df
 
     _DOCUMENTING = False
@@ -135,16 +135,13 @@ def arg_where(mask: "pli.Series") -> "pli.Series":
 def date_range(
     low: datetime,
     high: datetime,
-    interval: timedelta,
-    closed: Optional[str] = None,
+    interval: Union[str, timedelta],
+    closed: Optional[str] = "both",
     name: Optional[str] = None,
 ) -> "pli.Series":
     """
     Create a date range of type `Datetime`.
 
-    .. warning::
-        This API is experimental and may change without it being considered a breaking change.
-
     Parameters
     ----------
     low
@@ -153,8 +150,10 @@ def date_range(
         Upper bound of the date range
     interval
         Interval periods
-    closed {None, 'left', 'right'}
-        Make the interval closed to the 'left', 'right', or both sides (None, the default).
+        A python timedelta object or a polars duration `str`
+        e.g.: "3d12h4m25s" # 3 days, 12 hours, 4 minutes, and 25 seconds
+    closed {None, 'left', 'right', 'both', 'none'}
+        Make the interval closed to the 'left', 'right', 'none' or 'both' sides.
     name
         Name of the output Series
 
@@ -164,10 +163,8 @@ def date_range(
 
     Examples
     --------
-    >>> from datetime import datetime, timedelta
-    >>> pl.date_range(
-    ...     datetime(1985, 1, 1), datetime(2015, 7, 1), timedelta(days=1, hours=12)
-    ... )
+    >>> from datetime import datetime
+    >>> pl.date_range(datetime(1985, 1, 1), datetime(2015, 7, 1), "1d12h")
     shape: (7426,)
     Series: '' [datetime]
     [
@@ -199,9 +196,11 @@ def date_range(
     ]
 
     """
-    values = np.arange(low, high, interval, dtype="datetime64[ns]")
-    if closed in (None, "right") and (high - low) % interval == timedelta(0):
-        values = np.append(values, np.array(high, dtype="datetime64[ns]"))
-    if closed == "right":
-        values = values[1:]
-    return pli.Series(name=name, values=values.astype(np.int64)).cast(Datetime)
+    if isinstance(interval, timedelta):
+        interval = _timedelta_to_pl_duration(interval)
+    start = _datetime_to_pl_timestamp(low)
+    stop = _datetime_to_pl_timestamp(high)
+    if name is None:
+        name = ""
+
+    return pli.wrap_s(_py_date_range(start, stop, interval, closed, name))

py-polars/polars/utils.py

@@ -1,6 +1,6 @@
 import ctypes
 import typing as tp
-from datetime import timedelta
+from datetime import datetime, timedelta, timezone
 from typing import Any, Dict, Tuple, Union
 
 import numpy as np
@@ -40,3 +40,10 @@ def _ptr_to_numpy(ptr: int, len: int, ptr_type: Any) -> np.ndarray:
 
 def _timedelta_to_pl_duration(td: timedelta) -> str:
     return f"{td.days}d{td.seconds}s{td.microseconds}us"
+
+
+def _datetime_to_pl_timestamp(dt: datetime) -> int:
+    """
+    Converts a python datetime to a timestamp in nanoseconds
+    """
+    return int(dt.replace(tzinfo=timezone.utc).timestamp() * 1e9)

py-polars/src/conversion.rs

@@ -209,6 +209,19 @@ impl ToPyObject for Wrap<DataType> {
     }
 }
 
+impl FromPyObject<'_> for Wrap<ClosedWindow> {
+    fn extract(ob: &'_ PyAny) -> PyResult<Self> {
+        let s = ob.extract::<&str>()?;
+        Ok(Wrap(match s {
+            "none" => ClosedWindow::None,
+            "both" => ClosedWindow::Both,
+            "left" => ClosedWindow::Left,
+            "right" => ClosedWindow::Right,
+            _ => panic!("{}", "closed should be any of {'none', 'left', 'right'}"),
+        }))
+    }
+}
+
 impl FromPyObject<'_> for Wrap<DataType> {
     fn extract(ob: &PyAny) -> PyResult<Self> {
         let dtype = match ob.repr().unwrap().to_str().unwrap() {

py-polars/src/lazy/dataframe.rs

@@ -306,16 +306,10 @@ impl PyLazyFrame {
         offset: &str,
         truncate: bool,
         include_boundaries: bool,
-        closed: &str,
+        closed: Wrap<ClosedWindow>,
         by: Vec<PyExpr>,
     ) -> PyLazyGroupBy {
-        let closed_window = match closed {
-            "none" => ClosedWindow::None,
-            "both" => ClosedWindow::Both,
-            "left" => ClosedWindow::Left,
-            "right" => ClosedWindow::Right,
-            _ => panic!("{}", "closed should be any of {'none', 'left', 'right'}"),
-        };
+        let closed_window = closed.0;
         let by = by
             .into_iter()
             .map(|pyexpr| pyexpr.inner)

py-polars/src/lib.rs

@@ -33,10 +33,11 @@ pub mod utils;
 use crate::conversion::{get_df, get_lf, get_pyseq, get_series, Wrap};
 use crate::error::PyPolarsEr;
 use crate::file::get_either_file;
-use crate::prelude::{DataType, PyDataType};
+use crate::prelude::{ClosedWindow, DataType, Duration, PyDataType};
 use mimalloc::MiMalloc;
 use polars::functions::diag_concat_df;
 use polars_core::export::arrow::io::ipc::read::read_file_metadata;
+use polars_core::prelude::IntoSeries;
 use pyo3::types::PyDict;
 
 #[global_allocator]
@@ -286,6 +287,19 @@ pub fn map_mul(
     lazy::map_mul(&pyexpr, py, lambda, output_type, apply_groups)
 }
 
+#[pyfunction]
+fn py_date_range(
+    start: i64,
+    stop: i64,
+    every: &str,
+    closed: Wrap<ClosedWindow>,
+    name: &str,
+) -> PySeries {
+    polars_core::time::date_range(start, stop, Duration::parse(every), closed.0, name)
+        .into_series()
+        .into()
+}
+
 #[pymodule]
 fn polars(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_class::<PySeries>().unwrap();
@@ -319,5 +333,6 @@ fn polars(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(map_mul)).unwrap();
     m.add_wrapped(wrap_pyfunction!(py_diag_concat_df)).unwrap();
     m.add_wrapped(wrap_pyfunction!(py_datetime)).unwrap();
+    m.add_wrapped(wrap_pyfunction!(py_date_range)).unwrap();
     Ok(())
 }

py-polars/tests/test_datelike.py

@@ -195,3 +195,14 @@ def test_truncate() -> None:
         assert out.dt[-3] == stop - timedelta(hours=1)
         assert out.dt[-2] == stop - timedelta(hours=1)
         assert out.dt[-1] == stop
+
+
+def test_date_range() -> None:
+    result = pl.date_range(
+        datetime(1985, 1, 1), datetime(2015, 7, 1), timedelta(days=1, hours=12)
+    )
+    assert len(result) == 7426
+    assert result.dt[0] == datetime(1985, 1, 1)
+    assert result.dt[1] == datetime(1985, 1, 2, 12, 0)
+    assert result.dt[2] == datetime(1985, 1, 4, 0, 0)
+    assert result.dt[-1] == datetime(2015, 6, 30, 12, 0)

py-polars/tests/test_series.py

@@ -1,4 +1,4 @@
-from datetime import date, datetime, timedelta
+from datetime import date
 
 import numpy as np
 import pandas as pd
@@ -790,17 +790,6 @@ def test_trigonometry_functions() -> None:
     assert np.allclose(srs_float.arctan(), np.array([0.785, 0.0, -0.785]), atol=0.01)
 
 
-def test_date_range() -> None:
-    result = pl.date_range(
-        datetime(1985, 1, 1), datetime(2015, 7, 1), timedelta(days=1, hours=12)
-    )
-    assert len(result) == 7426
-    assert result.dt[0] == datetime(1985, 1, 1)
-    assert result.dt[1] == datetime(1985, 1, 2, 12, 0)
-    assert result.dt[2] == datetime(1985, 1, 4, 0, 0)
-    assert result.dt[-1] == datetime(2015, 6, 30, 12, 0)
-
-
 def test_abs() -> None:
     # ints
     s = pl.Series([1, -2, 3, -4])