Implement substring/3 (#669)

lib/explorer/backend/lazy_series.ex

@@ -117,6 +117,7 @@ defmodule Explorer.Backend.LazySeries do
     trim: 2,
     upcase: 1,
     downcase: 1,
+    substring: 3,
     # Float round
     round: 2,
     floor: 1,
@@ -892,6 +893,13 @@ defmodule Explorer.Backend.LazySeries do
     Backend.Series.new(data, :string)
   end
 
+  @impl true
+  def substring(series, offset, length) do
+    data = new(:substring, [lazy_series!(series), offset, length])
+
+    Backend.Series.new(data, :string)
+  end
+
   @impl true
   def round(series, decimals) when is_integer(decimals) and decimals >= 0 do
     data = new(:round, [lazy_series!(series), decimals])

lib/explorer/backend/series.ex

@@ -247,6 +247,7 @@ defmodule Explorer.Backend.Series do
   @callback trim(s, String.t() | nil) :: s
   @callback trim_leading(s) :: s
   @callback trim_trailing(s) :: s
+  @callback substring(s, integer(), non_neg_integer() | nil) :: s
 
   # Date / DateTime
 

lib/explorer/polars_backend/expression.ex

@@ -125,7 +125,8 @@ defmodule Explorer.PolarsBackend.Expression do
     trim_leading: 1,
     trim_trailing: 1,
     downcase: 1,
-    upcase: 1
+    upcase: 1,
+    substring: 3
   ]
 
   @custom_expressions [

lib/explorer/polars_backend/native.ex

@@ -336,6 +336,7 @@ defmodule Explorer.PolarsBackend.Native do
   def s_unordered_distinct(_s), do: err()
   def s_frequencies(_s), do: err()
   def s_cut(_s, _bins, _labels, _break_point_label, _category_label), do: err()
+  def s_substring(_s, _offset, _length), do: err()
 
   def s_qcut(_s, _quantiles, _labels, _break_point_label, _category_label),
     do: err()

lib/explorer/polars_backend/series.ex

@@ -631,6 +631,10 @@ defmodule Explorer.PolarsBackend.Series do
   def trim_trailing(series),
     do: Shared.apply_series(series, :s_trim_trailing)
 
+  @impl true
+  def substring(series, offset, length),
+    do: Shared.apply_series(series, :s_substring, [offset, length])
+
   # Float round
   @impl true
   def round(series, decimals),

lib/explorer/series.ex

@@ -4395,6 +4395,69 @@ defmodule Explorer.Series do
 
   def trim_trailing(%Series{dtype: dtype}), do: dtype_error("trim_trailing/1", dtype, [:string])
 
+  @doc """
+  Returns a string sliced from the offset to the end of the string, supporting
+  negative indexing
+
+  ## Examples
+
+      iex> s = Explorer.Series.from_list(["earth", "mars", "neptune"])
+      iex> Explorer.Series.substring(s, -3)
+      #Explorer.Series<
+        Polars[3]
+        string ["rth", "ars", "une"]
+      >
+
+      iex> s = Explorer.Series.from_list(["earth", "mars", "neptune"])
+      iex> Explorer.Series.substring(s, 1)
+      #Explorer.Series<
+        Polars[3]
+        string ["arth", "ars", "eptune"]
+      >
+  """
+  @doc type: :string_wise
+  @spec substring(Series.t(), integer()) :: Series.t()
+  def substring(%Series{dtype: :string} = series, offset) when is_integer(offset),
+    do: apply_series(series, :substring, [offset, nil])
+
+  @doc """
+  Returns a string sliced from the offset to the length provided, supporting
+  negative indexing
+
+  ## Examples
+
+      iex> s = Explorer.Series.from_list(["earth", "mars", "neptune"])
+      iex> Explorer.Series.substring(s, -3, 2)
+      #Explorer.Series<
+        Polars[3]
+        string ["rt", "ar", "un"]
+      >
+
+      iex> s = Explorer.Series.from_list(["earth", "mars", "neptune"])
+      iex> Explorer.Series.substring(s, 1, 5)
+      #Explorer.Series<
+        Polars[3]
+        string ["arth", "ars", "eptun"]
+      >
+
+      iex> d = Explorer.Series.from_list(["こんにちは世界", "مرحبًا", "안녕하세요"])
+      iex> Explorer.Series.substring(d, 1, 3)
+      #Explorer.Series<
+        Polars[3]
+        string ["んにち", "رحب", "녕하세"]
+      >
+  """
+  @doc type: :string_wise
+  @spec substring(Series.t(), integer(), integer()) :: Series.t()
+  def substring(%Series{dtype: :string} = series, offset, length)
+      when is_integer(offset)
+      when is_integer(length)
+      when length >= 0,
+      do: apply_series(series, :substring, [offset, length])
+
+  def substring(%Series{dtype: dtype}, _offset, _length),
+    do: dtype_error("substring/3", dtype, [:string])
+
   # Float
 
   @doc """

native/explorer/src/expressions.rs

@@ -765,6 +765,12 @@ pub fn expr_trim_trailing(expr: ExExpr) -> ExExpr {
     ExExpr::new(expr.str().rstrip(None))
 }
 
+#[rustler::nif]
+pub fn expr_substring(expr: ExExpr, offset: i64, length: Option<u64>) -> ExExpr {
+    let expr = expr.clone_inner();
+    ExExpr::new(expr.str().str_slice(offset, length))
+}
+
 #[rustler::nif]
 pub fn expr_round(expr: ExExpr, decimals: u32) -> ExExpr {
     let expr = expr.clone_inner();

native/explorer/src/lib.rs

@@ -251,6 +251,7 @@ rustler::init!(
         expr_trim,
         expr_trim_leading,
         expr_trim_trailing,
+        expr_substring,
         // float round expressions
         expr_round,
         expr_floor,
@@ -403,6 +404,7 @@ rustler::init!(
         s_standard_deviation,
         s_tan,
         s_trim,
+        s_substring,
         s_subtract,
         s_sum,
         s_tail,

native/explorer/src/series.rs

@@ -1377,6 +1377,17 @@ pub fn s_trim_trailing(s1: ExSeries) -> Result<ExSeries, ExplorerError> {
     Ok(ExSeries::new(s1.utf8()?.replace(r#"[ \s]+$"#, "")?.into()))
 }
 
+#[rustler::nif(schedule = "DirtyCpu")]
+pub fn s_substring(
+    s1: ExSeries,
+    offset: i64,
+    length: Option<u64>,
+) -> Result<ExSeries, ExplorerError> {
+    Ok(ExSeries::new(
+        s1.utf8()?.str_slice(offset, length)?.into_series(),
+    ))
+}
+
 #[rustler::nif(schedule = "DirtyCpu")]
 pub fn s_round(s: ExSeries, decimals: u32) -> Result<ExSeries, ExplorerError> {
     Ok(ExSeries::new(s.round(decimals)?.into_series()))

test/explorer/data_frame_test.exs

@@ -1517,6 +1517,39 @@ defmodule Explorer.DataFrameTest do
       end
     end
 
+    test "slice strings" do
+      df =
+        DF.new(
+          a: ["_hello", "_world", "_foo", "_bar"],
+          b: ["venus", "earth", "mars", "jupiter"],
+          c: ["_foo", "_bar", "_baz", "_quox"],
+          d: ["_foo", "_bar", "_baz", "_quox"],
+          e: ["_foo", "_bar", "_baz", "_quox"]
+        )
+
+      df1 =
+        DF.mutate(df,
+          f: substring(a, 1),
+          g: substring(b, 2, 5),
+          h: substring(c, -3),
+          i: substring(d, 6, 10),
+          j: substring(e, -15, 2)
+        )
+
+      assert DF.to_columns(df1, atom_keys: true) == %{
+               a: ["_hello", "_world", "_foo", "_bar"],
+               b: ["venus", "earth", "mars", "jupiter"],
+               c: ["_foo", "_bar", "_baz", "_quox"],
+               d: ["_foo", "_bar", "_baz", "_quox"],
+               e: ["_foo", "_bar", "_baz", "_quox"],
+               f: ["hello", "world", "foo", "bar"],
+               g: ["nus", "rth", "rs", "piter"],
+               h: ["foo", "bar", "baz", "uox"],
+               i: ["", "", "", ""],
+               j: ["_f", "_b", "_b", "_q"]
+             }
+    end
+
     test "trim characters from string" do
       df =
         DF.new(

test/explorer/series_test.exs

@@ -3980,6 +3980,36 @@ defmodule Explorer.SeriesTest do
     end
   end
 
+  describe "string_slicing" do
+    test "string_slice/2 positive offset" do
+      series = Series.from_list(["earth", "mars", "neptune"])
+
+      assert Series.substring(series, 2) |> Series.to_list() == ["rth", "rs", "ptune"]
+      assert Series.substring(series, 20) |> Series.to_list() == ["", "", ""]
+    end
+
+    test "string_slice/2 negative offset" do
+      series = Series.from_list(["earth", "mars", "neptune"])
+
+      assert Series.substring(series, -3) |> Series.to_list() == ["rth", "ars", "une"]
+      assert Series.substring(series, -9) |> Series.to_list() == ["earth", "mars", "neptune"]
+    end
+
+    test "string_slice/3 positive offset" do
+      series = Series.from_list(["earth", "mars", "neptune"])
+
+      assert Series.substring(series, 2, 3) |> Series.to_list() == ["rth", "rs", "ptu"]
+      assert Series.substring(series, 12, 13) |> Series.to_list() == ["", "", ""]
+    end
+
+    test "string_slice/3 negative offset" do
+      series = Series.from_list(["earth", "mars", "neptune"])
+
+      assert Series.substring(series, -4, 4) |> Series.to_list() == ["arth", "mars", "tune"]
+      assert Series.substring(series, -20, 4) |> Series.to_list() == ["eart", "mars", "nept"]
+    end
+  end
+
   describe "strptime/2 and strftime/2" do
     test "parse datetime from string" do
       series = Series.from_list(["2023-01-05 12:34:56", "XYZ", nil])