Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 64 additions & 0 deletions .github/workflows/pyspark-client-ci.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
name: Python thin client CI
on: [push, pull_request]
permissions:
contents: read
jobs:
test:
strategy:
fail-fast: false
matrix:
include:
- spark-version: 4.0.3
scala-version: 2.13
python-version: "3.12"
java-version: 17
- spark-version: 4.1.2
scala-version: 2.13
python-version: "3.13"
java-version: 17
runs-on: ubuntu-latest
env:
# define Java options for both official sbt and sbt-extras
JAVA_OPTS: -Xms2048M -Xmx2048M -Xss6M -XX:ReservedCodeCacheSize=256M -Dfile.encoding=UTF-8
JVM_OPTS: -Xms2048M -Xmx2048M -Xss6M -XX:ReservedCodeCacheSize=256M -Dfile.encoding=UTF-8
SPARK_VERSION: ${{ matrix.spark-version }}
SCALA_VERSION: ${{ matrix.scala-version }}
PIP_REQUESTS_TIMEOUT: 100

steps:
- uses: actions/checkout@v6
- uses: actions/setup-java@v5
with:
distribution: "zulu"
java-version: "${{ matrix.java-version }}"
- uses: actions/cache@v5
with:
path: |
~/.ivy2/cache
key: sbt-ivy-cache-spark-${{ matrix.spark-version}}-scala-${{ matrix.scala-version }}
- uses: actions/setup-python@v6
with:
python-version: ${{ matrix.python-version }}
- name: Install and configure Poetry
uses: snok/install-poetry@v1
with:
version: 2.1.3
virtualenvs-create: true
virtualenvs-in-project: false
installer-parallel: true
- name: Build Python package and its dependencies
working-directory: ./python
run: |
poetry install --with=connect
poetry run pip install pyspark-client==${{ matrix.spark-version }}

- name: Test SparkConnect
env:
SPARK_CLIENT_MODE_ENABLED: 1
working-directory: ./python
run: |
poetry run python dev/run_connect.py ${{ matrxi.spark-version }}
Comment thread
SemyonSinchenko marked this conversation as resolved.
Outdated
poetry run python -m pytest
poetry run python dev/stop_connect.py ${{ matrxi.spark-version }}
Comment thread
SemyonSinchenko marked this conversation as resolved.
Outdated


163 changes: 163 additions & 0 deletions python/dev/run_connect.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,163 @@
#!/usr/bin/python
Comment thread
james-willis marked this conversation as resolved.

import os
import shutil
import subprocess
import sys
from pathlib import Path

import argparse

SPARK_ARCHIVE_LINK = "https://dlcdn.apache.org/spark/spark-{}/spark-{}-bin-hadoop3-connect.tgz"


if __name__ == "__main__":
parser = argparse.ArgumentParser()
_ = parser.add_argument("spark", type=str, default="4.1.2")
Comment thread
SemyonSinchenko marked this conversation as resolved.
Outdated
args = parser.parse_args()
spark = str(args.spark)

spark_full_link = SPARK_ARCHIVE_LINK.format(spark, spark)

prj_root = Path(__file__).parent.parent.parent

print("Build Graphframes...")
os.chdir(prj_root)

build_command = ["./build/sbt", f"-Dspark.version={spark}", "clean", "+", "package"]
build_sbt = subprocess.run(
build_command,
stdout=subprocess.PIPE,
Comment thread
SemyonSinchenko marked this conversation as resolved.
universal_newlines=True,
)

if build_sbt.returncode == 0:
print("Done.")
else:
print(f"SBT build return an error: {build_sbt.returncode}")
print("stdout: ", build_sbt.stdout)
print("stderr: ", build_sbt.stderr)

sys.exit(1)

tmp_dir = prj_root.joinpath("tmp")
tmp_dir.mkdir(exist_ok=True)
os.chdir(tmp_dir)

spark_archive = spark_full_link.split("/")[-1]
unpackaed_spark_binary = spark_archive[:-4]
Comment thread
SemyonSinchenko marked this conversation as resolved.
Outdated

if not tmp_dir.joinpath(unpackaed_spark_binary).exists():
print(f"Downloading spark {spark}...")
if tmp_dir.joinpath(spark_archive).exists():
shutil.rmtree(
tmp_dir.joinpath(spark_archive),
ignore_errors=True,
)
print(f"Link is resolved to {spark_full_link}")

get_spark = subprocess.run(
[
"wget",
"--no-verbose",
spark_full_link,
],
stdout=subprocess.PIPE,
Comment thread
SemyonSinchenko marked this conversation as resolved.
universal_newlines=True,
)
if get_spark.returncode == 0:
print("Done.")
else:
print("Downlading failed.")
Comment thread
SemyonSinchenko marked this conversation as resolved.
Outdated
print("stdout: ", get_spark.stdout)
print("stdeerr: ", get_spark.stderr)
Comment thread
SemyonSinchenko marked this conversation as resolved.
Outdated
sys.exit(1)

print("Unpacking Spark...")
unpack_spark = subprocess.run(
[
"tar",
"-xzf",
spark_archive,
],
stdout=subprocess.PIPE,
Comment thread
SemyonSinchenko marked this conversation as resolved.
Outdated
universal_newlines=True,
)
if unpack_spark.returncode == 0:
print("Done.")
else:
print("Unpacking failed.")
print("stdout: ", unpack_spark.stdout)
print("stdeerr: ", unpack_spark.stderr)
Comment thread
SemyonSinchenko marked this conversation as resolved.
Outdated
sys.exit(1)

spark_home = tmp_dir.joinpath(unpackaed_spark_binary)
os.chdir(spark_home)

connect_jar = None
target_dir = prj_root.joinpath("connect").joinpath("target").joinpath("scala-2.13")
print(f"looking for the connect JAR in {target_dir.absolute()}")
for ff in target_dir.glob("graphframes-connect-spark4*"):
connect_jar = ff
break

if connect_jar is None:
raise ValueError("faile to locate connect JAR")
Comment thread
SemyonSinchenko marked this conversation as resolved.
Outdated

graphx_jar = None
target_dir = prj_root.joinpath("graphx").joinpath("target").joinpath("scala-2.13")
print(f"looking for the graphx JAR in {target_dir.absolute()}")
for ff in target_dir.glob("graphframes-graphx-spark4*"):
graphx_jar = ff
break

if graphx_jar is None:
raise ValueError("faile to locate graphx JAR")
Comment thread
SemyonSinchenko marked this conversation as resolved.
Outdated

core_jar = None
target_dir = prj_root.joinpath("core").joinpath("target").joinpath("scala-2.13")
print(f"looking for the core JAR in {target_dir.absolute()}")
for ff in target_dir.glob("graphframes-spark4*"):
core_jar = ff
break

if core_jar is None:
raise ValueError("faile to locate core JAR")
Comment thread
SemyonSinchenko marked this conversation as resolved.
Outdated


_ = shutil.copyfile(core_jar, spark_home.joinpath(core_jar.name))
_ = shutil.copyfile(graphx_jar, spark_home.joinpath(graphx_jar.name))
_ = shutil.copyfile(connect_jar, spark_home.joinpath(connect_jar.name))
checkpoint_dir = Path("/tmp/GFTestsCheckpointDir")
if checkpoint_dir.exists():
shutil.rmtree(checkpoint_dir.absolute().__str__(), ignore_errors=True)

checkpoint_dir.mkdir(exist_ok=True, parents=True)

run_connect_command = [
"./sbin/start-connect-server.sh",
"--jars",
f"{core_jar.name},{graphx_jar.name},{connect_jar.name}",
"--conf",
"spark.connect.extensions.relation.classes=org.apache.spark.sql.graphframes.GraphFramesConnect",
"--conf",
"spark.checkpoint.dir=/tmp/GFTestsCheckpointDir",
"--conf",
"spark.driver.memory=6g",
"--conf",
"spark.sql.shuffle.partitions=4",
]

print("Starting SparkConnect Server...")
spark_connect = subprocess.run(
run_connect_command,
stdout=subprocess.PIPE,
universal_newlines=True,
)

if spark_connect.returncode == 0:
print("Done.")
sys.exit(0)
else:
print("Failed to start SparkConnect server.")
sys.exit(1)
44 changes: 44 additions & 0 deletions python/dev/stop_connect.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
#!/usr/bin/python


import os
import shutil
import subprocess
import sys
from pathlib import Path

import argparse


if __name__ == "__main__":
parser = argparse.ArgumentParser()
_ = parser.add_argument("spark", type=str, default="4.1.2")
args = parser.parse_args()
spark = str(args.spark)

prj_root = Path(__file__).parent.parent.parent
scala_root = prj_root.joinpath("connect")
tmp_dir = prj_root.joinpath("tmp")
unpackaed_spark_binary = f"spark-{spark}-bin-hadoop3-connect"
spark_home = tmp_dir.joinpath(unpackaed_spark_binary)

os.chdir(spark_home)

checkpoint_dir = Path("/tmp/GFTestsCheckpointDir")

stop_connect_cmd = ["./sbin/stop-connect-server.sh"]
print("Stopping SparkConnect Server...")
spark_connect_stop = subprocess.run(
stop_connect_cmd,
stdout=subprocess.PIPE,
universal_newlines=True,
)

if spark_connect_stop.returncode == 0:
print("Done.")
else:
print("Something goes wrong...")
sys.exit(1)

shutil.rmtree(checkpoint_dir.absolute().__str__(), ignore_errors=True)
sys.exit(0)
3 changes: 3 additions & 0 deletions python/graphframes/connect/graphframes_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -637,6 +637,9 @@ def plan(self, session: SparkConnectClient) -> proto.Relation:
edge_filter=edge_filter,
max_path_length=max_path_length,
is_directed=is_directed,
checkpoint_interval=checkpoint_interval,
use_local_checkpoints=use_local_checkpoints,
storage_level=storage_level,
),
self._spark,
)
Expand Down
6 changes: 4 additions & 2 deletions python/graphframes/graphframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,17 +86,17 @@ def is_remote() -> bool:
return False


from graphframes.classic.graphframe import GraphFrame as GraphFrameClassic
from graphframes.internal.utils import (
_HASH2VEC_DECAY_FUNCTIONS,
_RandomWalksEmbeddingsParameters,
)
from graphframes.lib import Pregel

if TYPE_CHECKING:
from pyspark.sql import Column, DataFrame

from graphframes.classic.graphframe import GraphFrame as GraphFrameClassic
from graphframes.connect.graphframes_client import GraphFrameConnect
from graphframes.lib import Pregel

"""Constant for the vertices ID column name."""
ID = "id"
Expand Down Expand Up @@ -177,6 +177,8 @@ def __init__(self, v: DataFrame, e: DataFrame) -> None:

self._impl = GraphFrameConnect(v, e) # ty: ignore[invalid-argument-type]
else:
from graphframes.classic.graphframe import GraphFrame as GraphFrameClassic

self._impl = GraphFrameClassic(v, e) # ty: ignore[invalid-argument-type]

@property
Expand Down
Loading
Loading