BUG: pivot with array-like index/columns labels (GH#35785)

In the no-values path, `cols + columns_listlike` combined the label
arguments elementwise, so an ndarray/Index/Series/ExtensionArray of valid
column labels raised instead of being used as labels; a Series or
ExtensionArray also reached `unstack` unconverted. Convert `index` and
`columns` to plain lists once and reuse them.

The label check added in GH#66311 converted `index`/`values` a second
time, exhausting a one-shot iterator so the body saw an empty sequence,
and it rejected an unhashable entry of the `index` list even though
`set_index` takes that as a column of level values.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: jbrockmendel

doc/source/whatsnew/v3.1.0.rst

@@ -594,6 +594,7 @@ Reshaping
 - Bug in :meth:`DataFrame.combine` raising ``OverflowError`` when the combining function returned a value too large for the columns' common integer dtype (e.g. ``2**64``) instead of keeping the result (:issue:`66394`)
 - Bug in :meth:`DataFrame.melt` where ``var_name`` colliding with an ``id_vars`` column or ``value_name`` silently overwrote the affected column data instead of raising (:issue:`65654`)
 - Bug in :meth:`DataFrame.pivot_table` with ``margins=True`` raising ``TypeError`` when ``values`` has an :class:`ExtensionDtype` that cannot hold ``NA`` (e.g. :class:`IntervalDtype` with an integer subtype) and no ``columns`` were specified (:issue:`55484`)
+- Bug in :meth:`DataFrame.pivot` and :func:`pivot` raising ``TypeError`` or ``ValueError`` when ``index`` or ``columns`` was given as an array-like of column labels, e.g. a ``numpy.ndarray``, :class:`Index`, or :class:`Series`, or when ``values`` was given as an iterator of column labels (:issue:`35785`)
 - Bug in :meth:`DataFrame.select_dtypes` where a nullable extension dtype name such as ``"Int64"`` or ``"boolean"`` also selected numpy columns sharing the same scalar type; it now selects only columns of that extension dtype (:issue:`40234`)
 - Bug in :meth:`DataFrame.select_dtypes` where an interval spec giving a subtype but no ``closed`` keyword, such as ``"interval[int64]"`` or ``pd.IntervalDtype("int64")``, selected no columns; it now selects every interval column with that subtype, for any ``closed`` value (:issue:`40234`)
 - Bug in :meth:`DataFrame.select_dtypes` where the string form of an Arrow dtype such as ``"int64[pyarrow]"`` selected no columns and ``"float64[pyarrow]"`` selected numpy float columns; an Arrow dtype string now selects only columns of that exact dtype (:issue:`59888`)

pandas/core/frame.py

@@ -13180,11 +13180,11 @@ def pivot(
 
         Parameters
         ----------
-        columns : Hashable or a sequence of the previous
+        columns : Hashable or a sequence of the previous, or array-like
             Column to use to make new frame's columns.
-        index : Hashable or a sequence of the previous, optional
+        index : Hashable or a sequence of the previous, or array-like, optional
             Column to use to make new frame's index. If not given, uses existing index.
-        values : Hashable or a sequence of the previous, optional
+        values : Hashable or a sequence of the previous, or array-like, optional
             Column(s) to use for populating new frame's values. If not
             specified, all remaining columns will be used and the result will
             have hierarchically indexed columns.
@@ -13211,6 +13211,12 @@ def pivot(
 
         Notes
         -----
+        An array-like ``index`` or ``columns`` is a sequence of column labels, unlike
+        :meth:`DataFrame.pivot_table`, which also accepts an array of values to group
+        by. When ``values`` is not given, an array nested inside the ``index`` list is
+        a column of level values, as in :meth:`DataFrame.set_index`. ``index``,
+        ``columns`` and ``values`` may also be given as an iterator of labels.
+
         For finer-tuned control, see hierarchical indexing documentation along
         with the related stack/unstack methods.
 

pandas/core/reshape/pivot.py

@@ -14,6 +14,7 @@
 
 from pandas.core.dtypes.cast import maybe_downcast_to_dtype
 from pandas.core.dtypes.common import (
+    is_hashable,
     is_list_like,
     is_nested_list_like,
     is_scalar,
@@ -38,6 +39,7 @@
     from collections.abc import (
         Callable,
         Hashable,
+        Iterable,
     )
 
     from pandas._typing import (
@@ -747,9 +749,9 @@ def _convert_by(by):
 def pivot(
     data: DataFrame,
     *,
-    columns: IndexLabel,
-    index: IndexLabel | lib.NoDefault = lib.no_default,
-    values: IndexLabel | lib.NoDefault = lib.no_default,
+    columns: IndexLabel | Iterable[Hashable],
+    index: IndexLabel | Iterable[Hashable] | lib.NoDefault = lib.no_default,
+    values: IndexLabel | Iterable[Hashable] | lib.NoDefault = lib.no_default,
 ) -> DataFrame:
     """
     Return reshaped DataFrame organized by given index / column values.
@@ -764,11 +766,11 @@ def pivot(
     ----------
     data : DataFrame
         Input pandas DataFrame object.
-    columns : Hashable or a sequence of the previous
+    columns : Hashable or a sequence of the previous, or array-like
         Column to use to make new frame's columns.
-    index : Hashable or a sequence of the previous, optional
+    index : Hashable or a sequence of the previous, or array-like, optional
         Column to use to make new frame's index. If not given, uses existing index.
-    values : Hashable or a sequence of the previous, optional
+    values : Hashable or a sequence of the previous, or array-like, optional
         Column(s) to use for populating new frame's values. If not
         specified, all remaining columns will be used and the result will
         have hierarchically indexed columns.
@@ -795,6 +797,12 @@ def pivot(
 
     Notes
     -----
+    An array-like ``index`` or ``columns`` is a sequence of column labels, unlike
+    :meth:`DataFrame.pivot_table`, which also accepts an array of values to group by.
+    When ``values`` is not given, an array nested inside the ``index`` list is a column
+    of level values, as in :meth:`DataFrame.set_index`. ``index``, ``columns`` and
+    ``values`` may also be given as an iterator of labels.
+
     For finer-tuned control, see hierarchical indexing documentation along
     with the related stack/unstack methods.
 
@@ -897,18 +905,45 @@ def pivot(
        ...
     ValueError: Index contains duplicate entries, cannot reshape
     """
-    columns_listlike = com.convert_to_list_like(columns)
-
-    # GH#35785 without this, downstream label arithmetic raises cryptically.
-    labels_to_check: list[tuple[str, list]] = [("columns", list(columns_listlike))]
+    # GH#35785 normalize index/columns to plain lists once and reuse them: array-likes
+    #  of valid labels are then used as labels instead of being combined elementwise,
+    #  and a one-shot iterator is not exhausted by the check below.
+    columns_listlike = list(com.convert_to_list_like(columns))
+    # GH#35785 check up front so the error names the offending parameter; the
+    #  downstream set_index/unstack failure does not say which argument was wrong.
+    labels_to_check: list[tuple[str, list]] = [("columns", columns_listlike)]
+
+    index_listlike: list = []
     if index is not lib.no_default:
-        labels_to_check.append(("index", list(com.convert_to_list_like(index))))
+        index_listlike = list(com.convert_to_list_like(index))
+        index_labels = index_listlike
+        if values is lib.no_default:
+            # GH#35785 the set_index below takes an Index/Series/ndarray/list/iterator
+            #  entry as a column of level values rather than as a label
+            index_labels = [
+                label
+                for label in index_listlike
+                if not isinstance(label, (Index, ABCSeries, np.ndarray, list))
+                and not lib.is_iterator(label)
+            ]
+        labels_to_check.append(("index", index_labels))
+
+    if lib.is_iterator(values):
+        # GH#35785 materialize before the check below consumes it. ``values`` is not
+        #  list-ified like index/columns: it is used for the lookup as-is, where a
+        #  named Index/Series names the resulting columns level.
+        values = list(cast("Iterable[Hashable]", values))
     if values is not lib.no_default and not isinstance(values, tuple):
         # GH#17160 a tuple ``values`` is a single (MultiIndex) label; the
         #  existing lookup already raises a KeyError naming it.
         labels_to_check.append(("values", list(com.convert_to_list_like(values))))
+
     for param_name, labels in labels_to_check:
-        missing = [label for label in labels if label not in data.columns]
+        missing = [
+            label
+            for label in labels
+            if not is_hashable(label) or label not in data.columns
+        ]
         if missing:
             raise KeyError(
                 f"The following '{param_name}' labels are not columns of the "
@@ -926,18 +961,8 @@ def pivot(
 
     indexed: DataFrame | Series
     if values is lib.no_default:
-        if index is not lib.no_default:
-            cols = com.convert_to_list_like(index)
-        else:
-            cols = []
-
         append = index is lib.no_default
-        # error: Unsupported operand types for + ("List[Any]" and "ExtensionArray")
-        # error: Unsupported left operand type for + ("ExtensionArray")
-        indexed = data.set_index(
-            cols + columns_listlike,  # type: ignore[operator]
-            append=append,
-        )
+        indexed = data.set_index(index_listlike + columns_listlike, append=append)
     else:
         index_list: list[Index] | list[Series]
         if index is lib.no_default:
@@ -951,7 +976,7 @@ def pivot(
                     data._constructor_sliced(data.index, name=data.index.name)
                 ]
         else:
-            index_list = [data[idx] for idx in com.convert_to_list_like(index)]
+            index_list = [data[idx] for idx in index_listlike]
 
         data_columns = [data[col] for col in columns_listlike]
         index_list.extend(data_columns)
@@ -966,11 +991,8 @@ def pivot(
             )
         else:
             indexed = data._constructor_sliced(data[values]._values, index=multiindex)
-    # error: Argument 1 to "unstack" of "DataFrame" has incompatible type "Union
-    # [List[Any], ExtensionArray, ndarray[Any, Any], Index, Series]"; expected
-    # "Hashable"
     # unstack with a MultiIndex returns a DataFrame
-    result = cast("DataFrame", indexed.unstack(columns_listlike))  # type: ignore[arg-type]
+    result = cast("DataFrame", indexed.unstack(columns_listlike))
     result.index.names = [
         name if name is not lib.no_default else None for name in result.index.names
     ]

pandas/tests/reshape/test_pivot.py

@@ -3280,3 +3280,130 @@ def test_pivot_non_column_label_raises_gh35785():
     )
     expected.columns.name = "bar"
     tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize(
+    "box",
+    [list, np.array, Series, Index, pd.array],
+    ids=["list", "ndarray", "Series", "Index", "ExtensionArray"],
+)
+def test_pivot_array_like_labels_gh35785(box):
+    # GH#35785 array-likes of valid column labels are labels, not values
+    df = DataFrame(
+        {"foo": ["one", "one", "two"], "bar": ["A", "B", "A"], "baz": [1, 2, 3]}
+    )
+
+    result = df.pivot(index=box(["foo"]), columns=box(["bar"]))
+    tm.assert_frame_equal(result, df.pivot(index="foo", columns="bar"))
+
+    result = df.pivot(columns=box(["bar"]))
+    tm.assert_frame_equal(result, df.pivot(columns="bar"))
+
+    # the values path is affected too; e.g. a Series ``columns`` reached unstack
+    result = df.pivot(index=box(["foo"]), columns=box(["bar"]), values=["baz"])
+    tm.assert_frame_equal(
+        result, df.pivot(index=["foo"], columns=["bar"], values=["baz"])
+    )
+
+
+@pytest.mark.parametrize(
+    "box",
+    [list, np.array, Series, Index, pd.array],
+    ids=["list", "ndarray", "Series", "Index", "ExtensionArray"],
+)
+def test_pivot_array_like_multiple_labels_gh35785(box):
+    # GH#35785 two labels take unstack's _unstack_multiple path, which an
+    #  unconverted array-like reaches as an ambiguous truth value
+    df = DataFrame(
+        {
+            "lev1": [1, 1, 1, 2, 2, 2],
+            "lev2": [1, 1, 2, 1, 1, 2],
+            "lev3": [1, 2, 1, 2, 1, 2],
+            "values": [0, 1, 2, 3, 4, 5],
+        }
+    )
+
+    result = df.pivot(
+        index=box(["lev1"]), columns=box(["lev2", "lev3"]), values="values"
+    )
+    tm.assert_frame_equal(
+        result, df.pivot(index="lev1", columns=["lev2", "lev3"], values="values")
+    )
+
+
+@pytest.mark.parametrize("box", [Index, Series])
+def test_pivot_named_values_names_columns_level_gh35785(box):
+    # GH#35785 ``values`` is looked up as-is, so a named Index/Series names the
+    #  resulting columns level; normalizing it to a list would drop the name
+    df = DataFrame(
+        {"foo": ["one", "one", "two"], "bar": ["A", "B", "A"], "baz": [1, 2, 3]}
+    )
+
+    result = df.pivot(index="foo", columns="bar", values=box(["baz"], name="V"))
+
+    expected = df.pivot(index="foo", columns="bar", values=["baz"])
+    expected.columns.names = ["V", "bar"]
+    tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize("box", [np.array, list, Series, Index, iter])
+def test_pivot_index_array_of_values_gh35785(box):
+    # GH#35785 an entry that set_index takes as level values is not a label
+    df = DataFrame(
+        {"foo": ["one", "one", "two"], "bar": ["A", "B", "A"], "baz": [1, 2, 3]}
+    )
+    values = ["x", "y", "z"]
+
+    result = df.pivot(index=[box(values)], columns="bar")
+    expected = df.set_index([np.array(values), "bar"]).unstack(["bar"])
+    tm.assert_frame_equal(result, expected)
+
+    # only set_index's array spellings are exempt from the label check, and only
+    #  on this path: with ``values`` the entries are looked up as columns
+    with pytest.raises(KeyError, match="'index' labels are not columns"):
+        df.pivot(index=[box(values)], columns="bar", values="baz")
+    with pytest.raises(KeyError, match="'index' labels are not columns"):
+        df.pivot(index=[{"a": 1}], columns="bar")
+
+
+def test_pivot_array_columns_entry_raises_gh35785():
+    # GH#35785 unlike index, a columns entry is always a label, so an array is
+    #  named by the KeyError instead of failing later in unstack
+    df = DataFrame(
+        {"foo": ["one", "one", "two"], "bar": ["A", "B", "A"], "baz": [1, 2, 3]}
+    )
+
+    with pytest.raises(KeyError, match="'columns' labels are not columns"):
+        df.pivot(columns=[np.array(["x", "y", "z"])])
+    with pytest.raises(KeyError, match="'values' labels are not columns"):
+        df.pivot(index="foo", columns="bar", values=[np.array(["x", "y", "z"])])
+
+
+def test_pivot_datetimelike_array_labels_gh35785():
+    # GH#35785 a datetimelike array-like of labels also reached the elementwise add
+    df = DataFrame({pd.Timestamp("2016-03-02"): ["one", "two"], "bar": ["A", "B"]})
+    result = df.pivot(columns=np.array(["2016-03-02"], dtype="M8[ns]"))
+    tm.assert_frame_equal(result, df.pivot(columns=[pd.Timestamp("2016-03-02")]))
+
+    df = DataFrame({pd.Timedelta("1D"): ["one", "two"], "bar": ["A", "B"]})
+    result = df.pivot(columns=np.array([86400000000000], dtype="m8[ns]"))
+    tm.assert_frame_equal(result, df.pivot(columns=[pd.Timedelta("1D")]))
+
+
+def test_pivot_iterator_labels_gh35785():
+    # GH#35785 a one-shot iterator must not be consumed by the label check
+    df = DataFrame(
+        {"foo": ["one", "one", "two"], "bar": ["A", "B", "A"], "baz": [1, 2, 3]}
+    )
+
+    result = df.pivot(index=(col for col in ["foo"]), columns=(col for col in ["bar"]))
+    tm.assert_frame_equal(result, df.pivot(index="foo", columns="bar"))
+
+    result = df.pivot(
+        index=(col for col in ["foo"]),
+        columns=(col for col in ["bar"]),
+        values=(col for col in ["baz"]),
+    )
+    tm.assert_frame_equal(
+        result, df.pivot(index=["foo"], columns=["bar"], values=["baz"])
+    )