diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a23967f3..85e5a5c4c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Added `Explorer.Series.rle_id/1`, which assigns a run-length encoding (RLE) ID + to each row, incrementing every time the value changes from one row to the next. + ## [v0.12.0] - 2026-07-05 ### Updated diff --git a/lib/explorer/backend/lazy_series.ex b/lib/explorer/backend/lazy_series.ex index 21a48c0b7..b5afc1d66 100644 --- a/lib/explorer/backend/lazy_series.ex +++ b/lib/explorer/backend/lazy_series.ex @@ -127,6 +127,7 @@ defmodule Explorer.Backend.LazySeries do all: 1, any: 1, row_index: 1, + rle_id: 1, # Strings contains: 2, re_contains: 2, @@ -1251,6 +1252,17 @@ defmodule Explorer.Backend.LazySeries do Backend.Series.new(data, {:u, 32}) end + @impl true + def rle_id(series) do + args = [lazy_series!(series)] + + if aggregations?(args), do: raise_agg_inside_window(:rle_id) + + data = new(:rle_id, args, {:u, 32}, false) + + Backend.Series.new(data, {:u, 32}) + end + @impl true def count_matches(series, substring) do data = new(:count_matches, [lazy_series!(series), substring], {:u, 32}) diff --git a/lib/explorer/backend/series.ex b/lib/explorer/backend/series.ex index 4871c43bd..92191d667 100644 --- a/lib/explorer/backend/series.ex +++ b/lib/explorer/backend/series.ex @@ -94,6 +94,7 @@ defmodule Explorer.Backend.Series do @callback all?(s) :: boolean() | lazy_s() @callback any?(s) :: boolean() | lazy_s() @callback row_index(s) :: s | lazy_s() + @callback rle_id(s) :: s | lazy_s() # Cumulative diff --git a/lib/explorer/polars_backend/expression.ex b/lib/explorer/polars_backend/expression.ex index ea094f527..6363b91a1 100644 --- a/lib/explorer/polars_backend/expression.ex +++ b/lib/explorer/polars_backend/expression.ex @@ -70,6 +70,7 @@ defmodule Explorer.PolarsBackend.Expression do quotient: 2, remainder: 2, reverse: 1, + rle_id: 1, floor: 1, ceil: 1, select: 3, diff --git a/lib/explorer/polars_backend/native.ex b/lib/explorer/polars_backend/native.ex index 9adbcfbaa..601c3e9de 100644 --- a/lib/explorer/polars_backend/native.ex +++ b/lib/explorer/polars_backend/native.ex @@ -426,6 +426,7 @@ defmodule Explorer.PolarsBackend.Native do def s_reverse(_s), do: err() def s_round(_s, _decimals), do: err() def s_row_index(_s), do: err() + def s_rle_id(_s), do: err() def s_floor(_s), do: err() def s_ceil(_s), do: err() def s_rstrip(_s, _string), do: err() diff --git a/lib/explorer/polars_backend/series.ex b/lib/explorer/polars_backend/series.ex index 72fd769c8..4f60a4064 100644 --- a/lib/explorer/polars_backend/series.ex +++ b/lib/explorer/polars_backend/series.ex @@ -253,6 +253,9 @@ defmodule Explorer.PolarsBackend.Series do @impl true def row_index(series), do: Shared.apply_series(series, :s_row_index) + @impl true + def rle_id(series), do: Shared.apply_series(series, :s_rle_id) + # Cumulative @impl true diff --git a/lib/explorer/series.ex b/lib/explorer/series.ex index 718f4915b..8b45d43e2 100644 --- a/lib/explorer/series.ex +++ b/lib/explorer/series.ex @@ -3101,6 +3101,61 @@ defmodule Explorer.Series do @spec row_index(Series.t()) :: Series.t() def row_index(%Series{} = series), do: apply_series(series, :row_index) + @doc """ + Returns a run-length encoding ID for each element in the series. + + The ID starts at 0 and is incremented every time the value changes + from one row to the next. All elements belonging to the same run + of consecutive equal values share the same ID. `nil` is treated as its + own value. + + Because this function depends on the order of the rows, it may be + necessary to sort your data frame prior to running this function to get the + results that you desire. When used inside `Explorer.DataFrame.mutate/2` on + a grouped data frame, the IDs restart at 0 for each group. + + ## Supported dtypes + + All except `:list` and `:struct`. + + ## Examples + + iex> s = Series.from_list(["a", "a", "b", "b", "b", "a"]) + iex> Series.rle_id(s) + #Explorer.Series< + Polars[6] + u32 [0, 0, 1, 1, 1, 2] + > + + `nil`s break a run and form their own run: + + iex> s = Series.from_list([1, 1, nil, nil, 1]) + iex> Series.rle_id(s) + #Explorer.Series< + Polars[5] + u32 [0, 0, 1, 1, 2] + > + + This is useful to group consecutive observations together, for example + to label each contiguous stretch of a sensor reading: + + iex> require Explorer.DataFrame, as: DF + iex> df = DF.new(state: ["on", "on", "off", "on"]) + iex> DF.mutate(df, run: rle_id(state)) + #Explorer.DataFrame< + Polars[4 x 2] + state string ["on", "on", "off", "on"] + run u32 [0, 0, 1, 2] + > + """ + @doc type: :window + @spec rle_id(series :: Series.t()) :: Series.t() + def rle_id(%Series{dtype: {composite, _} = dtype}) when K.in(composite, [:list, :struct]), + do: dtype_error("rle_id/1", dtype, Shared.dtypes() -- [{:list, :any}, {:struct, :any}]) + + def rle_id(%Series{} = series), + do: apply_series(series, :rle_id) + # Cumulative @doc """ diff --git a/native/explorer/Cargo.toml b/native/explorer/Cargo.toml index 4a4007f9d..9bb21b719 100644 --- a/native/explorer/Cargo.toml +++ b/native/explorer/Cargo.toml @@ -71,6 +71,7 @@ features = [ "range", "rank", "regex", + "rle", "rolling_window", "round_series", "rows", @@ -88,7 +89,7 @@ features = [ [dependencies.polars-ops] version = "0.52" -features = ["abs", "ewma", "cum_agg", "cov", "index_of"] +features = ["abs", "ewma", "cum_agg", "cov", "index_of", "rle"] [features] default = ["ndjson", "cloud", "nif_version_2_15"] diff --git a/native/explorer/src/expressions.rs b/native/explorer/src/expressions.rs index abf80da20..ec8579413 100644 --- a/native/explorer/src/expressions.rs +++ b/native/explorer/src/expressions.rs @@ -779,6 +779,13 @@ pub fn expr_reverse(expr: ExExpr) -> ExExpr { ExExpr::new(expr.reverse()) } +#[rustler::nif] +pub fn expr_rle_id(expr: ExExpr) -> ExExpr { + let expr = expr.clone_inner(); + + ExExpr::new(expr.rle_id()) +} + #[rustler::nif] pub fn expr_sort( expr: ExExpr, diff --git a/native/explorer/src/series.rs b/native/explorer/src/series.rs index bafacb909..ab87d0abc 100644 --- a/native/explorer/src/series.rs +++ b/native/explorer/src/series.rs @@ -1968,6 +1968,12 @@ pub fn s_row_index(series: ExSeries) -> Result { Ok(ExSeries::new(s)) } +#[rustler::nif(schedule = "DirtyCpu")] +pub fn s_rle_id(series: ExSeries) -> Result { + let column = polars_ops::prelude::rle_id(&series.clone_inner().into_column())?; + Ok(ExSeries::new(column.take_materialized_series())) +} + #[rustler::nif(schedule = "DirtyCpu")] pub fn s_count_matches( s1: ExSeries, diff --git a/test/explorer/data_frame_test.exs b/test/explorer/data_frame_test.exs index 03eb1ff38..aa14eca87 100644 --- a/test/explorer/data_frame_test.exs +++ b/test/explorer/data_frame_test.exs @@ -4853,6 +4853,53 @@ defmodule Explorer.DataFrameTest do end end + describe "rle_id/1" do + test "should assign an id to each run of consecutive equal values" do + df = + %{ + state: ["on", "on", "off", "on"] + } + |> DF.new() + |> DF.mutate(run: rle_id(state)) + + assert DF.to_columns(df, atom_keys: true) == %{ + state: ["on", "on", "off", "on"], + run: [0, 0, 1, 2] + } + + assert df.dtypes["run"] == {:u, 32} + end + + test "should restart at zero for each group" do + df = + %{ + sensor: ["a", "a", "a", "b", "b", "a", "b"], + state: ["on", "off", "off", "on", "on", "on", "off"] + } + |> DF.new() + |> DF.group_by("sensor") + |> DF.mutate(run: rle_id(state)) + |> DF.ungroup() + + assert DF.to_columns(df, atom_keys: true) == %{ + sensor: ["a", "a", "a", "b", "b", "a", "b"], + state: ["on", "off", "off", "on", "on", "on", "off"], + run: [0, 1, 1, 0, 0, 2, 1] + } + end + + test "should work when combined with other expressions" do + df = + %{ + a: [1, 1, 2] + } + |> DF.new() + |> DF.mutate(run: rle_id(a) + 1) + + assert DF.to_columns(df, atom_keys: true) == %{a: [1, 1, 2], run: [1, 1, 2]} + end + end + describe "row_index/1" do test "works as row_count(), including offset" do df = DF.new(a: [1, 3, 5], b: [2, 4, 6]) diff --git a/test/explorer/series_test.exs b/test/explorer/series_test.exs index 2b3efce07..9a64e8374 100644 --- a/test/explorer/series_test.exs +++ b/test/explorer/series_test.exs @@ -3,10 +3,11 @@ defmodule Explorer.SeriesTest do # Note that for the `{:list, _}` and `{:struct, _}` dtypes, we have a separated file for the tests. - alias Explorer.Series - import ExUnit.CaptureLog + alias Explorer.DataFrame + alias Explorer.Series + doctest Explorer.Series test "defines doc metadata" do @@ -6213,14 +6214,14 @@ defmodule Explorer.SeriesTest do category_label: "cat" ) - assert Explorer.DataFrame.names(df) == ["values", "bp", "cat"] + assert DataFrame.names(df) == ["values", "bp", "cat"] end test "cut/3 with include breaks" do series = Series.from_list([1.0, 2.0, 3.0]) df = Series.cut(series, [1.5, 2.5], include_breaks: true) - assert Explorer.DataFrame.to_columns(df, atom_keys: true) == %{ + assert DataFrame.to_columns(df, atom_keys: true) == %{ category: ["(-inf, 1.5]", "(1.5, 2.5]", "(2.5, inf]"], break_point: [1.5, 2.5, :infinity], values: [1.0, 2.0, 3.0] @@ -6260,7 +6261,7 @@ defmodule Explorer.SeriesTest do series = Enum.to_list(-5..3) |> Series.from_list() df = Series.qcut(series, [0.0, 0.25, 0.75], include_breaks: false) - assert Explorer.DataFrame.to_columns(df, atom_keys: true) == %{ + assert DataFrame.to_columns(df, atom_keys: true) == %{ category: [ "(-inf, -5]", "(-5, -3]", @@ -6685,7 +6686,7 @@ defmodule Explorer.SeriesTest do assert Series.dtype(df[:values]) == {:s, 64} assert Series.dtype(df[:counts]) == {:u, 32} - assert Explorer.DataFrame.to_columns(df, atom_keys: true) == %{ + assert DataFrame.to_columns(df, atom_keys: true) == %{ values: [1, 2, 3, 4, 5, 6], counts: [4, 2, 2, 1, 1, 1] } @@ -6699,7 +6700,7 @@ defmodule Explorer.SeriesTest do assert Series.dtype(df[:values]) == :string assert Series.dtype(df[:counts]) == {:u, 32} - assert Explorer.DataFrame.to_columns(df, atom_keys: true) == %{ + assert DataFrame.to_columns(df, atom_keys: true) == %{ values: ["c", "a", "b"], counts: [3, 2, 1] } @@ -6713,7 +6714,7 @@ defmodule Explorer.SeriesTest do assert Series.dtype(df[:values]) == {:list, {:s, 64}} assert Series.dtype(df[:counts]) == {:u, 32} - assert Explorer.DataFrame.to_columns(df, atom_keys: true) == %{ + assert DataFrame.to_columns(df, atom_keys: true) == %{ values: [[1, 2], [4, 1], [3, 1, 3], [5, 6]], counts: [2, 2, 1, 1] } @@ -6727,13 +6728,90 @@ defmodule Explorer.SeriesTest do assert Series.dtype(df[:values]) == {:list, :string} assert Series.dtype(df[:counts]) == {:u, 32} - assert Explorer.DataFrame.to_columns(df, atom_keys: true) == %{ + assert DataFrame.to_columns(df, atom_keys: true) == %{ values: [["c"], ["a"], ["a", "b"]], counts: [3, 1, 1] } end end + describe "rle_id/1" do + test "should work with basic data types" do + data_types = [ + {:string, ["a", "a", "b", "b", "b", "a"], [0, 0, 1, 1, 1, 2]}, + {:integer, [1, 1, 2, 3, 3], [0, 0, 1, 2, 2]}, + {:boolean, [true, true, false, true], [0, 0, 1, 2]}, + {:date, + [ + ~D[2024-01-01], + ~D[2024-01-01], + ~D[2024-06-13], + ~D[2024-01-01], + ~D[2024-01-01], + ~D[2024-01-01], + ~D[2025-01-01] + ], [0, 0, 1, 2, 2, 2, 3]}, + {{:naive_datetime, :millisecond}, + [ + ~N[2024-01-01 00:00:00.0], + ~N[2024-01-01 00:00:00.0], + ~N[2024-01-01 12:30:00.0], + ~N[2024-01-01 12:30:00.0], + ~N[2024-01-01 00:00:00.0] + ], [0, 0, 1, 1, 2]}, + {:category, ["a", "a", "b", nil, nil, "a"], [0, 0, 1, 2, 2, 3]}, + {{:decimal, 38, 2}, [1, 1, 2, 2, 1], [0, 0, 1, 1, 2]}, + {:u8, [1, 1, 2, 3, 3], [0, 0, 1, 2, 2]} + ] + + Enum.each(data_types, fn {dtype, input, expected_output} -> + s = Series.from_list(input, dtype: dtype) + ids = Series.rle_id(s) + + assert Series.dtype(ids) == {:u, 32} + assert Series.to_list(ids) == expected_output + end) + end + + test "should treat nil as its own value" do + s = Series.from_list([1, 1, nil, nil, 1]) + + assert Series.to_list(Series.rle_id(s)) == [0, 0, 1, 1, 2] + end + + test "should work with an empty series" do + s = Series.from_list([], dtype: :integer) + ids = Series.rle_id(s) + + assert Series.dtype(ids) == {:u, 32} + assert Series.to_list(ids) == [] + end + + test "should raise for unsupported dtypes" do + s = Series.from_list([[1, 2], [3]]) + + assert_raise ArgumentError, + ~r/Explorer\.Series\.rle_id\/1 not implemented for dtype \{:list, \{:s, 64\}\}/, + fn -> Series.rle_id(s) end + end + + test "should raise when given an aggregation" do + message = + "it's not possible to have an aggregation operation inside :rle_id, " <> + "which is a window function" + + assert_raise RuntimeError, message, fn -> + %{ + a: [1, 1, 2] + } + |> DataFrame.new() + |> DataFrame.summarise_with(fn ldf -> + [run: Series.rle_id(Series.max(ldf["a"]))] + end) + end + end + end + describe "peaks/1" do test "max with signed integers" do s = Series.from_list([1, 2, 4, 1, 4])