diff --git a/CHANGELOG.md b/CHANGELOG.md index 00b618295..6444aef58 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed + +- Plain values in expressions on grouped dataframes, such as the branches of `if` + or the `1` in `sum(a) + 1`, are broadcast to every group again. Since v0.12.0 they were + given to Polars as series literals of size 1, which Polars does not broadcast over groups, + so `mutate/2` raised on mismatched lengths and `summarise/2` returned lists. + ## [v0.13.0] - 2026-09-28 ### Added diff --git a/lib/explorer/polars_backend/expression.ex b/lib/explorer/polars_backend/expression.ex index 6363b91a1..ada5ecfaf 100644 --- a/lib/explorer/polars_backend/expression.ex +++ b/lib/explorer/polars_backend/expression.ex @@ -205,6 +205,11 @@ defmodule Explorer.PolarsBackend.Expression do Native.expr_fill_missing_with_strategy(expr, Atom.to_string(strategy)) end + # A series of size 1 is how Explorer passes a plain value, such as the `1` in `a + 1`. + def to_expr(%LazySeries{op: :from_list, args: [[value], dtype]}) do + scalar_expr(value, dtype) + end + def to_expr(%LazySeries{op: :from_list, args: [list, dtype]}) do series = Explorer.PolarsBackend.Shared.from_list(list, dtype) Native.expr_series(series) @@ -351,8 +356,12 @@ defmodule Explorer.PolarsBackend.Expression do def to_expr(%DateTime{} = datetime), do: Native.expr_datetime(datetime) def to_expr(%Explorer.Duration{} = duration), do: Native.expr_duration(duration) - def to_expr(%Explorer.Series{data: %PolarsSeries{} = polars_series}), - do: Native.expr_series(polars_series) + def to_expr(%Explorer.Series{data: %PolarsSeries{} = polars_series, dtype: dtype} = series) do + case Explorer.Series.size(series) do + 1 -> series |> Explorer.Series.at(0) |> scalar_expr(dtype) + _size -> Native.expr_series(polars_series) + end + end def to_expr(%_{} = struct), do: raise("unsupported struct in expression: #{inspect(struct)}") @@ -378,4 +387,36 @@ defmodule Explorer.PolarsBackend.Expression do defp dtype(%LazySeries{dtype: dtype}), do: dtype defp dtype(%Explorer.Series{dtype: dtype}), do: dtype + + # Polars broadcasts a literal over the groups of a grouped dataframe, as one value + # per group, only when the literal is a scalar. A series literal is not a scalar, + # even when its size is 1. Values that have no scalar literal are lowered as the + # first value of a series literal, which Polars also takes as one value per group. + defp scalar_expr(value, dtype) do + if scalar_literal?(value) do + expr = to_expr(value) + + if Explorer.Shared.dtype_from_list!([value]) == dtype, + do: expr, + else: Native.expr_cast(expr, dtype) + else + [value] + |> Explorer.PolarsBackend.Shared.from_list(dtype) + |> Native.expr_series() + |> Native.expr_first() + end + end + + @min_s64 -Bitwise.bsl(1, 63) + @max_s64 Bitwise.bsl(1, 63) - 1 + + defp scalar_literal?(value) when is_integer(value), do: value in @min_s64..@max_s64 + defp scalar_literal?(value) when is_float(value), do: true + defp scalar_literal?(value) when is_binary(value), do: true + defp scalar_literal?(value) when is_atom(value), do: true + defp scalar_literal?(%Date{}), do: true + defp scalar_literal?(%NaiveDateTime{}), do: true + defp scalar_literal?(%DateTime{}), do: true + defp scalar_literal?(%Explorer.Duration{}), do: true + defp scalar_literal?(_value), do: false end diff --git a/test/explorer/data_frame/grouped_test.exs b/test/explorer/data_frame/grouped_test.exs index 118b80513..dd7a52d61 100644 --- a/test/explorer/data_frame/grouped_test.exs +++ b/test/explorer/data_frame/grouped_test.exs @@ -309,6 +309,60 @@ defmodule Explorer.DataFrame.GroupedTest do end end + describe "summarise/2 with plain values" do + test "returns a scalar per group for arithmetic and comparisons with plain values" do + df = DF.new(a: [1, 2, 4], b: ["x", "y", "z"], c: [1, 1, 2]) + + df1 = + df + |> DF.group_by(:c, stable: true) + |> DF.summarise( + plus: sum(a) + 1, + times: mean(a) * 2, + greater: sum(a) > 3, + equal: first(b) == "x" + ) + + assert DF.dtypes(df1) == %{ + "c" => {:s, 64}, + "plus" => {:s, 64}, + "times" => {:f, 64}, + "greater" => :boolean, + "equal" => :boolean + } + + assert DF.to_columns(df1, atom_keys: true) == %{ + c: [1, 2], + plus: [4, 5], + times: [3.0, 8.0], + greater: [false, true], + equal: [true, false] + } + end + end + + describe "summarise/2 with if/2" do + test "returns a scalar per group when the branches are aggregations and scalars" do + df = DF.new(a: [1, 2, 3], w: [1, 0, 0], c: [1, 1, 2]) + + df1 = + df + |> DF.group_by(:c, stable: true) + |> DF.summarise( + total: if(sum(w) == 0, do: nil, else: sum(w)), + weighted: sum(a * w) / if(sum(w) == 0, do: nil, else: sum(w)) + ) + + assert DF.dtypes(df1) == %{"c" => {:s, 64}, "total" => {:s, 64}, "weighted" => {:f, 64}} + + assert DF.to_columns(df1, atom_keys: true) == %{ + c: [1, 2], + total: [1, nil], + weighted: [1.0, nil] + } + end + end + describe "summarise_with/2" do test "with one group and one column with aggregations", %{df: df} do df1 = @@ -562,6 +616,42 @@ defmodule Explorer.DataFrame.GroupedTest do assert df2.dtypes == %{"a" => {:s, 64}, "b" => :string, "c" => {:s, 64}, "d" => {:f, 64}} assert df2.groups.columns == ["c"] end + + test "broadcasts scalar branches of if/2 when there is a group" do + df = DF.new(a: [1, 2, 4], c: [1, 1, 2]) + + df1 = DF.group_by(df, :c) + + df2 = + DF.mutate(df1, + d: if(a > 1, do: 1, else: 0), + e: if(a > 1, do: "big", else: "small"), + f: if(sum(a) > 3, do: sum(a), else: nil) + ) + + assert DF.to_columns(df2, atom_keys: true) == %{ + a: [1, 2, 4], + c: [1, 1, 2], + d: [0, 1, 1], + e: ["small", "big", "big"], + f: [nil, nil, 4] + } + + assert df2.groups.columns == ["c"] + end + + test "broadcasts a series of size 1 given to coalesce/2 when there is a group" do + df = DF.new(a: [1, 2, 3], c: [1, 1, 2]) + + df1 = DF.group_by(df, :c) + df2 = DF.mutate(df1, d: coalesce(shift(a, 1), Series.from_list([0]))) + + assert DF.to_columns(df2, atom_keys: true) == %{ + a: [1, 2, 3], + c: [1, 1, 2], + d: [0, 1, 0] + } + end end describe "mutate_with/2" do