Plain values in grouped expressions are not broadcast since Polars 0.52

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

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

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