Skip to content

Commit 50171fc

Browse files
authored
TST: slow pivot_table test (#56042)
1 parent 2393a23 commit 50171fc

File tree

1 file changed

+21
-22
lines changed

1 file changed

+21
-22
lines changed

pandas/tests/reshape/test_pivot.py

+21-22
Original file line numberDiff line numberDiff line change
@@ -1715,39 +1715,38 @@ def test_pivot_dtaccessor(self):
17151715
)
17161716
tm.assert_frame_equal(result, expected)
17171717

1718-
@pytest.mark.parametrize("i", range(1, 367))
1719-
def test_daily(self, i):
1718+
def test_daily(self):
17201719
rng = date_range("1/1/2000", "12/31/2004", freq="D")
1721-
ts = Series(np.random.default_rng(2).standard_normal(len(rng)), index=rng)
1720+
ts = Series(np.arange(len(rng)), index=rng)
17221721

1723-
annual = pivot_table(
1722+
result = pivot_table(
17241723
DataFrame(ts), index=ts.index.year, columns=ts.index.dayofyear
17251724
)
1726-
annual.columns = annual.columns.droplevel(0)
1725+
result.columns = result.columns.droplevel(0)
17271726

17281727
doy = np.asarray(ts.index.dayofyear)
17291728

1730-
subset = ts[doy == i]
1731-
subset.index = subset.index.year
1732-
1733-
result = annual[i].dropna()
1734-
tm.assert_series_equal(result, subset, check_names=False)
1735-
assert result.name == i
1729+
expected = {}
1730+
for y in ts.index.year.unique().values:
1731+
mask = ts.index.year == y
1732+
expected[y] = Series(ts.values[mask], index=doy[mask])
1733+
expected = DataFrame(expected, dtype=float).T
1734+
tm.assert_frame_equal(result, expected)
17361735

1737-
@pytest.mark.parametrize("i", range(1, 13))
1738-
def test_monthly(self, i):
1736+
def test_monthly(self):
17391737
rng = date_range("1/1/2000", "12/31/2004", freq="ME")
1740-
ts = Series(np.random.default_rng(2).standard_normal(len(rng)), index=rng)
1738+
ts = Series(np.arange(len(rng)), index=rng)
17411739

1742-
annual = pivot_table(DataFrame(ts), index=ts.index.year, columns=ts.index.month)
1743-
annual.columns = annual.columns.droplevel(0)
1740+
result = pivot_table(DataFrame(ts), index=ts.index.year, columns=ts.index.month)
1741+
result.columns = result.columns.droplevel(0)
17441742

1745-
month = ts.index.month
1746-
subset = ts[month == i]
1747-
subset.index = subset.index.year
1748-
result = annual[i].dropna()
1749-
tm.assert_series_equal(result, subset, check_names=False)
1750-
assert result.name == i
1743+
month = np.asarray(ts.index.month)
1744+
expected = {}
1745+
for y in ts.index.year.unique().values:
1746+
mask = ts.index.year == y
1747+
expected[y] = Series(ts.values[mask], index=month[mask])
1748+
expected = DataFrame(expected, dtype=float).T
1749+
tm.assert_frame_equal(result, expected)
17511750

17521751
def test_pivot_table_with_iterator_values(self, data):
17531752
# GH 12017

0 commit comments

Comments
 (0)