hzmsh opened a new pull request, #59538: URL: https://github.com/apache/spark/pull/59538
### What changes were proposed in this pull request? Fixes [SPARK-37553](https://issues.apache.org/jira/browse/SPARK-37553). Fix: - `DataFrame.pivot_table` no longer splits the Spark column names on `_` to get the pivot value and the value column back. Instead it pairs each output column with its value column by position. - This works because Spark's `PivotTransformer` emits the columns in a fixed order: for each pivot value, one column per aggregation, in the order the aggregations were passed ([L141-L150](https://github.com/apache/spark/blob/d18b51873dbcb8f18677b9a5860b3131f16b1d8e/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/PivotTransformer.scala#L141-L150); the other branch at [L153-L183](https://github.com/apache/spark/blob/d18b51873dbcb8f18677b9a5860b3131f16b1d8e/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/PivotTransformer.scala#L153-L183) uses the same order; names are built in [`outputName`](https://github.com/apache/spark/blob/d18b51873dbcb8f18677b9a5860b3131f16b1d8e/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/PivotTransformer.scala#L188-L208)). The pivot value is the name with the `_<value column>` suffix removed. - The change is in `python/pyspark/pandas/frame.py` only. Spark Connect builds the same `Pivot` plan, so both classic and Connect are fixed. - An earlier fix, [#34812](https://github.com/apache/spark/pull/34812) by the reporter of SPARK-37553, got one round of review in December 2021 and was then closed by the stale bot in April 2022. This PR follows [the review suggestion there](https://github.com/apache/spark/pull/34812#discussion_r762637825) to work from the original column names, and is smaller (+14/-3 in `frame.py`, against +28/-10 in that PR). It fixes the same cases. ### Why are the changes needed? `pivot_table` with more than one `values` column breaks whenever a value column name or a pivot value contains `_`. Column names like `unit_price` are common. ```python import pandas as pd import pyspark.pandas as ps pdf = pd.DataFrame({"a": [4, 2, 3], "b_b": [1, 2, 2], "e": [10, 20, 20], "c": [1, 2, 9]}) ps.from_pandas(pdf).pivot_table(index=["c"], columns="a", values=["b_b", "e"], aggfunc="sum") # master and 4.2.0: KeyError: 'b' # expected: the same as pdf.pivot_table(...), columns (b_b, 2), (b_b, 3), ... (e, 4) ``` When it started: - This code has been there since pandas API on Spark was ported from Koalas in Spark 3.2.0 ([SPARK-34890](https://issues.apache.org/jira/browse/SPARK-34890)). Users reported it in 2021, and it still fails on 4.2.0 and master. Reports from users: - [SPARK-37553](https://issues.apache.org/jira/browse/SPARK-37553) (Spark 3.2.0): the same `KeyError` with `values=['b_b', 'e']`, and with pivot values like `'4_4'`. - [databricks/koalas#2189](https://github.com/databricks/koalas/issues/2189): `values=['baz_d', 'ar_a']` on Databricks, `KeyError: 'ar'`. What goes wrong: - Spark names each output column `<pivot value>_<value column>`, and `DataFrame.pivot_table` splits the name on `_` to get the two parts back ([frame.py L7157-L7166](https://github.com/apache/spark/blob/d18b51873dbcb8f18677b9a5860b3131f16b1d8e/python/pyspark/pandas/frame.py#L7157-L7166)). `name.split("_")` turns `"4_b_b"` into `["4", "b", "b"]`. The code takes `"4"` as the pivot value and `"b"` as the value column, and there is no column `"b"`. A pivot value such as `"a_1"` fails the same way (`KeyError: '1'`). - If that piece happens to be the name of another column, there is no `KeyError`. With a pivot value `"x_e"` and `values=["b", "e"]`, both `x_e_b` and `x_e_e` get the label `(e, x)`, and the call fails with an `AssertionError` (`(4, 3)` for pivot values `"x_e"` and `"y"`). With `values=["b_b", "e"]` on a DataFrame that also has a column `b`, the `b_b` numbers come back under the label `b`, with no error. ### Does this PR introduce _any_ user-facing change? Yes, relative to 4.2.0 as well as master. Behavior change: - `pivot_table` with several `values` columns now works when a value column name or a pivot value contains `_`, and gives the same result as the same call on names without `_` (in the tested cases, the same as pandas). Before, it raised `KeyError` or `AssertionError`, or, when the DataFrame also had a column named like the first part (`b` next to `b_b`), put the `b_b` numbers under the label `b`. - `DataFrame.pivot` with a list of `values`, and `DataFrame.unstack()` on a DataFrame with a MultiIndex and two or more columns, call `pivot_table` the same way. With `_` in a name they raised the same `KeyError` in 4.2.0, and they work now. - When no value column name or pivot value contains `_`, the result is the same as before (same column order and labels). Not covered: - Float pivot values such as `1.5` with several `values` still fail with `UNRESOLVED_COLUMN`, because of the `.` in the generated column name. Same before and after this PR. - Two different (pivot value, value column) pairs that produce the same Spark column name, for example `x` + `1_b` and `x_1` + `b`, still fail with `AMBIGUOUS_REFERENCE`. Same before and after this PR. - Differences from pandas that do not come from `_` are unchanged: boolean pivot values still fail with `ValueError: Level values must be unique` (with `b_b` in `values` this used to be `KeyError: 'b'`), and a null pivot value still adds a `null` column that pandas drops. ### How was this patch tested? Tests: - Added two cases to `PivotTableMixin.test_pivot_table`, compared against pandas: - `values=["b_b", "e"]` with pivot values `"x_1"`, `"y"`, `"y_2"` and `aggfunc="sum"`. The data also has a column `b` that is not in `values`, so the test also covers the silent case where a wrong split finds an existing column and returns wrong labels instead of raising. - The same data with `values=["e", "b_b"]` and `aggfunc={"b_b": "max", "e": "sum"}`, so the dict order differs from `values` and pairing by the wrong list would be caught. - Without the fix, the new test fails in both classic and Connect with `KeyError: '1'`. - The repro above was also run against PySpark 4.2.0 from PyPI (`KeyError: 'b'`). - Commands: ``` python/run-tests --python-executables=python3 --testnames 'pyspark.pandas.tests.computation.test_pivot_table,pyspark.pandas.tests.connect.computation.test_parity_pivot_table,pyspark.pandas.tests.computation.test_pivot_table_adv,pyspark.pandas.tests.connect.computation.test_parity_pivot_table_adv,pyspark.pandas.tests.computation.test_pivot_table_multi_idx,pyspark.pandas.tests.connect.computation.test_parity_pivot_table_multi_idx,pyspark.pandas.tests.computation.test_pivot_table_multi_idx_adv,pyspark.pandas.tests.connect.computation.test_parity_pivot_table_multi_idx_adv,pyspark.pandas.tests.computation.test_pivot,pyspark.pandas.tests.connect.computation.test_parity_pivot,pyspark.pandas.tests.frame.test_reshaping,pyspark.pandas.tests.connect.frame.test_parity_reshaping' dev/lint-python ``` ### Was this patch authored or co-authored using generative AI tooling? Co-developed with Claude Code (Claude Opus 5.5): the fix, the tests and this description. The checks above were run locally. Generated-by: Claude Code (Claude Opus 5.5) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
