enh: add LazyFrame.pivot

[MarcoGorelli]

This is something I'd quite like to see. In Polars too, but it's easier to implement here 😄

There's an example in the Polars docs about how to do a lazy pivot: https://docs.pola.rs/api/python/stable/reference/dataframe/api/polars.DataFrame.pivot.html

API-wise, we could have:

    def pivot(
        self: Self,
        on: str | list[str],
        *,
        column_names: list[str],
        index: str | list[str] | None = None,
        values: str | list[str] | None = None,
        aggregate_function: Literal[
            "min", "max", "sum", "mean", "median", "len"
        ]
        | None = None,
        separator: str = "_",
    ) -> Self:

This is very similar to DataFrame.pivot, but requires specifying the output column names explicitly so that the computation can stay completely lazy

[MarcoGorelli]

This is now available in Polars: pola-rs/polars#25016

[dangotbanned]

Just following up here with what I've got in #3373, but for DataFrame.pivot:

narwhals/narwhals/_plan/dataframe.py

Lines 473 to 484 in 859ae00

	# TODO @dangotbanned: (Follow-up) Accept selectors in `on`, `index`, `values`
	def pivot(
	self,
	on: OneOrIterable[str],
	on_columns: Sequence[str] | Series | Self | None = None,
	*,
	index: OneOrIterable[str] | None = None,
	values: OneOrIterable[str] | None = None,
	aggregate_function: PivotAgg | None = None,
	sort_columns: bool = False,
	separator: str = "_",
	) -> Self:

For pyarrow, we also need to provide key_names - which I'm assuming is the same as column_names in (#1901 (comment))?

(pola-rs/polars#25016) explains what on_columns is for and that's what I ended up implementing.
I'm not sure how that would translate to our lazy backends though 🤔

DuckDB seems to have two versions of PIVOT by the way https://duckdb.org/docs/stable/sql/statements/pivot

[MarcoGorelli]

I think we should be able to re- implement it as a group by

[dangotbanned]

I'm not sure how that would translate to our lazy backends though

Sorry for being unclear, what I meant by this was Series | DataFrame as input for lazy besides polars probably doesn't make much sense

I think we should be able to re- implement it as a group by

Yeah, IIRC it would be:

• Use on_columns to determine the unique values for on
• Iterate over them to build up predicates for each column
• Aggregate on each predicate

Which should probably be fine for lazy backends - but for pyarrow that kind of operation is super expensive (if you remember partition_by 😅)

[FBruzzesi]

One good use case: pymc-labs/pymc-marketing#2224 (comment)

[MarcoGorelli]

but for pyarrow that kind of operation is super expensive (if you remember partition_by 😅)

That's different, here we're itereating over values which end up becoming the result dataframe's columns, and in general we're assuming that it's OK to iterate over columns (but never over rows)

[dangotbanned]

That's different, here we're itereating over values which end up becoming the result dataframe's columns

I could be totally wrong here, but I was connecting the dots between:

• (feat(expr-ir): Support over(*partition_by) #3224 (comment))
• and how the filter predicates are built in
  • polars (rust)
  • polars (old python docs)

For pyarrow, I was able to create an ExecPlan that used filtered projections on known unique values.
(hint: this is what I thought was similar to lazy pivot)

But the execution engine didn't do any predicate pushdown - so it performed terribly 😭

Can this be translated to Acero?

Well, yeah kinda.
...

Show Version 2

Although this specific case has an equivalent result - performance is still an order of magnitude slower than native grouped aggregates 😢

<pyarrow.acero.Declaration>
ExecPlan with 10 nodes:
:ConsumingSinkNode{}
  :UnionNode{}
    :ScalarAggregateNode{aggregates=[
	min(a, {skip_nulls=true, min_count=1}),
	mean(b, {skip_nulls=true, min_count=1}),
]}
      :ProjectNode{projection=[a, "b": subtract(cast(b, {to_type=double, ...}), c)]}
        :FilterNode{filter=(nwbead2bb446a61b == 1)}
          :TableSourceNode{}
    :ScalarAggregateNode{aggregates=[
	min(a, {skip_nulls=true, min_count=1}),
	mean(b, {skip_nulls=true, min_count=1}),
]}
      :ProjectNode{projection=[a, "b": subtract(cast(b, {to_type=double, ...}), c)]}
        :FilterNode{filter=(nwbead2bb446a61b == 0)}
          :TableSourceNode{}

Anyways, the important part for this was (#1901 (comment))

Which should probably be fine for lazy backends -...

All lazy backends (I would hope) can do predicate pushdown - making this a non-issue 🙂