snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/pandas/tests/groupby/test_missing_data.py

@@ -0,0 +1,421 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from distutils.version import LooseVersion
+import unittest
+
+import numpy as np
+import pandas as pd
+
+from pyspark import pandas as ps
+from pyspark.testing.pandasutils import ComparisonTestBase
+from pyspark.testing.sqlutils import SQLTestUtils
+
+
+class GroupbyMissingDataMixin:
+    @property
+    def pdf(self):
+        return pd.DataFrame(
+            {
+                "A": [1, 2, 1, 2],
+                "B": [3.1, 4.1, 4.1, 3.1],
+                "C": ["a", "b", "b", "a"],
+                "D": [True, False, False, True],
+            }
+        )
+
+    @property
+    def psdf(self):
+        return ps.from_pandas(self.pdf)
+
+    def test_fillna(self):
+        pdf = pd.DataFrame(
+            {
+                "A": [1, 1, 2, 2] * 3,
+                "B": [2, 4, None, 3] * 3,
+                "C": [None, None, None, 1] * 3,
+                "D": [0, 1, 5, 4] * 3,
+            }
+        )
+        psdf = ps.from_pandas(pdf)
+
+        self.assert_eq(
+            psdf.groupby("A").fillna(0).sort_index(), pdf.groupby("A").fillna(0).sort_index()
+        )
+        self.assert_eq(
+            psdf.groupby("A")["C"].fillna(0).sort_index(),
+            pdf.groupby("A")["C"].fillna(0).sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby("A")[["C"]].fillna(0).sort_index(),
+            pdf.groupby("A")[["C"]].fillna(0).sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby("A").fillna(method="bfill").sort_index(),
+            pdf.groupby("A").fillna(method="bfill").sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby("A")["C"].fillna(method="bfill").sort_index(),
+            pdf.groupby("A")["C"].fillna(method="bfill").sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby("A")[["C"]].fillna(method="bfill").sort_index(),
+            pdf.groupby("A")[["C"]].fillna(method="bfill").sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby("A").fillna(method="ffill").sort_index(),
+            pdf.groupby("A").fillna(method="ffill").sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby("A")["C"].fillna(method="ffill").sort_index(),
+            pdf.groupby("A")["C"].fillna(method="ffill").sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby("A")[["C"]].fillna(method="ffill").sort_index(),
+            pdf.groupby("A")[["C"]].fillna(method="ffill").sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby(psdf.A // 5).fillna(method="bfill").sort_index(),
+            pdf.groupby(pdf.A // 5).fillna(method="bfill").sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby(psdf.A // 5)["C"].fillna(method="bfill").sort_index(),
+            pdf.groupby(pdf.A // 5)["C"].fillna(method="bfill").sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby(psdf.A // 5)[["C"]].fillna(method="bfill").sort_index(),
+            pdf.groupby(pdf.A // 5)[["C"]].fillna(method="bfill").sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby(psdf.A // 5).fillna(method="ffill").sort_index(),
+            pdf.groupby(pdf.A // 5).fillna(method="ffill").sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby(psdf.A // 5)["C"].fillna(method="ffill").sort_index(),
+            pdf.groupby(pdf.A // 5)["C"].fillna(method="ffill").sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby(psdf.A // 5)[["C"]].fillna(method="ffill").sort_index(),
+            pdf.groupby(pdf.A // 5)[["C"]].fillna(method="ffill").sort_index(),
+        )
+        self.assert_eq(
+            psdf.C.rename().groupby(psdf.A).fillna(0).sort_index(),
+            pdf.C.rename().groupby(pdf.A).fillna(0).sort_index(),
+        )
+        self.assert_eq(
+            psdf.C.groupby(psdf.A.rename()).fillna(0).sort_index(),
+            pdf.C.groupby(pdf.A.rename()).fillna(0).sort_index(),
+        )
+        self.assert_eq(
+            psdf.C.rename().groupby(psdf.A.rename()).fillna(0).sort_index(),
+            pdf.C.rename().groupby(pdf.A.rename()).fillna(0).sort_index(),
+        )
+
+        # multi-index columns
+        columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C"), ("Z", "D")])
+        pdf.columns = columns
+        psdf.columns = columns
+
+        self.assert_eq(
+            psdf.groupby(("X", "A")).fillna(0).sort_index(),
+            pdf.groupby(("X", "A")).fillna(0).sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby(("X", "A")).fillna(method="bfill").sort_index(),
+            pdf.groupby(("X", "A")).fillna(method="bfill").sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby(("X", "A")).fillna(method="ffill").sort_index(),
+            pdf.groupby(("X", "A")).fillna(method="ffill").sort_index(),
+        )
+
+    def test_ffill(self):
+        idx = np.random.rand(4 * 3)
+        pdf = pd.DataFrame(
+            {
+                "A": [1, 1, 2, 2] * 3,
+                "B": [2, 4, None, 3] * 3,
+                "C": [None, None, None, 1] * 3,
+                "D": [0, 1, 5, 4] * 3,
+            },
+            index=idx,
+        )
+        psdf = ps.from_pandas(pdf)
+
+        self.assert_eq(
+            psdf.groupby("A").ffill().sort_index(), pdf.groupby("A").ffill().sort_index()
+        )
+        self.assert_eq(
+            psdf.groupby("A")[["B"]].ffill().sort_index(),
+            pdf.groupby("A")[["B"]].ffill().sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby("A")["B"].ffill().sort_index(), pdf.groupby("A")["B"].ffill().sort_index()
+        )
+        self.assert_eq(
+            psdf.groupby("A")["B"].ffill()[idx[6]], pdf.groupby("A")["B"].ffill()[idx[6]]
+        )
+
+        # multi-index columns
+        columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C"), ("Z", "D")])
+        pdf.columns = columns
+        psdf.columns = columns
+
+        self.assert_eq(
+            psdf.groupby(("X", "A")).ffill().sort_index(),
+            pdf.groupby(("X", "A")).ffill().sort_index(),
+        )
+
+    def test_bfill(self):
+        idx = np.random.rand(4 * 3)
+        pdf = pd.DataFrame(
+            {
+                "A": [1, 1, 2, 2] * 3,
+                "B": [2, 4, None, 3] * 3,
+                "C": [None, None, None, 1] * 3,
+                "D": [0, 1, 5, 4] * 3,
+            },
+            index=idx,
+        )
+        psdf = ps.from_pandas(pdf)
+
+        self.assert_eq(
+            psdf.groupby("A").bfill().sort_index(), pdf.groupby("A").bfill().sort_index()
+        )
+        self.assert_eq(
+            psdf.groupby("A")[["B"]].bfill().sort_index(),
+            pdf.groupby("A")[["B"]].bfill().sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby("A")["B"].bfill().sort_index(),
+            pdf.groupby("A")["B"].bfill().sort_index(),
+        )
+        self.assert_eq(
+            psdf.groupby("A")["B"].bfill()[idx[6]], pdf.groupby("A")["B"].bfill()[idx[6]]
+        )
+
+        # multi-index columns
+        columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C"), ("Z", "D")])
+        pdf.columns = columns
+        psdf.columns = columns
+
+        self.assert_eq(
+            psdf.groupby(("X", "A")).bfill().sort_index(),
+            pdf.groupby(("X", "A")).bfill().sort_index(),
+        )
+
+    def test_dropna(self):
+        pdf = pd.DataFrame(
+            {"A": [None, 1, None, 1, 2], "B": [1, 2, 3, None, None], "C": [4, 5, 6, 7, None]}
+        )
+        psdf = ps.from_pandas(pdf)
+
+        # pd.DataFrame.groupby with dropna parameter is implemented since pandas 1.1.0
+        if LooseVersion(pd.__version__) >= LooseVersion("1.1.0"):
+            for dropna in [True, False]:
+                for as_index in [True, False]:
+                    if as_index:
+
+                        def sort(df):
+                            return df.sort_index()
+
+                    else:
+
+                        def sort(df):
+                            return df.sort_values("A").reset_index(drop=True)
+
+                    self.assert_eq(
+                        sort(psdf.groupby("A", as_index=as_index, dropna=dropna).std()),
+                        sort(pdf.groupby("A", as_index=as_index, dropna=dropna).std()),
+                    )
+
+                    self.assert_eq(
+                        sort(psdf.groupby("A", as_index=as_index, dropna=dropna).B.std()),
+                        sort(pdf.groupby("A", as_index=as_index, dropna=dropna).B.std()),
+                    )
+                    self.assert_eq(
+                        sort(psdf.groupby("A", as_index=as_index, dropna=dropna)["B"].std()),
+                        sort(pdf.groupby("A", as_index=as_index, dropna=dropna)["B"].std()),
+                    )
+
+                    self.assert_eq(
+                        sort(
+                            psdf.groupby("A", as_index=as_index, dropna=dropna).agg(
+                                {"B": "min", "C": "std"}
+                            )
+                        ),
+                        sort(
+                            pdf.groupby("A", as_index=as_index, dropna=dropna).agg(
+                                {"B": "min", "C": "std"}
+                            )
+                        ),
+                    )
+
+            for dropna in [True, False]:
+                for as_index in [True, False]:
+                    if as_index:
+
+                        def sort(df):
+                            return df.sort_index()
+
+                    else:
+
+                        def sort(df):
+                            return df.sort_values(["A", "B"]).reset_index(drop=True)
+
+                    self.assert_eq(
+                        sort(
+                            psdf.groupby(["A", "B"], as_index=as_index, dropna=dropna).agg(
+                                {"C": ["min", "std"]}
+                            )
+                        ),
+                        sort(
+                            pdf.groupby(["A", "B"], as_index=as_index, dropna=dropna).agg(
+                                {"C": ["min", "std"]}
+                            )
+                        ),
+                        almost=True,
+                    )
+
+            # multi-index columns
+            columns = pd.MultiIndex.from_tuples([("X", "A"), ("X", "B"), ("Y", "C")])
+            pdf.columns = columns
+            psdf.columns = columns
+
+            for dropna in [True, False]:
+                for as_index in [True, False]:
+                    if as_index:
+
+                        def sort(df):
+                            return df.sort_index()
+
+                    else:
+
+                        def sort(df):
+                            return df.sort_values(("X", "A")).reset_index(drop=True)
+
+                    sorted_stats_psdf = sort(
+                        psdf.groupby(("X", "A"), as_index=as_index, dropna=dropna).agg(
+                            {("X", "B"): "min", ("Y", "C"): "std"}
+                        )
+                    )
+                    sorted_stats_pdf = sort(
+                        pdf.groupby(("X", "A"), as_index=as_index, dropna=dropna).agg(
+                            {("X", "B"): "min", ("Y", "C"): "std"}
+                        )
+                    )
+                    self.assert_eq(sorted_stats_psdf, sorted_stats_pdf)
+        else:
+            # Testing dropna=True (pandas default behavior)
+            for as_index in [True, False]:
+                if as_index:
+
+                    def sort(df):
+                        return df.sort_index()
+
+                else:
+
+                    def sort(df):
+                        return df.sort_values("A").reset_index(drop=True)
+
+                self.assert_eq(
+                    sort(psdf.groupby("A", as_index=as_index, dropna=True)["B"].min()),
+                    sort(pdf.groupby("A", as_index=as_index)["B"].min()),
+                )
+
+                if as_index:
+
+                    def sort(df):
+                        return df.sort_index()
+
+                else:
+
+                    def sort(df):
+                        return df.sort_values(["A", "B"]).reset_index(drop=True)
+
+                self.assert_eq(
+                    sort(
+                        psdf.groupby(["A", "B"], as_index=as_index, dropna=True).agg(
+                            {"C": ["min", "std"]}
+                        )
+                    ),
+                    sort(pdf.groupby(["A", "B"], as_index=as_index).agg({"C": ["min", "std"]})),
+                    almost=True,
+                )
+
+            # Testing dropna=False
+            index = pd.Index([1.0, 2.0, np.nan], name="A")
+            expected = pd.Series([2.0, np.nan, 1.0], index=index, name="B")
+            result = psdf.groupby("A", as_index=True, dropna=False)["B"].min().sort_index()
+            self.assert_eq(expected, result)
+
+            expected = pd.DataFrame({"A": [1.0, 2.0, np.nan], "B": [2.0, np.nan, 1.0]})
+            result = (
+                psdf.groupby("A", as_index=False, dropna=False)["B"]
+                .min()
+                .sort_values("A")
+                .reset_index(drop=True)
+            )
+            self.assert_eq(expected, result)
+
+            index = pd.MultiIndex.from_tuples(
+                [(1.0, 2.0), (1.0, None), (2.0, None), (None, 1.0), (None, 3.0)], names=["A", "B"]
+            )
+            expected = pd.DataFrame(
+                {
+                    ("C", "min"): [5.0, 7.0, np.nan, 4.0, 6.0],
+                    ("C", "std"): [np.nan, np.nan, np.nan, np.nan, np.nan],
+                },
+                index=index,
+            )
+            result = (
+                psdf.groupby(["A", "B"], as_index=True, dropna=False)
+                .agg({"C": ["min", "std"]})
+                .sort_index()
+            )
+            self.assert_eq(expected, result)
+
+            expected = pd.DataFrame(
+                {
+                    ("A", ""): [1.0, 1.0, 2.0, np.nan, np.nan],
+                    ("B", ""): [2.0, np.nan, np.nan, 1.0, 3.0],
+                    ("C", "min"): [5.0, 7.0, np.nan, 4.0, 6.0],
+                    ("C", "std"): [np.nan, np.nan, np.nan, np.nan, np.nan],
+                }
+            )
+            result = (
+                psdf.groupby(["A", "B"], as_index=False, dropna=False)
+                .agg({"C": ["min", "std"]})
+                .sort_values(["A", "B"])
+                .reset_index(drop=True)
+            )
+            self.assert_eq(expected, result)
+
+
+class GroupbyMissingDataTests(GroupbyMissingDataMixin, ComparisonTestBase, SQLTestUtils):
+    pass
+
+
+if __name__ == "__main__":
+    from pyspark.pandas.tests.groupby.test_missing_data import *  # noqa: F401
+
+    try:
+        import xmlrunner
+
+        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
+    except ImportError:
+        testRunner = None
+    unittest.main(testRunner=testRunner, verbosity=2)

snowflake/snowpark_connect/includes/python/pyspark/pandas/tests/groupby/test_split_apply.py

@@ -0,0 +1,187 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from distutils.version import LooseVersion
+import unittest
+
+import pandas as pd
+
+from pyspark import pandas as ps
+from pyspark.testing.pandasutils import ComparisonTestBase
+from pyspark.testing.sqlutils import SQLTestUtils
+
+
+class GroupbySplitApplyMixin:
+    @property
+    def pdf(self):
+        return pd.DataFrame(
+            {
+                "A": [1, 2, 1, 2],
+                "B": [3.1, 4.1, 4.1, 3.1],
+                "C": ["a", "b", "b", "a"],
+                "D": [True, False, False, True],
+            }
+        )
+
+    @property
+    def psdf(self):
+        return ps.from_pandas(self.pdf)
+
+    @unittest.skipIf(
+        LooseVersion(pd.__version__) >= LooseVersion("2.0.0"),
+        "TODO(SPARK-43445): Enable GroupBySlowTests.test_split_apply_combine_on_series "
+        "for pandas 2.0.0.",
+    )
+    def test_split_apply_combine_on_series(self):
+        pdf = pd.DataFrame(
+            {
+                "a": [1, 2, 6, 4, 4, 6, 4, 3, 7],
+                "b": [4, 2, 7, 3, 3, 1, 1, 1, 2],
+                "c": [4, 2, 7, 3, None, 1, 1, 1, 2],
+                "d": list("abcdefght"),
+            },
+            index=[0, 1, 3, 5, 6, 8, 9, 9, 9],
+        )
+        psdf = ps.from_pandas(pdf)
+
+        funcs = [
+            ((True, False), ["sum", "min", "max", "count", "first", "last"]),
+            ((True, True), ["mean"]),
+            ((False, False), ["var", "std", "skew"]),
+        ]
+        funcs = [(check_exact, almost, f) for (check_exact, almost), fs in funcs for f in fs]
+
+        for as_index in [True, False]:
+            if as_index:
+
+                def sort(df):
+                    return df.sort_index()
+
+            else:
+
+                def sort(df):
+                    return df.sort_values(list(df.columns)).reset_index(drop=True)
+
+            for check_exact, almost, func in funcs:
+                for kkey, pkey in [("b", "b"), (psdf.b, pdf.b)]:
+                    with self.subTest(as_index=as_index, func=func, key=pkey):
+                        if as_index is True or func != "std":
+                            self.assert_eq(
+                                sort(getattr(psdf.groupby(kkey, as_index=as_index).a, func)()),
+                                sort(getattr(pdf.groupby(pkey, as_index=as_index).a, func)()),
+                                check_exact=check_exact,
+                                almost=almost,
+                            )
+                            self.assert_eq(
+                                sort(getattr(psdf.groupby(kkey, as_index=as_index), func)()),
+                                sort(getattr(pdf.groupby(pkey, as_index=as_index), func)()),
+                                check_exact=check_exact,
+                                almost=almost,
+                            )
+                        else:
+                            # seems like a pandas' bug for as_index=False and func == "std"?
+                            self.assert_eq(
+                                sort(getattr(psdf.groupby(kkey, as_index=as_index).a, func)()),
+                                sort(pdf.groupby(pkey, as_index=True).a.std().reset_index()),
+                                check_exact=check_exact,
+                                almost=almost,
+                            )
+                            self.assert_eq(
+                                sort(getattr(psdf.groupby(kkey, as_index=as_index), func)()),
+                                sort(pdf.groupby(pkey, as_index=True).std().reset_index()),
+                                check_exact=check_exact,
+                                almost=almost,
+                            )
+
+                for kkey, pkey in [(psdf.b + 1, pdf.b + 1), (psdf.copy().b, pdf.copy().b)]:
+                    with self.subTest(as_index=as_index, func=func, key=pkey):
+                        self.assert_eq(
+                            sort(getattr(psdf.groupby(kkey, as_index=as_index).a, func)()),
+                            sort(getattr(pdf.groupby(pkey, as_index=as_index).a, func)()),
+                            check_exact=check_exact,
+                            almost=almost,
+                        )
+                        self.assert_eq(
+                            sort(getattr(psdf.groupby(kkey, as_index=as_index), func)()),
+                            sort(getattr(pdf.groupby(pkey, as_index=as_index), func)()),
+                            check_exact=check_exact,
+                            almost=almost,
+                        )
+
+            for check_exact, almost, func in funcs:
+                for i in [0, 4, 7]:
+                    with self.subTest(as_index=as_index, func=func, i=i):
+                        self.assert_eq(
+                            sort(getattr(psdf.groupby(psdf.b > i, as_index=as_index).a, func)()),
+                            sort(getattr(pdf.groupby(pdf.b > i, as_index=as_index).a, func)()),
+                            check_exact=check_exact,
+                            almost=almost,
+                        )
+                        self.assert_eq(
+                            sort(getattr(psdf.groupby(psdf.b > i, as_index=as_index), func)()),
+                            sort(getattr(pdf.groupby(pdf.b > i, as_index=as_index), func)()),
+                            check_exact=check_exact,
+                            almost=almost,
+                        )
+
+        for check_exact, almost, func in funcs:
+            for kkey, pkey in [
+                (psdf.b, pdf.b),
+                (psdf.b + 1, pdf.b + 1),
+                (psdf.copy().b, pdf.copy().b),
+                (psdf.b.rename(), pdf.b.rename()),
+            ]:
+                with self.subTest(func=func, key=pkey):
+                    self.assert_eq(
+                        getattr(psdf.a.groupby(kkey), func)().sort_index(),
+                        getattr(pdf.a.groupby(pkey), func)().sort_index(),
+                        check_exact=check_exact,
+                        almost=almost,
+                    )
+                    self.assert_eq(
+                        getattr((psdf.a + 1).groupby(kkey), func)().sort_index(),
+                        getattr((pdf.a + 1).groupby(pkey), func)().sort_index(),
+                        check_exact=check_exact,
+                        almost=almost,
+                    )
+                    self.assert_eq(
+                        getattr((psdf.b + 1).groupby(kkey), func)().sort_index(),
+                        getattr((pdf.b + 1).groupby(pkey), func)().sort_index(),
+                        check_exact=check_exact,
+                        almost=almost,
+                    )
+                    self.assert_eq(
+                        getattr(psdf.a.rename().groupby(kkey), func)().sort_index(),
+                        getattr(pdf.a.rename().groupby(pkey), func)().sort_index(),
+                        check_exact=check_exact,
+                        almost=almost,
+                    )
+
+
+class GroupbySplitApplyTests(GroupbySplitApplyMixin, ComparisonTestBase, SQLTestUtils):
+    pass
+
+
+if __name__ == "__main__":
+    from pyspark.pandas.tests.groupby.test_split_apply import *  # noqa: F401
+
+    try:
+        import xmlrunner
+
+        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
+    except ImportError:
+        testRunner = None
+    unittest.main(testRunner=testRunner, verbosity=2)