graphframes/python/tests/conftest.py at v0.12.2 · graphframes/graphframes · GitHub

History

147 lines (116 loc) · 4.91 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

from __future__ import annotations

import os

import pathlib

import tempfile

import warnings

import pytest

from pyspark.sql import SparkSession

from pyspark.version import __version__

from graphframes import GraphFrame

if __version__[:3] >= "3.4":

from pyspark.sql.utils import is_remote

else:

def is_remote() -> bool:

return False

spark_major_version = __version__[:1]

scala_version = os.environ.get("SCALA_VERSION", "2.12" if __version__ < "4" else "2.13")

is_thin_client = os.environ.get("SPARK_CLIENT_MODE_ENABLED", None) is not None

def get_gf_jar_locations() -> tuple[str, str, str]:

"""

Returns a location of the GraphFrames JAR and GraphFrames Connect JAR.

In the case your version of PySpark is not compatible with the version of GraphFrames,

this function will raise an exception!

"""

if is_thin_client:

return "", "", ""

project_root = pathlib.Path(__file__).parent.parent.parent

graphx_dir = project_root / "graphx" / "target" / f"scala-{scala_version}"

core_dir = project_root / "core" / "target" / f"scala-{scala_version}"

connect_dir = project_root / "connect" / "target" / f"scala-{scala_version}"

graphx_jar: str | None = None

core_jar: str | None = None

connect_jar: str | None = None

for pp in graphx_dir.glob(f"graphframes-graphx-spark{spark_major_version}*.jar"):

assert isinstance(pp, pathlib.PosixPath) # type checking

graphx_jar = str(pp.absolute())

if graphx_jar is None:

raise ValueError(

f"Failed to find graphframes jar for Spark {spark_major_version} in {graphx_dir}"

)

for pp in core_dir.glob(f"graphframes-spark{spark_major_version}*.jar"):

assert isinstance(pp, pathlib.PosixPath) # type checking

core_jar = str(pp.absolute())

if core_jar is None:

raise ValueError(

f"Failed to find graphframes jar for Spark {spark_major_version} in {core_dir}"

)

for pp in connect_dir.glob(f"graphframes-connect-spark{spark_major_version}*.jar"):

assert isinstance(pp, pathlib.PosixPath) # type checking

connect_jar = str(pp.absolute())

if connect_jar is None:

raise ValueError(

f"Failed to find graphframes connect jar for Spark {spark_major_version} in {connect_dir}"

)

return core_jar, connect_jar, graphx_jar

@pytest.fixture(scope="module")

def spark():

warnings.filterwarnings("ignore", category=ResourceWarning)

warnings.filterwarnings("ignore", category=DeprecationWarning)

(core_jar, connect_jar, graphx_jar) = get_gf_jar_locations()

with tempfile.TemporaryDirectory() as tmp_dir:

if is_thin_client:

spark = (

SparkSession.Builder()

.appName("GraphFramesTest")

.config("spark.sql.shuffle.partitions", 4)

.remote("sc://localhost:15002")

.getOrCreate()

)

yield spark

else:

builder = (

SparkSession.Builder()

.appName("GraphFramesTest")

.config("spark.sql.shuffle.partitions", 4)

.config("spark.checkpoint.dir", tmp_dir)

.config("spark.jars", f"{core_jar},{connect_jar},{graphx_jar}")

.config("spark.driver.memory", "6g")

)

if spark_major_version == "3":

# Spark 3 does not include connect by default

builder = builder.config(

"spark.jars.packages",

f"org.apache.spark:spark-connect_{scala_version}:{__version__}",

)

if is_remote():

builder = builder.remote("local[4]").config(

"spark.connect.extensions.relation.classes",

"org.apache.spark.sql.graphframes.GraphFramesConnect",

)

else:

builder = builder.master("local[4]")

spark = builder.getOrCreate()

yield spark

spark.stop()

@pytest.fixture(scope="module")

def local_g(spark: SparkSession):

localVertices = [(1, "A"), (2, "B"), (3, "C")]

localEdges = [(1, 2, "love"), (2, 1, "hate"), (2, 3, "follow")]

v = spark.createDataFrame(localVertices, ["id", "name"])

e = spark.createDataFrame(localEdges, ["src", "dst", "action"])

yield GraphFrame(v, e)

@pytest.fixture(scope="module")

def examples(spark: SparkSession):

if is_remote():

# TODO: We should update tests to be able to run all of them on Spark Connect

# At the moment the problem is that examples API is py4j based.

yield None

else:

from graphframes.classic.graphframe import _java_api

japi = _java_api(spark._sc)

assert japi is not None

examples = japi.examples()

assert examples is not None

from py4j.java_gateway import JavaObject

assert isinstance(examples, JavaObject)

yield examples

graphframes/python/tests/conftest.py at v0.12.2 · graphframes/graphframes · GitHub

Uh oh!

Expand file tree

Breadcrumbs

conftest.py

Latest commit

History

Breadcrumbs

conftest.py

File metadata and controls