1253 lines
41 KiB
Python
1253 lines
41 KiB
Python
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)
|