Add day_of_year and week_of_year (#717)

lib/explorer/backend/lazy_series.ex

@@ -127,6 +127,8 @@ defmodule Explorer.Backend.LazySeries do
     ceil: 1,
     # Date functions
     day_of_week: 1,
+    day_of_year: 1,
+    week_of_year: 1,
     month: 1,
     year: 1,
     hour: 1,
@@ -563,6 +565,20 @@ defmodule Explorer.Backend.LazySeries do
     Backend.Series.new(data, :integer)
   end
 
+  @impl true
+  def day_of_year(%Series{} = s) do
+    data = new(:day_of_year, [lazy_series!(s)], :integer)
+
+    Backend.Series.new(data, :integer)
+  end
+
+  @impl true
+  def week_of_year(%Series{} = s) do
+    data = new(:week_of_year, [lazy_series!(s)], :integer)
+
+    Backend.Series.new(data, :integer)
+  end
+
   @impl true
   def month(%Series{} = s) do
     data = new(:month, [lazy_series!(s)], :integer)

lib/explorer/backend/series.ex

@@ -258,6 +258,8 @@ defmodule Explorer.Backend.Series do
   # Date / DateTime
 
   @callback day_of_week(s) :: s
+  @callback day_of_year(s) :: s
+  @callback week_of_year(s) :: s
   @callback month(s) :: s
   @callback year(s) :: s
   @callback hour(s) :: s

lib/explorer/polars_backend/expression.ex

@@ -22,6 +22,8 @@ defmodule Explorer.PolarsBackend.Expression do
     coalesce: 2,
     count: 1,
     day_of_week: 1,
+    day_of_year: 1,
+    week_of_year: 1,
     month: 1,
     year: 1,
     hour: 1,

lib/explorer/polars_backend/native.ex

@@ -353,6 +353,8 @@ defmodule Explorer.PolarsBackend.Native do
   def s_ewm_mean(_s, _alpha, _adjust, _min_periods, _ignore_nils), do: err()
   def s_in(_s, _other), do: err()
   def s_day_of_week(_s), do: err()
+  def s_day_of_year(_s), do: err()
+  def s_week_of_year(_s), do: err()
   def s_month(_s), do: err()
   def s_year(_s), do: err()
   def s_hour(_s), do: err()

lib/explorer/polars_backend/series.ex

@@ -656,6 +656,14 @@ defmodule Explorer.PolarsBackend.Series do
   def day_of_week(series),
     do: Shared.apply_series(series, :s_day_of_week)
 
+  @impl true
+  def day_of_year(series),
+    do: Shared.apply_series(series, :s_day_of_year)
+
+  @impl true
+  def week_of_year(series),
+    do: Shared.apply_series(series, :s_week_of_year)
+
   @impl true
   def month(series),
     do: Shared.apply_series(series, :s_month)

lib/explorer/series.ex

@@ -5033,6 +5033,75 @@ defmodule Explorer.Series do
   def day_of_week(%Series{dtype: dtype}),
     do: dtype_error("day_of_week/1", dtype, @date_or_datetime_dtypes)
 
+  @doc """
+  Returns the day-of-year number starting from 1.
+  
+  The return value ranges from 1 to 366 (the last day of year differs by years).
+
+  ## Examples
+
+      iex> s = Explorer.Series.from_list([~D[2023-01-01], ~D[2023-01-02], ~D[2023-02-01], nil])
+      iex> Explorer.Series.day_of_year(s)
+      #Explorer.Series<
+        Polars[4]
+        integer [1, 2, 32, nil]
+      >
+
+  It can also be called on a datetime series.
+
+      iex> s = Explorer.Series.from_list(
+      ...>   [~N[2023-01-01 00:00:00], ~N[2023-01-02 00:00:00], ~N[2023-02-01 23:59:59], nil]
+      ...> )
+      iex> Explorer.Series.day_of_year(s)
+      #Explorer.Series<
+        Polars[4]
+        integer [1, 2, 32, nil]
+      >
+  """
+  @doc type: :datetime_wise
+  @spec day_of_year(Series.t()) :: Series.t()
+  def day_of_year(%Series{dtype: dtype} = series) when K.in(dtype, @date_or_datetime_dtypes),
+    do: apply_series_list(:day_of_year, [series])
+
+  def day_of_year(%Series{dtype: dtype}),
+    do: dtype_error("day_of_year/1", dtype, @date_or_datetime_dtypes)
+
+  @doc """
+  Returns the week-of-year number.
+
+  The return value ranges from 1 to 53 (the last week of year differs by years).
+
+  Weeks start on Monday and end on Sunday. If the final week of a year does not end on Sunday, the
+  first days of the following year will have a week number of 52 (or 53 for a leap year).
+
+  ## Examples
+
+      iex> s = Explorer.Series.from_list([~D[2023-01-01], ~D[2023-01-02], ~D[2023-02-01], nil])
+      iex> Explorer.Series.week_of_year(s)
+      #Explorer.Series<
+        Polars[4]
+        integer [52, 1, 5, nil]
+      >
+
+  It can also be called on a datetime series.
+
+      iex> s = Explorer.Series.from_list(
+      ...>   [~N[2023-01-01 00:00:00], ~N[2023-01-02 00:00:00], ~N[2023-02-01 23:59:59], nil]
+      ...> )
+      iex> Explorer.Series.week_of_year(s)
+      #Explorer.Series<
+        Polars[4]
+        integer [52, 1, 5, nil]
+      >
+  """
+  @doc type: :datetime_wise
+  @spec week_of_year(Series.t()) :: Series.t()
+  def week_of_year(%Series{dtype: dtype} = series) when K.in(dtype, @date_or_datetime_dtypes),
+    do: apply_series_list(:week_of_year, [series])
+
+  def week_of_year(%Series{dtype: dtype}),
+    do: dtype_error("week_of_year/1", dtype, @date_or_datetime_dtypes)
+
   @deprecated "Use cast(:date) instead"
   def to_date(%Series{dtype: dtype} = series) when K.in(dtype, @datetime_dtypes),
     do: cast(series, :date)

native/explorer/src/expressions.rs

@@ -903,6 +903,20 @@ pub fn expr_day_of_week(expr: ExExpr) -> ExExpr {
     ExExpr::new(expr.dt().weekday().cast(DataType::Int64))
 }
 
+#[rustler::nif]
+pub fn expr_day_of_year(expr: ExExpr) -> ExExpr {
+    let expr = expr.clone_inner();
+
+    ExExpr::new(expr.dt().ordinal_day().cast(DataType::Int64))
+}
+
+#[rustler::nif]
+pub fn expr_week_of_year(expr: ExExpr) -> ExExpr {
+    let expr = expr.clone_inner();
+
+    ExExpr::new(expr.dt().week().cast(DataType::Int64))
+}
+
 #[rustler::nif]
 pub fn expr_month(expr: ExExpr) -> ExExpr {
     let expr = expr.clone_inner();

native/explorer/src/lib.rs

@@ -147,6 +147,8 @@ rustler::init!(
         expr_datetime,
         expr_duration,
         expr_day_of_week,
+        expr_day_of_year,
+        expr_week_of_year,
         expr_month,
         expr_year,
         expr_hour,
@@ -312,6 +314,8 @@ rustler::init!(
         s_cos,
         s_upcase,
         s_day_of_week,
+        s_day_of_year,
+        s_week_of_year,
         s_month,
         s_year,
         s_hour,

native/explorer/src/series.rs

@@ -1480,6 +1480,20 @@ pub fn s_day_of_week(s: ExSeries) -> Result<ExSeries, ExplorerError> {
     Ok(ExSeries::new(s1))
 }
 
+#[rustler::nif(schedule = "DirtyCpu")]
+pub fn s_day_of_year(s: ExSeries) -> Result<ExSeries, ExplorerError> {
+    let s1 = s.ordinal_day()?.cast(&DataType::Int64)?;
+
+    Ok(ExSeries::new(s1))
+}
+
+#[rustler::nif(schedule = "DirtyCpu")]
+pub fn s_week_of_year(s: ExSeries) -> Result<ExSeries, ExplorerError> {
+    let s1 = s.week()?.cast(&DataType::Int64)?;
+
+    Ok(ExSeries::new(s1))
+}
+
 #[rustler::nif(schedule = "DirtyCpu")]
 pub fn s_month(s: ExSeries) -> Result<ExSeries, ExplorerError> {
     let s1 = s.month()?.cast(&DataType::Int64)?;

test/explorer/data_frame_test.exs

@@ -1713,7 +1713,11 @@ defmodule Explorer.DataFrameTest do
           h: year(b),
           i: hour(b),
           j: minute(b),
-          k: second(b)
+          k: second(b),
+          l: day_of_year(a),
+          m: day_of_year(b),
+          n: week_of_year(a),
+          o: week_of_year(b)
         )
 
       assert DF.to_columns(df1, atom_keys: true) == %{
@@ -1732,7 +1736,11 @@ defmodule Explorer.DataFrameTest do
                h: [2023, 2022, 2021, nil],
                i: [1, 2, 3, nil],
                j: [1, 2, 3, nil],
-               k: [1, 2, 3, nil]
+               k: [1, 2, 3, nil],
+               l: [15, 47, 79, nil],
+               m: [15, 47, 79, nil],
+               n: [2, 7, 11, nil],
+               o: [2, 7, 11, nil]
              }
 
       assert df1.dtypes == %{
@@ -1746,7 +1754,11 @@ defmodule Explorer.DataFrameTest do
                "h" => :integer,
                "i" => :integer,
                "j" => :integer,
-               "k" => :integer
+               "k" => :integer,
+               "l" => :integer,
+               "m" => :integer,
+               "n" => :integer,
+               "o" => :integer
              }
     end