snowpark-connect 0.24.0__py3-none-any.whl → 0.25.0__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/pandas/tests/indexes/test_indexing.py

@@ -1,432 +0,0 @@
-#
-# 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.
-#
-import unittest
-from distutils.version import LooseVersion
-
-import numpy as np
-import pandas as pd
-
-from pyspark import pandas as ps
-from pyspark.pandas.config import option_context
-from pyspark.testing.pandasutils import ComparisonTestBase
-from pyspark.testing.sqlutils import SQLTestUtils
-
-
-# This file contains test cases for 'Indexing, Iteration'
-# https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/frame.html#indexing-iteration
-class FrameIndexingMixin:
-    @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_head(self):
-        pdf, psdf = self.df_pair
-
-        self.assert_eq(psdf.head(2), pdf.head(2))
-        self.assert_eq(psdf.head(3), pdf.head(3))
-        self.assert_eq(psdf.head(0), pdf.head(0))
-        self.assert_eq(psdf.head(-3), pdf.head(-3))
-        self.assert_eq(psdf.head(-10), pdf.head(-10))
-        with option_context("compute.ordered_head", True):
-            self.assert_eq(psdf.head(), pdf.head())
-
-    @unittest.skipIf(
-        LooseVersion(pd.__version__) >= LooseVersion("2.0.0"),
-        "TODO(SPARK-43559): Enable DataFrameSlowTests.test_iteritems for pandas 2.0.0.",
-    )
-    def test_iteritems(self):
-        pdf = pd.DataFrame(
-            {"species": ["bear", "bear", "marsupial"], "population": [1864, 22000, 80000]},
-            index=["panda", "polar", "koala"],
-            columns=["species", "population"],
-        )
-        psdf = ps.from_pandas(pdf)
-
-        for (p_name, p_items), (k_name, k_items) in zip(pdf.iteritems(), psdf.iteritems()):
-            self.assert_eq(p_name, k_name)
-            self.assert_eq(p_items, k_items)
-
-    def test_keys(self):
-        pdf = pd.DataFrame(
-            [[1, 2], [4, 5], [7, 8]],
-            index=["cobra", "viper", "sidewinder"],
-            columns=["max_speed", "shield"],
-        )
-        psdf = ps.from_pandas(pdf)
-
-        self.assert_eq(psdf.keys(), pdf.keys())
-
-    def test_tail(self):
-        pdf = pd.DataFrame({"x": range(1000)})
-        psdf = ps.from_pandas(pdf)
-
-        self.assert_eq(pdf.tail(), psdf.tail())
-        self.assert_eq(pdf.tail(10), psdf.tail(10))
-        self.assert_eq(pdf.tail(-990), psdf.tail(-990))
-        self.assert_eq(pdf.tail(0), psdf.tail(0))
-        self.assert_eq(pdf.tail(-1001), psdf.tail(-1001))
-        self.assert_eq(pdf.tail(1001), psdf.tail(1001))
-        self.assert_eq((pdf + 1).tail(), (psdf + 1).tail())
-        self.assert_eq((pdf + 1).tail(10), (psdf + 1).tail(10))
-        self.assert_eq((pdf + 1).tail(-990), (psdf + 1).tail(-990))
-        self.assert_eq((pdf + 1).tail(0), (psdf + 1).tail(0))
-        self.assert_eq((pdf + 1).tail(-1001), (psdf + 1).tail(-1001))
-        self.assert_eq((pdf + 1).tail(1001), (psdf + 1).tail(1001))
-        with self.assertRaisesRegex(TypeError, "bad operand type for unary -: 'str'"):
-            psdf.tail("10")
-
-    def test_xs(self):
-        d = {
-            "num_legs": [4, 4, 2, 2],
-            "num_wings": [0, 0, 2, 2],
-            "class": ["mammal", "mammal", "mammal", "bird"],
-            "animal": ["cat", "dog", "bat", "penguin"],
-            "locomotion": ["walks", "walks", "flies", "walks"],
-        }
-        pdf = pd.DataFrame(data=d)
-        pdf = pdf.set_index(["class", "animal", "locomotion"])
-        psdf = ps.from_pandas(pdf)
-
-        self.assert_eq(psdf.xs("mammal"), pdf.xs("mammal"))
-        self.assert_eq(psdf.xs(("mammal",)), pdf.xs(("mammal",)))
-        self.assert_eq(psdf.xs(("mammal", "dog", "walks")), pdf.xs(("mammal", "dog", "walks")))
-        self.assert_eq(
-            ps.concat([psdf, psdf]).xs(("mammal", "dog", "walks")),
-            pd.concat([pdf, pdf]).xs(("mammal", "dog", "walks")),
-        )
-        self.assert_eq(psdf.xs("cat", level=1), pdf.xs("cat", level=1))
-        self.assert_eq(psdf.xs("flies", level=2), pdf.xs("flies", level=2))
-        self.assert_eq(psdf.xs("mammal", level=-3), pdf.xs("mammal", level=-3))
-
-        msg = 'axis should be either 0 or "index" currently.'
-        with self.assertRaisesRegex(NotImplementedError, msg):
-            psdf.xs("num_wings", axis=1)
-        with self.assertRaises(KeyError):
-            psdf.xs(("mammal", "dog", "walk"))
-        msg = r"'Key length \(4\) exceeds index depth \(3\)'"
-        with self.assertRaisesRegex(KeyError, msg):
-            psdf.xs(("mammal", "dog", "walks", "foo"))
-        msg = "'key' should be a scalar value or tuple that contains scalar values"
-        with self.assertRaisesRegex(TypeError, msg):
-            psdf.xs(["mammal", "dog", "walks", "foo"])
-
-        self.assertRaises(IndexError, lambda: psdf.xs("foo", level=-4))
-        self.assertRaises(IndexError, lambda: psdf.xs("foo", level=3))
-
-        self.assertRaises(KeyError, lambda: psdf.xs(("dog", "walks"), level=1))
-
-        # non-string names
-        pdf = pd.DataFrame(data=d)
-        pdf = pdf.set_index(["class", "animal", "num_legs", "num_wings"])
-        psdf = ps.from_pandas(pdf)
-
-        self.assert_eq(psdf.xs(("mammal", "dog", 4)), pdf.xs(("mammal", "dog", 4)))
-        self.assert_eq(psdf.xs(2, level=2), pdf.xs(2, level=2))
-
-        self.assert_eq((psdf + "a").xs(("mammal", "dog", 4)), (pdf + "a").xs(("mammal", "dog", 4)))
-        self.assert_eq((psdf + "a").xs(2, level=2), (pdf + "a").xs(2, level=2))
-
-    def test_where(self):
-        pdf, psdf = self.df_pair
-
-        # pandas requires `axis` argument when the `other` is Series.
-        # `axis` is not fully supported yet in pandas-on-Spark.
-        self.assert_eq(
-            psdf.where(psdf > 2, psdf.a + 10, axis=0), pdf.where(pdf > 2, pdf.a + 10, axis=0)
-        )
-
-        with self.assertRaisesRegex(TypeError, "type of cond must be a DataFrame or Series"):
-            psdf.where(1)
-        with self.assertRaisesRegex(
-            NotImplementedError, 'axis should be either 0 or "index" currently.'
-        ):
-            psdf.where(psdf > 2, psdf.a + 10, axis=1)
-
-    def test_mask(self):
-        psdf = ps.from_pandas(self.pdf)
-
-        with self.assertRaisesRegex(TypeError, "type of cond must be a DataFrame or Series"):
-            psdf.mask(1)
-
-    def test_query(self):
-        pdf = pd.DataFrame({"A": range(1, 6), "B": range(10, 0, -2), "C": range(10, 5, -1)})
-        psdf = ps.from_pandas(pdf)
-
-        exprs = ("A > B", "A < C", "C == B")
-        for expr in exprs:
-            self.assert_eq(psdf.query(expr), pdf.query(expr))
-
-        # test `inplace=True`
-        for expr in exprs:
-            dummy_psdf = psdf.copy()
-            dummy_pdf = pdf.copy()
-
-            pser = dummy_pdf.A
-            psser = dummy_psdf.A
-            dummy_pdf.query(expr, inplace=True)
-            dummy_psdf.query(expr, inplace=True)
-
-            self.assert_eq(dummy_psdf, dummy_pdf)
-            self.assert_eq(psser, pser)
-
-        # invalid values for `expr`
-        invalid_exprs = (1, 1.0, (exprs[0],), [exprs[0]])
-        for expr in invalid_exprs:
-            with self.assertRaisesRegex(
-                TypeError,
-                "expr must be a string to be evaluated, {} given".format(type(expr).__name__),
-            ):
-                psdf.query(expr)
-
-        # invalid values for `inplace`
-        invalid_inplaces = (1, 0, "True", "False")
-        for inplace in invalid_inplaces:
-            with self.assertRaisesRegex(
-                TypeError,
-                'For argument "inplace" expected type bool, received type {}.'.format(
-                    type(inplace).__name__
-                ),
-            ):
-                psdf.query("a < b", inplace=inplace)
-
-        # doesn't support for MultiIndex columns
-        columns = pd.MultiIndex.from_tuples([("A", "Z"), ("B", "X"), ("C", "C")])
-        psdf.columns = columns
-        with self.assertRaisesRegex(TypeError, "Doesn't support for MultiIndex columns"):
-            psdf.query("('A', 'Z') > ('B', 'X')")
-
-    def test_insert(self):
-        #
-        # Basic DataFrame
-        #
-        pdf = pd.DataFrame([1, 2, 3])
-        psdf = ps.from_pandas(pdf)
-
-        psdf.insert(1, "b", 10)
-        pdf.insert(1, "b", 10)
-        self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True)
-        psdf.insert(2, "c", 0.1)
-        pdf.insert(2, "c", 0.1)
-        self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True)
-        psdf.insert(3, "d", psdf.b + 1)
-        pdf.insert(3, "d", pdf.b + 1)
-        self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True)
-
-        psser = ps.Series([4, 5, 6])
-        self.assertRaises(ValueError, lambda: psdf.insert(0, "y", psser))
-        self.assertRaisesRegex(
-            ValueError, "cannot insert b, already exists", lambda: psdf.insert(1, "b", 10)
-        )
-        self.assertRaisesRegex(
-            TypeError,
-            '"column" should be a scalar value or tuple that contains scalar values',
-            lambda: psdf.insert(0, list("abc"), psser),
-        )
-        self.assertRaisesRegex(
-            TypeError,
-            "loc must be int",
-            lambda: psdf.insert((1,), "b", 10),
-        )
-        self.assertRaisesRegex(
-            NotImplementedError,
-            "Assigning column name as tuple is only supported for MultiIndex columns for now.",
-            lambda: psdf.insert(0, ("e",), 10),
-        )
-
-        self.assertRaises(ValueError, lambda: psdf.insert(0, "e", [7, 8, 9, 10]))
-        self.assertRaises(ValueError, lambda: psdf.insert(0, "f", ps.Series([7, 8])))
-        self.assertRaises(AssertionError, lambda: psdf.insert(100, "y", psser))
-        self.assertRaises(AssertionError, lambda: psdf.insert(1, "y", psser, allow_duplicates=True))
-
-        #
-        # DataFrame with MultiIndex as columns
-        #
-        pdf = pd.DataFrame({("x", "a", "b"): [1, 2, 3]})
-        psdf = ps.from_pandas(pdf)
-
-        psdf.insert(1, "b", 10)
-        pdf.insert(1, "b", 10)
-        self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True)
-        psdf.insert(2, "c", 0.1)
-        pdf.insert(2, "c", 0.1)
-        self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True)
-        psdf.insert(3, "d", psdf.b + 1)
-        pdf.insert(3, "d", pdf.b + 1)
-        self.assert_eq(psdf.sort_index(), pdf.sort_index(), almost=True)
-
-        self.assertRaisesRegex(
-            ValueError, "cannot insert d, already exists", lambda: psdf.insert(4, "d", 11)
-        )
-        self.assertRaisesRegex(
-            ValueError,
-            r"cannot insert \('x', 'a', 'b'\), already exists",
-            lambda: psdf.insert(4, ("x", "a", "b"), 11),
-        )
-        self.assertRaisesRegex(
-            ValueError,
-            '"column" must have length equal to number of column levels.',
-            lambda: psdf.insert(4, ("e",), 11),
-        )
-
-    def test_itertuples(self):
-        pdf = pd.DataFrame({"num_legs": [4, 2], "num_wings": [0, 2]}, index=["dog", "hawk"])
-        psdf = ps.from_pandas(pdf)
-
-        for ptuple, ktuple in zip(
-            pdf.itertuples(index=False, name="Animal"), psdf.itertuples(index=False, name="Animal")
-        ):
-            self.assert_eq(ptuple, ktuple)
-        for ptuple, ktuple in zip(pdf.itertuples(name=None), psdf.itertuples(name=None)):
-            self.assert_eq(ptuple, ktuple)
-        for ptuple, ktuple in zip(
-            pdf.itertuples(index=False, name=None), psdf.itertuples(index=False, name=None)
-        ):
-            self.assert_eq(ptuple, ktuple)
-
-        pdf.index = pd.MultiIndex.from_arrays(
-            [[1, 2], ["black", "brown"]], names=("count", "color")
-        )
-        psdf = ps.from_pandas(pdf)
-        for ptuple, ktuple in zip(pdf.itertuples(name="Animal"), psdf.itertuples(name="Animal")):
-            self.assert_eq(ptuple, ktuple)
-
-        pdf.columns = pd.MultiIndex.from_arrays(
-            [["CA", "WA"], ["age", "children"]], names=("origin", "info")
-        )
-        psdf = ps.from_pandas(pdf)
-        for ptuple, ktuple in zip(pdf.itertuples(name="Animal"), psdf.itertuples(name="Animal")):
-            self.assert_eq(ptuple, ktuple)
-
-        pdf = pd.DataFrame([1, 2, 3])
-        psdf = ps.from_pandas(pdf)
-        for ptuple, ktuple in zip(
-            (pdf + 1).itertuples(name="num"), (psdf + 1).itertuples(name="num")
-        ):
-            self.assert_eq(ptuple, ktuple)
-
-        # DataFrames with a large number of columns (>254)
-        pdf = pd.DataFrame(np.random.random((1, 255)))
-        psdf = ps.from_pandas(pdf)
-        for ptuple, ktuple in zip(pdf.itertuples(name="num"), psdf.itertuples(name="num")):
-            self.assert_eq(ptuple, ktuple)
-
-    def test_iterrows(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)
-
-        for (pdf_k, pdf_v), (psdf_k, psdf_v) in zip(pdf.iterrows(), psdf.iterrows()):
-            self.assert_eq(pdf_k, psdf_k)
-            self.assert_eq(pdf_v, psdf_v)
-
-        # MultiIndex
-        pmidx = pd.Index([(1, 2), (3, 4), (5, 6)])
-        pdf.index = pmidx
-        psdf = ps.from_pandas(pdf)
-
-        for (pdf_k, pdf_v), (psdf_k, psdf_v) in zip(pdf.iterrows(), psdf.iterrows()):
-            self.assert_eq(pdf_k, psdf_k)
-            self.assert_eq(pdf_v, psdf_v)
-
-    def test_multiindex_column_access(self):
-        columns = pd.MultiIndex.from_tuples(
-            [
-                ("a", "", "", "b"),
-                ("c", "", "d", ""),
-                ("e", "", "f", ""),
-                ("e", "g", "", ""),
-                ("", "", "", "h"),
-                ("i", "", "", ""),
-            ]
-        )
-
-        pdf = pd.DataFrame(
-            [
-                (1, "a", "x", 10, 100, 1000),
-                (2, "b", "y", 20, 200, 2000),
-                (3, "c", "z", 30, 300, 3000),
-            ],
-            columns=columns,
-            index=np.random.rand(3),
-        )
-        psdf = ps.from_pandas(pdf)
-
-        self.assert_eq(psdf, pdf)
-        self.assert_eq(psdf["a"], pdf["a"])
-        self.assert_eq(psdf["a"]["b"], pdf["a"]["b"])
-        self.assert_eq(psdf["c"], pdf["c"])
-        self.assert_eq(psdf["c"]["d"], pdf["c"]["d"])
-        self.assert_eq(psdf["e"], pdf["e"])
-        self.assert_eq(psdf["e"][""]["f"], pdf["e"][""]["f"])
-        self.assert_eq(psdf["e"]["g"], pdf["e"]["g"])
-        self.assert_eq(psdf[""], pdf[""])
-        self.assert_eq(psdf[""]["h"], pdf[""]["h"])
-        self.assert_eq(psdf["i"], pdf["i"])
-
-        self.assert_eq(psdf[["a", "e"]], pdf[["a", "e"]])
-        self.assert_eq(psdf[["e", "a"]], pdf[["e", "a"]])
-
-        self.assert_eq(psdf[("a",)], pdf[("a",)])
-        self.assert_eq(psdf[("e", "g")], pdf[("e", "g")])
-        # self.assert_eq(psdf[("i",)], pdf[("i",)])
-        self.assert_eq(psdf[("i", "")], pdf[("i", "")])
-
-        self.assertRaises(KeyError, lambda: psdf[("a", "b")])
-
-    def test_getitem_with_none_key(self):
-        psdf = self.psdf
-
-        with self.assertRaisesRegex(KeyError, "none key"):
-            psdf[None]
-
-    def test_iter_dataframe(self):
-        pdf, psdf = self.df_pair
-
-        for value_psdf, value_pdf in zip(psdf, pdf):
-            self.assert_eq(value_psdf, value_pdf)
-
-
-class FrameIndexingTests(FrameIndexingMixin, ComparisonTestBase, SQLTestUtils):
-    pass
-
-
-if __name__ == "__main__":
-    from pyspark.pandas.tests.indexes.test_indexing 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)