[C++][Python] Incorrect results from hash_pivot_wider #48679

from __future__ import annotations


def pivot_wider(on: str, index: str, values: str) -> pa.Table:
    import pyarrow as pa
    import pyarrow.compute as pc  # noqa: F811
    from pyarrow.acero import AggregateNodeOptions, Declaration, TableSourceNodeOptions

    data = {
        "foo": ["A", "A", "B", "B", "C"],
        "bar": ["k", "l", "m", "n", "o"],
        "N1": [1, 2, 2, 4, 2],
        "N2": [1, 2, 2, 4, 2],
    }

    table = pa.Table.from_pydict(data)
    source = Declaration("table_source", options=TableSourceNodeOptions(table))
    key_names = table.column(on).unique().to_pylist()
    options = pc.PivotWiderOptions(key_names, unexpected_key_behavior="raise")
    aggregates = [([on, values], "hash_pivot_wider", options, values)]
    pivot = Declaration("aggregate", AggregateNodeOptions(aggregates, [index]))
    result = Declaration.from_sequence([source, pivot]).to_table()
    return result.flatten()

print(pivot_wider(on="bar", index="foo", values="N1").to_pydict())

print(pivot_wider(on="foo", index="bar", values="N1").to_pydict())

{'foo': ['A', 'B', 'C'],
 'N1.k': [1, None, None],
 'N1.l': [2, None, None],
 'N1.m': [None, 2, None],
 'N1.n': [None, None, 2],
 'N1.o': [None, 4, None]}

{'bar': ['k', 'l', 'm', 'o', 'n'],
 'N1.A': [1, 2, None, None, None],
 'N1.B': [None, None, 2, None, 4],
 'N1.C': [None, None, None, 2, None]}

import polars as pl

data = {
        "foo": ["A", "A", "B", "B", "C"],
        "bar": ["k", "l", "m", "n", "o"],
        "N1": [1, 2, 2, 4, 2],
        "N2": [1, 2, 2, 4, 2],
    }
df = pl.DataFrame(data)

print(df.pivot(on="bar", index="foo", values="N1").to_dict(as_series=False))

{'foo': ['A', 'B', 'C'],
 'k': [1, None, None],
 'l': [2, None, None],
 'm': [None, 2, None],
 'n': [None, 4, None],
 'o': [None, None, 2]}

print(df.pivot(on="foo", index="bar", values="N1").to_dict(as_series=False))

{'bar': ['k', 'l', 'm', 'n', 'o'],
 'A': [1, 2, None, None, None],
 'B': [None, None, 2, 4, None],
 'C': [None, None, None, None, 2]}