feat(pyspark): support ibis.pyspark.connect() (#8515)

ibis-project · Mar 1, 2024 · 0f663e6 · 0f663e6
1 parent ef39aab
commit 0f663e6
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/ibis/backends/pyspark/__init__.py b/ibis/backends/pyspark/__init__.py
@@ -130,7 +130,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._cached_dataframes = {}
 
-    def do_connect(self, session: SparkSession) -> None:
+    def do_connect(self, session: SparkSession | None = None) -> None:
         """Create a PySpark `Backend` for use with Ibis.
 
         Parameters
@@ -147,6 +147,11 @@ def do_connect(self, session: SparkSession) -> None:
         <ibis.backends.pyspark.Backend at 0x...>
 
         """
+        if session is None:
+            from pyspark.sql import SparkSession
+
+            session = SparkSession.builder.getOrCreate()
+
         self._context = session.sparkContext
         self._session = session
 

diff --git a/ibis/backends/pyspark/tests/test_basic.py b/ibis/backends/pyspark/tests/test_basic.py
@@ -125,3 +125,9 @@ def test_string_literal_backslash_escaping(con):
     expr = ibis.literal("\\d\\e")
     result = con.execute(expr)
     assert result == "\\d\\e"
+
+
+def test_connect_without_explicit_session():
+    con = ibis.pyspark.connect()
+    result = con.sql("SELECT CAST(1 AS BIGINT) as foo").to_pandas()
+    tm.assert_frame_equal(result, pd.DataFrame({"foo": [1]}))