feat(expr-ir): Add BaseFrame.unpivot (#3368)

Author: dangotbanned

narwhals/_plan/arrow/dataframe.py

@@ -153,6 +153,27 @@ def _unique(
     unique = _unique
     unique_by = _unique
 
+    def unpivot(
+        self,
+        on: Sequence[str] | None,
+        index: Sequence[str] | None,
+        *,
+        variable_name: str = "variable",
+        value_name: str = "value",
+    ) -> Self:
+        n = len(self)
+        index = [] if index is None else list(index)
+        on_ = (c for c in self.columns if c not in index) if on is None else iter(on)
+        index_cols = self.native.select(index)
+        column = self.native.column
+        tables = (
+            index_cols.append_column(variable_name, fn.repeat(name, n)).append_column(
+                value_name, column(name)
+            )
+            for name in on_
+        )
+        return self._with_native(fn.concat_tables(tables, "permissive"))
+
     def with_row_index(self, name: str) -> Self:
         return self._with_native(self.native.add_column(0, name, fn.int_range(len(self))))
 

narwhals/_plan/compliant/dataframe.py

@@ -79,6 +79,14 @@ def unique_by(
         order_by: Sequence[str],
         keep: UniqueKeepStrategy = "any",
     ) -> Self: ...
+    def unpivot(
+        self,
+        on: Sequence[str] | None,
+        index: Sequence[str] | None,
+        *,
+        variable_name: str = "variable",
+        value_name: str = "value",
+    ) -> Self: ...
     def with_columns(self, irs: Seq[NamedIR]) -> Self: ...
     def with_row_index_by(
         self, name: str, order_by: Sequence[str], *, nulls_last: bool = False

narwhals/_plan/dataframe.py

@@ -159,6 +159,29 @@ def rename(self, mapping: Mapping[str, str]) -> Self:
     def collect_schema(self) -> Schema:
         return self.schema
 
+    def unpivot(
+        self,
+        on: OneOrIterable[ColumnNameOrSelector] | None = None,
+        *,
+        index: OneOrIterable[ColumnNameOrSelector] | None = None,
+        variable_name: str = "variable",
+        value_name: str = "value",
+    ) -> Self:
+        on_: Seq[str] | None = None
+        index_: Seq[str] | None = None
+        schema = self.schema
+        if on is not None:
+            s_irs = _parse.parse_into_seq_of_selector_ir(on)
+            on_ = expand_selector_irs_names(s_irs, schema=schema, require_any=True)
+        if index is not None:
+            s_irs = _parse.parse_into_seq_of_selector_ir(index)
+            index_ = expand_selector_irs_names(s_irs, schema=schema, require_any=True)
+        return self._with_compliant(
+            self._compliant.unpivot(
+                on_, index_, variable_name=variable_name, value_name=value_name
+            )
+        )
+
     def with_row_index(
         self, name: str = "index", *, order_by: OneOrIterable[ColumnNameOrSelector]
     ) -> Self:

tests/plan/unpivot_test.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Final
+
+import pytest
+
+import narwhals as nw
+from narwhals import _plan as nwp
+from narwhals._plan import selectors as ncs
+from tests.plan.utils import assert_equal_data, dataframe
+from tests.utils import PYARROW_VERSION
+
+if TYPE_CHECKING:
+    from narwhals._plan.typing import ColumnNameOrSelector, OneOrIterable
+    from tests.conftest import Data
+
+
+@pytest.fixture(scope="module")
+def data() -> Data:
+    return {"a": [7, 8, 9], "b": [1, 3, 5], "c": [2, 4, 6]}
+
+
+A: Final = [7, 8, 9]
+B: Final = [1, 3, 5]
+C: Final = [2, 4, 6]
+
+VAR = "variable"
+VALUE = "value"
+
+a = ncs.first()
+b = ncs.by_name("b")
+c = ncs.last()
+
+
+@pytest.mark.parametrize(
+    ("on", "index", "expected"),
+    [
+        ("b", [a], {"a": A, VAR: ["b", "b", "b"], VALUE: B}),
+        (
+            ["b", c],
+            a,
+            {"a": [*A, *A], VAR: ["b", "b", "b", "c", "c", "c"], VALUE: [*B, *C]},
+        ),
+        (
+            None,
+            ["a"],
+            {"a": [*A, *A], VAR: ["b", "b", "b", "c", "c", "c"], VALUE: [*B, *C]},
+        ),
+        ([b | c], None, {VAR: ["b", "b", "b", "c", "c", "c"], VALUE: [*B, *C]}),
+        (
+            None,
+            None,
+            {VAR: ["a", "a", "a", "b", "b", "b", "c", "c", "c"], VALUE: [*A, *B, *C]},
+        ),
+    ],
+)
+def test_unpivot(
+    data: Data,
+    on: OneOrIterable[ColumnNameOrSelector] | None,
+    index: OneOrIterable[ColumnNameOrSelector] | None,
+    expected: Data,
+) -> None:
+    sort_columns = [VAR] if index is None else [VAR, "a"]
+    result = dataframe(data).unpivot(on, index=index).sort(sort_columns)
+    assert_equal_data(result, expected)
+
+
+@pytest.mark.parametrize(
+    ("variable_name", "value_name"),
+    [
+        ("", "custom_value_name"),
+        ("custom_variable_name", ""),
+        ("custom_variable_name", "custom_value_name"),
+    ],
+)
+def test_unpivot_var_value_names(data: Data, variable_name: str, value_name: str) -> None:
+    result = dataframe(data).unpivot(
+        ~ncs.first(), index=["a"], variable_name=variable_name, value_name=value_name
+    )
+    assert result.collect_schema().names()[-2:] == [variable_name, value_name]
+
+
+def test_unpivot_default_var_value_names(data: Data) -> None:
+    result = dataframe(data).unpivot(nwp.nth(1, 2).meta.as_selector(), index=ncs.first())
+    assert result.collect_schema().names()[-2:] == [VAR, VALUE]
+
+
+@pytest.mark.xfail(PYARROW_VERSION < (14, 0, 0), reason="pyarrow<14")
+def test_unpivot_mixed_types() -> None:
+    df = dataframe({"idx": [0, 1], "a": [1, 2], "b": [1.5, 2.5]})
+    result = df.unpivot(["a", "b"], index="idx")
+    assert result.collect_schema().dtypes() == [nw.Int64(), nw.String(), nw.Float64()]