This is an automated email from the ASF dual-hosted git repository.
weichenxu123 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 753864fedee [SPARK-42732][PYSPARK][CONNECT] Support spark connect
session getActiveSession method
753864fedee is described below
commit 753864fedee62f638354040063d95b2b3ba93d46
Author: Weichen Xu <weichen...@databricks.com>
AuthorDate: Tue Mar 14 18:33:24 2023 +0800
[SPARK-42732][PYSPARK][CONNECT] Support spark connect session
getActiveSession method
### What changes were proposed in this pull request?
Support spark connect session getActiveSession method.
Spark connect ML needs this API to get active session in some cases (e.g.
fetching model attributes from server side).
### Why are the changes needed?
Manually.
### Does this PR introduce _any_ user-facing change?
Yes. Implemented
`pyspark.sql.connect.session.SparkSession.getActiveSession` API.
### How was this patch tested?
N/A
Closes #40353 from WeichenXu123/spark-connect-get-active-session.
Authored-by: Weichen Xu <weichen...@databricks.com>
Signed-off-by: Weichen Xu <weichen...@databricks.com>
---
python/pyspark/sql/connect/session.py | 13 ++++++++++++-
1 file changed, 12 insertions(+), 1 deletion(-)
diff --git a/python/pyspark/sql/connect/session.py
b/python/pyspark/sql/connect/session.py
index 5e7c8361d80..ffa139eba3e 100644
--- a/python/pyspark/sql/connect/session.py
+++ b/python/pyspark/sql/connect/session.py
@@ -74,6 +74,11 @@ if TYPE_CHECKING:
from pyspark.sql.connect.udf import UDFRegistration
+# `_active_spark_session` stores the active spark connect session created by
+# `SparkSession.builder.getOrCreate`. It is used by ML code.
+_active_spark_session = None
+
+
class SparkSession:
class Builder:
"""Builder for :class:`SparkSession`."""
@@ -119,7 +124,11 @@ class SparkSession:
raise NotImplementedError("enableHiveSupport not implemented for
Spark Connect")
def getOrCreate(self) -> "SparkSession":
- return SparkSession(connectionString=self._options["spark.remote"])
+ global _active_spark_session
+ if _active_spark_session is not None:
+ return _active_spark_session
+ _active_spark_session =
SparkSession(connectionString=self._options["spark.remote"])
+ return _active_spark_session
_client: SparkConnectClient
@@ -434,7 +443,9 @@ class SparkSession:
# specifically in Spark Connect the Spark Connect server is designed
for
# multi-tenancy - the remote client side cannot just stop the server
and stop
# other remote clients being used from other users.
+ global _active_spark_session
self.client.close()
+ _active_spark_session = None
if "SPARK_LOCAL_REMOTE" in os.environ:
# When local mode is in use, follow the regular Spark session's
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@spark.apache.org
For additional commands, e-mail: commits-h...@spark.apache.org
