diff --git a/src/narwhals/_arrow/dataframe.py b/src/narwhals/_arrow/dataframe.py index ccbf0f17be..08e89dbb67 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,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) diff --git a/tests/frame/pivot_test.py b/tests/frame/pivot_test.py index 260006555e..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], @@ -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) 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()