snowpark-connect 0.20.2__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/pandas/tests/data_type_ops/test_udt_ops.py

@@ -0,0 +1,192 @@
+#
+# 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 pandas as pd
+
+import pyspark.pandas as ps
+from pyspark.ml.linalg import SparseVector
+from pyspark.pandas.tests.data_type_ops.testing_utils import OpsTestBase
+
+
+class UDTOpsTestsMixin:
+    @property
+    def pser(self):
+        sparse_values = {0: 0.1, 1: 1.1}
+        return pd.Series([SparseVector(len(sparse_values), sparse_values)])
+
+    @property
+    def psser(self):
+        return ps.from_pandas(self.pser)
+
+    @property
+    def udt_pdf(self):
+        sparse_values = {0: 0.2, 1: 1.0}
+        psers = {
+            "this": self.pser,
+            "that": pd.Series([SparseVector(len(sparse_values), sparse_values)]),
+        }
+        return pd.concat(psers, axis=1)
+
+    @property
+    def udt_psdf(self):
+        return ps.from_pandas(self.udt_pdf)
+
+    def test_add(self):
+        self.assertRaises(TypeError, lambda: self.psser + "x")
+        self.assertRaises(TypeError, lambda: self.psser + 1)
+
+        for psser in self.pssers:
+            self.assertRaises(TypeError, lambda: self.psser + psser)
+
+    def test_sub(self):
+        self.assertRaises(TypeError, lambda: self.psser - "x")
+        self.assertRaises(TypeError, lambda: self.psser - 1)
+
+        for psser in self.pssers:
+            self.assertRaises(TypeError, lambda: self.psser - psser)
+
+    def test_mul(self):
+        self.assertRaises(TypeError, lambda: self.psser * "x")
+        self.assertRaises(TypeError, lambda: self.psser * 1)
+
+        for psser in self.pssers:
+            self.assertRaises(TypeError, lambda: self.psser * psser)
+
+    def test_truediv(self):
+        self.assertRaises(TypeError, lambda: self.psser / "x")
+        self.assertRaises(TypeError, lambda: self.psser / 1)
+
+        for psser in self.pssers:
+            self.assertRaises(TypeError, lambda: self.psser / psser)
+
+    def test_floordiv(self):
+        self.assertRaises(TypeError, lambda: self.psser // "x")
+        self.assertRaises(TypeError, lambda: self.psser // 1)
+
+        for psser in self.pssers:
+            self.assertRaises(TypeError, lambda: self.psser // psser)
+
+    def test_mod(self):
+        self.assertRaises(TypeError, lambda: self.psser % "x")
+        self.assertRaises(TypeError, lambda: self.psser % 1)
+
+        for psser in self.pssers:
+            self.assertRaises(TypeError, lambda: self.psser % psser)
+
+    def test_pow(self):
+        self.assertRaises(TypeError, lambda: self.psser ** "x")
+        self.assertRaises(TypeError, lambda: self.psser**1)
+
+        for psser in self.pssers:
+            self.assertRaises(TypeError, lambda: self.psser**psser)
+
+    def test_radd(self):
+        self.assertRaises(TypeError, lambda: "x" + self.psser)
+        self.assertRaises(TypeError, lambda: 1 + self.psser)
+
+    def test_rsub(self):
+        self.assertRaises(TypeError, lambda: "x" - self.psser)
+        self.assertRaises(TypeError, lambda: 1 - self.psser)
+
+    def test_rmul(self):
+        self.assertRaises(TypeError, lambda: "x" * self.psser)
+        self.assertRaises(TypeError, lambda: 2 * self.psser)
+
+    def test_rtruediv(self):
+        self.assertRaises(TypeError, lambda: "x" / self.psser)
+        self.assertRaises(TypeError, lambda: 1 / self.psser)
+
+    def test_rfloordiv(self):
+        self.assertRaises(TypeError, lambda: "x" // self.psser)
+        self.assertRaises(TypeError, lambda: 1 // self.psser)
+
+    def test_rmod(self):
+        self.assertRaises(TypeError, lambda: 1 % self.psser)
+
+    def test_rpow(self):
+        self.assertRaises(TypeError, lambda: "x" ** self.psser)
+        self.assertRaises(TypeError, lambda: 1**self.psser)
+
+    def test_from_to_pandas(self):
+        sparse_values = {0: 0.1, 1: 1.1}
+        sparse_vector = SparseVector(len(sparse_values), sparse_values)
+        pser = pd.Series([sparse_vector])
+        psser = ps.Series([sparse_vector])
+        self.assert_eq(pser, psser._to_pandas())
+        self.assert_eq(ps.from_pandas(pser), psser)
+
+    def test_isnull(self):
+        self.assert_eq(self.pser.isnull(), self.psser.isnull())
+
+    def test_astype(self):
+        self.assertRaises(TypeError, lambda: self.psser.astype(str))
+
+    def test_neg(self):
+        self.assertRaises(TypeError, lambda: -self.psser)
+
+    def test_abs(self):
+        self.assertRaises(TypeError, lambda: abs(self.psser))
+
+    def test_invert(self):
+        self.assertRaises(TypeError, lambda: ~self.psser)
+
+    def test_eq(self):
+        pdf, psdf = self.udt_pdf, self.udt_psdf
+        self.assert_eq(pdf["this"] == pdf["this"], psdf["this"] == psdf["this"])
+        self.assert_eq(pdf["this"] == pdf["that"], psdf["this"] == psdf["that"])
+
+    def test_ne(self):
+        pdf, psdf = self.udt_pdf, self.udt_psdf
+        self.assert_eq(pdf["this"] != pdf["this"], psdf["this"] != psdf["this"])
+        self.assert_eq(pdf["this"] != pdf["that"], psdf["this"] != psdf["that"])
+
+    def test_lt(self):
+        self.assertRaisesRegex(
+            TypeError, "< can not be applied to", lambda: self.psser < self.psser
+        )
+
+    def test_le(self):
+        self.assertRaisesRegex(
+            TypeError, "<= can not be applied to", lambda: self.psser <= self.psser
+        )
+
+    def test_gt(self):
+        self.assertRaisesRegex(
+            TypeError, "> can not be applied to", lambda: self.psser > self.psser
+        )
+
+    def test_ge(self):
+        self.assertRaisesRegex(
+            TypeError, ">= can not be applied to", lambda: self.psser >= self.psser
+        )
+
+
+class UDTOpsTests(UDTOpsTestsMixin, OpsTestBase):
+    pass
+
+
+if __name__ == "__main__":
+    import unittest
+    from pyspark.pandas.tests.data_type_ops.test_udt_ops 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/data_type_ops/testing_utils.py

@@ -0,0 +1,228 @@
+#
+# 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 datetime
+import decimal
+from distutils.version import LooseVersion
+
+import numpy as np
+import pandas as pd
+
+import pyspark.pandas as ps
+from pyspark.pandas.typedef import extension_dtypes
+
+from pyspark.pandas.typedef.typehints import (
+    extension_dtypes_available,
+    extension_float_dtypes_available,
+    extension_object_dtypes_available,
+)
+
+from pyspark.testing.pandasutils import ComparisonTestBase
+
+if extension_dtypes_available:
+    from pandas import Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype
+
+if extension_float_dtypes_available:
+    from pandas import Float32Dtype, Float64Dtype
+
+if extension_object_dtypes_available:
+    from pandas import BooleanDtype, StringDtype
+
+
+class OpsTestBase(ComparisonTestBase):
+    """The test base for arithmetic operations of different data types."""
+
+    @property
+    def numeric_pdf(self):
+        dtypes = [np.int32, int, np.float32, float]
+        sers = [pd.Series([1, 2, 3], dtype=dtype) for dtype in dtypes]
+        sers.append(pd.Series([decimal.Decimal(1), decimal.Decimal(2), decimal.Decimal(3)]))
+        sers.append(pd.Series([1, 2, np.nan], dtype=float))
+        # Skip decimal_nan test before v1.3.0, it not supported by pandas on spark yet.
+        if LooseVersion(pd.__version__) >= LooseVersion("1.3.0"):
+            sers.append(
+                pd.Series([decimal.Decimal(1), decimal.Decimal(2), decimal.Decimal(np.nan)])
+            )
+        pdf = pd.concat(sers, axis=1)
+        if LooseVersion(pd.__version__) >= LooseVersion("1.3.0"):
+            pdf.columns = [dtype.__name__ for dtype in dtypes] + [
+                "decimal",
+                "float_nan",
+                "decimal_nan",
+            ]
+        else:
+            pdf.columns = [dtype.__name__ for dtype in dtypes] + ["decimal", "float_nan"]
+        return pdf
+
+    @property
+    def numeric_psdf(self):
+        return ps.from_pandas(self.numeric_pdf)
+
+    @property
+    def numeric_df_cols(self):
+        return self.numeric_pdf.columns
+
+    @property
+    def integral_pdf(self):
+        return pd.DataFrame({"this": [1, 2, 3], "that": [2, 2, 1]})
+
+    @property
+    def integral_psdf(self):
+        return ps.from_pandas(self.integral_pdf)
+
+    @property
+    def non_numeric_pdf(self):
+        psers = {
+            "string": pd.Series(["x", "y", "z"]),
+            "bool": pd.Series([True, True, False]),
+            "date": pd.Series(
+                [datetime.date(1994, 1, 1), datetime.date(1994, 1, 2), datetime.date(1994, 1, 3)]
+            ),
+            "datetime": pd.to_datetime(pd.Series([1, 2, 3])),
+            "timedelta": pd.Series(
+                [datetime.timedelta(1), datetime.timedelta(hours=2), datetime.timedelta(weeks=3)]
+            ),
+            "categorical": pd.Series(["a", "b", "a"], dtype="category"),
+        }
+        return pd.concat(psers, axis=1)
+
+    @property
+    def non_numeric_psdf(self):
+        return ps.from_pandas(self.non_numeric_pdf)
+
+    @property
+    def non_numeric_df_cols(self):
+        return self.non_numeric_pdf.columns
+
+    @property
+    def pdf(self):
+        return pd.concat([self.numeric_pdf, self.non_numeric_pdf], axis=1)
+
+    @property
+    def df_cols(self):
+        return self.pdf.columns
+
+    @property
+    def numeric_psers(self):
+        dtypes = [np.float32, float, int, np.int32]
+        sers = [pd.Series([1, 2, 3], dtype=dtype) for dtype in dtypes]
+        sers.append(pd.Series([decimal.Decimal(1), decimal.Decimal(2), decimal.Decimal(3)]))
+        return sers
+
+    @property
+    def numeric_pssers(self):
+        return [ps.from_pandas(pser) for pser in self.numeric_psers]
+
+    @property
+    def numeric_pser_psser_pairs(self):
+        return zip(self.numeric_psers, self.numeric_pssers)
+
+    @property
+    def non_numeric_psers(self):
+        psers = {
+            "string": pd.Series(["x", "y", "z"]),
+            "datetime": pd.to_datetime(pd.Series([1, 2, 3])),
+            "bool": pd.Series([True, True, False]),
+            "date": pd.Series(
+                [datetime.date(1994, 1, 1), datetime.date(1994, 1, 2), datetime.date(1994, 1, 3)]
+            ),
+            "categorical": pd.Series(["a", "b", "a"], dtype="category"),
+        }
+        return psers
+
+    @property
+    def non_numeric_pssers(self):
+        pssers = {}
+
+        for k, v in self.non_numeric_psers.items():
+            pssers[k] = ps.from_pandas(v)
+        return pssers
+
+    @property
+    def non_numeric_pser_psser_pairs(self):
+        return zip(self.non_numeric_psers.values(), self.non_numeric_pssers.values())
+
+    @property
+    def pssers(self):
+        return self.numeric_pssers + list(self.non_numeric_pssers.values())
+
+    @property
+    def psers(self):
+        return self.numeric_psers + list(self.non_numeric_psers.values())
+
+    @property
+    def pser_psser_pairs(self):
+        return zip(self.psers, self.pssers)
+
+    @property
+    def string_extension_dtype(self):
+        return ["string", StringDtype()] if extension_object_dtypes_available else []
+
+    @property
+    def object_extension_dtypes(self):
+        return (
+            ["boolean", "string", BooleanDtype(), StringDtype()]
+            if extension_object_dtypes_available
+            else []
+        )
+
+    @property
+    def fractional_extension_dtypes(self):
+        return (
+            ["Float32", "Float64", Float32Dtype(), Float64Dtype()]
+            if extension_float_dtypes_available
+            else []
+        )
+
+    @property
+    def integral_extension_dtypes(self):
+        return (
+            [
+                "Int8",
+                "Int16",
+                "Int32",
+                "Int64",
+                Int8Dtype(),
+                Int16Dtype(),
+                Int32Dtype(),
+                Int64Dtype(),
+            ]
+            if extension_dtypes_available
+            else []
+        )
+
+    @property
+    def extension_dtypes(self):
+        return (
+            self.object_extension_dtypes
+            + self.fractional_extension_dtypes
+            + self.integral_extension_dtypes
+        )
+
+    def check_extension(self, left, right):
+        """
+        Compare `psser` and `pser` of numeric ExtensionDtypes.
+
+        This utility is to adjust an issue for comparing numeric ExtensionDtypes in specific
+        pandas versions. Please refer to https://github.com/pandas-dev/pandas/issues/39410.
+        """
+        if LooseVersion("1.1") <= LooseVersion(pd.__version__) < LooseVersion("1.2.2"):
+            self.assert_eq(left, right, check_exact=False)
+            self.assertTrue(isinstance(left.dtype, extension_dtypes))
+            self.assertTrue(isinstance(right.dtype, extension_dtypes))
+        else:
+            self.assert_eq(left, right)

snowflake/snowpark_connect/includes/python/pyspark/pandas/tests/diff_frames_ops/__init__.py

@@ -0,0 +1,16 @@
+#
+# 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.
+#

snowflake/snowpark_connect/includes/python/pyspark/pandas/tests/diff_frames_ops/test_align.py

@@ -0,0 +1,118 @@
+#
+# 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 pandas as pd
+
+from pyspark import pandas as ps
+from pyspark.pandas.config import set_option, reset_option
+from pyspark.testing.pandasutils import PandasOnSparkTestCase
+from pyspark.testing.sqlutils import SQLTestUtils
+
+
+class DiffFramesAlignMixin:
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        set_option("compute.ops_on_diff_frames", True)
+
+    @classmethod
+    def tearDownClass(cls):
+        reset_option("compute.ops_on_diff_frames")
+        super().tearDownClass()
+
+    def test_align(self):
+        pdf1 = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}, index=[10, 20, 30])
+        pdf2 = pd.DataFrame({"a": [4, 5, 6], "c": ["d", "e", "f"]}, index=[10, 11, 12])
+        psdf1 = ps.from_pandas(pdf1)
+        psdf2 = ps.from_pandas(pdf2)
+
+        for join in ["outer", "inner", "left", "right"]:
+            for axis in [None, 0]:
+                psdf_l, psdf_r = psdf1.align(psdf2, join=join, axis=axis)
+                pdf_l, pdf_r = pdf1.align(pdf2, join=join, axis=axis)
+                self.assert_eq(psdf_l.sort_index(), pdf_l.sort_index())
+                self.assert_eq(psdf_r.sort_index(), pdf_r.sort_index())
+
+        pser1 = pd.Series([7, 8, 9], index=[10, 11, 12])
+        pser2 = pd.Series(["g", "h", "i"], index=[10, 20, 30])
+        psser1 = ps.from_pandas(pser1)
+        psser2 = ps.from_pandas(pser2)
+
+        for join in ["outer", "inner", "left", "right"]:
+            psser_l, psser_r = psser1.align(psser2, join=join)
+            pser_l, pser_r = pser1.align(pser2, join=join)
+            self.assert_eq(psser_l.sort_index(), pser_l.sort_index())
+            self.assert_eq(psser_r.sort_index(), pser_r.sort_index())
+
+            psdf_l, psser_r = psdf1.align(psser1, join=join, axis=0)
+            pdf_l, pser_r = pdf1.align(pser1, join=join, axis=0)
+            self.assert_eq(psdf_l.sort_index(), pdf_l.sort_index())
+            self.assert_eq(psser_r.sort_index(), pser_r.sort_index())
+
+            psser_l, psdf_r = psser1.align(psdf1, join=join)
+            pser_l, pdf_r = pser1.align(pdf1, join=join)
+            self.assert_eq(psser_l.sort_index(), pser_l.sort_index())
+            self.assert_eq(psdf_r.sort_index(), pdf_r.sort_index())
+
+        # multi-index columns
+        pdf3 = pd.DataFrame(
+            {("x", "a"): [4, 5, 6], ("y", "c"): ["d", "e", "f"]}, index=[10, 11, 12]
+        )
+        psdf3 = ps.from_pandas(pdf3)
+        pser3 = pdf3[("y", "c")]
+        psser3 = psdf3[("y", "c")]
+
+        for join in ["outer", "inner", "left", "right"]:
+            psdf_l, psdf_r = psdf1.align(psdf3, join=join, axis=0)
+            pdf_l, pdf_r = pdf1.align(pdf3, join=join, axis=0)
+            self.assert_eq(psdf_l.sort_index(), pdf_l.sort_index())
+            self.assert_eq(psdf_r.sort_index(), pdf_r.sort_index())
+
+            psser_l, psser_r = psser1.align(psser3, join=join)
+            pser_l, pser_r = pser1.align(pser3, join=join)
+            self.assert_eq(psser_l.sort_index(), pser_l.sort_index())
+            self.assert_eq(psser_r.sort_index(), pser_r.sort_index())
+
+            psdf_l, psser_r = psdf1.align(psser3, join=join, axis=0)
+            pdf_l, pser_r = pdf1.align(pser3, join=join, axis=0)
+            self.assert_eq(psdf_l.sort_index(), pdf_l.sort_index())
+            self.assert_eq(psser_r.sort_index(), pser_r.sort_index())
+
+            psser_l, psdf_r = psser3.align(psdf1, join=join)
+            pser_l, pdf_r = pser3.align(pdf1, join=join)
+            self.assert_eq(psser_l.sort_index(), pser_l.sort_index())
+            self.assert_eq(psdf_r.sort_index(), pdf_r.sort_index())
+
+        self.assertRaises(ValueError, lambda: psdf1.align(psdf3, axis=None))
+        self.assertRaises(ValueError, lambda: psdf1.align(psdf3, axis=1))
+
+
+class DiffFramesAlignTests(DiffFramesAlignMixin, PandasOnSparkTestCase, SQLTestUtils):
+    pass
+
+
+if __name__ == "__main__":
+    import unittest
+    from pyspark.pandas.tests.diff_frames_ops.test_align 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)