Add Series.sort_by and Series.sort_with

lib/explorer/series.ex

@@ -1631,6 +1631,28 @@ defmodule Explorer.Series do
     |> Explorer.DataFrame.pull(:series)
   end
 
+  @doc type: :element_wise
+  defmacro arrange(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 type: :element_wise
+  def arrange_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.
 

test/explorer/series_test.exs

@@ -2848,6 +2848,66 @@ defmodule Explorer.SeriesTest do
     end
   end
 
+  describe "arrange/2" do
+    test "ascending order (default)" do
+      require Explorer.Series
+
+      s1 = Series.from_list([1, 2, 3])
+      result = Series.arrange(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.arrange(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.arrange(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.arrange(s1, remainder(_, 3), nils: :first)
+      assert Series.to_list(result) == [nil, 3, 1, 2]
+    end
+  end
+
+  describe "arrange_with/2" do
+    test "ascending order (default)" do
+      s1 = Series.from_list([1, 2, 3])
+      result = Series.arrange_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.arrange_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.arrange_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.arrange_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()