Add Series.sort_by and Series.sort_with (#762)

* pre-docs commit * switch to sort_by and sort_with * add docs * rename describe block 1 Co-authored-by: José Valim <[email protected]> * rename describe block 2 Co-authored-by: José Valim <[email protected]> --------- Co-authored-by: José Valim <[email protected]>
elixir-explorer · Dec 10, 2023 · a92cfee · a92cfee
1 parent 8d2473b
commit a92cfee
Show file tree

Hide file tree

Showing 2 changed files with 174 additions and 0 deletions.
diff --git a/lib/explorer/series.ex b/lib/explorer/series.ex
@@ -1631,6 +1631,117 @@ defmodule Explorer.Series do
     |> Explorer.DataFrame.pull(:series)
   end
 
+  @doc """
+  Sorts the series based on an expression.
+
+  > #### Notice {: .notice}
+  >
+  > This is a macro.
+  >
+  >   * You must `require Explorer.Series` before using it.
+  >   * You must use the special `_` syntax. See the moduledoc for details.
+
+  See `sort_with/3` for the callback-based version of this function.
+
+  ## Options
+
+    * `:direction` - `:asc` or `:desc`, meaning "ascending" or "descending", respectively.
+      By default it sorts in ascending order.
+
+    * `:nils` - `:first` or `:last`.
+      By default it is `:last` if direction is `:asc`, and `:first` otherwise.
+
+    * `:stable` - `true` or `false`.
+      Determines if the sorting is stable (ties are guaranteed to maintain their order) or not.
+      Unstable sorting may be more performant.
+      By default it is `true`.
+
+  ## Examples
+
+      iex> s = Explorer.Series.from_list([1, 2, 3])
+      iex> Explorer.Series.sort_by(s, remainder(_, 3))
+      #Explorer.Series<
+        Polars[3]
+        integer [3, 1, 2]
+      >
+
+      iex> s = Explorer.Series.from_list([1, 2, 3])
+      iex> Explorer.Series.sort_by(s, remainder(_, 3), direction: :desc)
+      #Explorer.Series<
+        Polars[3]
+        integer [2, 1, 3]
+      >
+
+      iex> s = Explorer.Series.from_list([1, nil, 2, 3])
+      iex> Explorer.Series.sort_by(s, -2 * _, nils: :first)
+      #Explorer.Series<
+        Polars[4]
+        integer [nil, 3, 2, 1]
+      >
+  """
+  @doc type: :shape
+  defmacro sort_by(series, query, opts \\ []) do
+    {direction, opts} = Keyword.pop(opts, :direction, :asc)
+
+    quote do
+      require Explorer.DataFrame
+
+      Explorer.DataFrame.new(_: unquote(series))
+      |> Explorer.DataFrame.arrange([{unquote(direction), unquote(query)}], unquote(opts))
+      |> Explorer.DataFrame.pull(:_)
+    end
+  end
+
+  @doc """
+  Sorts the series based on a callback that returns a lazy series.
+
+  See `sort_by/3` for the expression-based version of this function.
+
+  ## Options
+
+    * `:direction` - `:asc` or `:desc`, meaning "ascending" or "descending", respectively.
+      By default it sorts in ascending order.
+
+    * `:nils` - `:first` or `:last`.
+      By default it is `:last` if direction is `:asc`, and `:first` otherwise.
+
+    * `:stable` - `true` or `false`.
+      Determines if the sorting is stable (ties are guaranteed to maintain their order) or not.
+      Unstable sorting may be more performant.
+      By default it is `true`.
+
+  ## Examples
+
+      iex> s = Explorer.Series.from_list([1, 2, 3])
+      iex> Explorer.Series.sort_with(s, &Explorer.Series.remainder(&1, 3))
+      #Explorer.Series<
+        Polars[3]
+        integer [3, 1, 2]
+      >
+
+      iex> s = Explorer.Series.from_list([1, 2, 3])
+      iex> Explorer.Series.sort_with(s, &Explorer.Series.remainder(&1, 3), direction: :desc)
+      #Explorer.Series<
+        Polars[3]
+        integer [2, 1, 3]
+      >
+
+      iex> s = Explorer.Series.from_list([1, nil, 2, 3])
+      iex> Explorer.Series.sort_with(s, &Explorer.Series.multiply(-2, &1), nils: :first)
+      #Explorer.Series<
+        Polars[4]
+        integer [nil, 3, 2, 1]
+      >
+  """
+  @doc type: :shape
+  def sort_with(%Series{} = series, fun, opts \\ []) do
+    {direction, opts} = Keyword.pop(opts, :direction, :asc)
+
+    Explorer.DataFrame.new(series: series)
+    |> Explorer.DataFrame.arrange_with(&[{direction, fun.(&1[:series])}], opts)
+    |> Explorer.DataFrame.pull(:series)
+  end
+
   @doc """
   Filters a series with a mask.
 
@@ -4184,6 +4295,9 @@ defmodule Explorer.Series do
 
   Sorting is stable by default.
 
+  See `sort_by/3` for an expression-based sorting function.
+  See `sort_with/3` for a callback-based sorting function.
+
   ## Options
 
     * `:direction` - `:asc` or `:desc`, meaning "ascending" or "descending", respectively.

diff --git a/test/explorer/series_test.exs b/test/explorer/series_test.exs
@@ -2848,6 +2848,66 @@ defmodule Explorer.SeriesTest do
     end
   end
 
+  describe "sort_by/2" do
+    test "ascending order (default)" do
+      require Explorer.Series
+
+      s1 = Series.from_list([1, 2, 3])
+      result = Series.sort_by(s1, remainder(_, 3))
+      assert Series.to_list(result) == [3, 1, 2]
+    end
+
+    test "descending order" do
+      require Explorer.Series
+
+      s1 = Series.from_list([1, 2, 3])
+      result = Series.sort_by(s1, remainder(_, 3), direction: :desc)
+      assert Series.to_list(result) == [2, 1, 3]
+    end
+
+    test "nils last (default)" do
+      require Explorer.Series
+
+      s1 = Series.from_list([1, nil, 2, 3])
+      result = Series.sort_by(s1, remainder(_, 3))
+      assert Series.to_list(result) == [3, 1, 2, nil]
+    end
+
+    test "nils first" do
+      require Explorer.Series
+
+      s1 = Series.from_list([1, nil, 2, 3])
+      result = Series.sort_by(s1, remainder(_, 3), nils: :first)
+      assert Series.to_list(result) == [nil, 3, 1, 2]
+    end
+  end
+
+  describe "sort_with/2" do
+    test "ascending order (default)" do
+      s1 = Series.from_list([1, 2, 3])
+      result = Series.sort_with(s1, &Series.remainder(&1, 3))
+      assert Series.to_list(result) == [3, 1, 2]
+    end
+
+    test "descending order" do
+      s1 = Series.from_list([1, 2, 3])
+      result = Series.sort_with(s1, &Series.remainder(&1, 3), direction: :desc)
+      assert Series.to_list(result) == [2, 1, 3]
+    end
+
+    test "nils last (default)" do
+      s1 = Series.from_list([1, nil, 2, 3])
+      result = Series.sort_with(s1, &Series.remainder(&1, 3))
+      assert Series.to_list(result) == [3, 1, 2, nil]
+    end
+
+    test "nils first" do
+      s1 = Series.from_list([1, nil, 2, 3])
+      result = Series.sort_with(s1, &Series.remainder(&1, 3), nils: :first)
+      assert Series.to_list(result) == [nil, 3, 1, 2]
+    end
+  end
+
   describe "sample/2" do
     test "sample taking 10 elements" do
       s = 1..100 |> Enum.to_list() |> Series.from_list()