66 lines
2.4 KiB
Python
66 lines
2.4 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
import pandas as pd
|
|
from pandas import _np_version_under1p17
|
|
import pandas.util.testing as tm
|
|
|
|
|
|
@pytest.mark.filterwarnings("ignore:Sparse:FutureWarning")
|
|
@pytest.mark.filterwarnings("ignore:Series.to_sparse:FutureWarning")
|
|
@pytest.mark.filterwarnings("ignore:DataFrame.to_sparse:FutureWarning")
|
|
class TestPivotTable:
|
|
def setup_method(self, method):
|
|
rs = np.random.RandomState(0)
|
|
self.dense = pd.DataFrame(
|
|
{
|
|
"A": ["foo", "bar", "foo", "bar", "foo", "bar", "foo", "foo"],
|
|
"B": ["one", "one", "two", "three", "two", "two", "one", "three"],
|
|
"C": rs.randn(8),
|
|
"D": rs.randn(8),
|
|
"E": [np.nan, np.nan, 1, 2, np.nan, 1, np.nan, np.nan],
|
|
}
|
|
)
|
|
self.sparse = self.dense.to_sparse()
|
|
|
|
def test_pivot_table(self):
|
|
res_sparse = pd.pivot_table(self.sparse, index="A", columns="B", values="C")
|
|
res_dense = pd.pivot_table(self.dense, index="A", columns="B", values="C")
|
|
tm.assert_frame_equal(res_sparse, res_dense)
|
|
|
|
res_sparse = pd.pivot_table(self.sparse, index="A", columns="B", values="E")
|
|
res_dense = pd.pivot_table(self.dense, index="A", columns="B", values="E")
|
|
tm.assert_frame_equal(res_sparse, res_dense)
|
|
|
|
res_sparse = pd.pivot_table(
|
|
self.sparse, index="A", columns="B", values="E", aggfunc="mean"
|
|
)
|
|
res_dense = pd.pivot_table(
|
|
self.dense, index="A", columns="B", values="E", aggfunc="mean"
|
|
)
|
|
tm.assert_frame_equal(res_sparse, res_dense)
|
|
|
|
def test_pivot_table_with_nans(self):
|
|
res_sparse = pd.pivot_table(
|
|
self.sparse, index="A", columns="B", values="E", aggfunc="sum"
|
|
)
|
|
res_dense = pd.pivot_table(
|
|
self.dense, index="A", columns="B", values="E", aggfunc="sum"
|
|
)
|
|
tm.assert_frame_equal(res_sparse, res_dense)
|
|
|
|
@pytest.mark.xfail(
|
|
not _np_version_under1p17,
|
|
reason="failing occasionally on numpy > 1.17",
|
|
strict=False,
|
|
)
|
|
def test_pivot_table_multi(self):
|
|
res_sparse = pd.pivot_table(
|
|
self.sparse, index="A", columns="B", values=["D", "E"]
|
|
)
|
|
res_dense = pd.pivot_table(
|
|
self.dense, index="A", columns="B", values=["D", "E"]
|
|
)
|
|
res_dense = res_dense.apply(lambda x: x.astype("Sparse[float64]"))
|
|
tm.assert_frame_equal(res_sparse, res_dense)
|