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

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

@@ -1,412 +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 unittest
-
-from pyspark import SparkConf, SparkContext
-from pyspark.sql import SparkSession, SQLContext, Row
-from pyspark.sql.functions import col
-from pyspark.testing.sqlutils import ReusedSQLTestCase
-from pyspark.testing.utils import PySparkTestCase
-
-
-class SparkSessionTests(ReusedSQLTestCase):
-    def test_sqlcontext_reuses_sparksession(self):
-        sqlContext1 = SQLContext(self.sc)
-        sqlContext2 = SQLContext(self.sc)
-        self.assertTrue(sqlContext1.sparkSession is sqlContext2.sparkSession)
-
-
-class SparkSessionTests1(ReusedSQLTestCase):
-
-    # We can't include this test into SQLTests because we will stop class's SparkContext and cause
-    # other tests failed.
-    def test_sparksession_with_stopped_sparkcontext(self):
-        self.sc.stop()
-        sc = SparkContext("local[4]", self.sc.appName)
-        spark = SparkSession.builder.getOrCreate()
-        try:
-            df = spark.createDataFrame([(1, 2)], ["c", "c"])
-            df.collect()
-        finally:
-            spark.stop()
-            sc.stop()
-
-
-class SparkSessionTests2(PySparkTestCase):
-
-    # This test is separate because it's closely related with session's start and stop.
-    # See SPARK-23228.
-    def test_set_jvm_default_session(self):
-        spark = SparkSession.builder.getOrCreate()
-        try:
-            self.assertTrue(spark._jvm.SparkSession.getDefaultSession().isDefined())
-        finally:
-            spark.stop()
-            self.assertTrue(spark._jvm.SparkSession.getDefaultSession().isEmpty())
-
-    def test_jvm_default_session_already_set(self):
-        # Here, we assume there is the default session already set in JVM.
-        jsession = self.sc._jvm.SparkSession(self.sc._jsc.sc())
-        self.sc._jvm.SparkSession.setDefaultSession(jsession)
-
-        spark = SparkSession.builder.getOrCreate()
-        try:
-            self.assertTrue(spark._jvm.SparkSession.getDefaultSession().isDefined())
-            # The session should be the same with the exiting one.
-            self.assertTrue(jsession.equals(spark._jvm.SparkSession.getDefaultSession().get()))
-        finally:
-            spark.stop()
-
-
-class SparkSessionTests3(unittest.TestCase):
-    def test_active_session(self):
-        spark = SparkSession.builder.master("local").getOrCreate()
-        try:
-            activeSession = SparkSession.getActiveSession()
-            df = activeSession.createDataFrame([(1, "Alice")], ["age", "name"])
-            self.assertEqual(df.collect(), [Row(age=1, name="Alice")])
-            with self.assertRaises(ValueError):
-                activeSession.createDataFrame(activeSession._sc.parallelize([[], []]))
-        finally:
-            spark.stop()
-
-    def test_get_active_session_when_no_active_session(self):
-        active = SparkSession.getActiveSession()
-        self.assertEqual(active, None)
-        spark = SparkSession.builder.master("local").getOrCreate()
-        active = SparkSession.getActiveSession()
-        self.assertEqual(active, spark)
-        spark.stop()
-        active = SparkSession.getActiveSession()
-        self.assertEqual(active, None)
-
-    def test_spark_session(self):
-        spark = SparkSession.builder.master("local").config("some-config", "v2").getOrCreate()
-        try:
-            self.assertEqual(spark.conf.get("some-config"), "v2")
-            self.assertEqual(spark.sparkContext._conf.get("some-config"), "v2")
-            self.assertEqual(spark.version, spark.sparkContext.version)
-            spark.sql("CREATE DATABASE test_db")
-            spark.catalog.setCurrentDatabase("test_db")
-            self.assertEqual(spark.catalog.currentDatabase(), "test_db")
-            spark.sql("CREATE TABLE table1 (name STRING, age INT) USING parquet")
-            self.assertEqual(spark.table("table1").columns, ["name", "age"])
-            self.assertEqual(spark.range(3).count(), 3)
-
-            # SPARK-37516: Only plain column references work as variable in SQL.
-            self.assertEqual(
-                spark.sql("select {c} from range(1)", c=col("id")).first(), spark.range(1).first()
-            )
-            with self.assertRaisesRegex(ValueError, "Column"):
-                spark.sql("select {c} from range(10)", c=col("id") + 1)
-        finally:
-            spark.sql("DROP DATABASE test_db CASCADE")
-            spark.stop()
-
-    def test_global_default_session(self):
-        spark = SparkSession.builder.master("local").getOrCreate()
-        try:
-            self.assertEqual(SparkSession.builder.getOrCreate(), spark)
-        finally:
-            spark.stop()
-
-    def test_default_and_active_session(self):
-        spark = SparkSession.builder.master("local").getOrCreate()
-        activeSession = spark._jvm.SparkSession.getActiveSession()
-        defaultSession = spark._jvm.SparkSession.getDefaultSession()
-        try:
-            self.assertEqual(activeSession, defaultSession)
-        finally:
-            spark.stop()
-
-    def test_config_option_propagated_to_existing_session(self):
-        session1 = SparkSession.builder.master("local").config("spark-config1", "a").getOrCreate()
-        self.assertEqual(session1.conf.get("spark-config1"), "a")
-        session2 = SparkSession.builder.config("spark-config1", "b").getOrCreate()
-        try:
-            self.assertEqual(session1, session2)
-            self.assertEqual(session1.conf.get("spark-config1"), "b")
-        finally:
-            session1.stop()
-
-    def test_new_session(self):
-        session = SparkSession.builder.master("local").getOrCreate()
-        newSession = session.newSession()
-        try:
-            self.assertNotEqual(session, newSession)
-        finally:
-            session.stop()
-            newSession.stop()
-
-    def test_create_new_session_if_old_session_stopped(self):
-        session = SparkSession.builder.master("local").getOrCreate()
-        session.stop()
-        newSession = SparkSession.builder.master("local").getOrCreate()
-        try:
-            self.assertNotEqual(session, newSession)
-        finally:
-            newSession.stop()
-
-    def test_active_session_with_None_and_not_None_context(self):
-        from pyspark.context import SparkContext
-        from pyspark.conf import SparkConf
-
-        sc = None
-        session = None
-        try:
-            sc = SparkContext._active_spark_context
-            self.assertEqual(sc, None)
-            activeSession = SparkSession.getActiveSession()
-            self.assertEqual(activeSession, None)
-            sparkConf = SparkConf()
-            sc = SparkContext.getOrCreate(sparkConf)
-            activeSession = sc._jvm.SparkSession.getActiveSession()
-            self.assertFalse(activeSession.isDefined())
-            session = SparkSession(sc)
-            activeSession = sc._jvm.SparkSession.getActiveSession()
-            self.assertTrue(activeSession.isDefined())
-            activeSession2 = SparkSession.getActiveSession()
-            self.assertNotEqual(activeSession2, None)
-        finally:
-            if session is not None:
-                session.stop()
-            if sc is not None:
-                sc.stop()
-
-
-class SparkSessionTests4(ReusedSQLTestCase):
-    def test_get_active_session_after_create_dataframe(self):
-        session2 = None
-        try:
-            activeSession1 = SparkSession.getActiveSession()
-            session1 = self.spark
-            self.assertEqual(session1, activeSession1)
-            session2 = self.spark.newSession()
-            activeSession2 = SparkSession.getActiveSession()
-            self.assertEqual(session1, activeSession2)
-            self.assertNotEqual(session2, activeSession2)
-            session2.createDataFrame([(1, "Alice")], ["age", "name"])
-            activeSession3 = SparkSession.getActiveSession()
-            self.assertEqual(session2, activeSession3)
-            session1.createDataFrame([(1, "Alice")], ["age", "name"])
-            activeSession4 = SparkSession.getActiveSession()
-            self.assertEqual(session1, activeSession4)
-        finally:
-            if session2 is not None:
-                session2.stop()
-
-
-class SparkSessionTests5(unittest.TestCase):
-    def setUp(self):
-        # These tests require restarting the Spark context so we set up a new one for each test
-        # rather than at the class level.
-        self.sc = SparkContext("local[4]", self.__class__.__name__, conf=SparkConf())
-        self.spark = SparkSession(self.sc)
-
-    def tearDown(self):
-        self.sc.stop()
-        self.spark.stop()
-
-    def test_sqlcontext_with_stopped_sparksession(self):
-        # SPARK-30856: test that SQLContext.getOrCreate() returns a usable instance after
-        # the SparkSession is restarted.
-        sql_context = SQLContext.getOrCreate(self.spark.sparkContext)
-        self.spark.stop()
-        spark = SparkSession.builder.master("local[4]").appName(self.sc.appName).getOrCreate()
-        new_sql_context = SQLContext.getOrCreate(spark.sparkContext)
-
-        self.assertIsNot(new_sql_context, sql_context)
-        self.assertIs(SQLContext.getOrCreate(spark.sparkContext).sparkSession, spark)
-        try:
-            df = spark.createDataFrame([(1, 2)], ["c", "c"])
-            df.collect()
-        finally:
-            spark.stop()
-            self.assertIsNone(SQLContext._instantiatedContext)
-
-    def test_sqlcontext_with_stopped_sparkcontext(self):
-        # SPARK-30856: test initialization via SparkSession when only the SparkContext is stopped
-        self.sc.stop()
-        spark = SparkSession.builder.master("local[4]").appName(self.sc.appName).getOrCreate()
-        self.sc = spark.sparkContext
-        self.assertIs(SQLContext.getOrCreate(self.sc).sparkSession, spark)
-
-    def test_get_sqlcontext_with_stopped_sparkcontext(self):
-        # SPARK-30856: test initialization via SQLContext.getOrCreate() when only the SparkContext
-        # is stopped
-        self.sc.stop()
-        self.sc = SparkContext("local[4]", self.sc.appName)
-        self.assertIs(SQLContext.getOrCreate(self.sc)._sc, self.sc)
-
-
-class SparkSessionBuilderTests(unittest.TestCase):
-    def test_create_spark_context_first_then_spark_session(self):
-        sc = None
-        session = None
-        try:
-            conf = SparkConf().set("key1", "value1")
-            sc = SparkContext("local[4]", "SessionBuilderTests", conf=conf)
-            session = SparkSession.builder.config("key2", "value2").getOrCreate()
-
-            self.assertEqual(session.conf.get("key1"), "value1")
-            self.assertEqual(session.conf.get("key2"), "value2")
-            self.assertEqual(session.sparkContext, sc)
-
-            self.assertFalse(sc.getConf().contains("key2"))
-            self.assertEqual(sc.getConf().get("key1"), "value1")
-        finally:
-            if session is not None:
-                session.stop()
-            if sc is not None:
-                sc.stop()
-
-    def test_another_spark_session(self):
-        session1 = None
-        session2 = None
-        try:
-            session1 = SparkSession.builder.config("key1", "value1").getOrCreate()
-            session2 = SparkSession.builder.config(
-                "spark.sql.codegen.comments", "true"
-            ).getOrCreate()
-
-            self.assertEqual(session1.conf.get("key1"), "value1")
-            self.assertEqual(session2.conf.get("key1"), "value1")
-            self.assertEqual(session1.conf.get("spark.sql.codegen.comments"), "false")
-            self.assertEqual(session2.conf.get("spark.sql.codegen.comments"), "false")
-            self.assertEqual(session1.sparkContext, session2.sparkContext)
-
-            self.assertEqual(session1.sparkContext.getConf().get("key1"), "value1")
-            self.assertFalse(session1.sparkContext.getConf().contains("key2"))
-        finally:
-            if session1 is not None:
-                session1.stop()
-            if session2 is not None:
-                session2.stop()
-
-    def test_create_spark_context_with_initial_session_options(self):
-        sc = None
-        session = None
-        try:
-            conf = SparkConf().set("key1", "value1")
-            sc = SparkContext("local[4]", "SessionBuilderTests", conf=conf)
-            session = (
-                SparkSession.builder.config("spark.sql.codegen.comments", "true")
-                .enableHiveSupport()
-                .getOrCreate()
-            )
-
-            self.assertEqual(session._jsparkSession.sharedState().conf().get("key1"), "value1")
-            self.assertEqual(
-                session._jsparkSession.sharedState().conf().get("spark.sql.codegen.comments"),
-                "true",
-            )
-            self.assertEqual(
-                session._jsparkSession.sharedState().conf().get("spark.sql.catalogImplementation"),
-                "hive",
-            )
-            self.assertEqual(session.sparkContext, sc)
-        finally:
-            if session is not None:
-                session.stop()
-            if sc is not None:
-                sc.stop()
-
-    def test_create_spark_context_with_initial_session_options_bool(self):
-        session = None
-        # Test if `True` is set as "true".
-        try:
-            session = SparkSession.builder.config(
-                "spark.sql.pyspark.jvmStacktrace.enabled", True
-            ).getOrCreate()
-            self.assertEqual(session.conf.get("spark.sql.pyspark.jvmStacktrace.enabled"), "true")
-        finally:
-            if session is not None:
-                session.stop()
-        # Test if `False` is set as "false".
-        try:
-            session = SparkSession.builder.config(
-                "spark.sql.pyspark.jvmStacktrace.enabled", False
-            ).getOrCreate()
-            self.assertEqual(session.conf.get("spark.sql.pyspark.jvmStacktrace.enabled"), "false")
-        finally:
-            if session is not None:
-                session.stop()
-
-
-class SparkExtensionsTest(unittest.TestCase):
-    # These tests are separate because it uses 'spark.sql.extensions' which is
-    # static and immutable. This can't be set or unset, for example, via `spark.conf`.
-
-    @classmethod
-    def setUpClass(cls):
-        import glob
-        from pyspark.find_spark_home import _find_spark_home
-
-        SPARK_HOME = _find_spark_home()
-        filename_pattern = (
-            "sql/core/target/scala-*/test-classes/org/apache/spark/sql/"
-            "SparkSessionExtensionSuite.class"
-        )
-        if not glob.glob(os.path.join(SPARK_HOME, filename_pattern)):
-            raise unittest.SkipTest(
-                "'org.apache.spark.sql.SparkSessionExtensionSuite' is not "
-                "available. Will skip the related tests."
-            )
-
-        # Note that 'spark.sql.extensions' is a static immutable configuration.
-        cls.spark = (
-            SparkSession.builder.master("local[4]")
-            .appName(cls.__name__)
-            .config("spark.sql.extensions", "org.apache.spark.sql.MyExtensions")
-            .getOrCreate()
-        )
-
-    @classmethod
-    def tearDownClass(cls):
-        cls.spark.stop()
-
-    def test_use_custom_class_for_extensions(self):
-        self.assertTrue(
-            self.spark._jsparkSession.sessionState()
-            .planner()
-            .strategies()
-            .contains(
-                self.spark._jvm.org.apache.spark.sql.MySparkStrategy(self.spark._jsparkSession)
-            ),
-            "MySparkStrategy not found in active planner strategies",
-        )
-        self.assertTrue(
-            self.spark._jsparkSession.sessionState()
-            .analyzer()
-            .extendedResolutionRules()
-            .contains(self.spark._jvm.org.apache.spark.sql.MyRule(self.spark._jsparkSession)),
-            "MyRule not found in extended resolution rules",
-        )
-
-
-if __name__ == "__main__":
-    from pyspark.sql.tests.test_session 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)