diff options
| -rw-r--r-- | python/pyspark/sql/session.py | 16 |
1 files changed, 16 insertions, 0 deletions
diff --git a/python/pyspark/sql/session.py b/python/pyspark/sql/session.py index 6edbd59856..0c8024e4a8 100644 --- a/python/pyspark/sql/session.py +++ b/python/pyspark/sql/session.py @@ -581,6 +581,22 @@ class SparkSession(object): """ self._sc.stop() + @since(2.0) + def __enter__(self): + """ + Enable 'with SparkSession.builder.(...).getOrCreate() as session: app' syntax. + """ + return self + + @since(2.0) + def __exit__(self, exc_type, exc_val, exc_tb): + """ + Enable 'with SparkSession.builder.(...).getOrCreate() as session: app' syntax. + + Specifically stop the SparkSession on exit of the with block. + """ + self.stop() + def _test(): import os |