test_melt.py

import re

import numpy as np
import pytest

import pandas as pd
from pandas import (
    DataFrame,
    Index,
    date_range,
    lreshape,
    melt,
    wide_to_long,
)
import pandas._testing as tm


@pytest.fixture
def df():
    res = DataFrame(
        np.random.default_rng(2).standard_normal((10, 4)),
        columns=Index(list("ABCD"), dtype=object),
        index=date_range("2000-01-01", periods=10, freq="B"),
    )
    res["id1"] = (res["A"] > 0).astype(np.int64)
    res["id2"] = (res["B"] > 0).astype(np.int64)
    return res


@pytest.fixture
def df1():
    res = DataFrame(
        [
            [1.067683, -1.110463, 0.20867],
            [-1.321405, 0.368915, -1.055342],
            [-0.807333, 0.08298, -0.873361],
        ]
    )
    res.columns = [list("ABC"), list("abc")]
    res.columns.names = ["CAP", "low"]
    return res


@pytest.fixture
def var_name():
    return "var"


@pytest.fixture
def value_name():
    return "val"


class TestMelt:
    def test_top_level_method(self, df):
        result = melt(df)
        assert result.columns.tolist() == ["variable", "value"]

    def test_method_signatures(self, df, df1, var_name, value_name):
        tm.assert_frame_equal(df.melt(), melt(df))

        tm.assert_frame_equal(
            df.melt(id_vars=["id1", "id2"], value_vars=["A", "B"]),
            melt(df, id_vars=["id1", "id2"], value_vars=["A", "B"]),
        )

        tm.assert_frame_equal(
            df.melt(var_name=var_name, value_name=value_name),
            melt(df, var_name=var_name, value_name=value_name),
        )

        tm.assert_frame_equal(df1.melt(col_level=0), melt(df1, col_level=0))

    def test_default_col_names(self, df):
        result = df.melt()
        assert result.columns.tolist() == ["variable", "value"]

        result1 = df.melt(id_vars=["id1"])
        assert result1.columns.tolist() == ["id1", "variable", "value"]

        result2 = df.melt(id_vars=["id1", "id2"])
        assert result2.columns.tolist() == ["id1", "id2", "variable", "value"]

    def test_value_vars(self, df):
        result3 = df.melt(id_vars=["id1", "id2"], value_vars="A")
        assert len(result3) == 10

        result4 = df.melt(id_vars=["id1", "id2"], value_vars=["A", "B"])
        expected4 = DataFrame(
            {
                "id1": df["id1"].tolist() * 2,
                "id2": df["id2"].tolist() * 2,
                "variable": ["A"] * 10 + ["B"] * 10,
                "value": (df["A"].tolist() + df["B"].tolist()),
            },
            columns=["id1", "id2", "variable", "value"],
        )
        tm.assert_frame_equal(result4, expected4)

    @pytest.mark.parametrize("type_", (tuple, list, np.array))
    def test_value_vars_types(self, type_, df):
        # GH 15348
        expected = DataFrame(
            {
                "id1": df["id1"].tolist() * 2,
                "id2": df["id2"].tolist() * 2,
                "variable": ["A"] * 10 + ["B"] * 10,
                "value": (df["A"].tolist() + df["B"].tolist()),
            },
            columns=["id1", "id2", "variable", "value"],
        )
        result = df.melt(id_vars=["id1", "id2"], value_vars=type_(("A", "B")))
        tm.assert_frame_equal(result, expected)

    def test_vars_work_with_multiindex(self, df1):
        expected = DataFrame(
            {
                ("A", "a"): df1[("A", "a")],
                "CAP": ["B"] * len(df1),
                "low": ["b"] * len(df1),
                "value": df1[("B", "b")],
            },
            columns=[("A", "a"), "CAP", "low", "value"],
        )

        result = df1.melt(id_vars=[("A", "a")], value_vars=[("B", "b")])
        tm.assert_frame_equal(result, expected)

    @pytest.mark.parametrize(
        "id_vars, value_vars, col_level, expected",
        [
            (
                ["A"],
                ["B"],
                0,
                DataFrame(
                    {
                        "A": {0: 1.067683, 1: -1.321405, 2: -0.807333},
                        "CAP": {0: "B", 1: "B", 2: "B"},
                        "value": {0: -1.110463, 1: 0.368915, 2: 0.08298},
                    }
                ),
            ),
            (
                ["a"],
                ["b"],
                1,
                DataFrame(
                    {
                        "a": {0: 1.067683, 1: -1.321405, 2: -0.807333},
                        "low": {0: "b", 1: "b", 2: "b"},
                        "value": {0: -1.110463, 1: 0.368915, 2: 0.08298},
                    }
                ),
            ),
        ],
    )
    def test_single_vars_work_with_multiindex(
        self, id_vars, value_vars, col_level, expected, df1
    ):
        result = df1.melt(id_vars, value_vars, col_level=col_level)
        tm.assert_frame_equal(result, expected)

    @pytest.mark.parametrize(
        "id_vars, value_vars",
        [
            [("A", "a"), [("B", "b")]],
            [[("A", "a")], ("B", "b")],
            [("A", "a"), ("B", "b")],
        ],
    )
    def test_tuple_vars_fail_with_multiindex(self, id_vars, value_vars, df1):
        # melt should fail with an informative error message if
        # the columns have a MultiIndex and a tuple is passed
        # for id_vars or value_vars.
        msg = r"(id|value)_vars must be a list of tuples when columns are a MultiIndex"
        with pytest.raises(ValueError, match=msg):
            df1.melt(id_vars=id_vars, value_vars=value_vars)

    def test_custom_var_name(self, df, var_name):
        result5 = df.melt(var_name=var_name)
        assert result5.columns.tolist() == ["var", "value"]

        result6 = df.melt(id_vars=["id1"], var_name=var_name)
        assert result6.columns.tolist() == ["id1", "var", "value"]

        result7 = df.melt(id_vars=["id1", "id2"], var_name=var_name)
        assert result7.columns.tolist() == ["id1", "id2", "var", "value"]

        result8 = df.melt(id_vars=["id1", "id2"], value_vars="A", var_name=var_name)
        assert result8.columns.tolist() == ["id1", "id2", "var", "value"]

        result9 = df.melt(
            id_vars=["id1", "id2"], value_vars=["A", "B"], var_name=var_name
        )
        expected9 = DataFrame(
            {
                "id1": df["id1"].tolist() * 2,
                "id2": df["id2"].tolist() * 2,
                var_name: ["A"] * 10 + ["B"] * 10,
                "value": (df["A"].tolist() + df["B"].tolist()),
            },
            columns=["id1", "id2", var_name, "value"],
        )
        tm.assert_frame_equal(result9, expected9)

    def test_custom_value_name(self, df, value_name):
        result10 = df.melt(value_name=value_name)
        assert result10.columns.tolist() == ["variable", "val"]

        result11 = df.melt(id_vars=["id1"], value_name=value_name)
        assert result11.columns.tolist() == ["id1", "variable", "val"]

        result12 = df.melt(id_vars=["id1", "id2"], value_name=value_name)
        assert result12.columns.tolist() == ["id1", "id2", "variable", "val"]

        result13 = df.melt(
            id_vars=["id1", "id2"], value_vars="A", value_name=value_name
        )
        assert result13.columns.tolist() == ["id1", "id2", "variable", "val"]

        result14 = df.melt(
            id_vars=["id1", "id2"], value_vars=["A", "B"], value_name=value_name
        )
        expected14 = DataFrame(
            {
                "id1": df["id1"].tolist() * 2,
                "id2": df["id2"].tolist() * 2,
                "variable": ["A"] * 10 + ["B"] * 10,
                value_name: (df["A"].tolist() + df["B"].tolist()),
            },
            columns=["id1", "id2", "variable", value_name],
        )
        tm.assert_frame_equal(result14, expected14)

    def test_custom_var_and_value_name(self, df, value_name, var_name):
        result15 = df.melt(var_name=var_name, value_name=value_name)
        assert result15.columns.tolist() == ["var", "val"]

        result16 = df.melt(id_vars=["id1"], var_name=var_name, value_name=value_name)
        assert result16.columns.tolist() == ["id1", "var", "val"]

        result17 = df.melt(
            id_vars=["id1", "id2"], var_name=var_name, value_name=value_name
        )
        assert result17.columns.tolist() == ["id1", "id2", "var", "val"]

        result18 = df.melt(
            id_vars=["id1", "id2"],
            value_vars="A",
            var_name=var_name,
            value_name=value_name,
        )
        assert result18.columns.tolist() == ["id1", "id2", "var", "val"]

        result19 = df.melt(
            id_vars=["id1", "id2"],
            value_vars=["A", "B"],
            var_name=var_name,
            value_name=value_name,
        )
        expected19 = DataFrame(
            {
                "id1": df["id1"].tolist() * 2,
                "id2": df["id2"].tolist() * 2,
                var_name: ["A"] * 10 + ["B"] * 10,
                value_name: (df["A"].tolist() + df["B"].tolist()),
            },
            columns=["id1", "id2", var_name, value_name],
        )
        tm.assert_frame_equal(result19, expected19)

        df20 = df.copy()
        df20.columns.name = "foo"
        result20 = df20.melt()
        assert result20.columns.tolist() == ["foo", "value"]

    @pytest.mark.parametrize("col_level", [0, "CAP"])
    def test_col_level(self, col_level, df1):
        res = df1.melt(col_level=col_level)
        assert res.columns.tolist() == ["CAP", "value"]

    def test_multiindex(self, df1):
        res = df1.melt()
        assert res.columns.tolist() == ["CAP", "low", "value"]

    @pytest.mark.parametrize(
        "col",
        [
            pd.Series(date_range("2010", periods=5, tz="US/Pacific")),
            pd.Series(["a", "b", "c", "a", "d"], dtype="category"),
            pd.Series([0, 1, 0, 0, 0]),
        ],
    )
    def test_pandas_dtypes(self, col):
        # GH 15785
        df = DataFrame(
            {"klass": range(5), "col": col, "attr1": [1, 0, 0, 0, 0], "attr2": col}
        )
        expected_value = pd.concat([pd.Series([1, 0, 0, 0, 0]), col], ignore_index=True)
        result = melt(
            df, id_vars=["klass", "col"], var_name="attribute", value_name="value"
        )
        expected = DataFrame(
            {
                0: list(range(5)) * 2,
                1: pd.concat([col] * 2, ignore_index=True),
                2: ["attr1"] * 5 + ["attr2"] * 5,
                3: expected_value,
            }
        )
        expected.columns = ["klass", "col", "attribute", "value"]
        tm.assert_frame_equal(result, expected)

    def test_preserve_category(self):
        # GH 15853
        data = DataFrame({"A": [1, 2], "B": pd.Categorical(["X", "Y"])})
        result = melt(data, ["B"], ["A"])
        expected = DataFrame(
            {"B": pd.Categorical(["X", "Y"]), "variable": ["A", "A"], "value": [1, 2]}
        )

        tm.assert_frame_equal(result, expected)

    def test_melt_missing_columns_raises(self):
        # GH-23575
        # This test is to ensure that pandas raises an error if melting is
        # attempted with column names absent from the dataframe

        # Generate data
        df = DataFrame(
            np.random.default_rng(2).standard_normal((5, 4)), columns=list("abcd")
        )

        # Try to melt with missing `value_vars` column name
        msg = "The following id_vars or value_vars are not present in the DataFrame:"
        with pytest.raises(KeyError, match=msg):
            df.melt(["a", "b"], ["C", "d"])

        # Try to melt with missing `id_vars` column name
        with pytest.raises(KeyError, match=msg):
            df.melt(["A", "b"], ["c", "d"])

        # Multiple missing
        with pytest.raises(
            KeyError,
            match=msg,
        ):
            df.melt(["a", "b", "not_here", "or_there"], ["c", "d"])

        # Multiindex melt fails if column is missing from multilevel melt
        multi = df.copy()
        multi.columns = [list("ABCD"), list("abcd")]
        with pytest.raises(KeyError, match=msg):
            multi.melt([("E", "a")], [("B", "b")])
        # Multiindex fails if column is missing from single level melt
        with pytest.raises(KeyError, match=msg):
            multi.melt(["A"], ["F"], col_level=0)

    def test_melt_mixed_int_str_id_vars(self):
        # GH 29718
        df = DataFrame({0: ["foo"], "a": ["bar"], "b": [1], "d": [2]})
        result = melt(df, id_vars=[0, "a"], value_vars=["b", "d"])
        expected = DataFrame(
            {0: ["foo"] * 2, "a": ["bar"] * 2, "variable": list("bd"), "value": [1, 2]}
        )
        tm.assert_frame_equal(result, expected)

    def test_melt_mixed_int_str_value_vars(self):
        # GH 29718
        df = DataFrame({0: ["foo"], "a": ["bar"]})
        result = melt(df, value_vars=[0, "a"])
        expected = DataFrame({"variable": [0, "a"], "value": ["foo", "bar"]})
        tm.assert_frame_equal(result, expected)

    def test_ignore_index(self):
        # GH 17440
        df = DataFrame({"foo": [0], "bar": [1]}, index=["first"])
        result = melt(df, ignore_index=False)
        expected = DataFrame(
            {"variable": ["foo", "bar"], "value": [0, 1]}, index=["first", "first"]
        )
        tm.assert_frame_equal(result, expected)

    def test_ignore_multiindex(self):
        # GH 17440
        index = pd.MultiIndex.from_tuples(
            [("first", "second"), ("first", "third")], names=["baz", "foobar"]
        )
        df = DataFrame({"foo": [0, 1], "bar": [2, 3]}, index=index)
        result = melt(df, ignore_index=False)

        expected_index = pd.MultiIndex.from_tuples(
            [("first", "second"), ("first", "third")] * 2, names=["baz", "foobar"]
        )
        expected = DataFrame(
            {"variable": ["foo"] * 2 + ["bar"] * 2, "value": [0, 1, 2, 3]},
            index=expected_index,
        )

        tm.assert_frame_equal(result, expected)

    def test_ignore_index_name_and_type(self):
        # GH 17440
        index = Index(["foo", "bar"], dtype="category", name="baz")
        df = DataFrame({"x": [0, 1], "y": [2, 3]}, index=index)
        result = melt(df, ignore_index=False)

        expected_index = Index(["foo", "bar"] * 2, dtype="category", name="baz")
        expected = DataFrame(
            {"variable": ["x", "x", "y", "y"], "value": [0, 1, 2, 3]},
            index=expected_index,
        )

        tm.assert_frame_equal(result, expected)

    def test_melt_with_duplicate_columns(self):
        # GH#41951
        df = DataFrame([["id", 2, 3]], columns=["a", "b", "b"])
        result = df.melt(id_vars=["a"], value_vars=["b"])
        expected = DataFrame(
            [["id", "b", 2], ["id", "b", 3]], columns=["a", "variable", "value"]
        )
        tm.assert_frame_equal(result, expected)

    @pytest.mark.parametrize("dtype", ["Int8", "Int64"])
    def test_melt_ea_dtype(self, dtype):
        # GH#41570
        df = DataFrame(
            {
                "a": pd.Series([1, 2], dtype="Int8"),
                "b": pd.Series([3, 4], dtype=dtype),
            }
        )
        result = df.melt()
        expected = DataFrame(
            {
                "variable": ["a", "a", "b", "b"],
                "value": pd.Series([1, 2, 3, 4], dtype=dtype),
            }
        )
        tm.assert_frame_equal(result, expected)

    def test_melt_ea_columns(self):
        # GH 54297
        df = DataFrame(
            {
                "A": {0: "a", 1: "b", 2: "c"},
                "B": {0: 1, 1: 3, 2: 5},
                "C": {0: 2, 1: 4, 2: 6},
            }
        )
        df.columns = df.columns.astype("string[python]")
        result = df.melt(id_vars=["A"], value_vars=["B"])
        expected = DataFrame(
            {
                "A": list("abc"),
                "variable": pd.Series(["B"] * 3, dtype="string[python]"),
                "value": [1, 3, 5],
            }
        )
        tm.assert_frame_equal(result, expected)

    def test_melt_preserves_datetime(self):
        df = DataFrame(
            data=[
                {
                    "type": "A0",
                    "start_date": pd.Timestamp("2023/03/01", tz="Asia/Tokyo"),
                    "end_date": pd.Timestamp("2023/03/10", tz="Asia/Tokyo"),
                },
                {
                    "type": "A1",
                    "start_date": pd.Timestamp("2023/03/01", tz="Asia/Tokyo"),
                    "end_date": pd.Timestamp("2023/03/11", tz="Asia/Tokyo"),
                },
            ],
            index=["aaaa", "bbbb"],
        )
        result = df.melt(
            id_vars=["type"],
            value_vars=["start_date", "end_date"],
            var_name="start/end",
            value_name="date",
        )
        expected = DataFrame(
            {
                "type": {0: "A0", 1: "A1", 2: "A0", 3: "A1"},
                "start/end": {
                    0: "start_date",
                    1: "start_date",
                    2: "end_date",
                    3: "end_date",
                },
                "date": {
                    0: pd.Timestamp("2023-03-01 00:00:00+0900", tz="Asia/Tokyo"),
                    1: pd.Timestamp("2023-03-01 00:00:00+0900", tz="Asia/Tokyo"),
                    2: pd.Timestamp("2023-03-10 00:00:00+0900", tz="Asia/Tokyo"),
                    3: pd.Timestamp("2023-03-11 00:00:00+0900", tz="Asia/Tokyo"),
                },
            }
        )
        tm.assert_frame_equal(result, expected)

    def test_melt_allows_non_scalar_id_vars(self):
        df = DataFrame(
            data={"a": [1, 2, 3], "b": [4, 5, 6]},
            index=["11", "22", "33"],
        )
        result = df.melt(
            id_vars="a",
            var_name=0,
            value_name=1,
        )
        expected = DataFrame({"a": [1, 2, 3], 0: ["b"] * 3, 1: [4, 5, 6]})
        tm.assert_frame_equal(result, expected)

    def test_melt_allows_non_string_var_name(self):
        df = DataFrame(
            data={"a": [1, 2, 3], "b": [4, 5, 6]},
            index=["11", "22", "33"],
        )
        result = df.melt(
            id_vars=["a"],
            var_name=0,
            value_name=1,
        )
        expected = DataFrame({"a": [1, 2, 3], 0: ["b"] * 3, 1: [4, 5, 6]})
        tm.assert_frame_equal(result, expected)

    def test_melt_non_scalar_var_name_raises(self):
        df = DataFrame(
            data={"a": [1, 2, 3], "b": [4, 5, 6]},
            index=["11", "22", "33"],
        )
        with pytest.raises(ValueError, match=r".* must be a scalar."):
            df.melt(id_vars=["a"], var_name=[1, 2])


class TestLreshape:
    def test_pairs(self):
        data = {
            "birthdt": [
                "08jan2009",
                "20dec2008",
                "30dec2008",
                "21dec2008",
                "11jan2009",
            ],
            "birthwt": [1766, 3301, 1454, 3139, 4133],
            "id": [101, 102, 103, 104, 105],
            "sex": ["Male", "Female", "Female", "Female", "Female"],
            "visitdt1": [
                "11jan2009",
                "22dec2008",
                "04jan2009",
                "29dec2008",
                "20jan2009",
            ],
            "visitdt2": ["21jan2009", np.nan, "22jan2009", "31dec2008", "03feb2009"],
            "visitdt3": ["05feb2009", np.nan, np.nan, "02jan2009", "15feb2009"],
            "wt1": [1823, 3338, 1549, 3298, 4306],
            "wt2": [2011.0, np.nan, 1892.0, 3338.0, 4575.0],
            "wt3": [2293.0, np.nan, np.nan, 3377.0, 4805.0],
        }

        df = DataFrame(data)

        spec = {
            "visitdt": [f"visitdt{i:d}" for i in range(1, 4)],
            "wt": [f"wt{i:d}" for i in range(1, 4)],
        }
        result = lreshape(df, spec)

        exp_data = {
            "birthdt": [
                "08jan2009",
                "20dec2008",
                "30dec2008",
                "21dec2008",
                "11jan2009",
                "08jan2009",
                "30dec2008",
                "21dec2008",
                "11jan2009",
                "08jan2009",
                "21dec2008",
                "11jan2009",
            ],
            "birthwt": [
                1766,
                3301,
                1454,
                3139,
                4133,
                1766,
                1454,
                3139,
                4133,
                1766,
                3139,
                4133,
            ],
            "id": [101, 102, 103, 104, 105, 101, 103, 104, 105, 101, 104, 105],
            "sex": [
                "Male",
                "Female",
                "Female",
                "Female",
                "Female",
                "Male",
                "Female",
                "Female",
                "Female",
                "Male",
                "Female",
                "Female",
            ],
            "visitdt": [
                "11jan2009",
                "22dec2008",
                "04jan2009",
                "29dec2008",
                "20jan2009",
                "21jan2009",
                "22jan2009",
                "31dec2008",
                "03feb2009",
                "05feb2009",
                "02jan2009",
                "15feb2009",
            ],
            "wt": [
                1823.0,
                3338.0,
                1549.0,
                3298.0,
                4306.0,
                2011.0,
                1892.0,
                3338.0,
                4575.0,
                2293.0,
                3377.0,
                4805.0,
            ],
        }
        exp = DataFrame(exp_data, columns=result.columns)
        tm.assert_frame_equal(result, exp)

        result = lreshape(df, spec, dropna=False)
        exp_data = {
            "birthdt": [
                "08jan2009",
                "20dec2008",
                "30dec2008",
                "21dec2008",
                "11jan2009",
                "08jan2009",
                "20dec2008",
                "30dec2008",
                "21dec2008",
                "11jan2009",
                "08jan2009",
                "20dec2008",
                "30dec2008",
                "21dec2008",
                "11jan2009",
            ],
            "birthwt": [
                1766,
                3301,
                1454,
                3139,
                4133,
                1766,
                3301,
                1454,
                3139,
                4133,
                1766,
                3301,
                1454,
                3139,
                4133,
            ],
            "id": [
                101,
                102,
                103,
                104,
                105,
                101,
                102,
                103,
                104,
                105,
                101,
                102,
                103,
                104,
                105,
            ],
            "sex": [
                "Male",
                "Female",
                "Female",
                "Female",
                "Female",
                "Male",
                "Female",
                "Female",
                "Female",
                "Female",
                "Male",
                "Female",
                "Female",
                "Female",
                "Female",
            ],
            "visitdt": [
                "11jan2009",
                "22dec2008",
                "04jan2009",
                "29dec2008",
                "20jan2009",
                "21jan2009",
                np.nan,
                "22jan2009",
                "31dec2008",
                "03feb2009",
                "05feb2009",
                np.nan,
                np.nan,
                "02jan2009",
                "15feb2009",
            ],
            "wt": [
                1823.0,
                3338.0,
                1549.0,
                3298.0,
                4306.0,
                2011.0,
                np.nan,
                1892.0,
                3338.0,
                4575.0,
                2293.0,
                np.nan,
                np.nan,
                3377.0,
                4805.0,
            ],
        }
        exp = DataFrame(exp_data, columns=result.columns)
        tm.assert_frame_equal(result, exp)

        spec = {
            "visitdt": [f"visitdt{i:d}" for i in range(1, 3)],
            "wt": [f"wt{i:d}" for i in range(1, 4)],
        }
        msg = "All column lists must be same length"
        with pytest.raises(ValueError, match=msg):
            lreshape(df, spec)


class TestWideToLong:
    def test_simple(self):
        x = np.random.default_rng(2).standard_normal(3)
        df = DataFrame(
            {
                "A1970": {0: "a", 1: "b", 2: "c"},
                "A1980": {0: "d", 1: "e", 2: "f"},
                "B1970": {0: 2.5, 1: 1.2, 2: 0.7},
                "B1980": {0: 3.2, 1: 1.3, 2: 0.1},
                "X": dict(zip(range(3), x)),
            }
        )
        df["id"] = df.index
        exp_data = {
            "X": x.tolist() + x.tolist(),
            "A": ["a", "b", "c", "d", "e", "f"],
            "B": [2.5, 1.2, 0.7, 3.2, 1.3, 0.1],
            "year": [1970, 1970, 1970, 1980, 1980, 1980],
            "id": [0, 1, 2, 0, 1, 2],
        }
        expected = DataFrame(exp_data)
        expected = expected.set_index(["id", "year"])[["X", "A", "B"]]
        result = wide_to_long(df, ["A", "B"], i="id", j="year")
        tm.assert_frame_equal(result, expected)

    def test_stubs(self):
        # GH9204 wide_to_long call should not modify 'stubs' list
        df = DataFrame([[0, 1, 2, 3, 8], [4, 5, 6, 7, 9]])
        df.columns = ["id", "inc1", "inc2", "edu1", "edu2"]
        stubs = ["inc", "edu"]

        wide_to_long(df, stubs, i="id", j="age")

        assert stubs == ["inc", "edu"]

    def test_separating_character(self):
        # GH14779

        x = np.random.default_rng(2).standard_normal(3)
        df = DataFrame(
            {
                "A.1970": {0: "a", 1: "b", 2: "c"},
                "A.1980": {0: "d", 1: "e", 2: "f"},
                "B.1970": {0: 2.5, 1: 1.2, 2: 0.7},
                "B.1980": {0: 3.2, 1: 1.3, 2: 0.1},
                "X": dict(zip(range(3), x)),
            }
        )
        df["id"] = df.index
        exp_data = {
            "X": x.tolist() + x.tolist(),
            "A": ["a", "b", "c", "d", "e", "f"],
            "B": [2.5, 1.2, 0.7, 3.2, 1.3, 0.1],
            "year": [1970, 1970, 1970, 1980, 1980, 1980],
            "id": [0, 1, 2, 0, 1, 2],
        }
        expected = DataFrame(exp_data)
        expected = expected.set_index(["id", "year"])[["X", "A", "B"]]
        result = wide_to_long(df, ["A", "B"], i="id", j="year", sep=".")
        tm.assert_frame_equal(result, expected)

    def test_escapable_characters(self):
        x = np.random.default_rng(2).standard_normal(3)
        df = DataFrame(
            {
                "A(quarterly)1970": {0: "a", 1: "b", 2: "c"},
                "A(quarterly)1980": {0: "d", 1: "e", 2: "f"},
                "B(quarterly)1970": {0: 2.5, 1: 1.2, 2: 0.7},
                "B(quarterly)1980": {0: 3.2, 1: 1.3, 2: 0.1},
                "X": dict(zip(range(3), x)),
            }
        )
        df["id"] = df.index
        exp_data = {
            "X": x.tolist() + x.tolist(),
            "A(quarterly)": ["a", "b", "c", "d", "e", "f"],
            "B(quarterly)": [2.5, 1.2, 0.7, 3.2, 1.3, 0.1],
            "year": [1970, 1970, 1970, 1980, 1980, 1980],
            "id": [0, 1, 2, 0, 1, 2],
        }
        expected = DataFrame(exp_data)
        expected = expected.set_index(["id", "year"])[
            ["X", "A(quarterly)", "B(quarterly)"]
        ]
        result = wide_to_long(df, ["A(quarterly)", "B(quarterly)"], i="id", j="year")
        tm.assert_frame_equal(result, expected)

    def test_unbalanced(self):
        # test that we can have a varying amount of time variables
        df = DataFrame(
            {
                "A2010": [1.0, 2.0],
                "A2011": [3.0, 4.0],
                "B2010": [5.0, 6.0],
                "X": ["X1", "X2"],
            }
        )
        df["id"] = df.index
        exp_data = {
            "X": ["X1", "X2", "X1", "X2"],
            "A": [1.0, 2.0, 3.0, 4.0],
            "B": [5.0, 6.0, np.nan, np.nan],
            "id": [0, 1, 0, 1],
            "year": [2010, 2010, 2011, 2011],
        }
        expected = DataFrame(exp_data)
        expected = expected.set_index(["id", "year"])[["X", "A", "B"]]
        result = wide_to_long(df, ["A", "B"], i="id", j="year")
        tm.assert_frame_equal(result, expected)

    def test_character_overlap(self):
        # Test we handle overlapping characters in both id_vars and value_vars
        df = DataFrame(
            {
                "A11": ["a11", "a22", "a33"],
                "A12": ["a21", "a22", "a23"],
                "B11": ["b11", "b12", "b13"],
                "B12": ["b21", "b22", "b23"],
                "BB11": [1, 2, 3],
                "BB12": [4, 5, 6],
                "BBBX": [91, 92, 93],
                "BBBZ": [91, 92, 93],
            }
        )
        df["id"] = df.index
        expected = DataFrame(
            {
                "BBBX": [91, 92, 93, 91, 92, 93],
                "BBBZ": [91, 92, 93, 91, 92, 93],
                "A": ["a11", "a22", "a33", "a21", "a22", "a23"],
                "B": ["b11", "b12", "b13", "b21", "b22", "b23"],
                "BB": [1, 2, 3, 4, 5, 6],
                "id": [0, 1, 2, 0, 1, 2],
                "year": [11, 11, 11, 12, 12, 12],
            }
        )
        expected = expected.set_index(["id", "year"])[["BBBX", "BBBZ", "A", "B", "BB"]]
        result = wide_to_long(df, ["A", "B", "BB"], i="id", j="year")
        tm.assert_frame_equal(result.sort_index(axis=1), expected.sort_index(axis=1))

    def test_invalid_separator(self):
        # if an invalid separator is supplied a empty data frame is returned
        sep = "nope!"
        df = DataFrame(
            {
                "A2010": [1.0, 2.0],
                "A2011": [3.0, 4.0],
                "B2010": [5.0, 6.0],
                "X": ["X1", "X2"],
            }
        )
        df["id"] = df.index
        exp_data = {
            "X": "",
            "A2010": [],
            "A2011": [],
            "B2010": [],
            "id": [],
            "year": [],
            "A": [],
            "B": [],
        }
        expected = DataFrame(exp_data).astype({"year": np.int64})
        expected = expected.set_index(["id", "year"])[
            ["X", "A2010", "A2011", "B2010", "A", "B"]
        ]
        expected.index = expected.index.set_levels([0, 1], level=0)
        result = wide_to_long(df, ["A", "B"], i="id", j="year", sep=sep)
        tm.assert_frame_equal(result.sort_index(axis=1), expected.sort_index(axis=1))

    def test_num_string_disambiguation(self):
        # Test that we can disambiguate number value_vars from
        # string value_vars
        df = DataFrame(
            {
                "A11": ["a11", "a22", "a33"],
                "A12": ["a21", "a22", "a23"],
                "B11": ["b11", "b12", "b13"],
                "B12": ["b21", "b22", "b23"],
                "BB11": [1, 2, 3],
                "BB12": [4, 5, 6],
                "Arating": [91, 92, 93],
                "Arating_old": [91, 92, 93],
            }
        )
        df["id"] = df.index
        expected = DataFrame(
            {
                "Arating": [91, 92, 93, 91, 92, 93],
                "Arating_old": [91, 92, 93, 91, 92, 93],
                "A": ["a11", "a22", "a33", "a21", "a22", "a23"],
                "B": ["b11", "b12", "b13", "b21", "b22", "b23"],
                "BB": [1, 2, 3, 4, 5, 6],
                "id": [0, 1, 2, 0, 1, 2],
                "year": [11, 11, 11, 12, 12, 12],
            }
        )
        expected = expected.set_index(["id", "year"])[
            ["Arating", "Arating_old", "A", "B", "BB"]
        ]
        result = wide_to_long(df, ["A", "B", "BB"], i="id", j="year")
        tm.assert_frame_equal(result.sort_index(axis=1), expected.sort_index(axis=1))

    def test_invalid_suffixtype(self):
        # If all stubs names end with a string, but a numeric suffix is
        # assumed,  an empty data frame is returned
        df = DataFrame(
            {
                "Aone": [1.0, 2.0],
                "Atwo": [3.0, 4.0],
                "Bone": [5.0, 6.0],
                "X": ["X1", "X2"],
            }
        )
        df["id"] = df.index
        exp_data = {
            "X": "",
            "Aone": [],
            "Atwo": [],
            "Bone": [],
            "id": [],
            "year": [],
            "A": [],
            "B": [],
        }
        expected = DataFrame(exp_data).astype({"year": np.int64})

        expected = expected.set_index(["id", "year"])
        expected.index = expected.index.set_levels([0, 1], level=0)
        result = wide_to_long(df, ["A", "B"], i="id", j="year")
        tm.assert_frame_equal(result.sort_index(axis=1), expected.sort_index(axis=1))

    def test_multiple_id_columns(self):
        # Taken from http://www.ats.ucla.edu/stat/stata/modules/reshapel.htm
        df = DataFrame(
            {
                "famid": [1, 1, 1, 2, 2, 2, 3, 3, 3],
                "birth": [1, 2, 3, 1, 2, 3, 1, 2, 3],
                "ht1": [2.8, 2.9, 2.2, 2, 1.8, 1.9, 2.2, 2.3, 2.1],
                "ht2": [3.4, 3.8, 2.9, 3.2, 2.8, 2.4, 3.3, 3.4, 2.9],
            }
        )
        expected = DataFrame(
            {
                "ht": [
                    2.8,
                    3.4,
                    2.9,
                    3.8,
                    2.2,
                    2.9,
                    2.0,
                    3.2,
                    1.8,
                    2.8,
                    1.9,
                    2.4,
                    2.2,
                    3.3,
                    2.3,
                    3.4,
                    2.1,
                    2.9,
                ],
                "famid": [1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3],
                "birth": [1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3, 1, 1, 2, 2, 3, 3],
                "age": [1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2],
            }
        )
        expected = expected.set_index(["famid", "birth", "age"])[["ht"]]
        result = wide_to_long(df, "ht", i=["famid", "birth"], j="age")
        tm.assert_frame_equal(result, expected)

    def test_non_unique_idvars(self):
        # GH16382
        # Raise an error message if non unique id vars (i) are passed
        df = DataFrame(
            {"A_A1": [1, 2, 3, 4, 5], "B_B1": [1, 2, 3, 4, 5], "x": [1, 1, 1, 1, 1]}
        )
        msg = "the id variables need to uniquely identify each row"
        with pytest.raises(ValueError, match=msg):
            wide_to_long(df, ["A_A", "B_B"], i="x", j="colname")

    def test_cast_j_int(self):
        df = DataFrame(
            {
                "actor_1": ["CCH Pounder", "Johnny Depp", "Christoph Waltz"],
                "actor_2": ["Joel David Moore", "Orlando Bloom", "Rory Kinnear"],
                "actor_fb_likes_1": [1000.0, 40000.0, 11000.0],
                "actor_fb_likes_2": [936.0, 5000.0, 393.0],
                "title": ["Avatar", "Pirates of the Caribbean", "Spectre"],
            }
        )

        expected = DataFrame(
            {
                "actor": [
                    "CCH Pounder",
                    "Johnny Depp",
                    "Christoph Waltz",
                    "Joel David Moore",
                    "Orlando Bloom",
                    "Rory Kinnear",
                ],
                "actor_fb_likes": [1000.0, 40000.0, 11000.0, 936.0, 5000.0, 393.0],
                "num": [1, 1, 1, 2, 2, 2],
                "title": [
                    "Avatar",
                    "Pirates of the Caribbean",
                    "Spectre",
                    "Avatar",
                    "Pirates of the Caribbean",
                    "Spectre",
                ],
            }
        ).set_index(["title", "num"])
        result = wide_to_long(
            df, ["actor", "actor_fb_likes"], i="title", j="num", sep="_"
        )

        tm.assert_frame_equal(result, expected)

    def test_identical_stubnames(self):
        df = DataFrame(
            {
                "A2010": [1.0, 2.0],
                "A2011": [3.0, 4.0],
                "B2010": [5.0, 6.0],
                "A": ["X1", "X2"],
            }
        )
        msg = "stubname can't be identical to a column name"
        with pytest.raises(ValueError, match=msg):
            wide_to_long(df, ["A", "B"], i="A", j="colname")

    def test_nonnumeric_suffix(self):
        df = DataFrame(
            {
                "treatment_placebo": [1.0, 2.0],
                "treatment_test": [3.0, 4.0],
                "result_placebo": [5.0, 6.0],
                "A": ["X1", "X2"],
            }
        )
        expected = DataFrame(
            {
                "A": ["X1", "X2", "X1", "X2"],
                "colname": ["placebo", "placebo", "test", "test"],
                "result": [5.0, 6.0, np.nan, np.nan],
                "treatment": [1.0, 2.0, 3.0, 4.0],
            }
        )
        expected = expected.set_index(["A", "colname"])
        result = wide_to_long(
            df, ["result", "treatment"], i="A", j="colname", suffix="[a-z]+", sep="_"
        )
        tm.assert_frame_equal(result, expected)

    def test_mixed_type_suffix(self):
        df = DataFrame(
            {
                "A": ["X1", "X2"],
                "result_1": [0, 9],
                "result_foo": [5.0, 6.0],
                "treatment_1": [1.0, 2.0],
                "treatment_foo": [3.0, 4.0],
            }
        )
        expected = DataFrame(
            {
                "A": ["X1", "X2", "X1", "X2"],
                "colname": ["1", "1", "foo", "foo"],
                "result": [0.0, 9.0, 5.0, 6.0],
                "treatment": [1.0, 2.0, 3.0, 4.0],
            }
        ).set_index(["A", "colname"])
        result = wide_to_long(
            df, ["result", "treatment"], i="A", j="colname", suffix=".+", sep="_"
        )
        tm.assert_frame_equal(result, expected)

    def test_float_suffix(self):
        df = DataFrame(
            {
                "treatment_1.1": [1.0, 2.0],
                "treatment_2.1": [3.0, 4.0],
                "result_1.2": [5.0, 6.0],
                "result_1": [0, 9],
                "A": ["X1", "X2"],
            }
        )
        expected = DataFrame(
            {
                "A": ["X1", "X2", "X1", "X2", "X1", "X2", "X1", "X2"],
                "colname": [1.2, 1.2, 1.0, 1.0, 1.1, 1.1, 2.1, 2.1],
                "result": [5.0, 6.0, 0.0, 9.0, np.nan, np.nan, np.nan, np.nan],
                "treatment": [np.nan, np.nan, np.nan, np.nan, 1.0, 2.0, 3.0, 4.0],
            }
        )
        expected = expected.set_index(["A", "colname"])
        result = wide_to_long(
            df, ["result", "treatment"], i="A", j="colname", suffix="[0-9.]+", sep="_"
        )
        tm.assert_frame_equal(result, expected)

    def test_col_substring_of_stubname(self):
        # GH22468
        # Don't raise ValueError when a column name is a substring
        # of a stubname that's been passed as a string
        wide_data = {
            "node_id": {0: 0, 1: 1, 2: 2, 3: 3, 4: 4},
            "A": {0: 0.80, 1: 0.0, 2: 0.25, 3: 1.0, 4: 0.81},
            "PA0": {0: 0.74, 1: 0.56, 2: 0.56, 3: 0.98, 4: 0.6},
            "PA1": {0: 0.77, 1: 0.64, 2: 0.52, 3: 0.98, 4: 0.67},
            "PA3": {0: 0.34, 1: 0.70, 2: 0.52, 3: 0.98, 4: 0.67},
        }
        wide_df = DataFrame.from_dict(wide_data)
        expected = wide_to_long(wide_df, stubnames=["PA"], i=["node_id", "A"], j="time")
        result = wide_to_long(wide_df, stubnames="PA", i=["node_id", "A"], j="time")
        tm.assert_frame_equal(result, expected)

    def test_raise_of_column_name_value(self):
        # GH34731, enforced in 2.0
        # raise a ValueError if the resultant value column name matches
        # a name in the dataframe already (default name is "value")
        df = DataFrame({"col": list("ABC"), "value": range(10, 16, 2)})

        with pytest.raises(
            ValueError, match=re.escape("value_name (value) cannot match")
        ):
            df.melt(id_vars="value", value_name="value")

    @pytest.mark.parametrize("dtype", ["O", "string"])
    def test_missing_stubname(self, dtype):
        # GH46044
        df = DataFrame({"id": ["1", "2"], "a-1": [100, 200], "a-2": [300, 400]})
        df = df.astype({"id": dtype})
        result = wide_to_long(
            df,
            stubnames=["a", "b"],
            i="id",
            j="num",
            sep="-",
        )
        index = Index(
            [("1", 1), ("2", 1), ("1", 2), ("2", 2)],
            name=("id", "num"),
        )
        expected = DataFrame(
            {"a": [100, 200, 300, 400], "b": [np.nan] * 4},
            index=index,
        )
        new_level = expected.index.levels[0].astype(dtype)
        expected.index = expected.index.set_levels(new_level, level=0)
        tm.assert_frame_equal(result, expected)


def test_wide_to_long_pyarrow_string_columns():
    # GH 57066
    pytest.importorskip("pyarrow")
    df = DataFrame(
        {
            "ID": {0: 1},
            "R_test1": {0: 1},
            "R_test2": {0: 1},
            "R_test3": {0: 2},
            "D": {0: 1},
        }
    )
    df.columns = df.columns.astype("string[pyarrow_numpy]")
    result = wide_to_long(
        df, stubnames="R", i="ID", j="UNPIVOTED", sep="_", suffix=".*"
    )
    expected = DataFrame(
        [[1, 1], [1, 1], [1, 2]],
        columns=Index(["D", "R"], dtype=object),
        index=pd.MultiIndex.from_arrays(
            [
                [1, 1, 1],
                Index(["test1", "test2", "test3"], dtype="string[pyarrow_numpy]"),
            ],
            names=["ID", "UNPIVOTED"],
        ),
    )
    tm.assert_frame_equal(result, expected)