630 lines
22 KiB
Python
630 lines
22 KiB
Python
|
from datetime import datetime
|
||
|
|
from itertools import product
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from pandas.core.dtypes.common import is_float_dtype, is_integer_dtype
|
||
|
|
|
||
|
|
import pandas as pd
|
||
|
|
from pandas import (
|
||
|
|
Categorical,
|
||
|
|
CategoricalIndex,
|
||
|
|
DataFrame,
|
||
|
|
Index,
|
||
|
|
IntervalIndex,
|
||
|
|
MultiIndex,
|
||
|
|
RangeIndex,
|
||
|
|
Series,
|
||
|
|
Timestamp,
|
||
|
|
date_range,
|
||
|
|
)
|
||
|
|
import pandas._testing as tm
|
||
|
|
|
||
|
|
|
||
|
|
class TestResetIndex:
|
||
|
|
def test_set_reset(self):
|
||
|
|
|
||
|
|
idx = Index([2 ** 63, 2 ** 63 + 5, 2 ** 63 + 10], name="foo")
|
||
|
|
|
||
|
|
# set/reset
|
||
|
|
df = DataFrame({"A": [0, 1, 2]}, index=idx)
|
||
|
|
result = df.reset_index()
|
||
|
|
assert result["foo"].dtype == np.dtype("uint64")
|
||
|
|
|
||
|
|
df = result.set_index("foo")
|
||
|
|
tm.assert_index_equal(df.index, idx)
|
||
|
|
|
||
|
|
def test_set_index_reset_index_dt64tz(self):
|
||
|
|
|
||
|
|
idx = Index(date_range("20130101", periods=3, tz="US/Eastern"), name="foo")
|
||
|
|
|
||
|
|
# set/reset
|
||
|
|
df = DataFrame({"A": [0, 1, 2]}, index=idx)
|
||
|
|
result = df.reset_index()
|
||
|
|
assert result["foo"].dtype == "datetime64[ns, US/Eastern]"
|
||
|
|
|
||
|
|
df = result.set_index("foo")
|
||
|
|
tm.assert_index_equal(df.index, idx)
|
||
|
|
|
||
|
|
def test_reset_index_tz(self, tz_aware_fixture):
|
||
|
|
# GH 3950
|
||
|
|
# reset_index with single level
|
||
|
|
tz = tz_aware_fixture
|
||
|
|
idx = date_range("1/1/2011", periods=5, freq="D", tz=tz, name="idx")
|
||
|
|
df = DataFrame({"a": range(5), "b": ["A", "B", "C", "D", "E"]}, index=idx)
|
||
|
|
|
||
|
|
expected = DataFrame(
|
||
|
|
{
|
||
|
|
"idx": [
|
||
|
|
datetime(2011, 1, 1),
|
||
|
|
datetime(2011, 1, 2),
|
||
|
|
datetime(2011, 1, 3),
|
||
|
|
datetime(2011, 1, 4),
|
||
|
|
datetime(2011, 1, 5),
|
||
|
|
],
|
||
|
|
"a": range(5),
|
||
|
|
"b": ["A", "B", "C", "D", "E"],
|
||
|
|
},
|
||
|
|
columns=["idx", "a", "b"],
|
||
|
|
)
|
||
|
|
expected["idx"] = expected["idx"].apply(lambda d: Timestamp(d, tz=tz))
|
||
|
|
tm.assert_frame_equal(df.reset_index(), expected)
|
||
|
|
|
||
|
|
@pytest.mark.parametrize("tz", ["US/Eastern", "dateutil/US/Eastern"])
|
||
|
|
def test_frame_reset_index_tzaware_index(self, tz):
|
||
|
|
dr = date_range("2012-06-02", periods=10, tz=tz)
|
||
|
|
df = DataFrame(np.random.randn(len(dr)), dr)
|
||
|
|
roundtripped = df.reset_index().set_index("index")
|
||
|
|
xp = df.index.tz
|
||
|
|
rs = roundtripped.index.tz
|
||
|
|
assert xp == rs
|
||
|
|
|
||
|
|
def test_reset_index_with_intervals(self):
|
||
|
|
idx = IntervalIndex.from_breaks(np.arange(11), name="x")
|
||
|
|
original = DataFrame({"x": idx, "y": np.arange(10)})[["x", "y"]]
|
||
|
|
|
||
|
|
result = original.set_index("x")
|
||
|
|
expected = DataFrame({"y": np.arange(10)}, index=idx)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
result2 = result.reset_index()
|
||
|
|
tm.assert_frame_equal(result2, original)
|
||
|
|
|
||
|
|
def test_reset_index(self, float_frame):
|
||
|
|
stacked = float_frame.stack()[::2]
|
||
|
|
stacked = DataFrame({"foo": stacked, "bar": stacked})
|
||
|
|
|
||
|
|
names = ["first", "second"]
|
||
|
|
stacked.index.names = names
|
||
|
|
deleveled = stacked.reset_index()
|
||
|
|
for i, (lev, level_codes) in enumerate(
|
||
|
|
zip(stacked.index.levels, stacked.index.codes)
|
||
|
|
):
|
||
|
|
values = lev.take(level_codes)
|
||
|
|
name = names[i]
|
||
|
|
tm.assert_index_equal(values, Index(deleveled[name]))
|
||
|
|
|
||
|
|
stacked.index.names = [None, None]
|
||
|
|
deleveled2 = stacked.reset_index()
|
||
|
|
tm.assert_series_equal(
|
||
|
|
deleveled["first"], deleveled2["level_0"], check_names=False
|
||
|
|
)
|
||
|
|
tm.assert_series_equal(
|
||
|
|
deleveled["second"], deleveled2["level_1"], check_names=False
|
||
|
|
)
|
||
|
|
|
||
|
|
# default name assigned
|
||
|
|
rdf = float_frame.reset_index()
|
||
|
|
exp = Series(float_frame.index.values, name="index")
|
||
|
|
tm.assert_series_equal(rdf["index"], exp)
|
||
|
|
|
||
|
|
# default name assigned, corner case
|
||
|
|
df = float_frame.copy()
|
||
|
|
df["index"] = "foo"
|
||
|
|
rdf = df.reset_index()
|
||
|
|
exp = Series(float_frame.index.values, name="level_0")
|
||
|
|
tm.assert_series_equal(rdf["level_0"], exp)
|
||
|
|
|
||
|
|
# but this is ok
|
||
|
|
float_frame.index.name = "index"
|
||
|
|
deleveled = float_frame.reset_index()
|
||
|
|
tm.assert_series_equal(deleveled["index"], Series(float_frame.index))
|
||
|
|
tm.assert_index_equal(deleveled.index, Index(range(len(deleveled))), exact=True)
|
||
|
|
|
||
|
|
# preserve column names
|
||
|
|
float_frame.columns.name = "columns"
|
||
|
|
resetted = float_frame.reset_index()
|
||
|
|
assert resetted.columns.name == "columns"
|
||
|
|
|
||
|
|
# only remove certain columns
|
||
|
|
df = float_frame.reset_index().set_index(["index", "A", "B"])
|
||
|
|
rs = df.reset_index(["A", "B"])
|
||
|
|
|
||
|
|
# TODO should reset_index check_names ?
|
||
|
|
tm.assert_frame_equal(rs, float_frame, check_names=False)
|
||
|
|
|
||
|
|
rs = df.reset_index(["index", "A", "B"])
|
||
|
|
tm.assert_frame_equal(rs, float_frame.reset_index(), check_names=False)
|
||
|
|
|
||
|
|
rs = df.reset_index(["index", "A", "B"])
|
||
|
|
tm.assert_frame_equal(rs, float_frame.reset_index(), check_names=False)
|
||
|
|
|
||
|
|
rs = df.reset_index("A")
|
||
|
|
xp = float_frame.reset_index().set_index(["index", "B"])
|
||
|
|
tm.assert_frame_equal(rs, xp, check_names=False)
|
||
|
|
|
||
|
|
# test resetting in place
|
||
|
|
df = float_frame.copy()
|
||
|
|
resetted = float_frame.reset_index()
|
||
|
|
return_value = df.reset_index(inplace=True)
|
||
|
|
assert return_value is None
|
||
|
|
tm.assert_frame_equal(df, resetted, check_names=False)
|
||
|
|
|
||
|
|
df = float_frame.reset_index().set_index(["index", "A", "B"])
|
||
|
|
rs = df.reset_index("A", drop=True)
|
||
|
|
xp = float_frame.copy()
|
||
|
|
del xp["A"]
|
||
|
|
xp = xp.set_index(["B"], append=True)
|
||
|
|
tm.assert_frame_equal(rs, xp, check_names=False)
|
||
|
|
|
||
|
|
def test_reset_index_name(self):
|
||
|
|
df = DataFrame(
|
||
|
|
[[1, 2, 3, 4], [5, 6, 7, 8]],
|
||
|
|
columns=["A", "B", "C", "D"],
|
||
|
|
index=Index(range(2), name="x"),
|
||
|
|
)
|
||
|
|
assert df.reset_index().index.name is None
|
||
|
|
assert df.reset_index(drop=True).index.name is None
|
||
|
|
return_value = df.reset_index(inplace=True)
|
||
|
|
assert return_value is None
|
||
|
|
assert df.index.name is None
|
||
|
|
|
||
|
|
def test_reset_index_level(self):
|
||
|
|
df = DataFrame([[1, 2, 3, 4], [5, 6, 7, 8]], columns=["A", "B", "C", "D"])
|
||
|
|
|
||
|
|
for levels in ["A", "B"], [0, 1]:
|
||
|
|
# With MultiIndex
|
||
|
|
result = df.set_index(["A", "B"]).reset_index(level=levels[0])
|
||
|
|
tm.assert_frame_equal(result, df.set_index("B"))
|
||
|
|
|
||
|
|
result = df.set_index(["A", "B"]).reset_index(level=levels[:1])
|
||
|
|
tm.assert_frame_equal(result, df.set_index("B"))
|
||
|
|
|
||
|
|
result = df.set_index(["A", "B"]).reset_index(level=levels)
|
||
|
|
tm.assert_frame_equal(result, df)
|
||
|
|
|
||
|
|
result = df.set_index(["A", "B"]).reset_index(level=levels, drop=True)
|
||
|
|
tm.assert_frame_equal(result, df[["C", "D"]])
|
||
|
|
|
||
|
|
# With single-level Index (GH 16263)
|
||
|
|
result = df.set_index("A").reset_index(level=levels[0])
|
||
|
|
tm.assert_frame_equal(result, df)
|
||
|
|
|
||
|
|
result = df.set_index("A").reset_index(level=levels[:1])
|
||
|
|
tm.assert_frame_equal(result, df)
|
||
|
|
|
||
|
|
result = df.set_index(["A"]).reset_index(level=levels[0], drop=True)
|
||
|
|
tm.assert_frame_equal(result, df[["B", "C", "D"]])
|
||
|
|
|
||
|
|
# Missing levels - for both MultiIndex and single-level Index:
|
||
|
|
for idx_lev in ["A", "B"], ["A"]:
|
||
|
|
with pytest.raises(KeyError, match=r"(L|l)evel \(?E\)?"):
|
||
|
|
df.set_index(idx_lev).reset_index(level=["A", "E"])
|
||
|
|
with pytest.raises(IndexError, match="Too many levels"):
|
||
|
|
df.set_index(idx_lev).reset_index(level=[0, 1, 2])
|
||
|
|
|
||
|
|
def test_reset_index_right_dtype(self):
|
||
|
|
time = np.arange(0.0, 10, np.sqrt(2) / 2)
|
||
|
|
s1 = Series(
|
||
|
|
(9.81 * time ** 2) / 2, index=Index(time, name="time"), name="speed"
|
||
|
|
)
|
||
|
|
df = DataFrame(s1)
|
||
|
|
|
||
|
|
resetted = s1.reset_index()
|
||
|
|
assert resetted["time"].dtype == np.float64
|
||
|
|
|
||
|
|
resetted = df.reset_index()
|
||
|
|
assert resetted["time"].dtype == np.float64
|
||
|
|
|
||
|
|
def test_reset_index_multiindex_col(self):
|
||
|
|
vals = np.random.randn(3, 3).astype(object)
|
||
|
|
idx = ["x", "y", "z"]
|
||
|
|
full = np.hstack(([[x] for x in idx], vals))
|
||
|
|
df = DataFrame(
|
||
|
|
vals,
|
||
|
|
Index(idx, name="a"),
|
||
|
|
columns=[["b", "b", "c"], ["mean", "median", "mean"]],
|
||
|
|
)
|
||
|
|
rs = df.reset_index()
|
||
|
|
xp = DataFrame(
|
||
|
|
full, columns=[["a", "b", "b", "c"], ["", "mean", "median", "mean"]]
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(rs, xp)
|
||
|
|
|
||
|
|
rs = df.reset_index(col_fill=None)
|
||
|
|
xp = DataFrame(
|
||
|
|
full, columns=[["a", "b", "b", "c"], ["a", "mean", "median", "mean"]]
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(rs, xp)
|
||
|
|
|
||
|
|
rs = df.reset_index(col_level=1, col_fill="blah")
|
||
|
|
xp = DataFrame(
|
||
|
|
full, columns=[["blah", "b", "b", "c"], ["a", "mean", "median", "mean"]]
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(rs, xp)
|
||
|
|
|
||
|
|
df = DataFrame(
|
||
|
|
vals,
|
||
|
|
MultiIndex.from_arrays([[0, 1, 2], ["x", "y", "z"]], names=["d", "a"]),
|
||
|
|
columns=[["b", "b", "c"], ["mean", "median", "mean"]],
|
||
|
|
)
|
||
|
|
rs = df.reset_index("a")
|
||
|
|
xp = DataFrame(
|
||
|
|
full,
|
||
|
|
Index([0, 1, 2], name="d"),
|
||
|
|
columns=[["a", "b", "b", "c"], ["", "mean", "median", "mean"]],
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(rs, xp)
|
||
|
|
|
||
|
|
rs = df.reset_index("a", col_fill=None)
|
||
|
|
xp = DataFrame(
|
||
|
|
full,
|
||
|
|
Index(range(3), name="d"),
|
||
|
|
columns=[["a", "b", "b", "c"], ["a", "mean", "median", "mean"]],
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(rs, xp)
|
||
|
|
|
||
|
|
rs = df.reset_index("a", col_fill="blah", col_level=1)
|
||
|
|
xp = DataFrame(
|
||
|
|
full,
|
||
|
|
Index(range(3), name="d"),
|
||
|
|
columns=[["blah", "b", "b", "c"], ["a", "mean", "median", "mean"]],
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(rs, xp)
|
||
|
|
|
||
|
|
def test_reset_index_multiindex_nan(self):
|
||
|
|
# GH#6322, testing reset_index on MultiIndexes
|
||
|
|
# when we have a nan or all nan
|
||
|
|
df = DataFrame(
|
||
|
|
{"A": ["a", "b", "c"], "B": [0, 1, np.nan], "C": np.random.rand(3)}
|
||
|
|
)
|
||
|
|
rs = df.set_index(["A", "B"]).reset_index()
|
||
|
|
tm.assert_frame_equal(rs, df)
|
||
|
|
|
||
|
|
df = DataFrame(
|
||
|
|
{"A": [np.nan, "b", "c"], "B": [0, 1, 2], "C": np.random.rand(3)}
|
||
|
|
)
|
||
|
|
rs = df.set_index(["A", "B"]).reset_index()
|
||
|
|
tm.assert_frame_equal(rs, df)
|
||
|
|
|
||
|
|
df = DataFrame({"A": ["a", "b", "c"], "B": [0, 1, 2], "C": [np.nan, 1.1, 2.2]})
|
||
|
|
rs = df.set_index(["A", "B"]).reset_index()
|
||
|
|
tm.assert_frame_equal(rs, df)
|
||
|
|
|
||
|
|
df = DataFrame(
|
||
|
|
{
|
||
|
|
"A": ["a", "b", "c"],
|
||
|
|
"B": [np.nan, np.nan, np.nan],
|
||
|
|
"C": np.random.rand(3),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
rs = df.set_index(["A", "B"]).reset_index()
|
||
|
|
tm.assert_frame_equal(rs, df)
|
||
|
|
|
||
|
|
def test_reset_index_with_datetimeindex_cols(self):
|
||
|
|
# GH#5818
|
||
|
|
df = DataFrame(
|
||
|
|
[[1, 2], [3, 4]],
|
||
|
|
columns=date_range("1/1/2013", "1/2/2013"),
|
||
|
|
index=["A", "B"],
|
||
|
|
)
|
||
|
|
|
||
|
|
result = df.reset_index()
|
||
|
|
expected = DataFrame(
|
||
|
|
[["A", 1, 2], ["B", 3, 4]],
|
||
|
|
columns=["index", datetime(2013, 1, 1), datetime(2013, 1, 2)],
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
def test_reset_index_range(self):
|
||
|
|
# GH#12071
|
||
|
|
df = DataFrame([[0, 0], [1, 1]], columns=["A", "B"], index=RangeIndex(stop=2))
|
||
|
|
result = df.reset_index()
|
||
|
|
assert isinstance(result.index, RangeIndex)
|
||
|
|
expected = DataFrame(
|
||
|
|
[[0, 0, 0], [1, 1, 1]],
|
||
|
|
columns=["index", "A", "B"],
|
||
|
|
index=RangeIndex(stop=2),
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
def test_reset_index_multiindex_columns(self):
|
||
|
|
levels = [["A", ""], ["B", "b"]]
|
||
|
|
df = DataFrame([[0, 2], [1, 3]], columns=MultiIndex.from_tuples(levels))
|
||
|
|
result = df[["B"]].rename_axis("A").reset_index()
|
||
|
|
tm.assert_frame_equal(result, df)
|
||
|
|
|
||
|
|
# GH#16120: already existing column
|
||
|
|
msg = r"cannot insert \('A', ''\), already exists"
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
df.rename_axis("A").reset_index()
|
||
|
|
|
||
|
|
# GH#16164: multiindex (tuple) full key
|
||
|
|
result = df.set_index([("A", "")]).reset_index()
|
||
|
|
tm.assert_frame_equal(result, df)
|
||
|
|
|
||
|
|
# with additional (unnamed) index level
|
||
|
|
idx_col = DataFrame(
|
||
|
|
[[0], [1]], columns=MultiIndex.from_tuples([("level_0", "")])
|
||
|
|
)
|
||
|
|
expected = pd.concat([idx_col, df[[("B", "b"), ("A", "")]]], axis=1)
|
||
|
|
result = df.set_index([("B", "b")], append=True).reset_index()
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
# with index name which is a too long tuple...
|
||
|
|
msg = "Item must have length equal to number of levels."
|
||
|
|
with pytest.raises(ValueError, match=msg):
|
||
|
|
df.rename_axis([("C", "c", "i")]).reset_index()
|
||
|
|
|
||
|
|
# or too short...
|
||
|
|
levels = [["A", "a", ""], ["B", "b", "i"]]
|
||
|
|
df2 = DataFrame([[0, 2], [1, 3]], columns=MultiIndex.from_tuples(levels))
|
||
|
|
idx_col = DataFrame(
|
||
|
|
[[0], [1]], columns=MultiIndex.from_tuples([("C", "c", "ii")])
|
||
|
|
)
|
||
|
|
expected = pd.concat([idx_col, df2], axis=1)
|
||
|
|
result = df2.rename_axis([("C", "c")]).reset_index(col_fill="ii")
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
# ... which is incompatible with col_fill=None
|
||
|
|
with pytest.raises(
|
||
|
|
ValueError,
|
||
|
|
match=(
|
||
|
|
"col_fill=None is incompatible with "
|
||
|
|
r"incomplete column name \('C', 'c'\)"
|
||
|
|
),
|
||
|
|
):
|
||
|
|
df2.rename_axis([("C", "c")]).reset_index(col_fill=None)
|
||
|
|
|
||
|
|
# with col_level != 0
|
||
|
|
result = df2.rename_axis([("c", "ii")]).reset_index(col_level=1, col_fill="C")
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
def test_reset_index_datetime(self, tz_naive_fixture):
|
||
|
|
# GH#3950
|
||
|
|
tz = tz_naive_fixture
|
||
|
|
idx1 = pd.date_range("1/1/2011", periods=5, freq="D", tz=tz, name="idx1")
|
||
|
|
idx2 = Index(range(5), name="idx2", dtype="int64")
|
||
|
|
idx = MultiIndex.from_arrays([idx1, idx2])
|
||
|
|
df = DataFrame(
|
||
|
|
{"a": np.arange(5, dtype="int64"), "b": ["A", "B", "C", "D", "E"]},
|
||
|
|
index=idx,
|
||
|
|
)
|
||
|
|
|
||
|
|
expected = DataFrame(
|
||
|
|
{
|
||
|
|
"idx1": [
|
||
|
|
datetime(2011, 1, 1),
|
||
|
|
datetime(2011, 1, 2),
|
||
|
|
datetime(2011, 1, 3),
|
||
|
|
datetime(2011, 1, 4),
|
||
|
|
datetime(2011, 1, 5),
|
||
|
|
],
|
||
|
|
"idx2": np.arange(5, dtype="int64"),
|
||
|
|
"a": np.arange(5, dtype="int64"),
|
||
|
|
"b": ["A", "B", "C", "D", "E"],
|
||
|
|
},
|
||
|
|
columns=["idx1", "idx2", "a", "b"],
|
||
|
|
)
|
||
|
|
expected["idx1"] = expected["idx1"].apply(lambda d: Timestamp(d, tz=tz))
|
||
|
|
|
||
|
|
tm.assert_frame_equal(df.reset_index(), expected)
|
||
|
|
|
||
|
|
idx3 = pd.date_range(
|
||
|
|
"1/1/2012", periods=5, freq="MS", tz="Europe/Paris", name="idx3"
|
||
|
|
)
|
||
|
|
idx = MultiIndex.from_arrays([idx1, idx2, idx3])
|
||
|
|
df = DataFrame(
|
||
|
|
{"a": np.arange(5, dtype="int64"), "b": ["A", "B", "C", "D", "E"]},
|
||
|
|
index=idx,
|
||
|
|
)
|
||
|
|
|
||
|
|
expected = DataFrame(
|
||
|
|
{
|
||
|
|
"idx1": [
|
||
|
|
datetime(2011, 1, 1),
|
||
|
|
datetime(2011, 1, 2),
|
||
|
|
datetime(2011, 1, 3),
|
||
|
|
datetime(2011, 1, 4),
|
||
|
|
datetime(2011, 1, 5),
|
||
|
|
],
|
||
|
|
"idx2": np.arange(5, dtype="int64"),
|
||
|
|
"idx3": [
|
||
|
|
datetime(2012, 1, 1),
|
||
|
|
datetime(2012, 2, 1),
|
||
|
|
datetime(2012, 3, 1),
|
||
|
|
datetime(2012, 4, 1),
|
||
|
|
datetime(2012, 5, 1),
|
||
|
|
],
|
||
|
|
"a": np.arange(5, dtype="int64"),
|
||
|
|
"b": ["A", "B", "C", "D", "E"],
|
||
|
|
},
|
||
|
|
columns=["idx1", "idx2", "idx3", "a", "b"],
|
||
|
|
)
|
||
|
|
expected["idx1"] = expected["idx1"].apply(lambda d: Timestamp(d, tz=tz))
|
||
|
|
expected["idx3"] = expected["idx3"].apply(
|
||
|
|
lambda d: Timestamp(d, tz="Europe/Paris")
|
||
|
|
)
|
||
|
|
tm.assert_frame_equal(df.reset_index(), expected)
|
||
|
|
|
||
|
|
# GH#7793
|
||
|
|
idx = MultiIndex.from_product(
|
||
|
|
[["a", "b"], pd.date_range("20130101", periods=3, tz=tz)]
|
||
|
|
)
|
||
|
|
df = DataFrame(
|
||
|
|
np.arange(6, dtype="int64").reshape(6, 1), columns=["a"], index=idx
|
||
|
|
)
|
||
|
|
|
||
|
|
expected = DataFrame(
|
||
|
|
{
|
||
|
|
"level_0": "a a a b b b".split(),
|
||
|
|
"level_1": [
|
||
|
|
datetime(2013, 1, 1),
|
||
|
|
datetime(2013, 1, 2),
|
||
|
|
datetime(2013, 1, 3),
|
||
|
|
]
|
||
|
|
* 2,
|
||
|
|
"a": np.arange(6, dtype="int64"),
|
||
|
|
},
|
||
|
|
columns=["level_0", "level_1", "a"],
|
||
|
|
)
|
||
|
|
expected["level_1"] = expected["level_1"].apply(
|
||
|
|
lambda d: Timestamp(d, freq="D", tz=tz)
|
||
|
|
)
|
||
|
|
result = df.reset_index()
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
def test_reset_index_period(self):
|
||
|
|
# GH#7746
|
||
|
|
idx = MultiIndex.from_product(
|
||
|
|
[pd.period_range("20130101", periods=3, freq="M"), list("abc")],
|
||
|
|
names=["month", "feature"],
|
||
|
|
)
|
||
|
|
|
||
|
|
df = DataFrame(
|
||
|
|
np.arange(9, dtype="int64").reshape(-1, 1), index=idx, columns=["a"]
|
||
|
|
)
|
||
|
|
expected = DataFrame(
|
||
|
|
{
|
||
|
|
"month": (
|
||
|
|
[pd.Period("2013-01", freq="M")] * 3
|
||
|
|
+ [pd.Period("2013-02", freq="M")] * 3
|
||
|
|
+ [pd.Period("2013-03", freq="M")] * 3
|
||
|
|
),
|
||
|
|
"feature": ["a", "b", "c"] * 3,
|
||
|
|
"a": np.arange(9, dtype="int64"),
|
||
|
|
},
|
||
|
|
columns=["month", "feature", "a"],
|
||
|
|
)
|
||
|
|
result = df.reset_index()
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
def test_reset_index_delevel_infer_dtype(self):
|
||
|
|
tuples = list(product(["foo", "bar"], [10, 20], [1.0, 1.1]))
|
||
|
|
index = MultiIndex.from_tuples(tuples, names=["prm0", "prm1", "prm2"])
|
||
|
|
df = DataFrame(np.random.randn(8, 3), columns=["A", "B", "C"], index=index)
|
||
|
|
deleveled = df.reset_index()
|
||
|
|
assert is_integer_dtype(deleveled["prm1"])
|
||
|
|
assert is_float_dtype(deleveled["prm2"])
|
||
|
|
|
||
|
|
def test_reset_index_with_drop(
|
||
|
|
self, multiindex_year_month_day_dataframe_random_data
|
||
|
|
):
|
||
|
|
ymd = multiindex_year_month_day_dataframe_random_data
|
||
|
|
|
||
|
|
deleveled = ymd.reset_index(drop=True)
|
||
|
|
assert len(deleveled.columns) == len(ymd.columns)
|
||
|
|
assert deleveled.index.name == ymd.index.name
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"ix_data, exp_data",
|
||
|
|
[
|
||
|
|
(
|
||
|
|
[(pd.NaT, 1), (pd.NaT, 2)],
|
||
|
|
{"a": [pd.NaT, pd.NaT], "b": [1, 2], "x": [11, 12]},
|
||
|
|
),
|
||
|
|
(
|
||
|
|
[(pd.NaT, 1), (Timestamp("2020-01-01"), 2)],
|
||
|
|
{"a": [pd.NaT, Timestamp("2020-01-01")], "b": [1, 2], "x": [11, 12]},
|
||
|
|
),
|
||
|
|
(
|
||
|
|
[(pd.NaT, 1), (pd.Timedelta(123, "d"), 2)],
|
||
|
|
{"a": [pd.NaT, pd.Timedelta(123, "d")], "b": [1, 2], "x": [11, 12]},
|
||
|
|
),
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_reset_index_nat_multiindex(self, ix_data, exp_data):
|
||
|
|
# GH#36541: that reset_index() does not raise ValueError
|
||
|
|
ix = MultiIndex.from_tuples(ix_data, names=["a", "b"])
|
||
|
|
result = DataFrame({"x": [11, 12]}, index=ix)
|
||
|
|
result = result.reset_index()
|
||
|
|
|
||
|
|
expected = DataFrame(exp_data)
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"codes", ([[0, 0, 1, 1], [0, 1, 0, 1]], [[0, 0, -1, 1], [0, 1, 0, 1]])
|
||
|
|
)
|
||
|
|
def test_rest_index_multiindex_categorical_with_missing_values(self, codes):
|
||
|
|
# GH#24206
|
||
|
|
|
||
|
|
index = MultiIndex(
|
||
|
|
[CategoricalIndex(["A", "B"]), CategoricalIndex(["a", "b"])], codes
|
||
|
|
)
|
||
|
|
data = {"col": range(len(index))}
|
||
|
|
df = DataFrame(data=data, index=index)
|
||
|
|
|
||
|
|
expected = DataFrame(
|
||
|
|
{
|
||
|
|
"level_0": Categorical.from_codes(codes[0], categories=["A", "B"]),
|
||
|
|
"level_1": Categorical.from_codes(codes[1], categories=["a", "b"]),
|
||
|
|
"col": range(4),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
res = df.reset_index()
|
||
|
|
tm.assert_frame_equal(res, expected)
|
||
|
|
|
||
|
|
# roundtrip
|
||
|
|
res = expected.set_index(["level_0", "level_1"]).reset_index()
|
||
|
|
tm.assert_frame_equal(res, expected)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.mark.parametrize(
|
||
|
|
"array, dtype",
|
||
|
|
[
|
||
|
|
(["a", "b"], object),
|
||
|
|
(
|
||
|
|
pd.period_range("12-1-2000", periods=2, freq="Q-DEC"),
|
||
|
|
pd.PeriodDtype(freq="Q-DEC"),
|
||
|
|
),
|
||
|
|
],
|
||
|
|
)
|
||
|
|
def test_reset_index_dtypes_on_empty_frame_with_multiindex(array, dtype):
|
||
|
|
# GH 19602 - Preserve dtype on empty DataFrame with MultiIndex
|
||
|
|
idx = MultiIndex.from_product([[0, 1], [0.5, 1.0], array])
|
||
|
|
result = DataFrame(index=idx)[:0].reset_index().dtypes
|
||
|
|
expected = Series({"level_0": np.int64, "level_1": np.float64, "level_2": dtype})
|
||
|
|
tm.assert_series_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_reset_index_empty_frame_with_datetime64_multiindex():
|
||
|
|
# https://github.com/pandas-dev/pandas/issues/35606
|
||
|
|
idx = MultiIndex(
|
||
|
|
levels=[[Timestamp("2020-07-20 00:00:00")], [3, 4]],
|
||
|
|
codes=[[], []],
|
||
|
|
names=["a", "b"],
|
||
|
|
)
|
||
|
|
df = DataFrame(index=idx, columns=["c", "d"])
|
||
|
|
result = df.reset_index()
|
||
|
|
expected = DataFrame(
|
||
|
|
columns=list("abcd"), index=RangeIndex(start=0, stop=0, step=1)
|
||
|
|
)
|
||
|
|
expected["a"] = expected["a"].astype("datetime64[ns]")
|
||
|
|
expected["b"] = expected["b"].astype("int64")
|
||
|
|
tm.assert_frame_equal(result, expected)
|
||
|
|
|
||
|
|
|
||
|
|
def test_reset_index_empty_frame_with_datetime64_multiindex_from_groupby():
|
||
|
|
# https://github.com/pandas-dev/pandas/issues/35657
|
||
|
|
df = DataFrame({"c1": [10.0], "c2": ["a"], "c3": pd.to_datetime("2020-01-01")})
|
||
|
|
df = df.head(0).groupby(["c2", "c3"])[["c1"]].sum()
|
||
|
|
result = df.reset_index()
|
||
|
|
expected = DataFrame(
|
||
|
|
columns=["c2", "c3", "c1"], index=RangeIndex(start=0, stop=0, step=1)
|
||
|
|
)
|
||
|
|
expected["c3"] = expected["c3"].astype("datetime64[ns]")
|
||
|
|
expected["c1"] = expected["c1"].astype("float64")
|
||
|
|
tm.assert_frame_equal(result, expected)
|