feat: implement pivot for the pyarrow backend

src/narwhals/_arrow/dataframe.py

@@ -19,6 +19,7 @@
     Implementation,
     check_column_names_are_unique,
     convert_str_slice_to_int_slice,
+    exclude_column_names,
     generate_temporary_column_name,
     not_implemented,
     parse_columns_to_drop,
@@ -55,6 +56,7 @@
     from narwhals.typing import (
         IntoSchema,
         JoinStrategy,
+        PivotAgg,
         SizedMultiIndexSelector,
         SizedMultiNameSelector,
         SizeUnit,
@@ -78,6 +80,43 @@
 
 MYPY: Final = False
 
+# narwhals aggregate -> pyarrow hash aggregate (len and None are handled separately).
+_PIVOT_TO_PYARROW: Final[Mapping[str, str]] = {
+    "min": "min",
+    "max": "max",
+    "first": "first",
+    "last": "last",
+    "sum": "sum",
+    "mean": "mean",
+    "median": "approximate_median",
+}
+
+
+def _pivot_format_value(value: Any, /, *, quote: bool) -> str:
+    # Format one `on` value like polars: null, true/false, strings quoted only inside {...}.
+    if value is None:
+        return "null"
+    if isinstance(value, bool):
+        return "true" if value else "false"
+    if quote and isinstance(value, str):
+        return f'"{value}"'
+    return str(value)
+
+
+def _pivot_column_name(
+    value: str, on_values: tuple[Any, ...], n_values: int, separator: str, /
+) -> str:
+    # One `on`: the value. Several: `{"a","b"}`, or `null` if any is null. Extra `values`
+    # prefix the name with `{value}{separator}`.
+    if len(on_values) == 1:
+        suffix = _pivot_format_value(on_values[0], quote=False)
+    elif any(v is None for v in on_values):
+        suffix = "null"
+    else:
+        body = ",".join(_pivot_format_value(v, quote=True) for v in on_values)
+        suffix = f"{{{body}}}"
+    return f"{value}{separator}{suffix}" if n_values > 1 else suffix
+
 
 class ArrowDataFrame(
     EagerDataFrame["ArrowSeries", "ArrowExpr", "pa.Table", "ChunkedArrayAny"]
@@ -827,4 +866,165 @@ def unpivot(
         )
         return self._with_native(concat_tables(tables, "permissive"))
 
-    pivot = not_implemented()
+    def _pivot_aggregate(
+        self, keys: list[str], values: list[str], aggregate_function: PivotAgg | None, /
+    ) -> tuple[pa.Table, dict[str, str]]:
+        # One row per index + on combination; pyarrow names each aggregate `{column}_{function}`.
+        native = self.native
+        ordered = False
+        if aggregate_function == "len":
+            specs: list[Any] = [
+                (value, "count", pc.CountOptions(mode="all")) for value in values
+            ]
+            names = {value: f"{value}_count" for value in values}
+        else:
+            if aggregate_function is None:
+                counts = native.group_by(keys).aggregate([([], "count_all")])
+                if counts.num_rows and pc.max(counts["count_all"]).as_py() > 1:
+                    msg = (
+                        "Found multiple elements for some combination of `index` and `on`.\n\n"
+                        "Please specify an `aggregate_function`."
+                    )
+                    raise ValueError(msg)
+                function = "first"
+            else:
+                function = _PIVOT_TO_PYARROW[aggregate_function]
+            if function in {"first", "last"}:
+                # polars keeps nulls in first/last (and the no-agg single value); pyarrow drops them.
+                keep_nulls = pc.ScalarAggregateOptions(skip_nulls=False)
+                specs = [(value, function, keep_nulls) for value in values]
+                ordered = True
+            else:
+                specs = [(value, function) for value in values]
+            names = {value: f"{value}_{function}" for value in values}
+        if ordered and self._backend_version < (14,):  # pragma: no cover
+            msg = (
+                "Using `first`/`last` or an unaggregated `pivot` with the pyarrow "
+                "backend requires 'pyarrow>=14.0.0'.\n\n"
+                "See https://github.com/apache/arrow/issues/36709"
+            )
+            raise NotImplementedError(msg)
+        # first/last are ordered aggregators; pyarrow computes them only single-threaded.
+        grouped = (
+            native.group_by(keys, use_threads=False) if ordered else native.group_by(keys)
+        )
+        return grouped.aggregate(specs), names
+
+    def _pivot_distinct(self, columns: list[str], /) -> pa.Table:
+        # Distinct `columns` combinations, ordered by first appearance (group_by loses order).
+        native = self.native
+        token = generate_temporary_column_name(8, native.column_names)
+        return (
+            native.select(columns)
+            .append_column(token, pa.array(range(native.num_rows)))
+            .group_by(columns)
+            .aggregate([(token, "min")])
+            .sort_by(f"{token}_min")
+            .select(columns)
+        )
+
+    def _pivot_combinations(
+        self, on: list[str], *, sort_columns: bool
+    ) -> list[tuple[Any, ...]]:
+        # One column per `on` combination that occurs (not the full product).
+        on_rows = self._pivot_distinct(on)
+        combinations = list(zip(*(on_rows[name].to_pylist() for name in on), strict=True))
+        if sort_columns:
+            # polars sorts ascending, nulls first.
+            combinations.sort(key=lambda combo: tuple((v is not None, v) for v in combo))
+        return combinations
+
+    def _pivot_reshape(
+        self,
+        grouped: pa.Table,
+        base: pa.Table,
+        combinations: list[tuple[Any, ...]],
+        *,
+        agg_names: dict[str, str],
+        index: list[str],
+        on: list[str],
+        values: list[str],
+        separator: str,
+    ) -> tuple[list[Any], list[str]]:
+        # Scatter each cell to its row and column. Tuples compare structurally, so nulls
+        # match here where a pyarrow join would not.
+        n_index = len(index)
+        height = base.num_rows
+        row_of = {
+            combo: position
+            for position, combo in enumerate(
+                zip(*(base[name].to_pylist() for name in index), strict=True)
+            )
+        }
+        block_of = {combo: position for position, combo in enumerate(combinations)}
+        keys = list(
+            zip(*(grouped[name].to_pylist() for name in (*index, *on)), strict=True)
+        )
+
+        arrays: list[Any] = [base.column(position) for position in range(n_index)]
+        output_names = list(index)
+        for value in values:
+            dtype = grouped.schema.field(agg_names[value]).type
+            blocks: list[list[Any]] = [[None] * height for _ in combinations]
+            for key, cell in zip(
+                keys, grouped[agg_names[value]].to_pylist(), strict=True
+            ):
+                blocks[block_of[key[n_index:]]][row_of[key[:n_index]]] = cell
+            for combination, block in zip(combinations, blocks, strict=True):
+                arrays.append(pa.array(block, type=dtype))
+                output_names.append(
+                    _pivot_column_name(value, combination, len(values), separator)
+                )
+        return arrays, output_names
+
+    @staticmethod
+    def _pivot_fill_empty(result: pa.Table, output_names: list[str], /) -> pa.Table:
+        for output in output_names:
+            position = result.column_names.index(output)
+            column = result.column(position)
+            filled = cast(
+                "ChunkedArrayAny", pc.fill_null(column, pa.scalar(0, column.type))
+            )
+            result = result.set_column(position, output, filled)
+        return result
+
+    def pivot(
+        self,
+        on: Sequence[str],
+        *,
+        index: Sequence[str] | None,
+        values: Sequence[str] | None,
+        aggregate_function: PivotAgg | None,
+        sort_columns: bool,
+        separator: str,
+    ) -> Self:
+        on = list(on)
+        index = index or (
+            exclude_column_names(self, {*on, *values})
+            if values
+            else exclude_column_names(self, on)
+        )
+        values = values or exclude_column_names(self, {*on, *index})
+        index, values = list(index), list(values)
+
+        grouped, agg_names = self._pivot_aggregate(
+            [*index, *on], values, aggregate_function
+        )
+        base = self._pivot_distinct(index)
+        combinations = self._pivot_combinations(on, sort_columns=sort_columns)
+        arrays, output_names = self._pivot_reshape(
+            grouped,
+            base,
+            combinations,
+            agg_names=agg_names,
+            index=index,
+            on=on,
+            values=values,
+            separator=separator,
+        )
+
+        result = pa.Table.from_arrays(arrays, names=output_names)
+        # Empty cells are null; polars fills them with 0 for sum and len, null otherwise.
+        if aggregate_function in {"sum", "len"}:
+            result = self._pivot_fill_empty(result, output_names[len(index) :])
+        return self._with_native(result, validate_column_names=False)

tests/frame/pivot_test.py

@@ -7,7 +7,12 @@
 
 import narwhals as nw
 from narwhals.exceptions import NarwhalsError
-from tests.utils import POLARS_VERSION, ConstructorEager, assert_equal_data
+from tests.utils import (
+    POLARS_VERSION,
+    PYARROW_VERSION,
+    ConstructorEager,
+    assert_equal_data,
+)
 
 data = {
     "ix": [1, 2, 1, 1, 2, 2],
@@ -120,7 +125,17 @@ def test_pivot(
     index: str | list[str],
     request: pytest.FixtureRequest,
 ) -> None:
-    if any(x in str(constructor_eager) for x in ("pyarrow_table", "modin")):
+    if "modin" in str(constructor_eager):
+        request.applymarker(pytest.mark.xfail)
+    if "pyarrow_table" in str(constructor_eager) and agg_func == "median":
+        # pyarrow only has an approximate hash median, like `group_by(...).agg(median())`.
+        request.applymarker(pytest.mark.xfail)
+    if (
+        "pyarrow_table" in str(constructor_eager)
+        and agg_func in {"first", "last"}
+        and PYARROW_VERSION < (14,)
+    ):
+        # first/last are ordered aggregators; pyarrow supports them only from 14.0.
         request.applymarker(pytest.mark.xfail)
     if "polars" in str(constructor_eager) and POLARS_VERSION < (1, 0):
         # not implemented
@@ -148,7 +163,14 @@ def test_pivot(
 def test_pivot_no_agg(
     request: Any, constructor_eager: ConstructorEager, data_: Any, context: Any
 ) -> None:
-    if any(x in str(constructor_eager) for x in ("pyarrow_table", "modin")):
+    if "modin" in str(constructor_eager):
+        request.applymarker(pytest.mark.xfail)
+    if (
+        "pyarrow_table" in str(constructor_eager)
+        and PYARROW_VERSION < (14,)
+        and isinstance(context, does_not_raise)
+    ):
+        # the no-dup path extracts the single value with `first` (pyarrow>=14).
         request.applymarker(pytest.mark.xfail)
     if "polars" in str(constructor_eager) and POLARS_VERSION < (1, 0):
         # not implemented
@@ -172,7 +194,7 @@ def test_pivot_sort_columns(
     sort_columns: Any,
     expected: list[str],
 ) -> None:
-    if any(x in str(constructor_eager) for x in ("pyarrow_table", "modin")):
+    if "modin" in str(constructor_eager):
         request.applymarker(pytest.mark.xfail)
     if "polars" in str(constructor_eager) and POLARS_VERSION < (1, 0):
         # not implemented
@@ -220,7 +242,7 @@ def test_pivot_sort_columns(
 def test_pivot_names_out(
     request: Any, constructor_eager: ConstructorEager, kwargs: Any, expected: list[str]
 ) -> None:
-    if any(x in str(constructor_eager) for x in ("pyarrow_table", "modin")):
+    if "modin" in str(constructor_eager):
         request.applymarker(pytest.mark.xfail)
     if "polars" in str(constructor_eager) and POLARS_VERSION < (1, 0):
         # not implemented
@@ -243,7 +265,7 @@ def test_pivot_no_index_no_values(constructor_eager: ConstructorEager) -> None:
 def test_pivot_no_index(
     constructor_eager: ConstructorEager, request: pytest.FixtureRequest
 ) -> None:
-    if any(x in str(constructor_eager) for x in ("pyarrow_table", "modin")):
+    if "modin" in str(constructor_eager):
         request.applymarker(pytest.mark.xfail)
     if "polars" in str(constructor_eager) and POLARS_VERSION < (1, 0):
         # not implemented
@@ -258,3 +280,49 @@ def test_pivot_no_index(
         "b": [None, 2.0, 4.0, None],
     }
     assert_equal_data(result, expected)
+
+
+@pytest.mark.parametrize(
+    ("data_", "expected"),
+    [
+        (
+            {"i": ["a", "b"], "o": [True, False], "v": [1, 2]},
+            {"i": ["a", "b"], "true": [1, 0], "false": [0, 2]},
+        ),
+        (
+            {"i": ["a", "b"], "o": ["x", None], "v": [1, 2]},
+            {"i": ["a", "b"], "x": [1, 0], "null": [0, 2]},
+        ),
+    ],
+)
+def test_pivot_on_bool_or_null(
+    constructor_eager: ConstructorEager,
+    data_: Any,
+    expected: dict[str, list[Any]],
+    request: pytest.FixtureRequest,
+) -> None:
+    # Only polars and pyarrow name boolean and null `on` values as `true`/`false`/`null`.
+    if not any(x in str(constructor_eager) for x in ("polars", "pyarrow_table")):
+        pytest.skip("boolean and null `on` naming is specific to polars and pyarrow")
+    if "polars" in str(constructor_eager) and POLARS_VERSION < (1, 0):
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(constructor_eager(data_), eager_only=True)
+    result = df.pivot(on="o", index="i", values="v", aggregate_function="sum")
+    assert_equal_data(result, expected)
+
+
+def test_pivot_multiple_on_with_null(
+    constructor_eager: ConstructorEager, request: pytest.FixtureRequest
+) -> None:
+    # polars collapses an `on` combination holding a null to `null`, not `{...}`.
+    if not any(x in str(constructor_eager) for x in ("polars", "pyarrow_table")):
+        pytest.skip("null `on` naming is specific to polars and pyarrow")
+    if "polars" in str(constructor_eager) and POLARS_VERSION < (1, 0):
+        request.applymarker(pytest.mark.xfail)
+
+    data_ = {"i": ["a", "b"], "o1": ["x", None], "o2": [1, 2], "v": [1, 2]}
+    df = nw.from_native(constructor_eager(data_), eager_only=True)
+    result = df.pivot(on=["o1", "o2"], index="i", values="v", aggregate_function="sum")
+    expected = {"i": ["a", "b"], '{"x",1}': [1, 0], "null": [0, 2]}
+    assert_equal_data(result, expected)

tests/modern_polars/pivot_test.py

@@ -11,7 +11,7 @@
 def test_pivot(
     constructor_eager: ConstructorEager, request: pytest.FixtureRequest
 ) -> None:
-    if any(x in str(constructor_eager) for x in ("pyarrow_table", "modin")):
+    if "modin" in str(constructor_eager):
         request.applymarker(pytest.mark.xfail)
     if "polars" in str(constructor_eager) and POLARS_VERSION < (1, 0):
         pytest.skip()