graphframes · SemyonSinchenko · Jun 29, 2025 · Apr 1, 2025 · Apr 3, 2025 · Apr 4, 2025
diff --git a/.github/workflows/python-ci.yml b/.github/workflows/python-ci.yml
@@ -6,9 +6,14 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          - spark-version: 3.5.4
-            scala-version: 2.12.18
-            python-version: 3.10.6
+          - spark-version: 3.5.5
+            scala-version: 2.12
+            python-version: "3.10"
+            java-version: 11
+          - spark-version: 4.0.0
+            scala-version: 2.13
+            python-version: 3.12
+            java-version: 17
     runs-on: ubuntu-22.04
     env:
       # define Java options for both official sbt and sbt-extras
@@ -21,7 +26,7 @@ jobs:
     - uses: actions/setup-java@v4
       with:
         distribution: "zulu"
-        java-version: "11"
+        java-version: "${{ matrix.java-version }}"
     - uses: actions/cache@v4
       with:
         path: |
@@ -41,23 +46,24 @@ jobs:
       working-directory: ./python
       run: |
         poetry install --with=dev
+        poetry run pip install pyspark[connect]==${{ matrix.spark-version }}
     - name: Code style
       working-directory: ./python
       run: |
         poetry run python -m black --check graphframes
         poetry run python -m flake8 graphframes
         poetry run python -m isort --check graphframes
-
+    - name: Build jar
+      working-directory: ./python
+      run: |
+        poetry run python dev/build_jar.py ${{ matrix.spark-version }}
     - name: Test
       working-directory: ./python
       run: |
         poetry run python -m pytest
-
     - name: Test SparkConnect
       env:
         SPARK_CONNECT_MODE_ENABLED: 1
       working-directory: ./python
       run: |
-        poetry run python dev/run_connect.py
         poetry run python -m pytest
-        poetry run python dev/stop_connect.py
diff --git a/.github/workflows/scala-ci.yml b/.github/workflows/scala-ci.yml
@@ -6,17 +6,11 @@ jobs:
       fail-fast: false
       matrix:
         include:
-          - spark-version: 3.5.4
-            scala-version: 2.13.12
+          - spark-version: 3.5.5
             java-version: 11
-          - spark-version: 3.5.4
-            scala-version: 2.13.12
+          - spark-version: 3.5.5
             java-version: 17
-          - spark-version: 3.5.4
-            scala-version: 2.12.18
-            java-version: 11
-          - spark-version: 3.5.4
-            scala-version: 2.12.18
+          - spark-version: 4.0.0
             java-version: 17
     runs-on: ubuntu-22.04
     env:
@@ -34,11 +28,11 @@ jobs:
       with:
         path: |
           ~/.ivy2/cache
-        key: sbt-ivy-cache-spark-${{ matrix.spark-version}}-scala-${{ matrix.scala-version }}-java-${{ matrix.java-version }}
+        key: sbt-ivy-cache-spark-${{ matrix.spark-version}}-java-${{ matrix.java-version }}
     - name: Check scalafmt
-      run: build/sbt scalafmtCheckAll
+      run: build/sbt -Dspark.version=${{ matrix.spark-version }} root/scalafmtCheckAll connect/scalafmtCheckAll
     - name: Check scalastyle
-      run: build/sbt scalafixAll
+      run: build/sbt -Dspark.version=${{ matrix.spark-version }} root/scalafixAll connect/scalafixAll
     - name: Build and Test
-      run: build/sbt -v ++${{ matrix.scala-version }} -Dspark.version=${{ matrix.spark-version }} coverage test coverageReport
+      run: build/sbt -v -Dspark.version=${{ matrix.spark-version }} coverage +test coverageReport "project connect" coverage +test coverageReport
     - uses: codecov/codecov-action@v3
diff --git a/.github/workflows/scala-publish.yml b/.github/workflows/scala-publish.yml
@@ -9,6 +9,14 @@ on:
 
 jobs:
   release:
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - spark-version: 3.5.5
+            java-version: 11
+          - spark-version: 4.0.0
+            java-version: 17
     runs-on: ubuntu-latest
     steps:
       - name: Check out code
@@ -20,12 +28,12 @@ jobs:
         if: startsWith(github.ref, 'refs/tags/')
         uses: actions/setup-java@v4
         with:
-          java-version: '11'
+          java-version: "${{ matrix.java-version }}"
           distribution: 'zulu'
 
       - name: Run the release of core
         if: startsWith(github.ref, 'refs/tags/')
-        run: build/sbt ci-release
+        run: build/sbt -Dspark.version=${{ matrix.spark-version }} +ci-release "project connect" +ci-release
         env:
           PGP_PASSPHRASE: ${{ secrets.PGP_PASSPHRASE }}
           PGP_SECRET: ${{ secrets.PGP_SECRET }}

diff --git a/.gitignore b/.gitignore
@@ -27,6 +27,12 @@ metals.sbt
 # VSCode
 .vscode
 
+# metals
+.ammonite/
+.bloop/
+.metals/
+metals.sbt
+
 # Helix
 .helix
 

diff --git a/build.sbt b/build.sbt
@@ -1,16 +1,21 @@
 import xerial.sbt.Sonatype.sonatypeCentralHost
 
 lazy val sparkVer = sys.props.getOrElse("spark.version", "3.5.5")
+lazy val sparkMajorVer = sparkVer.substring(0, 1)
 lazy val sparkBranch = sparkVer.substring(0, 3)
-lazy val defaultScalaVer = sparkBranch match {
-  case "3.5" => "2.12.18"
-  case "3.4" => "2.12.17"
+lazy val scalaVersions = sparkMajorVer match {
+  case "4" => Seq("2.13.12")
+  case "3" => Seq("2.12.18", "2.13.12")
   case _ => throw new IllegalArgumentException(s"Unsupported Spark version: $sparkVer.")
 }
-lazy val scalaVer = sys.props.getOrElse("scala.version", defaultScalaVer)
-lazy val defaultScalaTestVer = scalaVer match {
-  case s if s.startsWith("2.12") || s.startsWith("2.13") => "3.0.8"
+lazy val scalaVer = sys.props.getOrElse("scala.version", scalaVersions(0))
+lazy val defaultScalaTestVer = "3.0.8"
+
+ThisBuild / version := {
+  val baseVersion = (ThisBuild / version).value
+  s"${baseVersion}-spark${sparkBranch}"
 }
+
 // Some vendors are using an own shading rule for protobuf
 lazy val protobufShadingPattern = sys.props.getOrElse("vendor.name", "oss") match {
   case "oss" => "org.sparkproject.connect.protobuf.@1"
@@ -19,6 +24,12 @@ lazy val protobufShadingPattern = sys.props.getOrElse("vendor.name", "oss") matc
     throw new IllegalArgumentException(s"Unsupported vendor name: $s; supported: 'oss', 'dbx'")
 }
 
+lazy val protocVersion = sparkMajorVer match {
+  case "4" => "4.29.3"
+  case "3" => "3.23.4"
+  case _ => throw new IllegalArgumentException(s"Unsupported Spark version: $sparkVer.")
+}
+
 ThisBuild / scalaVersion := scalaVer
 ThisBuild / organization := "org.graphframes"
 ThisBuild / homepage := Some(url("https://graphframes.io/"))
@@ -41,7 +52,7 @@ ThisBuild / developers := List(
 ThisBuild / sonatypeCredentialHost := "s01.oss.sonatype.org"
 ThisBuild / sonatypeRepository := "https://s01.oss.sonatype.org/service/local"
 ThisBuild / sonatypeProfileName := "io.graphframes"
-ThisBuild / crossScalaVersions := Seq("2.12.18", "2.13.12")
+ThisBuild / crossScalaVersions := scalaVersions
 
 // Scalafix configuration
 ThisBuild / semanticdbEnabled := true
@@ -52,7 +63,7 @@ lazy val commonSetting = Seq(
     "org.apache.spark" %% "spark-graphx" % sparkVer % "provided" cross CrossVersion.for3Use2_13,
     "org.apache.spark" %% "spark-sql" % sparkVer % "provided" cross CrossVersion.for3Use2_13,
     "org.apache.spark" %% "spark-mllib" % sparkVer % "provided" cross CrossVersion.for3Use2_13,
-    "org.slf4j" % "slf4j-api" % "2.0.16",
+    "org.slf4j" % "slf4j-api" % "2.0.16" % "provided",
     "org.scalatest" %% "scalatest" % defaultScalaTestVer % Test,
     "com.github.zafarkhaja" % "java-semver" % "0.10.2" % Test),
   Compile / scalacOptions ++= Seq("-deprecation", "-feature"),
@@ -96,15 +107,18 @@ lazy val root = (project in file("."))
   .settings(
     commonSetting,
     name := "graphframes",
-    moduleName := s"${name.value}-spark${sparkBranch}",
+    moduleName := s"${name.value}-spark$sparkMajorVer",
 
     // Global settings
     Global / concurrentRestrictions := Seq(Tags.limitAll(1)),
     autoAPIMappings := true,
     coverageHighlighting := false,
 
+    Compile / unmanagedSourceDirectories += (Compile / baseDirectory).value / "src" / "main" / s"scala-spark-$sparkMajorVer",
+
     // Assembly settings
     assembly / test := {}, // No tests in assembly
+    assemblyPackageScala / assembleArtifact := false,
     assembly / assemblyMergeStrategy := {
       case PathList("META-INF", xs @ _*) => MergeStrategy.discard
       case x if x.endsWith("module-info.class") => MergeStrategy.discard
@@ -123,16 +137,18 @@ lazy val connect = (project in file("graphframes-connect"))
   .dependsOn(root)
   .settings(
     commonSetting,
-    name := "graphframes-connect",
-    moduleName := s"${name.value}-spark${sparkBranch}",
+    name := s"graphframes-connect",
+    moduleName := s"${name.value}-spark${sparkMajorVer}",
+    Compile / unmanagedSourceDirectories += (Compile / baseDirectory).value / "src" / "main" / s"scala-spark-$sparkMajorVer",
     Compile / PB.targets := Seq(PB.gens.java -> (Compile / sourceManaged).value),
     Compile / PB.includePaths ++= Seq(file("src/main/protobuf")),
-    PB.protocVersion := "3.23.4", // Spark 3.5 branch
+    PB.protocVersion := protocVersion,
     libraryDependencies ++= Seq(
       "org.apache.spark" %% "spark-connect" % sparkVer % "provided" cross CrossVersion.for3Use2_13),
 
     // Assembly and shading
     assembly / test := {},
+    assemblyPackageScala / assembleArtifact := false,
     assembly / assemblyShadeRules := Seq(
       ShadeRule.rename("com.google.protobuf.**" -> protobufShadingPattern).inAll),
     assembly / assemblyMergeStrategy := {

diff --git a/.../sql/graphframes/GraphFramesConnect.scala → .../sql/graphframes/GraphFramesConnect.scala b/.../sql/graphframes/GraphFramesConnect.scala → .../sql/graphframes/GraphFramesConnect.scala
diff --git a/...-connect/src/main/scala-spark-4/org/apache/spark/sql/graphframes/GraphFramesConnect.scala b/...-connect/src/main/scala-spark-4/org/apache/spark/sql/graphframes/GraphFramesConnect.scala
@@ -0,0 +1,27 @@
+package org.apache.spark.sql.graphframes
+
+import java.util.Optional
+
+import org.graphframes.connect.proto.GraphFramesAPI
+
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.connect.planner.SparkConnectPlanner
+import org.apache.spark.sql.connect.plugin.RelationPlugin
+
+import com.google.protobuf
+
+class GraphFramesConnect extends RelationPlugin {
+  override def transform(
+      relation: Array[Byte],
+      planner: SparkConnectPlanner): Optional[LogicalPlan] = {
+    val rel = protobuf.Any.parseFrom(relation)
+    if (rel.is(classOf[GraphFramesAPI])) {
+      val protoCall = rel.unpack(classOf[GraphFramesAPI])
+      // Because the plugins API is changed in spark 4.0 it makes sense to separate plugin impl from the parsing logic
+      val result = GraphFramesConnectUtils.parseAPICall(protoCall, planner)
+      Optional.of(SparkShims.planFromDataFrame(result))
+    } else {
+      Optional.empty()
+    }
+  }
+}
diff --git a/...mes-connect/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala b/...mes-connect/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala
@@ -5,7 +5,6 @@ package org.apache.spark.sql.graphframes
 import com.google.protobuf.ByteString
 import org.apache.spark.sql.Column
 import org.apache.spark.sql.DataFrame
-import org.apache.spark.sql.Dataset
 import org.apache.spark.sql.connect.planner.SparkConnectPlanner
 import org.apache.spark.sql.functions.expr
 import org.apache.spark.sql.functions.lit
@@ -26,7 +25,7 @@ object GraphFramesConnectUtils {
       planner: SparkConnectPlanner): Column = {
     colOrExpr.getColOrExprCase match {
       case ColOrExprCase.COL =>
-        Column(
+        SparkShims.createColumn(
           planner.transformExpression(
             org.apache.spark.connect.proto.Expression.parseFrom(colOrExpr.getCol.toByteArray)))
       case ColOrExprCase.EXPR => expr(colOrExpr.getExpr)
@@ -53,7 +52,7 @@ object GraphFramesConnectUtils {
       throw new IllegalArgumentException(
         "Expected a serialized DataFrame but got an empty ByteString.")
     }
-    Dataset.ofRows(
+    SparkShims.createDataFrame(
       planner.sessionHolder.session,
       planner.transformRelation(
         org.apache.spark.connect.proto.Plan.parseFrom(data.toByteArray).getRoot))
@@ -190,7 +189,7 @@ object GraphFramesConnectUtils {
       case MethodCase.SHORTEST_PATHS => {
         graphFrame.shortestPaths
           .landmarks(
-            apiMessage.getShortestPaths.getLandmarksList.asScala.map(parseLongOrStringID))
+            apiMessage.getShortestPaths.getLandmarksList.asScala.map(parseLongOrStringID).toSeq)
           .run()
       }
       case MethodCase.STRONGLY_CONNECTED_COMPONENTS => {

diff --git a/python/dev/build_jar.py b/python/dev/build_jar.py
@@ -1,4 +1,3 @@
-import shutil
 import subprocess
 import sys
 from collections.abc import Sequence
@@ -9,14 +8,17 @@ def build(spark_versions: Sequence[str] = ["3.5.5"]):
     for spark_version in spark_versions:
         print("Building GraphFrames JAR...")
         print(f"SPARK_VERSION: {spark_version[:3]}")
-        assert spark_version[:3] in {"3.5",}, "Unsopported spark version!"
+        assert spark_version[:3] in {"3.5", "4.0"}, "Unsupported spark version!"
+
         project_root = Path(__file__).parent.parent.parent
         sbt_executable = project_root.joinpath("build").joinpath("sbt").absolute().__str__()
         sbt_build_command = [
             sbt_executable,
             f"-Dspark.version={spark_version}",
             "clean",
             "assembly",
+            "connect/clean",
+            "connect/assembly"
         ]
         sbt_build = subprocess.Popen(
             sbt_build_command,
@@ -37,22 +39,6 @@ def build(spark_versions: Sequence[str] = ["3.5.5"]):
         else:
             print("Building DONE successfully!")
 
-        python_resources = (
-            project_root.joinpath("python").joinpath("graphframes").joinpath("resources")
-        )
-        target_dir = project_root.joinpath("target").joinpath("scala-2.12")
-        gf_jar = None
-
-        for pp in target_dir.glob("*.jar"):
-            if "graphframes-assembly" in pp.name:
-                gf_jar = pp
-                break
-
-        assert gf_jar is not None, "Missing JAR!"
-        shutil.rmtree(python_resources, ignore_errors=True)
-        python_resources.mkdir(parents=True, exist_ok=True)
-        shutil.copy(gf_jar, python_resources.joinpath(f"spark-{spark_version}-{gf_jar.name}"))
-
 
 if __name__ == "__main__":
     if len(sys.argv) > 1: