Added date_ranges method to Polars - closes #69

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## 0.10.1 (unreleased)
 
+- Added `date_ranges` method to `Polars`
 - Added support for keyword arguments to `agg` methods
 - Aliased `apply` to `map_rows` for `DataFrame`
 

ext/polars/src/functions/range.rs

@@ -39,3 +39,19 @@ pub fn date_range(
     let time_unit = time_unit.map(|x| x.0);
     dsl::date_range(start, end, every, closed, time_unit, time_zone).into()
 }
+
+pub fn date_ranges(
+    start: &RbExpr,
+    end: &RbExpr,
+    every: String,
+    closed: Wrap<ClosedWindow>,
+    time_unit: Option<Wrap<TimeUnit>>,
+    time_zone: Option<TimeZone>,
+) -> RbExpr {
+    let start = start.inner.clone();
+    let end = end.inner.clone();
+    let every = Duration::parse(&every);
+    let closed = closed.0;
+    let time_unit = time_unit.map(|x| x.0);
+    dsl::date_ranges(start, end, every, closed, time_unit, time_zone).into()
+}

ext/polars/src/lib.rs

@@ -611,6 +611,7 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     )?;
     class.define_singleton_method("collect_all", function!(functions::lazy::collect_all, 1))?;
     class.define_singleton_method("date_range", function!(functions::range::date_range, 6))?;
+    class.define_singleton_method("date_ranges", function!(functions::range::date_ranges, 6))?;
     class.define_singleton_method(
         "dtype_str_repr",
         function!(functions::misc::dtype_str_repr, 1),

lib/polars/functions/range/date_range.rb

@@ -99,5 +99,97 @@ def date_range(
 
       result
     end
+
+    # Create a column of date ranges.
+    #
+    # @param start [Object]
+    #   Lower bound of the date range.
+    # @param stop [Object]
+    #   Upper bound of the date range.
+    # @param interval [Object]
+    #   Interval of the range periods, specified using the Polars duration string language (see "Notes" section below).
+    # @param closed ["both", "left", "right", "none"]
+    #   Define which sides of the range are closed (inclusive).
+    # @param time_unit [nil, "ns", "us", "ms"]
+    #   Time unit of the resulting `Datetime` data type.
+    #   Only takes effect if the output column is of type `Datetime`.
+    # @param time_zone [String]
+    #   Time zone of the resulting `Datetime` data type.
+    #   Only takes effect if the output column is of type `Datetime`.
+    # @param eager [Boolean]
+    #   Evaluate immediately and return a `Series`.
+    #   If set to `false` (default), return an expression instead.
+    #
+    # @return [Object]
+    #
+    # @note
+    #   `interval` is created according to the following string language:
+    #
+    #   - 1ns   (1 nanosecond)
+    #   - 1us   (1 microsecond)
+    #   - 1ms   (1 millisecond)
+    #   - 1s    (1 second)
+    #   - 1m    (1 minute)
+    #   - 1h    (1 hour)
+    #   - 1d    (1 calendar day)
+    #   - 1w    (1 calendar week)
+    #   - 1mo   (1 calendar month)
+    #   - 1q    (1 calendar quarter)
+    #   - 1y    (1 calendar year)
+    #
+    #   Or combine them:
+    #   "3d12h4m25s" # 3 days, 12 hours, 4 minutes, and 25 seconds
+    #
+    #   By "calendar day", we mean the corresponding time on the next day (which may
+    #   not be 24 hours, due to daylight savings). Similarly for "calendar week",
+    #   "calendar month", "calendar quarter", and "calendar year".
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "start" => [Date.new(2022, 1, 1), Date.new(2022, 1, 2)],
+    #       "end" => Date.new(2022, 1, 3)
+    #     }
+    #   )
+    #   df.with_columns(date_range: Polars.date_ranges("start", "end"))
+    #   # =>
+    #   # shape: (2, 3)
+    #   # ┌────────────┬────────────┬───────────────────────────────────┐
+    #   # │ start      ┆ end        ┆ date_range                        │
+    #   # │ ---        ┆ ---        ┆ ---                               │
+    #   # │ date       ┆ date       ┆ list[date]                        │
+    #   # ╞════════════╪════════════╪═══════════════════════════════════╡
+    #   # │ 2022-01-01 ┆ 2022-01-03 ┆ [2022-01-01, 2022-01-02, 2022-01… │
+    #   # │ 2022-01-02 ┆ 2022-01-03 ┆ [2022-01-02, 2022-01-03]          │
+    #   # └────────────┴────────────┴───────────────────────────────────┘
+    def date_ranges(
+      start,
+      stop,
+      interval = "1d",
+      closed: "both",
+      time_unit: nil,
+      time_zone: nil,
+      eager: false
+    )
+      interval = Utils.parse_interval_argument(interval)
+      if time_unit.nil? && interval.include?("ns")
+        time_unit = "ns"
+      end
+
+      start_rbexpr = Utils.parse_as_expression(start)
+      end_rbexpr = Utils.parse_as_expression(stop)
+
+      result = Utils.wrap_expr(
+        Plr.date_ranges(
+          start_rbexpr, end_rbexpr, interval, closed, time_unit, time_zone
+        )
+      )
+
+      if eager
+        return F.select(result).to_series
+      end
+
+      result
+    end
   end
 end

lib/polars/utils.rb

@@ -418,5 +418,12 @@ def self._combine_predicates(predicates)
     def self.parse_when_inputs(*predicates, **constraints)
       parse_predicates_constraints_as_expression(*predicates, **constraints)
     end
+
+    def self.parse_interval_argument(interval)
+      if interval.include?(" ")
+        interval = interval.gsub(" ", "")
+      end
+      interval.downcase
+    end
   end
 end