From 7ac42f929ff01d239d65ca7636a6d9c73f52209c Mon Sep 17 00:00:00 2001 From: mkzung <103102868+mkzung@users.noreply.github.com> Date: Fri, 24 Jul 2026 22:54:10 +0500 Subject: [PATCH 1/2] feat: implement pivot for the pyarrow backend Closes #2179. ArrowDataFrame.pivot was not_implemented. pyarrow has no native pivot, so this groups by index + on, then scatters each aggregated cell into its output row and column, matched on Python tuples so it stays correct when index or on hold nulls (a pyarrow join cannot match null keys). It follows polars for column names, order, null handling, and the empty-cell fill. median maps to pyarrow's approximate hash median, so its exact-value test stays xfail for pyarrow (same as group_by median). --- src/narwhals/_arrow/dataframe.py | 184 +++++++++++++++++++++++++++++- tests/frame/pivot_test.py | 59 +++++++++- tests/modern_polars/pivot_test.py | 2 +- 3 files changed, 238 insertions(+), 7 deletions(-) diff --git a/src/narwhals/_arrow/dataframe.py b/src/narwhals/_arrow/dataframe.py index ccbf0f17be..c9771109e6 100644 --- a/src/narwhals/_arrow/dataframe.py +++ b/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,147 @@ 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 + 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, use_threads=False).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] + else: + specs = [(value, function) for value in values] + names = {value: f"{value}_{function}" for value in values} + return native.group_by(keys, use_threads=False).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, use_threads=False) + .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, + agg_names: dict[str, str], + base: pa.Table, + index: list[str], + on: list[str], + values: list[str], + combinations: list[tuple[Any, ...]], + 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, agg_names, base, index, on, values, combinations, 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) diff --git a/tests/frame/pivot_test.py b/tests/frame/pivot_test.py index 260006555e..722a47994f 100644 --- a/tests/frame/pivot_test.py +++ b/tests/frame/pivot_test.py @@ -120,7 +120,10 @@ 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 "polars" in str(constructor_eager) and POLARS_VERSION < (1, 0): # not implemented @@ -148,7 +151,7 @@ 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 "polars" in str(constructor_eager) and POLARS_VERSION < (1, 0): # not implemented @@ -172,7 +175,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 +223,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 +246,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 +261,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) diff --git a/tests/modern_polars/pivot_test.py b/tests/modern_polars/pivot_test.py index b6e6b16dde..bfe39441bc 100644 --- a/tests/modern_polars/pivot_test.py +++ b/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() From e36a699c2cca61a693a5a28a2775fa7ef53c572f Mon Sep 17 00:00:00 2001 From: mkzung <103102868+mkzung@users.noreply.github.com> Date: Tue, 28 Jul 2026 14:36:43 +0500 Subject: [PATCH 2/2] fix: pyarrow pivot first/last on pyarrow<14 and PLR0917 first/last are ordered aggregators: pyarrow computes them only in single-threaded execution (use_threads=False), which Table.group_by does not accept before 14.0. Pass the flag only when it is needed and available, mirroring _configure_grouped in group_by.py, and xfail the ordered pivot cases on pyarrow<14. Make _pivot_reshape's schema arguments keyword-only so it stays under the positional-argument limit (PLR0917). --- src/narwhals/_arrow/dataframe.py | 36 ++++++++++++++++++++++++-------- tests/frame/pivot_test.py | 21 ++++++++++++++++++- 2 files changed, 47 insertions(+), 10 deletions(-) diff --git a/src/narwhals/_arrow/dataframe.py b/src/narwhals/_arrow/dataframe.py index c9771109e6..08e89dbb67 100644 --- a/src/narwhals/_arrow/dataframe.py +++ b/src/narwhals/_arrow/dataframe.py @@ -871,6 +871,7 @@ def _pivot_aggregate( ) -> 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 @@ -878,9 +879,7 @@ def _pivot_aggregate( names = {value: f"{value}_count" for value in values} else: if aggregate_function is None: - counts = native.group_by(keys, use_threads=False).aggregate( - [([], "count_all")] - ) + 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" @@ -894,10 +893,22 @@ def _pivot_aggregate( # 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} - return native.group_by(keys, use_threads=False).aggregate(specs), names + 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). @@ -906,7 +917,7 @@ def _pivot_distinct(self, columns: list[str], /) -> pa.Table: return ( native.select(columns) .append_column(token, pa.array(range(native.num_rows))) - .group_by(columns, use_threads=False) + .group_by(columns) .aggregate([(token, "min")]) .sort_by(f"{token}_min") .select(columns) @@ -926,14 +937,14 @@ def _pivot_combinations( def _pivot_reshape( self, grouped: pa.Table, - agg_names: dict[str, str], base: pa.Table, + combinations: list[tuple[Any, ...]], + *, + agg_names: dict[str, str], index: list[str], on: list[str], values: list[str], - combinations: list[tuple[Any, ...]], 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. @@ -1002,7 +1013,14 @@ def pivot( base = self._pivot_distinct(index) combinations = self._pivot_combinations(on, sort_columns=sort_columns) arrays, output_names = self._pivot_reshape( - grouped, agg_names, base, index, on, values, combinations, separator + grouped, + base, + combinations, + agg_names=agg_names, + index=index, + on=on, + values=values, + separator=separator, ) result = pa.Table.from_arrays(arrays, names=output_names) diff --git a/tests/frame/pivot_test.py b/tests/frame/pivot_test.py index 722a47994f..c037e5047d 100644 --- a/tests/frame/pivot_test.py +++ b/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], @@ -125,6 +130,13 @@ def test_pivot( 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 request.applymarker(pytest.mark.xfail) @@ -153,6 +165,13 @@ def test_pivot_no_agg( ) -> None: 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 request.applymarker(pytest.mark.xfail)