8th day of python challenges 111-117

venv/lib/python3.6/site-packages/pandas/tests/sparse/test_pivot.py

@@ -0,0 +1,65 @@
+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)