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

snowflake/snowpark_connect/includes/python/pyspark/pandas/tests/test_extension.py

@@ -1,152 +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 contextlib
-
-import numpy as np
-import pandas as pd
-
-from pyspark import pandas as ps
-from pyspark.testing.pandasutils import assert_produces_warning, ComparisonTestBase
-from pyspark.pandas.extensions import (
-    register_dataframe_accessor,
-    register_series_accessor,
-    register_index_accessor,
-)
-
-
-@contextlib.contextmanager
-def ensure_removed(obj, attr):
-    """
-    Ensure attribute attached to 'obj' during testing is removed in the end
-    """
-    try:
-        yield
-
-    finally:
-        try:
-            delattr(obj, attr)
-        except AttributeError:
-            pass
-
-
-class CustomAccessor:
-    def __init__(self, obj):
-        self.obj = obj
-        self.item = "item"
-
-    @property
-    def prop(self):
-        return self.item
-
-    def method(self):
-        return self.item
-
-    def check_length(self, col=None):
-        if type(self.obj) == ps.DataFrame or col is not None:
-            return len(self.obj[col])
-        else:
-            try:
-                return len(self.obj)
-            except Exception as e:
-                raise ValueError(str(e))
-
-
-class ExtensionTestsMixin:
-    @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 accessor(self):
-        return CustomAccessor(self.psdf)
-
-    def test_setup(self):
-        self.assertEqual("item", self.accessor.item)
-
-    def test_dataframe_register(self):
-        with ensure_removed(ps.DataFrame, "test"):
-            register_dataframe_accessor("test")(CustomAccessor)
-            assert self.psdf.test.prop == "item"
-            assert self.psdf.test.method() == "item"
-            assert len(self.psdf["a"]) == self.psdf.test.check_length("a")
-
-    def test_series_register(self):
-        with ensure_removed(ps.Series, "test"):
-            register_series_accessor("test")(CustomAccessor)
-            assert self.psdf.a.test.prop == "item"
-            assert self.psdf.a.test.method() == "item"
-            assert self.psdf.a.test.check_length() == len(self.psdf["a"])
-
-    def test_index_register(self):
-        with ensure_removed(ps.Index, "test"):
-            register_index_accessor("test")(CustomAccessor)
-            assert self.psdf.index.test.prop == "item"
-            assert self.psdf.index.test.method() == "item"
-            assert self.psdf.index.test.check_length() == self.psdf.index.size
-
-    def test_accessor_works(self):
-        register_series_accessor("test")(CustomAccessor)
-
-        s = ps.Series([1, 2])
-        assert s.test.obj is s
-        assert s.test.prop == "item"
-        assert s.test.method() == "item"
-
-    def test_overwrite_warns(self):
-        mean = ps.Series.mean
-        try:
-            with assert_produces_warning(UserWarning, raise_on_extra_warnings=False) as w:
-                register_series_accessor("mean")(CustomAccessor)
-                s = ps.Series([1, 2])
-                assert s.mean.prop == "item"
-            msg = str(w[0].message)
-            assert "mean" in msg
-            assert "CustomAccessor" in msg
-            assert "Series" in msg
-        finally:
-            ps.Series.mean = mean
-
-    def test_raises_attr_error(self):
-        with ensure_removed(ps.Series, "bad"):
-
-            class Bad:
-                def __init__(self, data):
-                    raise AttributeError("whoops")
-
-            with self.assertRaises(AttributeError):
-                ps.Series([1, 2], dtype=object).bad
-
-
-class ExtensionTests(ExtensionTestsMixin, ComparisonTestBase):
-    pass
-
-
-if __name__ == "__main__":
-    import unittest
-    from pyspark.pandas.tests.test_extension 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/test_frame_spark.py

@@ -1,162 +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 os
-
-import pandas as pd
-
-from pyspark import pandas as ps
-from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils
-from pyspark.testing.sqlutils import SQLTestUtils
-
-
-class SparkFrameMethodsTestsMixin:
-    def test_frame_apply_negative(self):
-        with self.assertRaisesRegex(
-            ValueError, "The output of the function.* pyspark.sql.DataFrame.*int"
-        ):
-            ps.range(10).spark.apply(lambda scol: 1)
-
-    def test_hint(self):
-        pdf1 = pd.DataFrame(
-            {"lkey": ["foo", "bar", "baz", "foo"], "value": [1, 2, 3, 5]}
-        ).set_index("lkey")
-        pdf2 = pd.DataFrame(
-            {"rkey": ["foo", "bar", "baz", "foo"], "value": [5, 6, 7, 8]}
-        ).set_index("rkey")
-        psdf1 = ps.from_pandas(pdf1)
-        psdf2 = ps.from_pandas(pdf2)
-
-        hints = ["broadcast", "merge", "shuffle_hash", "shuffle_replicate_nl"]
-
-        for hint in hints:
-            self.assert_eq(
-                pdf1.merge(pdf2, left_index=True, right_index=True).sort_values(
-                    ["value_x", "value_y"]
-                ),
-                psdf1.merge(psdf2.spark.hint(hint), left_index=True, right_index=True).sort_values(
-                    ["value_x", "value_y"]
-                ),
-                almost=True,
-            )
-            self.assert_eq(
-                pdf1.merge(pdf2 + 1, left_index=True, right_index=True).sort_values(
-                    ["value_x", "value_y"]
-                ),
-                psdf1.merge(
-                    (psdf2 + 1).spark.hint(hint), left_index=True, right_index=True
-                ).sort_values(["value_x", "value_y"]),
-                almost=True,
-            )
-
-    def test_repartition(self):
-        psdf = ps.DataFrame({"age": [5, 5, 2, 2], "name": ["Bob", "Bob", "Alice", "Alice"]})
-        num_partitions = psdf.to_spark().rdd.getNumPartitions() + 1
-
-        num_partitions += 1
-        new_psdf = psdf.spark.repartition(num_partitions)
-        self.assertEqual(new_psdf.to_spark().rdd.getNumPartitions(), num_partitions)
-        self.assert_eq(psdf.sort_index(), new_psdf.sort_index())
-
-        # Reserves Index
-        psdf = psdf.set_index("age")
-        num_partitions += 1
-        new_psdf = psdf.spark.repartition(num_partitions)
-        self.assertEqual(new_psdf.to_spark().rdd.getNumPartitions(), num_partitions)
-        self.assert_eq(psdf.sort_index(), new_psdf.sort_index())
-
-        # Reflects internal changes
-        psdf = psdf.reset_index()
-        psdf = psdf.set_index("name")
-        psdf2 = psdf + 1
-        num_partitions += 1
-        self.assert_eq(
-            psdf2.sort_index(), (psdf + 1).spark.repartition(num_partitions).sort_index()
-        )
-
-        # Reserves MultiIndex
-        psdf = ps.DataFrame({"a": ["a", "b", "c"]}, index=[[1, 2, 3], [4, 5, 6]])
-        num_partitions = psdf.to_spark().rdd.getNumPartitions() + 1
-        new_psdf = psdf.spark.repartition(num_partitions)
-        self.assertEqual(new_psdf.to_spark().rdd.getNumPartitions(), num_partitions)
-        self.assert_eq(psdf.sort_index(), new_psdf.sort_index())
-
-    def test_coalesce(self):
-        num_partitions = 10
-        psdf = ps.DataFrame({"age": [5, 5, 2, 2], "name": ["Bob", "Bob", "Alice", "Alice"]})
-        psdf = psdf.spark.repartition(num_partitions)
-
-        num_partitions -= 1
-        new_psdf = psdf.spark.coalesce(num_partitions)
-        self.assertEqual(new_psdf.to_spark().rdd.getNumPartitions(), num_partitions)
-        self.assert_eq(psdf.sort_index(), new_psdf.sort_index())
-
-        # Reserves Index
-        psdf = psdf.set_index("age")
-        num_partitions -= 1
-        new_psdf = psdf.spark.coalesce(num_partitions)
-        self.assertEqual(new_psdf.to_spark().rdd.getNumPartitions(), num_partitions)
-        self.assert_eq(psdf.sort_index(), new_psdf.sort_index())
-
-        # Reflects internal changes
-        psdf = psdf.reset_index()
-        psdf = psdf.set_index("name")
-        psdf2 = psdf + 1
-        num_partitions -= 1
-        self.assert_eq(psdf2.sort_index(), (psdf + 1).spark.coalesce(num_partitions).sort_index())
-
-        # Reserves MultiIndex
-        psdf = ps.DataFrame({"a": ["a", "b", "c"]}, index=[[1, 2, 3], [4, 5, 6]])
-        num_partitions -= 1
-        psdf = psdf.spark.repartition(num_partitions)
-
-        num_partitions -= 1
-        new_psdf = psdf.spark.coalesce(num_partitions)
-        self.assertEqual(new_psdf.to_spark().rdd.getNumPartitions(), num_partitions)
-        self.assert_eq(psdf.sort_index(), new_psdf.sort_index())
-
-    def test_checkpoint(self):
-        with self.temp_dir() as tmp:
-            self.spark.sparkContext.setCheckpointDir(tmp)
-            psdf = ps.DataFrame({"a": ["a", "b", "c"]})
-            new_psdf = psdf.spark.checkpoint()
-            self.assertIsNotNone(os.listdir(tmp))
-            self.assert_eq(psdf, new_psdf)
-
-    def test_local_checkpoint(self):
-        psdf = ps.DataFrame({"a": ["a", "b", "c"]})
-        new_psdf = psdf.spark.local_checkpoint()
-        self.assert_eq(psdf, new_psdf)
-
-
-class SparkFrameMethodsTests(
-    SparkFrameMethodsTestsMixin, PandasOnSparkTestCase, SQLTestUtils, TestUtils
-):
-    pass
-
-
-if __name__ == "__main__":
-    import unittest
-    from pyspark.pandas.tests.test_frame_spark 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/test_generic_functions.py

@@ -1,234 +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 numpy as np
-import pandas as pd
-
-import pyspark.pandas as ps
-from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils
-
-
-class GenericFunctionsTestsMixin:
-    def test_interpolate_error(self):
-        psdf = ps.range(10)
-
-        with self.assertRaisesRegex(
-            NotImplementedError, "interpolate currently works only for method='linear'"
-        ):
-            psdf.interpolate(method="quadratic")
-
-        with self.assertRaisesRegex(
-            NotImplementedError, "interpolate currently works only for method='linear'"
-        ):
-            psdf.id.interpolate(method="quadratic")
-
-        with self.assertRaisesRegex(ValueError, "limit must be > 0"):
-            psdf.interpolate(limit=0)
-
-        with self.assertRaisesRegex(ValueError, "limit must be > 0"):
-            psdf.id.interpolate(limit=0)
-
-        with self.assertRaisesRegex(ValueError, "invalid limit_direction"):
-            psdf.interpolate(limit_direction="jump")
-
-        with self.assertRaisesRegex(ValueError, "invalid limit_direction"):
-            psdf.id.interpolate(limit_direction="jump")
-
-        with self.assertRaisesRegex(ValueError, "invalid limit_area"):
-            psdf.interpolate(limit_area="jump")
-
-        with self.assertRaisesRegex(ValueError, "invalid limit_area"):
-            psdf.id.interpolate(limit_area="jump")
-
-    def _test_interpolate(self, pobj):
-        psobj = ps.from_pandas(pobj)
-        self.assert_eq(psobj.interpolate(), pobj.interpolate())
-        for limit in range(1, 5):
-            for limit_direction in [None, "forward", "backward", "both"]:
-                for limit_area in [None, "inside", "outside"]:
-                    self.assert_eq(
-                        psobj.interpolate(
-                            limit=limit, limit_direction=limit_direction, limit_area=limit_area
-                        ),
-                        pobj.interpolate(
-                            limit=limit, limit_direction=limit_direction, limit_area=limit_area
-                        ),
-                    )
-
-    def test_interpolate(self):
-        pser = pd.Series(
-            [
-                1,
-                np.nan,
-                3,
-            ],
-            name="a",
-        )
-        self._test_interpolate(pser)
-
-        pser = pd.Series(
-            [
-                np.nan,
-                np.nan,
-                np.nan,
-            ],
-            name="a",
-        )
-        self._test_interpolate(pser)
-
-        pser = pd.Series(
-            [
-                np.nan,
-                np.nan,
-                np.nan,
-                0,
-                1,
-                np.nan,
-                np.nan,
-                np.nan,
-                np.nan,
-                3,
-                np.nan,
-                np.nan,
-                np.nan,
-            ],
-            name="a",
-        )
-        self._test_interpolate(pser)
-
-        pdf = pd.DataFrame(
-            [
-                (1, 0.0, np.nan),
-                (2, np.nan, 2.0),
-                (3, 2.0, 3.0),
-                (4, np.nan, 4.0),
-                (5, np.nan, 1.0),
-            ],
-            columns=list("abc"),
-        )
-        self._test_interpolate(pdf)
-
-        pdf = pd.DataFrame(
-            [
-                (0.0, np.nan, -1.0, 1.0, np.nan),
-                (np.nan, 2.0, np.nan, np.nan, np.nan),
-                (2.0, 3.0, np.nan, 9.0, np.nan),
-                (np.nan, 4.0, -4.0, 16.0, np.nan),
-                (np.nan, 1.0, np.nan, 7.0, np.nan),
-            ],
-            columns=list("abcde"),
-        )
-        self._test_interpolate(pdf)
-
-        pdf = pd.DataFrame(
-            [
-                (0.0, np.nan, -1.0, False, np.nan),
-                (np.nan, 2.0, np.nan, True, np.nan),
-                (2.0, 3.0, np.nan, True, np.nan),
-                (np.nan, 4.0, -4.0, False, np.nan),
-                (np.nan, 1.0, np.nan, True, np.nan),
-            ],
-            columns=list("abcde"),
-        )
-        self._test_interpolate(pdf)
-
-    def _test_stat_functions(self, stat_func):
-        pdf = pd.DataFrame({"a": [np.nan, np.nan, np.nan], "b": [1, np.nan, 2], "c": [1, 2, 3]})
-        psdf = ps.from_pandas(pdf)
-        self.assert_eq(stat_func(pdf.a), stat_func(psdf.a))
-        self.assert_eq(stat_func(pdf.b), stat_func(psdf.b))
-        self.assert_eq(stat_func(pdf), stat_func(psdf))
-
-    # Fix skew and kurtosis and re-enable tests below
-    def test_stat_functions(self):
-        self._test_stat_functions(lambda x: x.sum())
-        self._test_stat_functions(lambda x: x.sum(skipna=False))
-        self._test_stat_functions(lambda x: x.mean())
-        self._test_stat_functions(lambda x: x.mean(skipna=False))
-        self._test_stat_functions(lambda x: x.product())
-        self._test_stat_functions(lambda x: x.product(skipna=False))
-        self._test_stat_functions(lambda x: x.min())
-        self._test_stat_functions(lambda x: x.min(skipna=False))
-        self._test_stat_functions(lambda x: x.max())
-        self._test_stat_functions(lambda x: x.max(skipna=False))
-        self._test_stat_functions(lambda x: x.std())
-        self._test_stat_functions(lambda x: x.std(skipna=False))
-        self._test_stat_functions(lambda x: x.std(ddof=2))
-        self._test_stat_functions(lambda x: x.var())
-        self._test_stat_functions(lambda x: x.var(ddof=2))
-        self._test_stat_functions(lambda x: x.sem())
-        self._test_stat_functions(lambda x: x.sem(skipna=False))
-        # self._test_stat_functions(lambda x: x.skew())
-        self._test_stat_functions(lambda x: x.skew(skipna=False))
-
-        # Test cases below return differently from pandas (either by design or to be fixed)
-        pdf = pd.DataFrame({"a": [np.nan, np.nan, np.nan], "b": [1, np.nan, 2], "c": [1, 2, 3]})
-        psdf = ps.from_pandas(pdf)
-
-        with self.assertRaisesRegex(TypeError, "ddof must be integer"):
-            psdf.std(ddof="ddof")
-        with self.assertRaisesRegex(TypeError, "ddof must be integer"):
-            psdf.a.std(ddof="ddof")
-
-        with self.assertRaisesRegex(TypeError, "ddof must be integer"):
-            psdf.var(ddof="ddof")
-        with self.assertRaisesRegex(TypeError, "ddof must be integer"):
-            psdf.a.var(ddof="ddof")
-
-        self.assert_eq(pdf.a.median(), psdf.a.median())
-        self.assert_eq(pdf.a.median(skipna=False), psdf.a.median(skipna=False))
-        self.assert_eq(1.0, psdf.b.median())
-        self.assert_eq(pdf.b.median(skipna=False), psdf.b.median(skipna=False))
-        self.assert_eq(pdf.c.median(), psdf.c.median())
-
-        self.assert_eq(pdf.a.kurtosis(skipna=False), psdf.a.kurtosis(skipna=False))
-        self.assert_eq(pdf.a.kurtosis(), psdf.a.kurtosis())
-        self.assert_eq(pdf.b.kurtosis(skipna=False), psdf.b.kurtosis(skipna=False))
-        self.assert_eq(pdf.b.kurtosis(), psdf.b.kurtosis())
-        self.assert_eq(pdf.c.kurtosis(), psdf.c.kurtosis())
-
-    def test_prod_precision(self):
-        pdf = pd.DataFrame(
-            {
-                "a": [np.nan, np.nan, np.nan, np.nan],
-                "b": [1, np.nan, np.nan, -4],
-                "c": [1, -2, 3, -4],
-                "d": [55108, 55108, 55108, 55108],
-            }
-        )
-        psdf = ps.from_pandas(pdf)
-
-        self.assert_eq(pdf.prod(), psdf.prod())
-        self.assert_eq(pdf.prod(skipna=False), psdf.prod(skipna=False))
-        self.assert_eq(pdf.prod(min_count=3), psdf.prod(min_count=3))
-        self.assert_eq(pdf.prod(skipna=False, min_count=3), psdf.prod(skipna=False, min_count=3))
-
-
-class GenericFunctionsTests(GenericFunctionsTestsMixin, PandasOnSparkTestCase, TestUtils):
-    pass
-
-
-if __name__ == "__main__":
-    import unittest
-    from pyspark.pandas.tests.test_generic_functions 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)