address comments

Author: huaxingao

python/pyspark/sql/functions.py

@@ -2714,20 +2714,22 @@ def from_csv(col, schema, options={}):
 
 
 @since(3.0)
-def getActiveSession():
+def _getActiveSession():
     """
     Returns the active SparkSession for the current thread
+    This method is not intended for user to call directly.
+    It is only used for getActiveSession method in session.py
     """
     from pyspark.sql import SparkSession
     sc = SparkContext._active_spark_context
     if sc is None:
-        sc = SparkContext()
-
-    if sc._jvm.SparkSession.getActiveSession().isDefined():
-        SparkSession(sc, sc._jvm.SparkSession.getActiveSession().get())
-        return SparkSession._activeSession
-    else:
         return None
+    else:
+        if sc._jvm.SparkSession.getActiveSession().isDefined():
+            SparkSession(sc, sc._jvm.SparkSession.getActiveSession().get())
+            return SparkSession._activeSession
+        else:
+            return None
 
 
 # ---------------------------- User Defined Function ----------------------------------

python/pyspark/sql/session.py

@@ -271,7 +271,7 @@ def getActiveSession(cls):
         [Row(age=1)]
         """
         from pyspark.sql import functions
-        return functions.getActiveSession()
+        return functions._getActiveSession()
 
     @property
     @since(2.0)
@@ -689,6 +689,8 @@ def createDataFrame(self, data, schema=None, samplingRatio=None, verifySchema=Tr
             ...
         Py4JJavaError: ...
         """
+        SparkSession._activeSession = self
+        self._jvm.SparkSession.setActiveSession(self._jsparkSession)
         if isinstance(data, DataFrame):
             raise TypeError("data is already a DataFrame")
 

python/pyspark/sql/tests.py

@@ -3965,6 +3965,42 @@ def test_create_new_session_if_old_session_stopped(self):
         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 = 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)
+
+
+class SparkSessionTests3(ReusedSQLTestCase):
+
+    def test_get_active_session_after_create_dataframe(self):
+        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)
+        session2.stop()
+
 
 class UDFInitializationTests(unittest.TestCase):
     def tearDown(self):