diff --git a/python/graphframes/graphframe.py b/python/graphframes/graphframe.py index 92833fab8..7145d3672 100644 --- a/python/graphframes/graphframe.py +++ b/python/graphframes/graphframe.py @@ -63,7 +63,7 @@ class GraphFrame(object): def __init__(self, v, e): self._vertices = v self._edges = e - self._spark = SparkSession.getActiveSession() + self._spark = v.sparkSession self._sc = self._spark._sc self._jvm_gf_api = _java_api(self._sc) diff --git a/python/graphframes/tests.py b/python/graphframes/tests.py index f09012d01..9a7ad1371 100644 --- a/python/graphframes/tests.py +++ b/python/graphframes/tests.py @@ -409,6 +409,28 @@ def test_triangle_counts(self): c = g.triangleCount() for row in c.select("id", "count").collect(): self.assertEqual(row.asDict()['count'], 1) + + def test_mutithreaded_sparksession_usage(self): + # Test that we can use the GraphFrame API from multiple threads + localVertices = [(1, "A"), (2, "B"), (3, "C")] + localEdges = [(1, 2, "love"), (2, 1, "hate"), (2, 3, "follow")] + v = self.spark.createDataFrame(localVertices, ["id", "name"]) + e = self.spark.createDataFrame(localEdges, ["src", "dst", "action"]) + + + exc = None + def run_graphframe() -> None: + try: + GraphFrame(v, e) + except Exception as _e: + nonlocal exc + exc = _e + + import threading + thread = threading.Thread(target=run_graphframe) + thread.start() + thread.join() + self.assertIsNone(exc, f"Exception was raised in thread: {exc}") class GraphFrameExamplesTest(GraphFrameTestCase):