snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/pandas/tests/frame/test_attrs.py

@@ -0,0 +1,384 @@
+#
+# 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 datetime import datetime
+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
+
+
+# This file contains test cases for 'Attributes and underlying data'
+# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#attributes-and-underlying-data
+class FrameAttrsMixin:
+    @property
+    def pdf(self):
+        return pd.DataFrame(
+            {"a": [1, 2, 3, 4, 5, 6, 7, 8, 9], "b": [4, 5, 6, 3, 2, 1, 0, 0, 0]},
+            index=np.random.rand(9),
+        )
+
+    @property
+    def df_pair(self):
+        pdf = self.pdf
+        psdf = ps.from_pandas(pdf)
+        return pdf, psdf
+
+    def test_column_names(self):
+        pdf, psdf = self.df_pair
+
+        self.assert_eq(psdf.columns, pdf.columns)
+        self.assert_eq(psdf[["b", "a"]].columns, pdf[["b", "a"]].columns)
+        self.assert_eq(psdf["a"].name, pdf["a"].name)
+        self.assert_eq((psdf["a"] + 1).name, (pdf["a"] + 1).name)
+
+        self.assert_eq((psdf.a + psdf.b).name, (pdf.a + pdf.b).name)
+        self.assert_eq((psdf.a + psdf.b.rename("a")).name, (pdf.a + pdf.b.rename("a")).name)
+        self.assert_eq((psdf.a + psdf.b.rename()).name, (pdf.a + pdf.b.rename()).name)
+        self.assert_eq((psdf.a.rename() + psdf.b).name, (pdf.a.rename() + pdf.b).name)
+        self.assert_eq(
+            (psdf.a.rename() + psdf.b.rename()).name, (pdf.a.rename() + pdf.b.rename()).name
+        )
+
+    def test_rename_columns(self):
+        pdf = pd.DataFrame(
+            {"a": [1, 2, 3, 4, 5, 6, 7], "b": [7, 6, 5, 4, 3, 2, 1]}, index=np.random.rand(7)
+        )
+        psdf = ps.from_pandas(pdf)
+
+        psdf.columns = ["x", "y"]
+        pdf.columns = ["x", "y"]
+        self.assert_eq(psdf.columns, pd.Index(["x", "y"]))
+        self.assert_eq(psdf, pdf)
+        self.assert_eq(psdf._internal.data_spark_column_names, ["x", "y"])
+        self.assert_eq(psdf.to_spark().columns, ["x", "y"])
+        self.assert_eq(psdf.to_spark(index_col="index").columns, ["index", "x", "y"])
+
+        columns = pdf.columns
+        columns.name = "lvl_1"
+
+        psdf.columns = columns
+        self.assert_eq(psdf.columns.names, ["lvl_1"])
+        self.assert_eq(psdf, pdf)
+
+        msg = "Length mismatch: Expected axis has 2 elements, new values have 4 elements"
+        with self.assertRaisesRegex(ValueError, msg):
+            psdf.columns = [1, 2, 3, 4]
+
+        # Multi-index columns
+        pdf = pd.DataFrame(
+            {("A", "0"): [1, 2, 2, 3], ("B", "1"): [1, 2, 3, 4]}, index=np.random.rand(4)
+        )
+        psdf = ps.from_pandas(pdf)
+
+        columns = pdf.columns
+        self.assert_eq(psdf.columns, columns)
+        self.assert_eq(psdf, pdf)
+
+        pdf.columns = ["x", "y"]
+        psdf.columns = ["x", "y"]
+        self.assert_eq(psdf.columns, pd.Index(["x", "y"]))
+        self.assert_eq(psdf, pdf)
+        self.assert_eq(psdf._internal.data_spark_column_names, ["x", "y"])
+        self.assert_eq(psdf.to_spark().columns, ["x", "y"])
+        self.assert_eq(psdf.to_spark(index_col="index").columns, ["index", "x", "y"])
+
+        pdf.columns = columns
+        psdf.columns = columns
+        self.assert_eq(psdf.columns, columns)
+        self.assert_eq(psdf, pdf)
+        self.assert_eq(psdf._internal.data_spark_column_names, ["(A, 0)", "(B, 1)"])
+        self.assert_eq(psdf.to_spark().columns, ["(A, 0)", "(B, 1)"])
+        self.assert_eq(psdf.to_spark(index_col="index").columns, ["index", "(A, 0)", "(B, 1)"])
+
+        columns.names = ["lvl_1", "lvl_2"]
+
+        psdf.columns = columns
+        self.assert_eq(psdf.columns.names, ["lvl_1", "lvl_2"])
+        self.assert_eq(psdf, pdf)
+        self.assert_eq(psdf._internal.data_spark_column_names, ["(A, 0)", "(B, 1)"])
+        self.assert_eq(psdf.to_spark().columns, ["(A, 0)", "(B, 1)"])
+        self.assert_eq(psdf.to_spark(index_col="index").columns, ["index", "(A, 0)", "(B, 1)"])
+
+    def test_multi_index_dtypes(self):
+        # SPARK-36930: Support ps.MultiIndex.dtypes
+        arrays = [[1, 1, 2, 2], ["red", "blue", "red", "blue"]]
+        pmidx = pd.MultiIndex.from_arrays(arrays, names=("number", "color"))
+        psmidx = ps.from_pandas(pmidx)
+
+        if LooseVersion(pd.__version__) >= LooseVersion("1.3"):
+            self.assert_eq(psmidx.dtypes, pmidx.dtypes)
+        else:
+            expected = pd.Series([np.dtype("int64"), np.dtype("O")], index=["number", "color"])
+            self.assert_eq(psmidx.dtypes, expected)
+
+        # multiple labels
+        pmidx = pd.MultiIndex.from_arrays(arrays, names=[("zero", "first"), ("one", "second")])
+        psmidx = ps.from_pandas(pmidx)
+
+        if LooseVersion(pd.__version__) >= LooseVersion("1.3"):
+            if LooseVersion(pd.__version__) not in (LooseVersion("1.4.1"), LooseVersion("1.4.2")):
+                self.assert_eq(psmidx.dtypes, pmidx.dtypes)
+        else:
+            expected = pd.Series(
+                [np.dtype("int64"), np.dtype("O")],
+                index=pd.Index([("zero", "first"), ("one", "second")]),
+            )
+            self.assert_eq(psmidx.dtypes, expected)
+
+    def test_multi_index_dtypes_not_unique_name(self):
+        # Regression test for https://github.com/pandas-dev/pandas/issues/45174
+        pmidx = pd.MultiIndex.from_arrays([[1], [2]], names=[1, 1])
+        psmidx = ps.from_pandas(pmidx)
+
+        if LooseVersion(pd.__version__) < LooseVersion("1.4"):
+            expected = pd.Series(
+                [np.dtype("int64"), np.dtype("int64")],
+                index=[1, 1],
+            )
+            self.assert_eq(psmidx.dtypes, expected)
+        else:
+            self.assert_eq(psmidx.dtypes, pmidx.dtypes)
+
+    def test_dtype(self):
+        pdf = pd.DataFrame(
+            {
+                "a": list("abc"),
+                "b": list(range(1, 4)),
+                "c": np.arange(3, 6).astype("i1"),
+                "d": np.arange(4.0, 7.0, dtype="float64"),
+                "e": [True, False, True],
+                "f": pd.date_range("20130101", periods=3),
+            },
+            index=np.random.rand(3),
+        )
+        psdf = ps.from_pandas(pdf)
+        self.assert_eq(psdf, pdf)
+        self.assertTrue((psdf.dtypes == pdf.dtypes).all())
+
+        # multi-index columns
+        columns = pd.MultiIndex.from_tuples(zip(list("xxxyyz"), list("abcdef")))
+        pdf.columns = columns
+        psdf.columns = columns
+        self.assertTrue((psdf.dtypes == pdf.dtypes).all())
+
+    def test_axes(self):
+        pdf = self.pdf
+        psdf = ps.from_pandas(pdf)
+        self.assert_eq(pdf.axes, psdf.axes)
+
+        # multi-index columns
+        columns = pd.MultiIndex.from_tuples([("x", "a"), ("y", "b")])
+        pdf.columns = columns
+        psdf.columns = columns
+        self.assert_eq(pdf.axes, psdf.axes)
+
+    def test_inplace(self):
+        pdf, psdf = self.df_pair
+
+        pser = pdf.a
+        psser = psdf.a
+
+        pdf["a"] = pdf["a"] + 10
+        psdf["a"] = psdf["a"] + 10
+
+        self.assert_eq(psdf, pdf)
+        # SPARK-38946: Since Spark 3.4, df.__setitem__ generate a new dataframe to follow
+        # pandas 1.4 behaviors
+        if LooseVersion(pd.__version__) >= LooseVersion("1.4.0"):
+            self.assert_eq(psser, pser)
+        else:
+            # Follow pandas latest behavior
+            with self.assertRaisesRegex(AssertionError, "Series are different"):
+                self.assert_eq(psser, pser)
+
+    def test_dataframe_multiindex_columns(self):
+        pdf = pd.DataFrame(
+            {
+                ("x", "a", "1"): [1, 2, 3],
+                ("x", "b", "2"): [4, 5, 6],
+                ("y.z", "c.d", "3"): [7, 8, 9],
+                ("x", "b", "4"): [10, 11, 12],
+            },
+            index=np.random.rand(3),
+        )
+        psdf = ps.from_pandas(pdf)
+
+        self.assert_eq(psdf, pdf)
+        self.assert_eq(psdf["x"], pdf["x"])
+        self.assert_eq(psdf["y.z"], pdf["y.z"])
+        self.assert_eq(psdf["x"]["b"], pdf["x"]["b"])
+        self.assert_eq(psdf["x"]["b"]["2"], pdf["x"]["b"]["2"])
+
+        self.assert_eq(psdf.x, pdf.x)
+        self.assert_eq(psdf.x.b, pdf.x.b)
+        self.assert_eq(psdf.x.b["2"], pdf.x.b["2"])
+
+        self.assertRaises(KeyError, lambda: psdf["z"])
+        self.assertRaises(AttributeError, lambda: psdf.z)
+
+        self.assert_eq(psdf[("x",)], pdf[("x",)])
+        self.assert_eq(psdf[("x", "a")], pdf[("x", "a")])
+        self.assert_eq(psdf[("x", "a", "1")], pdf[("x", "a", "1")])
+
+    def test_dataframe_column_level_name(self):
+        column = pd.Index(["A", "B", "C"], name="X")
+        pdf = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=column, index=np.random.rand(2))
+        psdf = ps.from_pandas(pdf)
+
+        self.assert_eq(psdf, pdf)
+        self.assert_eq(psdf.columns.names, pdf.columns.names)
+        self.assert_eq(psdf._to_pandas().columns.names, pdf.columns.names)
+
+    def test_dataframe_multiindex_names_level(self):
+        columns = pd.MultiIndex.from_tuples(
+            [("X", "A", "Z"), ("X", "B", "Z"), ("Y", "C", "Z"), ("Y", "D", "Z")],
+            names=["lvl_1", "lvl_2", "lv_3"],
+        )
+        pdf = pd.DataFrame(
+            [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16], [17, 18, 19, 20]],
+            columns=columns,
+            index=np.random.rand(5),
+        )
+        psdf = ps.from_pandas(pdf)
+
+        self.assert_eq(psdf.columns.names, pdf.columns.names)
+        self.assert_eq(psdf._to_pandas().columns.names, pdf.columns.names)
+
+        psdf1 = ps.from_pandas(pdf)
+        self.assert_eq(psdf1.columns.names, pdf.columns.names)
+
+        self.assertRaises(
+            AssertionError,
+            lambda: ps.DataFrame(psdf1._internal.copy(column_label_names=("level",))),
+        )
+
+        self.assert_eq(psdf["X"], pdf["X"])
+        self.assert_eq(psdf["X"].columns.names, pdf["X"].columns.names)
+        self.assert_eq(psdf["X"]._to_pandas().columns.names, pdf["X"].columns.names)
+        self.assert_eq(psdf["X"]["A"], pdf["X"]["A"])
+        self.assert_eq(psdf["X"]["A"].columns.names, pdf["X"]["A"].columns.names)
+        self.assert_eq(psdf["X"]["A"]._to_pandas().columns.names, pdf["X"]["A"].columns.names)
+        self.assert_eq(psdf[("X", "A")], pdf[("X", "A")])
+        self.assert_eq(psdf[("X", "A")].columns.names, pdf[("X", "A")].columns.names)
+        self.assert_eq(psdf[("X", "A")]._to_pandas().columns.names, pdf[("X", "A")].columns.names)
+        self.assert_eq(psdf[("X", "A", "Z")], pdf[("X", "A", "Z")])
+
+    def test_repr_cache_invalidation(self):
+        # If there is any cache, inplace operations should invalidate it.
+        df = ps.range(10)
+        df.__repr__()
+        df["a"] = df["id"]
+        self.assertEqual(df.__repr__(), df._to_pandas().__repr__())
+
+    def test_repr_html_cache_invalidation(self):
+        # If there is any cache, inplace operations should invalidate it.
+        df = ps.range(10)
+        df._repr_html_()
+        df["a"] = df["id"]
+        self.assertEqual(df._repr_html_(), df._to_pandas()._repr_html_())
+
+    def test_assign(self):
+        pdf, psdf = self.df_pair
+
+        psdf["w"] = 1.0
+        pdf["w"] = 1.0
+
+        self.assert_eq(psdf, pdf)
+
+        psdf.w = 10.0
+        pdf.w = 10.0
+
+        self.assert_eq(psdf, pdf)
+
+        psdf[1] = 1.0
+        pdf[1] = 1.0
+
+        self.assert_eq(psdf, pdf)
+
+        psdf = psdf.assign(a=psdf["a"] * 2)
+        pdf = pdf.assign(a=pdf["a"] * 2)
+
+        self.assert_eq(psdf, pdf)
+
+        # multi-index columns
+        columns = pd.MultiIndex.from_tuples([("x", "a"), ("x", "b"), ("y", "w"), ("y", "v")])
+        pdf.columns = columns
+        psdf.columns = columns
+
+        psdf[("a", "c")] = "def"
+        pdf[("a", "c")] = "def"
+
+        self.assert_eq(psdf, pdf)
+
+        psdf = psdf.assign(Z="ZZ")
+        pdf = pdf.assign(Z="ZZ")
+
+        self.assert_eq(psdf, pdf)
+
+        psdf["x"] = "ghi"
+        pdf["x"] = "ghi"
+
+        self.assert_eq(psdf, pdf)
+
+    def test_attributes(self):
+        psdf = self.psdf
+
+        self.assertIn("a", dir(psdf))
+        self.assertNotIn("foo", dir(psdf))
+        self.assertRaises(AttributeError, lambda: psdf.foo)
+
+        psdf = ps.DataFrame({"a b c": [1, 2, 3]})
+        self.assertNotIn("a b c", dir(psdf))
+        psdf = ps.DataFrame({"a": [1, 2], 5: [1, 2]})
+        self.assertIn("a", dir(psdf))
+        self.assertNotIn(5, dir(psdf))
+
+    def test_empty_timestamp(self):
+        pdf = pd.DataFrame(
+            {
+                "t": [
+                    datetime(2019, 1, 1, 0, 0, 0),
+                    datetime(2019, 1, 2, 0, 0, 0),
+                    datetime(2019, 1, 3, 0, 0, 0),
+                ]
+            },
+            index=np.random.rand(3),
+        )
+        psdf = ps.from_pandas(pdf)
+        self.assert_eq(psdf[psdf["t"] != psdf["t"]], pdf[pdf["t"] != pdf["t"]])
+        self.assert_eq(psdf[psdf["t"] != psdf["t"]].dtypes, pdf[pdf["t"] != pdf["t"]].dtypes)
+
+
+class FrameAttrsTests(FrameAttrsMixin, ComparisonTestBase, SQLTestUtils):
+    pass
+
+
+if __name__ == "__main__":
+    from pyspark.pandas.tests.frame.test_attrs 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)