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

snowflake/snowpark_connect/includes/python/pyspark/sql/tests/test_column.py

@@ -1,225 +0,0 @@
-# -*- encoding: utf-8 -*-
-#
-# 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 pyspark.sql import Column, Row
-from pyspark.sql.types import StructType, StructField, LongType
-from pyspark.errors import AnalysisException, PySparkTypeError
-from pyspark.testing.sqlutils import ReusedSQLTestCase
-
-
-class ColumnTestsMixin:
-    def test_column_name_encoding(self):
-        """Ensure that created columns has `str` type consistently."""
-        columns = self.spark.createDataFrame([("Alice", 1)], ["name", "age"]).columns
-        self.assertEqual(columns, ["name", "age"])
-        self.assertTrue(isinstance(columns[0], str))
-        self.assertTrue(isinstance(columns[1], str))
-
-    def test_and_in_expression(self):
-        self.assertEqual(4, self.df.filter((self.df.key <= 10) & (self.df.value <= "2")).count())
-        self.assertRaises(ValueError, lambda: (self.df.key <= 10) and (self.df.value <= "2"))
-        self.assertEqual(14, self.df.filter((self.df.key <= 3) | (self.df.value < "2")).count())
-        self.assertRaises(ValueError, lambda: self.df.key <= 3 or self.df.value < "2")
-        self.assertEqual(99, self.df.filter(~(self.df.key == 1)).count())
-        self.assertRaises(ValueError, lambda: not self.df.key == 1)
-
-    def test_validate_column_types(self):
-        from pyspark.sql.functions import udf, to_json
-        from pyspark.sql.column import _to_java_column
-
-        self.assertTrue("Column" in _to_java_column("a").getClass().toString())
-        self.assertTrue("Column" in _to_java_column("a").getClass().toString())
-        self.assertTrue("Column" in _to_java_column(self.spark.range(1).id).getClass().toString())
-
-        with self.assertRaises(PySparkTypeError) as pe:
-            _to_java_column(1)
-
-        self.check_error(
-            exception=pe.exception,
-            error_class="NOT_COLUMN_OR_STR",
-            message_parameters={"arg_name": "col", "arg_type": "int"},
-        )
-
-        class A:
-            pass
-
-        self.assertRaises(TypeError, lambda: _to_java_column(A()))
-        self.assertRaises(TypeError, lambda: _to_java_column([]))
-
-        with self.assertRaises(PySparkTypeError) as pe:
-            udf(lambda x: x)(None)
-
-        self.check_error(
-            exception=pe.exception,
-            error_class="NOT_COLUMN_OR_STR",
-            message_parameters={"arg_name": "col", "arg_type": "NoneType"},
-        )
-        self.assertRaises(TypeError, lambda: to_json(1))
-
-    def test_column_operators(self):
-        ci = self.df.key
-        cs = self.df.value
-        ci == cs
-        self.assertTrue(isinstance((-ci - 1 - 2) % 3 * 2.5 / 3.5, Column))
-        rcc = (1 + ci), (1 - ci), (1 * ci), (1 / ci), (1 % ci), (1**ci), (ci**1)
-        self.assertTrue(all(isinstance(c, Column) for c in rcc))
-        cb = [ci == 5, ci != 0, ci > 3, ci < 4, ci >= 0, ci <= 7]
-        self.assertTrue(all(isinstance(c, Column) for c in cb))
-        cbool = (ci & ci), (ci | ci), (~ci)
-        self.assertTrue(all(isinstance(c, Column) for c in cbool))
-        css = (
-            cs.contains("a"),
-            cs.like("a"),
-            cs.rlike("a"),
-            cs.ilike("A"),
-            cs.asc(),
-            cs.desc(),
-            cs.startswith("a"),
-            cs.endswith("a"),
-            ci.eqNullSafe(cs),
-        )
-        self.assertTrue(all(isinstance(c, Column) for c in css))
-        self.assertTrue(isinstance(ci.cast(LongType()), Column))
-        self.assertRaisesRegex(
-            ValueError, "Cannot apply 'in' operator against a column", lambda: 1 in cs
-        )
-
-    def test_column_accessor(self):
-        from pyspark.sql.functions import col
-
-        self.assertIsInstance(col("foo")[1:3], Column)
-        self.assertIsInstance(col("foo")[0], Column)
-        self.assertIsInstance(col("foo")["bar"], Column)
-        self.assertRaises(ValueError, lambda: col("foo")[0:10:2])
-
-    def test_column_select(self):
-        df = self.df
-        self.assertEqual(self.testData, df.select("*").collect())
-        self.assertEqual(self.testData, df.select(df.key, df.value).collect())
-        self.assertEqual([Row(value="1")], df.where(df.key == 1).select(df.value).collect())
-
-    def test_access_column(self):
-        df = self.df
-        self.assertTrue(isinstance(df.key, Column))
-        self.assertTrue(isinstance(df["key"], Column))
-        self.assertTrue(isinstance(df[0], Column))
-        self.assertRaises(IndexError, lambda: df[2])
-        self.assertRaises(AnalysisException, lambda: df["bad_key"])
-        self.assertRaises(TypeError, lambda: df[{}])
-
-    def test_column_name_with_non_ascii(self):
-        columnName = "数量"
-        self.assertTrue(isinstance(columnName, str))
-        schema = StructType([StructField(columnName, LongType(), True)])
-        df = self.spark.createDataFrame([(1,)], schema)
-        self.assertEqual(schema, df.schema)
-        self.assertEqual("DataFrame[数量: bigint]", str(df))
-        self.assertEqual([("数量", "bigint")], df.dtypes)
-        self.assertEqual(1, df.select("数量").first()[0])
-        self.assertEqual(1, df.select(df["数量"]).first()[0])
-        self.assertTrue(columnName in repr(df[columnName]))
-
-    def test_field_accessor(self):
-        df = self.spark.createDataFrame([Row(l=[1], r=Row(a=1, b="b"), d={"k": "v"})])
-        self.assertEqual(1, df.select(df.l[0]).first()[0])
-        self.assertEqual(1, df.select(df.r["a"]).first()[0])
-        self.assertEqual(1, df.select(df["r.a"]).first()[0])
-        self.assertEqual("b", df.select(df.r["b"]).first()[0])
-        self.assertEqual("b", df.select(df["r.b"]).first()[0])
-        self.assertEqual("v", df.select(df.d["k"]).first()[0])
-
-    def test_bitwise_operations(self):
-        from pyspark.sql import functions
-
-        row = Row(a=170, b=75)
-        df = self.spark.createDataFrame([row])
-        result = df.select(df.a.bitwiseAND(df.b)).collect()[0].asDict()
-        self.assertEqual(170 & 75, result["(a & b)"])
-        result = df.select(df.a.bitwiseOR(df.b)).collect()[0].asDict()
-        self.assertEqual(170 | 75, result["(a | b)"])
-        result = df.select(df.a.bitwiseXOR(df.b)).collect()[0].asDict()
-        self.assertEqual(170 ^ 75, result["(a ^ b)"])
-        result = df.select(functions.bitwiseNOT(df.b)).collect()[0].asDict()
-        self.assertEqual(~75, result["~b"])
-        result = df.select(functions.bitwise_not(df.b)).collect()[0].asDict()
-        self.assertEqual(~75, result["~b"])
-
-    def test_with_field(self):
-        from pyspark.sql.functions import lit, col
-
-        df = self.spark.createDataFrame([Row(a=Row(b=1, c=2))])
-        self.assertIsInstance(df["a"].withField("b", lit(3)), Column)
-        self.assertIsInstance(df["a"].withField("d", lit(3)), Column)
-        result = df.withColumn("a", df["a"].withField("d", lit(3))).collect()[0].asDict()
-        self.assertEqual(3, result["a"]["d"])
-        result = df.withColumn("a", df["a"].withField("b", lit(3))).collect()[0].asDict()
-        self.assertEqual(3, result["a"]["b"])
-
-        with self.assertRaises(PySparkTypeError) as pe:
-            df["a"].withField("b", 3)
-
-        self.check_error(
-            exception=pe.exception,
-            error_class="NOT_COLUMN",
-            message_parameters={"arg_name": "col", "arg_type": "int"},
-        )
-
-        with self.assertRaises(PySparkTypeError) as pe:
-            df["a"].withField(col("b"), lit(3))
-
-        self.check_error(
-            exception=pe.exception,
-            error_class="NOT_STR",
-            message_parameters={"arg_name": "fieldName", "arg_type": "Column"},
-        )
-
-    def test_drop_fields(self):
-        df = self.spark.createDataFrame([Row(a=Row(b=1, c=2, d=Row(e=3, f=4)))])
-        self.assertIsInstance(df["a"].dropFields("b"), Column)
-        self.assertIsInstance(df["a"].dropFields("b", "c"), Column)
-        self.assertIsInstance(df["a"].dropFields("d.e"), Column)
-
-        result = (
-            df.select(
-                df["a"].dropFields("b").alias("a1"),
-                df["a"].dropFields("d.e").alias("a2"),
-            )
-            .first()
-            .asDict(True)
-        )
-
-        self.assertTrue("b" not in result["a1"] and "c" in result["a1"] and "d" in result["a1"])
-
-        self.assertTrue("e" not in result["a2"]["d"] and "f" in result["a2"]["d"])
-
-
-class ColumnTests(ColumnTestsMixin, ReusedSQLTestCase):
-    pass
-
-
-if __name__ == "__main__":
-    import unittest
-    from pyspark.sql.tests.test_column 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/sql/tests/test_conf.py

@@ -1,83 +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.
-#
-from decimal import Decimal
-
-from pyspark.errors import IllegalArgumentException
-from pyspark.testing.sqlutils import ReusedSQLTestCase
-
-
-class ConfTestsMixin:
-    def test_conf(self):
-        spark = self.spark
-        spark.conf.set("bogo", "sipeo")
-        self.assertEqual(spark.conf.get("bogo"), "sipeo")
-        spark.conf.set("bogo", "ta")
-        self.assertEqual(spark.conf.get("bogo"), "ta")
-        self.assertEqual(spark.conf.get("bogo", "not.read"), "ta")
-        self.assertEqual(spark.conf.get("not.set", "ta"), "ta")
-        self.assertRaisesRegex(Exception, "not.set", lambda: spark.conf.get("not.set"))
-        spark.conf.unset("bogo")
-        self.assertEqual(spark.conf.get("bogo", "colombia"), "colombia")
-
-        self.assertEqual(spark.conf.get("hyukjin", None), None)
-
-        # This returns 'STATIC' because it's the default value of
-        # 'spark.sql.sources.partitionOverwriteMode', and `defaultValue` in
-        # `spark.conf.get` is unset.
-        self.assertEqual(spark.conf.get("spark.sql.sources.partitionOverwriteMode"), "STATIC")
-
-        # This returns None because 'spark.sql.sources.partitionOverwriteMode' is unset, but
-        # `defaultValue` in `spark.conf.get` is set to None.
-        self.assertEqual(spark.conf.get("spark.sql.sources.partitionOverwriteMode", None), None)
-
-        self.assertTrue(spark.conf.isModifiable("spark.sql.execution.arrow.maxRecordsPerBatch"))
-        self.assertFalse(spark.conf.isModifiable("spark.sql.warehouse.dir"))
-
-    def test_conf_with_python_objects(self):
-        spark = self.spark
-
-        for value, expected in [(True, "true"), (False, "false")]:
-            spark.conf.set("foo", value)
-            self.assertEqual(spark.conf.get("foo"), expected)
-
-        spark.conf.set("foo", 1)
-        self.assertEqual(spark.conf.get("foo"), "1")
-
-        with self.assertRaises(IllegalArgumentException):
-            spark.conf.set("foo", None)
-
-        with self.assertRaises(Exception):
-            spark.conf.set("foo", Decimal(1))
-
-        spark.conf.unset("foo")
-
-
-class ConfTests(ConfTestsMixin, ReusedSQLTestCase):
-    pass
-
-
-if __name__ == "__main__":
-    import unittest
-    from pyspark.sql.tests.test_conf 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/sql/tests/test_context.py

@@ -1,201 +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 shutil
-import sys
-import tempfile
-import unittest
-from importlib import reload
-
-import py4j
-
-from pyspark import SparkContext, SQLContext
-from pyspark.sql import Row, SparkSession
-from pyspark.sql.types import StructType, StringType, StructField
-from pyspark.testing.utils import ReusedPySparkTestCase
-
-
-class HiveContextSQLTests(ReusedPySparkTestCase):
-    @classmethod
-    def setUpClass(cls):
-        ReusedPySparkTestCase.setUpClass()
-        cls.tempdir = tempfile.NamedTemporaryFile(delete=False)
-        cls.hive_available = True
-        cls.spark = None
-        try:
-            cls.sc._jvm.org.apache.hadoop.hive.conf.HiveConf()
-        except py4j.protocol.Py4JError:
-            cls.tearDownClass()
-            cls.hive_available = False
-        except TypeError:
-            cls.tearDownClass()
-            cls.hive_available = False
-        if cls.hive_available:
-            cls.spark = SparkSession.builder.enableHiveSupport().getOrCreate()
-
-        os.unlink(cls.tempdir.name)
-        if cls.hive_available:
-            cls.testData = [Row(key=i, value=str(i)) for i in range(100)]
-            cls.df = cls.sc.parallelize(cls.testData).toDF()
-
-    def setUp(self):
-        if not self.hive_available:
-            self.skipTest("Hive is not available.")
-
-    @classmethod
-    def tearDownClass(cls):
-        ReusedPySparkTestCase.tearDownClass()
-        shutil.rmtree(cls.tempdir.name, ignore_errors=True)
-        if cls.spark is not None:
-            cls.spark.stop()
-            cls.spark = None
-
-    def test_save_and_load_table(self):
-        df = self.df
-        tmpPath = tempfile.mkdtemp()
-        shutil.rmtree(tmpPath)
-        df.write.saveAsTable("savedJsonTable", "json", "append", path=tmpPath)
-        actual = self.spark.catalog.createTable("externalJsonTable", tmpPath, "json")
-        self.assertEqual(
-            sorted(df.collect()), sorted(self.spark.sql("SELECT * FROM savedJsonTable").collect())
-        )
-        self.assertEqual(
-            sorted(df.collect()),
-            sorted(self.spark.sql("SELECT * FROM externalJsonTable").collect()),
-        )
-        self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
-        self.spark.sql("DROP TABLE externalJsonTable")
-
-        df.write.saveAsTable("savedJsonTable", "json", "overwrite", path=tmpPath)
-        schema = StructType([StructField("value", StringType(), True)])
-        actual = self.spark.catalog.createTable(
-            "externalJsonTable",
-            source="json",
-            schema=schema,
-            path=tmpPath,
-            noUse="this options will not be used",
-        )
-        self.assertEqual(
-            sorted(df.collect()), sorted(self.spark.sql("SELECT * FROM savedJsonTable").collect())
-        )
-        self.assertEqual(
-            sorted(df.select("value").collect()),
-            sorted(self.spark.sql("SELECT * FROM externalJsonTable").collect()),
-        )
-        self.assertEqual(sorted(df.select("value").collect()), sorted(actual.collect()))
-        self.spark.sql("DROP TABLE savedJsonTable")
-        self.spark.sql("DROP TABLE externalJsonTable")
-
-        defaultDataSourceName = self.spark.conf.get(
-            "spark.sql.sources.default", "org.apache.spark.sql.parquet"
-        )
-        self.spark.sql("SET spark.sql.sources.default=org.apache.spark.sql.json")
-        df.write.saveAsTable("savedJsonTable", path=tmpPath, mode="overwrite")
-        actual = self.spark.catalog.createTable("externalJsonTable", path=tmpPath)
-        self.assertEqual(
-            sorted(df.collect()), sorted(self.spark.sql("SELECT * FROM savedJsonTable").collect())
-        )
-        self.assertEqual(
-            sorted(df.collect()),
-            sorted(self.spark.sql("SELECT * FROM externalJsonTable").collect()),
-        )
-        self.assertEqual(sorted(df.collect()), sorted(actual.collect()))
-        self.spark.sql("DROP TABLE savedJsonTable")
-        self.spark.sql("DROP TABLE externalJsonTable")
-        self.spark.sql("SET spark.sql.sources.default=" + defaultDataSourceName)
-
-        shutil.rmtree(tmpPath)
-
-    def test_limit_and_take(self):
-        df = self.spark.range(1, 1000, numPartitions=10)
-
-        def assert_runs_only_one_job_stage_and_task(job_group_name, f):
-            tracker = self.sc.statusTracker()
-            self.sc.setJobGroup(job_group_name, description="")
-            f()
-            jobs = tracker.getJobIdsForGroup(job_group_name)
-            self.assertEqual(1, len(jobs))
-            stages = tracker.getJobInfo(jobs[0]).stageIds
-            self.assertEqual(1, len(stages))
-            self.assertEqual(1, tracker.getStageInfo(stages[0]).numTasks)
-
-        # Regression test for SPARK-10731: take should delegate to Scala implementation
-        assert_runs_only_one_job_stage_and_task("take", lambda: df.take(1))
-        # Regression test for SPARK-17514: limit(n).collect() should the perform same as take(n)
-        assert_runs_only_one_job_stage_and_task("collect_limit", lambda: df.limit(1).collect())
-
-    def test_unbounded_frames(self):
-        from pyspark.sql import functions as F
-        from pyspark.sql import window
-
-        df = self.spark.range(0, 3)
-
-        def rows_frame_match():
-            return (
-                "ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING"
-                in df.select(
-                    F.count("*").over(window.Window.rowsBetween(-sys.maxsize, sys.maxsize))
-                ).columns[0]
-            )
-
-        def range_frame_match():
-            return (
-                "RANGE BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING"
-                in df.select(
-                    F.count("*").over(window.Window.rangeBetween(-sys.maxsize, sys.maxsize))
-                ).columns[0]
-            )
-
-        for new_maxsize in [2**31 - 1, 2**63 - 1, 2**127 - 1]:
-            old_maxsize = sys.maxsize
-            sys.maxsize = new_maxsize
-            try:
-                # Manually reload window module to use monkey-patched sys.maxsize.
-                reload(window)
-                self.assertTrue(rows_frame_match())
-                self.assertTrue(range_frame_match())
-            finally:
-                sys.maxsize = old_maxsize
-
-        reload(window)
-
-
-class SQLContextTests(unittest.TestCase):
-    def test_get_or_create(self):
-        sc = None
-        sql_context = None
-        try:
-            sc = SparkContext("local[4]", "SQLContextTests")
-            sql_context = SQLContext.getOrCreate(sc)
-            assert isinstance(sql_context, SQLContext)
-        finally:
-            if sql_context is not None:
-                sql_context.sparkSession.stop()
-            if sc is not None:
-                sc.stop()
-
-
-if __name__ == "__main__":
-    from pyspark.sql.tests.test_context 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)