Return series rather than value

elixir-explorer · Dec 17, 2022 · 846b404 · 846b404
1 parent fbb5e3f
commit 846b404
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 10 deletions.
diff --git a/lib/explorer/backend/series.ex b/lib/explorer/backend/series.ex
@@ -53,7 +53,7 @@ defmodule Explorer.Backend.Series do
   @callback max(s) :: number() | Date.t() | NaiveDateTime.t() | lazy_s() | nil
   @callback mean(s) :: float() | lazy_s() | nil
   @callback median(s) :: float() | lazy_s() | nil
-  @callback mode(s) :: number() | lazy_s() | nil
+  @callback mode(s) :: s | lazy_s()
   @callback variance(s) :: float() | lazy_s() | nil
   @callback standard_deviation(s) :: float() | lazy_s() | nil
   @callback quantile(s, float()) :: number | Date.t() | NaiveDateTime.t() | lazy_s() | nil

diff --git a/lib/explorer/series.ex b/lib/explorer/series.ex
@@ -41,6 +41,7 @@ defmodule Explorer.Series do
   defguardp numeric_dtype?(dtype) when K.in(dtype, [:float, :integer])
   defguardp numeric_or_bool_dtype?(dtype) when K.in(dtype, [:float, :integer, :boolean])
   defguardp numeric_or_date_dtype?(dtype) when K.in(dtype, [:float, :integer, :date, :datetime])
+  defguardp numeric_or_string_dtype?(dtype) when K.in(dtype, [:float, :integer, :string])
 
   @impl true
   def fetch(series, idx) when is_integer(idx), do: {:ok, fetch!(series, idx)}
@@ -1253,14 +1254,38 @@ defmodule Explorer.Series do
 
       iex> s = Explorer.Series.from_list([1, 2, 2, 3])
       iex> Explorer.Series.mode(s)
-      2
+      #Explorer.Series<
+        Polars[1]
+        integer [2]
+      >
+
+      iex> s = Explorer.Series.from_list([1, 2, 3, 4])
+      iex> Explorer.Series.mode(s)
+      #Explorer.Series<
+        Polars[4]
+        integer [1, 2, 3, 4]
+      >
+
+      iex> s = Explorer.Series.from_list([1, 1, 2, 2])
+      iex> Explorer.Series.mode(s)
+      #Explorer.Series<
+        Polars[2]
+        integer [1, 2]
+      >
+
+      iex> s = Explorer.Series.from_list(["a", "b", "b", "c"])
+      iex> Explorer.Series.mode(s)
+      #Explorer.Series<
+        Polars[1]
+        string ["b"]
+      >
   """
   @doc type: :aggregation
-  @spec mode(series :: Series.t()) :: number() | nil
-  def mode(%Series{dtype: dtype} = series) when numeric_dtype?(dtype),
+  @spec mode(series :: Series.t()) :: Series.t() | nil
+  def mode(%Series{dtype: dtype} = series) when numeric_or_string_dtype?(dtype),
     do: Shared.apply_impl(series, :mode)
 
-  def mode(%Series{dtype: dtype}), do: dtype_error("mode/1", dtype, [:integer, :float])
+  def mode(%Series{dtype: dtype}), do: dtype_error("mode/1", dtype, [:integer, :float, :string])
 
   @doc """
   Gets the median value of the series.

diff --git a/native/explorer/src/series.rs b/native/explorer/src/series.rs
@@ -881,10 +881,7 @@ pub fn s_trim_trailing(data: ExSeries) -> Result<ExSeries, ExplorerError> {
 }
 
 #[rustler::nif(schedule = "DirtyCpu")]
-pub fn s_mode(env: Env, data: ExSeries) -> Result<Term, ExplorerError> {
+pub fn s_mode(data: ExSeries) -> Result<ExSeries, ExplorerError> {
     let s = &data.resource.0;
-    match s.dtype() {
-        DataType::UInt32 | DataType::Int64 | DataType::Float64 => Ok(s.mode()?.encode(env)),
-        dt => panic!("mode/1 not implemented for {:?}", dt),
-    }
+    Ok(ExSeries::new(s.mode()?))
 }