fixes

graphframes · SemyonSinchenko · Oct 6, 2025 · Sep 24, 2025 · Sep 25, 2025 · Sep 26, 2025
commit b43b52c0ca4638865cdfd3c281e698a2bc3ed40e
diff --git a/python/graphframes/classic/graphframe.py b/python/graphframes/classic/graphframe.py
@@ -19,10 +19,16 @@
 from typing import final
 
 from py4j.java_gateway import JavaObject
-from pyspark import SparkContext
+from pyspark import SparkContext, __version__
 from pyspark.sql import SparkSession
-from pyspark.sql.classic.column import Column, _to_seq
-from pyspark.sql.classic.dataframe import DataFrame
+
+if __version__.startswith("4"):
+    from pyspark.sql.classic.column import Column, _to_seq
+    from pyspark.sql.classic.dataframe import DataFrame
+else:
+    from pyspark.sql.column import Column, _to_seq
+    from pyspark.sql import DataFrame
+
 from pyspark.storagelevel import StorageLevel
 
 from graphframes.classic.utils import storage_level_to_jvm

diff --git a/python/graphframes/connect/graphframes_client.py b/python/graphframes/connect/graphframes_client.py
@@ -1061,4 +1061,6 @@ def plan(self, session: SparkConnectClient) -> proto.Relation:
                 plan.extension.Pack(graphframes_api_call)
                 return plan
 
-        return _dataframe_from_plan(TriangleCount(self._vertices, self._edges), self._spark)
+        return _dataframe_from_plan(
+            TriangleCount(self._vertices, self._edges, storage_level), self._spark
+        )