From 2dc0a88efeed00b660cfac935fc90266e1a02eb7 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Wed, 26 Aug 2026 03:04:42 +0000 Subject: [PATCH 1/8] [GRAPHFRAMES] Add GraphFrames core module --- assembly/pom.xml | 5 + graphframes/pom.xml | 77 +++++ .../apache/spark/graphframes/GraphFrame.scala | 280 ++++++++++++++++++ .../graphframes/InvalidGraphException.scala | 21 ++ .../spark/graphframes/GraphFrameSuite.scala | 115 +++++++ pom.xml | 1 + project/SparkBuild.scala | 8 +- python/mypy.ini | 3 + python/pyspark/graphframes/__init__.py | 20 ++ python/pyspark/graphframes/graphframe.py | 237 +++++++++++++++ python/pyspark/graphframes/tests/__init__.py | 16 + .../graphframes/tests/connect/__init__.py | 16 + .../tests/connect/test_parity_graphframe.py | 46 +++ .../graphframes/tests/test_graphframe.py | 124 ++++++++ 14 files changed, 966 insertions(+), 3 deletions(-) create mode 100644 graphframes/pom.xml create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/GraphFrame.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/InvalidGraphException.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameSuite.scala create mode 100644 python/pyspark/graphframes/__init__.py create mode 100644 python/pyspark/graphframes/graphframe.py create mode 100644 python/pyspark/graphframes/tests/__init__.py create mode 100644 python/pyspark/graphframes/tests/connect/__init__.py create mode 100644 python/pyspark/graphframes/tests/connect/test_parity_graphframe.py create mode 100644 python/pyspark/graphframes/tests/test_graphframe.py diff --git a/assembly/pom.xml b/assembly/pom.xml index b1d00776bdf68..90805c4432cab 100644 --- a/assembly/pom.xml +++ b/assembly/pom.xml @@ -64,6 +64,11 @@ spark-graphx_${scala.binary.version} ${project.version} + + org.apache.spark + spark-graphframes_${scala.binary.version} + ${project.version} + org.apache.spark spark-sql_${scala.binary.version} diff --git a/graphframes/pom.xml b/graphframes/pom.xml new file mode 100644 index 0000000000000..0a0bed887a114 --- /dev/null +++ b/graphframes/pom.xml @@ -0,0 +1,77 @@ + + + + + 4.0.0 + + org.apache.spark + spark-parent_2.13 + 5.0.0-SNAPSHOT + ../pom.xml + + + spark-graphframes_2.13 + + graphframes + + jar + Spark Project GraphFrames + https://spark.apache.org/ + + + + org.apache.spark + spark-sql_${scala.binary.version} + ${project.version} + provided + + + org.apache.spark + spark-core_${scala.binary.version} + ${project.version} + test-jar + test + + + org.apache.spark + spark-catalyst_${scala.binary.version} + ${project.version} + test-jar + test + + + org.apache.spark + spark-sql_${scala.binary.version} + ${project.version} + test-jar + test + + + org.apache.spark + spark-tags_${scala.binary.version} + + + + + target/scala-${scala.binary.version}/classes + target/scala-${scala.binary.version}/test-classes + + diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFrame.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFrame.scala new file mode 100644 index 0000000000000..bb6c5e684fca3 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFrame.scala @@ -0,0 +1,280 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes + +import org.apache.spark.internal.Logging +import org.apache.spark.sql.{Column, DataFrame} +import org.apache.spark.sql.functions.{array, col, count, countDistinct, explode, expr, struct} +import org.apache.spark.storage.StorageLevel + +/** + * A graph whose vertices and edges are represented by Spark DataFrames. + * + * Vertices must contain a unique `id` column. Edges must contain `src` and `dst` columns whose + * values identify the source and destination vertices. Additional columns are graph attributes. + * + * Use [[GraphFrame.apply]] to construct a graph. + */ +class GraphFrame private ( + @transient private val vertexDataFrame: DataFrame, + @transient private val edgeDataFrame: DataFrame) + extends Serializable { + + import GraphFrame._ + + /** Persist the vertex and edge DataFrames with their default storage level. */ + def cache(): this.type = { + vertices.cache() + edges.cache() + this + } + + /** Persist the vertex and edge DataFrames with their default storage level. */ + def persist(): this.type = { + vertices.persist() + edges.persist() + this + } + + /** Persist the vertex and edge DataFrames with `storageLevel`. */ + def persist(storageLevel: StorageLevel): this.type = { + vertices.persist(storageLevel) + edges.persist(storageLevel) + this + } + + /** Remove the vertex and edge DataFrames from the cache. */ + def unpersist(): this.type = unpersist(blocking = false) + + /** Remove the vertex and edge DataFrames from the cache. */ + def unpersist(blocking: Boolean): this.type = { + vertices.unpersist(blocking) + edges.unpersist(blocking) + this + } + + /** The graph's vertex DataFrame. */ + def vertices: DataFrame = { + requireDriverDataFrame(vertexDataFrame) + vertexDataFrame + } + + /** The graph's edge DataFrame. */ + def edges: DataFrame = { + requireDriverDataFrame(edgeDataFrame) + edgeDataFrame + } + + /** + * Return `(source vertex)-[edge]->(destination vertex)` triplets. + * + * The result has `src`, `edge`, and `dst` struct columns containing the complete corresponding + * input rows. + */ + @transient lazy val triplets: DataFrame = { + val sourceVertices = vertices.select( + col(quoted(ID)).alias("__graphframes_src_id"), + nested(vertices, SRC)) + val graphEdges = edges.select( + col(quoted(SRC)).alias("__graphframes_edge_src"), + col(quoted(DST)).alias("__graphframes_edge_dst"), + nested(edges, EDGE)) + val destinationVertices = vertices.select( + col(quoted(ID)).alias("__graphframes_dst_id"), + nested(vertices, DST)) + + sourceVertices + .join( + graphEdges, + col("__graphframes_src_id") === col("__graphframes_edge_src")) + .join( + destinationVertices, + col("__graphframes_dst_id") === col("__graphframes_edge_dst")) + .select(col(SRC), col(EDGE), col(DST)) + } + + /** Return the out-degree of every vertex having at least one outgoing edge. */ + @transient lazy val outDegrees: DataFrame = { + edges + .groupBy(col(quoted(SRC)).alias(ID)) + .agg(count("*").cast("int").alias(OUT_DEGREE)) + } + + /** Return the in-degree of every vertex having at least one incoming edge. */ + @transient lazy val inDegrees: DataFrame = { + edges + .groupBy(col(quoted(DST)).alias(ID)) + .agg(count("*").cast("int").alias(IN_DEGREE)) + } + + /** Return the total degree of every vertex incident to at least one edge. */ + @transient lazy val degrees: DataFrame = { + edges + .select(explode(array(col(quoted(SRC)), col(quoted(DST)))).alias(ID)) + .groupBy(ID) + .agg(count("*").cast("int").alias(DEGREE)) + } + + /** Return a graph with the direction of every edge reversed. */ + def reverse: GraphFrame = { + val attributes = edges.columns + .filterNot(name => name == SRC || name == DST) + .map(name => col(quoted(name))) + val reversedColumns = Seq( + col(quoted(DST)).alias(SRC), + col(quoted(SRC)).alias(DST)) ++ attributes + val reversedEdges = edges.select(reversedColumns: _*) + GraphFrame(vertices, reversedEdges) + } + + /** Return an undirected graph by adding a reversed copy of every edge. */ + def asUndirected: GraphFrame = GraphFrame(vertices, edges.unionByName(reverse.edges)) + + /** Filter vertices and remove edges incident to any removed vertex. */ + def filterVertices(condition: Column): GraphFrame = { + val filteredVertices = vertices.filter(condition) + val vertexIds = filteredVertices.select(col(quoted(ID))) + val filteredEdges = edges + .join(vertexIds, col(quoted(SRC)) === vertexIds(ID), "left_semi") + .join(vertexIds, col(quoted(DST)) === vertexIds(ID), "left_semi") + GraphFrame(filteredVertices, filteredEdges) + } + + /** Filter vertices using a SQL expression. */ + def filterVertices(condition: String): GraphFrame = filterVertices(expr(condition)) + + /** Filter edges while keeping all vertices. */ + def filterEdges(condition: Column): GraphFrame = GraphFrame(vertices, edges.filter(condition)) + + /** Filter edges using a SQL expression. */ + def filterEdges(condition: String): GraphFrame = filterEdges(expr(condition)) + + /** Return a graph without vertices that are not incident to an edge. */ + def dropIsolatedVertices(): GraphFrame = { + val incidentIds = edges.select(explode(array(col(quoted(SRC)), col(quoted(DST)))).alias(ID)) + GraphFrame(vertices.join(incidentIds, Seq(ID), "left_semi"), edges) + } + + /** + * Validate vertex uniqueness and ensure that every edge endpoint is present in `vertices`. + * + * This method runs Spark jobs and throws [[InvalidGraphException]] for an invalid graph. + */ + def validate(): Unit = { + val persistedVertices = vertices.persist(StorageLevel.MEMORY_AND_DISK) + try { + val vertexCount = persistedVertices.count() + val distinctVertexCount = persistedVertices.select(countDistinct(col(quoted(ID)))).head() + .getLong(0) + if (vertexCount != distinctVertexCount) { + throw new InvalidGraphException( + s"Graph contains ${vertexCount - distinctVertexCount} duplicate vertices") + } + + val endpoints = edges + .select(col(quoted(SRC)).alias(ID)) + .union(edges.select(col(quoted(DST)).alias(ID))) + .distinct() + val missingEndpointCount = endpoints.join(persistedVertices, Seq(ID), "left_anti").count() + if (missingEndpointCount > 0) { + throw new InvalidGraphException( + s"Graph contains $missingEndpointCount edge endpoints without matching vertices") + } + } finally { + persistedVertices.unpersist() + } + } + + override def toString: String = { + val vertexColumns = ID +: vertices.columns.filterNot(_ == ID).toSeq + val edgeColumns = SRC +: DST +: edges.columns.filterNot(c => c == SRC || c == DST).toSeq + val orderedVertices = vertices.select(vertexColumns.map(name => col(quoted(name))): _*) + val orderedEdges = edges.select(edgeColumns.map(name => col(quoted(name))): _*) + s"GraphFrame(v:$orderedVertices, e:$orderedEdges)" + } + + private def requireDriverDataFrame(dataFrame: DataFrame): Unit = { + if (dataFrame == null) { + throw new IllegalStateException("GraphFrame objects cannot be used inside Spark closures") + } + } +} + +object GraphFrame extends Logging { + + val ID: String = "id" + val SRC: String = "src" + val DST: String = "dst" + val EDGE: String = "edge" + val DEGREE: String = "degree" + val IN_DEGREE: String = "inDegree" + val OUT_DEGREE: String = "outDegree" + + /** Create a GraphFrame from vertex and edge DataFrames. */ + def apply(vertices: DataFrame, edges: DataFrame): GraphFrame = { + requireColumn(vertices, ID, "Vertex ID") + requireColumn(edges, SRC, "Source vertex ID") + requireColumn(edges, DST, "Destination vertex ID") + require( + vertices.sparkSession eq edges.sparkSession, + "Vertex and edge DataFrames must belong to the same SparkSession") + new GraphFrame(vertices, edges) + } + + /** + * Create a GraphFrame from an edge DataFrame, deriving and persisting its distinct vertices. + */ + def fromEdges(edges: DataFrame): GraphFrame = { + fromEdges(edges, StorageLevel.MEMORY_AND_DISK) + } + + /** + * Create a GraphFrame from an edge DataFrame, deriving its distinct vertices. + * + * The caller is responsible for unpersisting the returned graph's vertex DataFrame. + */ + def fromEdges(edges: DataFrame, storageLevel: StorageLevel): GraphFrame = { + requireColumn(edges, SRC, "Source vertex ID") + requireColumn(edges, DST, "Destination vertex ID") + logWarning( + s"GraphFrame.fromEdges persists derived vertices with storage level $storageLevel; " + + "call vertices.unpersist() when the graph is no longer needed") + val vertices = edges + .select(col(quoted(SRC)).alias(ID)) + .union(edges.select(col(quoted(DST)).alias(ID))) + .distinct() + .persist(storageLevel) + GraphFrame(vertices, edges) + } + + private def requireColumn(dataFrame: DataFrame, columnName: String, label: String): Unit = { + require( + dataFrame.columns.contains(columnName), + s"$label column '$columnName' is missing; available columns: " + + dataFrame.columns.mkString(", ")) + } + + private def nested(dataFrame: DataFrame, name: String): Column = { + val columns = dataFrame.columns.map(columnName => col(quoted(columnName))).toIndexedSeq + struct(columns: _*).alias(name) + } + + private def quoted(columnName: String): String = { + s"`${columnName.replace("`", "``")}`" + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/InvalidGraphException.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/InvalidGraphException.scala new file mode 100644 index 0000000000000..a0fe72d3ca1dd --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/InvalidGraphException.scala @@ -0,0 +1,21 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes + +/** Thrown when a GraphFrame contains duplicate vertices or unknown edge endpoints. */ +class InvalidGraphException(message: String) extends IllegalArgumentException(message) diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameSuite.scala new file mode 100644 index 0000000000000..fb3cb63a64d59 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameSuite.scala @@ -0,0 +1,115 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes + +import org.apache.spark.sql.{QueryTest, Row} +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.test.SharedSparkSession + +class GraphFrameSuite extends QueryTest with SharedSparkSession { + + import testImplicits._ + + private def graph: GraphFrame = { + val vertices = Seq((1L, "a"), (2L, "b"), (3L, "c"), (4L, "isolated")) + .toDF("id", "name") + val edges = Seq((1L, 2L, "friend"), (2L, 3L, "follow"), (2L, 1L, "friend")) + .toDF("src", "dst", "relationship") + GraphFrame(vertices, edges) + } + + test("construction validates required columns") { + val vertices = Seq((1L, "a")).toDF("id", "name") + val edges = Seq((1L, 1L)).toDF("src", "dst") + GraphFrame(vertices, edges) + + val missingId = intercept[IllegalArgumentException] { + GraphFrame(vertices.withColumnRenamed("id", "vertex"), edges) + } + assert(missingId.getMessage.contains("Vertex ID column 'id' is missing")) + + val missingDestination = intercept[IllegalArgumentException] { + GraphFrame(vertices, edges.withColumnRenamed("dst", "destination")) + } + assert(missingDestination.getMessage.contains("Destination vertex ID column 'dst' is missing")) + } + + test("degree DataFrames preserve GraphFrames schemas") { + checkAnswer(graph.outDegrees, Seq(Row(1L, 1), Row(2L, 2))) + checkAnswer(graph.inDegrees, Seq(Row(1L, 1), Row(2L, 1), Row(3L, 1))) + checkAnswer(graph.degrees, Seq(Row(1L, 2), Row(2L, 3), Row(3L, 1))) + } + + test("triplets contain complete source, edge, and destination rows") { + val result = graph.triplets.select( + col("src.id"), + col("src.name"), + col("edge.relationship"), + col("dst.id"), + col("dst.name")) + + checkAnswer( + result, + Seq( + Row(1L, "a", "friend", 2L, "b"), + Row(2L, "b", "follow", 3L, "c"), + Row(2L, "b", "friend", 1L, "a"))) + } + + test("relational graph transforms retain attributes") { + val filtered = graph.filterVertices(col("id") <= 2L) + checkAnswer(filtered.vertices, Seq(Row(1L, "a"), Row(2L, "b"))) + checkAnswer( + filtered.edges, + Seq(Row(1L, 2L, "friend"), Row(2L, 1L, "friend"))) + + checkAnswer( + graph.filterEdges(col("relationship") === "follow").edges, + Seq(Row(2L, 3L, "follow"))) + checkAnswer(graph.dropIsolatedVertices().vertices.select("id"), Seq(Row(1L), Row(2L), Row(3L))) + checkAnswer( + graph.reverse.edges, + Seq( + Row(2L, 1L, "friend"), + Row(3L, 2L, "follow"), + Row(1L, 2L, "friend"))) + } + + test("validate rejects duplicate vertices and unknown edge endpoints") { + graph.validate() + + val duplicateVertices = Seq((1L, "a"), (1L, "duplicate")).toDF("id", "name") + val noEdges = Seq.empty[(Long, Long)].toDF("src", "dst") + assertThrows[InvalidGraphException](GraphFrame(duplicateVertices, noEdges).validate()) + + val vertices = Seq(1L).toDF("id") + val unknownEndpoint = Seq((1L, 2L)).toDF("src", "dst") + assertThrows[InvalidGraphException](GraphFrame(vertices, unknownEndpoint).validate()) + } + + test("fromEdges derives distinct vertices") { + val edges = Seq((1L, 2L), (2L, 3L), (1L, 2L)).toDF("src", "dst") + val derived = GraphFrame.fromEdges(edges) + try { + checkAnswer(derived.vertices, Seq(Row(1L), Row(2L), Row(3L))) + checkAnswer(derived.edges, edges.collect().toSeq) + } finally { + derived.vertices.unpersist() + } + } +} diff --git a/pom.xml b/pom.xml index e40270844945e..5130e89d34df7 100644 --- a/pom.xml +++ b/pom.xml @@ -89,6 +89,7 @@ sql/connect/shims core graphx + graphframes mllib mllib-local tools diff --git a/project/SparkBuild.scala b/project/SparkBuild.scala index 6fffcaf47cb5b..24c69a46244d4 100644 --- a/project/SparkBuild.scala +++ b/project/SparkBuild.scala @@ -63,10 +63,12 @@ object BuildCommons { Seq("udf-worker-proto", "udf-worker-core", "udf-worker-grpc").map(ProjectRef(buildLocation, _)) val allProjects@Seq( - core, graphx, mllib, mllibLocal, repl, networkCommon, networkShuffle, launcher, unsafe, tags, sketch, kvstore, + core, graphx, graphframes, mllib, mllibLocal, repl, networkCommon, networkShuffle, + launcher, unsafe, tags, sketch, kvstore, commonUtils, commonUtilsJava, variant, pipelines, sparkConfig, _* ) = Seq( - "core", "graphx", "mllib", "mllib-local", "repl", "network-common", "network-shuffle", "launcher", "unsafe", + "core", "graphx", "graphframes", "mllib", "mllib-local", "repl", "network-common", + "network-shuffle", "launcher", "unsafe", "tags", "sketch", "kvstore", "common-utils", "common-utils-java", "variant", "pipelines", "config" ).map(ProjectRef(buildLocation, _)) ++ sqlProjects ++ streamingProjects ++ connectProjects ++ udfWorkerProjects @@ -414,7 +416,7 @@ object SparkBuild extends PomBuild { val mimaProjects = allProjects.filterNot { x => Seq( - spark, hive, hiveThriftServer, repl, networkCommon, networkShuffle, networkYarn, + spark, graphframes, hive, hiveThriftServer, repl, networkCommon, networkShuffle, networkYarn, unsafe, tags, tokenProviderKafka010, sqlKafka010, pipelines, connectCommon, connect, connectJdbc, connectClient, variant, connectShims, profiler, credentialAws, commonUtilsJava, sparkConfig, diff --git a/python/mypy.ini b/python/mypy.ini index 5baa77c370c51..f3e3e8d3b2cae 100644 --- a/python/mypy.ini +++ b/python/mypy.ini @@ -106,6 +106,9 @@ ignore_errors = True [mypy-pyspark.errors.tests.*] ignore_errors = True +[mypy-pyspark.graphframes.tests.*] +ignore_errors = True + [mypy-pyspark.logger.tests.*] ignore_errors = True diff --git a/python/pyspark/graphframes/__init__.py b/python/pyspark/graphframes/__init__.py new file mode 100644 index 0000000000000..32ed934e15314 --- /dev/null +++ b/python/pyspark/graphframes/__init__.py @@ -0,0 +1,20 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from pyspark.graphframes.graphframe import GraphFrame + +__all__ = ["GraphFrame"] diff --git a/python/pyspark/graphframes/graphframe.py b/python/pyspark/graphframes/graphframe.py new file mode 100644 index 0000000000000..edbe2de924910 --- /dev/null +++ b/python/pyspark/graphframes/graphframe.py @@ -0,0 +1,237 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from __future__ import annotations + +from typing import Union + +from pyspark.sql import Column, DataFrame +from pyspark.sql import functions as F +from pyspark.storagelevel import StorageLevel + + +class GraphFrame: + """A graph whose vertices and edges are represented by Spark DataFrames. + + The vertex DataFrame must contain a unique ``id`` column. The edge DataFrame must contain + ``src`` and ``dst`` columns identifying its source and destination vertices. All additional + columns are retained as graph attributes. + + The initial in-tree API consists entirely of DataFrame operations and therefore supports both + classic Spark and Spark Connect. + """ + + ID = "id" + SRC = "src" + DST = "dst" + EDGE = "edge" + + def __init__(self, vertices: DataFrame, edges: DataFrame) -> None: + self._require_column(vertices, self.ID, "Vertex ID") + self._require_column(edges, self.SRC, "Source vertex ID") + self._require_column(edges, self.DST, "Destination vertex ID") + self._vertices = vertices + self._edges = edges + + @property + def vertices(self) -> DataFrame: + """The graph's vertex DataFrame.""" + return self._vertices + + @property + def nodes(self) -> DataFrame: + """An alias for :attr:`vertices`.""" + return self.vertices + + @property + def edges(self) -> DataFrame: + """The graph's edge DataFrame.""" + return self._edges + + @property + def triplets(self) -> DataFrame: + """Return ``(source vertex)-[edge]->(destination vertex)`` triplets.""" + source_vertices = self.vertices.select( + self.vertices[self.ID].alias("__graphframes_src_id"), + self._nested(self.vertices, self.SRC), + ) + graph_edges = self.edges.select( + self.edges[self.SRC].alias("__graphframes_edge_src"), + self.edges[self.DST].alias("__graphframes_edge_dst"), + self._nested(self.edges, self.EDGE), + ) + destination_vertices = self.vertices.select( + self.vertices[self.ID].alias("__graphframes_dst_id"), + self._nested(self.vertices, self.DST), + ) + return ( + source_vertices.join( + graph_edges, + F.col("__graphframes_src_id") == F.col("__graphframes_edge_src"), + ) + .join( + destination_vertices, + F.col("__graphframes_dst_id") == F.col("__graphframes_edge_dst"), + ) + .select(self.SRC, self.EDGE, self.DST) + ) + + @property + def outDegrees(self) -> DataFrame: + """Return the out-degree of vertices having at least one outgoing edge.""" + return self.edges.groupBy(self.edges[self.SRC].alias(self.ID)).agg( + F.count("*").cast("int").alias("outDegree") + ) + + @property + def inDegrees(self) -> DataFrame: + """Return the in-degree of vertices having at least one incoming edge.""" + return self.edges.groupBy(self.edges[self.DST].alias(self.ID)).agg( + F.count("*").cast("int").alias("inDegree") + ) + + @property + def degrees(self) -> DataFrame: + """Return the total degree of vertices incident to at least one edge.""" + return ( + self.edges.select( + F.explode(F.array(self.edges[self.SRC], self.edges[self.DST])).alias(self.ID) + ) + .groupBy(self.ID) + .agg(F.count("*").cast("int").alias("degree")) + ) + + def cache(self) -> "GraphFrame": + """Persist the vertex and edge DataFrames with their default storage level.""" + self.vertices.cache() + self.edges.cache() + return self + + def persist( + self, storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER + ) -> "GraphFrame": + """Persist the vertex and edge DataFrames with ``storage_level``.""" + self.vertices.persist(storage_level) + self.edges.persist(storage_level) + return self + + def unpersist(self, blocking: bool = False) -> "GraphFrame": + """Remove the vertex and edge DataFrames from the cache.""" + self.vertices.unpersist(blocking) + self.edges.unpersist(blocking) + return self + + def filterVertices(self, condition: Union[Column, str]) -> "GraphFrame": + """Filter vertices and remove edges incident to any removed vertex.""" + filtered_vertices = self.vertices.filter(condition) + vertex_ids = filtered_vertices.select(filtered_vertices[self.ID]) + filtered_edges = self.edges.join( + vertex_ids, + self.edges[self.SRC] == vertex_ids[self.ID], + "left_semi", + ).join( + vertex_ids, + self.edges[self.DST] == vertex_ids[self.ID], + "left_semi", + ) + return GraphFrame(filtered_vertices, filtered_edges) + + def filterEdges(self, condition: Union[Column, str]) -> "GraphFrame": + """Filter edges while keeping all vertices.""" + return GraphFrame(self.vertices, self.edges.filter(condition)) + + def dropIsolatedVertices(self) -> "GraphFrame": + """Return a graph without vertices that are not incident to an edge.""" + incident_ids = self.edges.select( + F.explode(F.array(self.edges[self.SRC], self.edges[self.DST])).alias(self.ID) + ) + return GraphFrame(self.vertices.join(incident_ids, self.ID, "left_semi"), self.edges) + + def as_reversed(self) -> "GraphFrame": + """Return a graph with the direction of every edge reversed.""" + attributes = [ + self.edges[name] for name in self.edges.columns if name not in {self.SRC, self.DST} + ] + reversed_edges = self.edges.select( + self.edges[self.DST].alias(self.SRC), + self.edges[self.SRC].alias(self.DST), + *attributes, + ) + return GraphFrame(self.vertices, reversed_edges) + + def as_undirected(self) -> "GraphFrame": + """Return an undirected graph by adding a reversed copy of every edge.""" + return GraphFrame(self.vertices, self.edges.unionByName(self.as_reversed().edges)) + + def validate(self) -> None: + """Run jobs that validate vertex uniqueness and edge endpoint integrity.""" + vertex_count = self.vertices.count() + distinct_vertex_count = self.vertices.select(self.ID).distinct().count() + if vertex_count != distinct_vertex_count: + raise ValueError( + f"Graph contains {vertex_count - distinct_vertex_count} duplicate vertices" + ) + + endpoints = ( + self.edges.select(self.edges[self.SRC].alias(self.ID)) + .union(self.edges.select(self.edges[self.DST].alias(self.ID))) + .distinct() + ) + missing_endpoint_count = endpoints.join(self.vertices, self.ID, "left_anti").count() + if missing_endpoint_count: + raise ValueError( + f"Graph contains {missing_endpoint_count} edge endpoints without matching vertices" + ) + + @classmethod + def from_edges( + cls, + edges: DataFrame, + storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> "GraphFrame": + """Create a graph by deriving and persisting distinct vertices from ``edges``.""" + cls._require_column(edges, cls.SRC, "Source vertex ID") + cls._require_column(edges, cls.DST, "Destination vertex ID") + vertices = ( + edges.select(edges[cls.SRC].alias(cls.ID)) + .union(edges.select(edges[cls.DST].alias(cls.ID))) + .distinct() + .persist(storage_level) + ) + return cls(vertices, edges) + + def __repr__(self) -> str: + vertex_columns = [self.ID] + [name for name in self.vertices.columns if name != self.ID] + edge_columns = [self.SRC, self.DST] + [ + name for name in self.edges.columns if name not in {self.SRC, self.DST} + ] + return ( + f"GraphFrame(v:{self.vertices.select(*vertex_columns)!r}, " + f"e:{self.edges.select(*edge_columns)!r})" + ) + + @staticmethod + def _nested(dataframe: DataFrame, name: str) -> Column: + return F.struct(*[dataframe[column] for column in dataframe.columns]).alias(name) + + @staticmethod + def _require_column(dataframe: DataFrame, column: str, label: str) -> None: + if column not in dataframe.columns: + available = ", ".join(dataframe.columns) + raise ValueError( + f"{label} column '{column}' is missing; available columns: {available}" + ) diff --git a/python/pyspark/graphframes/tests/__init__.py b/python/pyspark/graphframes/tests/__init__.py new file mode 100644 index 0000000000000..cce3acad34a49 --- /dev/null +++ b/python/pyspark/graphframes/tests/__init__.py @@ -0,0 +1,16 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# diff --git a/python/pyspark/graphframes/tests/connect/__init__.py b/python/pyspark/graphframes/tests/connect/__init__.py new file mode 100644 index 0000000000000..cce3acad34a49 --- /dev/null +++ b/python/pyspark/graphframes/tests/connect/__init__.py @@ -0,0 +1,16 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# diff --git a/python/pyspark/graphframes/tests/connect/test_parity_graphframe.py b/python/pyspark/graphframes/tests/connect/test_parity_graphframe.py new file mode 100644 index 0000000000000..ca485b569ff5a --- /dev/null +++ b/python/pyspark/graphframes/tests/connect/test_parity_graphframe.py @@ -0,0 +1,46 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from pyspark.graphframes import GraphFrame +from pyspark.sql import Row +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class GraphFrameParityTests(ReusedConnectTestCase): + def test_dataframe_operations(self) -> None: + vertices = self.spark.createDataFrame( + [(1, "a"), (2, "b"), (3, "c")], + ["id", "name"], + ) + edges = self.spark.createDataFrame( + [(1, 2, "friend"), (2, 3, "follow"), (2, 1, "friend")], + ["src", "dst", "relationship"], + ) + graph = GraphFrame(vertices, edges) + + self.assertEqual( + sorted(graph.outDegrees.collect()), + [Row(id=1, outDegree=1), Row(id=2, outDegree=2)], + ) + self.assertEqual(graph.triplets.count(), 3) + self.assertEqual(graph.filterVertices("id <= 2").edges.count(), 2) + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/graphframes/tests/test_graphframe.py b/python/pyspark/graphframes/tests/test_graphframe.py new file mode 100644 index 0000000000000..339f109d70e76 --- /dev/null +++ b/python/pyspark/graphframes/tests/test_graphframe.py @@ -0,0 +1,124 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from pyspark.graphframes import GraphFrame +from pyspark.sql import Row +from pyspark.sql import functions as F +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +class GraphFrameTests(ReusedSQLTestCase): + def setUp(self) -> None: + super().setUp() + vertices = self.spark.createDataFrame( + [(1, "a"), (2, "b"), (3, "c"), (4, "isolated")], + ["id", "name"], + ) + edges = self.spark.createDataFrame( + [(1, 2, "friend"), (2, 3, "follow"), (2, 1, "friend")], + ["src", "dst", "relationship"], + ) + self.graph = GraphFrame(vertices, edges) + + def test_degrees(self) -> None: + self.assertEqual( + sorted(self.graph.outDegrees.collect()), + [Row(id=1, outDegree=1), Row(id=2, outDegree=2)], + ) + self.assertEqual( + sorted(self.graph.inDegrees.collect()), + [Row(id=1, inDegree=1), Row(id=2, inDegree=1), Row(id=3, inDegree=1)], + ) + self.assertEqual( + sorted(self.graph.degrees.collect()), + [Row(id=1, degree=2), Row(id=2, degree=3), Row(id=3, degree=1)], + ) + + def test_triplets(self) -> None: + rows = self.graph.triplets.select( + F.col("src.id").alias("src_id"), + F.col("src.name").alias("src_name"), + F.col("edge.relationship").alias("relationship"), + F.col("dst.id").alias("dst_id"), + F.col("dst.name").alias("dst_name"), + ).collect() + self.assertEqual( + sorted(rows), + [ + Row( + src_id=1, + src_name="a", + relationship="friend", + dst_id=2, + dst_name="b", + ), + Row( + src_id=2, + src_name="b", + relationship="follow", + dst_id=3, + dst_name="c", + ), + Row( + src_id=2, + src_name="b", + relationship="friend", + dst_id=1, + dst_name="a", + ), + ], + ) + + def test_relational_transforms(self) -> None: + filtered = self.graph.filterVertices(F.col("id") <= 2) + self.assertEqual( + sorted(filtered.vertices.collect()), + [Row(id=1, name="a"), Row(id=2, name="b")], + ) + self.assertEqual( + sorted(filtered.edges.collect()), + [ + Row(src=1, dst=2, relationship="friend"), + Row(src=2, dst=1, relationship="friend"), + ], + ) + self.assertEqual( + sorted(self.graph.dropIsolatedVertices().vertices.select("id").collect()), + [Row(id=1), Row(id=2), Row(id=3)], + ) + + def test_validate(self) -> None: + self.graph.validate() + + vertices = self.spark.createDataFrame([(1,), (1,)], ["id"]) + edges = self.spark.createDataFrame([], "src long, dst long") + with self.assertRaisesRegex(ValueError, "duplicate vertices"): + GraphFrame(vertices, edges).validate() + + +if __name__ == "__main__": + from pyspark.graphframes.tests.test_graphframe import * # noqa: F403 + + try: + import xmlrunner + + testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) + except ImportError: + testRunner = None + import unittest + + unittest.main(testRunner=testRunner, verbosity=2) From 553929d9a3c3b418f801032bc08d1ea9809220a8 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Wed, 26 Aug 2026 04:40:26 +0000 Subject: [PATCH 2/8] [GRAPHFRAMES] Add algorithms and built-in Connect support --- graphframes/pom.xml | 12 + .../apache/spark/graphframes/GraphFrame.scala | 1698 +++++++++++-- ...eption.scala => GraphFramePythonAPI.scala} | 17 +- .../apache/spark/graphframes/Logging.scala | 47 + .../convolutions/SamplingConvolution.scala | 194 ++ .../graphframes/embeddings/Hash2Vec.scala | 601 +++++ .../embeddings/RandomWalkEmbeddings.scala | 384 +++ .../apache/spark/graphframes/exceptions.scala | 76 + .../graphframes/lib/AggregateMessages.scala | 207 ++ .../graphframes/lib/AggregateNeighbors.scala | 459 ++++ .../spark/graphframes/lib/AllPaths.scala | 209 ++ .../apache/spark/graphframes/lib/BFS.scala | 231 ++ .../graphframes/lib/ConnectedComponents.scala | 222 ++ .../graphframes/lib/DetectingCycles.scala | 122 + .../graphframes/lib/GraphXConversions.scala | 206 ++ .../spark/graphframes/lib/HyperANF.scala | 237 ++ .../apache/spark/graphframes/lib/KCore.scala | 126 + .../graphframes/lib/LabelPropagation.scala | 146 ++ .../lib/MaximalIndependentSet.scala | 242 ++ .../spark/graphframes/lib/PageRank.scala | 181 ++ .../lib/ParallelPersonalizedPageRank.scala | 129 + .../apache/spark/graphframes/lib/Pregel.scala | 658 ++++++ .../lib/RandomizedContraction.scala | 291 +++ .../spark/graphframes/lib/SVDPlusPlus.scala | 257 ++ .../spark/graphframes/lib/ShortestPaths.scala | 265 +++ .../lib/StronglyConnectedComponents.scala | 58 + .../lib/StructureAwareLabelPropagation.scala | 291 +++ .../spark/graphframes/lib/TriangleCount.scala | 194 ++ .../spark/graphframes/lib/TwoPhase.scala | 622 +++++ .../org/apache/spark/graphframes/mixins.scala | 248 ++ .../spark/graphframes/pattern/patterns.scala | 300 +++ .../spark/graphframes/rw/RandomWalkBase.scala | 434 ++++ .../rw/RandomWalkWithRestart.scala | 103 + .../sql/graphframes/GraphFrameInternals.scala | 137 ++ .../sql/graphframes/GraphFramesConf.scala | 144 ++ .../expressions/FiniteAXPlusB.scala | 102 + .../graphframes/expressions/KCoreMerge.scala | 118 + .../expressions/KMinSampling.scala | 182 ++ .../GraphFrameInternalsSuite.scala | 214 ++ .../spark/graphframes/GraphFrameSuite.scala | 805 ++++++- .../GraphFrameTestSparkContext.scala | 86 + .../spark/graphframes/PatternMatchSuite.scala | 875 +++++++ .../spark/graphframes/SparkFunSuite.scala | 45 + .../apache/spark/graphframes/TestUtils.scala | 128 + .../SamplingConvolutionSuite.scala | 135 ++ .../embeddings/Hash2VecSuite.scala | 388 +++ .../spark/graphframes/examples/Graphs.scala | 248 ++ .../lib/AggregateMessagesSuite.scala | 180 ++ .../lib/AggregateNeighborsSuite.scala | 463 ++++ .../spark/graphframes/lib/AllPathsSuite.scala | 151 ++ .../spark/graphframes/lib/BFSSuite.scala | 182 ++ .../lib/ConnectedComponentsSuite.scala | 419 ++++ .../lib/DetectingCyclesSuite.scala | 80 + .../spark/graphframes/lib/HyperANFSuite.scala | 152 ++ .../spark/graphframes/lib/KCoreSuite.scala | 342 +++ .../lib/LabelPropagationSuite.scala | 44 + .../lib/MaximalIndependentSetSuite.scala | 138 ++ .../spark/graphframes/lib/PageRankSuite.scala | 82 + .../ParallelPersonalizedPageRankSuite.scala | 102 + .../spark/graphframes/lib/PregelSuite.scala | 590 +++++ .../lib/RandomizedContractionSuite.scala | 294 +++ .../graphframes/lib/SVDPlusPlusSuite.scala | 104 + .../graphframes/lib/ShortestPathsSuite.scala | 180 ++ .../StronglyConnectedComponentsSuite.scala | 42 + .../lib/StructureAwareLabelPropagation.scala | 360 +++ .../graphframes/lib/TriangleCountSuite.scala | 218 ++ .../graphframes/pattern/PatternSuite.scala | 282 +++ .../rw/RandomWalkWithRestartSuite.scala | 178 ++ .../pyspark/graphframes/classic/__init__.py | 16 + .../pyspark/graphframes/classic/graphframe.py | 603 +++++ python/pyspark/graphframes/classic/utils.py | 31 + .../pyspark/graphframes/connect/__init__.py | 16 + .../graphframes/connect/graphframes_client.py | 1642 +++++++++++++ python/pyspark/graphframes/connect/utils.py | 78 + python/pyspark/graphframes/graphframe.py | 1557 ++++++++++-- .../pyspark/graphframes/internal/__init__.py | 16 + python/pyspark/graphframes/internal/utils.py | 74 + python/pyspark/graphframes/lib/__init__.py | 20 + .../graphframes/lib/aggregate_messages.py | 60 + python/pyspark/graphframes/lib/pregel.py | 342 +++ .../tests/connect/test_all_algorithms.py | 221 ++ .../sql/connect/proto/graphframes_pb2.py | 110 + .../sql/connect/proto/graphframes_pb2.pyi | 2090 +++++++++++++++++ .../sql/connect/proto/relations_pb2.py | 359 +-- .../sql/connect/proto/relations_pb2.pyi | 10 + .../protobuf/spark/connect/graphframes.proto | 331 +++ .../protobuf/spark/connect/relations.proto | 2 + sql/connect/server/pom.xml | 6 + .../connect/planner/SparkConnectPlanner.scala | 6 + .../graphframes/GraphFramesConnectUtils.scala | 639 +++++ 90 files changed, 25316 insertions(+), 570 deletions(-) rename graphframes/src/main/scala/org/apache/spark/graphframes/{InvalidGraphException.scala => GraphFramePythonAPI.scala} (63%) create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/Logging.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/convolutions/SamplingConvolution.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/Hash2Vec.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/RandomWalkEmbeddings.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/exceptions.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateMessages.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateNeighbors.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/AllPaths.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/BFS.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/ConnectedComponents.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/DetectingCycles.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/GraphXConversions.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/HyperANF.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/KCore.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/LabelPropagation.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/MaximalIndependentSet.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/PageRank.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRank.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/Pregel.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/RandomizedContraction.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/SVDPlusPlus.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/ShortestPaths.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponents.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/TriangleCount.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/lib/TwoPhase.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/mixins.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/pattern/patterns.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkBase.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestart.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFrameInternals.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/FiniteAXPlusB.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KCoreMerge.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KMinSampling.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameInternalsSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameTestSparkContext.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/PatternMatchSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/SparkFunSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/TestUtils.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/convolutions/SamplingConvolutionSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/embeddings/Hash2VecSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/examples/Graphs.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateMessagesSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateNeighborsSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/AllPathsSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/BFSSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/ConnectedComponentsSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/DetectingCyclesSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/HyperANFSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/KCoreSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/LabelPropagationSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/MaximalIndependentSetSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/PageRankSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRankSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/PregelSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/RandomizedContractionSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/SVDPlusPlusSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/ShortestPathsSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponentsSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/lib/TriangleCountSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/pattern/PatternSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestartSuite.scala create mode 100644 python/pyspark/graphframes/classic/__init__.py create mode 100644 python/pyspark/graphframes/classic/graphframe.py create mode 100644 python/pyspark/graphframes/classic/utils.py create mode 100644 python/pyspark/graphframes/connect/__init__.py create mode 100644 python/pyspark/graphframes/connect/graphframes_client.py create mode 100644 python/pyspark/graphframes/connect/utils.py create mode 100644 python/pyspark/graphframes/internal/__init__.py create mode 100644 python/pyspark/graphframes/internal/utils.py create mode 100644 python/pyspark/graphframes/lib/__init__.py create mode 100644 python/pyspark/graphframes/lib/aggregate_messages.py create mode 100644 python/pyspark/graphframes/lib/pregel.py create mode 100644 python/pyspark/graphframes/tests/connect/test_all_algorithms.py create mode 100644 python/pyspark/sql/connect/proto/graphframes_pb2.py create mode 100644 python/pyspark/sql/connect/proto/graphframes_pb2.pyi create mode 100644 sql/connect/common/src/main/protobuf/spark/connect/graphframes.proto create mode 100644 sql/connect/server/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala diff --git a/graphframes/pom.xml b/graphframes/pom.xml index 0a0bed887a114..762377a681756 100644 --- a/graphframes/pom.xml +++ b/graphframes/pom.xml @@ -37,6 +37,18 @@ https://spark.apache.org/ + + org.apache.spark + spark-graphx_${scala.binary.version} + ${project.version} + provided + + + org.apache.spark + spark-mllib_${scala.binary.version} + ${project.version} + provided + org.apache.spark spark-sql_${scala.binary.version} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFrame.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFrame.scala index bb6c5e684fca3..cc7f6ac32865c 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFrame.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFrame.scala @@ -17,264 +17,1650 @@ package org.apache.spark.graphframes -import org.apache.spark.internal.Logging -import org.apache.spark.sql.{Column, DataFrame} -import org.apache.spark.sql.functions.{array, col, count, countDistinct, explode, expr, struct} +import org.apache.spark.graphx.Edge +import org.apache.spark.graphx.Graph +import org.apache.spark.ml.clustering.PowerIterationClustering +import org.apache.spark.sql._ +import org.apache.spark.sql.functions.array +import org.apache.spark.sql.functions.broadcast +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.count +import org.apache.spark.sql.functions.countDistinct +import org.apache.spark.sql.functions.explode +import org.apache.spark.sql.functions.expr +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.functions.monotonically_increasing_id +import org.apache.spark.sql.functions.struct +import org.apache.spark.sql.types._ import org.apache.spark.storage.StorageLevel +import org.apache.spark.graphframes.embeddings.RandomWalkEmbeddings +import org.apache.spark.graphframes.lib._ +import org.apache.spark.graphframes.pattern._ + +import java.util.Random +import scala.reflect.runtime.universe.TypeTag /** - * A graph whose vertices and edges are represented by Spark DataFrames. - * - * Vertices must contain a unique `id` column. Edges must contain `src` and `dst` columns whose - * values identify the source and destination vertices. Additional columns are graph attributes. + * A representation of a graph using `DataFrame`s. * - * Use [[GraphFrame.apply]] to construct a graph. + * @groupname structure Structure information + * @groupname conversions Conversions + * @groupname stdlib Standard graph algorithms + * @groupname subgraph Subgraph selection + * @groupname degree Graph topology + * @groupname motif Motif finding + * @groupname gml Graph Machine Learning + * @groupname utils Utility Methods */ class GraphFrame private ( - @transient private val vertexDataFrame: DataFrame, - @transient private val edgeDataFrame: DataFrame) - extends Serializable { + @transient private val _vertices: DataFrame, + @transient private val _edges: DataFrame) + extends Logging + with Serializable { import GraphFrame._ - /** Persist the vertex and edge DataFrames with their default storage level. */ + /** Default constructor is provided to support serialization */ + protected def this() = this(null, null) + + /** + * Return a string representation of the GraphFrame + * + * @group utils + */ + override def toString: String = { + // We call select on the vertices and edges to ensure that ID, SRC, DST always come first + // in the printed schema. + val vCols = (ID +: vertices.columns.filter(_ != ID).toIndexedSeq).map(quote).map(col) + val eCols = + (SRC +: DST +: edges.columns.filter(c => c != SRC && c != DST).toIndexedSeq) + .map(quote) + .map(col) + val v = vertices.select(vCols.toSeq: _*).toString + val e = edges.select(eCols.toSeq: _*).toString + "GraphFrame(v:" + v + ", e:" + e + ")" + } + + /** + * Persist the dataframe representation of vertices and edges of the graph with the default + * + * @group utils + * storage level. + */ def cache(): this.type = { - vertices.cache() - edges.cache() - this + persist() } - /** Persist the vertex and edge DataFrames with their default storage level. */ + /** + * Persist the dataframe representation of vertices and edges of the graph with the default + * + * @group utils + * storage level. + */ def persist(): this.type = { vertices.persist() edges.persist() this } - /** Persist the vertex and edge DataFrames with `storageLevel`. */ - def persist(storageLevel: StorageLevel): this.type = { - vertices.persist(storageLevel) - edges.persist(storageLevel) + /** + * Persist the dataframe representation of vertices and edges of the graph with the given + * storage level. + * @param newLevel + * One of: `MEMORY_ONLY`, `MEMORY_AND_DISK`, `MEMORY_ONLY_SER`, `MEMORY_AND_DISK_SER`, + * + * @group utils + * `DISK_ONLY`, `MEMORY_ONLY_2`, `MEMORY_AND_DISK_2`, etc.. + */ + def persist(newLevel: StorageLevel): this.type = { + vertices.persist(newLevel) + edges.persist(newLevel) this } - /** Remove the vertex and edge DataFrames from the cache. */ - def unpersist(): this.type = unpersist(blocking = false) + /** + * Mark the dataframe representation of vertices and edges of the graph as non-persistent, and + * remove all blocks for it from memory and disk. + * + * @group utils + */ + def unpersist(): this.type = { + vertices.unpersist() + edges.unpersist() + this + } - /** Remove the vertex and edge DataFrames from the cache. */ + /** + * Mark the dataframe representation of vertices and edges of the graph as non-persistent, and + * remove all blocks for it from memory and disk. + * @param blocking + * Whether to block until all blocks are deleted. + * + * @group utils + */ def unpersist(blocking: Boolean): this.type = { vertices.unpersist(blocking) edges.unpersist(blocking) this } - /** The graph's vertex DataFrame. */ + /** + * Validates the consistency and integrity of a graph by performing checks on the vertices and + * edges. + * + * @return + * Unit, as the method, performs validation checks and throws an exception if validation + * fails. + * @throws InvalidGraphException + * if there are any inconsistencies in the graph, such as duplicate vertices, mismatched + * vertices between edges and vertex DataFrames or missing connections. + * + * @group utils + */ + def validate(): Unit = + validate(checkVertices = true, intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK) + + /** + * Validates the consistency and integrity of a graph by performing checks on the vertices and + * edges. + * + * @param checkVertices + * a flag to indicate whether additional vertex consistency checks should be performed. If + * true, the method will verify that all vertices in the vertex DataFrame are represented in + * the edge DataFrame and vice versa. It is slow on big graphs. + * @param intermediateStorageLevel + * the storage level to be used when persisting intermediate DataFrame computations during the + * validation process. + * @return + * Unit, as the method, performs validation checks and throws an exception if validation + * fails. + * @throws InvalidGraphException + * if there are any inconsistencies in the graph, such as duplicate vertices, mismatched + * vertices between edges and vertex DataFrames or missing connections. + * + * @group utils + */ + def validate(checkVertices: Boolean, intermediateStorageLevel: StorageLevel): Unit = { + val persistedVertices = vertices.persist(intermediateStorageLevel) + val countDistinctVertices = persistedVertices.select(countDistinct(ID)).first().getLong(0) + val verticesCount = persistedVertices.count() + if (countDistinctVertices != verticesCount) { + throw new InvalidGraphException( + s"Graph contains (${verticesCount - countDistinctVertices}) duplicate vertices.") + } + if (checkVertices) { + val verticesSetFromEdges = edges + .select(col(SRC).alias(ID)) + .union(edges.select(col(DST).alias(ID))) + .distinct() + .persist(intermediateStorageLevel) + val countVerticesFromEdges = verticesSetFromEdges.count() + if (countVerticesFromEdges > countDistinctVertices) { + throw new InvalidGraphException( + s"Graph is inconsistent: edges has ${countVerticesFromEdges} " + + s"vertices, but vertices has ${countDistinctVertices} vertices.") + } + + val combined = verticesSetFromEdges.join(vertices, ID, "left_anti") + val countOfBadVertices = combined.count() + if (countOfBadVertices > 0) { + throw new InvalidGraphException( + "Vertices DataFrame does not contain all edges src/dst. " + + s"Found ${countOfBadVertices} edges src/dst that are not in the vertices DataFrame.") + } + persistedVertices.unpersist() + verticesSetFromEdges.unpersist() + () + } + } + + /** + * Converts the directed graph into an undirected graph by ensuring that all directed edges are + * bidirectional. For every directed edge (src, dst), a corresponding edge (dst, src) is added. + * + * @return + * a new GraphFrame representing the undirected graph. + * + * @group utils + */ + def asUndirected(): GraphFrame = { + val newEdges = edges + .select(col(SRC), col(DST), nestAsCol(edges, ATTR)) + .union(edges + .select(col(DST).alias(SRC), col(SRC).alias(DST), nestAsCol(edges, ATTR))) + .select(SRC, DST, ATTR) + val newColumns = Seq(col(SRC), col(DST)) ++ edges.columns + .filter(c => (c != SRC) && (c != DST)) + .map(c => col(ATTR).getField(c).alias(c)) + .toSeq + GraphFrame(vertices, newEdges.select(newColumns: _*)) + } + + /** + * Reverses the direction of all edges in the graph. For every directed edge (src, dst), the + * resulting graph will contain an edge (dst, src) with the same attributes. + * + * @return + * a new GraphFrame with all edge directions reversed. + * + * @group utils + */ + def asReversed(): GraphFrame = { + val newEdges = edges + .select(col(DST).alias(SRC), col(SRC).alias(DST), nestAsCol(edges, ATTR)) + .select(SRC, DST, ATTR) + val newColumns = Seq(col(SRC), col(DST)) ++ edges.columns + .filter(c => (c != SRC) && (c != DST)) + .map(c => col(ATTR).getField(c).alias(c)) + .toSeq + GraphFrame(vertices, newEdges.select(newColumns: _*)) + } + + // ============== Basic structural methods ============ + + /** + * The dataframe representation of the vertices of the graph. + * + * It contains a column called [[GraphFrame.ID]] with the id of the vertex, and various other + * user-defined attributes with other attributes. + * + * The order of the columns is available in [[vertexColumns]]. + * + * @group structure + */ def vertices: DataFrame = { - requireDriverDataFrame(vertexDataFrame) - vertexDataFrame + if (_vertices == null) { + throw new Exception("You cannot use GraphFrame objects within a Spark closure") + } + _vertices } - /** The graph's edge DataFrame. */ + /** + * The dataframe representation of the edges of the graph. + * + * It contains two columns called [[GraphFrame.SRC]] and [[GraphFrame.DST]] that contain the ids + * of the source vertex and the destination vertex of each edge, respectively. It may also + * contain various other columns with user-defined attributes for each edge. + * + * For symmetric graphs, both pairs src -> dst and dst -> src are present with the same + * attributes for each pair. + * + * The order of the columns is available in [[edgeColumns]]. + * + * @group structure + */ + // TODO(tjhunter) eventually clarify the treatment of duplicate edges def edges: DataFrame = { - requireDriverDataFrame(edgeDataFrame) - edgeDataFrame + if (_edges == null) { + throw new Exception("You cannot use GraphFrame objects within a Spark closure") + } + _edges } /** - * Return `(source vertex)-[edge]->(destination vertex)` triplets. + * Returns triplets: (source vertex)-[edge]->(destination vertex) for all edges in the graph. + * The DataFrame returned has 3 columns, with names: [[GraphFrame.SRC]], [[GraphFrame.EDGE]], + * and [[GraphFrame.DST]]. Each column is a struct. The 2 vertex columns have schema matching + * [[GraphFrame.vertices]], and the edge column has a schema matching [[GraphFrame.edges]]. For + * example, `triplets.select(col(SRC)(ID))` selects ID of the source column. * - * The result has `src`, `edge`, and `dst` struct columns containing the complete corresponding - * input rows. + * @group structure */ - @transient lazy val triplets: DataFrame = { - val sourceVertices = vertices.select( - col(quoted(ID)).alias("__graphframes_src_id"), - nested(vertices, SRC)) - val graphEdges = edges.select( - col(quoted(SRC)).alias("__graphframes_edge_src"), - col(quoted(DST)).alias("__graphframes_edge_dst"), - nested(edges, EDGE)) - val destinationVertices = vertices.select( - col(quoted(ID)).alias("__graphframes_dst_id"), - nested(vertices, DST)) - - sourceVertices + lazy val triplets: DataFrame = { + vertices + .select(col(ID).alias("src_id"), nestAsCol(vertices, SRC)) .join( - graphEdges, - col("__graphframes_src_id") === col("__graphframes_edge_src")) + edges + .select(col(SRC).alias("edge_src"), col(DST).alias("edge_dst"), nestAsCol(edges, EDGE)), + col("src_id") === col("edge_src")) .join( - destinationVertices, - col("__graphframes_dst_id") === col("__graphframes_edge_dst")) - .select(col(SRC), col(EDGE), col(DST)) + vertices.select(col(ID).alias("dst_id"), nestAsCol(vertices, DST)), + col("dst_id") === col("edge_dst")) + .drop("src_id", "edge_src", "dst_id", "edge_dst") + } + + // ============================ Conversions ======================================== + + /** + * Converts this [[GraphFrame]] instance to a GraphX `Graph`. Vertex and edge attributes are the + * original rows in [[vertices]] and [[edges]], respectively. + * + * Note that vertex (and edge) attributes include vertex IDs (and source, destination IDs) in + * order to support non-Long vertex IDs. If the vertex IDs are not convertible to Long values, + * then the values are indexed in order to generate corresponding Long vertex IDs (which is an + * expensive operation). + * + * The column ordering of the returned `Graph` vertex and edge attributes are specified by + * [[vertexColumns]] and [[edgeColumns]], respectively. + * + * @group conversions + */ + def toGraphX: Graph[Row, Row] = { + if (hasIntegralIdType) { + val vv = vertices.select(col(ID).cast(LongType), nestAsCol(vertices, ATTR)).rdd.map { + case Row(id: Long, attr: Row) => (id, attr) + case Row(null, _) => + throw new IllegalArgumentException( + s"Vertex ID cannot be null. Found null in column '$ID'.") + case _ => throw new GraphFramesUnreachableException() + } + val ee = edges + .select(col(SRC).cast(LongType), col(DST).cast(LongType), nestAsCol(edges, ATTR)) + .rdd + .map { + case Row(srcId: Long, dstId: Long, attr: Row) => Edge(srcId, dstId, attr) + case Row(null, _, _) | Row(_, null, _) => + throw new IllegalArgumentException(s"Edge '$SRC' and '$DST' cannot be null.") + case _ => throw new GraphFramesUnreachableException() + } + Graph[Row, Row](vv, ee) + } else { + // Compute Long vertex IDs + val vv = indexedVertices.select(LONG_ID, ATTR).rdd.map { + case Row(long_id: Long, attr: Row) => (long_id, attr) + case _ => throw new GraphFramesUnreachableException() + } + val ee = indexedEdges.select(LONG_SRC, LONG_DST, ATTR).rdd.map { + case Row(long_src: Long, long_dst: Long, attr: Row) => + Edge(long_src, long_dst, attr) + case _ => throw new GraphFramesUnreachableException() + } + Graph[Row, Row](vv, ee) + } } - /** Return the out-degree of every vertex having at least one outgoing edge. */ + /** + * The column names in the [[vertices]] DataFrame, in order. + * + * Helper method for [[toGraphX]] which specifies the schema of vertex attributes. The vertex + * attributes of the returned `Graph` are given as a `Row`, and this method defines the column + * ordering in that `Row`. + * + * @group conversions + */ + def vertexColumns: Array[String] = vertices.columns + + /** + * Version of [[vertexColumns]] which maps column names to indices in the Rows. + * + * @group conversions + */ + def vertexColumnMap: Map[String, Int] = vertexColumns.zipWithIndex.toMap + + /** + * The vertex names in the [[vertices]] DataFrame, in order. + * + * Helper method for [[toGraphX]] which specifies the schema of edge attributes. The edge + * attributes of the returned `edges` are given as a `Row`, and this method defines the column + * ordering in that `Row`. + * + * @group conversions + */ + def edgeColumns: Array[String] = edges.columns + + /** + * Version of [[edgeColumns]] which maps column names to indices in the Rows. + * + * @group conversions + */ + def edgeColumnMap: Map[String, Int] = edgeColumns.zipWithIndex.toMap + + // ============================ Degree metrics ======================================= + + /** + * The out-degree of each vertex in the graph, returned as a DataFrame with two columns: + * - [[GraphFrame.ID]] the ID of the vertex + * - "outDegree" (integer) storing the out-degree of the vertex Note that vertices with 0 + * out-edges are not returned in the result. + * + * @group degree + */ @transient lazy val outDegrees: DataFrame = { - edges - .groupBy(col(quoted(SRC)).alias(ID)) - .agg(count("*").cast("int").alias(OUT_DEGREE)) + edges.groupBy(edges(SRC).as(ID)).agg(count("*").cast("int").as("outDegree")) } - /** Return the in-degree of every vertex having at least one incoming edge. */ + /** + * The in-degree of each vertex in the graph, returned as a DataFame with two columns: + * - [[GraphFrame.ID]] the ID of the vertex "- "inDegree" (int) storing the in-degree of the + * vertex Note that vertices with 0 in-edges are not returned in the result. + * + * @group degree + */ @transient lazy val inDegrees: DataFrame = { - edges - .groupBy(col(quoted(DST)).alias(ID)) - .agg(count("*").cast("int").alias(IN_DEGREE)) + edges.groupBy(edges(DST).as(ID)).agg(count("*").cast("int").as("inDegree")) } - /** Return the total degree of every vertex incident to at least one edge. */ + /** + * The degree of each vertex in the graph, returned as a DataFrame with two columns: + * - [[GraphFrame.ID]] the ID of the vertex + * - 'degree' (integer) the degree of the vertex Note that vertices with 0 edges are not + * returned in the result. + * + * @group degree + */ @transient lazy val degrees: DataFrame = { edges - .select(explode(array(col(quoted(SRC)), col(quoted(DST)))).alias(ID)) + .select(explode(array(SRC, DST)).as(ID)) .groupBy(ID) - .agg(count("*").cast("int").alias(DEGREE)) + .agg(count("*").cast("int").as("degree")) + } + + /** + * The out-degree of each vertex per edge type, returned as a DataFrame with two columns: + * - [[GraphFrame.ID]] the ID of the vertex + * - "outDegrees" a struct with a field for each edge type, storing the out-degree count + * + * @param edgeTypeCol + * Name of the column in edges DataFrame that contains edge types + * @param edgeTypes + * Optional sequence of edge type values. If None, edge types will be discovered + * automatically. + * @group degree + */ + def typeOutDegree(edgeTypeCol: String, edgeTypes: Option[Seq[Any]] = None): DataFrame = { + val pivotDF = edgeTypes match { + case Some(types) => + edges.groupBy(col(SRC).as(ID)).pivot(edgeTypeCol, types) + case None => + edges.groupBy(col(SRC).as(ID)).pivot(edgeTypeCol) + } + val countDF = pivotDF.agg(count(lit(1))).na.fill(0) + val structCols = countDF.columns + .filter(_ != ID) + .map { colName => + col(colName).cast("int").as(colName) + } + .toSeq + countDF.select(col(ID), struct(structCols: _*).as("outDegrees")) + } + + /** + * The in-degree of each vertex per edge type, returned as a DataFrame with two columns: + * - [[GraphFrame.ID]] the ID of the vertex + * - "inDegrees" a struct with a field for each edge type, storing the in-degree count + * + * @param edgeTypeCol + * Name of the column in edges DataFrame that contains edge types + * @param edgeTypes + * Optional sequence of edge type values. If None, edge types will be discovered + * automatically. + * @group degree + */ + def typeInDegree(edgeTypeCol: String, edgeTypes: Option[Seq[Any]] = None): DataFrame = { + val pivotDF = edgeTypes match { + case Some(types) => + edges.groupBy(col(DST).as(ID)).pivot(edgeTypeCol, types) + case None => + edges.groupBy(col(DST).as(ID)).pivot(edgeTypeCol) + } + val countDF = pivotDF.agg(count(lit(1))).na.fill(0) + val structCols = countDF.columns + .filter(_ != ID) + .map { colName => + col(colName).cast("int").as(colName) + } + .toSeq + countDF.select(col(ID), struct(structCols: _*).as("inDegrees")) + } + + /** + * The total degree of each vertex per edge type (both in and out), returned as a DataFrame with + * two columns: + * - [[GraphFrame.ID]] the ID of the vertex + * - "degrees" a struct with a field for each edge type, storing the total degree count + * + * @param edgeTypeCol + * Name of the column in edges DataFrame that contains edge types + * @param edgeTypes + * Optional sequence of edge type values. If None, edge types will be discovered + * automatically. + * @group degree + */ + def typeDegree(edgeTypeCol: String, edgeTypes: Option[Seq[Any]] = None): DataFrame = { + val explodedEdges = edges.select(explode(array(col(SRC), col(DST))).as(ID), col(edgeTypeCol)) + + val pivotDF = edgeTypes match { + case Some(types) => + explodedEdges.groupBy(ID).pivot(edgeTypeCol, types) + case None => + explodedEdges.groupBy(ID).pivot(edgeTypeCol) + } + val countDF = pivotDF.agg(count(lit(1))).na.fill(0) + val structCols = countDF.columns + .filter(_ != ID) + .map { colName => + col(colName).cast("int").as(colName) + } + .toSeq + + countDF.select(col(ID), struct(structCols: _*).as("degrees")) + } + + // ============================ Motif finding ======================================== + + /** + * Motif finding: Searching the graph for structural patterns + * + * Motif finding uses a simple Domain-Specific Language (DSL) for expressing structural queries. + * For example, `graph.find("(a)-[e]->(b); (b)-[e2]->(a)")` will search for pairs of vertices + * `a,b` connected by edges in both directions. It will return a `DataFrame` of all such + * structures in the graph, with columns for each of the named elements (vertices or edges) in + * the motif. In this case, the returned columns will be in order of the pattern: "a, e, b, e2." + * + * DSL for expressing structural patterns: + * - The basic unit of a pattern is an edge. For example, `"(a)-[e]->(b)"` expresses an edge + * `e` from vertex `a` to vertex `b`. Note that vertices are denoted by parentheses `(a)`, + * while edges are denoted by square brackets `[e]`. + * - A pattern is expressed as a union of edges. Edge patterns can be joined with semicolons. + * Motif `"(a)-[e]->(b); (b)-[e2]->(c)"` specifies two edges from `a` to `b` to `c`. + * - Within a pattern, names can be assigned to vertices and edges. For example, + * `"(a)-[e]->(b)"` has three named elements: vertices `a,b` and edge `e`. These names serve + * two purposes: + * - The names can identify common elements among edges. For example, `"(a)-[e]->(b); + * (b)-[e2]->(c)"` specifies that the same vertex `b` is the destination of edge `e` and + * source of edge `e2`. + * - The names are used as column names in the result `DataFrame`. If a motif contains named + * vertex `a`, then the result `DataFrame` will contain a column "a" which is a + * `StructType` with sub-fields equivalent to the schema (columns) of + * [[GraphFrame.vertices]]. Similarly, an edge `e` in a motif will produce a column "e" in + * the result `DataFrame` with sub-fields equivalent to the schema (columns) of + * [[GraphFrame.edges]]. + * - Be aware that names do *not* identify *distinct* elements: two elements with different + * names may refer to the same graph element. For example, in the motif `"(a)-[e]->(b); + * (b)-[e2]->(c)"`, the names `a` and `c` could refer to the same vertex. To restrict + * named elements to be distinct vertices or edges, use post-hoc filters such as + * `resultDataframe.filter("a.id != c.id")`. + * - It is acceptable to omit names for vertices or edges in motifs when not needed. E.g., + * `"(a)-[]->(b)"` expresses an edge between vertices `a,b` but does not assign a name to + * the edge. There will be no column for the anonymous edge in the result `DataFrame`. + * Similarly, `"(a)-[e]->()"` indicates an out-edge of vertex `a` but does not name the + * destination vertex. These are called *anonymous* vertices and edges. + * - An edge can be negated to indicate that the edge should *not* be present in the graph. + * E.g., `"(a)-[]->(b); !(b)-[]->(a)"` finds edges from `a` to `b` for which there is *no* + * edge from `b` to `a`. + * + * Restrictions: + * - Motifs are not allowed to contain edges without any named elements: `"()-[]->()"` and + * `"!()-[]->()"` are prohibited terms. + * - Motifs are not allowed to contain named edges within negated terms (since these named + * edges would never appear within results). E.g., `"!(a)-[ab]->(b)"` is invalid, but + * `"!(a)-[]->(b)"` is valid. + * + * More complex queries, such as queries which operate on vertex or edge attributes, can be + * expressed by applying filters to the result `DataFrame`. + * + * This can return duplicate rows. E.g., a query `"(u)-[]->()"` will return a result for each + * matching edge, even if those edges share the same vertex `u`. + * + * ==Performance== + * Motif finding translates patterns into a series of joins. Enabling Spark's Cost-Based + * Optimizer (CBO) and join reordering can significantly improve performance by letting Spark + * choose more efficient join orderings based on table statistics: + * {{{ + * spark.conf.set("spark.sql.cbo.enabled", "true") + * spark.conf.set("spark.sql.cbo.joinReorder.enabled", "true") + * }}} + * The join reorder algorithm is bounded by `spark.sql.cbo.joinReorder.dp.threshold` (default: + * `12`). If the estimated number of joins in your motif exceeds this threshold, increase it + * accordingly: + * {{{ + * spark.conf.set("spark.sql.cbo.joinReorder.dp.threshold", "20") + * }}} + * CBO relies on table statistics, so run `ANALYZE TABLE COMPUTE STATISTICS` on the + * vertices and edges tables to ensure accurate statistics are available. + * + * @param pattern + * Pattern specifying a motif to search for. + * @return + * `DataFrame` containing all instances of the motif. + * @group motif + */ + def find(pattern: String): DataFrame = { + val VarLengthPattern = """\((\w*)\)-\[(\w*)\*(\d*)\.\.(\d*)\]-(>?)\((\w*)\)""".r + val FixedLengthUndirectedPattern = """\((\w*)\)-\[(\w*)\*(\d*)\]-\((\w*)\)""".r + + pattern match { + case VarLengthPattern(src, name, min, max, direction, dst) => + if (min.isEmpty || max.isEmpty) { + throw new InvalidParseException( + s"Unbounded length pattern ${pattern} is not supported! " + + "Please a pattern of defined length.") + } + findVarLengthPattern(src, name, min.toInt, max.toInt, direction, dst) + + case FixedLengthUndirectedPattern(src, name, hop, dst) => + if (hop.isEmpty) { + throw new InvalidParseException("Missing hop!") + } + findVarLengthPattern(src, name, hop.toInt, hop.toInt, "", dst) + + case _ => + findAugmentedPatterns(pattern) + } + } + + def findVarLengthPattern( + src: String, + name: String, + min: Int, + max: Int, + direction: String, + dst: String): DataFrame = { + val strToSeq: Seq[(Int, String)] = (min to max).reverse.map { hop => + (hop, s"($src)-[$name*$hop]->($dst)") + } + val strToSeqReverse: Seq[(Int, String)] = if (direction.isEmpty) { + (min to max).reverse.map(hop => (hop, s"($src)<-[$name*$hop]-($dst)")) + } else { + Seq.empty[(Int, String)] + } + + val out: Seq[DataFrame] = strToSeq.map { case (hop, patternStr) => + findAugmentedPatterns(patternStr) + .withColumn("_hop", lit(hop)) + .withColumn("_pattern", lit(patternStr)) + .withColumn("_direction", lit("out")) + } + + val in: Seq[DataFrame] = strToSeqReverse.map { case (hop, patternStr) => + findAugmentedPatterns(patternStr) + .withColumn("_hop", lit(hop)) + .withColumn("_pattern", lit(patternStr)) + .withColumn("_direction", lit("in")) + } + + val ret = (out ++ in).reduce((a, b) => a.unionByName(b, allowMissingColumns = true)) + ret.orderBy("_hop", "_direction") } - /** Return a graph with the direction of every edge reversed. */ - def reverse: GraphFrame = { - val attributes = edges.columns - .filterNot(name => name == SRC || name == DST) - .map(name => col(quoted(name))) - val reversedColumns = Seq( - col(quoted(DST)).alias(SRC), - col(quoted(SRC)).alias(DST)) ++ attributes - val reversedEdges = edges.select(reversedColumns: _*) - GraphFrame(vertices, reversedEdges) + def findAugmentedPatterns(pattern: String): DataFrame = { + val patterns = Pattern.parse(pattern) + + // For each named vertex appearing only in a negated term, we augment the positive terms + // with the vertex as a standalone term `(v)`. + // See https://github.com/graphframes/graphframes/issues/276 + val namedVerticesOnlyInNegatedTerms = Pattern.findNamedVerticesOnlyInNegatedTerms(patterns) + val extraPositivePatterns = namedVerticesOnlyInNegatedTerms.map(v => NamedVertex(v)) + val augmentedPatterns = extraPositivePatterns ++ patterns + val df = findSimple(augmentedPatterns) + + val names = Pattern + .findNamedElementsInOrder(patterns, includeEdges = true) + .filter(x => !x.startsWith("__tmpv")) + if (names.isEmpty) df else df.select(quote(names.head), names.tail.map(quote): _*) } - /** Return an undirected graph by adding a reversed copy of every edge. */ - def asUndirected: GraphFrame = GraphFrame(vertices, edges.unionByName(reverse.edges)) + // ======================== Other queries =================================== + + /** + * Breadth-first search (BFS) + * + * Refer to the documentation of [[org.apache.spark.graphframes.lib.BFS]] for the description of + * the output. + * + * @group stdlib + */ + def bfs: BFS = new BFS(this) + + /** + * Enumerate all paths between source and destination vertices. + * + * See [[org.apache.spark.graphframes.lib.AllPaths]] for details. + * + * @group stdlib + */ + def allPaths: AllPaths = new AllPaths(this) + + /** + * Aggregate information from neighboring vertices and edges through a controlled traversal. + * + * This method provides a flexible way to perform graph traversals while accumulating state at + * each vertex. It can be used to implement various graph algorithms that require propagating + * information through the graph, such as influence propagation, belief propagation, or custom + * message-passing algorithms. + * + * The traversal starts from a set of starting vertices (by default all vertices) and proceeds + * for up to a specified number of hops. At each step, accumulators are updated based on + * neighboring vertices and edges. The traversal can be stopped early based on conditions, and + * results can be collected when target conditions are met. + * + * Key features: + * - Configurable starting vertices via `setStartingVertices()` + * - Maximum number of hops via `setMaxHops()` + * - Accumulators to maintain state during traversal via `setAccumulators()` or + * `addAccumulator()` + * - Stopping conditions to terminate traversal early via `setStoppingCondition()` + * - Target conditions to collect results when specific conditions are met via + * `setTargetCondition()` + * - Edge filtering via `setEdgeFilter()` + * - Control over intermediate storage and checkpointing + * + * The algorithm works as follows: + * 1. Initialize accumulators for starting vertices + * 2. For each iteration up to maxHops: + * - Join current frontier with edges to get neighbors + * - Update accumulators using the provided update expressions + * - Apply stopping conditions to determine which vertices should stop + * - Apply target conditions to determine which stopped vertices should be collected + * - Continue with vertices that haven't stopped + * 3. Return collected results as a DataFrame + * + * The result DataFrame contains: + * - The accumulators' final values for collected vertices + * - The vertex ID (in column "id") + * - The number of hops taken (in column "hop") + * + * Note: This is a stateful iterative algorithm that may be performance-intensive for large + * graphs or large maxHops values. Consider using appropriate storage levels and checkpoint + * intervals for stability. + * + * @see + * [[org.apache.spark.graphframes.lib.AggregateNeighbors]] for implementation details + * @return + * an [[org.apache.spark.graphframes.lib.AggregateNeighbors]] instance for configuration + * @group stdlib + */ + def aggregateNeighbors: AggregateNeighbors = new AggregateNeighbors(this) + + /** + * This is a primitive for implementing graph algorithms. This method aggregates values from the + * neighboring edges and vertices of each vertex. See + * [[org.apache.spark.graphframes.lib.AggregateMessages AggregateMessages]] for detailed + * documentation. + * + * @group stdlib + */ + def aggregateMessages: AggregateMessages = new AggregateMessages(this) - /** Filter vertices and remove edges incident to any removed vertex. */ + /** + * Filter the vertices according to Column expression, remove edges containing any dropped + * vertices. + * @group subgraph + */ def filterVertices(condition: Column): GraphFrame = { - val filteredVertices = vertices.filter(condition) - val vertexIds = filteredVertices.select(col(quoted(ID))) - val filteredEdges = edges - .join(vertexIds, col(quoted(SRC)) === vertexIds(ID), "left_semi") - .join(vertexIds, col(quoted(DST)) === vertexIds(ID), "left_semi") - GraphFrame(filteredVertices, filteredEdges) + val vv = vertices.filter(condition) + val ee = edges + .join(vv, vv(ID) === edges(SRC), "left_semi") + .join(vv, vv(ID) === edges(DST), "left_semi") + GraphFrame(vv, ee) } - /** Filter vertices using a SQL expression. */ - def filterVertices(condition: String): GraphFrame = filterVertices(expr(condition)) + /** + * Filter the vertices according to String expression, remove edges containing any dropped + * vertices. + * @group subgraph + */ + def filterVertices(conditionExpr: String): GraphFrame = filterVertices(expr(conditionExpr)) - /** Filter edges while keeping all vertices. */ - def filterEdges(condition: Column): GraphFrame = GraphFrame(vertices, edges.filter(condition)) + /** + * Filter the edges according to Column expression, keep all vertices. + * @group subgraph + */ + def filterEdges(condition: Column): GraphFrame = { + val vv = vertices + val ee = edges.filter(condition) + GraphFrame(vv, ee) + } - /** Filter edges using a SQL expression. */ - def filterEdges(condition: String): GraphFrame = filterEdges(expr(condition)) + /** + * Filter the edges according to String expression. + * @group subgraph + */ + def filterEdges(conditionExpr: String): GraphFrame = filterEdges(expr(conditionExpr)) - /** Return a graph without vertices that are not incident to an edge. */ + /** + * Drop isolated vertices, vertices not contained in any edges. + * @group subgraph + */ def dropIsolatedVertices(): GraphFrame = { - val incidentIds = edges.select(explode(array(col(quoted(SRC)), col(quoted(DST)))).alias(ID)) - GraphFrame(vertices.join(incidentIds, Seq(ID), "left_semi"), edges) + val ee = edges + val e1 = ee.withColumn(ID, explode(array(col(SRC), col(DST)))) + val vv = vertices.join(e1, Seq(ID), "left_semi") + GraphFrame(vv, ee) } + // **** Standard library **** + /** - * Validate vertex uniqueness and ensure that every edge endpoint is present in `vertices`. + * Connected component algorithm. + * + * See [[org.apache.spark.graphframes.lib.ConnectedComponents]] for more details. * - * This method runs Spark jobs and throws [[InvalidGraphException]] for an invalid graph. + * @group stdlib */ - def validate(): Unit = { - val persistedVertices = vertices.persist(StorageLevel.MEMORY_AND_DISK) - try { - val vertexCount = persistedVertices.count() - val distinctVertexCount = persistedVertices.select(countDistinct(col(quoted(ID)))).head() - .getLong(0) - if (vertexCount != distinctVertexCount) { - throw new InvalidGraphException( - s"Graph contains ${vertexCount - distinctVertexCount} duplicate vertices") + def connectedComponents: ConnectedComponents = new ConnectedComponents(this) + + /** + * Label propagation algorithm. + * + * See [[org.apache.spark.graphframes.lib.LabelPropagation]] for more details. + * + * @group stdlib + */ + def labelPropagation: LabelPropagation = new LabelPropagation(this) + + /** + * Mix of label- and structure propagation. + * + * See [[org.apache.spark.graphframes.lib.StructureAwareLabelPropagation]] for more details. + * + * @group stdlib + */ + def structureAwareLabelPropagation: StructureAwareLabelPropagation = + new StructureAwareLabelPropagation(this) + + /** + * PageRank algorithm. + * + * See [[org.apache.spark.graphframes.lib.PageRank]] for more details. + * + * @group stdlib + */ + def pageRank: PageRank = new PageRank(this) + + /** + * Parallel personalized PageRank algorithm. + * + * See [[org.apache.spark.graphframes.lib.ParallelPersonalizedPageRank]] for more details. + * + * @group stdlib + */ + def parallelPersonalizedPageRank: ParallelPersonalizedPageRank = + new ParallelPersonalizedPageRank(this) + + /** + * Pregel algorithm. + * + * @see + * [[org.apache.spark.graphframes.lib.Pregel]] + * @group stdlib + */ + def pregel = new Pregel(this) + + /** + * Shortest paths algorithm. + * + * See [[org.apache.spark.graphframes.lib.ShortestPaths]] for more details. + * + * @group stdlib + */ + def shortestPaths: ShortestPaths = new ShortestPaths(this) + + /** + * Strongly connected components algorithm. + * + * See [[org.apache.spark.graphframes.lib.StronglyConnectedComponents]] for more details. + * + * @group stdlib + */ + def stronglyConnectedComponents: StronglyConnectedComponents = + new StronglyConnectedComponents(this) + + /** + * SVD++ algorithm. + * + * See [[org.apache.spark.graphframes.lib.SVDPlusPlus]] for more details. + * + * @group stdlib + */ + def svdPlusPlus: SVDPlusPlus = new SVDPlusPlus(this) + + /** + * Triangle count algorithm. + * + * See [[org.apache.spark.graphframes.lib.TriangleCount]] for more details. + * + * @group stdlib + */ + def triangleCount: TriangleCount = new TriangleCount(this) + + /** + * Power Iteration Clustering (PIC), a scalable graph clustering algorithm developed by Lin and + * Cohen. From the abstract: PIC finds a very low-dimensional embedding of a dataset using + * truncated power iteration on a normalized pair-wise similarity matrix of the data. + * + * PowerIterationClustering algorithm. + * @param k + * The number of clusters to create (k). + * @param maxIter + * Param for maximum number of iterations (>= 0). + * @param weightCol + * Param for weight column name. + * + * @group stdlib + */ + def powerIterationClustering(k: Int, maxIter: Int, weightCol: Option[String]): DataFrame = { + val integralTypeEdges = if (hasIntegralIdType) { + edges + } else { + val pureIds = + indexedEdges.drop(SRC, DST).withColumnsRenamed(Map(LONG_SRC -> SRC, LONG_DST -> DST)) + if (weightCol.isDefined) { + pureIds.select( + col(SRC), + col(DST), + col("attr").getField(weightCol.get).alias(weightCol.get)) + } else { + pureIds } + } + val powerIterationClustering = + new PowerIterationClustering().setK(k).setMaxIter(maxIter).setDstCol(DST).setSrcCol(SRC) + val result = weightCol match { + case Some(col) => + powerIterationClustering.setWeightCol(col).assignClusters(integralTypeEdges) + case None => + powerIterationClustering + .setWeightCol("_weight") + .assignClusters(integralTypeEdges.withColumn("_weight", lit(1.0))) + } - val endpoints = edges - .select(col(quoted(SRC)).alias(ID)) - .union(edges.select(col(quoted(DST)).alias(ID))) - .distinct() - val missingEndpointCount = endpoints.join(persistedVertices, Seq(ID), "left_anti").count() - if (missingEndpointCount > 0) { - throw new InvalidGraphException( - s"Graph contains $missingEndpointCount edge endpoints without matching vertices") + if (hasIntegralIdType) { + result + } else { + result + .join( + indexedVertices.select(col(LONG_ID).alias(ID), col(ID).alias("_ID")), + Seq(ID), + "inner") + .select(col("_ID").alias(ID), col("cluster")) + } + } + + /** + * K-Core decomposition. + * + * See [[org.apache.spark.graphframes.lib.KCore]] for more details. + * + * @group stdlib + */ + def kCore: KCore = new KCore(this) + + /** + * Find all cycles in the graph. An implementation of the Rocha–Thatte cycle detection + * algorithm. + * + * Rocha, Rodrigo Caetano, and Bhalchandra D. Thatte. "Distributed cycle detection in + * large-scale sparse graphs." Proceedings of Simpósio Brasileiro de Pesquisa Operacional + * (SBPO’15) (2015): 1-11. + * + * Returns a DataFrame with unique cycles. + * + * @return + * an instance of DetectingCycles initialized with the current context + * + * @group stdlib + */ + def detectingCycles: DetectingCycles = new DetectingCycles(this) + + /** + * Maximal Independent Set algorithm. + * + * See [[org.apache.spark.graphframes.lib.MaximalIndependentSet]] for more details. + * + * @group stdlib + */ + def maximalIndependentSet: MaximalIndependentSet = new MaximalIndependentSet(this) + + /** + * Run an approximate neighbor function backed by the HLL-sketches. + * + * See [[org.apache.spark.graphframes.lib.HyperANF]] for more details. + * + * @group stdlib + */ + def hyperANF: HyperANF = new HyperANF(this) + + // ========= Graph Machine Learning ========== + + /** + * Random Walks Based node embeddings. + * + * See [[org.apache.spark.graphframes.embeddings.RandomWalkEmbeddings]] for more details. + * + * @group gml + */ + def randomWalksBasedEmbedding: RandomWalkEmbeddings = new RandomWalkEmbeddings(this) + + // ========= Motif finding (private) ========= + + /** + * Primary method implementing motif finding. This iterative method handles one pattern (via + * [[findIncremental()]] on each iteration, augmenting the `DataFrame` in prevDF with each new + * pattern. + * + * @return + * `DataFrame` containing all instances of the motif specified by the given patterns + */ + private def findSimple(patterns: Seq[Pattern]): DataFrame = { + val (_, finalDFOpt, _) = + patterns.foldLeft((Seq.empty[Pattern], Option.empty[DataFrame], Seq.empty[String])) { + case ((handledPatterns, dfOpt, names), cur) => + val (nextDF, nextNames) = findIncremental(this, handledPatterns, dfOpt, names, cur) + (handledPatterns :+ cur, nextDF, nextNames) } - } finally { - persistedVertices.unpersist() + finalDFOpt.getOrElse(spark.emptyDataFrame) + } + + // ========= Other private methods =========== + + private[graphframes] def spark: SparkSession = vertices.sparkSession + + /** + * True if the id type can be cast to Long. + * + * This is important for performance reasons. The underlying graphx implementation only deals + * with Long types. + */ + private[graphframes] lazy val hasIntegralIdType: Boolean = { + vertices.schema(ID).dataType match { + case _ @(ByteType | IntegerType | LongType | ShortType) => true + case _ => false } } - override def toString: String = { - val vertexColumns = ID +: vertices.columns.filterNot(_ == ID).toSeq - val edgeColumns = SRC +: DST +: edges.columns.filterNot(c => c == SRC || c == DST).toSeq - val orderedVertices = vertices.select(vertexColumns.map(name => col(quoted(name))): _*) - val orderedEdges = edges.select(edgeColumns.map(name => col(quoted(name))): _*) - s"GraphFrame(v:$orderedVertices, e:$orderedEdges)" + /** + * Vertices with each vertex assigned a unique long ID. If the vertex ID type is integral, this + * casts the original IDs to long. + * + * Columns: + * - $LONG_ID: the new ID of LongType + * - $ORIGINAL_ID: the ID provided by the user + * - $ATTR: all the original vertex attributes + */ + private[graphframes] lazy val indexedVertices: DataFrame = { + if (hasIntegralIdType) { + val indexedVertices = vertices.select(nestAsCol(vertices, ATTR)) + indexedVertices.select( + col(ATTR + "." + ID).cast("long").as(LONG_ID), + col(ATTR + "." + ID).as(ID), + col(ATTR)) + } else { + val withLongIds = vertices + .select(ID) + .repartition(col(ID)) + .sortWithinPartitions(ID) + .withColumn(LONG_ID, monotonically_increasing_id()) + .persist(StorageLevel.MEMORY_AND_DISK) + vertices + .select(col(ID), nestAsCol(vertices, ATTR)) + .join(withLongIds, ID) + .select(LONG_ID, ID, ATTR) + } } - private def requireDriverDataFrame(dataFrame: DataFrame): Unit = { - if (dataFrame == null) { - throw new IllegalStateException("GraphFrame objects cannot be used inside Spark closures") + /** + * Columns: + * - $SRC + * - $LONG_SRC + * - $DST + * - $LONG_DST + * - $ATTR + */ + private[graphframes] lazy val indexedEdges: DataFrame = { + val packedEdges = edges.select(col(SRC), col(DST), nestAsCol(edges, ATTR)) + if (hasIntegralIdType) { + packedEdges.select( + col(SRC), + col(SRC).cast("long").as(LONG_SRC), + col(DST), + col(DST).cast("long").as(LONG_DST), + col(ATTR)) + } else { + val indexedSourceEdges = + packedEdges.join(indexedVertices.select(col(ID).as(SRC), col(LONG_ID).as(LONG_SRC)), SRC) + val indexedEdges = indexedSourceEdges.join( + indexedVertices.select(col(ID).as(DST), col(LONG_ID).as(LONG_DST)), + DST) + indexedEdges.select(SRC, LONG_SRC, DST, LONG_DST, ATTR) } } + + /** + * A cached conversion of this graph to the GraphX structure. All the data is stripped away. + */ + @transient lazy private[graphframes] val cachedTopologyGraphX: Graph[Unit, Unit] = { + cachedGraphX.mapVertices((_, _) => ()).mapEdges(_ => ()) + } + + /** + * A cached conversion of this graph to the GraphX structure, with the data stored for each edge + * and vertex. + */ + @transient private lazy val cachedGraphX: Graph[Row, Row] = { toGraphX } + } -object GraphFrame extends Logging { +object GraphFrame extends Serializable with Logging { + /** + * Implements `a.join(b, joinCol)`, handling skew in the join keys. + * @param a + * DataFrame which may have multiple rows with the same key in `joinCol` + * @param b + * DataFrame which has exactly 1 row for every key in `a.joinCol`. + * @param joinCol + * Name of column on which to do join + * @param hubs + * Set of join keys which are high-degree (skewed) + * @param logPrefix + * Prefix for logging, e.g., name of algorithm doing the join + * @return + * `a.join(b, joinCol)` + * @tparam T + * DataType for join key + */ + private[graphframes] def skewedJoin[T]( + a: DataFrame, + b: DataFrame, + joinCol: String, + hubs: Set[T], + logPrefix: String): DataFrame = { + if (hubs.isEmpty) { + // No skew. Do regular join. + a.join(b, joinCol) + } else { + logDebug(s"$logPrefix Skewed join with ${hubs.size} high-degree keys.") + val isHub = (c: Column) => c.isInCollection(hubs) + val hashJoined = a + .filter(!isHub(col(joinCol))) + .join(b.filter(!isHub(col(joinCol))), joinCol) + val broadcastJoined = a + .filter(isHub(col(joinCol))) + .join(broadcast(b.filter(isHub(col(joinCol)))), joinCol) + hashJoined.unionAll(broadcastJoined) + } + } + + /** + * Column name for vertex IDs in [[GraphFrame.vertices]] Note that GraphFrame assigns a unique + * long ID to each vertex, If the vertex ID type is one of byte / int / long / short type, + * GraphFrame casts the original IDs to long as the unique long ID, otherwise GraphFrame + * generates the unique long ID by Spark function ``monotonically_increasing_id`` which is less + * performant. + */ val ID: String = "id" + + /** + * Column name for source vertices of edges. + * - In [[GraphFrame.edges]], this is a column of vertex IDs. + * - In [[GraphFrame.triplets]], this is a column of vertices with schema matching + * [[GraphFrame.vertices]]. + */ val SRC: String = "src" + + /** + * Column name for destination vertices of edges. + * - In [[GraphFrame.edges]], this is a column of vertex IDs. + * - In [[GraphFrame.triplets]], this is a column of vertices with schema matching + * [[GraphFrame.vertices]]. + */ val DST: String = "dst" + + /** + * Column name for edge in [[GraphFrame.triplets]]. In [[GraphFrame.triplets]], this is a column + * of edges with schema matching [[GraphFrame.edges]]. + */ val EDGE: String = "edge" - val DEGREE: String = "degree" - val IN_DEGREE: String = "inDegree" - val OUT_DEGREE: String = "outDegree" - /** Create a GraphFrame from vertex and edge DataFrames. */ + /** + * Column name representing the weight attribute of edges in a graph. + * + * This field is used to identify and represent the weight associated with edges in a + * GraphFrame. The weight generally encodes the strength or importance of the connection between + * two nodes in a graph. + */ + val WEIGHT: String = "weight" + + // ============================ Constructors and converters ================================= + + /** + * Create a new [[GraphFrame]] from vertex and edge `DataFrame`s. + * + * @param vertices + * Vertex DataFrame. This must include a column "id" containing unique vertex IDs. All other + * columns are treated as vertex attributes. + * @param edges + * Edge DataFrame. This must include columns "src" and "dst" containing source and destination + * vertex IDs. All other columns are treated as edge attributes. + * @return + * New [[GraphFrame]] instance + */ def apply(vertices: DataFrame, edges: DataFrame): GraphFrame = { - requireColumn(vertices, ID, "Vertex ID") - requireColumn(edges, SRC, "Source vertex ID") - requireColumn(edges, DST, "Destination vertex ID") require( - vertices.sparkSession eq edges.sparkSession, - "Vertex and edge DataFrames must belong to the same SparkSession") + vertices.columns.contains(ID), + s"Vertex ID column '$ID' missing from vertex DataFrame, which has columns: " + + vertices.columns.mkString(",")) + require( + edges.columns.contains(SRC), + s"Source vertex ID column '$SRC' missing from edge DataFrame, which has columns: " + + edges.columns.mkString(",")) + require( + edges.columns.contains(DST), + s"Destination vertex ID column '$DST' missing from edge DataFrame, which has columns: " + + edges.columns.mkString(",")) + new GraphFrame(vertices, edges) } /** - * Create a GraphFrame from an edge DataFrame, deriving and persisting its distinct vertices. + * Create a new [[GraphFrame]] from an edge `DataFrame`. The resulting [[GraphFrame]] will have + * [[GraphFrame.vertices]] with a single "id" column. + * + * Note: The [[GraphFrame.vertices]] DataFrame will be persisted at level + * `StorageLevel.MEMORY_AND_DISK`. + * @param e + * Edge DataFrame. This must include columns "src" and "dst" containing source and destination + * vertex IDs. All other columns are treated as edge attributes. + * @return + * New [[GraphFrame]] instance + * + * @group conversions */ - def fromEdges(edges: DataFrame): GraphFrame = { - fromEdges(edges, StorageLevel.MEMORY_AND_DISK) + def fromEdges(e: DataFrame): GraphFrame = { + fromEdges(e, StorageLevel.MEMORY_AND_DISK) } /** - * Create a GraphFrame from an edge DataFrame, deriving its distinct vertices. + * Create a new [[GraphFrame]] from an edge `DataFrame`. The resulting [[GraphFrame]] will have + * [[GraphFrame.vertices]] with a single "id" column. * - * The caller is responsible for unpersisting the returned graph's vertex DataFrame. + * Note: The [[GraphFrame.vertices]] DataFrame will be persisted at level + * `StorageLevel.MEMORY_AND_DISK`. + * @param e + * Edge DataFrame. This must include columns "src" and "dst" containing source and destination + * vertex IDs. All other columns are treated as edge attributes. + * @param storageLevel + * StorageLevel to persist the graph vertices + * @return + * New [[GraphFrame]] instance + * + * @group conversions */ - def fromEdges(edges: DataFrame, storageLevel: StorageLevel): GraphFrame = { - requireColumn(edges, SRC, "Source vertex ID") - requireColumn(edges, DST, "Destination vertex ID") - logWarning( - s"GraphFrame.fromEdges persists derived vertices with storage level $storageLevel; " + - "call vertices.unpersist() when the graph is no longer needed") - val vertices = edges - .select(col(quoted(SRC)).alias(ID)) - .union(edges.select(col(quoted(DST)).alias(ID))) - .distinct() - .persist(storageLevel) - GraphFrame(vertices, edges) + def fromEdges(e: DataFrame, storageLevel: StorageLevel): GraphFrame = { + logWarn( + s"this method persists graph vertices with storage level ${storageLevel.toString()}, users should manually unpersist it when the graph is not needed!") + val srcs = e.select(e("src").as("id")) + val dsts = e.select(e("dst").as("id")) + val v = srcs.unionAll(dsts).distinct().persist(storageLevel) + apply(v, e) } - private def requireColumn(dataFrame: DataFrame, columnName: String, label: String): Unit = { - require( - dataFrame.columns.contains(columnName), - s"$label column '$columnName' is missing; available columns: " + - dataFrame.columns.mkString(", ")) + /** + * Converts a GraphX `Graph` instance into a [[GraphFrame]]. + * + * This converts each `org.apache.spark.rdd.RDD` in the `Graph` to a `DataFrame` using schema + * inference. + * + * Vertex ID column names will be converted to "id" for the vertex DataFrame, and to "src" and + * "dst" for the edge DataFrame. + * + * @group conversions + */ + def fromGraphX[VD: TypeTag, ED: TypeTag](graph: Graph[VD, ED]): GraphFrame = { + val spark = SparkSession.builder().getOrCreate() + val vv = spark.createDataFrame(graph.vertices).toDF(ID, ATTR) + val ee = spark.createDataFrame(graph.edges).toDF(SRC, DST, ATTR) + GraphFrame(vv, ee) + } + + /** + * Given: + * - a GraphFrame `originalGraph` + * - a GraphX graph derived from the GraphFrame using [[GraphFrame.toGraphX]] this method + * merges attributes from the GraphX graph into the original GraphFrame. + * + * This method is useful for doing computations using the GraphX API and then merging the + * results with a GraphFrame. For example, given: + * - GraphFrame `originalGraph` + * - GraphX Graph[String, Int] `graph` with a String vertex attribute we want to call + * "category" and an Int edge attribute we want to call "count" We can call + * `fromGraphX(originalGraph, graph, Seq("category"), Seq("count"))` to produce a new + * GraphFrame. The new GraphFrame will be an augmented version of `originalGraph`, with new + * [[GraphFrame.vertices]] column "category" and new [[GraphFrame.edges]] column "count" + * added. + * + * See [[org.apache.spark.graphframes.examples.BeliefPropagation]] for example usage. + * + * @param originalGraph + * Original GraphFrame used to compute the GraphX graph. + * @param graph + * GraphX graph. Vertex and edge attributes, if any, will be merged into the original graph as + * new columns. If the attributes are `Product` types such as tuples, then each element of the + * `Product` will be put in a separate column. If the attributes are other types, then the + * entire GraphX attribute will become a single new column. + * @param vertexNames + * Column name(s) for vertex attributes in the GraphX graph. If there is no vertex attribute, + * this should be empty. If there is a singleton attribute, this should have a single column + * name. If the attribute is a `Product` type, this should be a list of names matching the + * order of the attribute elements. + * @param edgeNames + * Column name(s) for edge attributes in the GraphX graph. If there is no edge attribute, this + * should be empty. If there is a singleton attribute, this should have a single column name. + * If the attribute is a `Product` type, this should be a list of names matching the order of + * the attribute elements. + * @tparam V + * the type of the vertex data + * @tparam E + * the type of the edge data + * @return + * original graph augmented with vertex and column attributes from the GraphX graph + * + * @group conversions + */ + def fromGraphX[V: TypeTag, E: TypeTag]( + originalGraph: GraphFrame, + graph: Graph[V, E], + vertexNames: Seq[String] = Nil, + edgeNames: Seq[String] = Nil): GraphFrame = { + GraphXConversions.fromGraphX[V, E](originalGraph, graph, vertexNames, edgeNames) + } + + // ============== Private constants ============== + + /** Default name for attribute columns when converting from GraphX [[Graph]] format */ + private[graphframes] val ATTR: String = "attr" + + /** + * The integral id that is used as a surrogate id when using graphX implementation + */ + private[graphframes] val LONG_ID: String = "new_id" + + private[graphframes] val LONG_SRC: String = "new_src" + private[graphframes] val LONG_DST: String = "new_dst" + private[graphframes] val GX_ATTR: String = "graphx_attr" + + /** + * Helper for column names containing a dot. Quotes the given column name with backticks to + * avoid further parsing. + * + * Note: This can be replaced with org.apache.spark.sql.catalyst.util.QuotingUtils.quoteIfNeeded + * once support for Spark 3 has been dropped + */ + private[graphframes] def quote(column: String): String = + s"`${column.replace("`", "``")}`" + + /** + * Helper for column names containing a dot. Quotes the given column name with backticks to + * avoid further parsing. The column name can be given in segments, e.g. quote("col", "field") + * representing column "col.field", which returns "`col`.`field`". + * + * Note: This can be replaced with org.apache.spark.sql.catalyst.util.QuotingUtils.quoted once + * support for Spark 3 has been dropped + */ + private[graphframes] def quote(columnSegments: String*): String = + columnSegments.map(quote).mkString(".") + + /** + * Helper for using [col].* in Spark 1.4. Returns sequence of [col].[field] for all fields. Both + * [col] and [field] are quoted with backticks to work with columns and fields containing dots. + */ + private[graphframes] def colStar(df: DataFrame, col: String): Seq[String] = { + df.schema(col).dataType match { + case s: StructType => + s.fieldNames.map(f => quote(col, f)).toIndexedSeq + case other => + throw new RuntimeException( + s"Unknown error in GraphFrame. Expected column $col to be" + + s" StructType, but found type: $other") + } + } + + /** Nest all columns within a single StructType column with the given name */ + private[graphframes] def nestAsCol(df: DataFrame, name: String): Column = { + struct(df.columns.map(quote).map(c => df(c)).toSeq: _*).as(name) + } + + // ========== Motif finding ========== + + private val random: Random = new Random(classOf[GraphFrame].getName.##.toLong) + + private def prefixWithName(name: String, col: String): String = name + "." + col + private def vId(name: String): String = prefixWithName(name, ID) + private def eSrcId(name: String): String = prefixWithName(name, SRC) + private def eDstId(name: String): String = prefixWithName(name, DST) + + private def maybeUnion(aOpt: Option[DataFrame], bOpt: Option[DataFrame]): Option[DataFrame] = { + (aOpt, bOpt) match { + case (Some(a), Some(b)) => + Some(a.unionByName(b, allowMissingColumns = true).orderBy("_direction")) + case (Some(a), None) => Some(a) + case (None, Some(b)) => Some(b) + case (None, None) => None + } + } + + private def maybeCrossJoin(aOpt: Option[DataFrame], b: DataFrame): DataFrame = { + aOpt match { + case Some(a) => a.crossJoin(b) + case None => b + } + } + + private def maybeJoin( + aOpt: Option[DataFrame], + b: DataFrame, + joinExprs: DataFrame => Column): DataFrame = { + aOpt match { + case Some(a) => a.join(b, joinExprs(a)) + case None => b + } } - private def nested(dataFrame: DataFrame, name: String): Column = { - val columns = dataFrame.columns.map(columnName => col(quoted(columnName))).toIndexedSeq - struct(columns: _*).alias(name) + /** Indicate whether a named vertex has been seen in any of the given patterns */ + private def seen(v: NamedVertex, patterns: Seq[Pattern]) = patterns.exists(p => seen1(v, p)) + + /** Indicate whether a named vertex has been seen in the given pattern */ + private def seen1(v: NamedVertex, pattern: Pattern): Boolean = pattern match { + case Negation(edge) => + seen1(v, edge) + case UndirectedEdge(edge) => + seen1(v, edge) + case AnonymousEdge(src, dst) => + seen1(v, src) || seen1(v, dst) + case NamedEdge(_, src, dst) => + seen1(v, src) || seen1(v, dst) + case v2 @ NamedVertex(_) => + v2 == v + case AnonymousVertex => + false } - private def quoted(columnName: String): String = { - s"`${columnName.replace("`", "``")}`" + /** + * Augment the given DataFrame based on a pattern. + * + * @param prevPatterns + * Patterns which have contributed to the given DataFrame + * @param prev + * Given DataFrame + * @param pattern + * Pattern to search for + * @return + * DataFrame augmented with the current search pattern + */ + private def findIncremental( + gf: GraphFrame, + prevPatterns: Seq[Pattern], + prev: Option[DataFrame], + prevNames: Seq[String], + pattern: Pattern): (Option[DataFrame], Seq[String]) = { + def nestE(name: String): DataFrame = gf.edges.select(nestAsCol(gf.edges, name)) + def nestV(name: String): DataFrame = gf.vertices.select(nestAsCol(gf.vertices, name)) + + pattern match { + + case AnonymousVertex => + (prev, prevNames) + + case v @ NamedVertex(name) => + if (seen(v, prevPatterns)) { + for (prev <- prev) assert(prev.columns.toSet.contains(name)) + (prev, prevNames) + } else { + (Some(maybeCrossJoin(prev, nestV(name))), prevNames :+ name) + } + + case UndirectedEdge(edge) => + val srcName: String = edge match { + case NamedEdge(_, NamedVertex(n), _) => n + case AnonymousEdge(NamedVertex(n), _) => n + case _ => "" + } + val dstName: String = edge match { + case NamedEdge(_, _, NamedVertex(n)) => n + case AnonymousEdge(_, NamedVertex(n)) => n + case _ => "" + } + val edgeName: String = edge match { + case NamedEdge(n, _, _) => n + case _ => "" + } + + val patternStr: String = s"($srcName)-[$edgeName]->($dstName)" + val reversedPatternStr: String = s"($srcName)<-[$edgeName]-($dstName)" + + val reversedEdge: Pattern = { + edge match { + case e: NamedEdge => + e.copy(src = e.dst, dst = e.src) + case e: AnonymousEdge => + e.copy(src = e.dst, dst = e.src) + case _ => edge + } + } + + val (dfIn, _) = findIncremental(gf, prevPatterns, prev, prevNames, reversedEdge) + val (dfOut, names) = findIncremental(gf, prevPatterns, prev, prevNames, edge) + + val df1 = dfIn match { + case Some(d) => + Some( + d.withColumn("_pattern", lit(reversedPatternStr)) + .withColumn("_direction", lit("in"))) + case None => None + } + + val df2 = dfOut match { + case Some(d) => + Some( + d.withColumn("_pattern", lit(patternStr)) + .withColumn("_direction", lit("out"))) + case None => None + } + + val df = maybeUnion(df1, df2) + (df, names :+ "_pattern" :+ "_direction") + + case NamedEdge(name, AnonymousVertex, AnonymousVertex) => + val eRen = nestE(name) + (Some(maybeCrossJoin(prev, eRen)), prevNames :+ name) + + case NamedEdge(name, AnonymousVertex, dst @ NamedVertex(dstName)) => + if (seen(dst, prevPatterns)) { + val eRen = nestE(name) + ( + Some(maybeJoin(prev, eRen, prev => eRen(eDstId(name)) === prev(vId(dstName)))), + prevNames :+ name) + } else { + val eRen = nestE(name) + val dstV = nestV(dstName) + ( + Some( + maybeCrossJoin(prev, eRen) + .join(dstV, eRen(eDstId(name)) === dstV(vId(dstName)))), + prevNames :+ name :+ dstName) + } + + case NamedEdge(name, src @ NamedVertex(srcName), AnonymousVertex) => + if (seen(src, prevPatterns)) { + val eRen = nestE(name) + ( + Some(maybeJoin(prev, eRen, prev => eRen(eSrcId(name)) === prev(vId(srcName)))), + prevNames :+ name) + } else { + val eRen = nestE(name) + val srcV = nestV(srcName) + ( + Some( + maybeCrossJoin(prev, eRen) + .join(srcV, eRen(eSrcId(name)) === srcV(vId(srcName)))), + prevNames :+ srcName :+ name) + } + + case NamedEdge(name, src @ NamedVertex(srcName), dst @ NamedVertex(dstName)) => + (seen(src, prevPatterns), seen(dst, prevPatterns)) match { + case (true, true) => + val eRen = nestE(name) + ( + Some( + maybeJoin( + prev, + eRen, + prev => + eRen(eSrcId(name)) === prev(vId(srcName)) && eRen(eDstId(name)) === prev( + vId(dstName)))), + prevNames :+ name) + + case (true, false) => + val eRen = nestE(name) + val dstV = nestV(dstName) + ( + Some( + maybeJoin(prev, eRen, prev => eRen(eSrcId(name)) === prev(vId(srcName))) + .join(dstV, eRen(eDstId(name)) === dstV(vId(dstName)))), + prevNames :+ name :+ dstName) + + case (false, true) => + val eRen = nestE(name) + val srcV = nestV(srcName) + ( + Some( + maybeJoin(prev, eRen, prev => eRen(eDstId(name)) === prev(vId(dstName))) + .join(srcV, eRen(eSrcId(name)) === srcV(vId(srcName)))), + prevNames :+ srcName :+ name) + + case (false, false) if srcName != dstName => + val eRen = nestE(name) + val srcV = nestV(srcName) + val dstV = nestV(dstName) + ( + Some( + maybeCrossJoin(prev, eRen) + .join(srcV, eRen(eSrcId(name)) === srcV(vId(srcName))) + .join(dstV, eRen(eDstId(name)) === dstV(vId(dstName)))), + prevNames :+ srcName :+ name :+ dstName) + // TODO: expose the plans from joining these in the opposite order + + case (false, false) if srcName == dstName => + val eRen = nestE(name) + val srcV = nestV(srcName) + ( + Some( + maybeCrossJoin(prev, eRen) + .join( + srcV, + eRen(eSrcId(name)) === srcV(vId(srcName)) && + eRen(eDstId(name)) === srcV(vId(srcName)))), + prevNames :+ srcName :+ name) + + case _ => throw new GraphFramesUnreachableException() + } + + case AnonymousEdge(src, dst) => + val tmpName = "__tmp" + random.nextLong.toString + val (df, names) = + findIncremental(gf, prevPatterns, prev, prevNames, NamedEdge(tmpName, src, dst)) + (df.map(_.drop(tmpName)), names.filter(_ != tmpName)) + + case Negation(edge) => + prev match { + case Some(p) => + val (df, names) = findIncremental(gf, prevPatterns, Some(p), prevNames, edge) + // TODO: _pattern. _direction columns should be ignored if it is impacting + (df.map(result => p.except(result)), names) + case None => + throw new InvalidPatternException + } + } } } diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/InvalidGraphException.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFramePythonAPI.scala similarity index 63% rename from graphframes/src/main/scala/org/apache/spark/graphframes/InvalidGraphException.scala rename to graphframes/src/main/scala/org/apache/spark/graphframes/GraphFramePythonAPI.scala index a0fe72d3ca1dd..f4b0676c01789 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/InvalidGraphException.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFramePythonAPI.scala @@ -17,5 +17,18 @@ package org.apache.spark.graphframes -/** Thrown when a GraphFrame contains duplicate vertices or unknown edge endpoints. */ -class InvalidGraphException(message: String) extends IllegalArgumentException(message) +import org.apache.spark.sql.DataFrame +import org.apache.spark.graphframes.lib.AggregateMessages + +private[graphframes] class GraphFramePythonAPI { + + def createGraph(v: DataFrame, e: DataFrame): GraphFrame = GraphFrame(v, e) + + val ID: String = GraphFrame.ID + val SRC: String = GraphFrame.SRC + val DST: String = GraphFrame.DST + val EDGE: String = GraphFrame.EDGE + val ATTR: String = GraphFrame.ATTR + + lazy val aggregateMessages: AggregateMessages.type = AggregateMessages +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/Logging.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/Logging.scala new file mode 100644 index 0000000000000..9b45139e04df8 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/Logging.scala @@ -0,0 +1,47 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes + +import org.slf4j.Logger +import org.slf4j.LoggerFactory + +// This needs to be accessible to org.apache.spark.graphx.lib.backport +private[org] trait Logging { + + @transient private lazy val logger: Logger = LoggerFactory.getLogger(getClass.getName) + + protected def logDebug(s: => String): Unit = { + if (logger.isDebugEnabled) logger.debug(s) + } + + protected def logWarn(s: => String): Unit = { + if (logger.isWarnEnabled) logger.warn(s) + } + + protected def logInfo(s: => String): Unit = { + if (logger.isInfoEnabled) logger.info(s) + } + + protected def logTrace(s: => String): Unit = { + if (logger.isTraceEnabled) logger.trace(s) + } + + protected def resultIsPersistent(): Unit = { + logWarn("Returned DataFrame is persistent and materialized!") + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/convolutions/SamplingConvolution.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/convolutions/SamplingConvolution.scala new file mode 100644 index 0000000000000..a4e4b1a90b79a --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/convolutions/SamplingConvolution.scala @@ -0,0 +1,194 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.convolutions + +import org.apache.spark.ml.functions._ +import org.apache.spark.ml.stat.Summarizer +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.graphframes.expressions.KMinSampling +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.Logging + +/** + * A convolution operation on graph data that aggregates features from sampled neighbors using + * min-hash sampling with a seed. + * + * For each vertex in the input GraphFrame, this class samples up to a maximum number of neighbors + * using a min-hash approach based on hashing the destination vertex IDs with a provided seed. The + * feature embeddings of these sampled neighbors are averaged, resulting in an aggregated neighbor + * embedding. Optionally, this aggregated embedding can be concatenated with the original vertex + * features to produce an updated feature vector. + * + * The graph is expected to have a standard structure with ID columns in vertices and edges (via + * GraphFrame conventions). + */ +class SamplingConvolution extends Serializable with Logging { + private var graph: GraphFrame = _ + private var featuresCol: String = "embedding" + private var maxNbrs: Int = 50 + private var useEdgeDirections: Boolean = false + private var seed: Long = 42L + private var concatEmbeddings: Boolean = true + + /** + * Specifies the GraphFrame to perform the convolution on. + * + * @param graph + * The input GraphFrame, which must contain vertices with an ID column and the features column + * as specified by setFeaturesCol. Feature columns should contain vector types (e.g., Vectors + * from MLlib). + * @return + * This SamplingConvolution instance for method chaining. + */ + def onGraph(graph: GraphFrame): this.type = { + this.graph = graph + this + } + + /** + * Sets the name of the column in the vertices DataFrame containing the feature vectors. + * + * @param value + * The string name of the column holding vertex embeddings as Vectors. + * @return + * This SamplingConvolution instance for method chaining. + */ + def setFeaturesCol(value: String): this.type = { + featuresCol = value + this + } + + /** + * Sets the maximum number of neighbors to sample per vertex. + * + * @param value + * The maximum number of neighbors (integer) to select using min-hash sampling. + * @return + * This SamplingConvolution instance for method chaining. + */ + def setMaxNbrs(value: Int): this.type = { + maxNbrs = value + this + } + + /** + * Sets whether to use directed edge directions for neighbor sampling. + * + * If true, only outgoing edges are considered for sampling. If false, edges are treated as + * undirected, and both incoming and outgoing directions are included. + * + * @param value + * Boolean indicating whether to consider edge directions. + * @return + * This SamplingConvolution instance for method chaining. + */ + def setUseEdgeDirections(value: Boolean): this.type = { + useEdgeDirections = value + this + } + + /** + * Sets the seed for the random hash function used in min-hash sampling. + * + * @param value + * The long integer seed value for the xxhash64 hash function, ensuring reproducibility of + * sampling. + * @return + * This SamplingConvolution instance for method chaining. + */ + def setSeed(value: Long): this.type = { + seed = value + this + } + + /** + * Sets whether to concatenate the aggregated neighbor features to the original vertex features. + * + * If true, the features column is updated by concatenating the aggregated neighbor embedding to + * the original features. If false, a new column "nbr_embedding" is added containing the + * aggregated neighbor features. + * + * @param value + * Boolean flag for concatenation. + * @return + * This SamplingConvolution instance for method chaining. + */ + def setConcatEmbeddings(value: Boolean): this.type = { + concatEmbeddings = value + this + } + + /** + * Executes the sampling convolution operation and returns the resulting DataFrame. + * + * The output DataFrame includes all original vertex columns. If concatEmbeddings is true, the + * features column is updated with the concatenated [original features, aggregated neighbor + * features]. If false, a new "nbr_embedding" column is added containing the aggregated neighbor + * features as a Vector. + * + * @return + * A DataFrame with the convoluted embeddings, retaining all original vertex columns plus + * modifications. + */ + def run(): DataFrame = { + // Min-Hash sampling based on ID + seed + val preAggs = (if (useEdgeDirections) { + graph.edges + .select(col(GraphFrame.SRC), col(GraphFrame.DST)) + } else { + graph.edges + .select(GraphFrame.SRC, GraphFrame.DST) + .union(graph.edges.select(GraphFrame.DST, GraphFrame.SRC)) + .distinct() + }) + .withColumn("hash", xxhash64(col(GraphFrame.DST), lit(seed))) + .groupBy(col(GraphFrame.SRC).alias(GraphFrame.ID)) + + val vertexDtype = graph.vertices.schema(GraphFrame.ID).dataType + val encoder = KMinSampling.getEncoder( + graph.vertices.sparkSession, + vertexDtype, + Seq(GraphFrame.ID, "hash")) + val samplingUDF = KMinSampling.fromSparkType(vertexDtype, maxNbrs, encoder) + + val sampledNbrs = preAggs + .agg(samplingUDF(col(GraphFrame.DST), col("hash")).alias("nbrs")) + .select(col(GraphFrame.ID), explode(col("nbrs")).alias("nbr")) + + val joined = + sampledNbrs.join( + graph.vertices.select(col(GraphFrame.ID).alias("nbr"), col(featuresCol)), + Seq("nbr"), + "left") + + val foldedNbrsFeatures = + joined.groupBy(GraphFrame.ID).agg(Summarizer.mean(col(featuresCol)).alias("nbr_embedding")) + + val originalAndNbrs = graph.vertices.join(foldedNbrsFeatures, Seq(GraphFrame.ID), "left") + + if (concatEmbeddings) { + originalAndNbrs.withColumn( + featuresCol, + array_to_vector( + concat(vector_to_array(col(featuresCol)), vector_to_array(col("nbr_embedding"))))) + } else { + originalAndNbrs + } + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/Hash2Vec.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/Hash2Vec.scala new file mode 100644 index 0000000000000..e443bec75fd0d --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/Hash2Vec.scala @@ -0,0 +1,601 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.embeddings + +import dev.ludovic.netlib.blas.BLAS +import org.apache.spark.ml.linalg +import org.apache.spark.ml.linalg.SQLDataTypes.VectorType +import org.apache.spark.ml.linalg.Vectors +import org.apache.spark.ml.stat.Summarizer +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.hash +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.functions.pmod +import org.apache.spark.sql.functions.transform +import org.apache.spark.sql.functions.udf +import org.apache.spark.sql.types.ArrayType +import org.apache.spark.sql.types.ByteType +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.types.LongType +import org.apache.spark.sql.types.ShortType +import org.apache.spark.sql.types.StringType +import org.apache.spark.sql.types.StructField +import org.apache.spark.sql.types.StructType +import org.apache.spark.unsafe.hash.Murmur3_x86_32._ +import org.apache.spark.graphframes.GraphFramesUnsupportedVertexTypeException +import org.apache.spark.graphframes.rw.RandomWalkBase + +import scala.reflect.ClassTag +import scala.util.hashing.MurmurHash3 + +/** + * Implementation of Hash2Vec, an efficient word embedding technique using feature hashing. Based + * on: Argerich, Luis, Joaquín Torré Zaffaroni, and Matías J. Cano. "Hash2vec, feature hashing for + * word embeddings." arXiv preprint arXiv:1608.08940 (2016). + * + * Produces embeddings for elements in sequences using a hash-based approach to avoid storing a + * vocabulary. Uses MurmurHash3 for hashing elements to embedding indices and signs. + * + * Output DataFrame has columns "id" (element identifier, same type as sequence elements) and + * "vector" (dense vector of doubles, summed across all occurrences). + * + * Tradeoffs: Higher numPartitions reduces local state and memory per partition but increases + * aggregation and merging overhead across partitions. Larger embeddingsDim provides richer + * representations but consumes more memory. Seeds control hashing for reproducibility. + */ +class Hash2Vec extends Serializable { + private def decayGaussian(d: Int, sigma: Double): Double = { + math.exp(-(d * d) / (sigma * sigma)) + } + private val possibleDecayFunctions: Seq[String] = Seq("gaussian", "constant") + + private var contextSize: Int = 5 + private var numPartitions: Int = 5 + private var embeddingsDim: Int = 512 + private var sequenceCol: String = RandomWalkBase.rwColName + private var decayFunction: String = "gaussian" + private var gaussianSigma: Double = 1.0 + private var hashingSeed: Int = 42 + private var signHashingSeed: Int = 18 + private var doL2Norm: Boolean = true + private var safeL2NormAsChannel: Boolean = true + private var maxVectorsPerPartition: Int = 100000 + + /** + * Sets whether final vectors are L2‑normalized after aggregation across partitions. When + * normalization is enabled, each vector is scaled to unit length (L2 norm = 1). + * + * When `safeNorm` is true (default), the method adds an extra channel to the vector equal to + * `log(L2‑norm + 1) / sqrt(dim)`. This preserves some information about the original magnitude + * while still making vectors comparable via cosine similarity. + * + * When `safeNorm` is false, normalizes without adding an extra channel, discarding magnitude + * entirely. + * + * This setting applies globally to all output vectors. + * + * @param doNorm + * If true, output vectors are normalized. + * @param safeNorm + * If true (and doNorm is true), retains magnitude information in an extra dimension. If + * false, performs standard L2‑normalization. + * @return + * This Hash2Vec instance for method chaining. + */ + def setDoNormalization(doNorm: Boolean, safeNorm: Boolean): this.type = { + doL2Norm = doNorm + safeL2NormAsChannel = safeNorm + this + } + + /** + * Limits the maximum number of distinct element vectors that can be processed inside a single + * partition before flushing intermediate results to the iterator. + * + * Partition processing uses a paged matrix to store vectors. When the number of allocated + * vectors reaches this limit within a partition, the current batch of vectors is returned (as + * part of the iterator) and a new empty batch is started for the remaining elements. + * + * This prevents a single partition from consuming unbounded memory while processing very large + * vocabularies, at the cost of producing multiple iterator batches per partition. + * + * The default value is 100000. + * + * @param value + * Upper bound on distinct vectors processed per batch inside a partition. + * @return + * This Hash2Vec instance for method chaining. + */ + def setMaxVectorsPerPartition(value: Int): this.type = { + maxVectorsPerPartition = value + this + } + + /** + * Convenience overload for `setDoNormalization(doNorm, safeNorm)` that uses safe‑mode (extra + * channel) by default. Equivalent to `setDoNormalization(value, true)`. + * + * @param value + * If true, output vectors are L2‑normalized with safe (extra‑channel) semantics. + * @return + * This Hash2Vec instance for method chaining. + */ + def setDoNormalization(value: Boolean): this.type = { + setDoNormalization(value, true) + this + } + + /** + * Sets the context window size around each element to consider during training. Larger values + * incorporate more distant elements but increase computation time. Default: 5. + */ + def setContextSize(value: Int): this.type = { + contextSize = value + this + } + + /** + * Sets the number of partitions for RDDs to parallelize computation. More partitions distribute + * workload and reduce memory per partition but complicate merging across partitions. Default: + * 5. + */ + def setNumPartitions(value: Int): this.type = { + numPartitions = value + this + } + + /** + * Sets the dimensionality of the dense embedding vectors. Larger dimensions allow richer + * representations but require more memory. Corresponds to the hash table size. Default: 512. + */ + def setEmbeddingsDim(value: Int): this.type = { + embeddingsDim = value + this + } + + /** + * Sets the column name containing sequences of elements (as arrays). Default: "random_walk". + */ + def setSequenceCol(value: String): this.type = { + sequenceCol = value + this + } + + /** + * Sets the decay function used to weight context elements by distance. Supported values: + * "gaussian", "constant". Default: "gaussian". + */ + def setDecayFunction(value: String): this.type = { + val sep = ", " + require( + possibleDecayFunctions.contains(value), + s"supported functions: ${possibleDecayFunctions.mkString(sep)}") + decayFunction = value + this + } + + /** + * Sets the sigma parameter for Gaussian decay weighting. Smaller values decay weights faster + * with distance. Default: 1.0. + */ + def setGaussianSigma(value: Double): this.type = { + gaussianSigma = value + this + } + + /** + * Sets the seed for hashing elements to embedding indices. Used for reproducibility of + * embeddings. Default: 42. + */ + def setHashingSeed(value: Int): this.type = { + hashingSeed = value + this + } + + /** + * Sets the seed for hashing elements to determine the sign of contributions. Used for + * reproducibility of embeddings. Default: 18. + */ + def setSignHashSeed(value: Int): this.type = { + signHashingSeed = value + this + } + + private def nonNegativeMod(x: Int, mod: Int): Int = { + val rawMod = x % mod + rawMod + (if (rawMod < 0) mod else 0) + } + + private var weightFunction: (Int) => Double = _ + + private def seededStringHashFunc(s: String, seed: Int, dim: Int): Int = { + var h = seed + var i = 0 + val len = s.length + while (i < len) { + h = MurmurHash3.mix(h, s.charAt(i).toInt) + i += 1 + } + nonNegativeMod(MurmurHash3.finalizeHash(h, len), dim) + } + + private def seededLongHashFunc(l: Long, seed: Int, dim: Int): Int = + nonNegativeMod(hashLong(l, seed), dim) + + private def normalize(vector: linalg.Vector, addChannel: Boolean): linalg.Vector = { + val blas = BLAS.getInstance() + val arr = vector.toArray + val norm = blas.dnrm2(arr.size, arr, 0, 1) + blas.dscal(arr.size, 1 / (norm + 1e-6), arr, 1) + if (addChannel) { + val scaledL2 = math.log(norm + 1) + val newChannel = scaledL2 / math.sqrt(vector.size.toDouble) + new linalg.DenseVector(arr :+ newChannel) + } else { + new linalg.DenseVector(arr) + } + } + + /** + * Runs the Hash2Vec algorithm on the input DataFrame containing sequences. The specified + * sequenceCol must contain arrays of elements (string or numeric). Produces a DataFrame with + * "id" (element ID, same type as elements) and "vector" (embedding vector, VectorType). + * Embeddings are summed across all partitions and occurrences. + */ + def run(rawData: DataFrame): DataFrame = { + val spark = rawData.sparkSession + require( + rawData.schema(sequenceCol).dataType.isInstanceOf[ArrayType], + "sequence should be array") + val elDataType = rawData.schema(sequenceCol).dataType.asInstanceOf[ArrayType].elementType + + val data = + if (elDataType.isInstanceOf[ByteType] || elDataType.isInstanceOf[ShortType] || elDataType + .isInstanceOf[IntegerType]) { + rawData.withColumn( + sequenceCol, + transform(col(sequenceCol), (f: Column) => f.cast(LongType))) + } else { + rawData + } + + weightFunction = decayFunction match { + case "gaussian" => (d: Int) => decayGaussian(d, gaussianSigma) + case "constant" => (_: Int) => 1.0 + case _ => throw new RuntimeException(s"unsupported decay functions $decayFunction") + } + + val (rowRDD, schema) = elDataType match { + case _: StringType => + ( + runTyped[String](data).map(f => Row(f._1, Vectors.dense(f._2))), + StructType(Seq(StructField("id", StringType), StructField("vector", VectorType)))) + case _: LongType => + ( + runTyped[Long](data).map(f => Row(f._1, Vectors.dense(f._2))), + StructType(Seq(StructField("id", LongType), StructField("vector", VectorType)))) + case _ => + throw new GraphFramesUnsupportedVertexTypeException( + s"Hash2vec supports only string or numeric types of elements but got ${elDataType.toString()}") + } + + val embeddings = spark + .createDataFrame(rowRDD, schema) + .groupBy("id") + .agg(Summarizer.sum(col("vector")).alias("vector")) + + val normalizer = (x: linalg.Vector) => normalize(x, safeL2NormAsChannel) + + if (doL2Norm) { + embeddings.withColumn("vector", udf(normalizer).apply(col("vector"))) + } else { + embeddings + } + } + + private def runTyped[T: ClassTag](data: DataFrame): RDD[(T, Array[Double])] = { + // we should put sequences starts from the same vertex + // to the same partition when possible + data + .withColumn("hash_id", pmod(hash(col(sequenceCol).getItem(0)), lit(numPartitions))) + .repartition(numPartitions, col("hash_id")) + .sortWithinPartitions(col("hash_id")) + .select(col(sequenceCol)) + .rdd + .map(_.getSeq[T](0)) + .mapPartitions { iter => + val elemType = implicitly[ClassTag[T]].runtimeClass + elemType match { + case clazz if clazz == classOf[String] => + processStringPartition(iter.asInstanceOf[Iterator[Seq[String]]]) + .asInstanceOf[Iterator[(T, Array[Double])]] + case clazz if clazz == classOf[Long] => + processLongPartition(iter.asInstanceOf[Iterator[Seq[Long]]]) + .asInstanceOf[Iterator[(T, Array[Double])]] + case _ => + throw new GraphFramesUnsupportedVertexTypeException( + s"Hash2vec does not support type ${elemType.getCanonicalName}") + } + } + } + + private def processStringPartition( + iter: Iterator[Seq[String]]): Iterator[(String, Array[Double])] = { + + val localHashSeed = hashingSeed + val localSignHashSeed = signHashingSeed + val localEmbeddingsDim = embeddingsDim + val localContextSize = contextSize + val localMaxVecrtors = maxVectorsPerPartition + + val weightCache = new Array[Double](localContextSize + 1) + for (d <- 1 to localContextSize) weightCache(d) = weightFunction(d) + val signs = Array[Double](-1.0, 1.0) + + new Iterator[(String, Array[Double])] { + + var currentBatchResult: Iterator[(String, Array[Double])] = Iterator.empty + + var vocabIndex: collection.mutable.HashMap[String, Int] = _ + var matrix: Hash2Vec.PagedMatrixDouble = _ + + override def hasNext: Boolean = { + if (currentBatchResult.hasNext) return true + if (!iter.hasNext) return false + + fetchNextBatch() + currentBatchResult.hasNext + } + + override def next(): (String, Array[Double]) = { + currentBatchResult.next() + } + + private def fetchNextBatch(): Unit = { + vocabIndex = new collection.mutable.HashMap[String, Int]() + vocabIndex.sizeHint(math.min(localMaxVecrtors, 50000)) + + matrix = new Hash2Vec.PagedMatrixDouble(localEmbeddingsDim) + + var currentBatchSize = 0 + + while (iter.hasNext && currentBatchSize < localMaxVecrtors) { + val seq = iter.next() + val currentSeqSize = seq.length + var idx = 0 + + while (idx < currentSeqSize) { + val currentWord = seq(idx) + + var vectorId = vocabIndex.getOrElse(currentWord, -1) + + if (vectorId == -1) { + vectorId = matrix.allocateVector() + vocabIndex.put(currentWord, vectorId) + currentBatchSize += 1 + } + + val start = math.max(0, idx - localContextSize) + val end = math.min(currentSeqSize - 1, idx + localContextSize) + var cIdx = start + + while (cIdx <= end) { + if (cIdx != idx) { + val word = seq(cIdx) + val embeddingIdx = seededStringHashFunc(word, localHashSeed, localEmbeddingsDim) + + val rawSignHash = seededStringHashFunc(word, localSignHashSeed, 65536) + val sign = signs(rawSignHash & 1) + + val weight = weightCache(math.abs(cIdx - idx)) + + matrix.add(vectorId, embeddingIdx, sign * weight) + } + cIdx += 1 + } + idx += 1 + } + } + + currentBatchResult = vocabIndex.iterator.map { case (word, id) => + (word, matrix.getVector(id)) + } + } + } + } + + // Specialized partition processing for Long + private def processLongPartition(iter: Iterator[Seq[Long]]): Iterator[(Long, Array[Double])] = { + + val localHashSeed = hashingSeed + val localSignHashSeed = signHashingSeed + val localEmbeddingsDim = embeddingsDim + val localContextSize = contextSize + val localMaxVectors = maxVectorsPerPartition + + val weightCache = new Array[Double](localContextSize + 1) + for (d <- 1 to localContextSize) weightCache(d) = weightFunction(d) + val signs = Array[Double](-1.0, 1.0) + + new Iterator[(Long, Array[Double])] { + + var currentBatchResult: Iterator[(Long, Array[Double])] = Iterator.empty + + var vocabIndex: collection.mutable.LongMap[Int] = _ + var matrix: Hash2Vec.PagedMatrixDouble = _ + + override def hasNext: Boolean = { + if (currentBatchResult.hasNext) return true + if (!iter.hasNext) return false + + fetchNextBatch() + currentBatchResult.hasNext + } + + override def next(): (Long, Array[Double]) = { + currentBatchResult.next() + } + + private def fetchNextBatch(): Unit = { + vocabIndex = new collection.mutable.LongMap[Int]() + vocabIndex.sizeHint(math.min(localMaxVectors, 100000)) + + matrix = new Hash2Vec.PagedMatrixDouble(localEmbeddingsDim) + + var currentBatchSize = 0 + + while (iter.hasNext && currentBatchSize < localMaxVectors) { + val seq = iter.next() + val currentSeqSize = seq.length + var idx = 0 + + while (idx < currentSeqSize) { + val currentWord = seq(idx) + + var vectorId = vocabIndex.getOrElse(currentWord, -1) + + if (vectorId == -1) { + vectorId = matrix.allocateVector() + vocabIndex.put(currentWord, vectorId) + currentBatchSize += 1 + } + + val start = math.max(0, idx - localContextSize) + val end = math.min(currentSeqSize - 1, idx + localContextSize) + var cIdx = start + + while (cIdx <= end) { + if (cIdx != idx) { + val word = seq(cIdx) + val embeddingIdx = seededLongHashFunc(word, localHashSeed, localEmbeddingsDim) + val sign = signs(seededLongHashFunc(word, localSignHashSeed, 2)) + val weight = weightCache(math.abs(cIdx - idx)) + + matrix.add(vectorId, embeddingIdx, sign * weight) + } + cIdx += 1 + } + idx += 1 + } + } + + currentBatchResult = vocabIndex.iterator.map { case (word, id) => + (word, matrix.getVector(id)) + } + } + } + } +} + +object Hash2Vec { + + /** + * A paged matrix of double-precision vectors that stores vectors contiguously in large + * fixed‑sized pages, each holding PAGE_SIZE (4096) vectors of dimension `dim`. + * + * This layout replaces a HashMap[T, Array[Double]] with two separate structures: + * 1. A mapping from element identifier (T) to a vector ID (Int), maintained by the caller. + * 2. The actual vector data stored in a few large arrays (pages) instead of many small + * per‑element arrays. + * + * Advantages over a HashMap-of-arrays: + * 1. Eliminates per‑vector Array object overhead (object + * header, reference, GC metadata). + * 2. Reduces GC pressure because the backing store is a small number of large long‑lived + * arrays, not many short‑lived small arrays that become garbage as the map is updated. + * 3. Better memory locality: vectors of the same dimension are stored consecutively, + * improving cache line utilisation during sequential access (e.g., inside a page). + * 4. Predictable memory growth: pages are allocated only when the current page is full, + * avoiding repeated resizing of a hash‑map and associated re‑hashing / copying. + * + * The cost is an extra indirection to compute the page index and offset, which is cheap (bit + * shifts and masks) compared to the GC and memory overhead it saves. + * + * Implementation notes: + * 1. PAGE_BITS = 12, PAGE_SIZE = 4096 (2^12). This keeps pageIdx = + * vectorId >>> PAGE_BITS and localRow = vectorId & PAGE_MASK cheap, while limiting page memory + * to PAGE_SIZE * dim doubles. + * 2. The first page is pre‑allocated in the constructor; subsequent + * pages are added on‑demand when allocateVector() crosses a page boundary. + * 3. allocateVector() + * returns a monotonically increasing integer ID, which is the index of the vector across all + * pages. The caller stores this ID in a HashMap[T, Int] instead of storing the whole array. + * 4. add() and getVector() compute the flat index inside the page as localRow * dim + offset. + * 5. Thread safety: not required; each partition processes its own local PagedMatrixDouble + * instance. + */ + private[graphframes] class PagedMatrixDouble(val dim: Int) { + private final val PAGE_BITS = 12 + private final val PAGE_SIZE = 1 << PAGE_BITS // 4096 -> 2^12 + private final val PAGE_MASK = PAGE_SIZE - 1 // 0xFFF + + private val pages = new collection.mutable.ArrayBuffer[Array[Double]]() + private var vectorCount = 0 + + addPage() + + private def addPage(): Unit = { + val size = PAGE_SIZE.toLong * dim + if (size > Int.MaxValue) { + throw new RuntimeException(s"Dimension $dim is too large for current Page Size.") + } + pages += new Array[Double](size.toInt) + } + + /** Allocate a new zero‑initialized vector and return its unique integer ID. */ + def allocateVector(): Int = { + val id = vectorCount + val localIdx = id & PAGE_MASK // ~id % 4096 + + if (localIdx == 0 && id > 0) { + addPage() + } + + vectorCount += 1 + id + } + + /** Accumulate `value` into the component `offset` (0‑based) of vector `vectorId`. */ + @inline + def add(vectorId: Int, offset: Int, value: Double): Unit = { + // vectorId / PAGE_SIZE using unsigned shift (page index) + val pageIdx = vectorId >>> PAGE_BITS + // vectorId % PAGE_SIZE (row inside the page) + val localRow = vectorId & PAGE_MASK + + val idx = (localRow * dim) + offset + pages(pageIdx)(idx) += value + } + + /** Return a fresh copy of the vector identified by `vectorId`. */ + def getVector(vectorId: Int): Array[Double] = { + val pageIdx = vectorId >>> PAGE_BITS + val localRow = vectorId & PAGE_MASK + val page = pages(pageIdx) + + val res = new Array[Double](dim) + val startPos = localRow * dim + System.arraycopy(page, startPos, res, 0, dim) + res + } + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/RandomWalkEmbeddings.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/RandomWalkEmbeddings.scala new file mode 100644 index 0000000000000..6772b8f7cccd8 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/RandomWalkEmbeddings.scala @@ -0,0 +1,384 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.embeddings + +import org.apache.spark.ml.feature.Word2Vec +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.transform +import org.apache.spark.sql.types.StringType +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFramesW2VException +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.convolutions.SamplingConvolution +import org.apache.spark.graphframes.embeddings.RandomWalkEmbeddings.rwModels +import org.apache.spark.graphframes.rw.RandomWalkBase +import org.apache.spark.graphframes.rw.RandomWalkWithRestart + +/** + * RandomWalkEmbeddings is a class for generating node embeddings in a graph using random walks + * and sequence-to-vector models. This implementation supports two types of embedding models: + * Word2Vec and Hash2Vec, each with different performance characteristics. + * + * Word2Vec is based on the skip-gram model, which typically provides higher quality embeddings + * due to its ability to capture semantic relationships through gradient descent optimization. + * However, it is computationally expensive, requires more memory, and scales to approximately 20 + * million vertices in a graph, as it depends on transforming sequences into a vocabulary. + * + * Hash2Vec uses random projection hashing, making it much faster and more memory-efficient, with + * excellent horizontal scaling properties. Its drawbacks include the need for wider embedding + * dimensions (typically 512 or more, depending on graph size) and generally lower quality due to + * its sparse nature. + * + * Additionally, this class supports optional neighbor aggregation, where embeddings from sampled + * neighbors are aggregated (using average) and concatenated with the node's own embedding. This + * technique leverages min-hash sampling and has shown to improve predictive power by over 20% in + * synthetic tests. It is particularly efficient for Hash2Vec, as Word2Vec already incorporates + * neighborhood information through random walks and skip-gram learning. + * + * This class provides also a way to run only embedding model (or sequnce2vec model) on top of + * cached RandomWalks. Users can provide a path to cached walks in parquet format. + * + * To use this class, instantiate with a GraphFrame, set the random walk generator, choose the + * sequence model (Word2Vec or Hash2Vec), and optionally configure other parameters like seed, + * edge direction usage, neighbor aggregation, and maximum neighbors for sampling. + */ +class RandomWalkEmbeddings private[graphframes] (private val graph: GraphFrame) + extends Serializable + with Logging + with WithIntermediateStorageLevel { + private var sequenceModel: Either[Word2Vec, Hash2Vec] = _ + private var randomWalks: RandomWalkBase = _ + private var aggregateNeighbors: Boolean = true + private var useEdgeDirections: Boolean = false + private var maxNbrs: Int = 50 + private var seed: Long = 42L + private var cachedRwPath: Option[String] = None + private var cleanUpAfterRun: Boolean = false + + /** + * Sets the sequence model to use for generating embeddings. This can be either a Word2Vec model + * (Left(Word2Vec)) or a Hash2Vec model (Right(Hash2Vec)). No default; this must be set before + * running. + * @param value + * The sequence model to use. + * @return + * This instance for method chaining. + */ + def setSequenceModel(value: Either[Word2Vec, Hash2Vec]): this.type = { + sequenceModel = value + this + } + + /** + * Sets the random walk generator to use. No default; this must be set before running. + * @param value + * The random walk generator instance. + * @return + * This instance for method chaining. + */ + def setRandomWalks(value: RandomWalkBase): this.type = { + randomWalks = value + this + } + + /** + * Sets the random seed for reproducibility. Default: 42L. + * @param value + * The random seed. + * @return + * This instance for method chaining. + */ + def setSeed(value: Long): this.type = { + seed = value + this + } + + /** + * Sets whether to use edge directions in random walks and neighbor aggregation. If true, + * considers directed edges; otherwise, treats the graph as undirected. Default: false. + * @param value + * Boolean flag for using edge directions. + * @return + * This instance for method chaining. + */ + def setUseEdgeDirections(value: Boolean): this.type = { + useEdgeDirections = value + this + } + + /** + * Sets whether to aggregate neighbor embeddings via min-hash sampling, concatenating the + * aggregated vector with the node's own embedding. This improves predictive power (e.g., +20% + * in tests) and is more efficient for Hash2Vec. For Word2Vec, this adds redundant information + * since it already learns neighborhood relations. Default: true. + * @param value + * Boolean flag for neighbor aggregation. + * @return + * This instance for method chaining. + */ + def setAggregateNeighbors(value: Boolean): this.type = { + aggregateNeighbors = value + this + } + + /** + * Sets the maximum number of neighbors to sample for aggregation. Used only if + * aggregateNeighbors is true. Default: 50. + * @param value + * Maximum neighbors to sample. + * @return + * This instance for method chaining. + */ + def setMaxNbrs(value: Int): this.type = { + maxNbrs = value + this + } + + /** + * Sets the path to the existing cached RandomWalks if you want to run only embeddings model and + * skip the sequences generation step. + * + * @param path + * to walks in parquet format + * @return + * This instance for method chaining. + */ + def useCachedRandomWalks(path: String): this.type = { + cachedRwPath = Some(path) + this + } + + /** + * Sets whether to clean up temporary random walk files after generating embeddings. Default: + * false. + * @param value + * Boolean flag for clean-up. + * @return + * This instance for method chaining. + */ + def setCleanUpAfterRun(value: Boolean): this.type = { + cleanUpAfterRun = value + this + } + + /** + * Executes the random walk embedding generation process. Requires that sequenceModel and + * randomWalks are set. The input GraphFrame must have valid vertex and edge DataFrames, with + * vertices containing an ID column. + * + * The process generates random walks, applies the chosen sequence model to produce initial + * embeddings, and optionally aggregates neighbor embeddings if aggregateNeighbors is enabled. + * + * @return + * A DataFrame containing the original vertex columns plus an additional "embedding" column + * (as defined by RandomWalkEmbeddings.embeddingColName) of type Vector containing the node + * embeddings. If aggregateNeighbors is true, the embedding will be a concatenation of the + * node's embedding and the averaged embeddings of sampled neighbors. + */ + def run(): DataFrame = { + if (rwModels == null) { + throw new GraphFramesW2VException("model should be set!") + } + val spark = graph.vertices.sparkSession + val walksGenerator = randomWalks.onGraph(graph).setUseEdgeDirection(useEdgeDirections) + val walks = if (cachedRwPath.isDefined) { + spark.read.parquet(cachedRwPath.get) + } else { + walksGenerator.run() + } + + val embeddings = sequenceModel match { + case Left(w2v) => { + val model = w2v.setInputCol(RandomWalkBase.rwColName) + val preProcessedSequences = + if (graph.vertices.schema(GraphFrame.ID).dataType != StringType) { + walks.withColumn( + RandomWalkBase.rwColName, + transform(col(RandomWalkBase.rwColName), (c: Column) => c.cast(StringType))) + } else { + walks + } + + val fittedW2V = model.fit(preProcessedSequences) + fittedW2V.getVectors.withColumnsRenamed( + Map("word" -> GraphFrame.ID, "vector" -> RandomWalkEmbeddings.embeddingColName)) + } + case Right(h2v) => { + h2v.run(walks).withColumnRenamed("vector", RandomWalkEmbeddings.embeddingColName) + } + } + + val persistedEmbeddings = if (aggregateNeighbors) { + embeddings.persist(intermediateStorageLevel) + } else { + embeddings + } + + // If requested, do the following: + // - sample neighbors up to maxNbrs + // - compute avg embedding of sample + // - concatenate self and aggregated + val aggregated = if (aggregateNeighbors) { + new SamplingConvolution() + .onGraph(GraphFrame(persistedEmbeddings, graph.edges)) + .setFeaturesCol(RandomWalkEmbeddings.embeddingColName) + .setMaxNbrs(maxNbrs) + .setSeed(seed) + .setUseEdgeDirections(useEdgeDirections) + .setConcatEmbeddings(true) + .run() + } else { + // we need to create a new DataFrame, so unpersisting peristedEmbeddings + // does not accidentally unpersist the result; + // dummy operations are cheap, but will create a different plan. + persistedEmbeddings + .withColumnRenamed(RandomWalkEmbeddings.embeddingColName, "x") + .withColumnRenamed("x", RandomWalkEmbeddings.embeddingColName) + } + + val persistedDF = aggregated.persist(intermediateStorageLevel) + + // materialize + persistedDF.count() + resultIsPersistent() + + // clean memory + persistedEmbeddings.unpersist() + + if (cleanUpAfterRun) { + walksGenerator.cleanUp() + } + + persistedDF + } +} + +/** + * Companion object for RandomWalkEmbeddings. + */ +object RandomWalkEmbeddings extends Serializable { + + /** Name of the embedding column in the output DataFrame. */ + val embeddingColName: String = "embedding" + + private val rwModels = Seq("rw_with_restart") + + /** + * While this API is public, it is not recommended to use it. The only purpose of this API is to + * provide a smooth way to initialize the whole embeddings pipeline with a single method call + * that is usable for Python API (py4j and Spark Connect). + * + * Instead of this API it is recommended to use new + setters of the class! + */ + def pythonAPI( + graph: GraphFrame, + useEdgeDirection: Boolean, + rwModel: String, + rwMaxNbrs: Int, + rwNumWalksPerNode: Int, + rwBatchSize: Int, + rwNumBatches: Int, + rwSeed: Long, + rwRestartProbability: Double, + rwTemporaryPrefix: String, + rwCachedWalks: String, + sequenceModel: String, + hash2vecContextSize: Int, + hash2vecNumPartitions: Int, + hash2vecEmbeddingsDim: Int, + hash2vecDecayFunction: String, + hash2vecGaussianSigma: Double, + hash2vecHashingSeed: Int, + hash2vecSignSeed: Int, + hash2vecDoL2Norm: Boolean, + hash2vecSafeL2: Boolean, + word2vecMaxIter: Int, + word2vecEmbeddingsDim: Int, + word2vecWindowSize: Int, + word2vecNumPartitions: Int, + word2vecMinCount: Int, + word2vecMaxSentenceLength: Int, + word2vecSeed: Long, + word2vecStepSize: Double, + aggregateNeighbors: Boolean, + aggregateNeighborsMaxNbrs: Int, + aggregateNeighborsSeed: Long, + cleanUpAfterRun: Boolean): DataFrame = { + val randomWalksModel: RandomWalkBase = rwModel match { + case "rw_with_restart" => + new RandomWalkWithRestart() + .setRestartProbability(rwRestartProbability) + .setBatchSize(rwBatchSize) + .setNumBatches(rwNumBatches) + .setMaxNbrsPerVertex(rwMaxNbrs) + .setNumWalksPerNode(rwNumWalksPerNode) + .setUseEdgeDirection(useEdgeDirection) + .setTemporaryPrefix(rwTemporaryPrefix) + .setGlobalSeed(rwSeed) + case _: String => + throw new GraphFramesW2VException( + s"unsupported RW $rwModel, supported: ${rwModels.mkString}") + } + + val embeddingsModel = sequenceModel match { + case "hash2vec" => + Right( + new Hash2Vec() + .setContextSize(hash2vecContextSize) + .setDecayFunction(hash2vecDecayFunction) + .setDoNormalization(hash2vecDoL2Norm, hash2vecSafeL2) + .setEmbeddingsDim(hash2vecEmbeddingsDim) + .setGaussianSigma(hash2vecGaussianSigma) + .setHashingSeed(hash2vecHashingSeed) + .setNumPartitions(hash2vecNumPartitions) + .setSignHashSeed(hash2vecSignSeed)) + case "word2vec" => + Left( + new Word2Vec() + .setMaxIter(word2vecMaxIter) + .setMaxSentenceLength(word2vecMaxSentenceLength) + .setMinCount(word2vecMinCount) + .setNumPartitions(word2vecNumPartitions) + .setSeed(word2vecSeed) + .setStepSize(word2vecStepSize) + .setVectorSize(word2vecEmbeddingsDim) + .setWindowSize(word2vecWindowSize)) + case _: String => + throw new GraphFramesW2VException( + s"unsupported sequence model $sequenceModel, supported are 'word2vec' and 'hash2vec'") + } + + val embeddingsGenerator = new RandomWalkEmbeddings(graph) + .setRandomWalks(randomWalksModel) + .setSequenceModel(embeddingsModel) + .setAggregateNeighbors(aggregateNeighbors) + .setMaxNbrs(aggregateNeighborsMaxNbrs) + .setUseEdgeDirections(useEdgeDirection) + .setSeed(aggregateNeighborsSeed) + .setCleanUpAfterRun(cleanUpAfterRun) + + if (rwCachedWalks == "") { + embeddingsGenerator.run() + } else { + embeddingsGenerator.useCachedRandomWalks(rwCachedWalks).run() + } + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/exceptions.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/exceptions.scala new file mode 100644 index 0000000000000..21f4134dd7bfe --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/exceptions.scala @@ -0,0 +1,76 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes + +// All the public exceptions thrown by GraphFrame methods + +/** + * Exception thrown when a pattern String for motif finding cannot be parsed. + */ +class InvalidParseException(message: String) extends Exception(message) + +/** + * Thrown when a GraphFrame algorithm is given a vertex ID which does not exist in the graph. + */ +class NoSuchVertexException(message: String) extends Exception(message) + +/** + * Exception thrown when a parsed pattern for motif finding cannot be translated into a DataFrame + * query. + */ +class InvalidPatternException() extends Exception() + +/** + * Exception that should not be reachable + */ +class GraphFramesUnreachableException() + extends Exception("This exception should not be reachable") + +/** + * Exception thrown when an invalid property group is encountered. + * + * This exception typically indicates that an operation or configuration is using a property group + * that is not supported, invalid, or improperly defined. + * + * @param message + * A detailed error message describing the issue. + */ +class InvalidPropertyGroupException(message: String) extends Exception(message) + +/** + * Exception thrown when the graph is invalid, e.g. duplicate vertices, inconsistency between + * vertex set and edges src / dst, etc. + * + * @param message + * A descriptive error message providing details about why the graph operation is invalid. + */ +class InvalidGraphException(message: String) extends Exception(message) + +class GraphFramesW2VException(message: String) extends Exception(message) + +class GraphFramesUnsupportedVertexTypeException(message: String) extends Exception(message) + +/** + * Exception thrown when a Spark version requirement is not met. + * + * @param version + * The minimum version of Apache Spark required. + */ +class GraphFramesSparkVersionException(version: String) + extends Exception( + s"Called GraphFrames feature requires at least $version or above version of Apache Spark") diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateMessages.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateMessages.scala new file mode 100644 index 0000000000000..fb2be56278fa9 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateMessages.scala @@ -0,0 +1,207 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.expr +import org.apache.spark.sql.functions.struct +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithIntermediateStorageLevel + +/** + * This is a primitive for implementing graph algorithms. This method aggregates messages from the + * neighboring edges and vertices of each vertex. + * + * For each triplet (source vertex, edge, destination vertex) in [[GraphFrame.triplets]], this can + * send a message to the source and/or destination vertices. + * - `AggregateMessages.sendToSrc()` sends a message to the source vertex of each triplet + * - `AggregateMessages.sendToDst()` sends a message to the destination vertex of each triplet + * - `AggregateMessages.agg` specifies an aggregation function for aggregating the messages sent + * to each vertex. It also runs the aggregation, computing a DataFrame with one row for each + * vertex which receives > 0 messages. The DataFrame has 2 columns: + * - vertex column ID (named [[GraphFrame.ID]]) + * - aggregate from messages sent to vertex (with the name given to the `Column` specified in + * `AggregateMessages.agg()`) + * + * When specifying the messages and aggregation function, the user may reference columns using: + * - [[AggregateMessages.src]]: column for source vertex of edge + * - [[AggregateMessages.edge]]: column for edge + * - [[AggregateMessages.dst]]: column for destination vertex of edge + * - [[AggregateMessages.msg]]: message sent to vertex (for aggregation function) + * + * Note: If you use this operation to write an iterative algorithm, you may want to use + * `checkpoint()` (`localCheckpoint()`) as a workaround for caching issues. + * + * @example + * We can use this function to compute the in-degree of each vertex + * {{{ + * val g: GraphFrame = Graph.textFile("twittergraph") + * val inDeg: DataFrame = + * g.aggregateMessages().sendToDst(lit(1)).agg(sum(AggregateMessagesBuilder.msg)) + * }}} + */ +class AggregateMessages private[graphframes] (private val g: GraphFrame) + extends Arguments + with Serializable + with WithIntermediateStorageLevel + with Logging { + + import org.apache.spark.graphframes.GraphFrame.DST + import org.apache.spark.graphframes.GraphFrame.ID + import org.apache.spark.graphframes.GraphFrame.SRC + + private var msgToSrc: Seq[Column] = Vector() + + /** Send message to source vertex */ + def sendToSrc(value: Column, values: Column*): this.type = { + msgToSrc = value +: values + this + } + + // Python API compatibility + def sendToSrc(value: Column): this.type = { + msgToSrc :+= value + this + } + def sendToSrc(value: String): this.type = { + msgToSrc :+= expr(value) + this + } + + /** Send message to source vertex, specifying SQL expression as a String */ + def sendToSrc(value: String, values: String*): this.type = + sendToSrc(expr(value), values.map(expr): _*) + + private var msgToDst: Seq[Column] = Vector() + + /** Send message to destination vertex */ + def sendToDst(value: Column, values: Column*): this.type = { + msgToDst = value +: values + this + } + + // Python API compatibility + def sendToDst(value: String): this.type = { + msgToDst :+= expr(value) + this + } + def sendToDst(value: Column): this.type = { + msgToDst :+= value + this + } + + /** Send message to destination vertex, specifying SQL expression as a String */ + def sendToDst(value: String, values: String*): this.type = + sendToDst(expr(value), values.map(expr): _*) + + /** + * Run the aggregation, returning the resulting DataFrame of aggregated messages. This is a lazy + * operation, so the DataFrame will not be materialized until an action is executed on it. + * + * This returns a DataFrame with schema: + * - column "id": vertex ID + * - aggCol: aggregate result + * - aggCols: one column with the result of each additional defined aggregation + * If you need to join this with the original [[GraphFrame.vertices]], you can run an inner join + * of the form: + * {{{ + * val g: GraphFrame = ... + * val aggResult = g.AggregateMessagesBuilder.sendToSrc(msg).agg(aggFunc) + * aggResult.join(g.vertices, ID) + * }}} + */ + def agg(aggCol: Column, aggCols: Column*): DataFrame = { + require( + msgToSrc.nonEmpty || msgToDst.nonEmpty, + "To run GraphFrame.aggregateMessages," + + " messages must be sent to src, dst, or both. Set using sendToSrc(), sendToDst().") + val triplets = g.triplets + + def msgColumn(columns: Seq[Column], idColumn: Column): DataFrame = columns match { + case Seq(c) => triplets.select(idColumn.as(ID), c.as(AggregateMessages.MSG_COL_NAME)) + case columns => + triplets.select(idColumn.as(ID), struct(columns: _*).as(AggregateMessages.MSG_COL_NAME)) + } + + val cachedVertices = g.vertices.persist(intermediateStorageLevel) + + val sentMsgsToSrc = msgToSrc.headOption.map { _ => + val msgsToSrc = msgColumn(msgToSrc, triplets(SRC)(ID)) + msgsToSrc + .join(cachedVertices, ID) + .select(msgsToSrc(AggregateMessages.MSG_COL_NAME), col(ID)) + } + val sentMsgsToDst = msgToDst.headOption.map { _ => + val msgsToDst = msgColumn(msgToDst, triplets(DST)(ID)) + + msgsToDst + .join(cachedVertices, ID) + .select(msgsToDst(AggregateMessages.MSG_COL_NAME), col(ID)) + } + val unionMsgs = (sentMsgsToSrc, sentMsgsToDst) match { + case (Some(toSrc), Some(toDst)) => + toSrc.unionAll(toDst) + case (Some(toSrc), None) => toSrc + case (None, Some(toDst)) => toDst + case _ => + // Should never happen. Specify this case to avoid compilation warnings. + throw new RuntimeException("AggregateMessages: No messages were specified to be sent.") + } + + val cachedResult = + unionMsgs.groupBy(ID).agg(aggCol, aggCols: _*).persist(intermediateStorageLevel) + // materialize + cachedResult.count() + cachedVertices.unpersist() + resultIsPersistent() + cachedResult + } + + // Python compatibility + def agg(aggCol: Column): DataFrame = agg(aggCol, Seq.empty[Column]: _*) + def agg(aggCol: String): DataFrame = agg(expr(aggCol), Seq.empty[Column]: _*) + + /** + * Run the aggregation, specifying SQL expression as a String + * + * See the overloaded method documentation for more details. + */ + def agg(aggCol: String, aggCols: String*): DataFrame = + agg(expr(aggCol), aggCols.map(expr(_)): _*) +} + +object AggregateMessages extends Logging with Serializable { + + /** Column name for aggregated messages, used in [[AggregateMessages.msg]] */ + val MSG_COL_NAME: String = "MSG" + + /** Reference for source column, used for specifying messages */ + def src: Column = col(GraphFrame.SRC) + + /** Reference for destination column, used for specifying messages */ + def dst: Column = col(GraphFrame.DST) + + /** Reference for edge column, used for specifying messages */ + def edge: Column = col(GraphFrame.EDGE) + + /** Reference for message column, used for specifying aggregation function */ + def msg: Column = col(MSG_COL_NAME) +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateNeighbors.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateNeighbors.scala new file mode 100644 index 0000000000000..ddf21ceb541f0 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateNeighbors.scala @@ -0,0 +1,459 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithCheckpointInterval +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLocalCheckpoints + +/** + * A class for performing multi-hop neighbor aggregation on a graph. + * + * AggregateNeighbors allows you to explore the graph up to a specified number of hops, + * accumulating values along paths using customizable accumulator expressions. It supports both + * stopping conditions (when to stop exploring) and target conditions (when to collect a result). + * If no target condition is provided, collect all traversals that reached stopping condition. The + * algorithm processes the graph in a breadth‑first manner. + * + * Use the builder pattern to configure parameters, then call `run()` to execute. + * + * @param graph + * the GraphFrame on which to perform aggregation + */ +class AggregateNeighbors private[graphframes] (graph: GraphFrame) + extends Serializable + with Logging + with WithIntermediateStorageLevel + with WithCheckpointInterval + with WithLocalCheckpoints { + + import AggregateNeighbors._ + + private var startingVertices: Column = lit(true) + + private var maxHops: Int = 3 + private var stoppingCondition: Option[Column] = None + private var targetCondition: Option[Column] = None + private var accumulatorsNames: Seq[String] = Seq.empty + private var accumulatorsInits: Seq[Column] = Seq.empty + private var accumulatorsUpdates: Seq[Column] = Seq.empty + + private var requiredVertexAttributes: Seq[String] = Seq.empty + private var requiredEdgeAttributes: Seq[String] = Seq.empty + + private var edgeFilter: Column = lit(true) + private var removeLoops: Boolean = false + + /** + * Specifies which vertices to start the aggregation from. + * + * Only vertices satisfying the given column expression will be used as seeds for the traversal. + * The default is `true` (all vertices are seeds). + * + * @param value + * a Boolean Column expression that selects seed vertices + * @return + * this AggregateNeighbors instance for method chaining + */ + def setStartingVertices(value: Column): this.type = { + this.startingVertices = value + this + } + + /** + * Sets the maximum number of hops to explore from the starting vertices. + * + * The algorithm will stop after this many iterations even if the stopping condition hasn't been + * reached for all paths. The default is 3. Be aware that high value of max hops on a dense + * graph will lead to a huge memory load and potential OOM errors. + * + * @param value + * positive integer for maximum hop count + * @return + * this AggregateNeighbors instance for method chaining + */ + def setMaxHops(value: Int): this.type = { + require(value > 0, "maxHops should be positive.") + this.maxHops = value + this + } + + /** + * Sets a condition that, when true for a vertex, stops further exploration along that path. + * + * The column expression can refer to: + * - accumulator columns + * - source vertex attributes via `srcAttr("attrName")` + * - destination vertex attributes via `dstAttr("attrName")` + * - edge attributes via `edgeAttr("attrName")` + * + * Either `setStoppingCondition` or `setTargetCondition` must be called. When both are provided + * only accumulators that reach `targetCondition` are saved to results. + * + * @param value + * a Boolean Column expression that triggers stopping + * @return + * this AggregateNeighbors instance for method chaining + */ + def setStoppingCondition(value: Column): this.type = { + this.stoppingCondition = Some(value) + this + } + + /** + * Sets a condition that, when true for a vertex, marks it as a target and saves the accumulator + * values. + * + * Either `setStoppingCondition` or `setTargetCondition` must be called. When both are provided + * only accumulators that reach `targetCondition` are saved to results. + * + * @param value + * a Boolean Column expression that defines a target vertex + * @return + * this AggregateNeighbors instance for method chaining + */ + def setTargetCondition(value: Column): this.type = { + this.targetCondition = Some(value) + this + } + + /** + * Defines multiple accumulators used to compute values along each path. + * + * Each accumulator has a name, an initial value expression, and an update expression. The + * update expression is evaluated at each hop and can refer to: + * - the accumulator's previous value (using the accumulator's column name) + * - source vertex attributes via `srcAttr("attrName")` + * - destination vertex attributes via `dstAttr("attrName")` + * - edge attributes via `edgeAttr("attrName")` + * + * @param names + * sequence of accumulator column names + * @param inits + * sequence of Column expressions that give the initial value for each accumulator (evaluated + * on starting vertices) + * @param updates + * sequence of Column expressions that define how to update the accumulator when traversing an + * edge + * @return + * this AggregateNeighbors instance for method chaining + */ + def setAccumulators(names: Seq[String], inits: Seq[Column], updates: Seq[Column]): this.type = { + require( + inits.size == updates.size && updates.size == names.size, + "Inits, updates and names must have the same size.") + this.accumulatorsNames = names + this.accumulatorsInits = inits + this.accumulatorsUpdates = updates + this + } + + /** + * Adds a single accumulator to those already configured. + * + * This method can be called multiple times to add several accumulators. + * + * @param name + * accumulator column name + * @param init + * Column expression for the accumulator's initial value + * @param update + * Column expression that updates the accumulator when traversing an edge + * @return + * this AggregateNeighbors instance for method chaining + */ + def addAccumulator(name: String, init: Column, update: Column): this.type = { + this.accumulatorsNames :+= name + this.accumulatorsInits :+= init + this.accumulatorsUpdates :+= update + this + } + + /** + * Specifies which vertex attributes should be carried through the traversal. + * + * By default, all vertex columns are carried. Specifying a subset can improve performance by + * reducing the amount of data shuffled. + * + * @param values + * sequence of vertex column names to keep + * @return + * this AggregateNeighbors instance for method chaining + */ + def setRequiredVertexAttributes(values: Seq[String]): this.type = { + this.requiredVertexAttributes = values + this + } + + /** + * Specifies which edge attributes should be carried through the traversal. + * + * By default, all edge columns are carried. Specifying a subset can improve performance by + * reducing the amount of data shuffled. + * + * @param values + * sequence of edge column names to keep + * @return + * this AggregateNeighbors instance for method chaining + */ + def setRequiredEdgeAttributes(values: Seq[String]): this.type = { + this.requiredEdgeAttributes = values + this + } + + /** + * Filters which edges can be traversed during the aggregation. + * + * Only edges satisfying the given column expression are considered. The default is `true` (all + * edges are traversable). + * + * @param value + * a Boolean Column expression that filters edges + * @return + * this AggregateNeighbors instance for method chaining + */ + def setEdgeFilter(value: Column): this.type = { + this.edgeFilter = value + this + } + + /** + * Controls whether self‑loops (edges where src == dst) are excluded. + * + * @param value + * if true, self‑loop edges are filtered out + * @return + * this AggregateNeighbors instance for method chaining + */ + def setRemoveLoops(value: Boolean): this.type = { + this.removeLoops = value + this + } + + /** + * Executes the configured neighbor aggregation and returns the result DataFrame. + * + * The result contains one row per (starting vertex, target vertex) pair that satisfied either + * the stopping condition or the target condition. Columns are: + * - `id`: the vertex ID of the target (or stopping) vertex + * - `hop`: the number of hops taken to reach it + * - one column for each accumulator, holding its final value + * + * @return + * DataFrame with aggregation results + */ + def run(): DataFrame = { + require(maxHops > 0, "maxHops must be greater than 0") + if (maxHops > 10) + logWarn(s"maxHops is very large ($maxHops). This might be performance-intensive.") + require(accumulatorsNames.nonEmpty, "At least one accumulator must be added") + require( + stoppingCondition.orElse(targetCondition).isDefined, + "Any of target or stopping conditions should be provided") + + val reqAttrs = (if (requiredVertexAttributes.isEmpty) { + graph.vertices.columns.toSeq + } else { + requiredVertexAttributes + }).map(col(_)) + + val reqEdgeAttr = (if (requiredEdgeAttributes.isEmpty) { + graph.edges.columns.toSeq + } else { + requiredEdgeAttributes + }).map(col(_)) + + val verticesWithAttributes = graph.vertices.select( + col(GraphFrame.ID).alias("dst_id"), + struct(reqAttrs: _*).alias(dstAttributes)) + + // "right" side of the join for each iteration + val edgesBase = graph.edges + val edgesFiltered = if (removeLoops) { + edgesBase.filter(col(GraphFrame.SRC) =!= col(GraphFrame.DST)) + } else { + edgesBase + } + val semiTriplets = edgesFiltered + .select( + col(GraphFrame.SRC), + col(GraphFrame.DST), + struct(reqEdgeAttr: _*).alias(edgeAttributes)) + .join(verticesWithAttributes, col("dst_id") === col(GraphFrame.DST), "left") + .repartition(col(GraphFrame.SRC)) // to avoid shuffle; + .persist(intermediateStorageLevel) + + // memory-tracking + val persistenceQueue = collection.mutable.Queue.empty[DataFrame] + + val statesColumns = + (accumulatorsNames ++ Seq( + srcAttributes, + "src_id", + currentPathLenColName, + stoppingCondColName)).map(col(_)) + val finishedColumns = + (accumulatorsNames ++ Seq("src_id", currentPathLenColName)).map(col(_)) + + // holder of the current state of accumulators + var states: DataFrame = graph.vertices + .filter(startingVertices) + .withColumns(accumulatorsNames.zip(accumulatorsInits).toMap) + .withColumn(srcAttributes, struct(reqAttrs: _*)) + .withColumnRenamed(GraphFrame.ID, "src_id") + .withColumn(currentPathLenColName, lit(0)) + .withColumn(stoppingCondColName, lit(false)) + .select(statesColumns: _*) + .persist(intermediateStorageLevel) + + // holder of the finished accumulators + var finished: DataFrame = states + .filter(col(stoppingCondColName)) + .select(finishedColumns: _*) + .withColumnRenamed("src_id", GraphFrame.ID) + .persist(intermediateStorageLevel) + + var collected = finished.count() + + persistenceQueue.enqueue(states) + persistenceQueue.enqueue(finished) + + var converged = states.isEmpty + var iter = 0 + + while ((!converged) && (iter < maxHops)) { + iter += 1 + // get full triplets by joining states (frontier) with semiTriplets + val fullTriplets = + states.join(semiTriplets, col("src_id") === col(GraphFrame.SRC)).filter(edgeFilter) + + var colsToSelect = accumulatorsUpdates + .zip(accumulatorsNames) + .map(r => r._1.alias(r._2)) + .toSeq + + // Build expressions for stopping and targeting using single Column conditions + val isTargetExpr = targetCondition.getOrElse(lit(false)) + // We are stopping if any of them are reached + val shouldStopExpr = stoppingCondition.getOrElse(lit(false)) || isTargetExpr + + colsToSelect = colsToSelect :+ shouldStopExpr.alias(stoppingCondColName) + colsToSelect = colsToSelect :+ isTargetExpr.alias("_is_target") + colsToSelect = colsToSelect :+ lit(iter).alias(currentPathLenColName) :+ + col(GraphFrame.DST).alias("src_id") :+ + col(dstAttributes).alias(srcAttributes) + val updatedStates = fullTriplets.select(colsToSelect: _*) + + var newStates = updatedStates.filter(!col(stoppingCondColName)).select(statesColumns: _*) + var newFinished = if (targetCondition.isDefined) { + finished.unionByName( + updatedStates + .filter(col("_is_target")) + .select(finishedColumns: _*) + .withColumnRenamed("src_id", GraphFrame.ID)) + } else { + finished.unionByName( + updatedStates + .filter(col(stoppingCondColName)) + .select(finishedColumns: _*) + .withColumnRenamed("src_id", GraphFrame.ID)) + } + + if ((checkpointInterval > 0) && (iter % checkpointInterval == 0)) { + if (useLocalCheckpoints) { + newStates = newStates.localCheckpoint() + newFinished = newFinished.localCheckpoint() + } else { + newStates = newStates.checkpoint() + newFinished = newFinished.checkpoint() + } + } + + newStates = newStates.persist(intermediateStorageLevel) + newFinished = newFinished.persist(intermediateStorageLevel) + + persistenceQueue.enqueue(newStates) + persistenceQueue.enqueue(newFinished) + + // materialize to unpersist + collected = newFinished.count() + converged = newStates.isEmpty + + // unpersist + persistenceQueue.dequeue().unpersist(true) + persistenceQueue.dequeue().unpersist(true) + + logInfo(s"iteration $iter, collected $collected rows") + + states = newStates + finished = newFinished + } + + persistenceQueue.dequeue().unpersist(true) // clear states + semiTriplets.unpersist(true) // clear triplets + + resultIsPersistent() + finished + } +} + +object AggregateNeighbors extends Serializable { + private val stoppingCondColName: String = "_stopped" + private val currentPathLenColName: String = "hop" + private val srcAttributes: String = "src_attributes" + private val dstAttributes: String = "dst_attributes" + private val edgeAttributes: String = "edge_attributes" + + /** + * Creates a column that references a source vertex attribute within accumulator update + * expressions, stopping conditions, or target conditions. + * + * @param name + * the name of the source vertex attribute + * @return + * a Column referencing the attribute + */ + def srcAttr(name: String): Column = col(srcAttributes).getField(name) + + /** + * Creates a column that references a destination vertex attribute within accumulator update + * expressions, stopping conditions, or target conditions. + * + * @param name + * the name of the destination vertex attribute + * @return + * a Column referencing the attribute + */ + def dstAttr(name: String): Column = col(dstAttributes).getField(name) + + /** + * Creates a column that references an edge attribute within accumulator update expressions, + * stopping conditions, or target conditions. + * + * @param name + * the name of the edge attribute + * @return + * a Column referencing the attribute + */ + def edgeAttr(name: String): Column = col(edgeAttributes).getField(name) +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AllPaths.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AllPaths.scala new file mode 100644 index 0000000000000..784438440201d --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AllPaths.scala @@ -0,0 +1,209 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.array +import org.apache.spark.sql.functions.array_contains +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.concat +import org.apache.spark.sql.functions.expr +import org.apache.spark.sql.graphframes.GraphFrameInternals +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.WithCheckpointInterval +import org.apache.spark.graphframes.WithDirection +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLocalCheckpoints + +/** + * Computes all simple paths between source and destination vertices. + * + * This algorithm enumerates paths up to `maxPathLength` hops. It supports directed and undirected + * traversal as well as optional edge filtering. It returns all simple paths between source and + * destination vertices. Here the term "simple" means no repeated vertices. For example, if there + * are paths A-B-C, A-D-C and the edge B-A, user asked to find all the paths between "A" and "C" + * only A-B-C and A-D-C will be returned, but not the A-B-A-D-C. The default value of the + * `maxPathLength` is `5`. Keep in mind that requesting `maxPathLength` of the scale of the graph + * diameter may tend this algorithm will try to return (almost) all simple paths in the graph that + * can create huge performance degradation or even OOM-like errors. Algorithm supports both + * directed and undirected graphs. + * + * Returned DataFrame schema: + * - `path`: array of vertex ids in traversal order + * - `len`: number of edges in the path (Long) + * + * Note: in the case of undirected graph an algorithm run on the internal graph made by union + * edges and reversed edges. It is assumed that graph does not have multi-edges. Results may be + * unstable and unpredictable for the graph with multi-edges. + */ +class AllPaths private[graphframes] (private val graph: GraphFrame) + extends Arguments + with Serializable + with WithDirection + with WithLocalCheckpoints + with WithCheckpointInterval + with WithIntermediateStorageLevel { + + private var maxPathLength: Int = 5 + private var fromExpression: Column = _ + private var toExpression: Column = _ + private var edgeFilterExpression: Option[Column] = None + + /** + * Sets the expression identifying the source (starting) vertices. + * + * @param value + * a Column expression evaluated against vertex attributes to select source vertices + * @return + * this instance for method chaining + */ + def fromExpr(value: Column): this.type = { + fromExpression = value + this + } + + /** + * Sets the expression identifying the source (starting) vertices. + * + * @param value + * a SQL expression string evaluated against vertex attributes to select source vertices + * @return + * this instance for method chaining + */ + def fromExpr(value: String): this.type = fromExpr(expr(value)) + + /** + * Sets the expression identifying the destination (target) vertices. + * + * @param value + * a Column expression evaluated against vertex attributes to select destination vertices + * @return + * this instance for method chaining + */ + def toExpr(value: Column): this.type = { + toExpression = value + this + } + + /** + * Sets the expression identifying the destination (target) vertices. + * + * @param value + * a SQL expression string evaluated against vertex attributes to select destination vertices + * @return + * this instance for method chaining + */ + def toExpr(value: String): this.type = toExpr(expr(value)) + + /** + * Sets the maximum path length (number of edges) for the enumerated paths. + * + * Setting a large value (e.g. on the scale of the graph diameter) may cause the algorithm to + * attempt to collect a very large number of paths, leading to severe performance degradation or + * out-of-memory errors. Use with caution on large or densely connected graphs. + * + * @param value + * the maximum number of edges in a path; must be greater than 0. Default is 5. + * @return + * this instance for method chaining + */ + def maxPathLength(value: Int): this.type = { + require(value > 0, s"AllPaths maxPathLength must be > 0, but was set to $value") + maxPathLength = value + this + } + + /** + * Sets an optional filter expression applied to edges during traversal. Only edges satisfying + * this condition will be considered. + * + * @param value + * a Column expression evaluated against edge attributes + * @return + * this instance for method chaining + */ + def edgeFilter(value: Column): this.type = { + edgeFilterExpression = Some(value) + this + } + + /** + * Sets an optional filter expression applied to edges during traversal. Only edges satisfying + * this condition will be considered. + * + * @param value + * a SQL expression string evaluated against edge attributes + * @return + * this instance for method chaining + */ + def edgeFilter(value: String): this.type = edgeFilter(expr(value)) + + /** + * Executes the AllPaths algorithm and returns all simple paths between the specified source and + * destination vertices. + * + * @return + * a DataFrame with the following columns: + * - `path`: an array of vertex ids in traversal order + * - `len`: the number of edges in the path (Long) + */ + def run(): DataFrame = { + require(fromExpression != null, "fromExpr is required.") + require(toExpression != null, "toExpr is required.") + require( + graph.vertices.columns.toSet.intersect(Set("hop", "path", "len")).isEmpty, + "columns `hop`, `path` and `len` are reserved by algorithm") + + val traversalGraph = if (isDirected) { + graph + } else { + val edgeColumns = graph.edges.columns.toSeq + val reversed = graph.edges.select( + (Seq( + col(GraphFrame.DST).alias(GraphFrame.SRC), + col(GraphFrame.SRC).alias(GraphFrame.DST)) ++ + edgeColumns.filterNot(c => c == GraphFrame.SRC || c == GraphFrame.DST).map(col)): _*) + GraphFrame(graph.vertices, graph.edges.unionByName(reversed)) + } + + val agg = traversalGraph.aggregateNeighbors + .setStartingVertices(fromExpression) + .setMaxHops(maxPathLength) + .setTargetCondition( + GraphFrameInternals.applyExprToCol(graph.spark, toExpression, "dst_attributes")) + .setStoppingCondition( + array_contains(col("path"), AggregateNeighbors.dstAttr(GraphFrame.ID))) + .addAccumulator( + "path", + array(col(GraphFrame.ID)), + concat(col("path"), array(AggregateNeighbors.dstAttr(GraphFrame.ID)))) + .setUseLocalCheckpoints(useLocalCheckpoints) + .setCheckpointInterval(checkpointInterval) + .setIntermediateStorageLevel(intermediateStorageLevel) + + edgeFilterExpression.foreach { ef => + agg.setEdgeFilter(GraphFrameInternals.applyExprToCol(graph.spark, ef, "edge_attributes")) + } + + agg + .run() + .select(col("path"), col("hop").alias("len")) + .distinct() + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/BFS.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/BFS.scala new file mode 100644 index 0000000000000..1eefebe749bd3 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/BFS.scala @@ -0,0 +1,231 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.expr +import org.apache.spark.sql.graphframes.GraphFrameInternals +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrame.nestAsCol +import org.apache.spark.graphframes.Logging + +/** + * Breadth-first search (BFS) + * + * This method returns a DataFrame of valid shortest paths from vertices matching `fromExpr` to + * vertices matching `toExpr`. If multiple paths are valid and have the same length, the DataFrame + * will return one Row for each path. If no paths are valid, the DataFrame will be empty. Note: + * "Shortest" means globally shortest path. I.e., if the shortest path between two vertices + * matching `fromExpr` and `toExpr` is length 5 (edges) but no path is shorter than 5, then all + * paths returned by BFS will have length 5. + * + * The returned DataFrame will have the following columns: + * - `from` start vertex of path + * - `e[i]` edge i in the path, indexed from 0 + * - `v[i]` intermediate vertex i in the path, indexed from 1 + * - `to` end vertex of path + * Each of these columns is a StructType whose fields are the same as the columns of + * [[GraphFrame.vertices]] or [[GraphFrame.edges]]. + * + * For example, suppose we have a graph g. Say the vertices DataFrame of g has columns "id" and + * "job", and the edges DataFrame of g has columns "src", "dst", and "relation". + * {{{ + * // Search from vertex "Joe" to find the closet vertices with attribute job = CEO. + * g.bfs(col("id") === "Joe", col("job") === "CEO").run() + * }}} + * If we found a path of 3 edges, each row would have columns: + * {{{from | e0 | v1 | e1 | v2 | e2 | to}}} In the above row, each vertex column (from, v1, v2, + * to) would have fields "id" and "job" (just like g.vertices). Each edge column (e0, e1, e2) + * would have fields "src", "dst", and "relation". + * + * If there are ties, then each of the equal paths will be returned as a separate Row. + * + * If one or more vertices match both the from and to conditions, then there is a 0-hop path. The + * returned DataFrame will have the "from" and "to" columns (as above); however, the "from" and + * "to" columns will be exactly the same. There will be one row for each vertex in + * [[GraphFrame.vertices]] matching both `fromExpr` and `toExpr`. + * + * Parameters: + * + * - `fromExpr` Spark SQL expression specifying valid starting vertices for the BFS. This + * condition will be matched against each vertex's id or attributes. To start from a specific + * vertex, this could be "id = [start vertex id]". To start from multiple valid vertices, this + * can operate on vertex attributes. + * - `toExpr` Spark SQL expression specifying valid target vertices for the BFS. This condition + * will be matched against each vertex's id or attributes. + * - `maxPathLength` Limit on the length of paths. If no valid paths of length <= maxPathLength + * are found, then the BFS is terminated. (default = 10) + * - `edgeFilter` Spark SQL expression specifying edges which may be used in the search. This + * allows the user to disallow crossing certain edges. Such filters can be applied post-hoc + * after BFS, run specifying the filter here is more efficient. + * + * Returns: + * - DataFrame of valid shortest paths found in the BFS + */ +class BFS private[graphframes] (private val graph: GraphFrame) + extends Arguments + with Serializable { + + private var maxPathLength: Int = 10 + private var edgeFilter: Option[Column] = None + private var fromExpr: Column = _ + private var toExpr: Column = _ + + def fromExpr(value: Column): this.type = { + fromExpr = value + this + } + + def fromExpr(value: String): this.type = fromExpr(expr(value)) + + def toExpr(value: Column): this.type = { + toExpr = value + this + } + + def toExpr(value: String): this.type = toExpr(expr(value)) + + def maxPathLength(value: Int): this.type = { + require(value >= 0, s"BFS maxPathLength must be >= 0, but was set to $value") + maxPathLength = value + this + } + + def edgeFilter(value: Column): this.type = { + edgeFilter = Some(value) + this + } + + def edgeFilter(value: String): this.type = edgeFilter(expr(value)) + + def run(): DataFrame = { + require(fromExpr != null, "fromExpr is required.") + require(toExpr != null, "toExpr is required.") + BFS.run(graph, fromExpr, toExpr, maxPathLength, edgeFilter) + } +} + +private object BFS extends Logging with Serializable { + + private def run( + g: GraphFrame, + from: Column, + to: Column, + maxPathLength: Int, + edgeFilter: Option[Column]): DataFrame = { + val fromDF = g.vertices.filter(from) + val toDF = g.vertices.filter(to) + if (fromDF.take(1).isEmpty || toDF.take(1).isEmpty) { + // Return empty DataFrame + return g.spark.createDataFrame( + g.spark.sparkContext.parallelize(Seq.empty[Row]), + g.vertices.schema) + } + + val fromEqualsToDF = fromDF.filter(to) + if (fromEqualsToDF.take(1).nonEmpty) { + // from == to, so return matching vertices + return fromEqualsToDF.select( + nestAsCol(fromEqualsToDF, "from"), + nestAsCol(fromEqualsToDF, "to")) + } + + // We handled edge cases above, so now we do BFS. + + // Edges a->b, to be reused for each iteration + val a2b: DataFrame = { + val a2b = g.find("(a)-[e]->(b)") + edgeFilter match { + case Some(ef) => + val efExpr = GraphFrameInternals.applyExprToCol(g.spark, ef, "e") + a2b.filter(efExpr) + case None => + a2b + } + } + + // We will always apply fromExpr to column "a" + val fromAExpr = GraphFrameInternals.applyExprToCol(g.spark, from, "a") + + // DataFrame of current search paths + var paths: DataFrame = null + + var iter = 0 + var foundPath = false + while (iter < maxPathLength && !foundPath) { + val nextVertex = s"v${iter + 1}" + val nextEdge = s"e$iter" + // Take another step + if (iter == 0) { + // Note: We could avoid this special case by initializing paths with just 1 "from" column, + // but that would create a longer lineage for the result DataFrame. + paths = a2b + .filter(fromAExpr) + .filter(col("a.id") =!= col("b.id")) // remove self-loops + .withColumnRenamed("a", "from") + .withColumnRenamed("e", nextEdge) + .withColumnRenamed("b", nextVertex) + } else { + val prevVertex = s"v$iter" + val nextLinks = a2b + .withColumnRenamed("a", prevVertex) + .withColumnRenamed("e", nextEdge) + .withColumnRenamed("b", nextVertex) + paths = paths + .join(nextLinks, paths(prevVertex + ".id") === nextLinks(prevVertex + ".id")) + .drop(paths(prevVertex)) + // Make sure we are not backtracking within each path. + // TODO: Avoid crossing paths; i.e., touch each vertex at most once. + val previousVertexChecks = Range(1, iter + 1) + .map(i => paths(s"v$i.id") =!= paths(nextVertex + ".id")) + .foldLeft(paths("from.id") =!= paths(nextVertex + ".id"))((c1, c2) => c1 && c2) + paths = paths.filter(previousVertexChecks) + } + // Check if done by applying toExpr to column nextVertex + val toVExpr = GraphFrameInternals.applyExprToCol(g.spark, to, nextVertex) + val foundPathDF = paths.filter(toVExpr) + if (foundPathDF.take(1).nonEmpty) { + // Found path + paths = foundPathDF.withColumnRenamed(nextVertex, "to") + foundPath = true + } + iter += 1 + } + if (foundPath) { + logInfo(s"GraphFrame.bfs found path of length $iter.") + def rank(c: String): Double = { + // from < e0 < v1 < e1 < ... < to + c match { + case "from" => 0.0 + case "to" => Double.PositiveInfinity + case _ if c.startsWith("e") => 0.6 + c.substring(1).toInt + case _ if c.startsWith("v") => 0.3 + c.substring(1).toInt + } + } + val ordered = paths.columns.sortBy(rank _) + paths.select(ordered.map(col).toSeq: _*) + } else { + logInfo(s"GraphFrame.bfs failed to find a path of length <= $maxPathLength.") + // Return empty DataFrame + g.spark.createDataFrame(g.spark.sparkContext.parallelize(Seq.empty[Row]), g.vertices.schema) + } + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ConnectedComponents.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ConnectedComponents.scala new file mode 100644 index 0000000000000..360b65776f6ed --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ConnectedComponents.scala @@ -0,0 +1,222 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.graphx +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.graphframes.GraphFramesConf +import org.apache.spark.storage.StorageLevel +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithBroadcastThreshold +import org.apache.spark.graphframes.WithCheckpointInterval +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLocalCheckpoints +import org.apache.spark.graphframes.WithMaxIter +import org.apache.spark.graphframes.WithUseLabelsAsComponents + +/** + * Connected Components algorithm. + * + * Computes the connected component membership of each vertex and returns a DataFrame of vertex + * information with each vertex assigned a component ID. + * + * The resulting DataFrame contains all the vertex information and one additional column: + * - component (`LongType`): unique ID for this component + */ +class ConnectedComponents private[graphframes] (private val graph: GraphFrame) + extends Arguments + with Logging + with WithCheckpointInterval + with WithBroadcastThreshold + with WithIntermediateStorageLevel + with WithUseLabelsAsComponents + with WithMaxIter + with WithLocalCheckpoints { + + import ConnectedComponents._ + + private var algorithm: String = GraphFramesConf.getConnectedComponentsAlgorithm + .getOrElse(ALGO_TWO_PHASE) + + private var isGraphPrepared: Boolean = false + + setCheckpointInterval( + GraphFramesConf.getConnectedComponentsCheckpointInterval.getOrElse(checkpointInterval)) + setBroadcastThreshold( + GraphFramesConf.getConnectedComponentsBroadcastThreshold.getOrElse(broadcastThreshold)) + setIntermediateStorageLevel( + GraphFramesConf.getConnectedComponentsStorageLevel.getOrElse(intermediateStorageLevel)) + setUseLabelsAsComponents( + GraphFramesConf.getUseLabelsAsComponents.getOrElse(useLabelsAsComponents)) + setUseLocalCheckpoints(GraphFramesConf.getUseLocalCheckpoints.getOrElse(useLocalCheckpoints)) + + /** + * Sets the algorithm to use for computing connected components. Supported values: + * - [[ConnectedComponents.ALGO_GRAPHX]]: use the GraphX implementation + * - [[ConnectedComponents.ALGO_GRAPHFRAMES]]: deprecated alias for + * [[ConnectedComponents.ALGO_TWO_PHASE]] + * - [[ConnectedComponents.ALGO_TWO_PHASE]]: use the two-phase label propagation + * implementation + * - [[ConnectedComponents.ALGO_RANDOMIZED_CONTRACTION]]: use the randomized contraction + * implementation + */ + def setAlgorithm(value: String): this.type = { + val normalized = value.toLowerCase + normalized match { + case ALGO_GRAPHX | ALGO_TWO_PHASE | ALGO_RANDOMIZED_CONTRACTION => + algorithm = normalized + case ALGO_GRAPHFRAMES => + logWarn( + s"Algorithm '$ALGO_GRAPHFRAMES' is deprecated and will be removed in a future release. " + + s"Using '$ALGO_TWO_PHASE' instead.") + algorithm = ALGO_TWO_PHASE + case _ => + throw new IllegalArgumentException( + s"Unsupported algorithm: '$value'. " + + s"Supported values are: $ALGO_GRAPHX, $ALGO_TWO_PHASE, " + + s"$ALGO_RANDOMIZED_CONTRACTION, $ALGO_GRAPHFRAMES (deprecated).") + } + this + } + + /** + * Gets the algorithm used for computing connected components. + */ + def getAlgorithm: String = algorithm + + /** + * !! WARNING: INTERNAL API — FOR VERY EXPERIENCED USERS ONLY !! + * + * Sets whether the graph has already been prepared before being passed to the algorithm, + * skipping the internal graph preparation step. The default is `false`, meaning the algorithm + * will always prepare the graph itself, which is the safe and recommended behaviour. + * + * Only set this to `true` if you have '''already performed all required preparation steps + * yourself''' and you fully understand what those steps are for the specific algorithm you are + * using. '''The preparation requirements differ significantly between algorithms:''' + * + * - `two_phase` and `randomized_contraction` each require their own distinct preparation + * steps. These are NOT interchangeable. You MUST study the internal source code of the + * algorithm you intend to use and replicate its exact preparation logic before enabling + * this flag. + * + * '''Incorrect use of this flag WILL produce silently wrong results with no error or warning at + * runtime.''' There is no validation that the graph has been correctly prepared. You are + * entirely responsible for ensuring correctness. + * + * @param value + * true if the graph is already prepared, false otherwise (default: false) + */ + def setIsGraphPrepared(value: Boolean): this.type = { + logWarn( + "INTERNAL API ONLY WAS CALLED. This is an internal option for advanced users who fully " + + "understand graph preparation internals. Misuse will produce silently wrong results.") + isGraphPrepared = value + this + } + + /** + * Runs the algorithm. + */ + def run(): DataFrame = { + algorithm match { + case ALGO_GRAPHX => + ConnectedComponents.runGraphX( + graph, + maxIter.getOrElse(Int.MaxValue), + intermediateStorageLevel) + case ALGO_TWO_PHASE => + if (broadcastThreshold == -1) { + TwoPhase.runAQE( + graph, + checkpointInterval = checkpointInterval, + intermediateStorageLevel = intermediateStorageLevel, + useLabelsAsComponents = useLabelsAsComponents, + useLocalCheckpoints = useLocalCheckpoints, + isGraphPrepared = isGraphPrepared) + } else { + TwoPhase.run( + graph, + broadcastThreshold = broadcastThreshold, + checkpointInterval = checkpointInterval, + intermediateStorageLevel = intermediateStorageLevel, + useLabelsAsComponents = useLabelsAsComponents, + useLocalCheckpoints = useLocalCheckpoints, + isGraphPrepared = isGraphPrepared) + } + case ALGO_RANDOMIZED_CONTRACTION => + RandomizedContraction.run( + graph, + useLabelsAsComponents = useLabelsAsComponents, + intermediateStorageLevel = intermediateStorageLevel, + useLocalCheckpoints = useLocalCheckpoints, + checkpointInterval = checkpointInterval, + isGraphPrepared = isGraphPrepared) + // the check is inside the setter + case _ => throw new GraphFramesUnreachableException() + } + } + + @deprecated("use graph.connectedComponents instead", "0.11.0") + def run(graph: GraphFrame): DataFrame = { + new ConnectedComponents(graph).run() + } +} + +object ConnectedComponents extends Logging { + + private[graphframes] val COMPONENT = "component" + private[graphframes] val ORIG_ID = "orig_id" + + val ALGO_GRAPHX = "graphx" + + /** + * @deprecated + * Use [[ALGO_TWO_PHASE]] instead. + */ + val ALGO_GRAPHFRAMES = "graphframes" + + val ALGO_TWO_PHASE = "two_phase" + val ALGO_RANDOMIZED_CONTRACTION = "randomized_contraction" + + /** + * Runs the GraphX connected components implementation. + */ + private[graphframes] def runGraphX( + graph: GraphFrame, + maxIter: Int, + intermediateStorageLevel: StorageLevel): DataFrame = { + val gx = graph.cachedTopologyGraphX + val components = + graphx.lib.ConnectedComponents.run(gx, maxIter) + val result = GraphXConversions + .fromGraphX(graph, components, vertexNames = Seq(ConnectedComponents.COMPONENT)) + .vertices + .persist(intermediateStorageLevel) + + val _ = result.count() + gx.unpersist() + components.unpersist() + + resultIsPersistent() + + result + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/DetectingCycles.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/DetectingCycles.scala new file mode 100644 index 0000000000000..0204abcf90fa2 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/DetectingCycles.scala @@ -0,0 +1,122 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types.ArrayType +import org.apache.spark.storage.StorageLevel +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithCheckpointInterval +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLocalCheckpoints + +class DetectingCycles private[graphframes] (private val graph: GraphFrame) + extends Arguments + with Serializable + with Logging + with WithIntermediateStorageLevel + with WithLocalCheckpoints + with WithCheckpointInterval { + import DetectingCycles._ + def run(): DataFrame = { + val rawRes = DetectingCycles.run( + graph, + useLocalCheckpoints, + checkpointInterval, + intermediateStorageLevel) + val explodedRes = rawRes + .select( + col(GraphFrame.ID), + filter(col(foundSeqCol), x => size(x) > lit(0)).alias(foundSeqCol)) + .filter(size(col(foundSeqCol)) > lit(0)) + .select( + // from vid -> [[cycle1, cycle2, ...]] + // to vid -> [cycle1], vid -> [cycle2], ... + explode(col(foundSeqCol)).alias(foundSeqCol)) + .persist(intermediateStorageLevel) + explodedRes.count() + resultIsPersistent() + rawRes.unpersist() + explodedRes + } +} + +object DetectingCycles { + private val storedSeqCol: String = "sequences" + val foundSeqCol: String = "found_cycles" + + def run( + graph: GraphFrame, + useLocalCheckpoints: Boolean, + checkpointInterval: Int, + intermediateStorageLevel: StorageLevel): DataFrame = { + val preparedGraph = GraphFrame( + graph.vertices.select(GraphFrame.ID), + graph.edges.select(GraphFrame.SRC, GraphFrame.DST)) + + val vertexDT = preparedGraph.vertices.schema(GraphFrame.ID).dataType + + // Each vertex stores sequences from the previous iteration, initial is just Array(Array(ID)) + val initSequences = array(array(col(GraphFrame.ID))) + // Each vertex stores all the found cycles + val foundSequences = array().cast(ArrayType(ArrayType(vertexDT))) + // Message is simply stored sequences + // Send only sequences if the starting vertex of them is less than the destination + val sentMessages = when( + size(Pregel.src(storedSeqCol)) =!= lit(0), + filter(Pregel.src(storedSeqCol), (x: Column) => x.getItem(0) <= Pregel.dst(GraphFrame.ID))) + .otherwise(lit(null).cast(ArrayType(ArrayType(vertexDT)))) + // If the sequence contains the current vertex ID somewhere in the middle, it is + // a previously detected cycle and a sequence should be discarded. + val filterOutSequences = flatten(collect_list(Pregel.msg)) + when(Pregel.msg.isNull, array(array()).cast(ArrayType(ArrayType(vertexDT)))) + .otherwise(filter(Pregel.msg, x => !(array_position(x, col(GraphFrame.ID)) > lit(1)))) + // update found sequences by appending all from messages that start from the current vertex ID + val updateFound = when(Pregel.msg.isNull, col(foundSeqCol)).otherwise( + array_union( + col(foundSeqCol), + transform( + filter(Pregel.msg, x => try_element_at(x, lit(1)) === col(GraphFrame.ID)), + x => array_append(x, col(GraphFrame.ID))))) + // update stored sequences by filtering out already added sequences + val updateSequences = transform( + filter(Pregel.msg, x => !array_contains(x, col(GraphFrame.ID))), + x => array_append(x, col(GraphFrame.ID))) + + preparedGraph.pregel + .setCheckpointInterval(checkpointInterval) + .setUseLocalCheckpoints(useLocalCheckpoints) + .setIntermediateStorageLevel(intermediateStorageLevel) + .setEarlyStopping(false) + .setSkipMessagesFromNonActiveVertices(true) + .setInitialActiveVertexExpression(lit(true)) + .sendMsgToDst(sentMessages) + .setUpdateActiveVertexExpression(Pregel.msg.isNotNull && (size(updateSequences) > lit(0))) + .withVertexColumn(storedSeqCol, initSequences, updateSequences) + .withVertexColumn(foundSeqCol, foundSequences, updateFound) + .aggMsgs(filterOutSequences) + // Memory optimization: only include required columns in triplets + // For cycle detection, we only need the sequences from source vertex + // and just the ID from destination vertex (ID is always included) + .requiredSrcColumns(storedSeqCol) + .run() + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/GraphXConversions.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/GraphXConversions.scala new file mode 100644 index 0000000000000..a4442a307f192 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/GraphXConversions.scala @@ -0,0 +1,206 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.graphx.Graph +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types.StructField +import org.apache.spark.sql.types.StructType +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.NoSuchVertexException + +import scala.reflect.runtime.universe._ + +/** + * Convenience functions to map GraphX graphs to GraphFrames, checking for the types expected by + * GraphX. + */ +private[graphframes] object GraphXConversions { + + import GraphFrame._ + + /** Indicates if T is a Unit type */ + private def isUnitType[T: TypeTag]: Boolean = { + val t = typeOf[T] + typeOf[Unit] =:= t + } + + /** Indicates if T is a Product type */ + private def isProductType[T: TypeTag]: Boolean = { + val t = typeOf[T] + // See http://stackoverflow.com/questions/21209006/how-to-check-if-reflected-type-represents-a-tuple + t.typeSymbol.fullName.startsWith("scala.Tuple") + } + + /** See [[GraphFrame.fromGraphX()]] for documentation */ + def fromGraphX[V: TypeTag, E: TypeTag]( + originalGraph: GraphFrame, + graph: Graph[V, E], + vertexNames: Seq[String] = Nil, + edgeNames: Seq[String] = Nil): GraphFrame = { + val spark = originalGraph.spark + // catalyst does not like the unit type, make sure to filter it first. + val vertexDF: DataFrame = if (isUnitType[V]) { + val vertexData = graph.vertices.map { case (vid, _) => Tuple1(vid) } + spark.createDataFrame(vertexData).toDF(LONG_ID) + } else if (isProductType[V]) { + val vertexData = graph.vertices.map { case (vid, data) => (vid, data) } + val vertexDF0 = spark.createDataFrame(vertexData).toDF(LONG_ID, GX_ATTR) + renameStructFields(vertexDF0, GX_ATTR, vertexNames) + } else { + // Assume it is just one field, and pack it in a tuple to have a structure. + val vertexData = graph.vertices.map { case (vid, data) => (vid, Tuple1(data)) } + val vertexDF0 = spark.createDataFrame(vertexData).toDF(LONG_ID, GX_ATTR) + renameStructFields(vertexDF0, GX_ATTR, vertexNames) + } + + val edgeDF: DataFrame = if (isUnitType[E]) { + val edgeData = graph.edges.map { e => (e.srcId, e.dstId) } + spark.createDataFrame(edgeData).toDF(LONG_SRC, LONG_DST) + } else if (isProductType[E]) { + val edgeData = graph.edges.map { e => (e.srcId, e.dstId, e.attr) } + val edgeDF0 = spark.createDataFrame(edgeData).toDF(LONG_SRC, LONG_DST, GX_ATTR) + renameStructFields(edgeDF0, GX_ATTR, edgeNames) + } else { + val edgeData = graph.edges.map { e => (e.srcId, e.dstId, Tuple1(e.attr)) } + val edgeDF0 = spark.createDataFrame(edgeData).toDF(LONG_SRC, LONG_DST, GX_ATTR) + renameStructFields(edgeDF0, GX_ATTR, edgeNames) + } + fromGraphX(originalGraph, vertexDF, edgeDF) + } + + /** + * Given the name of a column (assumed to contain a struct), renames all the fields of this + * struct. + * + * @param structName + * Struct name whose fields will be renamed. This method assumes this field exists and will + * not check for errors. + * @param fieldNames + * List of new field names corresponding to all fields in the struct col. + */ + private[lib] def renameStructFields( + df: DataFrame, + structName: String, + fieldNames: Seq[String]): DataFrame = { + // TODO(tjh) this looses metadata and other info in the process + val origSubfields = colStar(df, structName).map(col) + val renamedSubfields = origSubfields.zip(fieldNames).map { case (orig, newName) => + orig.as(newName) + } + val otherFields = df.schema.fieldNames.filter(_ != structName).map(quote).map(col) + if (renamedSubfields.isEmpty) { + // Do not attempt to add an empty structure. + df.select(otherFields.toSeq: _*) + } else { + val renamedStruct = struct(renamedSubfields.toSeq: _*).as(structName) + df.select((renamedStruct +: otherFields).toSeq: _*) + } + } + + private def drop(df: DataFrame, cols: String*): DataFrame = { + val remainingCols = df.schema.map(_.name).filterNot(cols.contains).map(quote).map(n => df(n)) + df.select(remainingCols: _*) + } + + /** Unpacks all struct fields and leaves other fields alone */ + private def unpackStructFields(df: DataFrame): DataFrame = { + val cols = df.schema.flatMap { + case StructField(fname, dt: StructType, nullable @ _, meta @ _) => + dt.iterator.map(sub => col(quote(fname, sub.name)).as(sub.name)) + case f => Seq(col(quote(f.name))) + } + df.select(cols: _*) + } + + /** + * Joins all the data from the original columns against the new data. Assumes the columns are + * not going to conflict. + * + * @param gxVertexData + * DataFrame with column [[LONG_ID]] and optionally column [[GX_ATTR]] + * @param gxEdgeData + * DataFrame with columns [[LONG_DST]], [[LONG_SRC]] and optionally column [[GX_ATTR]] + */ + private def fromGraphX( + originalGraph: GraphFrame, + gxVertexData: DataFrame, + gxEdgeData: DataFrame): GraphFrame = { + // The ID is going to be unpacked from the attr field + val packedVertices = drop(originalGraph.indexedVertices, ID).join(gxVertexData, LONG_ID) + val vertexDF = unpackStructFields(drop(packedVertices, LONG_ID)) + + val packedEdges = { + val indexedEdges = originalGraph.indexedEdges + // Handle 2 cases: GraphX edge has attr, or not. + val hasGxAttr = gxEdgeData.schema.exists(_.name == GX_ATTR) + val gxCol = if (hasGxAttr) { Seq(col(GX_ATTR)) } + else { Seq() } + val sel1 = Seq(col(LONG_SRC), col(LONG_DST)) ++ gxCol + val gxe = gxEdgeData.select(sel1: _*) + val sel3 = Seq(col(ATTR)) ++ gxCol + // TODO: CHECK IN UNIT TESTS: Drop the src and dst columns from the index, they are already + // in the attributes and will be unpacked with the rest of the user columns. + // TODO(tjh) 2-step join? + gxe + .join( + indexedEdges.select(indexedEdges(LONG_SRC), indexedEdges(LONG_DST), indexedEdges(ATTR)), + (gxe(LONG_SRC) === indexedEdges(LONG_SRC)) && (gxe(LONG_DST) === indexedEdges( + LONG_DST))) + .select(sel3: _*) + } + val edgeDF = unpackStructFields(drop(packedEdges, LONG_SRC, LONG_DST)) + + GraphFrame(vertexDF, edgeDF) + } + + /** + * Given a graph and an object, gets the the corresponding integral id in the internal + * representation. + */ + private[graphframes] def integralId(graph: GraphFrame, vertexId: Any): Long = { + // Check if we can directly convert it + vertexId match { + case x: Int => return x.toLong + case x: Long => return x.toLong + case x: Short => return x.toLong + case x: Byte => return x.toLong + case _ => + } + // If the vertex is a non-integral type such as a String, we need to use the translation table. + val longIdRow: Array[Row] = graph.indexedVertices + .filter(col(GraphFrame.ID) === vertexId) + .select(GraphFrame.LONG_ID) + .take(1) + if (longIdRow.isEmpty) { + throw new NoSuchVertexException( + "GraphFrame algorithm given vertex ID which does not exist" + + s" in Graph. Vertex ID $vertexId not contained in $graph") + } + // TODO(tjh): could do more informative message + longIdRow.head.getLong(0) + } +} + +private[lib] trait Arguments { + private[lib] def check[A](a: Option[A], name: String): A = { + a.getOrElse(throw new IllegalArgumentException(s"Param $name is required.")) + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/HyperANF.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/HyperANF.scala new file mode 100644 index 0000000000000..fa4d8629c54fb --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/HyperANF.scala @@ -0,0 +1,237 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.hll_sketch_agg +import org.apache.spark.sql.functions.hll_union_agg +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.functions.udf +import org.apache.spark.sql.types.ByteType +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.types.LongType +import org.apache.spark.sql.types.ShortType +import org.apache.spark.sql.types.StringType +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFramesUnsupportedVertexTypeException +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithCheckpointInterval +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLocalCheckpoints + +/** + * HyperANF-style approximation of the neighbourhood function on top of GraphFrames. + * + * This implementation is inspired by + * [[https://arxiv.org/pdf/1011.5599 Vigna, Paolo; Boldi, Marco; Rosa, Sebastiano. "HyperANF: Approximating the Neighbourhood Function of Very Large Graphs on a Budget." arXiv preprint arXiv:1011.5599 (2010)]]. + * + * The input graph is treated as directed: for each vertex, reachability is computed by following + * outgoing edges from `src` to `dst`. + * + * Compared with the cumulative neighbourhood-function presentation in the paper, this + * implementation returns one column per hop, `hop_0`, `hop_1`, `hop_2`, ..., `hop_N`. The `hop_0` + * column contains a HyperLogLog sketch of the source vertex itself, and each `hop_k` column for + * `k >= 1` contains a HyperLogLog sketch of the set of vertices reachable in exactly `k` hops. To + * derive the cumulative approximate neighbourhood function for distances up to some hop `k`, a + * user can combine `hop_0` through `hop_k` with `hll_union` and then apply `hll_sketch_estimate` + * to the merged sketch. + * + * The computation can also be restricted to a subgraph by supplying an edge filter expression via + * [[setEdgesFilterExpression]]. A common use case is to filter on `src`, for example + * `src IN (...)`, to obtain sketches only for a selected set of starting vertices. + * + * @param graph + * input graph whose directed edges are used for reachability expansion + */ +class HyperANF private[graphframes] (graph: GraphFrame) + extends Serializable + with Logging + with WithCheckpointInterval + with WithIntermediateStorageLevel + with WithLocalCheckpoints { + private var nHops: Int = 3 + private var edgesFilterExpression: Column = lit(true) + private var lgNomEntries: Int = 12 + + /** + * Sets the log2 of nominal entries used by HLL sketch aggregations. + */ + def setLgNomEntries(value: Int): this.type = { + require((value >= 4) && (value <= 21), "lgNomEntries must be between 4 and 21") + lgNomEntries = value + this + } + + /** + * Sets the edge filter expression used before running the computation. + * + * Only edges satisfying this predicate participate in the directed reachability expansion. This + * effectively runs the algorithm on the subgraph induced by the filtered edge set. + * + * A common use case is filtering on `src`, for example `src IN (...)`, to limit the result to a + * chosen set of starting vertices. + * + * @param value + * filter expression applied to `graph.edges` + * @return + * this HyperANF instance + */ + def setEdgesFilterExpression(value: Column): this.type = { + edgesFilterExpression = value + this + } + + /** + * Sets the maximum hop distance to compute. + * + * The result will contain `hop_0`, `hop_1`, `hop_2`, ..., `hop_N`, where `N` is the configured + * number of hops. + * + * @param value + * positive number of hops to compute + * @return + * this HyperANF instance + */ + def setNHops(value: Int): this.type = { + require(value > 0, "n-hops cannot be negative or zero") + nHops = value + this + } + + /** + * Runs the HyperANF-style computation. + * + * The returned `DataFrame` has one row per source vertex present in the filtered edge set. It + * contains the vertex id column `id` and one sketch column per hop: `hop_0`, `hop_1`, `hop_2`, + * ..., `hop_N`. The `hop_0` column stores a HyperLogLog sketch containing `id` itself. Each + * `hop_k` column for `k >= 1` stores a HyperLogLog sketch for the set of vertices reachable + * from `id` in exactly `k` directed hops. + * + * To obtain an approximate cumulative neighbourhood size up to hop `k`, union `hop_0` through + * `hop_k` with `hll_union` and then apply `hll_sketch_estimate`. + * + * @return + * a `DataFrame` with exact-hop HyperLogLog sketches per source vertex + */ + def run(): DataFrame = { + val edges = + graph.edges + .filter(edgesFilterExpression) + .select(GraphFrame.SRC, GraphFrame.DST) + .persist(intermediateStorageLevel) + var hop = 1 + + val hop0Func = graph.vertices.schema(GraphFrame.ID).dataType match { + case IntegerType => udf(HyperANF.hllInt(lgNomEntries)) + case LongType => udf(HyperANF.hllLong(lgNomEntries)) + case StringType => udf(HyperANF.hllString(lgNomEntries)) + case ShortType => udf(HyperANF.hllShort(lgNomEntries)) + case ByteType => udf(HyperANF.hllByte(lgNomEntries)) + case _ => + throw new GraphFramesUnsupportedVertexTypeException( + s"Unsupported vertex ID type: ${graph.vertices.schema(GraphFrame.ID).dataType}") + } + var state = edges + .groupBy(col(GraphFrame.SRC).alias(GraphFrame.ID)) + .agg(hll_sketch_agg(GraphFrame.DST, lgNomEntries).alias("hop_1")) + .select(col(GraphFrame.ID), hop0Func(col(GraphFrame.ID)).alias("hop_0"), col("hop_1")) + .persist(intermediateStorageLevel) + + // materialize + val cnt = state.count() + logInfo(s"found $cnt vertices with at least one outgoing edge") + + val shouldCheckpoint = (checkpointInterval > 0) && (checkpointInterval < nHops) + + while (hop < nHops) { + hop += 1 + + val nState = edges + .join( + state.select(GraphFrame.ID, s"hop_${hop - 1}"), + col(GraphFrame.DST) === col(GraphFrame.ID), + "left") + .groupBy(col(GraphFrame.SRC).alias(GraphFrame.ID)) + .agg(hll_union_agg(s"hop_${hop - 1}").alias(s"hop_${hop}")) + + // standard GF persist-unpersist-checkpoint flow + state = { + val stateToPersist = state.join(nState, GraphFrame.ID) + if (shouldCheckpoint && hop % checkpointInterval == 0) { + if (useLocalCheckpoints) { + stateToPersist.localCheckpoint(eager = false) + } else { + stateToPersist.checkpoint(eager = false) + } + } else { + stateToPersist.persist(intermediateStorageLevel) + // materialize + stateToPersist.count() + + state.unpersist() + stateToPersist + } + } + + logInfo(s"hop $hop / $nHops was computed") + } + + // state is already persisted at the moment + resultIsPersistent() + edges.unpersist() + + state + } +} + +private object HyperANF extends Serializable { + // If you are confusing to see 5 almost identical functions: + // it was intentional. HLL does not have `update(Object)`. + + def hllInt(lgNomEntries: Int): Int => Array[Byte] = (id) => { + val sketch = new org.apache.datasketches.hll.HllSketch(lgNomEntries) + sketch.update(id.toLong) + sketch.toCompactByteArray() + } + + def hllLong(lgNomEntries: Int): Long => Array[Byte] = (id) => { + val sketch = new org.apache.datasketches.hll.HllSketch(lgNomEntries) + sketch.update(id) + sketch.toCompactByteArray() + } + + def hllString(lgNomEntries: Int): String => Array[Byte] = (id) => { + val sketch = new org.apache.datasketches.hll.HllSketch(lgNomEntries) + sketch.update(id) + sketch.toCompactByteArray() + } + + def hllShort(lgNomEntries: Int): Short => Array[Byte] = (id) => { + val sketch = new org.apache.datasketches.hll.HllSketch(lgNomEntries) + sketch.update(id.toLong) + sketch.toCompactByteArray() + } + + def hllByte(lgNomEntries: Int): Byte => Array[Byte] = (id) => { + val sketch = new org.apache.datasketches.hll.HllSketch(lgNomEntries) + sketch.update(id.toLong) + sketch.toCompactByteArray() + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/KCore.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/KCore.scala new file mode 100644 index 0000000000000..9227601377fbe --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/KCore.scala @@ -0,0 +1,126 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.catalyst.FunctionIdentifier +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.functions.call_function +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.collect_list +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.functions.when +import org.apache.spark.sql.graphframes.expressions.KCoreMerge +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.storage.StorageLevel +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithCheckpointInterval +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLocalCheckpoints + +/** + * K-Core decomposition algorithm implementation for GraphFrames. + * + * This object provides the `run` method to compute the k-core decomposition of a graph, which + * assigns each vertex the maximum k such that the vertex is part of a k-core. A k-core is a + * maximal connected subgraph in which every vertex has degree at least k. + * + * The algorithm is based on the distributed k-core decomposition approach described in: + * + * Mandal, Aritra, and Mohammad Al Hasan. "A distributed k-core decomposition algorithm on spark." + * 2017 IEEE International Conference on Big Data (Big Data). IEEE, 2017. + * + * '''Edge representation''': K-core decomposition is defined for undirected graphs. Since + * GraphFrames represents edges as directed, each undirected edge `{u, v}` should be supplied as a + * single directed edge in either direction — the algorithm symmetrizes internally. Supplying both + * `(u, v)` and `(v, u)` will double-count the edge and produce incorrect results. + */ +class KCore private[graphframes] (private val graph: GraphFrame) + extends Serializable + with WithIntermediateStorageLevel + with WithCheckpointInterval + with WithLocalCheckpoints { + import org.apache.spark.graphframes.lib.KCore.kCoreColumnName + def run(): DataFrame = { + val result = + KCore.run(graph, intermediateStorageLevel, checkpointInterval, useLocalCheckpoints) + val allVertices = graph.vertices + .select(GraphFrame.ID) + .join(result, Seq(GraphFrame.ID), "left") + .withColumn( + kCoreColumnName, + when(col(kCoreColumnName).isNull, lit(0)).otherwise(col(kCoreColumnName))) + .persist(intermediateStorageLevel) + + // materialize + allVertices.count() + result.unpersist() + allVertices + } +} + +object KCore extends Serializable with Logging { + val kCoreColumnName = "kcore" + def run( + graph: GraphFrame, + storageLevel: StorageLevel, + checkpointInterval: Int, + useLocalCheckpoints: Boolean): DataFrame = { + val degrees = graph.degrees + val preparedGraph = GraphFrame( + degrees.withColumn("degree", col("degree").cast(IntegerType)), + graph.edges.select(GraphFrame.SRC, GraphFrame.DST)) + + val functionRegistry = graph.vertices.sparkSession.sessionState.functionRegistry + functionRegistry.registerFunction( + new FunctionIdentifier("_kcoreMerge", Some("builtin"), Some("system")), + (children: Seq[Expression]) => KCoreMerge(children(0), children(1)), + "scala_udf") + + try { + val pregel = preparedGraph.pregel + .setMaxIter(Int.MaxValue) + .setIntermediateStorageLevel(storageLevel) + .setCheckpointInterval(checkpointInterval) + .withVertexColumn( + kCoreColumnName, + col("degree"), + call_function("_kcoreMerge", Pregel.msg, col(kCoreColumnName))) + .sendMsgToSrc(Pregel.dst(kCoreColumnName)) + .sendMsgToDst(Pregel.src(kCoreColumnName)) + .setInitialActiveVertexExpression(lit(true)) + .setUpdateActiveVertexExpression( + col(kCoreColumnName) =!= call_function("_kcoreMerge", Pregel.msg, col(kCoreColumnName))) + .setEarlyStopping(false) + .setStopIfAllNonActiveVertices(true) + .setSkipMessagesFromNonActiveVertices(false) + .setUseLocalCheckpoints(useLocalCheckpoints) + .aggMsgs(collect_list(Pregel.msg)) + + pregel.run() + } finally { + val dereg = functionRegistry.dropFunction( + new FunctionIdentifier("_kcoreMerge", Some("builtin"), Some("system"))) + if (!dereg) { + logWarn( + "graphframes faced an internal error and was not able to de-register function _kcoreMerge; Spark' functionRegistry is in a bad state") + } + } + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/LabelPropagation.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/LabelPropagation.scala new file mode 100644 index 0000000000000..becc77e6ff78b --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/LabelPropagation.scala @@ -0,0 +1,146 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.graphx +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.types.MapType +import org.apache.spark.storage.StorageLevel +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithAlgorithmChoice +import org.apache.spark.graphframes.WithCheckpointInterval +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLocalCheckpoints +import org.apache.spark.graphframes.WithMaxIter + +/** + * Run static Label Propagation for detecting communities in networks. + * + * Each node in the network is initially assigned to its own community. At every iteration, nodes + * send their community affiliation to all neighbors and update their state to the mode community + * affiliation of incoming messages. + * + * LPA is a standard community detection algorithm for graphs. It is very inexpensive + * computationally, although (1) convergence is not guaranteed and (2) one can end up with trivial + * solutions (all nodes are identified into a single community). + * + * The resulting DataFrame contains all the original vertex information and one additional column: + * - label (`LongType`): label of community affiliation + */ +class LabelPropagation private[graphframes] (private val graph: GraphFrame) + extends Arguments + with WithAlgorithmChoice + with WithCheckpointInterval + with WithMaxIter + with WithLocalCheckpoints + with WithIntermediateStorageLevel + with Logging { + + def run(): DataFrame = { + val maxIterChecked = check(maxIter, "maxIter") + val res = algorithm match { + case "graphx" => LabelPropagation.runInGraphX(graph, maxIterChecked) + case "graphframes" => + LabelPropagation.runInGraphFrames( + graph, + maxIterChecked, + checkpointInterval, + useLocalCheckpoints = useLocalCheckpoints, + intermediateStorageLevel = intermediateStorageLevel) + } + resultIsPersistent() + res + } +} + +private object LabelPropagation { + private def runInGraphX(graph: GraphFrame, maxIter: Int): DataFrame = { + val gx = graphx.lib.LabelPropagation.run(graph.cachedTopologyGraphX, maxIter) + val res = GraphXConversions.fromGraphX(graph, gx, vertexNames = Seq(LABEL_ID)).vertices + res.persist(StorageLevel.MEMORY_AND_DISK_SER) + res.count() + gx.unpersist() + res + } + + private def keyWithMaxValue(column: Column): Column = { + // Get the key with the highest value, using the key to break a tie. To do this, simply get + // map entries, swap the value and key columns to create the natural ordering, multiply key by -1 and then + // take the key from the max entry (multiply it by -1 again to get the original key). + array_max( + transform( + map_entries(column), + x => struct(x.getField("value"), (lit(-1) * x.getField("key")).alias("key")))) + .getField("key") * lit(-1) + } + + private def runInGraphFrames( + graph: GraphFrame, + maxIter: Int, + checkpointInterval: Int, + isDirected: Boolean = true, + useLocalCheckpoints: Boolean, + intermediateStorageLevel: StorageLevel): DataFrame = { + // Overall: + // - Initial labels - IDs + // - Active vertex col (halt voting) - did the label changed? + // - Choosing a new label - top across neighbours (tie-braking is deterministic) + + val preparedGraph = GraphFrame( + graph.vertices.select(GraphFrame.ID), + graph.edges.select(GraphFrame.SRC, GraphFrame.DST)) + + var pregel = preparedGraph.pregel + .withVertexColumn(LABEL_ID, col(GraphFrame.ID).alias(LABEL_ID), keyWithMaxValue(Pregel.msg)) + .setMaxIter(maxIter) + .setStopIfAllNonActiveVertices(true) + .setEarlyStopping(false) + .setCheckpointInterval(checkpointInterval) + .setSkipMessagesFromNonActiveVertices(false) + .setUpdateActiveVertexExpression(col(LABEL_ID) =!= keyWithMaxValue(Pregel.msg)) + .setUseLocalCheckpoints(useLocalCheckpoints) + .setIntermediateStorageLevel(intermediateStorageLevel) + // Memory optimization: only include required columns in triplets + .requiredSrcColumns(LABEL_ID) + .requiredDstColumns(LABEL_ID) + + if (isDirected) { + pregel = pregel.sendMsgToDst(Pregel.src(LABEL_ID)) + } else { + pregel = pregel.sendMsgToDst(Pregel.src(LABEL_ID)).sendMsgToSrc(Pregel.dst(LABEL_ID)) + } + + pregel = pregel.aggMsgs( + reduce( + collect_list(Pregel.msg), + map().cast(MapType(graph.vertices.schema(GraphFrame.ID).dataType, IntegerType)), + (acc, x) => + map_zip_with( + acc, + map(x, lit(1)), + (_, left, right) => coalesce(left, lit(0)) + coalesce(right, lit(0))))) + + pregel.run() + } + + private val LABEL_ID = "label" +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/MaximalIndependentSet.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/MaximalIndependentSet.scala new file mode 100644 index 0000000000000..86785afcff2a0 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/MaximalIndependentSet.scala @@ -0,0 +1,242 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types.DoubleType +import org.apache.spark.storage.StorageLevel +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithCheckpointInterval +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLocalCheckpoints + +import java.io.IOException + +/** + * This class implements a distributed algorithm for finding a Maximal Independent Set (MIS) in a + * graph. + * + * An MIS is a set of vertices such that no two vertices in the set are adjacent (i.e., there is + * no edge between any two vertices in the set), and the set is maximal, meaning that adding any + * other vertex to the set would violate the independence property. Note that this implementation + * finds a maximal (but not necessarily maximum) independent set; that is, it ensures no more + * vertices can be added to the set, but does not guarantee that the set has the largest possible + * number of vertices among all possible independent sets in the graph. + * + * The algorithm implemented here is based on the paper: Ghaffari, Mohsen. "An improved + * distributed algorithm for maximal independent set." Proceedings of the twenty-seventh annual + * ACM-SIAM symposium on Discrete algorithms. Society for Industrial and Applied Mathematics, + * 2016. + * + * Note: This is a randomized, non-deterministic algorithm. The result may vary between runs even + * if a fixed random seed is provided because how Apache Spark works. + * + * @param graph + */ +class MaximalIndependentSet private[graphframes] (private val graph: GraphFrame) + extends Serializable + with WithIntermediateStorageLevel + with WithCheckpointInterval + with WithLocalCheckpoints { + def run(seed: Long): DataFrame = { + MaximalIndependentSet.run( + graph, + checkpointInterval, + useLocalCheckpoints, + intermediateStorageLevel, + seed) + } +} + +object MaximalIndependentSet extends Serializable with Logging { + private val probCol = "prob" + private val degCol = "effectiveDegree" + private val isNominated = "isNominated" + private val notJoinedMISCol = "notJoinMIS" + private val isMIS = "isMIS" + + private def run( + graph: GraphFrame, + checkpointInterval: Int, + useLocalCheckpoints: Boolean, + storageLevel: StorageLevel, + seed: Long): DataFrame = { + // initial p = 1/2 + var vertices = + graph.vertices + .select(col(GraphFrame.ID), lit(0.5).cast(DoubleType).alias(probCol)) + .persist(storageLevel) + + // make edges undirected and de-duplicate + // persist() for future usage + val edges = graph.edges + .select(GraphFrame.SRC, GraphFrame.DST) + .union( + graph.edges.select( + col(GraphFrame.DST).alias(GraphFrame.SRC), + col(GraphFrame.SRC).alias(GraphFrame.DST))) + .filter(col(GraphFrame.SRC) =!= col(GraphFrame.DST)) + .distinct() + .persist(storageLevel) + + var misDF = graph.vertices.select(col(GraphFrame.ID), lit(false).alias(isMIS)) + + var i = 0 + var converged = false + val spark = graph.vertices.sparkSession + + val shouldCheckpoint = checkpointInterval > 0 + if (!useLocalCheckpoints && spark.sparkContext.getCheckpointDir.isEmpty) { + // Spark-Connect workaround + spark.sparkContext + .setCheckpointDir(spark.conf + .getOption("spark.checkpoint.dir") match { + case Some(d) => d + case None => + throw new IOException( + "Checkpoint directory is not set. Please set it first using sc.setCheckpointDir()" + + "or by specifying the conf 'spark.checkpoint.dir'.") + }) + } + + val rng = new util.Random(seed) + + // randomized algorithms are not working with AQE well + val originalAQE = spark.conf.get("spark.sql.adaptive.enabled") + try { + spark.conf.set("spark.sql.adaptive.enabled", "false") + + while (!converged) { + val iterSeed = rng.nextLong() + // compute effective degree as a sum of nbrs p + val effectiveDegrees = + edges + .join(vertices, col(GraphFrame.ID) === col(GraphFrame.DST)) + .groupBy(GraphFrame.SRC) + .agg(sum(col(probCol)).alias(degCol)) + + // update p per vertex by condition: + // if effective degree >= 2 then p / 2 + // else min(2p, 1/2) + // + // + mark vertices based on p + val probs = vertices + .join(effectiveDegrees, col(GraphFrame.ID) === col(GraphFrame.SRC)) + .drop(GraphFrame.SRC) + .withColumn( + probCol, + when(col(degCol) >= lit(2), col(probCol) / lit(2.0)).otherwise( + when(lit(2) * col(probCol) <= lit(0.5), lit(2) * col(probCol)).otherwise(lit(0.5)))) + .withColumn(isNominated, col(probCol) >= rand(iterSeed)) + .select(GraphFrame.ID, isNominated, probCol) + .persist(storageLevel) + + val isolatedVertices = + vertices + .join(probs.select(col(GraphFrame.ID)), Seq(GraphFrame.ID), "left_anti") + .select(GraphFrame.ID) + + // if no nbr of v is marked and v is marked, + // v is joined MIS and removed with all it's nbrs + val isJoinedMIS = probs + .join( + edges + .join(probs, col(GraphFrame.ID) === col(GraphFrame.DST)) + .groupBy(GraphFrame.SRC) + .agg(bool_or(col(isNominated)).alias(notJoinedMISCol)), + col(GraphFrame.SRC) === col(GraphFrame.ID)) + .select(GraphFrame.ID, probCol, isNominated, notJoinedMISCol) + + val joinedMIS = + isJoinedMIS.filter((!col(notJoinedMISCol)) && col(isNominated)).select(GraphFrame.ID) + + // update current MIS + val updatedMIS = misDF + .join( + isolatedVertices.select(col(GraphFrame.ID), lit(true).alias("f")), + Seq(GraphFrame.ID), + "left") + .select(col(GraphFrame.ID), (col(isMIS) || col("f")).alias(isMIS)) + .join( + joinedMIS.select(col(GraphFrame.ID), lit(true).alias("f")), + Seq(GraphFrame.ID), + "left") + .select(col(GraphFrame.ID), (col(isMIS) || col("f")).alias(isMIS)) + .persist(storageLevel) + + // We cannot not checkpoint current MIS, otherwise it is almost not working. + if (useLocalCheckpoints) { + val newMis = updatedMIS.localCheckpoint(eager = true) + newMis.count() + misDF.unpersist() + misDF = newMis + } else { + val newMis = updatedMIS.checkpoint(eager = true) + newMis.count() + misDF.unpersist() + misDF = newMis + } + + val neighborsOfMIS = edges + .join(joinedMIS, col(GraphFrame.ID) === col(GraphFrame.DST)) + .select(col(GraphFrame.SRC)) + + val updatedVertices = probs + .join(joinedMIS, Seq(GraphFrame.ID), "left_anti") + .join(neighborsOfMIS, col(GraphFrame.ID) === col(GraphFrame.SRC), "left_anti") + .select(GraphFrame.ID, probCol) + + // checkpointing of vertices + if (shouldCheckpoint && (i % checkpointInterval == 0)) { + if (useLocalCheckpoints) { + vertices = updatedVertices.localCheckpoint(eager = true) + } else { + vertices = updatedVertices.checkpoint(eager = true) + } + } else { + vertices = updatedVertices + } + + // algorithm stops if no more vertex left + converged = vertices.isEmpty + + updatedVertices.unpersist() + probs.unpersist() + + logInfo(s"iteration $i finished, vertices left: ${vertices.count()}") + i += 1 + } + + vertices.unpersist(true) + edges.unpersist(true) + + val mis = misDF.filter(col(isMIS)).select(GraphFrame.ID).persist(storageLevel) + // materialize + mis.count() + resultIsPersistent() + misDF.unpersist(true) + + mis + } finally { + // Restore original AQE setting + spark.conf.set("spark.sql.adaptive.enabled", originalAQE) + } + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/PageRank.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/PageRank.scala new file mode 100644 index 0000000000000..faf92ed4cfcc9 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/PageRank.scala @@ -0,0 +1,181 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.graphx.{lib => graphxlib} +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.Logging + +/** + * PageRank algorithm implementation. There are two implementations of PageRank. + * + * The first one uses the `org.apache.spark.graphx.graph` interface with `aggregateMessages` and + * runs PageRank for a fixed number of iterations. This can be executed by setting `maxIter`. + * Conceptually, the algorithm does the following: + * {{{ + * var PR = Array.fill(n)( 1.0 ) + * val oldPR = Array.fill(n)( 1.0 ) + * for( iter <- 0 until maxIter ) { + * swap(oldPR, PR) + * for( i <- 0 until n ) { + * PR[i] = alpha + (1 - alpha) * inNbrs[i].map(j => oldPR[j] / outDeg[j]).sum + * } + * } + * }}} + * + * The second implementation uses the `org.apache.spark.graphx.Pregel` interface and runs PageRank + * until convergence and this can be run by setting `tol`. Conceptually, the algorithm does the + * following: + * {{{ + * var PR = Array.fill(n)( 1.0 ) + * val oldPR = Array.fill(n)( 0.0 ) + * while( max(abs(PR - oldPr)) > tol ) { + * swap(oldPR, PR) + * for( i <- 0 until n if abs(PR[i] - oldPR[i]) > tol ) { + * PR[i] = alpha + (1 - \alpha) * inNbrs[i].map(j => oldPR[j] / outDeg[j]).sum + * } + * } + * }}} + * + * `alpha` is the random reset probability (typically 0.15), `inNbrs[i]` is the set of neighbors + * which link to `i` and `outDeg[j]` is the out degree of vertex `j`. + * + * Note that this is not the "normalized" PageRank and as a consequence pages that have no inlinks + * will have a PageRank of alpha. In particular, the pageranks may have some values greater than 1. + * + * The resulting vertices DataFrame contains one additional column: + * - pagerank (`DoubleType`): the pagerank of this vertex + * + * The resulting edges DataFrame contains one additional column: + * - weight (`DoubleType`): the normalized weight of this edge after running PageRank + */ +class PageRank private[graphframes] (private val graph: GraphFrame) + extends Arguments + with Logging { + + private var tol: Option[Double] = None + private var resetProb: Option[Double] = Some(0.15) + private var maxIter: Option[Int] = None + private var srcId: Option[Any] = None + + /** Source vertex for a Personalized Page Rank (optional) */ + def sourceId(value: Any): this.type = { + this.srcId = Some(value) + this + } + + /** Reset probability "alpha" */ + def resetProbability(value: Double): this.type = { + resetProb = Some(value) + this + } + + def tol(value: Double): this.type = { + tol = Some(value) + this + } + + def maxIter(value: Int): this.type = { + maxIter = Some(value) + this + } + + def run(): GraphFrame = { + val res = tol match { + case Some(t) => + assert(maxIter.isEmpty, "You cannot specify maxIter() and tol() at the same time.") + PageRank.runUntilConvergence(graph, t, resetProb.get, srcId) + case None => + PageRank.run(graph, check(maxIter, "maxIter"), resetProb.get, srcId) + } + resultIsPersistent() + res + } +} + +// TODO: srcID's type should be checked. The most futureproof check would be Encoder because it is +// compatible with Datasets after that. +private object PageRank { + + /** + * Run PageRank for a fixed number of iterations returning a graph with vertex attributes + * containing the PageRank and edge attributes the normalized edge weight. + * + * @param graph + * the graph on which to compute PageRank + * @param maxIter + * the number of iterations of PageRank to run + * @param resetProb + * the random reset probability (alpha) + * @return + * the graph containing with each vertex containing the PageRank and each edge containing the + * normalized weight. + */ + def run( + graph: GraphFrame, + maxIter: Int, + resetProb: Double = 0.15, + srcId: Option[Any] = None): GraphFrame = { + val longSrcId = srcId.map(GraphXConversions.integralId(graph, _)) + val gx = + graphxlib.PageRank.runWithOptions(graph.cachedTopologyGraphX, maxIter, resetProb, longSrcId) + val res = GraphXConversions + .fromGraphX(graph, gx, vertexNames = Seq(PAGERANK), edgeNames = Seq(WEIGHT)) + .persist() + res.vertices.count() + res.edges.count() + gx.unpersist() + res + } + + /** + * Run a dynamic version of PageRank returning a graph with vertex attributes containing the + * PageRank and edge attributes containing the normalized edge weight. + * + * @param graph + * the graph on which to compute PageRank + * @param tol + * the tolerance allowed at convergence (smaller => more accurate). + * @param resetProb + * the random reset probability (alpha) + * @param srcId + * the source vertex for a Personalized Page Rank (optional) + * @return + * the graph containing with each vertex containing the PageRank and each edge containing the + * normalized weight. + */ + def runUntilConvergence( + graph: GraphFrame, + tol: Double, + resetProb: Double = 0.15, + srcId: Option[Any] = None): GraphFrame = { + val longSrcId = srcId.map(GraphXConversions.integralId(graph, _)) + val gx = graphxlib.PageRank.runUntilConvergenceWithOptions( + graph.cachedTopologyGraphX, + tol, + resetProb, + longSrcId) + GraphXConversions.fromGraphX(graph, gx, vertexNames = Seq(PAGERANK), edgeNames = Seq(WEIGHT)) + } + + /** Default name for the pagerank column. */ + private val PAGERANK = "pagerank" + + /** Default name for the weight column. */ + private val WEIGHT = "weight" +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRank.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRank.scala new file mode 100644 index 0000000000000..6e183eae9ef24 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRank.scala @@ -0,0 +1,129 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.graphx +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithMaxIter + +/** + * Parallel Personalized PageRank algorithm implementation. + * + * This implementation uses the standalone [[GraphFrame]] interface and runs personalized PageRank + * in parallel for a fixed number of iterations. This can be run by setting `maxIter`. The source + * vertex Ids are set in `sourceIds`. A simple local implementation of this algorithm is as + * follows. + * {{{ + * var oldPR = Array.fill(n)( 1.0 ) + * val PR = (0 until n).map(i => if sourceIds.contains(i) alpha else 0.0) + * for( iter <- 0 until maxIter ) { + * swap(oldPR, PR) + * for( i <- 0 until n ) { + * PR[i] = (1 - alpha) * inNbrs[i].map(j => oldPR[j] / outDeg[j]).sum + * if (sourceIds.contains(i)) PR[i] += alpha + * } + * } + * }}} + * + * `alpha` is the random reset probability (typically 0.15), `inNbrs[i]` is the set of neighbors + * which link to `i` and `outDeg[j]` is the out degree of vertex `j`. + * + * Note that this is not the "normalized" PageRank and as a consequence pages that have no inlinks + * will have a PageRank of alpha. In particular, the pageranks may have some values greater than 1. + * + * The resulting vertices DataFrame contains one additional column: + * - pageranks (`VectorType`): the pageranks of this vertex from all input source vertices + * + * The resulting edges DataFrame contains one additional column: + * - weight (`DoubleType`): the normalized weight of this edge after running PageRank + */ +class ParallelPersonalizedPageRank private[graphframes] (private val graph: GraphFrame) + extends Arguments + with WithMaxIter + with Logging { + + private var resetProb: Option[Double] = Some(0.15) + private var srcIds: Array[Any] = Array() + + /** Source vertices for a Personalized Page Rank */ + def sourceIds(values: Array[Any]): this.type = { + this.srcIds = values + this + } + + /** Reset probability "alpha" */ + def resetProbability(value: Double): this.type = { + resetProb = Some(value) + this + } + + def run(): GraphFrame = { + require(maxIter != None, "Max number of iterations maxIter() must be provided") + require(srcIds.nonEmpty, "Source vertices Ids sourceIds() must be provided") + val res = ParallelPersonalizedPageRank.run(graph, maxIter.get, resetProb.get, srcIds) + resultIsPersistent() + res + } +} + +private object ParallelPersonalizedPageRank { + + /** Default name for the pageranks column. */ + private val PAGERANKS = "pageranks" + + /** Default name for the weight column. */ + private val WEIGHT = "weight" + + /** + * Run Personalized PageRank for a fixed number of iterations, for a set of starting nodes in + * parallel. Returns a graph with vertex attributes containing the pageranks relative to all + * starting nodes (as a vector) and edge attributes the normalized edge weight + * + * @param graph + * The graph on which to compute personalized pagerank + * @param maxIter + * The number of iterations to run + * @param resetProb + * The random reset probability + * @param sourceIds + * The list of sources to compute personalized pagerank from + * @return + * the graph with vertex attributes containing the pageranks relative to all starting nodes as + * a vector and edge attributes the normalized edge weight + */ + def run( + graph: GraphFrame, + maxIter: Int, + resetProb: Double, + sourceIds: Array[Any]): GraphFrame = { + val longSrcIds = sourceIds.map(GraphXConversions.integralId(graph, _)) + val gx = graphx.lib.PageRank.runParallelPersonalizedPageRank( + graph.cachedTopologyGraphX, + maxIter, + resetProb, + longSrcIds) + val gf = GraphXConversions + .fromGraphX(graph, gx, vertexNames = Seq(PAGERANKS), edgeNames = Seq(WEIGHT)) + .persist() + gf.vertices.count() + gf.edges.count() + gx.unpersist() + gf + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/Pregel.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/Pregel.scala new file mode 100644 index 0000000000000..3d4d492bea7c8 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/Pregel.scala @@ -0,0 +1,658 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.array +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.explode +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.functions.struct +import org.apache.spark.sql.graphframes.GraphFrameInternals +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrame._ +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLocalCheckpoints + +import java.io.IOException +import scala.util.control.Breaks.break +import scala.util.control.Breaks.breakable + +/** + * Implements a Pregel-like bulk-synchronous message-passing API based on DataFrame operations. + * + * See Malewicz et al., Pregel: a system for + * large-scale graph processing for a detailed description of the Pregel algorithm. + * + * You can construct a Pregel instance using either this constructor or + * [[org.apache.spark.graphframes.GraphFrame#pregel]], then use builder pattern to describe the + * operations, and then call [[run]] to start a run. It returns a DataFrame of vertices from the + * last iteration. + * + * When a run starts, it expands the vertices DataFrame using column expressions defined by + * [[withVertexColumn]]. Those additional vertex properties can be changed during Pregel + * iterations. In each Pregel iteration, there are three phases: + * - Given each edge triplet, generate messages and specify target vertices to send, described + * by [[sendMsgToDst]] and [[sendMsgToSrc]]. + * - Aggregate messages by target vertex IDs, described by [[aggMsgs]]. + * - Update additional vertex properties based on aggregated messages and states from previous + * iteration, described by [[withVertexColumn]]. + * + * Please find what columns you can reference at each phase in the method API docs. + * + * You can control the number of iterations by [[setMaxIter]] and check API docs for advanced + * controls. + * + * Example code for Page Rank: + * + * {{{ + * val edges = ... + * val vertices = GraphFrame.fromEdges(edges).outDegrees.cache() + * val numVertices = vertices.count() + * val graph = GraphFrame(vertices, edges) + * val alpha = 0.15 + * val ranks = graph.pregel + * .withVertexColumn("rank", lit(1.0 / numVertices), + * coalesce(Pregel.msg, lit(0.0)) * (1.0 - alpha) + alpha / numVertices) + * .sendMsgToDst(Pregel.src("rank") / Pregel.src("outDegree")) + * .aggMsgs(sum(Pregel.msg)) + * .run() + * }}} + * + * Migration note: pre 0.12 users that used edge columns in Pregel expressions should explicitly + * specify these columns using [[org.apache.spark.graphframes.lib.Pregel#requiredDstColumns]]. In + * 0.11 and earlier there was an unspecified bug that leads all the edge columns are always kept + * and persisted that created a bug memory pressure (2 columns in O(|E|) rows in the form of + * `StructType`). That behavior is considered as bug and starting from 0.12 edge columns are not + * kept by default. + * + * @param graph + * The graph that Pregel will run on. + * @see + * [[org.apache.spark.graphframes.GraphFrame#pregel]] + * @see + * Malewicz et al., Pregel: a system for + * large-scale graph processing. + */ +class Pregel(val graph: GraphFrame) + extends Logging + with WithLocalCheckpoints + with WithIntermediateStorageLevel { + + private val withVertexColumnList = collection.mutable.ListBuffer.empty[(String, Column, Column)] + + private var maxIter: Int = 10 + private var checkpointInterval = 2 + private var earlyStopping = false + private var stopIfAllNonActiveVertices = false + private var skipMessagesFromNonActiveVertices = false + private var initialActiveVertexExpression = lit(true) + private var updateActiveVertexExpression = lit(true) + + private val sendMsgs = collection.mutable.ListBuffer.empty[(Column, Column)] + private var aggMsgsCol: Column = null + + // Required columns for source and destination vertices in triplets + // When empty, all columns are selected (default behavior) + private val requiredSrcColumnsList = collection.mutable.ListBuffer.empty[String] + private val requiredDstColumnsList = collection.mutable.ListBuffer.empty[String] + + // Required columns for edges + // When empty, only src and dst are selected + private val requiredEdgeColumnsList = collection.mutable.ListBuffer.empty[String] + + /** Sets the max number of iterations (default: 10). */ + def setMaxIter(value: Int): this.type = { + maxIter = value + this + } + + /** + * Sets the number of iterations between two checkpoints (default: 2). + * + * This is an advanced control to balance query plan optimization and checkpoint data I/O cost. + * In most cases, you should keep the default value. + * + * Checkpoint is disabled if this is set to 0. + */ + def setCheckpointInterval(value: Int): this.type = { + checkpointInterval = value + this + } + + /** + * Should Pregel stop earlier in case of no new messages to send? + * + * Early stopping allows to terminate Pregel before reaching maxIter by checking is there any + * non-null message or not. While in some cases it may gain significant performance boost, it + * other cases it can tend to performance degradation, because checking is messages DataFrame is + * empty or not is an action and requires materialization of the Spark Plan with some additional + * computations. + * + * In the case when user can assume a good value of maxIter it is recommended to leave this + * value to the default "false". In the case when it is hard to estimate an amount of iterations + * required for convergence, it is recommended to set this value to "false" to avoid iterating + * over convergence until reaching maxIter. When this value is "true", maxIter can be set to a + * bigger value without risks. + * + * @param value + * should Pregel checks for the termination condition on each step + * @return + */ + def setEarlyStopping(value: Boolean): this.type = { + earlyStopping = value + this + } + + /** + * Should Pregel stop earlier in case all the vertices are marked as non active. + * + * This feature allows to terminate Pregel before reaching maxIter by checking are there active + * vertex left. A good example of activity check is PageRank: (see Malewicz, Grzegorz, et al. + * "Pregel: a system for large-scale graph processing." Proceedings of the 2010 ACM SIGMOD + * International Conference on Management of data. 2010., a part about voting to halt) + * - after each iteration we are checking is the change in rank less than tolerance and if so, + * we can mark vertex as non active + * - if all the vertices are non active, we stop iterations + * @param value + * should Pregel stop earlier by vertices voting + * @return + */ + def setStopIfAllNonActiveVertices(value: Boolean): this.type = { + stopIfAllNonActiveVertices = value + this + } + + /** + * Set the initial expression for the active/non-active flag per vertex. + * + * In most of the cases the default expression (true for all the vertices) should works fine. + * For some cases it makes sense to set a custom expression. A good example is + * multiple-landmarks shortest-paths algorithm: + * - the only initially active vertices in that case should be landmarkds, because only this + * vertices initially have non-null distances but all the other vertices have null distances + * and there is no reason to mark them active initially. + * @param expression + * an initial expression that will be used to create an active-flag vertex column + * @return + */ + def setInitialActiveVertexExpression(expression: Column): this.type = { + initialActiveVertexExpression = expression + this + } + + /** + * Set an expression that will be used after each superstep to update the active-flag vertex + * column. + * + * An example is PageRank algorithm: in that case such an expression may looks like abs(old_rank - + * new_rank) >= tolerance + * + * @param expression + * an expression, that will be used after each superstep to update the active-flag vertex + * column + * @return + */ + def setUpdateActiveVertexExpression(expression: Column): this.type = { + updateActiveVertexExpression = expression + this + } + + /** + * With a true value, Pregel will not generate messages from vertices, marked as non active. + * + * For example, for Shortest Paths, there is no reason to pass distances from vertices, for that + * these distances did not change at the latest iteration. It allows significantly reduce an + * amount of generated messages. + * + * Be careful, for algorithms like Label Propagation or Pregel, even if the vertex is not + * active, we still need to generate messages, otherwise algorithm will return an incorrect + * result! + * + * @param value + * should Pregel skip generation of messages for non active vertices. + * @return + */ + def setSkipMessagesFromNonActiveVertices(value: Boolean): this.type = { + skipMessagesFromNonActiveVertices = value + this + } + + /** + * Defines an additional vertex column at the start of run and how to update it in each + * iteration. + * + * You can call it multiple times to add more than one additional vertex columns. + * + * @param colName + * the name of the additional vertex column. It cannot be an existing vertex column in the + * graph. + * @param initialExpr + * the expression to initialize the additional vertex column. You can reference all original + * vertex columns in this expression. + * @param updateAfterAggMsgsExpr + * the expression to update the additional vertex column after messages aggregation. You can + * reference all original vertex columns, additional vertex columns, and the aggregated + * message column using [[Pregel$#msg]]. If the vertex received no messages, the message + * column would be null. + */ + def withVertexColumn( + colName: String, + initialExpr: Column, + updateAfterAggMsgsExpr: Column): this.type = { + // TODO: check if this column exists. + require( + colName != null && colName != ID && colName != Pregel.MSG_COL_NAME, + "additional column name cannot be null and cannot be the same name with ID column or " + + "msg column.") + require(initialExpr != null, "additional column should provide a nonnull initial expression.") + require( + updateAfterAggMsgsExpr != null, + "additional column should provide a nonnull " + + "updateAfterAggMsgs expression.") + withVertexColumnList += Tuple3(colName, initialExpr, updateAfterAggMsgsExpr) + this + } + + /** + * Defines a message to send to the source vertex of each edge triplet. + * + * You can call it multiple times to send more than one messages. + * + * @param msgExpr + * the expression of the message to send to the source vertex given a (src, edge, dst) + * triplet. Source/destination vertex properties and edge properties are nested under columns + * `src`, `dst`, and `edge`, respectively. You can reference them using [[Pregel$#src]], + * [[Pregel$#dst]], and [[Pregel$#edge]]. Null messages are not included in message + * aggregation. + * @see + * [[sendMsgToDst]] + */ + def sendMsgToSrc(msgExpr: Column): this.type = { + sendMsgs += Tuple2(Pregel.src(ID), msgExpr) + this + } + + /** + * Defines a message to send to the destination vertex of each edge triplet. + * + * You can call it multiple times to send more than one messages. + * + * @param msgExpr + * the message expression to send to the destination vertex given a (`src`, `edge`, `dst`) + * triplet. Source/destination vertex properties and edge properties are nested under columns + * `src`, `dst`, and `edge`, respectively. You can reference them using [[Pregel$#src]], + * [[Pregel$#dst]], and [[Pregel$#edge]]. Null messages are not included in message + * aggregation. + * @see + * [[sendMsgToSrc]] + */ + def sendMsgToDst(msgExpr: Column): this.type = { + sendMsgs += Tuple2(Pregel.dst(ID), msgExpr) + this + } + + /** + * Specifies which source vertex columns are required when constructing triplets. + * + * By default, all source vertex columns are included in triplets, which can create large + * intermediate datasets for algorithms with significant state (e.g., cycle detection, random + * walks). Use this method to reduce memory usage by specifying only the columns that are + * actually needed by the sendMsgToSrc and sendMsgToDst expressions. + * + * The ID column and the active flag column (if used) are always included automatically. + * + * @param colName + * the first required source vertex column name + * @param colNames + * additional required source vertex column names + * @see + * [[requiredDstColumns]] + */ + def requiredSrcColumns(colName: String, colNames: String*): this.type = { + requiredSrcColumnsList.clear() + requiredSrcColumnsList += colName + requiredSrcColumnsList ++= colNames + this + } + + /** + * Specifies which destination vertex columns are required when constructing triplets. + * + * By default, all destination vertex columns are included in triplets, which can create large + * intermediate datasets for algorithms with significant state (e.g., cycle detection, random + * walks). Use this method to reduce memory usage by specifying only the columns that are + * actually needed by the sendMsgToSrc and sendMsgToDst expressions. + * + * The ID column and the active flag column (if used) are always included automatically. + * + * @param colName + * the first required destination vertex column name + * @param colNames + * additional required destination vertex column names + * @see + * [[requiredSrcColumns]] + */ + def requiredDstColumns(colName: String, colNames: String*): this.type = { + requiredDstColumnsList.clear() + requiredDstColumnsList += colName + requiredDstColumnsList ++= colNames + this + } + + /** + * Specifies which edge columns are required when constructing triplets. + * + * By default, only the source and destination ID columns from edges are included in triplets. + * Use this method to include additional edge properties that are needed by the sendMsgToSrc and + * sendMsgToDst expressions. + * + * @param colName + * the first required edge column name + * @param colNames + * additional required edge column names + * @see + * [[requiredSrcColumns]] and [[requiredDstColumns]] + */ + def requiredEdgeColumns(colName: String, colNames: String*): this.type = { + requiredEdgeColumnsList.clear() + requiredEdgeColumnsList += colName + requiredEdgeColumnsList ++= colNames + this + } + + /** + * Defines how messages are aggregated after grouped by target vertex IDs. + * + * @param aggExpr + * the message aggregation expression, such as `sum(Pregel.msg)`. You can reference the + * message column by [[Pregel$#msg]] and the vertex ID by [[GraphFrame$#ID]], while the latter + * is usually not used. + */ + def aggMsgs(aggExpr: Column): this.type = { + aggMsgsCol = aggExpr + this + } + + /** + * Runs the defined Pregel algorithm. + * + * @return + * the result vertex DataFrame from the final iteration including both original and additional + * columns. + */ + def run(): DataFrame = { + require( + sendMsgs.length > 0, + "We need to set at least one message expression for pregel running.") + require(aggMsgsCol != null, "We need to set aggMsgs for pregel running.") + require(maxIter >= 1, "The max iteration number should be >= 1.") + require( + checkpointInterval >= 0, + "The checkpoint interval should be >= 0, 0 indicates no checkpoint.") + require( + withVertexColumnList.size > 0, + "There should be at least one additional vertex columns for updating.") + + val sendMsgsColList = sendMsgs.toList.map { case (id, msg) => + struct(id.as(ID), msg.as("msg")) + } + + val initVertexCols = withVertexColumnList.toList.map { case (colName, initExpr, _) => + initExpr.as(colName) + } + val updateVertexCols = withVertexColumnList.toList.map { case (colName, _, updateExpr) => + updateExpr.as(colName) + } + + var lastRoundPersistent: scala.collection.mutable.Queue[DataFrame] = + scala.collection.mutable.Queue[DataFrame]() + + val initialAttributes = graph.vertices.columns.map(col).toSeq + + var currentVertices = graph.vertices.select( + ((initialAttributes :+ initialActiveVertexExpression.alias( + Pregel.ACTIVE_FLAG_COL)) ++ initVertexCols): _*) + + // Automatic optimization: detect if destination vertex state is needed by analyzing + // the MESSAGE expressions only (not the target ID expressions, since dst.id is always + // available from the edge). If no message expression references dst.* columns, + // we can skip the second join entirely. + // Additionally, if the only dst field referenced is "id", we can still skip since + // dst.id is available from the edge's dst column. + val messageExpressions = sendMsgs.toList.map { case (_, msgExpr) => msgExpr } + val allDstRefs = messageExpressions.flatMap { expr => + GraphFrameInternals.extractColumnReferences(graph.spark, expr).get(DST) + } + val dstPrefixReferenced = allDstRefs.nonEmpty + val dstFieldsReferenced = allDstRefs.flatten.toSet + + // We need the dst join if dst is referenced AND fields other than just "id" are accessed + val needsDstState = + dstPrefixReferenced && (dstFieldsReferenced.isEmpty || dstFieldsReferenced != Set(ID)) + if (!needsDstState) { + logDebug( + "Optimization: skipping second join (dst state not required by message expressions)") + } + + val edges = (if (requiredEdgeColumnsList.isEmpty) { + graph.edges + .select(col(SRC).alias("edge_src"), col(DST).alias("edge_dst")) + } else { + graph.edges + .select( + col(SRC).alias("edge_src"), + col(DST).alias("edge_dst"), + struct( + requiredEdgeColumnsList.head, + requiredEdgeColumnsList.tail.toSeq: _*).as(EDGE)) + }).repartition(col("edge_src")).persist(intermediateStorageLevel) + + var iteration = 1 + + val shouldCheckpoint = checkpointInterval > 0 + + if (shouldCheckpoint && graph.spark.sparkContext.getCheckpointDir.isEmpty && !useLocalCheckpoints) { + // Spark Connect workaround + graph.spark.conf.getOption("spark.checkpoint.dir") match { + case Some(d) => graph.spark.sparkContext.setCheckpointDir(d) + case None => + throw new IOException( + "Checkpoint directory is not set. Please set it first using sc.setCheckpointDir()" + + "or by specifying the conf 'spark.checkpoint.dir'.") + } + } + + // Columns to include in triplet structs (ID + active flag always included if specified) + val srcCols = + if (requiredSrcColumnsList.isEmpty) Seq(col("*")) + else (Seq(ID, Pregel.ACTIVE_FLAG_COL) ++ requiredSrcColumnsList).distinct.map(col) + val dstCols = + if (requiredDstColumnsList.isEmpty) Seq(col("*")) + else (Seq(ID, Pregel.ACTIVE_FLAG_COL) ++ requiredDstColumnsList).distinct.map(col) + + breakable { + while (iteration <= maxIter) { + logInfo(s"start Pregel iteration $iteration / $maxIter") + val currRoundPersistent = scala.collection.mutable.Queue[DataFrame]() + currRoundPersistent.enqueue(currentVertices.persist(intermediateStorageLevel)) + + // Prune non-active vertices early if skipMessagesFromNonActiveVertices + // is enabled and we don't need the dst state. + val srcVertices = + if (!needsDstState && skipMessagesFromNonActiveVertices) + currentVertices.filter(col(Pregel.ACTIVE_FLAG_COL)) + else currentVertices + + // Build triplets: start with src vertex state joined with edges + val srcWithEdges = srcVertices + .select(struct(srcCols: _*).as(SRC)) + .join(edges, Pregel.src(ID) === col("edge_src")) + + // Only perform the second join (adding dst vertex state) if needed + var tripletsDF = if (needsDstState) { + srcWithEdges + .join( + currentVertices.select(struct(dstCols: _*).as(DST)), + col("edge_dst") === Pregel.dst(ID)) + .drop(col("edge_src"), col("edge_dst")) + } else { + // Skip second join - dst state not needed by any message expression. + // Create a minimal dst struct with just the id from edge_dst for sendMsgToDst to work. + srcWithEdges + .withColumn(DST, struct(col("edge_dst").as(ID))) + .drop(col("edge_src"), col("edge_dst")) + } + + // Only prune here if we didn't prune above. + if (needsDstState && skipMessagesFromNonActiveVertices) { + tripletsDF = tripletsDF.filter( + Pregel.src(Pregel.ACTIVE_FLAG_COL) || Pregel.dst(Pregel.ACTIVE_FLAG_COL)) + } + + val msgDF: DataFrame = tripletsDF + .select(explode(array(sendMsgsColList: _*)).as("msg")) + .select(col("msg.id"), col("msg.msg").as(Pregel.MSG_COL_NAME)) + .filter(Pregel.msg.isNotNull) + + if (earlyStopping && msgDF.isEmpty) { + logInfo( + s"there are no more non-null messages; Pregel stops earlier at iteration $iteration") + while (lastRoundPersistent.nonEmpty) { + lastRoundPersistent.dequeue().unpersist() + } + lastRoundPersistent = currRoundPersistent + break() + } + + val newAggMsgDF = msgDF + .groupBy(ID) + .agg(aggMsgsCol.as(Pregel.MSG_COL_NAME)) + + val verticesWithMsg = currentVertices.join(newAggMsgDF, Seq(ID), "left_outer") + + currentVertices = verticesWithMsg.select( + ((initialAttributes :+ updateActiveVertexExpression.alias( + Pregel.ACTIVE_FLAG_COL)) ++ updateVertexCols): _*) + + if (shouldCheckpoint && iteration % checkpointInterval == 0) { + if (useLocalCheckpoints) { + currentVertices = currentVertices.localCheckpoint(eager = false) + } else { + currentVertices = currentVertices.checkpoint(eager = false) + } + } else { + // checkpointing do persistence and we do not need to do it again + currRoundPersistent.enqueue(currentVertices.persist(intermediateStorageLevel)) + } + + if (stopIfAllNonActiveVertices) { + if (currentVertices.filter(col(Pregel.ACTIVE_FLAG_COL)).isEmpty) { + logInfo( + s"all the verties are non-active; Pregel stops earlier at iteration $iteration") + while (lastRoundPersistent.nonEmpty) { + lastRoundPersistent.dequeue().unpersist() + } + lastRoundPersistent = currRoundPersistent + break() + } + } + + if (!earlyStopping && !stopIfAllNonActiveVertices) { + // we need to call materialize + currentVertices.count() + } + + while (lastRoundPersistent.nonEmpty) { + lastRoundPersistent.dequeue().unpersist() + } + lastRoundPersistent = currRoundPersistent + + iteration += 1 + } + } + + val res = currentVertices.persist(intermediateStorageLevel) + res.count() + while (lastRoundPersistent.nonEmpty) { + lastRoundPersistent.dequeue().unpersist() + } + edges.unpersist() + System.gc() + res + } + +} + +/** + * Constants and utilities for the Pregel algorithm. + */ +object Pregel extends Serializable { + + /** + * A constant column name for generated and aggregated messages. + * + * The vertices DataFrame must not contain this column. + */ + val MSG_COL_NAME = "_pregel_msg_" + + /** + * A constant column name for active vertex flag. + */ + val ACTIVE_FLAG_COL = "_pregel_is_active" + + /** + * References the message column in aggregating messages and updating additional vertex columns. + * + * @see + * [[Pregel.aggMsgs]] and [[Pregel.withVertexColumn]] + */ + val msg: Column = col(MSG_COL_NAME) + + /** + * References a source vertex column in generating messages to send. + * + * @param colName + * the vertex column name. + * @see + * [[Pregel.sendMsgToSrc]] and [[Pregel.sendMsgToDst]] + */ + def src(colName: String): Column = col(GraphFrame.SRC + "." + colName) + + /** + * References a destination vertex column in generating messages to send. + * + * @param colName + * the vertex column name. + * @see + * [[Pregel.sendMsgToSrc]] and [[Pregel.sendMsgToDst]] + */ + def dst(colName: String): Column = col(GraphFrame.DST + "." + colName) + + /** + * References an edge column in generating messages to send. + * + * @param colName + * the edge column name. + * @see + * [[Pregel.sendMsgToSrc]] and [[Pregel.sendMsgToDst]] + */ + def edge(colName: String): Column = col(GraphFrame.EDGE + "." + colName) +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/RandomizedContraction.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/RandomizedContraction.scala new file mode 100644 index 0000000000000..e42c1cba259fe --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/RandomizedContraction.scala @@ -0,0 +1,291 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.hadoop.fs.Path +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.catalyst.FunctionIdentifier +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.graphframes.expressions.FiniteAXPlusB +import org.apache.spark.storage.StorageLevel +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrame.DST +import org.apache.spark.graphframes.GraphFrame.ID +import org.apache.spark.graphframes.GraphFrame.LONG_DST +import org.apache.spark.graphframes.GraphFrame.LONG_ID +import org.apache.spark.graphframes.GraphFrame.LONG_SRC +import org.apache.spark.graphframes.GraphFrame.SRC +import org.apache.spark.graphframes.Logging + +import java.io.IOException +import java.util.UUID +import scala.collection.mutable.Stack +import scala.util.Random + +/** + * Implementation of parallel connected components algorithm using randomized contraction, based + * on Bögeholz, Harald, Michael Brand, and Radu-Alexandru Todor. "In-database connected component + * analysis." 2020 IEEE 36th International Conference on Data Engineering (ICDE). IEEE, 2020. + * + * The algorithm contracts the graph iteratively using random linear functions, until no edges + * remain, then reconstructs the component identifiers. + */ +private[graphframes] object RandomizedContraction extends Logging with Serializable { + private val CHECKPOINT_NAME_PREFIX = "randomized-contraction" + + private def prepare(graph: GraphFrame): GraphFrame = { + val vertices = graph.indexedVertices + .select(col(LONG_ID).as(ID)) + + val edges = graph.indexedEdges + .select(col(LONG_SRC).as(SRC), col(LONG_DST).as(DST)) + val symmetricEdges = edges + .union(edges.select(col(DST).alias(SRC), col(SRC).alias(DST))) + .distinct() + GraphFrame(vertices, symmetricEdges) + + } + + def run( + inputGraph: GraphFrame, + useLabelsAsComponents: Boolean, + intermediateStorageLevel: StorageLevel, + useLocalCheckpoints: Boolean, + checkpointInterval: Int, + isGraphPrepared: Boolean): DataFrame = { + val spark = inputGraph.vertices.sparkSession + val sc = spark.sparkContext + val runId = UUID.randomUUID().toString.takeRight(8) + val logPrefix = s"[CC $runId]" + + val checkpointDir = sc.getCheckpointDir + .map { d => + new Path(d, s"$CHECKPOINT_NAME_PREFIX-$runId").toString + } + .getOrElse { + // Spark-Connect workaround + spark.conf.getOption("spark.checkpoint.dir") match { + case Some(d) => new Path(d, s"$CHECKPOINT_NAME_PREFIX-$runId").toString + case None => + throw new IOException( + "Checkpoint directory is not set. Please set it first using sc.setCheckpointDir()" + + "or by specifying the conf 'spark.checkpoint.dir'.") + } + } + logInfo(s"$logPrefix Using $checkpointDir for storing intermediate tables.") + + val functionRegistry = spark.sessionState.functionRegistry + functionRegistry.registerFunction( + new FunctionIdentifier("_axpb", Some("builtin"), Some("system")), + (children: Seq[Expression]) => FiniteAXPlusB(children(0), children(1), children(2)), + "scala_udf") + + val random = new Random() + random.setSeed(42L) + val stackA = Stack.empty[Long] + val stackB = Stack.empty[Long] + var iter = 0 + + def tableName(iter: Int): String = s"${checkpointDir}/ccreps-${iter}" + + val graph = if (isGraphPrepared) { + inputGraph + } else { + prepare(inputGraph) + } + + var edges = + graph.edges.select(SRC, DST).persist(intermediateStorageLevel) + + def axpb(a: Long, x: Column, b: Long): Column = call_function("_axpb", lit(a), x, lit(b)) + + try { + var rA = 0L + var graphSize = edges.count() + var ccRepresentatives: DataFrame = null + + // "no edges graph" + if (graphSize == 0L) { + val result = inputGraph.vertices + .select(col(ID), col(ID).alias(ConnectedComponents.COMPONENT)) + .persist(intermediateStorageLevel) + result.count() + edges.unpersist() + + return result + } + + while (graphSize > 0) { + logInfo(s"iteration ${iter}, edges left ${graphSize}") + iter += 1 + rA = 0L + while (rA == 0L) { + rA = random.nextLong() + } + val rB = random.nextLong() + stackA.push(rA) + stackB.push(rB) + + ccRepresentatives = edges + .groupBy(SRC) + .agg(min(axpb(rA, col(DST), rB)).alias("rep")) + .select(col(SRC).alias("v"), least(axpb(rA, col(SRC), rB), col("rep")).alias("rep")) + + // "free" checkpointing + ccRepresentatives.write.parquet(tableName(iter)) + ccRepresentatives = spark.read.parquet(tableName(iter)) + + val edges2 = edges + .join(ccRepresentatives, col(SRC) === col("v")) + .select(col("rep").alias(SRC), col(DST)) + + // save ref to unpersist + val oldEdges = edges + + edges = { + val te = edges2 + .alias("e") + .join( + ccRepresentatives.alias("r2"), + col(s"e.$DST") === col("r2.v") && + col(s"e.$SRC") =!= col("r2.rep")) + .select(col(s"e.$SRC").alias(SRC), col("r2.rep").alias(DST)) + .distinct() + + if ((iter > 0) && (iter % checkpointInterval == 0)) { + if (useLocalCheckpoints) { + te.localCheckpoint() + } else { + te.checkpoint() + } + } else { + te.persist(intermediateStorageLevel) + } + } + + graphSize = edges.count() + oldEdges.unpersist() + } + + logInfo(s"graph was successfully contracted for $iter iterations") + logInfo("start reverse transformation") + + var accA = 1L + var accB = 0L + + edges.unpersist(true) + + while (iter > 1) { + iter -= 1 + val poppedA = stackA.pop() + val poppedB = stackB.pop() + + val oldAccA = accA + accA = FiniteAXPlusB.axpb(oldAccA, poppedA, 0L) + accB = FiniteAXPlusB.axpb(oldAccA, poppedB, accB) + + val ccRepsR = tableName(iter) + val ccRepsR1 = tableName(iter + 1) + + val result = spark.read + .parquet(ccRepsR) + .alias("r1") + .join( + spark.read.parquet(ccRepsR1).alias("r2"), + col("r1.rep") === col("r2.v"), + "left_outer") + .select( + col("r1.v"), + coalesce(col("r2.rep"), axpb(accA, col("r1.rep"), accB)).alias("rep")) + .persist(intermediateStorageLevel) + + result.write.mode("overwrite").parquet(ccRepsR) + + result.unpersist() + val oldPath = new Path(ccRepsR1) + val fs = oldPath.getFileSystem(sc.hadoopConfiguration) + + if (fs.exists(oldPath)) { + fs.delete(oldPath, true) + } + } + + val finalReps = spark.read + .parquet(tableName(1)) + .select(col("v").alias(ID), col("rep").alias(ConnectedComponents.COMPONENT)) + + val outputComponents = if (useLabelsAsComponents && (!inputGraph.hasIntegralIdType)) { + val labels = inputGraph.indexedVertices + .withColumnRenamed(ID, ConnectedComponents.ORIG_ID) + .join(finalReps, col(ID) === col(LONG_ID)) + .groupBy(ConnectedComponents.COMPONENT) + .agg(min(ConnectedComponents.ORIG_ID).alias("new_component")) + + inputGraph.indexedVertices + .withColumnRenamed(ID, ConnectedComponents.ORIG_ID) + .join(finalReps, col(ID) === col(LONG_ID), "left") + .join(labels, ConnectedComponents.COMPONENT, "left") + .select( + col(ConnectedComponents.ORIG_ID).alias(ID), + coalesce(col("new_component"), col(ConnectedComponents.ORIG_ID)) + .alias(ConnectedComponents.COMPONENT)) + } else if (useLabelsAsComponents) { + val labels = + finalReps.groupBy(ConnectedComponents.COMPONENT).agg(min(ID).alias("new_component")) + inputGraph.vertices + .join(finalReps, ID, "left") + .join(labels, ConnectedComponents.COMPONENT, "left") + .select( + col(ID), + coalesce(col("new_component"), col(ID)) + .alias(ConnectedComponents.COMPONENT)) + } else { + inputGraph.vertices + .join(finalReps, ID, "left") + .select( + col(ID), + coalesce(col(ConnectedComponents.COMPONENT), col(ID)) + .alias(ConnectedComponents.COMPONENT)) + } + + outputComponents.persist(intermediateStorageLevel) + // materialize to be able to clean up everything + outputComponents.count() + + // clean-up + val chDirPath = new Path(checkpointDir) + val fs = chDirPath.getFileSystem(sc.hadoopConfiguration) + if (fs.exists(chDirPath)) { + fs.delete(chDirPath, true) + } + + outputComponents + } finally { + // to be 100% sure; + edges.unpersist() + val dereg = functionRegistry.dropFunction( + new FunctionIdentifier("_axpb", Some("builtin"), Some("system"))) + if (!dereg) { + logWarn( + "graphframes faced an internal error and was not able to de-register function _axpb; Spark' functionRegistry is in a bad state") + } + } + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/SVDPlusPlus.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/SVDPlusPlus.scala new file mode 100644 index 0000000000000..5b61a7393c54d --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/SVDPlusPlus.scala @@ -0,0 +1,257 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.graphx.Edge +import org.apache.spark.graphx.{lib => graphxlib} +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.col +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.InvalidGraphException +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithMaxIter + +/** + * Arguments for SVD++ algorithm. + * + * This class implements the SVD++ algorithm for Collaborative Filtering, primarily used for + * Recommender Systems (Link Prediction). + * + * Based on the paper "Factorization Meets the Neighborhood: a Multifaceted Collaborative + * Filtering Model" by Yehuda Koren (2008), available at + * [[https://dl.acm.org/citation.cfm?id=1401944]]. + * + * ==Problem Definition== + * The algorithm predicts unknown ratings in a user-item system. It accounts for: + * - Explicit preferences (user ratings). + * - Implicit feedback (the history of items a user has interacted with). + * - User and Item biases. + * + * The prediction rule for a rating `r_ui` (user `u`, item `i`) is: + * {{{ + * r_ui = µ + b_u + b_i + q_i^T * (p_u + |N(u)|^-0.5 * sum(y_j for j in N(u))) + * }}} + * Where `N(u)` is the set of items user `u` has interacted with (implicit feedback). + * + * ==Input Requirements== + * !!! IMPORTANT !!! The input graph MUST be a **Directed Bipartite Graph** representing + * interactions: + * - **Vertices**: A mix of Users and Items. + * - **Edges**: Directed strictly from **User (src) -> Item (dst)**. + * - **Edge Attribute**: A numeric column (default "weight") representing the rating. + * + * DO NOT use this on general/undirected graphs (e.g., social networks), as the algorithm relies + * on the asymmetry between Users (who provide feedback) and Items (who receive it). + * + * ==Output Model (Node Embeddings)== + * The algorithm returns a DataFrame of vertices with the trained model parameters. These + * parameters function as embeddings: + * + * - `column1` (Array[Double]): **Primary Latent Factors (Explicit Embedding)**. + * - For Users: Represents preferences (`p_u`). + * - For Items: Represents characteristics (`q_i`). + * - `column2` (Array[Double]): **Implicit Factors (Implicit Embedding)**. + * - For Items: Represents the influence of the item (`y_i`) on a user's profile based on + * viewing history. + * - For Users: Generally unused/zero. + * - `column3` (Double): **Bias**. + * - For Users: User bias (`b_u`). + * - For Items: Item bias (`b_i`). + * - `column4` (Double): **Implicit Normalization Term**. + * - For Users: Precomputed `|N(u)|^-0.5`. + * - For Items: Unused. + * + * ==Parameter Tuning Guide== + * + * Constraints: + * - `minValue` / `maxValue`: Hard bounds for predicted ratings. Predictions outside this range + * are clipped. Set these to your rating scale limits (e.g., 1.0 and 5.0). + * + * Learning Rates (Step sizes for Gradient Descent): + * - `gamma1`: Learning rate for **Biases** (`b_u`, `b_i`). + * - `gamma2`: Learning rate for **Embeddings/Factors** (`p_u`, `q_i`, `y_j`). > Tip: Increase + * if convergence is too slow. Decrease if the loss explodes (NaN). + * + * Regularization (Preventing Overfitting): + * - `gamma6`: Regularization for **Biases**. + * - `gamma7`: Regularization for **Embeddings/Factors**. > Tip: Increase these if the model + * performs well on training data but poorly on test data. + */ +class SVDPlusPlus private[graphframes] (private val graph: GraphFrame) + extends Arguments + with WithMaxIter + with Logging { + private var _rank: Int = 10 + private var _minVal: Double = 0.0 + private var _maxVal: Double = 5.0 + private var _gamma1: Double = 0.007 + private var _gamma2: Double = 0.007 + private var _gamma6: Double = 0.005 + private var _gamma7: Double = 0.015 + + private var _loss: Option[Double] = None + + def rank(value: Int): this.type = { + _rank = value + this + } + + def minValue(value: Double): this.type = { + _minVal = value + this + } + + def maxValue(value: Double): this.type = { + _maxVal = value + this + } + + def gamma1(value: Double): this.type = { + _gamma1 = value + this + } + + def gamma2(value: Double): this.type = { + _gamma2 = value + this + } + + def gamma6(value: Double): this.type = { + _gamma6 = value + this + } + + def gamma7(value: Double): this.type = { + _gamma7 = value + this + } + + def run(): DataFrame = { + import SVDPlusPlus.COLUMN_WEIGHT + + if (!graph.edges.columns.contains(COLUMN_WEIGHT)) { + throw new InvalidGraphException(s"SVD++ requires a weight column $COLUMN_WEIGHT") + } + val conf = new graphxlib.SVDPlusPlus.Conf( + rank = _rank, + maxIters = maxIter.getOrElse(2), + minVal = _minVal, + maxVal = _maxVal, + gamma1 = _gamma1, + gamma2 = _gamma2, + gamma6 = _gamma6, + gamma7 = _gamma7) + + val g = if (graph.hasIntegralIdType) { + graph + } else { + val iVertices = graph.indexedVertices + val iEdges = graph.indexedEdges.select( + col(GraphFrame.LONG_SRC).alias(GraphFrame.SRC), + col(GraphFrame.LONG_DST).alias(GraphFrame.DST), + col(GraphFrame.ATTR).getField(COLUMN_WEIGHT).alias(COLUMN_WEIGHT)) + + GraphFrame(iVertices, iEdges) + } + + val (df, l) = SVDPlusPlus.run(g, conf) + val result = if (graph.hasIntegralIdType) { + df.persist() + } else { + val iV = graph.indexedVertices + df.withColumnRenamed(GraphFrame.ID, GraphFrame.LONG_ID) + .join(iV, GraphFrame.LONG_ID) + .drop(GraphFrame.LONG_ID) + .persist() + } + _loss = Some(l) + + // materialize + result.count() + + // unpersist + df.unpersist() + resultIsPersistent() + result + } + + def loss: Double = { + // We could use types instead to make sure that it is never accessed before being run. + _loss.getOrElse(throw new Exception("The algorithm has not been run yet")) + } +} + +object SVDPlusPlus { + + private def run(graph: GraphFrame, conf: graphxlib.SVDPlusPlus.Conf): (DataFrame, Double) = { + val edges = graph.edges.select(GraphFrame.SRC, GraphFrame.DST, COLUMN_WEIGHT).rdd.map { row => + val src = row.getAs[Number](0).longValue() + val dst = row.getAs[Number](1).longValue() + val w = row.getAs[Number](2).doubleValue() + Edge(src, dst, w) + } + val (gx, res) = graphxlib.SVDPlusPlus.run(edges, conf) + val gf = GraphXConversions.fromGraphX( + graph, + gx, + vertexNames = Seq(COLUMN1, COLUMN2, COLUMN3, COLUMN4)) + val vertices = gf.vertices.persist() + vertices.count() + gx.unpersist() + (vertices, res) + } + + /** + * Name for input edge DataFrame column containing edge weights. + * + * Note: This column name may change in the future! + */ + val COLUMN_WEIGHT = "weight" + + /** + * Name for output vertexDataFrame column containing first parameter of learned model, of type + * `Array[Double]`. + * + * Note: This column name may change in the future! + */ + val COLUMN1 = "column1" + + /** + * Name for output vertexDataFrame column containing second parameter of learned model, of type + * `Array[Double]`. + * + * Note: This column name may change in the future! + */ + val COLUMN2 = "column2" + + /** + * Name for output vertexDataFrame column containing third parameter of learned model, of type + * `Double`. + * + * Note: This column name may change in the future! + */ + val COLUMN3 = "column3" + + /** + * Name for output vertexDataFrame column containing fourth parameter of learned model, of type + * `Double`. + * + * Note: This column name may change in the future! + */ + val COLUMN4 = "column4" +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ShortestPaths.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ShortestPaths.scala new file mode 100644 index 0000000000000..40acfc814721b --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ShortestPaths.scala @@ -0,0 +1,265 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.graphx +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.collect_list +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.functions.map +import org.apache.spark.sql.functions.map_values +import org.apache.spark.sql.functions.map_zip_with +import org.apache.spark.sql.functions.reduce +import org.apache.spark.sql.functions.transform_keys +import org.apache.spark.sql.functions.transform_values +import org.apache.spark.sql.functions.when +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.types.MapType +import org.apache.spark.storage.StorageLevel +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrame.quote +import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithAlgorithmChoice +import org.apache.spark.graphframes.WithCheckpointInterval +import org.apache.spark.graphframes.WithDirection +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLocalCheckpoints + +import java.util +import scala.jdk.CollectionConverters._ + +/** + * Computes shortest paths from every vertex to the given set of landmark vertices. Note that this + * takes edge direction into account. + * + * The returned DataFrame contains all the original vertex information as well as one additional + * column: + * - distances (`MapType[vertex ID type, IntegerType]`): For each vertex v, a map containing the + * shortest-path distance to each reachable landmark vertex. + */ +class ShortestPaths private[graphframes] (private val graph: GraphFrame) + extends Arguments + with WithAlgorithmChoice + with WithCheckpointInterval + with WithLocalCheckpoints + with WithIntermediateStorageLevel + with WithDirection { + import org.apache.spark.graphframes.lib.ShortestPaths._ + + private var lmarks: Option[Seq[Any]] = None + + /** + * The list of landmark vertex ids. Shortest paths will be computed to each landmark. + */ + def landmarks(value: Seq[Any]): this.type = { + // TODO(tjh) do some initial checks here, without running queries. + lmarks = Some(value) + this + } + + /** + * The list of landmark vertex ids. Shortest paths will be computed to each landmark. + */ + def landmarks(value: util.ArrayList[Any]): this.type = { + landmarks(value.asScala.toSeq) + } + + def run(): DataFrame = { + val lmarksChecked = check(lmarks, "landmarks") + val res = algorithm match { + case ALGO_GRAPHX => runInGraphX(graph, lmarksChecked, isDirected) + case ALGO_GRAPHFRAMES => + runInGraphFrames( + graph, + lmarksChecked, + checkpointInterval, + useLocalCheckpoints = useLocalCheckpoints, + intermediateStorageLevel = intermediateStorageLevel, + isDirected = isDirected) + case _ => throw new GraphFramesUnreachableException() + } + resultIsPersistent() + res + } +} + +private object ShortestPaths extends Logging { + + private def runInGraphX( + graph: GraphFrame, + landmarks: Seq[Any], + isDirected: Boolean): DataFrame = { + val longIdToLandmark = landmarks.map(l => GraphXConversions.integralId(graph, l) -> l).toMap + val topology = if (isDirected) { + graph.cachedTopologyGraphX + } else { + val undirectedEdges = graph.cachedTopologyGraphX.edges.flatMap { edge => + Iterator(edge, graphx.Edge(edge.dstId, edge.srcId, ())) + } + graphx.Graph.fromEdges(undirectedEdges, ()) + } + val gx = graphx.lib.ShortestPaths + .run(topology, longIdToLandmark.keys.toSeq.sorted) + val g = GraphXConversions.fromGraphX(graph, gx, vertexNames = Seq(DISTANCE_ID)) + val distanceCol: Column = if (graph.hasIntegralIdType) { + g.vertices(DISTANCE_ID) + } else { + val longIdToLandmarkFlatten: Seq[Column] = longIdToLandmark.flatMap { + case (k: Long, v: Any) => Seq(lit(k), lit(v)) + }.toSeq + val longIdToLandmarkColumn = map(longIdToLandmarkFlatten: _*) + transform_keys(col(DISTANCE_ID), (longId: Column, _) => longIdToLandmarkColumn(longId)) + } + val cols = graph.vertices.columns.map(quote).map(col) :+ distanceCol.as(DISTANCE_ID) + val res = g.vertices.select(cols.toSeq: _*) + res.persist(StorageLevel.MEMORY_AND_DISK_SER) + res.count() + gx.unpersist() + res + } + + private def runInGraphFrames( + graph: GraphFrame, + landmarks: Seq[Any], + checkpointInterval: Int, + isDirected: Boolean, + useLocalCheckpoints: Boolean, + intermediateStorageLevel: StorageLevel): DataFrame = { + val vertexType = graph.vertices.schema(GraphFrame.ID).dataType + + // For landmark vertices the initial distance to itself is set to 0 + // Example: graph with vertices a, b, c, d; landmarks = (c, d) + // we should init the following: + // (a, Map()), (b, Map()), (c, Map(c -> 0)), (d, Map(d -> 0)) + // + // Inside the following function it is done by applying multiple case-when + // because we know exactly that only one landmark could be equal to the nodeId. + // For example, for vertex c it will be: + // when(id == "a", Map(a -> 0)) + // .when(id == "b", Map(b -> 0)) + // .when(id == "c", Map(c -> 0)) --> this one is the only true + // .when(id == "d", Map(d -> 0)) + def initDistancesMap(vertexId: Column): Column = { + val firstLmarkCol = lit(landmarks.head) + var initCol = when(vertexId === firstLmarkCol, map(firstLmarkCol, lit(0))) + for (lmark <- landmarks.tail) { + initCol = initCol.when(vertexId === lit(lmark), map(lit(lmark), lit(0))) + } + initCol + } + + // Concatenations of two distance maps: + // If one map is null just take another. + // In case both maps are not null: + // - iterate over keys + // - if value in the left map is null or greater than value from the right map take right one + // else take left one + def concatMaps(distancesLeft: Column, distancesRight: Column): Column = + when(distancesLeft.isNull, distancesRight) + .when(distancesRight.isNull, distancesLeft) + .otherwise(map_zip_with( + distancesLeft, + distancesRight, + (_, leftDistance, rightDistance) => { + when(leftDistance.isNull || (leftDistance > rightDistance), rightDistance) + .otherwise(leftDistance) + })) + + // If distance is null, result of d + 1 will be null too + def incrementDistances(distancesMap: Column): Column = + transform_values(distancesMap, (_, distance) => distance + lit(1)) + + // Takes an array of distance maps and reduce them with concatMaps + def aggregateArrayOfDistanceMaps(arrayCol: Column): Column = + reduce(arrayCol, lit(null).cast(MapType(vertexType, IntegerType)), concatMaps) + + // Checks that a sent distances map can change the destination distances. + // Evaluation would be "true" in case in the new distances map + // for one of keys present a non-null value but in the old distances map it is null + // or new distance is less than old one. + def isDistanceImprovedWithMessage(newMap: Column, oldMap: Column): Column = reduce( + map_values( + map_zip_with( + newMap, + oldMap, + (_, newDistance, rightDistance) => + (newDistance.isNotNull && rightDistance.isNull) || (newDistance < rightDistance))), + lit(false), + (left, right) => left || right) + + val srcDistanceCol = Pregel.src(DISTANCE_ID) + val dstDistanceCol = Pregel.dst(DISTANCE_ID) + + // Initial active-vertex col expression: only landmarks + val initialActiveVerticesExpr = col(GraphFrame.ID).isInCollection(landmarks) + + // Mark vertex as active only in the case idstance changed + val updateActiveVierticesExpr = isDistanceImprovedWithMessage(Pregel.msg, col(DISTANCE_ID)) + + val preparedGraph = GraphFrame( + graph.vertices.select(GraphFrame.ID), + graph.edges.select(GraphFrame.SRC, GraphFrame.DST)) + + // Overall: + // 1. Initialize distances + // 2. If new message can improve distances send it + // 3. Collect and aggregate messages + val pregel = preparedGraph.pregel + .setIntermediateStorageLevel(intermediateStorageLevel) + .setMaxIter(Int.MaxValue) // That is how the GraphX implementation works + .withVertexColumn( + DISTANCE_ID, + when(col(GraphFrame.ID).isInCollection(landmarks), initDistancesMap(col(GraphFrame.ID))) + .otherwise(map().cast(MapType(vertexType, IntegerType))), + concatMaps(col(DISTANCE_ID), Pregel.msg)) + .sendMsgToSrc(when( + isDistanceImprovedWithMessage(incrementDistances(dstDistanceCol), srcDistanceCol), + incrementDistances(dstDistanceCol))) + .aggMsgs(aggregateArrayOfDistanceMaps(collect_list(Pregel.msg))) + .setEarlyStopping(true) + .setInitialActiveVertexExpression(initialActiveVerticesExpr) + .setUpdateActiveVertexExpression(updateActiveVierticesExpr) + .setStopIfAllNonActiveVertices(true) + .setSkipMessagesFromNonActiveVertices(true) + .setCheckpointInterval(checkpointInterval) + .setUseLocalCheckpoints(useLocalCheckpoints) + // Memory optimization: only include required columns in triplets + .requiredSrcColumns(DISTANCE_ID) + .requiredDstColumns(DISTANCE_ID) + + // Experimental feature + if (isDirected) { + pregel.run() + } else { + // For consider edges as undirected, + // it is enough to send messages in both directions + pregel + .sendMsgToDst( + when( + isDistanceImprovedWithMessage(incrementDistances(srcDistanceCol), dstDistanceCol), + incrementDistances(srcDistanceCol))) + .run() + } + + } + + private val DISTANCE_ID = "distances" +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponents.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponents.scala new file mode 100644 index 0000000000000..aa11d4ec3c666 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponents.scala @@ -0,0 +1,58 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.graphx.{lib => graphxlib} +import org.apache.spark.sql.DataFrame +import org.apache.spark.storage.StorageLevel +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithMaxIter + +/** + * Compute the strongly connected component (SCC) of each vertex and return a DataFrame with each + * vertex assigned to the SCC containing that vertex. + * + * The resulting DataFrame contains all the original vertex information and one additional column: + * - component (`LongType`): unique ID for this component + */ +class StronglyConnectedComponents private[graphframes] (private val graph: GraphFrame) + extends Arguments + with WithMaxIter + with Logging { + + def run(): DataFrame = { + val res = StronglyConnectedComponents.run(graph, check(maxIter, "maxIter")) + resultIsPersistent() + res + } +} + +/** Strongly connected components algorithm implementation. */ +private object StronglyConnectedComponents { + private def run(graph: GraphFrame, numIter: Int): DataFrame = { + val gx = graphxlib.StronglyConnectedComponents.run(graph.cachedTopologyGraphX, numIter) + val res = GraphXConversions.fromGraphX(graph, gx, vertexNames = Seq(COMPONENT_ID)).vertices + res.persist(StorageLevel.MEMORY_AND_DISK_SER) + res.count() + gx.unpersist() + res + } + + private[graphframes] val COMPONENT_ID = "component" +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala new file mode 100644 index 0000000000000..724b98191ec7d --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala @@ -0,0 +1,291 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types._ +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrame._ +import org.apache.spark.graphframes.GraphFramesSparkVersionException +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithCheckpointInterval +import org.apache.spark.graphframes.WithDirection +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLgNomEntries +import org.apache.spark.graphframes.WithLocalCheckpoints +import org.apache.spark.graphframes.WithMaxIter + +/** + * Neighborhood-aware community detection via weighted label propagation. + * + * This algorithm is a Label Propagation variant where each incoming label vote is weighted by a + * combination of: + * - optional direct-link baseline strength (enabled unless `ignoreDirectLinks = true`), and + * - neighborhood-overlap strength (`structuralSimilarityMultiplier * commonNeighbors`). + * + * Intuitively, labels from neighbors that are structurally similar to the destination (many + * common neighbors) can be amplified, instead of treating all edges equally. + * + * At each iteration, every vertex aggregates weighted incoming votes by label and picks the label + * with maximum total weight. + * + * Main hyperparameters: + * - `maxIter` (required): maximum number of propagation rounds. + * - `ignoreDirectLinks` (default `false`): whether to drop direct-link baseline vote mass. + * - `structuralSimilarityMultiplier` (default `0.5`): scales neighborhood-overlap contribution. + * + * Edge-weight regimes: + * - `ignoreDirectLinks = false`: + * {{{ + * edgeWeight(src, dst) = 1 + structuralSimilarityMultiplier * commonNeighbors(src, dst) + * }}} + * - `ignoreDirectLinks = true`: + * {{{ + * edgeWeight(src, dst) = structuralSimilarityMultiplier * commonNeighbors(src, dst) + * }}} + * + * This implementation is inspired by neighborhood-strength-driven label propagation ideas from: + * Xie, Jierui, and Boleslaw K. Szymanski. "Community detection using a neighborhood strength + * driven label propagation algorithm." 2011 IEEE Network Science Workshop. IEEE, 2011. + * + * Note: this implementation does not strictly reproduce the paper; it adopts the core idea of + * modulating label votes with a common-neighbor term within the GraphFrames/Pregel design. + */ +class StructureAwareLabelPropagation private[graphframes] (private val graph: GraphFrame) + extends Arguments + with WithCheckpointInterval + with WithMaxIter + with WithLocalCheckpoints + with WithIntermediateStorageLevel + with WithDirection + with WithLgNomEntries + with Logging { + + private var structuralSimilarityMultiplier: Double = 0.5 + private var ignoreDirectLinks: Boolean = false + private var initialLabelCol: Option[String] = None + + import StructureAwareLabelPropagation._ + + /** + * Sets whether direct-link baseline vote mass is ignored. + * + * If `false` (default), each existing edge contributes a baseline of `1.0` before structural + * overlap is added. If `true`, only structural overlap contributes vote mass. + */ + def setIgnoreDirectLinks(value: Boolean): this.type = { + ignoreDirectLinks = value + this + } + + /** + * Sets multiplier for the neighborhood-overlap signal (common neighbors). + * + * Edge weighting is: + * - when direct links are included: + * {{{ + * edgeWeight(src, dst) = 1 + structuralSimilarityMultiplier * commonNeighbors(src, dst) + * }}} + * - when direct links are ignored: + * {{{ + * edgeWeight(src, dst) = structuralSimilarityMultiplier * commonNeighbors(src, dst) + * }}} + * + * `commonNeighbors(src, dst)` is the (approximate) number of shared out-neighbors between + * source and destination. + * + * The value must be non-negative. + */ + def setStructuralSimilarityMultiplier(value: Double): this.type = { + require(value >= 0.0, "structuralSimilarityMultiplier must be >= 0") + structuralSimilarityMultiplier = value + this + } + + /** + * Sets an explicit vertex column to use as initial labels. + * + * By default, each vertex starts with its own `id` as label. When this setter is used, the + * algorithm initializes labels from the provided attribute column instead, enabling + * attribute-guided label propagation (attribute propagation): labels can start from domain + * values such as categories, types, or seeds and then propagate through the graph structure. + * + * The output `label` column keeps the data type of the provided column. + */ + def setInitialLabelCol(col: String): this.type = { + require( + graph.vertices.columns.contains(col), + s"Initial label column '$col' does not exist in vertex columns: ${graph.vertices.columns.mkString(", ")}") + initialLabelCol = Some(col) + this + } + + def run(): DataFrame = { + // Validate parameters + val maxIterChecked = check(maxIter, "maxIter") + require( + !(ignoreDirectLinks && structuralSimilarityMultiplier == 0.0), + "structuralSimilarityMultiplier must be > 0 when ignoreDirectLinks is true") + + // Sketch-based features require Spark >= 4.1 + val sparkVersion = graph.vertices.sparkSession.version + if (sparkVersion.substring(0, 3) < "4.1") { + throw new GraphFramesSparkVersionException("4.1.0") + } + + val edges = if (isDirected) { + graph.edges.select(SRC, DST) + } else { + graph.edges + .select(col(DST).alias(SRC), col(SRC).alias(DST)) + .union(graph.edges.select(SRC, DST)) + .distinct() + } + + val directLinkScale = if (ignoreDirectLinks) 0.0 else 1.0 + + // Compute approximate common neighbor counts on edges and materialize. + val enrichedEdges = + computeEdgeApproxCommonNeighbors( + edges, + lgNomEntries, + structuralSimilarityMultiplier, + directLinkScale, + isDirected) + + val vertices = if (initialLabelCol.isDefined) { + graph.vertices.select(col(GraphFrame.ID), col(initialLabelCol.get).alias(INITIAL_LABEL_COL)) + } else { + graph.vertices.select(col(GraphFrame.ID), col(GraphFrame.ID).alias(INITIAL_LABEL_COL)) + } + + val preparedGraph = GraphFrame(vertices, enrichedEdges) + + val pregel = preparedGraph.pregel + + pregel + .setMaxIter(maxIterChecked) + .setCheckpointInterval(checkpointInterval) + .setUseLocalCheckpoints(useLocalCheckpoints) + .setIntermediateStorageLevel(intermediateStorageLevel) + .requiredEdgeColumns(EDGE_WEIGHT_COL) + .sendMsgToDst(struct(Pregel.src(LABEL_COL), Pregel.edge(EDGE_WEIGHT_COL))) + .aggMsgs(aggregateMessages(Pregel.msg, vertices.schema(INITIAL_LABEL_COL).dataType)) + .withVertexColumn( + LABEL_COL, + col(INITIAL_LABEL_COL), + coalesce(keyWithMaxValue(Pregel.msg), col(LABEL_COL))) + .setSkipMessagesFromNonActiveVertices(false) + .setUpdateActiveVertexExpression( + col(LABEL_COL) =!= coalesce(keyWithMaxValue(Pregel.msg), col(LABEL_COL))) + .setStopIfAllNonActiveVertices(true) + .setEarlyStopping(false) + + val result = pregel.run().drop(INITIAL_LABEL_COL) + resultIsPersistent() + + result + } +} + +object StructureAwareLabelPropagation extends Logging { + + val LABEL_COL = "label" + private val INITIAL_LABEL_COL = "initial_label" + private val EDGE_WEIGHT_COL = "edge_weight" + + private def aggregateMessages(msgCol: Column, idType: DataType): Column = reduce( + collect_list(msgCol), + map().cast(MapType(idType, DoubleType)), + (acc, x) => + map_zip_with( + acc, + map(x.getField(LABEL_COL), x.getField(EDGE_WEIGHT_COL)), + (_, left, right) => coalesce(left, lit(0.0)) + coalesce(right, lit(0.0)))) + + private def keyWithMaxValue(column: Column): Column = array_max( + transform( + map_entries(column), + x => struct(x.getField("value"), x.getField("key").alias("key")))) + .getField("key") + + private def computeEdgeApproxCommonNeighbors( + edges: DataFrame, + lgNomEntries: Int, + structuralSimilarityMultiplier: Double, + directLinkScale: Double, + isDirected: Boolean): DataFrame = { + + // For directed graphs we use weak neighborhoods for structural similarity: + // N(v) = In(v) union Out(v). + // + // Label propagation still follows edge direction via sendMsgToDst, but the + // common-neighbor term ignores orientation to avoid making the structural + // signal too sparse on low-degree directed graphs. + def thetaSketchAggExpr = (c: String) => expr(s"theta_sketch_agg($c, $lgNomEntries)") + + var vertexSketches = edges + .groupBy(col(SRC).alias(ID)) + .agg(thetaSketchAggExpr(DST).alias("nbr_theta_sketch")) + + if (isDirected) { + val thetaSketchUnion = (left: String, right: String) => expr(s"theta_union($left, $right)") + + vertexSketches = vertexSketches + .join( + edges + .groupBy(col(DST).alias(ID)) + .agg(thetaSketchAggExpr(SRC).alias("nbr_theta_sketch_dst")), + Seq(ID), + "full") + .withColumn( + "nbr_theta_sketch", + when(col("nbr_theta_sketch").isNull, col("nbr_theta_sketch_dst")) + .when(col("nbr_theta_sketch_dst").isNull, col("nbr_theta_sketch")) + .otherwise(thetaSketchUnion("nbr_theta_sketch", "nbr_theta_sketch_dst"))) + } + + val thetaSketchIntersect = (left: String, right: String) => + expr(s"theta_sketch_estimate(theta_intersection($left, $right))") + + // Prepare sketches for join (id, nbr_theta_sketch) + val srcSketch = vertexSketches.select( + col(ID).alias("sk_src_id"), + col("nbr_theta_sketch").alias("src_nbr_sketch")) + val dstSketch = vertexSketches.select( + col(ID).alias("sk_dst_id"), + col("nbr_theta_sketch").alias("dst_nbr_sketch")) + + // Join edges with sketches. Use left_outer so edges without sketches will have null sketches. + val e = edges + .select(SRC, DST) + .join(srcSketch, col(SRC) === col("sk_src_id"), "left_outer") + .join(dstSketch, col(DST) === col("sk_dst_id"), "left_outer") + + // Compute approximate intersection and coalesce nulls to 0.0 + val weightCol = lit(directLinkScale) + lit(structuralSimilarityMultiplier) * coalesce( + thetaSketchIntersect("src_nbr_sketch", "dst_nbr_sketch"), + lit(0.0)) + val edgesWithOverlap = e + .select(col(SRC), col(DST), weightCol.alias(EDGE_WEIGHT_COL)) + + edgesWithOverlap + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/TriangleCount.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/TriangleCount.scala new file mode 100644 index 0000000000000..069c994c4dc83 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/TriangleCount.scala @@ -0,0 +1,194 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ +import org.apache.spark.storage.StorageLevel +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFramesSparkVersionException +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLgNomEntries + +/** + * Triangle count implementation. + * + * This class provides two algorithms for counting triangles: + * - A direct version that computes exact triangle counts using set intersection of neighbor + * lists. + * - An approximate version based on the DataSketches library (Theta sketches), which trades off + * accuracy for performance on large-scale graphs. + * + * The output DataFrame contains two columns: + * - "id": the vertex id + * - "count": the number of triangles passing through the vertex + */ +class TriangleCount private[graphframes] (private val graph: GraphFrame) + extends Arguments + with Serializable + with WithIntermediateStorageLevel + with WithLgNomEntries { + + private var algorithm: String = "exact" + private val supportedAlgorithms: Set[String] = Set("exact", "approx") + + private def supportedAlgorithmsRepr: String = supportedAlgorithms.mkString(", ") + + /** + * Sets the triangle counting algorithm. Options are "exact" (default) or "approx". + */ + def setAlgorithm(value: String): this.type = { + require( + supportedAlgorithms.contains(value), + s"supported algorithms: ${supportedAlgorithmsRepr}") + algorithm = value + this + } + + def run(): DataFrame = { + if (algorithm == "exact") { + TriangleCount.run(graph, intermediateStorageLevel) + } else { + TriangleCount.approximateRun(graph, intermediateStorageLevel, lgNomEntries) + } + } +} + +private object TriangleCount extends Logging { + import org.apache.spark.graphframes.GraphFrame._ + + private def prepareGraph(graph: GraphFrame): GraphFrame = { + // Dedup edges by flipping them to have SRC < DST + // Remove self-loops + val dedupedE = graph.edges + .filter(col(SRC) =!= col(DST)) + .select( + when(col(SRC) < col(DST), col(SRC)).otherwise(col(DST)).as(SRC), + when(col(SRC) < col(DST), col(DST)).otherwise(col(SRC)).as(DST)) + .distinct() + + // Prepare the graph with no isolated vertices. + GraphFrame(graph.vertices.select(ID), dedupedE).dropIsolatedVertices() + } + + private def approximateRun( + graph: GraphFrame, + intermediateStorageLevel: StorageLevel, + lgNomEntries: Int): DataFrame = { + val spark = graph.vertices.sparkSession + val sparkVersion = spark.version + + if (sparkVersion.substring(0, 3) < "4.1") { + throw new GraphFramesSparkVersionException("4.1.0") + } + + val thetaSketchAgg = (colName: String) => expr(s"theta_sketch_agg($colName, $lgNomEntries)") + val thetaSketchIntersect = (colLeft: String, colRight: String) => + expr(s"theta_sketch_estimate(theta_intersection($colLeft, $colRight))") + + val g2 = prepareGraph(graph) + + val verticesWithNeighbors = g2.aggregateMessages + .setIntermediateStorageLevel(intermediateStorageLevel) + .sendToSrc(AggregateMessages.dst(ID)) + .sendToDst(AggregateMessages.src(ID)) + .agg(thetaSketchAgg(AggregateMessages.MSG_COL_NAME).alias("neighbors")) + .persist(intermediateStorageLevel) + + val triangles = verticesWithNeighbors + .select(col(ID), col("neighbors").alias("src_set")) + .join(g2.edges, col(ID) === col(SRC)) + .drop(ID) + .join( + verticesWithNeighbors.select(col(ID), col("neighbors").alias("dst_set")), + col(ID) === col(DST)) + .drop(ID) + // Count of common neighbors of SRC and DST + .withColumn("triplets", thetaSketchIntersect("src_set", "dst_set")) + .filter(col("triplets") > lit(0)) + .persist(intermediateStorageLevel) + + val srcTriangles = triangles.groupBy(SRC).agg(sum(col("triplets")).alias("src_triplets")) + val dstTriangles = triangles.groupBy(DST).agg(sum(col("triplets")).alias("dst_triplets")) + + val result = graph.vertices + .join(srcTriangles, col(ID) === col(SRC), "left_outer") + .join(dstTriangles, col(ID) === col(DST), "left_outer") + // Each triangle counted twice, so divide by 2. + .withColumn( + COUNT_ID, + floor( + (coalesce(col("src_triplets"), lit(0)) + coalesce(col("dst_triplets"), lit(0))) / lit( + 2))) + .select(col(ID), col(COUNT_ID)) + + result.persist(intermediateStorageLevel) + result.count() + verticesWithNeighbors.unpersist() + triangles.unpersist() + resultIsPersistent() + result + } + + private def run(graph: GraphFrame, intermediateStorageLevel: StorageLevel): DataFrame = { + val g2 = prepareGraph(graph) + + val verticesWithNeighbors = g2.aggregateMessages + .setIntermediateStorageLevel(intermediateStorageLevel) + .sendToSrc(AggregateMessages.dst(ID)) + .sendToDst(AggregateMessages.src(ID)) + .agg(collect_set(AggregateMessages.msg).alias("neighbors")) + .persist(intermediateStorageLevel) + + val triangles = verticesWithNeighbors + .select(col(ID), col("neighbors").alias("src_set")) + .join(g2.edges, col(ID) === col(SRC)) + .drop(ID) + .join( + verticesWithNeighbors.select(col(ID), col("neighbors").alias("dst_set")), + col(ID) === col(DST)) + .drop(ID) + // Count of common neighbors of SRC and DST + .withColumn("triplets", array_size(array_intersect(col("src_set"), col("dst_set")))) + .filter(col("triplets") > lit(0)) + .persist(intermediateStorageLevel) + + val srcTriangles = triangles.groupBy(SRC).agg(sum(col("triplets")).alias("src_triplets")) + val dstTriangles = triangles.groupBy(DST).agg(sum(col("triplets")).alias("dst_triplets")) + + val result = graph.vertices + .join(srcTriangles, col(ID) === col(SRC), "left_outer") + .join(dstTriangles, col(ID) === col(DST), "left_outer") + // Each triangle counted twice, so divide by 2. + .withColumn( + COUNT_ID, + floor( + (coalesce(col("src_triplets"), lit(0)) + coalesce(col("dst_triplets"), lit(0))) / lit( + 2))) + + result.persist(intermediateStorageLevel) + result.count() + verticesWithNeighbors.unpersist() + triangles.unpersist() + resultIsPersistent() + result + } + + private val COUNT_ID = "count" +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/TwoPhase.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/TwoPhase.scala new file mode 100644 index 0000000000000..4cf3e70fb11eb --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/TwoPhase.scala @@ -0,0 +1,622 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.hadoop.fs.Path +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types.DecimalType +import org.apache.spark.storage.StorageLevel +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrame.ATTR +import org.apache.spark.graphframes.GraphFrame.DST +import org.apache.spark.graphframes.GraphFrame.ID +import org.apache.spark.graphframes.GraphFrame.LONG_DST +import org.apache.spark.graphframes.GraphFrame.LONG_ID +import org.apache.spark.graphframes.GraphFrame.LONG_SRC +import org.apache.spark.graphframes.GraphFrame.SRC +import org.apache.spark.graphframes.Logging + +import java.io.IOException +import java.math.BigDecimal +import java.util.UUID + +/** + * Two-phase label propagation implementation of connected components. + * + * This is the primary GraphFrames-native implementation. It iteratively applies large-star and + * small-star steps until convergence, using checkpointing to manage query plan growth. + */ +private[graphframes] object TwoPhase extends Logging { + + private val CHECKPOINT_NAME_PREFIX = "connected-components" + private val MIN_NBR = "min_nbr" + private val CNT = "cnt" + + /** + * Returns the symmetric directed graph of the graph specified by input edges. + * @param ee + * non-bidirectional edges + */ + private def symmetrize(ee: DataFrame): DataFrame = { + val EDGE = "_edge" + ee.select(explode( + array(struct(col(SRC), col(DST)), struct(col(DST).as(SRC), col(SRC).as(DST)))).as(EDGE)) + .select(col(s"$EDGE.$SRC").as(SRC), col(s"$EDGE.$DST").as(DST)) + } + + /** + * Prepares the input graph for computing connected components by: + * - de-duplicating vertices and assigning unique long IDs to each, + * - changing edge directions to have increasing long IDs from src to dst, + * - de-duplicating edges and removing self-loops. + * + * In the returned GraphFrame, the vertex DataFrame has two columns: + * - column `id` stores a long ID assigned to the vertex, + * - column `attr` stores the original vertex attributes. + * + * The edge DataFrame has two columns: + * - column `src` stores the long ID of the source vertex, + * - column `dst` stores the long ID of the destination vertex, where we always have `src` < + * `dst`. + */ + private def prepare(graph: GraphFrame): GraphFrame = { + val vertices = graph.indexedVertices + .select(col(LONG_ID).as(ID), col(ATTR)) + val edges = graph.indexedEdges + .select(col(LONG_SRC).as(SRC), col(LONG_DST).as(DST)) + val orderedEdges = edges + .filter(col(SRC) =!= col(DST)) + .select(minValue(col(SRC), col(DST)).as(SRC), maxValue(col(SRC), col(DST)).as(DST)) + .distinct() + GraphFrame(vertices, orderedEdges) + } + + /** + * Returns the min vertex among each vertex and its neighbors in a DataFrame with three columns: + * - `src`, the ID of the vertex + * - `min_nbr`, the min vertex ID among itself and its neighbors + * - `cnt`, the total number of neighbors + */ + private def minNbrs(ee: DataFrame): DataFrame = { + symmetrize(ee) + .groupBy(SRC) + .agg(min(col(DST)).as(MIN_NBR), count("*").as(CNT)) + .withColumn(MIN_NBR, minValue(col(SRC), col(MIN_NBR))) + } + + private def minValue(x: Column, y: Column): Column = { + when(x < y, x).otherwise(y) + } + + private def maxValue(x: Column, y: Column): Column = { + when(x > y, x).otherwise(y) + } + + /** + * Computes the sum of all `min_nbr` values and the undirected edge count of the given DataFrame + * in a single Spark job. The sum is cast to DecimalType(38, 0) for high precision. Used to + * detect convergence between iterations and to check graph sparsity for the pruning + * optimization. The edge count is derived as `sum(cnt) / 2` since each undirected edge appears + * once in each direction in the symmetrized graph. + */ + private def calcMinNbrSum(minNbrsDF: DataFrame): (BigDecimal, Long) = { + val row = minNbrsDF + .select( + sum(col(MIN_NBR).cast(DecimalType(38, 0))), + coalesce((sum(col(CNT)) / 2).cast("long"), lit(0L))) + .first() + (row.getAs[BigDecimal](0), row.getLong(1)) + } + + /** + * Builds the output DataFrame by joining the indexed vertices with the final edge assignments + * and resolving component labels. + */ + private def buildOutput( + graph: GraphFrame, + vv: DataFrame, + ee: DataFrame, + useLabelsAsComponents: Boolean): DataFrame = { + val indexedLabel = vv + .join(ee, vv(ID) === ee(DST), "left_outer") + .select( + vv(ATTR), + when(ee(SRC).isNull, vv(ID)).otherwise(ee(SRC)).as(ConnectedComponents.COMPONENT), + col(ATTR + "." + ID).as(ID)) + + if (graph.hasIntegralIdType || !useLabelsAsComponents) { + indexedLabel + .select(col(s"$ATTR.*"), col(ConnectedComponents.COMPONENT)) + } else { + indexedLabel + .join( + indexedLabel + .groupBy(col(ConnectedComponents.COMPONENT)) + .agg(min(col(ID)).as(ConnectedComponents.ORIG_ID)) + .select(col(ConnectedComponents.COMPONENT), col(ConnectedComponents.ORIG_ID)), + ConnectedComponents.COMPONENT) + .select( + col(s"$ATTR.*"), + col(ConnectedComponents.ORIG_ID).as(ConnectedComponents.COMPONENT)) + } + } + + /** + * Performs a possibly skewed join between edges and current component assignments. The skew + * join is done by broadcast join for frequent keys and normal join for the rest. + */ + private def skewedJoin( + edges: DataFrame, + minNbrsDF: DataFrame, + broadcastThreshold: Int, + logPrefix: String): DataFrame = { + import edges.sparkSession.implicits._ + val hubs = minNbrsDF + .filter(col(CNT) > broadcastThreshold) + .select(SRC) + .as[Long] + .collect() + .toSet + GraphFrame.skewedJoin(edges, minNbrsDF, SRC, hubs, logPrefix) + } + + /** + * Prunes leaf nodes (vertices with out-degree 0 and in-degree 1) from the graph to create a + * smaller shrunken graph. Returns Some((vertices, edges, nodeCount)) if the shrunken graph is + * significantly smaller than the original (by shrinkageThreshold), None otherwise. + */ + private[graphframes] def pruneLeafNodes( + edges: DataFrame, + intermediateStorageLevel: StorageLevel, + numNodes: Long, + shrinkageThreshold: Double): Option[(DataFrame, DataFrame, Long)] = { + + // vertices whose indegree > 1 + val v1 = edges + .groupBy(DST) + .agg(count("*").as(CNT)) + .filter(col(CNT) > 1) + .select(col(DST).as(ID)) + + // vertices whose outdegree > 0 or indegree > 1 + val newVV = edges + .select(col(SRC).as(ID)) + .union(v1) + .distinct() + .persist(intermediateStorageLevel) + val newVVCnt = newVV.count() + + if (newVVCnt * shrinkageThreshold < numNodes) { + val newEE = edges + .join(newVV.withColumnRenamed(ID, DST), DST) + .persist(intermediateStorageLevel) + Some((newVV, newEE, newVVCnt)) + } else { + newVV.unpersist(blocking = false) + None + } + } + + /** + * Given the vertices and converged edges of the shrunken graph, joins back to reconstruct the + * converged edges of the original graph. + */ + private[graphframes] def joinBack( + vertices: DataFrame, + edges: DataFrame, + edgesBeforePruning: DataFrame): DataFrame = { + + val cc = vertices + .as("vertices") + .join(edges.as("edges"), col(s"vertices.$ID") === col(s"edges.$DST"), "left_outer") + .select( + when(col(s"edges.$SRC").isNull, col(s"vertices.$ID")) + .otherwise(col(s"edges.$SRC")) + .as(SRC), + col(s"vertices.$ID").as(DST)) + + cc.as("cc") + .join( + edgesBeforePruning.as("edgesBeforePruning"), + col(s"cc.$DST") === col(s"edgesBeforePruning.$SRC")) + .select(col(s"cc.$SRC"), col(s"edgesBeforePruning.$DST")) + .union(cc) + .distinct() + } + + /** + * Runs the two-phase label propagation connected components algorithm. + */ + private[graphframes] def run( + graph: GraphFrame, + broadcastThreshold: Int, + checkpointInterval: Int, + intermediateStorageLevel: StorageLevel, + useLabelsAsComponents: Boolean, + useLocalCheckpoints: Boolean, + isGraphPrepared: Boolean, + optStartIter: Int = 2, + sparsityThreshold: Double = 2.0, + shrinkageThreshold: Double = 2.0): DataFrame = { + + val spark = graph.spark + val sc = spark.sparkContext + val originalAQE = spark.conf.get("spark.sql.adaptive.enabled") + + try { + spark.conf.set("spark.sql.adaptive.enabled", "false") + + val runId = UUID.randomUUID().toString.takeRight(8) + val logPrefix = s"[CC $runId]" + logInfo(s"$logPrefix Start connected components with run ID $runId.") + + val shouldCheckpoint = checkpointInterval > 0 + val checkpointDir: Option[String] = if (useLocalCheckpoints) { None } + else if (shouldCheckpoint) { + val dir = sc.getCheckpointDir + .map { d => + new Path(d, s"$CHECKPOINT_NAME_PREFIX-$runId").toString + } + .getOrElse { + spark.conf.getOption("spark.checkpoint.dir") match { + case Some(d) => new Path(d, s"$CHECKPOINT_NAME_PREFIX-$runId").toString + case None => + throw new IOException( + "Checkpoint directory is not set. Please set it first using sc.setCheckpointDir()" + + "or by specifying the conf 'spark.checkpoint.dir'.") + } + } + logInfo(s"$logPrefix Using $dir for checkpointing with interval $checkpointInterval.") + Some(dir) + } else { + logInfo( + s"$logPrefix Checkpointing is disabled because checkpointInterval=$checkpointInterval.") + None + } + + logInfo(s"$logPrefix Preparing the graph for connected component computation ...") + val g = if (isGraphPrepared) graph else prepare(graph) + val vv = g.vertices + var ee = g.edges.persist(intermediateStorageLevel) // src < dst + logInfo(s"$logPrefix Found ${ee.count()} edges after preparation.") + var numNodes = vv.count() + logInfo(s"$logPrefix Found $numNodes nodes after preparation.") + + var converged = false + var iteration = 1 + var isOptimized = false + var triedToOptimize = false + var shouldKeepCheckpoint = false + var edgesBeforePruning: DataFrame = null + var shrunkenGraphNodes: DataFrame = null + + var minNbrs1: DataFrame = minNbrs(ee) // src >= min_nbr + .persist(intermediateStorageLevel) + + var (prevSum, _) = calcMinNbrSum(minNbrs1) + + var lastRoundPersistedDFs = Seq[DataFrame](ee, minNbrs1) + while (!converged) { + var currRoundPersistedDFs = Seq[DataFrame]() + + // large-star step + // connect all strictly larger neighbors to the min neighbor (including self) + ee = skewedJoin(ee, minNbrs1, broadcastThreshold, logPrefix) + .select(col(DST).as(SRC), col(MIN_NBR).as(DST)) // src > dst + .distinct() + .persist(intermediateStorageLevel) + currRoundPersistedDFs = currRoundPersistedDFs :+ ee + + // small-star step + // compute min neighbors (excluding self-min) + val minNbrs2 = ee + .groupBy(col(SRC)) + .agg(min(col(DST)).as(MIN_NBR), count("*").as(CNT)) // src > min_nbr + .persist(intermediateStorageLevel) + currRoundPersistedDFs = currRoundPersistedDFs :+ minNbrs2 + + // connect all smaller neighbors to the min neighbor + ee = skewedJoin(ee, minNbrs2, broadcastThreshold, logPrefix) + .select(col(MIN_NBR).as(SRC), col(DST)) // src <= dst + .filter(col(SRC) =!= col(DST)) // src < dst + // connect self to the min neighbor + ee = ee + .union(minNbrs2.select(col(MIN_NBR).as(SRC), col(SRC).as(DST))) // src < dst + .distinct() + + // checkpointing + if (shouldCheckpoint && (iteration % checkpointInterval == 0)) { + if (useLocalCheckpoints) { + ee = ee.localCheckpoint(eager = true) + } else { + val out = s"${checkpointDir.get}/$iteration" + ee.write.parquet(out) + ee = spark.read.parquet(out) + + if (iteration > checkpointInterval) { + val path = new Path(s"${checkpointDir.get}/${iteration - checkpointInterval}") + // keep the checkpoint when edgesBeforePruning points to it + if (!shouldKeepCheckpoint) { + path.getFileSystem(sc.hadoopConfiguration).delete(path, true) + } else { + shouldKeepCheckpoint = false + } + } + + System.gc() + } + } + + ee.persist(intermediateStorageLevel) + currRoundPersistedDFs = currRoundPersistedDFs :+ ee + + minNbrs1 = minNbrs(ee) // src >= min_nbr + .persist(intermediateStorageLevel) + currRoundPersistedDFs = currRoundPersistedDFs :+ minNbrs1 + + // test convergence + val (currSum, edgeCnt) = calcMinNbrSum(minNbrs1) + logInfo(s"$logPrefix Sum of assigned components in iteration $iteration: $currSum.") + + // Pruning Node Optimization: construct a new small graph with fewer nodes, + // and find connected components of the shrunken graph, then join back to get the + // connected components of the original graph. + + // If the graph becomes sparse and current iteration >= $optStartIter, we start to + // try such optimization. However, the optimization is only performed if the shrunken + // graph is much smaller than the original graph, otherwise we do not perform it ( + // in this case, the only additional cost is to determine the size of shrunken graph). + // In current implementation, we only try such optimization one time and it is + // performed at most one time. So the additional cost is bounded. + + // According to such heuristic rule, we can determine when and whether we should + // perform the optimization. For the sparse graphs (defined by sparsityThreshold), + // we will try such optimization at the end of $optStartIter iteration (default is 2). + // For the dense graph, its edges will be pruned at each large/small star join iteration, + // and we will try the optimization once the graph becomes sparse. + if ((edgeCnt < sparsityThreshold * numNodes) && (edgeCnt > 0) + && (iteration >= optStartIter) && (!triedToOptimize)) { + edgesBeforePruning = ee + pruneLeafNodes(ee, intermediateStorageLevel, numNodes, shrinkageThreshold) match { + case Some(r) => + shrunkenGraphNodes = r._1 + ee = r._2 + currRoundPersistedDFs = currRoundPersistedDFs :+ ee + numNodes = r._3 + isOptimized = true + shouldKeepCheckpoint = true + logInfo(s"$logPrefix Pruning node optimization performed in iteration $iteration.") + logInfo(s"$logPrefix Shrunken graph node count: $numNodes.") + case None => + logInfo(s"$logPrefix Pruning node optimization not performed.") + } + triedToOptimize = true + } + + if (currSum == prevSum) { + converged = true + } else { + prevSum = currSum + } + + for (persistedDF <- lastRoundPersistedDFs) { + persistedDF.unpersist() + } + lastRoundPersistedDFs = currRoundPersistedDFs + iteration += 1 + } + + if (isOptimized) { + ee = joinBack(shrunkenGraphNodes, ee, edgesBeforePruning) + } + + logInfo(s"$logPrefix Connected components converged in ${iteration - 1} iterations.") + logInfo(s"$logPrefix Join and return component assignments with original vertex IDs.") + + val output = buildOutput(graph, vv, ee, useLabelsAsComponents) + .persist(intermediateStorageLevel) + + output.count() + + for (persistedDF <- lastRoundPersistedDFs) { + persistedDF.unpersist() + } + if (shrunkenGraphNodes != null) { + shrunkenGraphNodes.unpersist() + } + + resultIsPersistent() + + output + } finally { + spark.conf.set("spark.sql.adaptive.enabled", originalAQE) + } + } + + /** + * Runs the two-phase label propagation connected components algorithm using Adaptive Query + * Execution (AQE). Unlike `run`, this method does not manipulate AQE settings, does not use + * skewed joins, and uses simpler checkpointing. + */ + private[graphframes] def runAQE( + graph: GraphFrame, + checkpointInterval: Int, + intermediateStorageLevel: StorageLevel, + useLabelsAsComponents: Boolean, + useLocalCheckpoints: Boolean, + isGraphPrepared: Boolean, + optStartIter: Int = 2, + sparsityThreshold: Double = 2.0, + shrinkageThreshold: Double = 2.0): DataFrame = { + + val runId = UUID.randomUUID().toString.takeRight(8) + val logPrefix = s"[CC $runId]" + logInfo(s"$logPrefix Start connected components with run ID $runId.") + + val shouldCheckpoint = checkpointInterval > 0 + + logInfo(s"$logPrefix Preparing the graph for connected component computation ...") + val g = if (isGraphPrepared) graph else prepare(graph) + val vv = g.vertices + var ee = g.edges.persist(intermediateStorageLevel) // src < dst + logInfo(s"$logPrefix Found ${ee.count()} edges after preparation.") + var numNodes = vv.count() + logInfo(s"$logPrefix Found $numNodes nodes after preparation.") + + var converged = false + var iteration = 1 + var isOptimized = false + var triedToOptimize = false + var edgesBeforePruning: DataFrame = null + var shrunkenGraphNodes: DataFrame = null + + var minNbrs1: DataFrame = symmetrize(ee) + .groupBy(SRC) + .agg(min(col(DST)).as(MIN_NBR), count("*").as(CNT)) + .withColumn(MIN_NBR, minValue(col(SRC), col(MIN_NBR))) + .persist(intermediateStorageLevel) + + var (prevSum, _) = calcMinNbrSum(minNbrs1) + + var lastRoundPersistedDFs = Seq[DataFrame](ee, minNbrs1) + while (!converged) { + var currRoundPersistedDFs = Seq[DataFrame]() + + // large-star step + // connect all strictly larger neighbors to the min neighbor (including self) + ee = ee + .join(minNbrs1, SRC) + .select(col(DST).as(SRC), col(MIN_NBR).as(DST)) // src > dst + .distinct() + .persist(intermediateStorageLevel) + currRoundPersistedDFs = currRoundPersistedDFs :+ ee + + // small-star step + // compute min neighbors (excluding self-min) + val minNbrs2 = ee + .groupBy(col(SRC)) + .agg(min(col(DST)).as(MIN_NBR)) // src > min_nbr + .persist(intermediateStorageLevel) + currRoundPersistedDFs = currRoundPersistedDFs :+ minNbrs2 + + // connect all smaller neighbors to the min neighbor + ee = ee + .join(minNbrs2, SRC) + .select(col(MIN_NBR).as(SRC), col(DST)) // src <= dst + .filter(col(SRC) =!= col(DST)) // src < dst + // connect self to the min neighbor + ee = ee + .union(minNbrs2.select(col(MIN_NBR).as(SRC), col(SRC).as(DST))) // src < dst + .distinct() + + // checkpointing + if (shouldCheckpoint && (iteration % checkpointInterval == 0)) { + if (useLocalCheckpoints) { + ee = ee.localCheckpoint(eager = true) + } else { + ee = ee.checkpoint(eager = true) + } + } + + ee.persist(intermediateStorageLevel) + currRoundPersistedDFs = currRoundPersistedDFs :+ ee + + minNbrs1 = symmetrize(ee) + .groupBy(SRC) + .agg(min(col(DST)).as(MIN_NBR), count("*").as(CNT)) + .withColumn(MIN_NBR, minValue(col(SRC), col(MIN_NBR))) + .persist(intermediateStorageLevel) + currRoundPersistedDFs = currRoundPersistedDFs :+ minNbrs1 + + // test convergence + val (currSum, edgeCnt) = calcMinNbrSum(minNbrs1) + logInfo(s"$logPrefix Sum of assigned components in iteration $iteration: $currSum.") + + // Pruning Node Optimization: construct a new small graph with fewer nodes, + // and find connected components of the shrunken graph, then join back to get the + // connected components of the original graph. + + // If the graph becomes sparse and current iteration >= $optStartIter, we start to + // try such optimization. However, the optimization is only performed if the shrunken + // graph is much smaller than the original graph, otherwise we do not perform it ( + // in this case, the only additional cost is to determine the size of shrunken graph). + // In current implementation, we only try such optimization one time and it is + // performed at most one time. So the additional cost is bounded. + + // According to such heuristic rule, we can determine when and whether we should + // perform the optimization. For the sparse graphs (defined by sparsityThreshold), + // we will try such optimization at the end of $optStartIter iteration (default is 2). + // For the dense graph, its edges will be pruned at each large/small star join iteration, + // and we will try the optimization once the graph becomes sparse. + if ((edgeCnt < sparsityThreshold * numNodes) && (edgeCnt > 0) + && (iteration >= optStartIter) && (!triedToOptimize)) { + edgesBeforePruning = ee + pruneLeafNodes(ee, intermediateStorageLevel, numNodes, shrinkageThreshold) match { + case Some(r) => + shrunkenGraphNodes = r._1 + ee = r._2 + currRoundPersistedDFs = currRoundPersistedDFs :+ ee + numNodes = r._3 + isOptimized = true + logInfo(s"$logPrefix Pruning node optimization performed in iteration $iteration.") + logInfo(s"$logPrefix Shrunken graph node count: $numNodes.") + case None => + logInfo(s"$logPrefix Pruning node optimization not performed.") + } + triedToOptimize = true + } + + if (currSum == prevSum) { + converged = true + } else { + prevSum = currSum + } + + for (persistedDF <- lastRoundPersistedDFs) { + persistedDF.unpersist() + } + lastRoundPersistedDFs = currRoundPersistedDFs + iteration += 1 + } + + if (isOptimized) { + ee = joinBack(shrunkenGraphNodes, ee, edgesBeforePruning) + } + + logInfo(s"$logPrefix Connected components converged in ${iteration - 1} iterations.") + logInfo(s"$logPrefix Join and return component assignments with original vertex IDs.") + + val output = buildOutput(graph, vv, ee, useLabelsAsComponents) + .persist(intermediateStorageLevel) + + output.count() + + for (persistedDF <- lastRoundPersistedDFs) { + persistedDF.unpersist() + } + if (shrunkenGraphNodes != null) { + shrunkenGraphNodes.unpersist() + } + + resultIsPersistent() + + output + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/mixins.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/mixins.scala new file mode 100644 index 0000000000000..431d734eb7331 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/mixins.scala @@ -0,0 +1,248 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes + +import org.apache.spark.storage.StorageLevel + +private[graphframes] trait WithAlgorithmChoice { + protected val ALGO_GRAPHX = "graphx" + protected val ALGO_GRAPHFRAMES = "graphframes" + protected var algorithm: String = ALGO_GRAPHX + val supportedAlgorithms: Array[String] = Array(ALGO_GRAPHX, ALGO_GRAPHFRAMES) + + /** + * Set an algorithm to use. Supported algorithms are "graphx" and "graphframes". + * + * @param value + * @return + */ + def setAlgorithm(value: String): this.type = { + require( + supportedAlgorithms.contains(value), + s"Supported algorithms are {${supportedAlgorithms.mkString(", ")}}, but got $value.") + algorithm = value + this + } + + def getAlgorithm: String = algorithm +} + +private[graphframes] trait WithCheckpointInterval extends Logging { + protected var checkpointInterval: Int = 2 + + /** + * Sets checkpoint interval in terms of number of iterations (default: 2). Checkpointing + * regularly helps recover from failures, clean shuffle files, shorten the lineage of the + * computation graph, and reduce the complexity of plan optimization. As of Spark 2.0, the + * complexity of plan optimization would grow exponentially without checkpointing. Hence, + * disabling or setting longer-than-default checkpoint intervals are not recommended. Checkpoint + * data is saved under `org.apache.spark.SparkContext.getCheckpointDir` with prefix of the + * algorithm name. If the checkpoint directory is not set, this throws a `java.io.IOException`. + * Set a nonpositive value to disable checkpointing. This parameter is only used when the + * algorithm is set to "graphframes". Its default value might change in the future. + * @see + * `org.apache.spark.SparkContext.setCheckpointDir` in Spark API doc + */ + def setCheckpointInterval(value: Int): this.type = { + if (value <= 0 || value > 2) { + logWarn( + s"Set checkpointInterval to $value. This would blow up the query plan and hang the " + + "driver for large graphs.") + } + checkpointInterval = value + this + } + + // python-friendly setter + private[graphframes] def setCheckpointInterval(value: java.lang.Integer): this.type = { + setCheckpointInterval(value.toInt) + } + + /** + * Gets checkpoint interval. + */ + def getCheckpointInterval: Int = checkpointInterval +} + +private[graphframes] trait WithBroadcastThreshold extends Logging { + protected var broadcastThreshold: Int = 1000000 + + /** + * Sets a broadcast threshold in propagating component assignments (default: 1,000,000). If a + * node degree is greater than this threshold at some iteration, its component assignment will + * be collected and then broadcasted back to propagate the assignment to its neighbors. + * Otherwise, the assignment propagation is done by a normal Spark join. This parameter is only + * used when the algorithm is set to "graphframes". If the value is -1, then the skewness + * problem is left to the Apache Spark AQE optimizer. + * + * **WARNING** using a broadcast threshold is non-free! Under the hood it is calling an action, + * and if a broadcast threshold is set, then AQE is disabled to avoid wrong results! If your + * graph does not contain gigantic components, it is strongly recommended to set this value to + * -1. On benchmarks setting it to -1 gains about x5 better results in performance. + * + * **WARNING** the current default value is 1,000,000. It is left for backward compatibility + * only. In the future versions it may be set to -1 as more reasonable for the most real-world + * cases (e.g., the data deduplication problem). + */ + def setBroadcastThreshold(value: Int): this.type = { + broadcastThreshold = value + this + } + + // python-friendly setter + private[graphframes] def setBroadcastThreshold(value: java.lang.Integer): this.type = { + setBroadcastThreshold(value.toInt) + } + + /** + * Gets broadcast threshold in propagating component assignment. + * @see + * [[org.apache.spark.graphframes.lib.ConnectedComponents.setBroadcastThreshold]] + */ + def getBroadcastThreshold: Int = broadcastThreshold +} + +private[graphframes] trait WithIntermediateStorageLevel extends Logging { + + protected var intermediateStorageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK + + /** + * Sets storage level for intermediate datasets that require multiple passes (default: + * ``MEMORY_AND_DISK``). + */ + def setIntermediateStorageLevel(value: StorageLevel): this.type = { + intermediateStorageLevel = value + this + } + + /** + * Gets storage level for intermediate datasets that require multiple passes. + */ + def getIntermediateStorageLevel: StorageLevel = intermediateStorageLevel + +} + +private[graphframes] trait WithLgNomEntries { + protected var lgNomEntries: Int = 12 + + /** + * Sets the log2 of nominal entries used by Theta sketch aggregations. + */ + def setLgNomEntries(value: Int): this.type = { + require((value >= 4) && (value <= 26), "lgNomEntries must be between 4 and 26") + lgNomEntries = value + this + } + + /** + * Gets log2 of nominal entries used by Theta sketch aggregations. + */ + def getLgNomEntries: Int = lgNomEntries +} + +private[graphframes] trait WithMaxIter { + protected var maxIter: Option[Int] = None + + /** + * The max number of iterations of algorithm to be performed. + */ + def maxIter(value: Int): this.type = { + maxIter = Some(value) + this + } +} + +private[graphframes] trait WithUseLabelsAsComponents { + protected var useLabelsAsComponents: Boolean = false + + /** + * Sets whether to use vertex labels as component identifiers (default: false). When true, + * vertex labels will be used as component identifiers instead of computing connected + * components. + */ + def setUseLabelsAsComponents(value: Boolean): this.type = { + useLabelsAsComponents = value + this + } + + /** + * Gets whether to use vertex labels as component identifiers. + */ + def getUseLabelsAsComponents: Boolean = useLabelsAsComponents +} + +/** + * Provides support for local checkpoints in Spark computations. + * + * Local checkpoints offer a faster alternative to regular checkpoints as they don't require + * configuration of checkpointDir in persistent storage (like HDFS or S3). While being more + * performant, local checkpoints are less reliable since they don't survive node failures and the + * data is not persisted across multiple nodes. + */ +private[graphframes] trait WithLocalCheckpoints { + protected var useLocalCheckpoints: Boolean = false + + /** + * Sets whether to use local checkpoints instead of regular checkpoints (default: false). Local + * checkpoints are faster but less reliable as they don't survive node failures. + * + * @param value + * true to use local checkpoints, false for regular checkpoints + * @return + * this instance + */ + def setUseLocalCheckpoints(value: Boolean): this.type = { + useLocalCheckpoints = value + this + } + + /** + * Gets whether local checkpoints are being used instead of regular checkpoints. + * + * @return + * true if local checkpoints are enabled, false otherwise + */ + def getUseLocalCheckpoints: Boolean = useLocalCheckpoints +} + +/** + * Provides support of graph directions for algorithms. + */ +private[graphframes] trait WithDirection { + protected var isDirected: Boolean = true + + /** + * Sets should graph be cosidered as directed. + * + * @param value + * true to handle graph as directed + * @return + */ + def setIsDirected(value: Boolean): this.type = { + isDirected = value + this + } + + /** + * Gets should graph be considered as directed. + * + * @return + * true if directed + */ + def getIsDirected: Boolean = isDirected +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/pattern/patterns.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/pattern/patterns.scala new file mode 100644 index 0000000000000..ee95049ec64ee --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/pattern/patterns.scala @@ -0,0 +1,300 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.pattern + +import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.graphframes.InvalidParseException + +import scala.collection.mutable +import scala.util.parsing.combinator._ + +/** + * Parser for graph patterns for motif finding. Copied from GraphFrames with minor modification. + */ +private[graphframes] object PatternParser extends RegexParsers { + private val vertexName: Parser[Vertex] = "[a-zA-Z0-9_]+".r ^^ { NamedVertex.apply } + private val anonymousVertex: Parser[Vertex] = "" ^^ { _ => AnonymousVertex } + private val vertex: Parser[Vertex] = "(" ~> (vertexName | anonymousVertex) <~ ")" + private val namedEdge: Parser[Edge] = + vertex ~ "-" ~ "[" ~ "[a-zA-Z0-9_]+".r ~ "]" ~ ("->" | "-") ~ vertex ^^ { + case src ~ "-" ~ "[" ~ name ~ "]" ~ "->" ~ dst => NamedEdge(name, src, dst) + case src ~ "-" ~ "[" ~ name ~ "]" ~ "-" ~ dst => UndirectedEdge(NamedEdge(name, src, dst)) + case _ => throw new GraphFramesUnreachableException() + } + val anonymousEdge: Parser[Edge] = + vertex ~ "-" ~ "[" ~ "]" ~ ("->" | "-") ~ vertex ^^ { + case src ~ "-" ~ "[" ~ "]" ~ "->" ~ dst => AnonymousEdge(src, dst) + case src ~ "-" ~ "[" ~ "]" ~ "-" ~ dst => UndirectedEdge(AnonymousEdge(src, dst)) + case _ => throw new GraphFramesUnreachableException() + } + private val edge: Parser[Edge] = namedEdge | anonymousEdge + private val negatedEdge: Parser[Pattern] = + "!" ~ edge ^^ { case _ ~ e => + Negation(e) + } + private val pattern: Parser[Pattern] = edge | vertex | negatedEdge + val patterns: Parser[List[Pattern]] = repsep(pattern, ";") +} + +private[graphframes] object Pattern { + def parse(s: String): Seq[Pattern] = { + import PatternParser._ + val rewrittenStr: String = rewriteFixedLengthPattern(rewriteIncomingEdges(s)) + val result = parseAll(patterns, rewrittenStr) match { + case result: Success[_] => + result.asInstanceOf[Success[Seq[Pattern]]].get + case result: NoSuccess => + throw new InvalidParseException( + s"Failed to parse bad motif string: '$s'. Returned message: ${result.msg}") + } + assertValidPatterns(result) + result + } + + /** + * Rewrite a motif string if there are incoming edges + */ + private[graphframes] def rewriteIncomingEdges(patterns: String): String = { + val reversedEdge = + """(!?)\(([a-zA-Z0-9_]*)\)<-\[([a-zA-Z0-9_.*]*)\]-\(([a-zA-Z0-9_]*)\)""".r + val bidirectionalEdge = + """(!?)\(([a-zA-Z0-9_]*)\)<-\[([a-zA-Z0-9_.*]*)\]->\(([a-zA-Z0-9_]*)\)""".r + + val outgoingEdges: Seq[String] = patterns.split(";").toSeq.map { pattern => + pattern.trim match { + case reversedEdge(negation, dst, edge, src) => + s"$negation($src)-[$edge]->($dst)" + case bidirectionalEdge(negation, src, edge, dst) => + if (!negation.isEmpty) { + throw new InvalidParseException( + s"Motif finding does not support negated bidirectional edge: '$pattern'.") + } + if (edge.isEmpty || edge.contains("*")) { + s"($src)-[$edge]->($dst);($dst)-[$edge]->($src)" + } else { + s"($src)-[${edge}1]->($dst);($dst)-[${edge}2]->($src)" + } + case original => original + } + } + + outgoingEdges.mkString(";") + } + + /** + * Rewrite fixed-length pattern + */ + private[graphframes] def rewriteFixedLengthPattern(patterns: String): String = { + val fixedLengthPattern = + """(!?)\(([a-zA-Z0-9_]*)\)-\[([a-zA-Z0-9_]*)\*([0-9]+)\]->\(([a-zA-Z0-9_]*)\)""".r + val expandedEdges: Seq[String] = patterns.split(";").toSeq.map { pattern => + pattern.trim match { + case fixedLengthPattern(negation, src, name, num, dst) => + val hop: Int = num.toInt + if (hop > 0) { + val midVertices = + if (src.isEmpty && dst.isEmpty) (1 until hop).map(i => s"__tmpv${i}") + else (1 until hop).map(i => s"_${src}${dst}${i}") + val vertices = src +: midVertices :+ dst + vertices + .sliding(2) + .zipWithIndex + .map { + case (Seq(v1, v2), i) => + if (name.isEmpty) s"${negation}(${v1})-[]->(${v2})" + else s"${negation}(${v1})-[_${name}${i + 1}]->(${v2})" + case _ => + throw new InvalidParseException( + s"Cannot rewrite fixed-length pattern as a chain: '$pattern'.") + } + .mkString(";") + } else { + throw new InvalidParseException(s"Hop must be greater than 0: '$pattern'.") + } + case original => original + } + } + + expandedEdges.mkString(";") + } + + /** + * Checks all Patterns for validity: + * - Disallow named edges within negated terms + * - Disallow term "()-[]->()" and its negation + * - Disallow name to be shared by a vertex and an edge + * @throws InvalidParseException + * if an negated terms contain named edges + */ + private def assertValidPatterns(patterns: Seq[Pattern]): Unit = { + + // vertexNames, edgeNames are used to check for duplicate names across vertices and edges + val vertexNames = mutable.HashSet.empty[String] + val edgeNames = mutable.HashSet.empty[String] + def addVertex(v: Vertex): Unit = v match { + case NamedVertex(name) => + if (edgeNames.contains(name)) { + throw new InvalidParseException( + s"Motif reused name '$name' for both a vertex and " + + "an edge, which is not allowed.") + } + vertexNames += name + case AnonymousVertex => // pass + } + def addEdge(e: Edge): Unit = e match { + case NamedEdge(name, src, dst) => + if (vertexNames.contains(name)) { + throw new InvalidParseException( + s"Motif reused name '$name' for both a vertex and " + + "an edge, which is not allowed.") + } + if (edgeNames.contains(name)) { + throw new InvalidParseException( + s"Motif reused name '$name' for multiple edges, " + + "which is not allowed.") + } + edgeNames += name + addVertex(src) + addVertex(dst) + case AnonymousEdge(src, dst) => + addVertex(src) + addVertex(dst) + case UndirectedEdge(edge) => + addEdge(edge) + } + + patterns.foreach { + case Negation(edge) => + edge match { + case NamedEdge(name, src, dst) => + throw new InvalidParseException( + "Motif finding does not support negated named " + + s"edges, but the given pattern contained: !($src)-[$name]->($dst)") + case AnonymousEdge(AnonymousVertex, AnonymousVertex) => + throw new InvalidParseException( + "Motif finding does not support completely " + + "anonymous negated edges !()-[]->(). Users can check for 0 edges in the graph " + + "using the edges DataFrame.") + case e @ UndirectedEdge(edge) => + edge match { + case AnonymousEdge(AnonymousVertex, AnonymousVertex) => + throw new InvalidParseException( + "Motif finding does not support completely " + + "anonymous negated edges !()-[]-(). Users can check for the existence of edges in the " + + "graph using the edges DataFrame.") + case _ => addEdge(e) + } + case e @ AnonymousEdge(_, _) => + addEdge(e) + } + case AnonymousEdge(AnonymousVertex, AnonymousVertex) => + throw new InvalidParseException( + "Motif finding does not support completely " + + "anonymous edges ()-[]->(). Users can check for the existence of edges in the " + + "graph using the edges DataFrame.") + case e @ UndirectedEdge(edge) => + edge match { + case AnonymousEdge(AnonymousVertex, AnonymousVertex) => + throw new InvalidParseException( + "Motif finding does not support completely " + + "anonymous edges ()-[]-(). Users can check for the existence of edges in the " + + "graph using the edges DataFrame.") + case _ => addEdge(e) + } + case e @ AnonymousEdge(_, _) => + addEdge(e) + case e @ NamedEdge(_, _, _) => + addEdge(e) + case AnonymousVertex => + throw new InvalidParseException( + "Motif finding does not allow a lone anonymous vertex " + + "\"()\" in a motif. Users can check for the existence of vertices in the graph " + + "using the vertices DataFrame.") + case v @ NamedVertex(_) => + addVertex(v) + } + } + + /** + * Return the set of named vertices which only appear in negated terms, in sorted order. + */ + private[graphframes] def findNamedVerticesOnlyInNegatedTerms( + patterns: Seq[Pattern]): Seq[String] = { + val vPos = findNamedElementsInOrder( + patterns.filter(p => !p.isInstanceOf[Negation]), + includeEdges = false).toSet + val vNeg = findNamedElementsInOrder( + patterns.filter(p => p.isInstanceOf[Negation]), + includeEdges = false).toSet + vNeg.diff(vPos).toSeq.sorted + } + + /** + * Return the set of named vertices (and optionally edges) appearing in the given patterns, in + * the order they first appear in the sequence of patterns. + * @param includeEdges + * If true, include named edges in the returned sequence. + */ + private[graphframes] def findNamedElementsInOrder( + patterns: Seq[Pattern], + includeEdges: Boolean): Seq[String] = { + val elementSet = mutable.LinkedHashSet.empty[String] + def findNamedElementsHelper(pattern: Pattern): Unit = pattern match { + case Negation(child) => + findNamedElementsHelper(child) + case UndirectedEdge(child) => + findNamedElementsHelper(child) + elementSet += "_pattern" + elementSet += "_direction" + case AnonymousVertex => // pass + case NamedVertex(name) => + if (!elementSet.contains(name)) { + elementSet += name + } + case AnonymousEdge(src, dst) => + findNamedElementsHelper(src) + findNamedElementsHelper(dst) + case NamedEdge(name, src, dst) => + findNamedElementsHelper(src) + if (includeEdges && !elementSet.contains(name)) { + elementSet += name + } + findNamedElementsHelper(dst) + } + patterns.foreach(findNamedElementsHelper) + elementSet.toSeq + } +} + +private[graphframes] sealed trait Pattern + +private[graphframes] case class Negation(child: Edge) extends Pattern + +private[graphframes] sealed trait Vertex extends Pattern + +private[graphframes] case object AnonymousVertex extends Vertex + +private[graphframes] case class NamedVertex(name: String) extends Vertex + +private[graphframes] sealed trait Edge extends Pattern + +private[graphframes] case class UndirectedEdge(edge: Edge) extends Edge + +private[graphframes] case class AnonymousEdge(src: Vertex, dst: Vertex) extends Edge + +private[graphframes] case class NamedEdge(name: String, src: Vertex, dst: Vertex) extends Edge diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkBase.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkBase.scala new file mode 100644 index 0000000000000..7aafb15e0add1 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkBase.scala @@ -0,0 +1,434 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.rw + +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.array +import org.apache.spark.sql.functions.array_sort +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.collect_list +import org.apache.spark.sql.functions.concat +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.functions.reduce +import org.apache.spark.sql.functions.struct +import org.apache.spark.sql.functions.when +import org.apache.spark.sql.functions.xxhash64 +import org.apache.spark.sql.graphframes.expressions.KMinSampling +import org.apache.spark.sql.types.ArrayType +import org.apache.spark.sql.types.DataType +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithIntermediateStorageLevel + +import scala.util.Random + +/** + * Base trait for implementing random walk algorithms on graph data. Provides common functionality + * for generating random walks across a graph structure. + */ +trait RandomWalkBase extends Serializable with Logging with WithIntermediateStorageLevel { + + /** Maximum number of neighbors to consider per vertex during random walks. */ + protected var maxNbrs: Int = 50 + + /** GraphFrame on which random walks are performed. */ + protected var graph: GraphFrame = null + + /** Number of random walks to generate per node. */ + protected var numWalksPerNode: Int = 5 + + /** Size of each batch in the random walk process. */ + protected var batchSize: Int = 10 + + /** Number of batches to run in the random walk process. */ + protected var numBatches: Int = 5 + + /** Whether to respect edge direction in the graph (true for directed graphs). */ + protected var useEdgeDirection: Boolean = false + + /** Global random seed for reproducibility. */ + protected var globalSeed: Long = 42L + + /** Optional prefix for temporary storage during random walks. */ + protected var temporaryPrefix: Option[String] = None + + /** Unique identifier for the current random walk run. */ + protected var runID: String = java.util.UUID.randomUUID().toString + + /** Starting batch index for continuous mode */ + protected var startingIteration: Int = 1 + + /** Internal handler of vertex data type */ + private var idDataType: DataType = null + + /** + * Sets the graph to perform random walks on. + * + * @param graph + * the GraphFrame to run random walks on + * @return + * this RandomWalkBase instance for chaining + */ + def onGraph(graph: GraphFrame): this.type = { + this.graph = graph + this.idDataType = graph.vertices.schema(GraphFrame.ID).dataType + this + } + + /** + * Sets the temporary prefix for storing intermediate results. + * + * @param value + * the prefix string + * @return + * this RandomWalkBase instance for chaining + */ + def setTemporaryPrefix(value: String): this.type = { + temporaryPrefix = Some(value) + this + } + + /** + * Sets the maximum number of neighbors per vertex. + * + * @param value + * the max number of neighbors + * @return + * this RandomWalkBase instance for chaining + */ + def setMaxNbrsPerVertex(value: Int): this.type = { + maxNbrs = value + this + } + + /** + * Sets the number of walks per node. + * + * @param value + * number of walks + * @return + * this RandomWalkBase instance for chaining + */ + def setNumWalksPerNode(value: Int): this.type = { + numWalksPerNode = value + this + } + + /** + * Sets the batch size. + * + * @param value + * batch size + * @return + * this RandomWalkBase instance for chaining + */ + def setBatchSize(value: Int): this.type = { + batchSize = value + this + } + + /** + * Sets the number of batches. + * + * @param value + * number of batches + * @return + * this RandomWalkBase instance for chaining + */ + def setNumBatches(value: Int): this.type = { + numBatches = value + this + } + + /** + * Sets whether to use edge direction. + * + * @param value + * true if the graph is directed + * @return + * this RandomWalkBase instance for chaining + */ + def setUseEdgeDirection(value: Boolean): this.type = { + useEdgeDirection = value + this + } + + /** + * Sets the global random seed. + * + * @param value + * the seed value + * @return + * this RandomWalkBase instance for chaining + */ + def setGlobalSeed(value: Long): this.type = { + globalSeed = value + this + } + + /** + * Sets the random walk runID. If provided, cached batches from existing random walk run will be + * reused. User should be careful, that temporary prefix points to the right direction as well + * the cached data starting from the set index exists. + * + * @param value + * @return + */ + def setRunId(value: String): this.type = { + require(value != "", "empty string is not supported as run ID") + runID = value + this + } + + /** + * Get the generated (or provided) runID. This method returns current runID! + * + * @return + */ + def getRunId(): String = runID + + /** + * Sets the startng batch index for the continuous mode. See @setWalkId comment for details. + * + * @param value + * @return + */ + def setStartingFromBatch(value: Int): this.type = { + require(value >= 1, s"batches are one-indexed but got $value") + startingIteration = value + this + } + + /** + * Generates a temporary path for a given iteration. + * + * @param iter + * iteration number + * @return + * path string + */ + private def iterationTmpPath(iter: Int): String = if (temporaryPrefix.get.endsWith("/")) { + s"${temporaryPrefix.get}${runID}_batch_${iter}" + } else { + s"${temporaryPrefix.get}/${runID}_batch_${iter}" + } + + private def sortAndConcat(arrCol: Column): Column = { + def ordering(left: Column, right: Column): Column = { + when(left < right, lit(-1)).when(left === right, lit(0)).otherwise(lit(1)) + } + val sorted = array_sort( + arrCol, + (left, right) => + ordering( + left.getField(RandomWalkBase.batchIDColName), + right.getField(RandomWalkBase.batchIDColName))) + + reduce( + sorted, + array().cast(ArrayType(idDataType)), + (left, right) => concat(left, right.getField(RandomWalkBase.rwColName))) + } + + /** + * Executes the random walk algorithm on the set graph. + * + * @return + * DataFrame containing the random walks + */ + def run(): DataFrame = { + if (graph == null) { + throw new IllegalArgumentException("Graph is not set") + } + if (temporaryPrefix.isEmpty) { + throw new IllegalArgumentException("Temporary prefix is required for random walks.") + } + + logInfo(s"Starting random walk with runID: $runID") + + val iterationsRng = new Random(globalSeed) + val spark = graph.vertices.sparkSession + + // If we're starting from a batch index > 1, we need to skip the seeds for previous batches + // to ensure the same sequence of random numbers as if we started from batch 1 + if (startingIteration > 1) { + logInfo(s"Skipping ${startingIteration - 1} seeds to maintain seed consistency") + for (_ <- 1 until startingIteration) { + iterationsRng.nextLong() + } + } + + for (i <- startingIteration to numBatches) { + logInfo(s"Starting batch $i of $numBatches") + val iterSeed = iterationsRng.nextLong() + val preparedGraph = prepareGraph(iterSeed) + val prevIterationDF = if (i == 1) { None } + else { + Some(spark.read.parquet(iterationTmpPath(i - 1))) + } + val iterationResult: DataFrame = runIter(preparedGraph, prevIterationDF, iterSeed) + .withColumn(RandomWalkBase.batchIDColName, lit(i)) + iterationResult.write.mode("overwrite").parquet(iterationTmpPath(i)) + } + + logInfo("Finished all batches, merging results.") + + val result = (1 to numBatches) + .map(i => spark.read.parquet(iterationTmpPath(i))) + .reduce((a, b) => a.union(b)) + .groupBy(RandomWalkBase.walkIdCol) + .agg(sortAndConcat( + collect_list(struct(RandomWalkBase.batchIDColName, RandomWalkBase.rwColName))) + .alias(RandomWalkBase.rwColName)) + .persist(intermediateStorageLevel) + + val cnt = result.count() + resultIsPersistent() + logInfo(s"$cnt random walks are returned") + result + } + + /** + * Deletes all temporary files associated with a given instance. This method uses Hadoop + * FileSystem to remove the directory containing batch files for the specified run ID. The + * temporary prefix must be set and accessible via the current SparkContext's Hadoop + * configuration. + */ + def cleanUp(): Unit = { + if (temporaryPrefix.isEmpty) { + throw new IllegalArgumentException("Temporary prefix is required for clean-up.") + } + val spark = graph.vertices.sparkSession + RandomWalkBase.cleanUp(temporaryPrefix.get, runID, numBatches, spark) + logInfo(s"Clean-up completed for run ID: $runID") + } + + /** + * Prepares the graph for random walk by limiting neighbors and handling direction. + * + * @return + * prepared GraphFrame + */ + protected def prepareGraph(iterationSeed: Long): GraphFrame = { + val preAggs = (if (useEdgeDirection) { + graph.edges + .select(col(GraphFrame.SRC), col(GraphFrame.DST)) + + } else { + graph.edges + .select(GraphFrame.SRC, GraphFrame.DST) + .union(graph.edges.select(GraphFrame.DST, GraphFrame.SRC)) + .distinct() + }) + // xxhash64(src, dst, seed) ~= fault tolerant random order + .withColumn( + "rand_rank", + xxhash64(col(GraphFrame.SRC), col(GraphFrame.DST), lit(iterationSeed))) + .groupBy(col(GraphFrame.SRC).alias(GraphFrame.ID)) + + // typed sampling aggregator + val dataType = graph.vertices.schema(GraphFrame.ID).dataType + val encoder = KMinSampling.getEncoder( + graph.vertices.sparkSession, + dataType, + Seq(GraphFrame.ID, "rand_rank")) + val kMinSamplingUDAF = + KMinSampling.fromSparkType(dataType, maxNbrs, encoder) + + // at most maxNbrs per vertex with a stable uniform sampling + val vertices = preAggs.agg( + kMinSamplingUDAF(col(GraphFrame.DST), col("rand_rank")) + .alias(RandomWalkBase.nbrsColName)) + + val edges = graph.edges + GraphFrame(vertices, edges) + } + + /** + * Runs a single iteration of the random walk. + * + * @param graph + * prepared graph + * @param prevIterationDF + * DataFrame from previous iteration (if any) + * @param iterSeed + * seed for this iteration + * @return + * DataFrame result of this iteration + */ + protected def runIter( + graph: GraphFrame, + prevIterationDF: Option[DataFrame], + iterSeed: Long): DataFrame +} + +object RandomWalkBase extends Serializable { + + /** Column name for the random walk array. */ + val rwColName: String = "random_walk" + + /** Column name for the unique walk ID. */ + val walkIdCol: String = "random_walk_uuid" + + /** Column name for neighbors list. */ + val nbrsColName: String = "random_walk_nbrs" + + /** Column name for the current visiting vertex. */ + val currVisitingVertexColName: String = "random_walk_curr_vertex" + + /** Column name for batch ID inside walk. */ + val batchIDColName: String = "random_walk_batch_it" + + /** + * Deletes all temporary files associated with a given runID. This method uses Hadoop FileSystem + * to remove the directory containing batch files for the specified run ID. The temporary prefix + * must be set and accessible via the current SparkContext's Hadoop configuration. + * + * @param temporaryPrefix + * the temporary prefix path + * @param runID + * the run ID used to generate batch directories + * @param numBatches + * the number of batch directories to look for + * @param spark + * the SparkSession used to get the Hadoop configuration + */ + def cleanUp( + temporaryPrefix: String, + runID: String, + numBatches: Int, + spark: org.apache.spark.sql.SparkSession): Unit = { + val sc = spark.sparkContext + val hadoopConf = sc.hadoopConfiguration + val fs = org.apache.hadoop.fs.FileSystem.get(hadoopConf) + val basePath = temporaryPrefix + val runPath = if (basePath.endsWith("/")) { + s"${basePath}${runID}_batch_" + } else { + s"${basePath}/${runID}_batch_" + } + // Delete all batch directories (1 to numBatches) + for (i <- 1 to numBatches) { + val path = new org.apache.hadoop.fs.Path(s"${runPath}${i}") + if (fs.exists(path)) { + fs.delete(path, true) // recursive delete + } + } + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestart.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestart.scala new file mode 100644 index 0000000000000..06302a15a8455 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestart.scala @@ -0,0 +1,103 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.rw + +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types.ArrayType +import org.apache.spark.graphframes.GraphFrame + +/** + * An implementation of random walk with restart. At each step of the walk, there is a probability + * (defined by restartProbability) to reset the walk to the original starting node, otherwise the + * walk continues to a random neighbor. + */ +/** + * An implementation of random walk with restart. At each step of the walk, there is a probability + * (defined by restartProbability) to reset the walk to the original starting node, otherwise the + * walk continues to a random neighbor. + */ +class RandomWalkWithRestart extends RandomWalkBase { + + /** The probability of restarting the walk at each step (resets to starting node). */ + private var restartProbability: Double = 0.1 + + /** + * Sets the restart probability for the random walk. + * + * @param value + * the probability value (between 0.0 and 1.0) + * @return + * this RandomWalkWithRestart instance for chaining + */ + def setRestartProbability(value: Double): this.type = { + restartProbability = value + this + } + + override protected def runIter( + graph: GraphFrame, + prevIterationDF: Option[DataFrame], + iterSeed: Long): DataFrame = { + val neighbors = graph.vertices.select(col(GraphFrame.ID), col(RandomWalkBase.nbrsColName)) + val walksDtype = ArrayType(graph.vertices.schema(GraphFrame.ID).dataType) + var walks = if (prevIterationDF.isEmpty) { + graph.vertices.select( + col(GraphFrame.ID).alias("startingNode"), + col(GraphFrame.ID).alias(RandomWalkBase.currVisitingVertexColName), + explode( + when( + array_size(col(RandomWalkBase.nbrsColName)) > lit(0), + array((0 until numWalksPerNode).map(_ => uuid()): _*)).otherwise(array())) + .alias(RandomWalkBase.walkIdCol), + array().cast(walksDtype).alias(RandomWalkBase.rwColName)) + } else { + prevIterationDF.get.select( + col("startingNode"), + col(RandomWalkBase.currVisitingVertexColName), + col(RandomWalkBase.walkIdCol), + array().cast(walksDtype).alias(RandomWalkBase.rwColName)) + } + + val localRandom = new util.Random(iterSeed) + + for (_ <- (0 until batchSize)) { + val currentSeed = localRandom.nextLong() + + walks = walks + .join( + neighbors, + col(GraphFrame.ID) === col(RandomWalkBase.currVisitingVertexColName), + "left") + .withColumn("doRestart", rand(currentSeed) <= lit(restartProbability)) + .withColumn( + "nextNode", + when(col("doRestart"), col("startingNode")).otherwise( + element_at(shuffle(col(RandomWalkBase.nbrsColName)), 1))) + .select( + col(RandomWalkBase.walkIdCol), + col("startingNode"), + col("nextNode").alias(RandomWalkBase.currVisitingVertexColName), + array_append( + col(RandomWalkBase.rwColName), + col(RandomWalkBase.currVisitingVertexColName)).alias(RandomWalkBase.rwColName)) + } + + walks + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFrameInternals.scala b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFrameInternals.scala new file mode 100644 index 0000000000000..308ade66076e8 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFrameInternals.scala @@ -0,0 +1,137 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.graphframes + +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute +import org.apache.spark.sql.catalyst.analysis.UnresolvedExtractValue +import org.apache.spark.sql.catalyst.expressions.AttributeReference +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.expressions.GetStructField +import org.apache.spark.sql.catalyst.expressions.Literal +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.classic.ClassicConversions._ +import org.apache.spark.sql.classic.{DataFrame => ClassicDataFrame} +import org.apache.spark.sql.classic.Dataset +import org.apache.spark.sql.classic.ExpressionUtils +import org.apache.spark.sql.classic.{SparkSession => ClassicSparkSession} + +import scala.collection.mutable + +object GraphFrameInternals { + + /** + * Extracts all column references from a Column expression, returning a map from top-level + * prefix to the set of nested field names accessed under that prefix. + * + * For nested column references like "src.id" or "edge.weight", this returns Map("src" -> + * Set("id"), "edge" -> Set("weight")). For top-level references like "src" (the whole struct), + * it returns Map("src" -> Set()). + * + * This handles both unresolved expressions (UnresolvedAttribute, UnresolvedExtractValue) and + * resolved expressions (AttributeReference, GetStructField). + * + * Note: Deeply nested struct access (e.g., "dst.location.city") is not fully parsed. In such + * cases, the prefix is recorded with an empty field set, which causes callers to conservatively + * assume the entire struct is needed. This is the safe/correct fallback behavior. + * + * @param spark + * the SparkSession (needed for expression conversion in Spark 4) + * @param expr + * the Column expression to analyze + * @return + * a Map from column prefix to the set of nested field names accessed + */ + def extractColumnReferences(spark: SparkSession, expr: Column): Map[String, Set[String]] = { + val refs = mutable.Map.empty[String, mutable.Set[String]] + + def addRef(prefix: String, field: Option[String]): Unit = { + val fields = refs.getOrElseUpdate(prefix, mutable.Set.empty[String]) + field.foreach(fields += _) + } + + val converted = spark.asInstanceOf[ClassicSparkSession].converter(expr.node) + converted.foreach { + // Unresolved: col("src.id") -> UnresolvedAttribute(Seq("src", "id")) + case UnresolvedAttribute(nameParts) if nameParts.nonEmpty => + addRef(nameParts.head, nameParts.lift(1)) + + // Unresolved: col("src")("id") -> UnresolvedExtractValue + case UnresolvedExtractValue(child, extraction) => + child match { + case UnresolvedAttribute(nameParts) if nameParts.nonEmpty => + extraction match { + case Literal(fieldName: String, _) => addRef(nameParts.head, Some(fieldName)) + case Literal(fieldName, _) if fieldName != null => + // Handle UTF8String (Spark's internal string representation) + addRef(nameParts.head, Some(fieldName.toString)) + case _ => addRef(nameParts.head, None) // Unknown field access + } + case _ => // Nested extraction we can't easily parse - conservative fallback + } + + // Resolved: AttributeReference for top-level columns + case attr: AttributeReference => + addRef(attr.name, None) + + // Resolved: GetStructField for nested field access like struct.field + // Note: Only handles single-level nesting; deeper nesting falls through to default case + case GetStructField(child, _, Some(fieldName)) => + child match { + case attr: AttributeReference => addRef(attr.name, Some(fieldName)) + case _ => // Deeply nested struct access - conservative fallback (join will be used) + } + + case _ => // ignore other expression types + } + + refs.map { case (k, v) => k -> v.toSet }.toMap + } + + /** + * Apply the given SQL expression (such as `id = 3`) to the field in a column, rather than to + * the column itself. + * + * @param expr + * SQL expression, such as `id = 3` + * @param colName + * Column name, such as `myVertex` + * @return + * SQL expression applied to the column fields, such as `myVertex.id = 3` + */ + def applyExprToCol(spark: SparkSession, expr: Column, colName: String): Column = { + val converted = spark.asInstanceOf[ClassicSparkSession].converter(expr.node) + ExpressionUtils.column(converted.transform { case UnresolvedAttribute(nameParts) => + UnresolvedAttribute(colName +: nameParts) + }) + } + + def createDataFrame(spark: SparkSession, plan: LogicalPlan): DataFrame = { + Dataset.ofRows(spark.asInstanceOf[ClassicSparkSession], plan) + } + + def planFromDataFrame(df: DataFrame): LogicalPlan = { + df.asInstanceOf[ClassicDataFrame].logicalPlan + } + + def createColumn(expr: Expression): Column = { + Column(expr) + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala new file mode 100644 index 0000000000000..38e3e19ee8e14 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala @@ -0,0 +1,144 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.graphframes + +import org.apache.spark.internal.config.ConfigEntry +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.storage.StorageLevel + +object GraphFramesConf { + private val USE_LOCAL_CHECKPOINTS = + SQLConf + .buildConf("spark.graphframes.useLocalCheckpoints") + .doc(""" Tells the connected components algorithm to use local checkpoints (default: "false"). + | If set to "true", iterative algorithm will use the checkpointing mechanism to the persistent storage. + | Local checkpoints are faster but can make the whole job less prone to errors. + | @note This option may become default "true" in the future. + |""".stripMargin) + .version("0.9.3") + .booleanConf + .createOptional + + private val USE_LABELS_AS_COMPONENTS = + SQLConf + .buildConf("spark.graphframes.useLabelsAsComponents") + .doc(""" Tells the connected components algorithm to use (default: "true") labels as components in the output + | DataFrame. If set to "false", randomly generated labels with the data type LONG will returned. + |""".stripMargin) + .version("0.9.0") + .booleanConf + .createOptional + + private val CONNECTED_COMPONENTS_ALGORITHM = + SQLConf + .buildConf("spark.graphframes.connectedComponents.algorithm") + .doc(""" Sets the connected components algorithm to use (default: "graphframes"). Supported algorithms + | - "two_phase": Uses alternating large star and small star iterations proposed in + | [[http://dx.doi.org/10.1145/2670979.2670997 Connected Components in MapReduce and Beyond]] + | - "randomized_contraction": Uses randomized algorithm proposed in + | [[https://arxiv.org/pdf/1802.09478 In-database connected component analysis]] + | - "graphframes": Deprecated alias for "two_phase" + | - "graphx": Converts the graph to a GraphX graph and then uses the connected components + | implementation in GraphX. + | @see org.apache.spark.graphframes.lib.ConnectedComponents.supportedAlgorithms""".stripMargin) + .version("0.9.0") + .stringConf + .createOptional + + private val CONNECTED_COMPONENTS_BROADCAST_THRESHOLD = + SQLConf + .buildConf("spark.graphframes.connectedComponents.broadcastthreshold") + .doc(""" Sets broadcast threshold in propagating component assignments (default: 1000000). If a node + | degree is greater than this threshold at some iteration, its component assignment will be + | collected and then broadcasted back to propagate the assignment to its neighbors. Otherwise, + | the assignment propagation is done by a normal Spark join. This parameter is only used when + | the algorithm is set to "graphframes".""".stripMargin) + .version("0.9.0") + .intConf + .createOptional + + private val CONNECTED_COMPONENTS_CHECKPOINT_INTERVAL = + SQLConf + .buildConf("spark.graphframes.connectedComponents.checkpointinterval") + .doc(""" Sets checkpoint interval in terms of number of iterations (default: 2). Checkpointing + | regularly helps recover from failures, clean shuffle files, shorten the lineage of the + | computation graph, and reduce the complexity of plan optimization. As of Spark 2.0, the + | complexity of plan optimization would grow exponentially without checkpointing. Hence, + | disabling or setting longer-than-default checkpoint intervals are not recommended. Checkpoint + | data is saved under `org.apache.spark.SparkContext.getCheckpointDir` with prefix + | "connected-components". If the checkpoint directory is not set, this throws a + | `java.io.IOException`. Set a nonpositive value to disable checkpointing. This parameter is + | only used when the algorithm is set to "graphframes". Its default value might change in the + | future. + | @see `org.apache.spark.SparkContext.setCheckpointDir` in Spark API doc""".stripMargin) + .version("0.9.0") + .intConf + .createOptional + + private val CONNECTED_COMPONENTS_INTERMEDIATE_STORAGE_LEVEL = + SQLConf + .buildConf("spark.graphframes.connectedComponents.intermediatestoragelevel") + .doc("Sets storage level for intermediate datasets that require multiple passes (default: ``MEMORY_AND_DISK``).") + .version("0.9.0") + .stringConf + .createOptional + + private def get(entry: ConfigEntry[_]): Option[String] = { + try { + Option(SparkSession.getActiveSession.get.conf.get(entry.key)) + } catch { + case _: NoSuchElementException => None + } + } + + def getConnectedComponentsAlgorithm: Option[String] = { + get(CONNECTED_COMPONENTS_ALGORITHM) match { + case Some(threshold) => Some(threshold.toLowerCase) + case _ => None + } + } + + def getConnectedComponentsBroadcastThreshold: Option[Int] = { + get(CONNECTED_COMPONENTS_BROADCAST_THRESHOLD) match { + case Some(threshold) => Some(threshold.toInt) + case _ => None + } + } + + def getConnectedComponentsCheckpointInterval: Option[Int] = { + get(CONNECTED_COMPONENTS_CHECKPOINT_INTERVAL) match { + case Some(interval) => Some(interval.toInt) + case _ => None + } + } + + def getConnectedComponentsStorageLevel: Option[StorageLevel] = { + get(CONNECTED_COMPONENTS_INTERMEDIATE_STORAGE_LEVEL) match { + case Some(level) => Some(StorageLevel.fromString(level.toUpperCase)) + case _ => None + } + } + + def getUseLabelsAsComponents: Option[Boolean] = get(USE_LABELS_AS_COMPONENTS) match { + case Some(use) => Some(use.toBoolean) + case _ => None + } + + def getUseLocalCheckpoints: Option[Boolean] = get(USE_LOCAL_CHECKPOINTS).map(_.toBoolean) +} diff --git a/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/FiniteAXPlusB.scala b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/FiniteAXPlusB.scala new file mode 100644 index 0000000000000..201882cd81dff --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/FiniteAXPlusB.scala @@ -0,0 +1,102 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.graphframes.expressions + +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.expressions.TernaryExpression +import org.apache.spark.sql.catalyst.expressions.codegen.Block._ +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback +import org.apache.spark.sql.catalyst.expressions.codegen.ExprCode +import org.apache.spark.sql.types.DataType +import org.apache.spark.sql.types.LongType + +case class FiniteAXPlusB(first: Expression, second: Expression, third: Expression) + extends TernaryExpression + with CodegenFallback { + override def dataType: DataType = LongType + + override protected def withNewChildrenInternal( + newFirst: Expression, + newSecond: Expression, + newThird: Expression): Expression = copy(newFirst, newSecond, newThird) + + override protected def nullSafeEval(input1: Any, input2: Any, input3: Any): Any = { + val a = input1.asInstanceOf[Long] + val x = input2.asInstanceOf[Long] + val b = input3.asInstanceOf[Long] + + FiniteAXPlusB.axpb(a, x, b) + } + + override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { + val a = ctx.freshName("a") + val x = ctx.freshName("x") + val b = ctx.freshName("b") + val r = ctx.freshName("r") + + val aGenCode = first.genCode(ctx) + val xGenCode = second.genCode(ctx) + val bGenCode = third.genCode(ctx) + + ev.copy(code = code""" + ${aGenCode.code} + ${xGenCode.code} + ${bGenCode.code} + long $a = ${aGenCode.value}; + long $x = ${xGenCode.value}; + long $b = ${bGenCode.value}; + long $r = 0L; + long irrpoly = 0x1bL; + while ($x != 0L) { + if (($x & 1L) != 0L) { + $r ^= $a; + } + $x = ($x >>> 1) & 0x7fffffffffffffffL; + if (($a & (1L << 63)) != 0L) { + $a = ($a << 1) ^ irrpoly; + } else { + $a <<= 1; + } + } + boolean ${ev.isNull} = false; + long ${ev.value} = $r ^ $b; + """) + } +} + +object FiniteAXPlusB extends Serializable { + def axpb(a: Long, x: Long, b: Long): Long = { + var r = 0L + val irrpoly = 0x1bL + var currentA = a + var currentX = x + while (currentX != 0L) { + if ((currentX & 1L) != 0L) { + r ^= currentA + } + currentX = (currentX >>> 1) & 0x7fffffffffffffffL + if ((currentA & (1L << 63)) != 0L) { + currentA = (currentA << 1) ^ irrpoly + } else { + currentA <<= 1 + } + } + r ^ b + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KCoreMerge.scala b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KCoreMerge.scala new file mode 100644 index 0000000000000..32d637c516d59 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KCoreMerge.scala @@ -0,0 +1,118 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.graphframes.expressions + +import org.apache.spark.sql.catalyst.expressions.BinaryExpression +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.expressions.codegen.Block._ +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenContext +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback +import org.apache.spark.sql.catalyst.expressions.codegen.ExprCode +import org.apache.spark.sql.catalyst.util.ArrayData +import org.apache.spark.sql.types.DataType +import org.apache.spark.sql.types.IntegerType + +/** + * Mandal, Aritra, and Mohammad Al Hasan. "A distributed k-core decomposition algorithm on spark." + * 2017 IEEE International Conference on Big Data (Big Data). IEEE, 2017. + * + * @param left + * array of nbrs cores + * @param right + * core of the vertex + */ +case class KCoreMerge(left: Expression, right: Expression) + extends BinaryExpression + with CodegenFallback { + override protected def withNewChildrenInternal( + newLeft: Expression, + newRight: Expression): Expression = copy(newLeft, newRight) + + override def dataType: DataType = IntegerType + + /** + * Each node initializes its core value with the degree of itself. Each node (say u) then sends + * messages to its neighbors v ∈ N (u) with the current estimate of its (u’s) core value. For an + * undirected graph with m edges, there can be at most a total of 2m messages that have been + * sent during a message passing session. Upon receiving all the messages from its neighbors, + * the vertex u computes the largest value l such that the number of neighbors of u whose + * current core value estimate is `l` or larger is equal or higher than `l` + */ + override protected def nullSafeEval(input1: Any, input2: Any): Any = { + val arrayOfElements = input1.asInstanceOf[ArrayData].toIntArray() + val currentCore = input2.asInstanceOf[Int] + + val counts = arrayOfElements.foldLeft(new Array[Int](currentCore + 1))((acc, el) => + if (el > currentCore) { + acc(currentCore) = acc(currentCore) + 1 + acc + } else { + acc(el) = acc(el) + 1 + acc + }) + + var currentWeight = 0 + for (i <- currentCore to 1 by -1) { + currentWeight += counts(i) + if (i <= currentWeight) { + return i + } + } + + return 0 + } + + override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { + val arrayOfElements = ctx.freshName("arrayOfElements") + val currentCore = ctx.freshName("currentCore") + val counts = ctx.freshName("counts") + val currentWeight = ctx.freshName("currentWeight") + val el = ctx.freshName("el") + val i = ctx.freshName("i") + + val leftGenCode = left.genCode(ctx) + val rightGenCode = right.genCode(ctx) + ev.copy(code""" + |${leftGenCode.code} + |${rightGenCode.code} + |int ${ev.value} = 0; + |boolean ${ev.isNull} = false; + |int[] $arrayOfElements = ${leftGenCode.value}.toIntArray(); + |int $currentCore = ${rightGenCode.value}; + | + |int[] $counts = new int[$currentCore + 1]; + |for (int $i = 0; $i < $arrayOfElements.length; $i++) { + | int $el = $arrayOfElements[$i]; + | if ($el > $currentCore) { + | $counts[$currentCore] += 1; + | } else { + | $counts[$el] += 1; + | } + |} + | + |int $currentWeight = 0; + |for (int $i = $currentCore; $i >= 1; $i--) { + | $currentWeight += $counts[$i]; + | if ($i <= $currentWeight) { + | ${ev.value} = $i; + | break; + | } + |} + """.stripMargin) + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KMinSampling.scala b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KMinSampling.scala new file mode 100644 index 0000000000000..0401ac8ff779e --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KMinSampling.scala @@ -0,0 +1,182 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.graphframes.expressions + +import org.apache.spark.sql.Encoder +import org.apache.spark.sql.Encoders +import org.apache.spark.sql.Row +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder +import org.apache.spark.sql.expressions.Aggregator +import org.apache.spark.sql.expressions.UserDefinedFunction +import org.apache.spark.sql.functions.udaf +import org.apache.spark.sql.types._ +import org.apache.spark.sql.types.DataType +import org.apache.spark.graphframes.GraphFramesUnsupportedVertexTypeException + +import scala.annotation.nowarn +import scala.reflect.ClassTag +import scala.reflect.runtime.universe.TypeTag + +case class KMinAccum[T](values: Array[T], weights: Array[Long], var cnt: Int) extends Serializable + +case class KMinSampling[T: ClassTag](size: Int)(implicit + @nowarn tag: TypeTag[T], + ord: Ordering[T]) + extends Aggregator[Row, KMinAccum[T], Seq[T]] + with Serializable { + + override def zero: KMinAccum[T] = KMinAccum(Array.ofDim[T](size), Array.ofDim[Long](size), 0) + + override def reduce(b: KMinAccum[T], a: Row): KMinAccum[T] = { + val newWeight = a.getLong(1) + val newValue = a.getAs[T](0) + // fast-path: buffer is already full of "strong" elements + // the case of "influencer" vertex + if (b.cnt == size) { + val lastWeight = b.weights.last + if ((lastWeight < newWeight) || ((lastWeight == newWeight) && (ord.compare( + newValue, + b.values.last) >= 0))) { + return b + } + } + + // slow-path: custom binary search for (Weight, Value) + // We want to find the first index where (b.w, b.v) > (newWeight, newValue) + var low = 0 + var high = b.cnt - 1 + var idx = b.cnt // Default insertion point is at the end + + while (low <= high) { + val mid = (low + high) / 2 + val midWeight = b.weights(mid) + + // Compare (midWeight, midValue) vs (newWeight, newValue) + val res = + if (midWeight < newWeight) -1 + else if (midWeight > newWeight) 1 + else ord.compare(b.values(mid), newValue) + + if (res <= 0) { + // mid is smaller or equal: we must insert after mid + low = mid + 1 + } else { + // mid is larger: potential insertion point here + idx = mid + high = mid - 1 + } + } + + if (idx < size) { + val newCount = math.min(b.cnt + 1, size) + if (idx < newCount - 1) { + // shift to the right if needed + System.arraycopy(b.weights, idx, b.weights, idx + 1, newCount - idx - 1) + System.arraycopy(b.values, idx, b.values, idx + 1, newCount - idx - 1) + } + + b.weights(idx) = newWeight + b.values(idx) = newValue + b.cnt = newCount + } + + b + } + + override def merge(b1: KMinAccum[T], b2: KMinAccum[T]): KMinAccum[T] = { + + if (b1.cnt == 0) { + return b2 + } + + if (b2.cnt == 0) { + return b1 + } + + val resultSize = math.min(b1.cnt + b2.cnt, size) + val newValues = Array.ofDim[T](resultSize) + val newWeights = Array.ofDim[Long](resultSize) + + var i = 0 + var j = 0 + var r = 0 + + while (r < resultSize) { + val useLeft = if (i >= b1.cnt) { + false + } else if (j >= b2.cnt) { + true + } else { + val wLeft = b1.weights(i) + val wRight = b2.weights(j) + + if (wLeft < wRight) { + true + } else if (wLeft > wRight) { + false + } else { + ord.compare(b1.values(i), b2.values(j)) <= 0 + } + } + + if (useLeft) { + newWeights(r) = b1.weights(i) + newValues(r) = b1.values(i) + i += 1 + } else { + newWeights(r) = b2.weights(j) + newValues(r) = b2.values(j) + j += 1 + } + + r += 1 + } + + KMinAccum(newValues, newWeights, resultSize) + } + + override def finish(reduction: KMinAccum[T]): Seq[T] = + reduction.values.slice(0, reduction.cnt).toSeq + // TODO: replace by Kryo after 4.0.2 is released, see SPARK-52819 + override def bufferEncoder: Encoder[KMinAccum[T]] = Encoders.product + override def outputEncoder: Encoder[Seq[T]] = ExpressionEncoder[Seq[T]]() +} + +object KMinSampling extends Serializable { + def getEncoder(spark: SparkSession, dataType: DataType, colNames: Seq[String]): Encoder[Row] = { + // That is very stupid way actually. But it is the only way with public API + spark + .createDataFrame( + java.util.Collections.emptyList[Row](), + StructType( + StructField(colNames(0), dataType) :: StructField(colNames(1), LongType) :: Nil)) + .encoder + } + + def fromSparkType(dataType: DataType, size: Int, encoder: Encoder[Row]): UserDefinedFunction = { + dataType match { + case StringType => udaf(KMinSampling[java.lang.String](size), encoder) + case ShortType => udaf(KMinSampling[java.lang.Short](size), encoder) + case ByteType => udaf(KMinSampling[java.lang.Byte](size), encoder) + case IntegerType => udaf(KMinSampling[java.lang.Integer](size), encoder) + case LongType => udaf(KMinSampling[java.lang.Long](size), encoder) + case _ => throw new GraphFramesUnsupportedVertexTypeException("unsupported vertex type") + } + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameInternalsSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameInternalsSuite.scala new file mode 100644 index 0000000000000..430cfe27588a9 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameInternalsSuite.scala @@ -0,0 +1,214 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes + +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.graphframes.GraphFrameInternals +import org.apache.spark.graphframes.lib.Pregel + +/** + * Unit tests for GraphFrameInternals.extractColumnReferences. + * + * These tests verify that column references are correctly extracted from various expression + * patterns, which is critical for the Pregel dst join optimization. The optimization skips a join + * when dst columns are not referenced, so we must correctly identify all column references. + */ +class GraphFrameInternalsSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + // ============================================================================ + // Basic Column References + // ============================================================================ + + test("extractColumnReferences - simple dot notation") { + val refs = GraphFrameInternals.extractColumnReferences(spark, col("src.id")) + assert(refs === Map("src" -> Set("id"))) + } + + test("extractColumnReferences - Pregel.src helper") { + val refs = GraphFrameInternals.extractColumnReferences(spark, Pregel.src("rank")) + assert(refs === Map("src" -> Set("rank"))) + } + + test("extractColumnReferences - Pregel.dst helper") { + val refs = GraphFrameInternals.extractColumnReferences(spark, Pregel.dst("value")) + assert(refs === Map("dst" -> Set("value"))) + } + + test("extractColumnReferences - Pregel.edge helper") { + val refs = GraphFrameInternals.extractColumnReferences(spark, Pregel.edge("weight")) + assert(refs === Map("edge" -> Set("weight"))) + } + + test("extractColumnReferences - whole struct reference") { + val refs = GraphFrameInternals.extractColumnReferences(spark, col("dst")) + assert(refs === Map("dst" -> Set())) + } + + // ============================================================================ + // Bracket Notation + // ============================================================================ + + test("extractColumnReferences - bracket notation") { + val refs = GraphFrameInternals.extractColumnReferences(spark, col("src")("id")) + assert(refs === Map("src" -> Set("id"))) + } + + // ============================================================================ + // Complex Expressions with Multiple References + // ============================================================================ + + test("extractColumnReferences - arithmetic with multiple refs from same prefix") { + val refs = + GraphFrameInternals.extractColumnReferences( + spark, + Pregel.src("rank") / Pregel.src("outDegree")) + assert(refs === Map("src" -> Set("rank", "outDegree"))) + } + + test("extractColumnReferences - expression with src, dst, and edge") { + val refs = GraphFrameInternals.extractColumnReferences( + spark, + Pregel.src("value") + Pregel.dst("value") + Pregel.edge("weight")) + assert(refs === Map("src" -> Set("value"), "dst" -> Set("value"), "edge" -> Set("weight"))) + } + + test("extractColumnReferences - when/case expression") { + val refs = GraphFrameInternals.extractColumnReferences( + spark, + when(Pregel.dst("value") > Pregel.src("value"), Pregel.edge("weight"))) + assert(refs.contains("dst"), "Should detect dst reference") + assert(refs.contains("src"), "Should detect src reference") + assert(refs.contains("edge"), "Should detect edge reference") + } + + test("extractColumnReferences - coalesce with multiple refs") { + val refs = + GraphFrameInternals.extractColumnReferences( + spark, + coalesce(col("dst.value"), col("src.default"))) + assert(refs.contains("dst")) + assert(refs.contains("src")) + } + + // ============================================================================ + // Column Used as Map/Array Key - Critical Cases! + // These verify that foreach traversal catches column refs used as arguments + // ============================================================================ + + test("extractColumnReferences - column used as map key via element_at") { + // element_at(col("edge.weights"), col("dst.name")) + // Should detect BOTH "edge" and "dst" references + val refs = + GraphFrameInternals.extractColumnReferences( + spark, + element_at(col("edge.weights"), col("dst.name"))) + assert(refs.contains("edge"), "Should detect edge reference (the map)") + assert(refs.contains("dst"), "Should detect dst reference (used as map key)") + } + + test("extractColumnReferences - column used as array index via element_at") { + // element_at(col("edge.values"), col("dst.index")) + // Should detect BOTH "edge" and "dst" references + val refs = + GraphFrameInternals.extractColumnReferences( + spark, + element_at(col("edge.values"), col("dst.index"))) + assert(refs.contains("edge"), "Should detect edge reference (the array)") + assert(refs.contains("dst"), "Should detect dst reference (used as array index)") + } + + test("extractColumnReferences - column in nested function call") { + // concat(col("src.prefix"), col("dst.suffix")) + val refs = + GraphFrameInternals.extractColumnReferences( + spark, + concat(col("src.prefix"), col("dst.suffix"))) + assert(refs.contains("src")) + assert(refs.contains("dst")) + } + + test("extractColumnReferences - column in aggregate-like expression") { + // greatest(col("src.value"), col("dst.value"), col("edge.weight")) + val refs = GraphFrameInternals.extractColumnReferences( + spark, + greatest(col("src.value"), col("dst.value"), col("edge.weight"))) + assert(refs.contains("src")) + assert(refs.contains("dst")) + assert(refs.contains("edge")) + } + + // ============================================================================ + // Deeply Nested Struct Access + // ============================================================================ + + test("extractColumnReferences - deeply nested struct via bracket notation") { + // col("dst")("location")("city") - three levels deep + // Should still detect "dst" as a prefix + val refs = GraphFrameInternals.extractColumnReferences(spark, col("dst")("location")("city")) + assert(refs.contains("dst"), "Should detect dst prefix even for deeply nested access") + } + + test("extractColumnReferences - two-level nesting via dot notation") { + // col("dst.location.city") - parsed as UnresolvedAttribute(Seq("dst", "location", "city")) + val refs = GraphFrameInternals.extractColumnReferences(spark, col("dst.location.city")) + assert(refs.contains("dst")) + // We should get "location" as the first-level field (we only parse one level deep) + assert(refs("dst").contains("location")) + } + + // ============================================================================ + // Edge Cases - No Column References + // ============================================================================ + + test("extractColumnReferences - literal only expression") { + val refs = GraphFrameInternals.extractColumnReferences(spark, lit(42)) + assert(refs.isEmpty) + } + + test("extractColumnReferences - null literal") { + val refs = GraphFrameInternals.extractColumnReferences(spark, lit(null)) + assert(refs.isEmpty) + } + + test("extractColumnReferences - complex math with no columns") { + val refs = GraphFrameInternals.extractColumnReferences(spark, lit(1) + lit(2) * lit(3)) + assert(refs.isEmpty) + } + + test("extractColumnReferences - string literal") { + val refs = GraphFrameInternals.extractColumnReferences(spark, lit("hello")) + assert(refs.isEmpty) + } + + // ============================================================================ + // Mixed Expressions - Columns and Literals + // ============================================================================ + + test("extractColumnReferences - column plus literal") { + val refs = GraphFrameInternals.extractColumnReferences(spark, col("src.value") + lit(10)) + assert(refs === Map("src" -> Set("value"))) + } + + test("extractColumnReferences - conditional with literal fallback") { + val refs = GraphFrameInternals.extractColumnReferences( + spark, + when(col("dst.flag"), col("src.value")).otherwise(lit(0))) + assert(refs.contains("dst")) + assert(refs.contains("src")) + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameSuite.scala index fb3cb63a64d59..159bbf88c6304 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameSuite.scala @@ -17,99 +17,760 @@ package org.apache.spark.graphframes -import org.apache.spark.sql.{QueryTest, Row} -import org.apache.spark.sql.functions.col -import org.apache.spark.sql.test.SharedSparkSession +import org.apache.commons.io.FileUtils +import org.apache.hadoop.fs.Path +import org.apache.spark.graphx.Edge +import org.apache.spark.graphx.Graph +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.types.LongType +import org.apache.spark.sql.types.StringType +import org.apache.spark.sql.types.StructField +import org.apache.spark.sql.types.StructType +import org.apache.spark.storage.StorageLevel +import org.apache.spark.graphframes.examples.Graphs -class GraphFrameSuite extends QueryTest with SharedSparkSession { +import java.io.File +import java.nio.file.Files - import testImplicits._ +class GraphFrameSuite extends SparkFunSuite with GraphFrameTestSparkContext { - private def graph: GraphFrame = { - val vertices = Seq((1L, "a"), (2L, "b"), (3L, "c"), (4L, "isolated")) - .toDF("id", "name") - val edges = Seq((1L, 2L, "friend"), (2L, 3L, "follow"), (2L, 1L, "friend")) - .toDF("src", "dst", "relationship") - GraphFrame(vertices, edges) + import GraphFrame._ + + var vertices: DataFrame = _ + val localVertices: Map[Long, String] = Map(1L -> "A", 2L -> "B", 3L -> "C") + val localEdges: Map[(Long, Long), String] = + Map((1L, 2L) -> "love", (2L, 1L) -> "hate", (2L, 3L) -> "follow") + var edges: DataFrame = _ + var tempDir: File = _ + + override def beforeAll(): Unit = { + super.beforeAll() + tempDir = Files.createTempDirectory(null).toFile() + vertices = spark.createDataFrame(localVertices.toSeq).toDF("id", "name") + edges = spark + .createDataFrame(localEdges.toSeq.map { case ((src, dst), action) => + (src, dst, action) + }) + .toDF("src", "dst", "action") } - test("construction validates required columns") { - val vertices = Seq((1L, "a")).toDF("id", "name") - val edges = Seq((1L, 1L)).toDF("src", "dst") - GraphFrame(vertices, edges) + override def afterAll(): Unit = { + FileUtils.deleteQuietly(tempDir) + super.afterAll() + } + + test("test validate") { + val goodG = GraphFrame( + spark.createDataFrame(Seq((1L, "a"), (2L, "b"), (3L, "c"))).toDF("id", "attr"), + spark.createDataFrame(Seq((1L, 2L), (2L, 1L), (2L, 3L))).toDF("src", "dst")) + goodG.validate() // no exception should be thrown + + val notDistinctVertices = GraphFrame( + spark.createDataFrame(Seq((1L, "a"), (2L, "b"), (3L, "c"), (1L, "d"))).toDF("id", "attr"), + spark.createDataFrame(Seq((1L, 2L), (2L, 1L), (2L, 3L))).toDF("src", "dst")) + assertThrows[InvalidGraphException](notDistinctVertices.validate()) - val missingId = intercept[IllegalArgumentException] { - GraphFrame(vertices.withColumnRenamed("id", "vertex"), edges) + val missingVertices = GraphFrame( + spark.createDataFrame(Seq((1L, "a"), (2L, "b"), (3L, "c"))).toDF("id", "attr"), + spark.createDataFrame(Seq((1L, 2L), (2L, 1L), (2L, 3L), (1L, 4L))).toDF("src", "dst")) + assertThrows[InvalidGraphException](missingVertices.validate()) + } + + test("construction from DataFrames") { + val g = GraphFrame(vertices, edges) + g.vertices.collect().foreach { + case Row(id: Long, name: String) => + assert(localVertices(id) === name) + case _: Row => throw new GraphFramesUnreachableException() + } + g.edges.collect().foreach { + case Row(src: Long, dst: Long, action: String) => + assert(localEdges((src, dst)) === action) + case _: Row => throw new GraphFramesUnreachableException() + } + intercept[IllegalArgumentException] { + val badVertices = vertices.select(col("id").as("uid"), col("name")) + GraphFrame(badVertices, edges) } - assert(missingId.getMessage.contains("Vertex ID column 'id' is missing")) + intercept[IllegalArgumentException] { + val badEdges = edges.select(col("src").as("srcId"), col("dst"), col("action")) + GraphFrame(vertices, badEdges) + } + intercept[IllegalArgumentException] { + val badEdges = edges.select(col("src"), col("dst").as("dstId"), col("action")) + GraphFrame(vertices, badEdges) + } + } - val missingDestination = intercept[IllegalArgumentException] { - GraphFrame(vertices, edges.withColumnRenamed("dst", "destination")) + test("construction from DataFrames with dots in column names") { + val g = GraphFrame( + vertices.withColumnRenamed("name", "a.name"), + edges.withColumnRenamed("action", "the.action")) + g.vertices.collect().foreach { + case Row(id: Long, name: String) => + assert(localVertices(id) === name) + case _: Row => throw new GraphFramesUnreachableException() + } + g.edges.collect().foreach { + case Row(src: Long, dst: Long, action: String) => + assert(localEdges((src, dst)) === action) + case _: Row => throw new GraphFramesUnreachableException() } - assert(missingDestination.getMessage.contains("Destination vertex ID column 'dst' is missing")) + g.pageRank.maxIter(10).run() } - test("degree DataFrames preserve GraphFrames schemas") { - checkAnswer(graph.outDegrees, Seq(Row(1L, 1), Row(2L, 2))) - checkAnswer(graph.inDegrees, Seq(Row(1L, 1), Row(2L, 1), Row(3L, 1))) - checkAnswer(graph.degrees, Seq(Row(1L, 2), Row(2L, 3), Row(3L, 1))) + test("construction from DataFrames with backquote in column names") { + val g = GraphFrame( + vertices.withColumnRenamed("name", "a `name`"), + edges.withColumnRenamed("action", "the `action`")) + g.vertices.collect().foreach { + case Row(id: Long, name: String) => + assert(localVertices(id) === name) + case _: Row => throw new GraphFramesUnreachableException() + } + g.edges.collect().foreach { + case Row(src: Long, dst: Long, action: String) => + assert(localEdges((src, dst)) === action) + case _: Row => throw new GraphFramesUnreachableException() + } + g.pageRank.maxIter(10).run() } - test("triplets contain complete source, edge, and destination rows") { - val result = graph.triplets.select( - col("src.id"), - col("src.name"), - col("edge.relationship"), - col("dst.id"), - col("dst.name")) + test("construction from edge DataFrame") { + val g = GraphFrame.fromEdges(edges) + assert(g.vertices.columns === Array("id")) + val idsFromVertices = g.vertices.select("id").rdd.map(_.getLong(0)).collect() + val idsFromVerticesSet = idsFromVertices.toSet + assert(idsFromVertices.length === idsFromVerticesSet.size) + val idsFromEdgesSet = g.edges + .select("src", "dst") + .rdd + .flatMap { + case Row(src: Long, dst: Long) => + Seq(src, dst) + case _: Row => throw new GraphFramesUnreachableException() + } + .collect() + .toSet + assert(idsFromVerticesSet === idsFromEdgesSet) + g.vertices.unpersist() + } - checkAnswer( - result, - Seq( - Row(1L, "a", "friend", 2L, "b"), - Row(2L, "b", "follow", 3L, "c"), - Row(2L, "b", "friend", 1L, "a"))) - } - - test("relational graph transforms retain attributes") { - val filtered = graph.filterVertices(col("id") <= 2L) - checkAnswer(filtered.vertices, Seq(Row(1L, "a"), Row(2L, "b"))) - checkAnswer( - filtered.edges, - Seq(Row(1L, 2L, "friend"), Row(2L, 1L, "friend"))) - - checkAnswer( - graph.filterEdges(col("relationship") === "follow").edges, - Seq(Row(2L, 3L, "follow"))) - checkAnswer(graph.dropIsolatedVertices().vertices.select("id"), Seq(Row(1L), Row(2L), Row(3L))) - checkAnswer( - graph.reverse.edges, - Seq( - Row(2L, 1L, "friend"), - Row(3L, 2L, "follow"), - Row(1L, 2L, "friend"))) + test("construction from edges DataFrame, string IDs") { + val edges = spark + .createDataFrame(Seq(("a", "b", "love"), ("b", "a", "hate"), ("b", "c", "follow"))) + .toDF("src", "dst", "action") + + val g = GraphFrame.fromEdges(edges) + assert(g.vertices.columns === Array("id")) + val idsFromVertices = g.vertices.select("id").rdd.map(_.getString(0)).collect() + val idsFromVerticesSet = idsFromVertices.toSet + assert(idsFromVertices.length === idsFromVerticesSet.size) + val idsFromEdgesSet = g.edges + .select("src", "dst") + .rdd + .flatMap { + case Row(src: String, dst: String) => + Seq(src, dst) + case _: Row => throw new GraphFramesUnreachableException() + } + .collect() + .toSet + assert(idsFromVerticesSet === idsFromEdgesSet) + g.vertices.unpersist() + } + + test("construction from edges DataFrame, different storage level") { + val edges = spark + .createDataFrame(Seq((1L, 2L, "love"), (2L, 1L, "hate"), (2L, 3L, "follow"))) + .toDF("src", "dst", "action") + + var g = GraphFrame.fromEdges(edges) + assert(g.vertices.storageLevel === StorageLevel.MEMORY_AND_DISK) + g.vertices.unpersist() + + g = GraphFrame.fromEdges(edges, StorageLevel.MEMORY_AND_DISK_SER) + assert(g.vertices.storageLevel === StorageLevel.MEMORY_AND_DISK_SER) + g.vertices.unpersist() } - test("validate rejects duplicate vertices and unknown edge endpoints") { - graph.validate() + test("construction from GraphX") { + val vv: RDD[(Long, String)] = vertices.rdd.map { + case Row(id: Long, name: String) => + (id, name) + case _: Row => throw new GraphFramesUnreachableException() + } + val ee: RDD[Edge[String]] = edges.rdd.map { + case Row(src: Long, dst: Long, action: String) => + Edge(src, dst, action) + case _: Row => throw new GraphFramesUnreachableException() + } + val g = Graph[String, String](vv, ee) + val gf = GraphFrame.fromGraphX(g) + gf.vertices.select("id", "attr").collect().foreach { + case Row(id: Long, name: String) => + assert(localVertices(id) === name) + case _: Row => throw new GraphFramesUnreachableException() + } + gf.edges.select("src", "dst", "attr").collect().foreach { + case Row(src: Long, dst: Long, action: String) => + assert(localEdges((src, dst)) === action) + case _: Row => throw new GraphFramesUnreachableException() + } + } - val duplicateVertices = Seq((1L, "a"), (1L, "duplicate")).toDF("id", "name") - val noEdges = Seq.empty[(Long, Long)].toDF("src", "dst") - assertThrows[InvalidGraphException](GraphFrame(duplicateVertices, noEdges).validate()) + test("convert to GraphX: Long IDs") { + val gf = GraphFrame(vertices, edges) + val g = gf.toGraphX + g.vertices.collect().foreach { + case (id0, Row(id1: Long, name: String)) => + assert(id0 === id1) + assert(localVertices(id0) === name) + case _ => throw new GraphFramesUnreachableException() + } + g.edges.collect().foreach { + case Edge(src0, dst0, Row(src1: Long, dst1: Long, action: String)) => + assert(src0 === src1) + assert(dst0 === dst1) + assert(localEdges((src0, dst0)) === action) + case _ => throw new GraphFramesUnreachableException() + } + } - val vertices = Seq(1L).toDF("id") - val unknownEndpoint = Seq((1L, 2L)).toDF("src", "dst") - assertThrows[InvalidGraphException](GraphFrame(vertices, unknownEndpoint).validate()) + test("convert to GraphX: Int IDs") { + val vv = vertices.select(col("id").cast(IntegerType).as("id"), col("name")) + val ee = edges.select( + col("src").cast(IntegerType).as("src"), + col("dst").cast(IntegerType).as("dst"), + col("action")) + val gf = GraphFrame(vv, ee) + val g = gf.toGraphX + // Int IDs should be directly cast to Long, so ID values should match. + val vCols = gf.vertexColumnMap + val eCols = gf.edgeColumnMap + g.vertices.collect().foreach { + case (id0: Long, attr: Row) => + val id1 = attr.getInt(vCols("id")) + val name = attr.getString(vCols("name")) + assert(id0 === id1) + assert(localVertices(id0) === name) + case _ => throw new GraphFramesUnreachableException() + } + g.edges.collect().foreach { + case Edge(src0: Long, dst0: Long, attr: Row) => + val src1 = attr.getInt(eCols("src")) + val dst1 = attr.getInt(eCols("dst")) + val action = attr.getString(eCols("action")) + assert(src0 === src1) + assert(dst0 === dst1) + assert(localEdges((src0, dst0)) === action) + case _ => throw new GraphFramesUnreachableException() + } } - test("fromEdges derives distinct vertices") { - val edges = Seq((1L, 2L), (2L, 3L), (1L, 2L)).toDF("src", "dst") - val derived = GraphFrame.fromEdges(edges) + test("convert to GraphX: String IDs") { try { - checkAnswer(derived.vertices, Seq(Row(1L), Row(2L), Row(3L))) - checkAnswer(derived.edges, edges.collect().toSeq) - } finally { - derived.vertices.unpersist() + val vv = vertices.select(col("id").cast(StringType).as("id"), col("name")) + val ee = edges.select( + col("src").cast(StringType).as("src"), + col("dst").cast(StringType).as("dst"), + col("action")) + val gf = GraphFrame(vv, ee) + val g = gf.toGraphX + // String IDs will be re-indexed, so ID values may not match. + val vCols = gf.vertexColumnMap + val eCols = gf.edgeColumnMap + // First, get index. + val new2oldID: Map[Long, String] = g.vertices + .map { case (id: Long, attr: Row) => + (id, attr.getString(vCols("id"))) + } + .collect() + .toMap + // Same as in test with Int IDs, but with re-indexing + g.vertices.collect().foreach { case (id0: Long, attr: Row) => + val id1 = attr.getString(vCols("id")) + val name = attr.getString(vCols("name")) + assert(new2oldID(id0) === id1) + assert(localVertices(new2oldID(id0).toLong) === name) + } + g.edges.collect().foreach { case Edge(src0: Long, dst0: Long, attr: Row) => + val src1 = attr.getString(eCols("src")) + val dst1 = attr.getString(eCols("dst")) + val action = attr.getString(eCols("action")) + assert(new2oldID(src0) === src1) + assert(new2oldID(dst0) === dst1) + assert(localEdges((new2oldID(src0).toLong, new2oldID(dst0).toLong)) === action) + } + } catch { + case e: Exception => + e.printStackTrace() + throw e + } + } + + test("save/load") { + val g0 = GraphFrame(vertices, edges) + val vPath = new Path(tempDir.getPath, "vertices").toString + val ePath = new Path(tempDir.getPath, "edges").toString + g0.vertices.write.parquet(vPath) + g0.edges.write.parquet(ePath) + + val v1 = spark.read.parquet(vPath) + val e1 = spark.read.parquet(ePath) + val g1 = GraphFrame(v1, e1) + + g1.vertices.collect().foreach { + case Row(id: Long, name: String) => + assert(localVertices(id) === name) + case _ => throw new GraphFramesUnreachableException() + } + g1.edges.collect().foreach { + case Row(src: Long, dst: Long, action: String) => + assert(localEdges((src, dst)) === action) + case _ => throw new GraphFramesUnreachableException() + } + } + + test("degree metrics") { + val g = GraphFrame(vertices, edges) + + assert(g.outDegrees.columns === Seq("id", "outDegree")) + val outDegrees = g.outDegrees + .collect() + .map { + case Row(id: Long, outDeg: Int) => + (id, outDeg) + case _ => throw new GraphFramesUnreachableException() + } + .toMap + assert(outDegrees === Map(1L -> 1, 2L -> 2)) + + assert(g.inDegrees.columns === Seq("id", "inDegree")) + val inDegrees = g.inDegrees + .collect() + .map { + case Row(id: Long, inDeg: Int) => + (id, inDeg) + case _ => throw new GraphFramesUnreachableException() + } + .toMap + assert(inDegrees === Map(1L -> 1, 2L -> 1, 3L -> 1)) + + assert(g.degrees.columns === Seq("id", "degree")) + val degrees = g.degrees + .collect() + .map { + case Row(id: Long, deg: Int) => + (id, deg) + case _ => throw new GraphFramesUnreachableException() + } + .toMap + assert(degrees === Map(1L -> 2, 2L -> 3, 3L -> 1)) + } + + test("type degree metrics") { + val g = GraphFrame(vertices, edges) + + assert(g.typeOutDegree("action").columns === Seq("id", "outDegrees")) + val typeOutDegrees = g.typeOutDegree("action").collect() + + val outDegreesSchema = + g.typeOutDegree("action").schema("outDegrees").dataType.asInstanceOf[StructType] + val outDegreesFieldNames = outDegreesSchema.fields.map(_.name).toSet + assert(outDegreesFieldNames === Set("love", "hate", "follow")) + + val typeOutDegMap = typeOutDegrees.map { row => + val id = row.getLong(0) + val degrees = row.getStruct(1) + (id, degrees) + }.toMap + + assert(typeOutDegMap(1L).getAs[Int]("love") === 1) + assert(typeOutDegMap(1L).getAs[Int]("hate") === 0) + assert(typeOutDegMap(1L).getAs[Int]("follow") === 0) + + assert(typeOutDegMap(2L).getAs[Int]("love") === 0) + assert(typeOutDegMap(2L).getAs[Int]("hate") === 1) + assert(typeOutDegMap(2L).getAs[Int]("follow") === 1) + + assert(g.typeInDegree("action").columns === Seq("id", "inDegrees")) + val typeInDegrees = g.typeInDegree("action").collect() + + val inDegreesSchema = + g.typeInDegree("action").schema("inDegrees").dataType.asInstanceOf[StructType] + val inDegreesFieldNames = inDegreesSchema.fields.map(_.name).toSet + assert(inDegreesFieldNames === Set("love", "hate", "follow")) + + val typeInDegMap = typeInDegrees.map { row => + val id = row.getLong(0) + val degrees = row.getStruct(1) + (id, degrees) + }.toMap + + assert(typeInDegMap(1L).getAs[Int]("love") === 0) + assert(typeInDegMap(1L).getAs[Int]("hate") === 1) + assert(typeInDegMap(1L).getAs[Int]("follow") === 0) + + assert(typeInDegMap(2L).getAs[Int]("love") === 1) + assert(typeInDegMap(2L).getAs[Int]("hate") === 0) + assert(typeInDegMap(2L).getAs[Int]("follow") === 0) + + assert(typeInDegMap(3L).getAs[Int]("love") === 0) + assert(typeInDegMap(3L).getAs[Int]("hate") === 0) + assert(typeInDegMap(3L).getAs[Int]("follow") === 1) + + assert(g.typeDegree("action").columns === Seq("id", "degrees")) + val typeDegrees = g.typeDegree("action").collect() + + val degreesSchema = g.typeDegree("action").schema("degrees").dataType.asInstanceOf[StructType] + val degreesFieldNames = degreesSchema.fields.map(_.name).toSet + assert(degreesFieldNames === Set("love", "hate", "follow")) + + val typeDegMap = typeDegrees.map { row => + val id = row.getLong(0) + val degrees = row.getStruct(1) + (id, degrees) + }.toMap + + assert(typeDegMap(1L).getAs[Int]("love") === 1) + assert(typeDegMap(1L).getAs[Int]("hate") === 1) + assert(typeDegMap(1L).getAs[Int]("follow") === 0) + + assert(typeDegMap(2L).getAs[Int]("love") === 1) + assert(typeDegMap(2L).getAs[Int]("hate") === 1) + assert(typeDegMap(2L).getAs[Int]("follow") === 1) + + assert(typeDegMap(3L).getAs[Int]("love") === 0) + assert(typeDegMap(3L).getAs[Int]("hate") === 0) + assert(typeDegMap(3L).getAs[Int]("follow") === 1) + } + + test("type degree metrics with explicit edge types") { + val g = GraphFrame(vertices, edges) + val edgeTypes = Seq("love", "hate", "follow") + + val typeOutDegrees = g.typeOutDegree("action", Some(edgeTypes)).collect() + + val typeOutDegMap = typeOutDegrees.map { row => + val id = row.getLong(0) + val degrees = row.getStruct(1) + (id, degrees) + }.toMap + + assert(typeOutDegMap(1L).getAs[Int]("love") === 1) + assert(typeOutDegMap(1L).getAs[Int]("hate") === 0) + assert(typeOutDegMap(1L).getAs[Int]("follow") === 0) + + assert(typeOutDegMap(2L).getAs[Int]("love") === 0) + assert(typeOutDegMap(2L).getAs[Int]("hate") === 1) + assert(typeOutDegMap(2L).getAs[Int]("follow") === 1) + + val typeInDegrees = g.typeInDegree("action", Some(edgeTypes)).collect() + val typeInDegMap = typeInDegrees.map { row => + val id = row.getLong(0) + val degrees = row.getStruct(1) + (id, degrees) + }.toMap + + assert(typeInDegMap(1L).getAs[Int]("love") === 0) + assert(typeInDegMap(1L).getAs[Int]("hate") === 1) + assert(typeInDegMap(1L).getAs[Int]("follow") === 0) + + assert(typeInDegMap(2L).getAs[Int]("love") === 1) + assert(typeInDegMap(2L).getAs[Int]("hate") === 0) + assert(typeInDegMap(2L).getAs[Int]("follow") === 0) + + assert(typeInDegMap(3L).getAs[Int]("love") === 0) + assert(typeInDegMap(3L).getAs[Int]("hate") === 0) + assert(typeInDegMap(3L).getAs[Int]("follow") === 1) + + val typeDegrees = g.typeDegree("action", Some(edgeTypes)).collect() + val typeDegMap = typeDegrees.map { row => + val id = row.getLong(0) + val degrees = row.getStruct(1) + (id, degrees) + }.toMap + + assert(typeDegMap(1L).getAs[Int]("love") === 1) + assert(typeDegMap(1L).getAs[Int]("hate") === 1) + assert(typeDegMap(1L).getAs[Int]("follow") === 0) + + assert(typeDegMap(2L).getAs[Int]("love") === 1) + assert(typeDegMap(2L).getAs[Int]("hate") === 1) + assert(typeDegMap(2L).getAs[Int]("follow") === 1) + + assert(typeDegMap(3L).getAs[Int]("love") === 0) + assert(typeDegMap(3L).getAs[Int]("hate") === 0) + assert(typeDegMap(3L).getAs[Int]("follow") === 1) + } + + test("cache") { + val g = GraphFrame(vertices, edges) + + g.persist(StorageLevel.MEMORY_ONLY) + + g.unpersist() + // org.apache.spark.sql.execution.columnar.InMemoryRelation is private and not accessible + // This has prevented us from validating DataFrame's are cached. + } + + test("basic operations on an empty graph") { + for (empty <- Seq(Graphs.empty[Int], Graphs.empty[Long], Graphs.empty[String])) { + assert(empty.inDegrees.count() === 0L) + assert(empty.outDegrees.count() === 0L) + assert(empty.degrees.count() === 0L) + assert(empty.triplets.count() === 0L) + } + } + + test("test triplets") { + // Basic triplets test + val g = GraphFrame(vertices, edges) + val triplets = g.triplets.collect() + assert(triplets.length === localEdges.size) + triplets.foreach { + case Row(src: Row, edge: Row, dst: Row) => + assert(src.getLong(0) === edge.getLong(0)) // src.id === edge.src + assert(dst.getLong(0) === edge.getLong(1)) // dst.id === edge.dst + assert(localEdges((edge.getLong(0), edge.getLong(1))) === edge.getString(2)) + case _ => throw new GraphFramesUnreachableException() + } + + // Test with attributes + val v2 = vertices.withColumn("age", lit(10)) + val e2 = edges.withColumn("weight", lit(2.0)) + val g2 = GraphFrame(v2, e2) + val triplets2 = g2.triplets.collect() + triplets2.foreach { + case Row(src: Row, edge: Row, _: Row) => + assert(src.getInt(2) === 10) // Check vertex attribute + assert(edge.getDouble(3) === 2.0) // Check edge attribute + case _ => throw new GraphFramesUnreachableException() + } + + // Test with dots in column names + val v3 = v2.withColumnRenamed("age", "person.age") + val e3 = e2.withColumnRenamed("weight", "edge.weight") + val g3 = GraphFrame(v3, e3) + val triplets3 = g3.triplets.collect() + triplets3.foreach { + case Row(src: Row, edge: Row, _: Row) => + assert(src.getInt(2) === 10) // Check vertex attribute + assert(edge.getDouble(3) === 2.0) // Check edge attribute + case _ => throw new GraphFramesUnreachableException() + } + } + + test("nestAsCol with dots in column names") { + val df = vertices.withColumnRenamed("name", "a.name") + val col = nestAsCol(df, "attr") + assert( + df.select(col).schema === StructType( + Seq( + StructField( + "attr", + StructType(Seq( + StructField("id", LongType, nullable = false), + StructField("a.name", StringType, nullable = true))), + nullable = false)))) + } + + test("nestAsCol with backquote in column names") { + val df = vertices.withColumnRenamed("name", "a `name`") + val col = nestAsCol(df, "attr") + assert( + df.select(col).schema === StructType( + Seq( + StructField( + "attr", + StructType(Seq( + StructField("id", LongType, nullable = false), + StructField("a `name`", StringType, nullable = true))), + nullable = false)))) + } + + test("power iteration clustering wrapper") { + val spark = this.spark + import spark.implicits._ + val edges = spark + .createDataFrame( + Seq( + (1, 0, 0.5), + (2, 0, 0.5), + (2, 1, 0.7), + (3, 0, 0.5), + (3, 1, 0.7), + (3, 2, 0.9), + (4, 0, 0.5), + (4, 1, 0.7), + (4, 2, 0.9), + (4, 3, 1.1), + (5, 0, 0.5), + (5, 1, 0.7), + (5, 2, 0.9), + (5, 3, 1.1), + (5, 4, 1.3))) + .toDF("src", "dst", "weight") + val vertices = Seq(0, 1, 2, 3, 4, 5).toDF("id") + val gf = GraphFrame(vertices, edges) + val clusters = gf + .powerIterationClustering(k = 2, maxIter = 40, weightCol = Some("weight")) + .collect() + .sortBy(_.getAs[Long]("id")) + .map(_.getAs[Int]("cluster")) + .toSeq + assert(Seq(0, 0, 0, 0, 1, 0) == clusters) + } + + test("power iteration clustering string ids") { + val spark = this.spark + import spark.implicits._ + val edges = spark + .createDataFrame( + Seq( + ("1", "0", 0.5), + ("2", "0", 0.5), + ("2", "1", 0.7), + ("3", "0", 0.5), + ("3", "1", 0.7), + ("3", "2", 0.9), + ("4", "0", 0.5), + ("4", "1", 0.7), + ("4", "2", 0.9), + ("4", "3", 1.1), + ("5", "0", 0.5), + ("5", "1", 0.7), + ("5", "2", 0.9), + ("5", "3", 1.1), + ("5", "4", 1.3))) + .toDF("src", "dst", "weight") + val vertices = Seq("0", "1", "2", "3", "4", "5").toDF("id") + val gf = GraphFrame(vertices, edges) + val clusters = gf + .powerIterationClustering(k = 2, maxIter = 40, weightCol = Some("weight")) + .collect() + .sortBy(_.getAs[String]("id")) + .map(_.getAs[Int]("cluster")) + .toSeq + assert(Seq(1, 1, 1, 1, 1, 0) == clusters) + } + + test("convert directed graph to undirected") { + val v = spark.createDataFrame(Seq((1L, "a"), (2L, "b"), (3L, "c"))).toDF("id", "name") + val e = spark.createDataFrame(Seq((1L, 2L), (2L, 3L))).toDF("src", "dst") + val g = GraphFrame(v, e) + val undirected = g.asUndirected() + + // Check edge count doubled + assert(undirected.edges.count() === 2 * g.edges.count()) + + // Verify reverse edges exist + val edges = undirected.edges.sort("src", "dst").collect() + assert(edges.length === 4) + assert(edges(0).getLong(0) === 1L) + assert(edges(0).getLong(1) === 2L) + assert(edges(1).getLong(0) === 2L) + assert(edges(1).getLong(1) === 1L) + assert(edges(2).getLong(0) === 2L) + assert(edges(2).getLong(1) === 3L) + assert(edges(3).getLong(0) === 3L) + assert(edges(3).getLong(1) === 2L) + } + + test("reverse directed graph edges") { + val v = spark.createDataFrame(Seq((1L, "a"), (2L, "b"), (3L, "c"))).toDF("id", "name") + val e = spark.createDataFrame(Seq((1L, 2L), (2L, 3L))).toDF("src", "dst") + val g = GraphFrame(v, e) + val reversed = g.asReversed() + + // Check edge count is the same + assert(reversed.edges.count() === g.edges.count()) + + // Verify edges are reversed + val edges = reversed.edges.sort("src", "dst").collect() + assert(edges.length === 2) + assert(edges(0).getLong(0) === 2L) + assert(edges(0).getLong(1) === 1L) + assert(edges(1).getLong(0) === 3L) + assert(edges(1).getLong(1) === 2L) + } + + test("reverse directed graph edges with attributes") { + val v = spark.createDataFrame(Seq((1L, "a"), (2L, "b"))).toDF("id", "name") + val e = spark.createDataFrame(Seq((1L, 2L, "edge1"))).toDF("src", "dst", "attr") + val g = GraphFrame(v, e) + val reversed = g.asReversed() + + val edges = reversed.edges.collect() + assert(edges.length === 1) + assert(edges(0).getLong(0) === 2L) + assert(edges(0).getLong(1) === 1L) + assert(edges(0).getString(2) === "edge1") + } + + test("toGraphX should throw IllegalArgumentException for null IDs") { + val schema = StructType( + Seq( + StructField("id", LongType, nullable = true), + StructField("attr", StringType, nullable = true))) + + val data = spark.sparkContext.parallelize(Seq(Row(1L, "a"), Row(null, "b"))) + + val vertices = spark.createDataFrame(data, schema) + val edges = spark.createDataFrame(Seq((1L, 1L, "friend"))).toDF("src", "dst", "relationship") + + val g = GraphFrame(vertices, edges) + + val e = intercept[org.apache.spark.SparkException] { + // Trigger an action to hit the lazy exception in the .map + g.toGraphX.vertices.collect() + } + + assert(e.getCause.isInstanceOf[IllegalArgumentException]) + assert(e.getMessage.contains("Vertex ID cannot be null")) + } + + test("toGraphX should throw IllegalArgumentException for null Edge Src/Dst") { + val vertices = spark.createDataFrame(Seq((1L, "a"))).toDF("id", "attr") + + val edgeSchema = StructType( + Seq( + StructField("src", LongType, nullable = true), + StructField("dst", LongType, nullable = true), + StructField("relationship", StringType, nullable = true))) + + val edgeData = spark.sparkContext.parallelize(Seq(Row(1L, null, "friend"))) + + val edges = spark.createDataFrame(edgeData, edgeSchema) + + val g = GraphFrame(vertices, edges) + + val e = intercept[org.apache.spark.SparkException] { + // Trigger action on edges + g.toGraphX.edges.collect() } + + assert(e.getCause.isInstanceOf[IllegalArgumentException]) + assert(e.getMessage.contains("Edge")) + assert(e.getMessage.contains("cannot be null")) + } + + test("convert directed graph with edge attributes to undirected") { + val v = spark.createDataFrame(Seq((1L, "a"), (2L, "b"))).toDF("id", "name") + val e = spark.createDataFrame(Seq((1L, 2L, "edge1"))).toDF("src", "dst", "attr") + val g = GraphFrame(v, e) + val undirected = g.asUndirected() + + val edges = undirected.edges.collect() + assert(edges.length === 2) + assert( + edges.exists(r => r.getLong(0) == 1L && r.getLong(1) == 2L && r.getString(2) == "edge1")) + assert( + edges.exists(r => r.getLong(0) == 2L && r.getLong(1) == 1L && r.getString(2) == "edge1")) } } diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameTestSparkContext.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameTestSparkContext.scala new file mode 100644 index 0000000000000..06dddd1387707 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameTestSparkContext.scala @@ -0,0 +1,86 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes + +import org.apache.commons.io.FileUtils +import org.apache.spark.SparkContext +import org.apache.spark.sql.SQLContext +import org.apache.spark.sql.SQLImplicits +import org.apache.spark.sql.SparkSession +import org.scalatest.BeforeAndAfterAll +import org.scalatest.Suite + +import java.io.File +import java.nio.file.Files + +trait GraphFrameTestSparkContext extends BeforeAndAfterAll { self: Suite => + @transient var spark: SparkSession = _ + @transient var sc: SparkContext = _ + @transient var sqlContext: SQLContext = _ + @transient var sparkMajorVersion: Int = _ + @transient var sparkMinorVersion: Int = _ + + // Inspired by https://stackoverflow.com/a/59377177 + protected def sparkSession: SparkSession = spark + protected lazy val sqlImplicits: SQLImplicits = self.sparkSession.implicits + + /** Check if current spark version is at least of the provided minimum version */ + def isLaterVersion(minVersion: String): Boolean = { + val (minMajorVersion, minMinorVersion) = TestUtils.majorMinorVersion(minVersion) + if (sparkMajorVersion != minMajorVersion) { + return sparkMajorVersion > minMajorVersion + } else { + return sparkMinorVersion >= minMinorVersion + } + } + + override def beforeAll(): Unit = { + super.beforeAll() + + spark = SparkSession + .builder() + .master("local[2]") + .appName("GraphFramesUnitTest") + .config("spark.sql.shuffle.partitions", 4) + .config("spark.sql.adaptive.enabled", "true") + .getOrCreate() + + val checkpointDir = Files.createTempDirectory(this.getClass.getName).toString + spark.sparkContext.setCheckpointDir(checkpointDir) + sc = spark.sparkContext + sqlContext = spark.sqlContext + + val (verMajor, verMinor) = TestUtils.majorMinorVersion(sc.version) + sparkMajorVersion = verMajor + sparkMinorVersion = verMinor + } + + override def afterAll(): Unit = { + val checkpointDir = sc.getCheckpointDir + if (spark != null) { + spark.stop() + } + spark = null + sc = null + + checkpointDir.foreach { dir => + FileUtils.deleteQuietly(new File(dir)) + } + super.afterAll() + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/PatternMatchSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/PatternMatchSuite.scala new file mode 100644 index 0000000000000..47d2700c3db0b --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/PatternMatchSuite.scala @@ -0,0 +1,875 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes + +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.functions.when + +/** + * Cases to go through: + * - Any negated terms? + * - Any anonymous vertices + * - in non-negated terms? + * - in negated terms? + * - # named vertices grounding a negated term to non-negated terms: 2, 1, 0 + * - Named edges? + */ +class PatternMatchSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + @transient var v: DataFrame = _ + @transient var e: DataFrame = _ + @transient var noEdges: DataFrame = _ + @transient var g: GraphFrame = _ + + override def beforeAll(): Unit = { + super.beforeAll() + + v = spark + .createDataFrame(List((0L, "a", "f"), (1L, "b", "m"), (2L, "c", "m"), (3L, "d", "f"))) + .toDF("id", "attr", "gender") + e = spark + .createDataFrame( + List( + (0L, 1L, "friend"), + (1L, 0L, "follow"), + (1L, 2L, "friend"), + (2L, 3L, "follow"), + (2L, 0L, "unknown"))) + .toDF("src", "dst", "relationship") + noEdges = v + .select(col("id").alias("src")) + .crossJoin(v.select(col("id").alias("dst"))) + .except(e.select("src", "dst")) + g = GraphFrame(v, e) + } + + override def afterAll(): Unit = { + v = null + e = null + g = null + super.afterAll() + } + + private def compareResultToExpected[A](result: Set[A], expected: Set[A]): Unit = { + if (result !== expected) { + throw new AssertionError( + "result !== expected.\n" + + s"Result contained additional values: ${result.diff(expected)}\n" + + s"Expected contained additional values: ${expected.diff(result)}\n" + + s"Result: $result\n" + + s"Expected: $expected") + } + } + + test("test compareResultToExpected") { + intercept[AssertionError] { + compareResultToExpected(Set(1, 2), Set(2, 3)) + } + } + + test("empty query should return nothing") { + val emptiness = g.find("") + assert(emptiness.count() === 0) + } + + test("filter edges and drop isolated vertices") { + // string expression + val s = "relationship = 'friend'" + // column expression + val c = col("relationship") === "friend" + // expected subgraph vertices + val expected_v = Set(Row(0L, "a", "f"), Row(1L, "b", "m"), Row(2L, "c", "m")) + // expected subgraph edges + val expected_e = Set(Row(0L, 1L, "friend"), Row(1L, 2L, "friend")) + + val res_s = g.filterEdges(s) + assert(res_s.vertices.collect().toSet === v.collect().toSet) + assert(res_s.edges.collect().toSet === expected_e) + + val res_c = g.filterEdges(c) + assert(res_c.vertices.collect().toSet === v.collect().toSet) + assert(res_c.edges.collect().toSet === expected_e) + + val res = res_s.dropIsolatedVertices() + assert(res.vertices.collect().toSet === expected_v) + assert(res.edges.collect().toSet === expected_e) + } + + test("filter vertices") { + // string expression + val s = "id > 0" + // column expression + val c = col("id") > 0 + // expected subgraph vertices + val expected_v = Set(Row(1L, "b", "m"), Row(2L, "c", "m"), Row(3L, "d", "f")) + // expected subgraph edges + val expected_e = Set(Row(1L, 2L, "friend"), Row(2L, 3L, "follow")) + + val res_s = g.filterVertices(s) + assert(res_s.vertices.collect().toSet === expected_v) + assert(res_s.edges.collect().toSet === expected_e) + + val res_c = g.filterVertices(c) + assert(res_c.vertices.collect().toSet === expected_v) + assert(res_c.edges.collect().toSet === expected_e) + } + + test("triangles") { + val triangles = g + .find("(a)-[]->(b); (b)-[]->(c); (c)-[]->(a)") + .select("a.id", "b.id", "c.id") + + assert(triangles.collect().toSet === Set(Row(0L, 1L, 2L), Row(2L, 0L, 1L), Row(1L, 2L, 0L))) + } + + /* ====================================== Vertex queries ===================================== */ + + test("single named vertex") { + val vertices = g.find("(a)") + + assert(vertices.columns === Array("a")) + val res = vertices.select("a.id", "a.attr").collect().toSet + compareResultToExpected(res, v.select("id", "attr").collect().toSet) + } + + /* =========================== Single-edge queries without negated terms ===================== */ + + test("triplet with anonymous edge") { + val triplets = g.find("(u)-[]->(v)") + + assert(triplets.columns === Array("u", "v")) + val res = triplets.select("u.id", "u.attr", "v.id", "v.attr").collect().toSet + compareResultToExpected( + res, + Set( + Row(0L, "a", 1L, "b"), + Row(1L, "b", 0L, "a"), + Row(1L, "b", 2L, "c"), + Row(2L, "c", 3L, "d"), + Row(2L, "c", 0L, "a"))) + } + + test("triplet with named edge") { + val triplets = g.find("(u)-[uv]->(v)") + + assert(triplets.columns === Array("u", "uv", "v")) + val res = triplets + .select("u.id", "u.attr", "uv.src", "uv.dst", "uv.relationship", "v.id", "v.attr") + .collect() + .toSet + compareResultToExpected( + res, + Set( + Row(0L, "a", 0L, 1L, "friend", 1L, "b"), + Row(1L, "b", 1L, 0L, "follow", 0L, "a"), + Row(1L, "b", 1L, 2L, "friend", 2L, "c"), + Row(2L, "c", 2L, 3L, "follow", 3L, "d"), + Row(2L, "c", 2L, 0L, "unknown", 0L, "a"))) + } + + test("triplet with anonymous vertex") { + val triplets = g.find("(u)-[]->()") + + assert(triplets.columns === Array("u")) + // Do not use compareResultToExpected since it uses sets, and we expect duplicates. + assert( + triplets.select("u.id", "u.attr").collect().sortBy(_.getLong(0)) === Array( + Row(0L, "a"), + Row(1L, "b"), + Row(1L, "b"), + Row(2L, "c"), + Row(2L, "c")).sortBy(_.getLong(0))) + } + + test("triplet with 2 anonymous vertices") { + val triplets = g.find("()-[uv]->()") + + assert(triplets.columns === Array("uv")) + val res = triplets.select("uv.src", "uv.dst", "uv.relationship").collect().toSet + val expected = e.select("src", "dst", "relationship").collect().toSet + compareResultToExpected(res, expected) + } + + test("self-loop") { + val myE = spark + .createDataFrame(List((1L, 1L, "self"), (3L, 3L, "self"))) + .toDF("src", "dst", "relationship") + .union(e) + val myG = GraphFrame(v, myE) + + val selfLoops = myG.find("(a)-[]->(a)") + assert(selfLoops.columns === Array("a")) + val res = selfLoops.select("a.id").collect().toSet + compareResultToExpected(res, Set(Row(1L), Row(3L))) + + val selfLoops2 = myG.find("(a)-[]->(b); (a)-[]->(a)") + assert(selfLoops2.columns === Array("a", "b")) + val res2 = selfLoops2 + .select("a.id", "b.id") + .where("a.id != b.id") + .collect() + .toSet + compareResultToExpected(res2, Set(Row(1L, 0L), Row(1L, 2L))) + } + + test("duplicate edges") { + val myE = spark + .createDataFrame(List((1L, 0L, "dup"), (1L, 2L, "dup"))) + .toDF("src", "dst", "relationship") + .union(e) + val myG = GraphFrame(v, myE) + + val edges = myG + .find("(a)-[]->(b)") + .where("a.id = 1") + val res = edges + .select("a.id", "b.id") + .collect() + .sortBy(_.getLong(1)) + val expected = Array(Row(1L, 0L), Row(1L, 0L), Row(1L, 2L), Row(1L, 2L)) + assert(res === expected) + } + + /* ======================== Multiple-edge queries without negated terms ===================== */ + + test("triangle cycles") { + val triangles = g.find("(a)-[]->(b); (b)-[]->(c); (c)-[]->(a)") + + assert(triangles.columns === Array("a", "b", "c")) + val res = triangles + .select("a.id", "b.id", "c.id") + .collect() + .toSet + compareResultToExpected(res, Set(Row(0L, 1L, 2L), Row(2L, 0L, 1L), Row(1L, 2L, 0L))) + } + + test("disconnected edges create an outer join") { + val edgePairs = g.find("(a)-[]->(b); (c)-[]->(d)") + + assert(edgePairs.columns === Array("a", "b", "c", "d")) + val res = edgePairs + .select("a.id", "b.id", "c.id", "d.id") + .collect() + .toSet + + val ab = e.select(col("src").alias("a"), col("dst").alias("b")) + val cd = e.select(col("src").alias("c"), col("dst").alias("d")) + val expected = ab + .crossJoin(cd) + .collect() + .toSet + compareResultToExpected(res, expected) + val numEdges = e.count() + assert(expected.size === numEdges * numEdges) + } + + /* ========== 2 named vertices grounding a negated term to non-negated terms =============== */ + + test("edges without back edges") { + val edges = g.find("(a)-[]->(b); !(b)-[]->(a)") + + assert(edges.columns === Array("a", "b")) + val res = edges + .select("a.id", "b.id") + .collect() + .toSet + compareResultToExpected(res, Set(Row(1L, 2L), Row(2L, 0L), Row(2L, 3L))) + } + + test("a->b->c but not c->a") { + val edges = g.find("(a)-[]->(b); (b)-[]->(c); !(c)-[]->(a)") + + assert(edges.columns === Array("a", "b", "c")) + val res = edges + .select("a.id", "b.id", "c.id") + .collect() + .toSet + assert(res === Set(Row(0L, 1L, 0L), Row(1L, 0L, 1L), Row(1L, 2L, 3L))) + } + + test("three connected vertices not in a triangle") { + val fof = g + .find("(u)-[]->(v); (v)-[]->(w); !(u)-[]->(w); !(w)-[]->(u)") + .select("u.id", "v.id", "w.id") + .collect() + .toSet + + compareResultToExpected(fof, Set(Row(1L, 0L, 1L), Row(0L, 1L, 0L), Row(1L, 2L, 3L))) + } + + /* ========== 1 named vertex grounding a negated term to non-negated terms =============== */ + + test("a->b but not b->c") { + val edges = g.find("(a)-[]->(b); !(b)-[]->(c)") + + assert(edges.columns === Array("a", "b", "c")) + val res = edges + .select("a.id", "b.id", "c.id") + .collect() + .toSet + compareResultToExpected( + res, + Set( + Row(0L, 1L, 1L), + Row(0L, 1L, 3L), + Row(1L, 0L, 0L), + Row(1L, 0L, 2L), + Row(1L, 0L, 3L), + Row(1L, 2L, 1L), + Row(1L, 2L, 2L), + Row(2L, 3L, 0L), + Row(2L, 3L, 1L), + Row(2L, 3L, 2L), + Row(2L, 3L, 3L), + Row(2L, 0L, 0L), + Row(2L, 0L, 2L), + Row(2L, 0L, 3L))) + } + + test("a->b where b has no out edges") { + val edges = g.find("(a)-[]->(b); !(b)-[]->()") + + assert(edges.columns === Array("a", "b")) + val res = edges + .select("a.id", "b.id") + .collect() + .toSet + compareResultToExpected(res, Set(Row(2L, 3L))) + } + + /* ========== 0 named vertices grounding a negated term to non-negated terms =============== */ + + test("a->b but not c->d") { + val edgePairs = g.find("(a)-[]->(b); !(c)-[]->(d)") + + assert(edgePairs.columns === Array("a", "b", "c", "d")) + val res = edgePairs + .select("a.id", "b.id", "c.id", "d.id") + .collect() + .toSet + val expected = e + .select(col("src").alias("a"), col("dst").alias("b")) + .crossJoin(noEdges.select(col("src").alias("c"), col("dst").alias("d"))) + .select("a", "b", "c", "d") + .collect() + .toSet + compareResultToExpected(res, expected) + assert(expected.size === noEdges.count() * e.count()) // make sure there are no duplicates + } + + test("a->b, c where c has no out edges") { + val triplets = g.find("(a)-[]->(b); !(c)-[]->()") + + assert(triplets.columns === Array("a", "b", "c")) + val res = triplets + .select("a.id", "b.id", "c.id") + .collect() + .toSet + val expected = + Set(Row(0L, 1L, 3L), Row(1L, 0L, 3L), Row(1L, 2L, 3L), Row(2L, 3L, 3L), Row(2L, 0L, 3L)) + compareResultToExpected(res, expected) + } + + /* ======= Varying # of named vertices grounding a negated term to non-negated terms ========= */ + + // Note: This is a deceptive query. + // Users may intend "there exists no such c connecting b->c->a," which is a different query. + test("a->b, c without edges b->c->a") { + val seq = g.find("(a)-[]->(b); !(b)-[]->(c); !(c)-[]->(a)") + + assert(seq.columns === Array("a", "b", "c")) + val res = seq + .select("a.id", "b.id", "c.id") + .collect() + .toSet + val expected = Set( + Row(0L, 1L, 3L), + Row(1L, 0L, 2L), + Row(1L, 0L, 3L), + Row(1L, 2L, 1L), + Row(1L, 2L, 2L), + Row(2L, 3L, 0L), + Row(2L, 3L, 2L), + Row(2L, 3L, 3L), + Row(2L, 0L, 0L), + Row(2L, 0L, 2L), + Row(2L, 0L, 3L)) + compareResultToExpected(res, expected) + } + + test("a->b, c, d with no edges a->c, c->d") { + val edgePairs = g.find("(a)-[]->(b); !(a)-[]->(c); !(c)-[]->(d)") + + assert(edgePairs.columns === Array("a", "b", "c", "d")) + val res = edgePairs + .select("a.id", "b.id", "c.id", "d.id") + .where("a.id = 0 AND a.id != b.id") // check subset for brevity + .collect() + .toSet + val expected = Set( + Row(0L, 1L, 0L, 0L), + Row(0L, 1L, 0L, 2L), + Row(0L, 1L, 0L, 3L), + Row(0L, 1L, 2L, 1L), + Row(0L, 1L, 2L, 2L), + Row(0L, 1L, 3L, 0L), + Row(0L, 1L, 3L, 1L), + Row(0L, 1L, 3L, 2L), + Row(0L, 1L, 3L, 3L)) + compareResultToExpected(res, expected) + } + + /* ============================== 0 non-negated terms ============================== */ + + test("query without non-negated terms, with one named vertex") { + val res = g + .find("!(v)-[]->()") + .select("v.id") + .collect() + .toSet + compareResultToExpected(res, Set(Row(3L))) + } + + test("query without non-negated terms, with two named vertices") { + val res = g + .find("!(u)-[]->(v)") + .select("u.id", "v.id") + .collect() + .toSet + val expected = noEdges + .select(col("src").alias("u"), col("dst").alias("v")) + .collect() + .toSet + compareResultToExpected(res, expected) + } + + /* ======================== Other corner cases and implementation checks ==================== */ + + test("named edges") { + // edges whose destination leads nowhere + val edges = g + .find("()-[e]->(v); !(v)-[]->()") + .select("e.src", "e.dst") + val res = edges.collect().toSet + compareResultToExpected(res, Set(Row(2L, 3L))) + } + + test("a->b but not a->b") { + val edges = g.find("(a)-[]->(b); !(a)-[]->(b)") + assert(edges.count() === 0) + + val edges2 = g.find("(a)-[ab]->(b); !(a)-[]->(b)") + assert(edges2.count() === 0) + } + + test("named edge __tmp") { + // named edge __tmp should not be removed if there is an anonymous edge + val edges = g.find("()-[__tmp]->(v); (v)-[]->(w)") + assert(edges.columns === Array("__tmp", "v", "w")) + } + + test("find column order") { + val fof = g + .find("(u)-[e]->(v); (v)-[]->(w); !(u)-[]->(w); !(w)-[]->(u)") + .where("u.id != v.id AND v.id != w.id AND u.id != w.id") + assert(fof.columns === Array("u", "e", "v", "w")) + compareResultToExpected( + fof.select("u.id", "v.id", "w.id").collect().toSet, + Set(Row(1L, 2L, 3L))) + + val fv = g.find("(u)") + assert(fv.columns === Array("u")) + + val fve = g.find("(u)-[e2]->()") + assert(fve.columns === Array("u", "e2")) + + val fed = g.find("()-[e]->(w)") + assert(fed.columns === Array("e", "w")) + } + + /* ================================= Invalid queries =================================== */ + + test("Disallow empty term ()-[]->()") { + intercept[InvalidParseException] { + g.find("()-[]->()") + } + } + + test("Disallow named edges in negated terms") { + intercept[InvalidParseException] { + g.find("!()-[ab]->()") + } + intercept[InvalidParseException] { + g.find("(u)-[]->(v); !(a)-[ab]->(b)") + } + intercept[InvalidParseException] { + g.find("(u)-[ab]->(v); !(a)-[ab]->(b)") + } + } + + test("Disallow using the same name for both a vertex and an edge") { + intercept[InvalidParseException] { + g.find("(a)-[a]->(b)") + } + intercept[InvalidParseException] { + g.find("(a)-[]->(b); (c)-[a]->(d)") + } + } + + test("Unbound variable-length pattern (u)->[*..5]->(v)") { + intercept[InvalidParseException] { + g.find("(u)-[*..5]->(v)") + } + } + + /* ============================= More complex use case examples ============================== */ + + test("triangles via post-hoc filter") { + val triangles = g + .find("(a)-[]->(b); (b)-[]->(c); (d)-[]->(e)") + .where("c.id = d.id AND e.id = a.id") + .select("a.id", "b.id", "c.id") + + val res = triangles.collect().toSet + compareResultToExpected(res, Set(Row(0L, 1L, 2L), Row(2L, 0L, 1L), Row(1L, 2L, 0L))) + } + + test("fixed-length 3") { + val fixedLengthEdge = g + .find("(u)-[*3]->(v)") + .where("u.id == 0") + .select("u.id", "_uv1.id", "_uv2.id", "v.id") + + val res = fixedLengthEdge.collect().toSet + val expected = Set(Row(0L, 1L, 2L, 0L), Row(0L, 1L, 2L, 3L), Row(0L, 1L, 0L, 1L)) + compareResultToExpected(res, expected) + } + + test("fixed-length 3 can be expressed with fixed-length 2 with a chain") { + val fixedLengthEdge1 = g + .find("(u)-[*2]->(v);(v)-[]->(k)") + .where("u.id == 0") + .select("u.id", "_uv1.id", "v.id", "k.id") + + val fixedLengthEdge2 = g + .find("(u)-[]->(v);(v)-[*2]->(k)") + .where("u.id == 0") + .select("u.id", "v.id", "_vk1.id", "k.id") + + val res1 = fixedLengthEdge1.collect().toSet + val res2 = fixedLengthEdge2.collect().toSet + + val expected = Set(Row(0L, 1L, 2L, 0L), Row(0L, 1L, 2L, 3L), Row(0L, 1L, 0L, 1L)) + compareResultToExpected(res1, expected) + compareResultToExpected(res2, expected) + } + + test("fixed-length 3 with named edge") { + val fixedLengthNamedEdge = g + .find("(u)-[e*3]->(v)") + .where("u.id == 0") + + val expectedCols = Seq("u", "_e1", "_uv1", "_e2", "_uv2", "_e3", "v") + + assert(fixedLengthNamedEdge.schema.map(_.name) == expectedCols) + } + + test("fixed-length 5") { + val fixedLengthEdge = g + .find("(u)-[*5]->(v)") + .where("u.id == 0") + .select("u.id", "_uv1.id", "_uv2.id", "_uv3.id", "_uv4.id", "v.id") + + val res = fixedLengthEdge.collect().toSet + val expected = Set( + Row(0L, 1L, 2L, 0L, 1L, 0L), + Row(0L, 1L, 0L, 1L, 0L, 1L), + Row(0L, 1L, 2L, 0L, 1L, 2L), + Row(0L, 1L, 0L, 1L, 2L, 0L), + Row(0L, 1L, 0L, 1L, 2L, 3L)) + compareResultToExpected(res, expected) + } + + test("fixed-length 5 can be expressed with a chain") { + val fixedLengthEdge1 = g + .find("(u)-[*5]->(v)") + .where("u.id == 0") + .select("u.id", "_uv1.id", "_uv2.id", "_uv3.id", "_uv4.id", "v.id") + + val fixedLengthEdge2 = g + .find("(u)-[*2]->(v);(v)-[*3]->(g)") + .where("u.id == 0") + .select("u.id", "_uv1.id", "v.id", "_vg1.id", "_vg2.id", "g.id") + + val fixedLengthEdge3 = g + .find("(u)-[*2]->(v);(v)-[*2]->(g);(g)-[e]->(k)") + .where("u.id == 0") + .select("u.id", "_uv1.id", "v.id", "_vg1.id", "g.id", "k.id") + + val res1 = fixedLengthEdge1.collect().toSet + val res2 = fixedLengthEdge2.collect().toSet + val res3 = fixedLengthEdge3.collect().toSet + + compareResultToExpected(res1, res2) + compareResultToExpected(res1, res3) + } + + test("var-length pattern 2..2") { + val varEdge = g + .find("(u)-[*2..2]->(v)") + .where("u.id == 0") + .drop("_hop", "_pattern", "_direction") + + val fixedEdge = g + .find("(u)-[*2]->(v)") + .where("u.id == 0") + + assert(varEdge.schema == fixedEdge.schema) + assert(varEdge.except(fixedEdge).isEmpty && fixedEdge.except(varEdge).isEmpty) + } + + test("var-length pattern 2..3") { + val varEdge = g + .find("(u)-[*2..3]->(v)") + .where("u.id == 0") + .drop("_hop", "_pattern", "_direction") + + val fixedEdge2 = g + .find("(u)-[*2]->(v)") + .where("u.id == 0") + + val fixedEdge3 = g + .find("(u)-[*3]->(v)") + .where("u.id == 0") + + val unionEdge = fixedEdge3 + .unionByName(fixedEdge2, allowMissingColumns = true) + + assert(varEdge.schema == unionEdge.schema) + assert(varEdge.except(unionEdge).isEmpty && unionEdge.except(varEdge).isEmpty) + } + + test("var-length pattern 2..3 with named edge") { + val varEdge = g + .find("(u)-[e*2..3]->(v)") + .where("u.id == 0") + + val expectedCols = + Seq("u", "_e1", "_uv1", "_e2", "_uv2", "_e3", "v", "_hop", "_pattern", "_direction") + + assert(varEdge.schema.map(_.name) == expectedCols) + } + + test("var-length pattern 3..5") { + val varEdge = g + .find("(u)-[*3..5]->(v)") + .where("u.id == 0") + .drop("_hop", "_pattern", "_direction") + + val fixedEdge3 = g + .find("(u)-[*3]->(v)") + .where("u.id == 0") + + val fixedEdge4 = g + .find("(u)-[*4]->(v)") + .where("u.id == 0") + + val fixedEdge5 = g + .find("(u)-[*5]->(v)") + .where("u.id == 0") + + val unionEdge = fixedEdge5 + .unionByName(fixedEdge4, allowMissingColumns = true) + .unionByName(fixedEdge3, allowMissingColumns = true) + + assert(varEdge.schema == unionEdge.schema) + assert(varEdge.except(unionEdge).isEmpty && unionEdge.except(varEdge).isEmpty) + } + + test("undirected edge") { + val res = g + .find("(u)-[]-(v)") + .where("u.id == 0") + .select("u.id", "v.id") + .collect() + .toSet + + val expected = Set(Row(0L, 1L), Row(0L, 2L)) + + compareResultToExpected(res, expected) + } + + test("undirected information column") { + val res1 = g + .find("(u)-[e1]-(v)") + .where("u.id == 0") + .select("_pattern", "_direction") + .collect() + .toSet + + val expected1 = Set(Row("(u)<-[e1]-(v)", "in"), Row("(u)-[e1]->(v)", "out")) + + compareResultToExpected(res1, expected1) + + val res2 = g + .find("(u)-[]-(v)") + .where("u.id == 0") + .select("_pattern", "_direction") + .collect() + .toSet + + val expected2 = Set(Row("(u)<-[]-(v)", "in"), Row("(u)-[]->(v)", "out")) + + compareResultToExpected(res2, expected2) + } + + test("undirected edge within a chain") { + val res = g + .find("(u)-[]-(v);(v)-[]->(k)") + .where("u.id == 0") + .select("u.id", "v.id", "k.id") + .collect() + .toSet + + val expected = Set(Row(0L, 1L, 2L), Row(0L, 1L, 0L), Row(0L, 2L, 0L), Row(0L, 2L, 3L)) + + compareResultToExpected(res, expected) + } + + test("undirected with edge name") { + val res = g + .find("(u)-[e]-(v)") + .where("u.id == 0") + .select("e.src", "e.dst", "e.relationship") + .collect() + .toSet + + val expected = Set(Row(0L, 1L, "friend"), Row(1L, 0L, "follow"), Row(2L, 0L, "unknown")) + + compareResultToExpected(res, expected) + } + + test("undirected var-length pattern") { + val res = g + .find("(u)-[e*1..3]-(v)") + .where("u.id == 2") + .drop("_pattern", "_direction") + + val df1 = g + .find("(u)-[e*1..3]->(v)") + .where("u.id == 2") + .drop("_pattern", "_direction") + + val df2 = g + .find("(v)-[e*1..3]->(u)") + .where("u.id == 2") + .drop("_pattern", "_direction") + + val expected = df1.unionByName(df2, allowMissingColumns = true) + + assert(res.schema === expected.schema) + assert(res.except(expected).isEmpty && expected.except(res).isEmpty) + } + + test("undirected fixed-length pattern") { + val res = g.find("(u)-[e*3]-(v)") + val expected = g.find("(u)-[e*3..3]-(v)") + + assert(res.schema === expected.schema) + assert(res.except(expected).isEmpty && expected.except(res).isEmpty) + } + + test("undirected edge without vertex name") { + val res = g.find("()-[e*3]-()").drop("_pattern").collect().toSet + val expected = + g.find("(u)-[e*3]-(v)").select("_e1", "_e2", "_e3", "_hop", "_direction").collect().toSet + + compareResultToExpected(res, expected) + } + + test("directed edge name without vertex name") { + val res = g.find("()-[e*3]->()").collect().toSet + val expected = g.find("(u)-[e*3]->(v)").select("_e1", "_e2", "_e3").collect().toSet + + compareResultToExpected(res, expected) + } + + test("stateful predicates via UDFs") { + val chain4 = g + .find("(a)-[ab]->(b); (b)-[bc]->(c); (c)-[cd]->(d)") + .where("a.id != b.id AND b.id != c.id AND c.id != a.id") + + // Using DataFrame operations, but not really operating in a stateful manner + val chainWith2Friends = chain4.where( + Seq("ab", "bc", "cd") + .map(e => when(col(e)("relationship") === "friend", 1).otherwise(0)) + .reduce(_ + _) >= 2) + + assert(chainWith2Friends.count() === 4) + chainWith2Friends + .select("ab.relationship", "bc.relationship", "cd.relationship") + .collect() + .foreach { + case Row(ab: String, bc: String, cd: String) => + val numFriends = Seq(ab, bc, cd).map(r => if (r == "friend") 1 else 0).sum + assert(numFriends >= 2) + case _ => throw new GraphFramesUnreachableException() + } + + // Operating in a stateful manner, where cnt is the state. + def sumFriends(cnt: Column, relationship: Column): Column = { + when(relationship === "friend", cnt + 1).otherwise(cnt) + } + val condition = + Seq("ab", "bc", "cd").foldLeft(lit(0))((cnt, e) => sumFriends(cnt, col(e)("relationship"))) + val chainWith2Friends2 = chain4.where(condition >= 2) + + compareResultToExpected(chainWith2Friends.collect().toSet, chainWith2Friends2.collect().toSet) + } + + /* ===================================== Join elimination =================================== */ + + /* + // Join elimination will not work without Ankur's improved indexing. + test("join elimination - simple") { + import org.apache.spark.sql.catalyst.plans.logical.Join + + val edges = g.find("(u)-[e]->(v)", _.select("e_src", "e_dst")) + val joins = edges.queryExecution.optimizedPlan.collect { + case j: Join => j + } + + assert(joins.isEmpty, s"joins was non-empty: ${joins.map(_.toString()).mkString("; ")}") + } + + test("join elimination - with aliases") { + import org.apache.spark.sql.catalyst.plans.logical.Join + + val edges = g.find("(u)-[]->(v)", _.select("u_id", "v_id")) + println(edges.queryExecution.optimizedPlan) + val joins = edges.queryExecution.optimizedPlan.collect { + case j: Join => j + } + assert(joins.isEmpty, s"joins was non-empty: ${joins.map(_.toString()).mkString("; ")}") + } + */ +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/SparkFunSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/SparkFunSuite.scala new file mode 100644 index 0000000000000..b0ed7c370a48a --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/SparkFunSuite.scala @@ -0,0 +1,45 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes + +import org.scalatest.Outcome +import org.scalatest.funsuite.AnyFunSuite + +/** + * Base abstract class for all unit tests in Spark for handling common functionality. + */ +abstract class SparkFunSuite extends AnyFunSuite with Logging { + + /** + * Log the suite name and the test name before and after each test. + * + * Subclasses should never override this method. If they wish to run custom code before and + * after each test, they should mix in the {{org.scalatest.BeforeAndAfter}} trait instead. + */ + final protected override def withFixture(test: NoArgTest): Outcome = { + val testName = test.text + val suiteName = this.getClass.getName + val shortSuiteName = suiteName.replaceAll("org.apache.spark", "o.a.s") + try { + logInfo(s"\n\n===== TEST OUTPUT FOR $shortSuiteName: '$testName' =====\n") + test() + } finally { + logInfo(s"\n\n===== FINISHED $shortSuiteName: '$testName' =====\n") + } + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/TestUtils.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/TestUtils.scala new file mode 100644 index 0000000000000..1302c424561c1 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/TestUtils.scala @@ -0,0 +1,128 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes + +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.types.DataType +import org.apache.spark.sql.types.StructType +import org.apache.spark.graphframes.GraphFrame._ + +object TestUtils { + + private[this] val majorMinorRegex = """^(\d+)\.(\d+)(\..*)?$""".r + + /** Extract major/minor version integer pairs from a version string */ + def majorMinorVersion(sparkVersion: String): (Int, Int) = { + majorMinorRegex.findFirstMatchIn(sparkVersion) match { + case Some(m) => + (m.group(1).toInt, m.group(2).toInt) + case None => + throw new IllegalArgumentException( + s"Spark tried to parse '$sparkVersion' as a Spark" + + " version string, but it could not find the major and minor version numbers.") + } + } + + /** Return true if the major and minor versions are greater or eq to constraints */ + def requireSparkVersionGE(major: Int, minor: Int, sparkVersion: String): Boolean = { + val (gotMajor, gotMinor) = TestUtils.majorMinorVersion(sparkVersion) + (gotMajor > major) || ((gotMajor == major) && (gotMinor >= minor)) + } + + /** + * Check whether the given schema contains a column of the required data type. + * + * @param colName + * column name + * @param dataType + * required column data type + */ + def checkColumnType( + schema: StructType, + colName: String, + dataType: DataType, + msg: String = ""): Unit = { + val actualDataType = schema(colName).dataType + val message = if (msg != null && msg.trim.length > 0) " " + msg else "" + require( + actualDataType.equals(dataType), + s"Column $colName must be of type $dataType but was actually $actualDataType.$message") + } + + /** Confirm ID, SRC, DST columns are present */ + def testSchemaInvariant(g: GraphFrame): Unit = { + val vCols = g.vertices.columns + val eCols = g.edges.columns + assert(vCols.contains(ID)) + assert(eCols.contains(SRC)) + assert(eCols.contains(DST)) + } + + /** + * Test validity of both GraphFrames. Also ensure that the GraphFrames match: + * - vertex column schema match + * - `before` columns are a subset of the `after` columns, and schema match + */ + def testSchemaInvariants(before: GraphFrame, after: GraphFrame): Unit = { + testSchemaInvariant(before) + testSchemaInvariant(after) + // The IDs, source and destination columns should be of the same type + // with the same metadata. + for (colName <- Seq(ID)) { + val b = before.vertices.schema(colName) + val a = after.vertices.schema(colName) + // TODO(tjh) check nullability and metadata + assert(a.dataType == b.dataType, (a, b)) + } + for (colName <- Seq(SRC, DST)) { + val b = before.edges.schema(colName) + val a = after.edges.schema(colName) + // TODO(tjh) check nullability and metadata + assert(a.dataType == b.dataType, (a, b)) + } + // All the columns before should be found after (with some extra columns, + // potentially). + val afterVNames = before.vertices.schema.fields.map(_.name) + for (f <- before.vertices.schema.iterator) { + if (!afterVNames.contains(f.name)) { + throw new Exception(s"vertex error: ${f.name} should be in ${afterVNames.mkString(", ")}") + } + assert( + before.vertices.schema(f.name) == after.vertices.schema(f.name), + s"${before.vertices.schema} != ${after.vertices.schema}") + } + + for (f <- before.edges.schema.iterator) { + val a = before.edges.schema(f.name) + val b = after.edges.schema(f.name) + assert( + a.dataType == b.dataType, + s"${before.edges.schema} not a subset of ${after.edges.schema}") + } + } + + /** + * Test validity of both GraphFrames. Also ensure that the GraphFrames match: + * - vertex column schema match + * - `before` columns are a subset of the `after` columns, and schema match + */ + def testSchemaInvariants(before: GraphFrame, afterVertices: DataFrame): Unit = { + testSchemaInvariants(before, GraphFrame(afterVertices, before.edges)) + } + +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/convolutions/SamplingConvolutionSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/convolutions/SamplingConvolutionSuite.scala new file mode 100644 index 0000000000000..4ddb1dd395957 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/convolutions/SamplingConvolutionSuite.scala @@ -0,0 +1,135 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.convolutions + +import org.apache.spark.ml.linalg.DenseVector +import org.apache.spark.ml.linalg.Vector +import org.apache.spark.ml.linalg.Vectors +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.scalatest.BeforeAndAfterAll + +class SamplingConvolutionSuite + extends SparkFunSuite + with GraphFrameTestSparkContext + with BeforeAndAfterAll { + + private var testGraph: GraphFrame = _ + + override def beforeAll(): Unit = { + super.beforeAll() + // Create test graph with 5 vertices and randomish edges + val vertices: DataFrame = spark + .createDataFrame( + Seq( + (0L, Vectors.dense(0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0)), + (1L, Vectors.dense(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0)), + (2L, Vectors.dense(2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0)), + (3L, Vectors.dense(3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0)), + (4L, Vectors.dense(4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0)))) + .toDF("id", "embedding") + + val edges: DataFrame = spark + .createDataFrame(Seq((0L, 1L), (0L, 2L), (1L, 2L), (1L, 3L), (2L, 3L), (2L, 4L), (3L, 4L))) + .toDF("src", "dst") + + testGraph = GraphFrame(vertices, edges) + } + + test("big maxNbrs: result is correct average of feature vectors") { + val conv = new SamplingConvolution() + .onGraph(testGraph) + .setFeaturesCol("embedding") + .setMaxNbrs(10) // Large number, more than neighbors + .setConcatEmbeddings(false) + .setSeed(42L) + + val result = conv.run() + // Collect nbr_embedding for vertex 0 (neighbors 1,2) + val row = result.where(col("id") === 0L).select("nbr_embedding").collect()(0) + val vec: Vector = row.getAs[Vector](0) + val expected = Vectors.dense( + (1 + 2) / 2.0, + (2 + 3) / 2.0, + (3 + 4) / 2.0, + (4 + 5) / 2.0, + (5 + 6) / 2.0, + (6 + 7) / 2.0, + (7 + 8) / 2.0, + (8 + 9) / 2.0, + (9 + 10) / 2.0, + (10 + 11) / 2.0) + assert(vec === expected) + } + + test("small maxNbrs: min-hash sampling is reproducible") { + val conv1 = new SamplingConvolution() + .onGraph(testGraph) + .setFeaturesCol("embedding") + .setMaxNbrs(1) + .setConcatEmbeddings(false) + .setSeed(100L) + + val result1 = conv1.run() + val conv2 = new SamplingConvolution() + .onGraph(testGraph) + .setFeaturesCol("embedding") + .setMaxNbrs(1) + .setConcatEmbeddings(false) + .setSeed(100L) + + val result2 = conv2.run() + // Both should have same nbr_embedding for same seeds + val rows1 = result1.select("id", "nbr_embedding").collect().sortBy(r => r.getLong(0)) + val rows2 = result2.select("id", "nbr_embedding").collect().sortBy(r => r.getLong(0)) + rows1.zip(rows2).foreach { case (r1, r2) => + val v1: Vector = r1.getAs[DenseVector](1) + val v2: Vector = r2.getAs[DenseVector](1) + assert(v1 === v2) + } + } + + test("concatenating features increases size correctly") { + val convConcat = new SamplingConvolution() + .onGraph(testGraph) + .setFeaturesCol("embedding") + .setMaxNbrs(5) + .setConcatEmbeddings(true) + .setSeed(42L) + + val resultConcat = convConcat.run() + val rowConcat = resultConcat.where(col("id") === 0L).select("embedding").collect()(0) + val vecConcat: Vector = rowConcat.getAs[Vector](0) + assert(vecConcat.size == 20) // Original 10 + 10 from nbr + + val convNoConcat = new SamplingConvolution() + .onGraph(testGraph) + .setFeaturesCol("embedding") + .setMaxNbrs(5) + .setConcatEmbeddings(false) + .setSeed(42L) + + val resultNoConcat = convNoConcat.run() + val rowNoConcat = resultNoConcat.where(col("id") === 0L).select("embedding").collect()(0) + val vecNoConcat: Vector = rowNoConcat.getAs[Vector](0) + assert(vecNoConcat.size == 10) // Only original + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/embeddings/Hash2VecSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/embeddings/Hash2VecSuite.scala new file mode 100644 index 0000000000000..daaaca149f866 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/embeddings/Hash2VecSuite.scala @@ -0,0 +1,388 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.embeddings + +import org.apache.spark.ml.linalg.DenseVector +import org.apache.spark.ml.linalg.SQLDataTypes.VectorType +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.types.LongType +import org.apache.spark.sql.types.StringType +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.scalatest.BeforeAndAfterAll + +import scala.util.Random + +class Hash2VecSuite extends SparkFunSuite with GraphFrameTestSparkContext with BeforeAndAfterAll { + private var longSequences: DataFrame = _ + private var stringSequences: DataFrame = _ + private var uniqueElementsCnt: Int = _ + + private def approxEqual(left: Double, right: Double, err: Double = 1e-6.toDouble): Boolean = + math.abs(left - right) < err + + override def beforeAll(): Unit = { + super.beforeAll() + + val rng = new Random(42L) + + val sequences = + (1 to 100).map(idx => (idx, (1 to 20).map(_ => rng.nextInt(30).toLong).toSeq)).toSeq + + uniqueElementsCnt = sequences.flatMap(f => f._2).distinct.length + + val strSequences = sequences.map(f => (f._1, f._2.map(_.toString()))).toSeq + + longSequences = spark.createDataFrame(sequences).toDF("id", "seq") + stringSequences = spark.createDataFrame(strSequences).toDF("id", "seq") + } + + test("hash2vec long input") { + val hash2vecResults = new Hash2Vec().setSequenceCol("seq").run(longSequences) + assert(hash2vecResults.schema.fields.length === 2) + assert(hash2vecResults.schema.fields.map(_.name).toSeq === Seq("id", "vector")) + assert(hash2vecResults.schema("id").dataType === LongType) + assert(hash2vecResults.schema("vector").dataType === VectorType) + val collected = hash2vecResults.collect() + assert(collected.length === uniqueElementsCnt) + } + + test("hash2vec string input") { + val hash2vecResults = new Hash2Vec().setSequenceCol("seq").run(stringSequences) + assert(hash2vecResults.schema.fields.length === 2) + assert(hash2vecResults.schema.fields.map(_.name).toSeq === Seq("id", "vector")) + assert(hash2vecResults.schema("id").dataType === StringType) + assert(hash2vecResults.schema("vector").dataType === VectorType) + val collected = hash2vecResults.collect() + assert(collected.length === uniqueElementsCnt) + } + + test("hash2vec reproducible with seed") { + val hash2vecResults = + new Hash2Vec().setSequenceCol("seq").setHashingSeed(42).run(longSequences) + val hash2vecResults2 = + new Hash2Vec().setSequenceCol("seq").setHashingSeed(42).run(longSequences) + val hash2vecResults3 = + new Hash2Vec().setSequenceCol("seq").setHashingSeed(43).run(longSequences) + + val shouldMatch = hash2vecResults + .withColumnRenamed("vector", "left") + .join(hash2vecResults2, Seq("id"), "inner") + + assert(shouldMatch.count() === hash2vecResults.count()) + assert(shouldMatch.filter(col("left") =!= col("vector")).count() === 0) + + val shouldNotMatch = hash2vecResults + .withColumnRenamed("vector", "left") + .join(hash2vecResults3, Seq("id"), "inner") + + assert(shouldNotMatch.count() === uniqueElementsCnt) + assert(shouldNotMatch.filter(col("left") =!= col("vector")).count() > 0) + } + + test("hash2vec L2") { + val hash2vecResults = + new Hash2Vec() + .setDoNormalization(true, false) + .setSequenceCol("seq") + .run(longSequences) + .collect() + + def naiveL2norm(vector: DenseVector): Double = { + val squaredSum = vector.values.map(el => el * el).sum[Double] + math.sqrt(squaredSum) + } + + assert( + hash2vecResults + .map(r => r.getAs[DenseVector](1)) + .map(v => naiveL2norm(v)) + .forall(f => approxEqual(f, 1.0))) + } + + test("hash2vec safe L2") { + val hash2vecResults = new Hash2Vec() + .setDoNormalization(true, true) + .setEmbeddingsDim(128) + .setSequenceCol("seq") + .run(longSequences) + .collect() + + assert(hash2vecResults.forall(r => r.getAs[DenseVector](1).size === 129)) + } + + test("constant decay") { + val hash2vecResults = new Hash2Vec() + .setDecayFunction("constant") + .setSequenceCol("seq") + .run(longSequences) + hash2vecResults.write.mode("overwrite").format("noop").save() + } + + test("context longer than sequence") { + val hash2vecResults = new Hash2Vec() + .setSequenceCol("seq") + .setContextSize(30) + .run(longSequences) + hash2vecResults.write.mode("overwrite").format("noop").save() + } + + test("PagedMatrixDouble helper - page extension") { + val dim = 50 + val matrix = new Hash2Vec.PagedMatrixDouble(dim) + // Allocate enough vectors to force a second page (PAGE_SIZE = 4096) + // We'll allocate two pages' worth to be sure. + // For simplicity, allocate 2 * PAGE_SIZE vectors. + val PAGE_SIZE = 1 << 12 // 4096 + val totalVectors = 2 * PAGE_SIZE + val ids = (0 until totalVectors).map { _ => + matrix.allocateVector() + } + // IDs should be 0,1,2,...,totalVectors-1 + assert(ids.toSeq == (0 until totalVectors).toSeq) + // Check that internal pages count grew + // Since internal structure is private, we just verify no exception occurred. + } + + test("PagedMatrixDouble helper - add and retrieve") { + val dim = 10 + val matrix = new Hash2Vec.PagedMatrixDouble(dim) + val id0 = matrix.allocateVector() + val id1 = matrix.allocateVector() + assert(id0 == 0) + assert(id1 == 1) + + // Add values to vector 0 at different offsets + matrix.add(id0, 0, 5.0) + matrix.add(id0, 3, 2.0) + matrix.add(id0, 9, -1.0) + + // Add values to vector 1 + matrix.add(id1, 0, 10.0) + matrix.add(id1, 5, 3.0) + + // Retrieve and verify + val vec0 = matrix.getVector(id0) + assert(vec0.length == dim) + assert(vec0(0) == 5.0) + assert(vec0(3) == 2.0) + assert(vec0(9) == -1.0) + // Other positions should be zero + (0 until dim).filterNot(idx => idx == 0 || idx == 3 || idx == 9).foreach { idx => + assert(vec0(idx) == 0.0) + } + + val vec1 = matrix.getVector(id1) + assert(vec1.length == dim) + assert(vec1(0) == 10.0) + assert(vec1(5) == 3.0) + (0 until dim).filterNot(idx => idx == 0 || idx == 5).foreach { idx => + assert(vec1(idx) == 0.0) + } + } + + test("PagedMatrixDouble helper - cross-page addressing") { + val dim = 20 + val matrix = new Hash2Vec.PagedMatrixDouble(dim) + val PAGE_SIZE = 1 << 12 + // Allocate vectors up to the end of first page + for (_ <- 0 until PAGE_SIZE) matrix.allocateVector() + val firstPageLastId = PAGE_SIZE - 1 + // Allocate first vector of second page + val secondPageFirstId = matrix.allocateVector() + assert(secondPageFirstId == PAGE_SIZE) + + // Add to the last vector of first page + matrix.add(firstPageLastId, 0, 100.0) + matrix.add(firstPageLastId, dim - 1, 200.0) + + // Add to the first vector of second page + matrix.add(secondPageFirstId, 0, 300.0) + matrix.add(secondPageFirstId, 10, 400.0) + + // Retrieve and verify + val vecFirstPage = matrix.getVector(firstPageLastId) + assert(vecFirstPage(0) == 100.0) + assert(vecFirstPage(dim - 1) == 200.0) + (1 until dim - 1).foreach { idx => + assert(vecFirstPage(idx) == 0.0) + } + + val vecSecondPage = matrix.getVector(secondPageFirstId) + assert(vecSecondPage(0) == 300.0) + assert(vecSecondPage(10) == 400.0) + (1 until dim).filterNot(_ == 10).foreach { idx => + assert(vecSecondPage(idx) == 0.0) + } + } + + test("Hash2Vec - cosine distances reflect co‑occurrence patterns") { + // Create a tiny dataset where some words co‑occur often, others rarely. + // We'll use string sequences for simplicity. + val sequences = Seq( + Seq("apple", "banana", "apple", "cherry", "banana"), + Seq("apple", "banana", "cherry", "banana"), + Seq("apple", "banana", "apple", "banana", "banana"), + Seq("cherry", "date", "cherry", "date"), + Seq("date", "elderberry", "date"), + Seq("elderberry", "fig", "elderberry"), + Seq("fig", "fig", "fig") // fig appears often alone + ) + + val df = spark.createDataFrame(sequences.map(Tuple1(_))).toDF("seq") + + val embeddings = new Hash2Vec() + .setSequenceCol("seq") + .setEmbeddingsDim(128) // enough dimensions to capture patterns + .setContextSize(2) + .setDecayFunction("constant") + .setHashingSeed(777) + .setSignHashSeed(888) + .run(df) + + // Collect embeddings into a local map + val embMap = embeddings + .collect() + .map { row => + val id = row.getString(0) + val vec = row.getAs[DenseVector](1) + id -> vec + } + .toMap + + // Helper to compute cosine similarity between two vectors + def cosineSimilarity(v1: DenseVector, v2: DenseVector): Double = { + val a = v1.values + val b = v2.values + require(a.length == b.length) + var dot = 0.0 + var norm1 = 0.0 + var norm2 = 0.0 + var i = 0 + while (i < a.length) { + dot += a(i) * b(i) + norm1 += a(i) * a(i) + norm2 += b(i) * b(i) + i += 1 + } + dot / (math.sqrt(norm1) * math.sqrt(norm2)) + } + + // apple and banana co‑occur very frequently → high similarity + val appleBananaSim = cosineSimilarity(embMap("apple"), embMap("banana")) + // cherry and date also co‑occur frequently (in the fourth sequence) + val cherryDateSim = cosineSimilarity(embMap("cherry"), embMap("date")) + // apple and fig almost never appear together → low similarity + val appleFigSim = cosineSimilarity(embMap("apple"), embMap("fig")) + // banana and fig also rarely together + val bananaFigSim = cosineSimilarity(embMap("banana"), embMap("fig")) + // elderberry and fig co‑occur (sixth sequence) + val elderberryFigSim = cosineSimilarity(embMap("elderberry"), embMap("fig")) + + // Assert ordering of similarities matches expected co‑occurrence patterns + // apple‑banana should be among the highest similarities + assert(appleBananaSim > 0.3, s"apple‑banana similarity $appleBananaSim should be > 0.3") + // apple‑fig should be low (close to zero or negative) + assert( + appleFigSim < appleBananaSim, + s"apple‑fig ($appleFigSim) should be < apple‑banana ($appleBananaSim)") + assert( + bananaFigSim < appleBananaSim, + s"banana‑fig ($bananaFigSim) should be < apple‑banana ($appleBananaSim)") + // cherry‑date similarity should be relatively high (they co‑occur exclusively) + assert(cherryDateSim > 0.2, s"cherry‑date similarity $cherryDateSim should be > 0.2") + // elderberry‑fig should be higher than apple‑fig (because they co‑occur) + assert( + elderberryFigSim > appleFigSim, + s"elderberry‑fig ($elderberryFigSim) should be > apple‑fig ($appleFigSim)") + + // Self‑similarity should be 1.0 (or close after normalization) + val appleSelf = cosineSimilarity(embMap("apple"), embMap("apple")) + assert(math.abs(appleSelf - 1.0) < 1e-6, s"self‑similarity should be ~1.0, got $appleSelf") + } + + test("Hash2Vec - long‑typed co‑occurrence") { + // Use numeric ids to test long sequences. + val sequences = Seq( + Seq(1L, 2L, 1L, 3L, 2L), // 1‑2 frequent, 3 appears with 2 + Seq(1L, 2L, 3L, 2L), + Seq(1L, 2L, 1L, 2L, 2L), + Seq(3L, 4L, 3L, 4L), // 3‑4 frequent pair + Seq(4L, 5L, 4L), + Seq(5L, 6L, 5L), + Seq(6L, 6L, 6L)) + + val df = spark.createDataFrame(sequences.map(Tuple1(_))).toDF("seq") + + val embeddings = new Hash2Vec() + .setSequenceCol("seq") + .setEmbeddingsDim(128) + .setContextSize(2) + .setDecayFunction("constant") + .setHashingSeed(777) + .setSignHashSeed(888) + .run(df) + + val embMap = embeddings + .collect() + .map { row => + val id = row.getLong(0) + val vec = row.getAs[DenseVector](1) + id -> vec + } + .toMap + + def cosineSimilarity(v1: DenseVector, v2: DenseVector): Double = { + val a = v1.values + val b = v2.values + var dot = 0.0 + var norm1 = 0.0 + var norm2 = 0.0 + var i = 0 + while (i < a.length) { + dot += a(i) * b(i) + norm1 += a(i) * a(i) + norm2 += b(i) * b(i) + i += 1 + } + dot / (math.sqrt(norm1) * math.sqrt(norm2)) + } + + val sim12 = cosineSimilarity(embMap(1L), embMap(2L)) + val sim13 = cosineSimilarity(embMap(1L), embMap(3L)) + val sim34 = cosineSimilarity(embMap(3L), embMap(4L)) + val sim16 = cosineSimilarity(embMap(1L), embMap(6L)) + val sim56 = cosineSimilarity(embMap(5L), embMap(6L)) + + // 1‑2 co‑occur very often + assert(sim12 > 0.3, s"1‑2 similarity $sim12 should be > 0.3") + // 1‑3 appear together less often than 1‑2 + assert(sim13 < sim12, s"1‑3 ($sim13) should be < 1‑2 ($sim12)") + // 3‑4 are exclusive pair + assert(sim34 > 0.2, s"3‑4 similarity $sim34 should be > 0.2") + // 1 and 6 almost never together + assert(sim16 < 0.1, s"1‑6 similarity $sim16 should be near zero") + // 5‑6 co‑occur in a sequence + assert(sim56 > sim16, s"5‑6 ($sim56) should be > 1‑6 ($sim16)") + + // Self similarity + val self = cosineSimilarity(embMap(1L), embMap(1L)) + assert(math.abs(self - 1.0) < 1e-6, s"self‑similarity should be ~1.0, got $self") + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/examples/Graphs.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/examples/Graphs.scala new file mode 100644 index 0000000000000..a0fa223398b70 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/examples/Graphs.scala @@ -0,0 +1,248 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.examples + +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.functions.randn +import org.apache.spark.sql.functions.udf +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrame._ + +import scala.reflect.runtime.universe.TypeTag + +class Graphs private[graphframes] () { + // Note: this cannot be values: we are creating and destroying spark contexts during the tests, + // and turning these into vals means we would hold onto a potentially destroyed spark context. + private def spark: SparkSession = SparkSession.builder().getOrCreate() + + /** + * Returns an empty GraphFrame of the given ID type. + */ + def empty[T: TypeTag]: GraphFrame = { + val _spark = spark + import _spark.implicits._ + val vertices = Seq.empty[Tuple1[T]].toDF(ID) + val edges = Seq.empty[(T, T)].toDF(SRC, DST) + GraphFrame(vertices, edges) + } + + /** + * Returns a chain graph of the given size with Long ID type. The vertex IDs are 0, 1, ..., n-1, + * and the edges are (0, 1), (1, 2), ...., (n-2, n-1). + */ + def chain(n: Long): GraphFrame = { + require(n >= 0, s"Chain graph size must be nonnegative but got $n.") + val vertices = spark.range(n).toDF(ID) + val edges = spark + .range(n - 1L) + .toDF(ID) + .select(col(ID).as(SRC), (col(ID) + 1L).as(DST)) + GraphFrame(vertices, edges) + } + + /** + * Graph of friends in a social network. + */ + def friends: GraphFrame = { + // For the same reason as above, this cannot be a value. + // Vertex DataFrame + val v = spark + .createDataFrame( + List( + ("a", "Alice", 34), + ("b", "Bob", 36), + ("c", "Charlie", 30), + ("d", "David", 29), + ("e", "Esther", 32), + ("f", "Fanny", 36), + ("g", "Gabby", 60))) + .toDF("id", "name", "age") + // Edge DataFrame + val e = spark + .createDataFrame( + List( + ("a", "b", "friend"), + ("b", "c", "follow"), + ("c", "b", "follow"), + ("f", "c", "follow"), + ("e", "f", "follow"), + ("e", "d", "friend"), + ("d", "a", "friend"), + ("a", "e", "friend"))) + .toDF("src", "dst", "relationship") + // Create a GraphFrame + GraphFrame(v, e) + } + + /** + * Two densely connected blobs (vertices 0->n-1 and n->2n-1) connected by a single edge (0->n) + * @param blobSize + * the size of each blob. + * @return + */ + def twoBlobs(blobSize: Int): GraphFrame = { + val n = blobSize + val edges1 = for (v1 <- 0 until n; v2 <- 0 until n) yield (v1.toLong, v2.toLong, s"$v1-$v2") + val edges2 = for { + v1 <- n until (2 * n) + v2 <- n until (2 * n) + } yield (v1.toLong, v2.toLong, s"$v1-$v2") + val edges = edges1 ++ edges2 ++ Seq((0L, n.toLong, s"0-$n")) + val vertices = (0 until (2 * n)).map { v => (v.toLong, s"$v", v) } + val e = spark.createDataFrame(edges).toDF("src", "dst", "e_attr1") + val v = spark.createDataFrame(vertices).toDF("id", "v_attr1", "v_attr2") + GraphFrame(v, e) + } + + /** + * Returns a star graph with Long ID type, consisting of a central element indexed 0 (the root) + * and the n other leaf vertices 1, 2, ..., n. + * @param n + * the number of leaves + */ + def star(n: Long): GraphFrame = { + require(n >= 0L) + val vertices = spark.range(n + 1L).toDF(ID) + val edges = spark.range(1L, n + 1L).toDF(DST).withColumn(SRC, lit(0L)) + GraphFrame(vertices, edges) + } + + /** + * Some synthetic data that sits in Spark. + * + * No description available. + * @return + */ + def ALSSyntheticData(): GraphFrame = { + val sc = spark.sparkContext + val data = sc.parallelize(als_data.toIndexedSeq).map { line => + val fields = line.split(",") + (fields(0).toLong * 2, fields(1).toLong * 2 + 1, fields(2).toDouble) + } + val edges = spark.createDataFrame(data).toDF("src", "dst", "weight") + val vs = + data.flatMap(r => r._1 :: r._2 :: Nil).collect().distinct.map(x => Tuple1(x)).toIndexedSeq + val vertices = spark.createDataFrame(vs).toDF("id") + GraphFrame(vertices, edges) + } + + private lazy val als_data = + """ + |1,1,5.0 + |1,2,1.0 + |1,3,5.0 + |1,4,1.0 + |2,1,5.0 + |2,2,1.0 + |2,3,5.0 + |2,4,1.0 + |3,1,1.0 + |3,2,5.0 + |3,3,1.0 + |3,4,5.0 + |4,1,1.0 + |4,2,5.0 + |4,3,1.0 + |4,4,5.0 + """.stripMargin.split("\n").map(_.trim).filterNot(_.isEmpty) + + /** + * This method generates a grid Ising model with random parameters. + * + * Ising models are probabilistic graphical models over binary variables x,,i,,. Each binary + * variable x,,i,, corresponds to one vertex, and it may take values -1 or +1. The probability + * distribution P(X) (over all x,,i,,) is parameterized by vertex factors a,,i,, and edge + * factors b,,ij,,: + * {{{ + * P(X) = (1/Z) * exp[ \sum_i a_i x_i + \sum_{ij} b_{ij} x_i x_j ] + * }}} + * where Z is the normalization constant (partition function). See + * [[https://en.wikipedia.org/wiki/Ising_model Wikipedia]] for more information on Ising models. + * + * Each vertex is parameterized by a single scalar a,,i,,. Each edge is parameterized by a + * single scalar b,,ij,,. + * + * @param n + * Length of one side of the grid. The grid will be of size n x n. + * @param vStd + * Standard deviation of normal distribution used to generate vertex factors "a". Default of + * 1.0. + * @param eStd + * Standard deviation of normal distribution used to generate edge factors "b". Default of + * 1.0. + * @return + * GraphFrame. Vertices have columns "id" and "a". Edges have columns "src", "dst", and "b". + * Edges are directed, but they should be treated as undirected in any algorithms run on this + * model. Vertex IDs are of the form "i,j". E.g., vertex "1,3" is in the second row and fourth + * column of the grid. + */ + def gridIsingModel(spark: SparkSession, n: Int, vStd: Double, eStd: Double): GraphFrame = { + require(n >= 1, s"Grid graph must have size >= 1, but was given invalid value n = $n") + + // To create grid + // Avoid Cartesian join due to SPARK-15425: use generator since n should be small + val res = for { x <- Range(0, n); y <- Range(0, n) } yield (x, y) + val coordinates = spark.createDataFrame(res).toDF("i", "j") + + // Create SQL expression for converting coordinates (i,j) to a string ID "i,j" + val toIDudf = udf { (i: Int, j: Int) => i.toString + "," + j.toString } + + // Create the vertex DataFrame + // Create SQL expression for converting coordinates (i,j) to a string ID "i,j" + val vIDcol = toIDudf(col("i"), col("j")) + // Add random parameters generated from a normal distribution + val seed = 12345 + val vertices = coordinates + .withColumn("id", vIDcol) // vertex IDs "i,j" + .withColumn("a", randn(seed.toLong) * vStd) // Ising parameter for vertex + + // Create the edge DataFrame + // Create SQL expression for converting coordinates (i,j+1) and (i+1,j) to string IDs + val rightIDcol = toIDudf(col("i"), col("j") + 1) + val downIDcol = toIDudf(col("i") + 1, col("j")) + val horizontalEdges = coordinates + .filter(col("j") =!= n - 1) + .select(vIDcol.as("src"), rightIDcol.as("dst")) + val verticalEdges = coordinates + .filter(col("i") =!= n - 1) + .select(vIDcol.as("src"), downIDcol.as("dst")) + val allEdges = horizontalEdges.union(verticalEdges) + // Add random parameters from a normal distribution + val edges = + allEdges.withColumn("b", randn(seed.toLong + 1L) * eStd) // Ising parameter for edge + + // Create the GraphFrame + val g = GraphFrame(vertices, edges) + + // Materialize graph as workaround for SPARK-13333 + g.vertices.cache().count() + g.edges.cache().count() + + g + } + + /** Version of `gridIsingModel` with vStd, eStd set to 1.0. */ + def gridIsingModel(spark: SparkSession, n: Int): GraphFrame = + gridIsingModel(spark, n, 1.0, 1.0) + +} + +/** Example GraphFrames for testing the API */ +object Graphs extends Graphs diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateMessagesSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateMessagesSuite.scala new file mode 100644 index 0000000000000..5ee33cd5530a9 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateMessagesSuite.scala @@ -0,0 +1,180 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types._ +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.TestUtils +import org.apache.spark.graphframes.examples.Graphs + +import scala.collection.mutable + +class AggregateMessagesSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + test("aggregateMessages") { + val AM = AggregateMessages + val g = Graphs.friends + // For each user, sum the ages of the adjacent users, + // plus 1 for the src's sum if the edge is "friend". + val msgToSrc = AM.dst("age") + + when(AM.edge("relationship") === "friend", lit(1)).otherwise(0) + val msgToDst = AM.src("age") + val agg = g.aggregateMessages + .sendToSrc(msgToSrc) + .sendToDst(msgToDst) + .agg(sum(AM.msg).as("summedAges")) + // Convert agg to a Map. + import org.apache.spark.sql._ + val aggMap: Map[String, Long] = agg + .select("id", "summedAges") + .collect() + .map { + case Row(id: String, s: Long) => + id -> s + case _: Row => throw new GraphFramesUnreachableException() + } + .toMap + agg.unpersist() + // Compute the truth via brute force for comparison. + val trueAgg: Map[String, Int] = { + val user2age = g.vertices + .select("id", "age") + .collect() + .map { + case Row(id: String, age: Int) => + id -> age + case _: Row => throw new GraphFramesUnreachableException() + } + .toMap + val a = mutable.HashMap.empty[String, Int] + g.edges.select("src", "dst", "relationship").collect().foreach { + case Row(src: String, dst: String, relationship: String) => + a.put( + src, + a.getOrElse(src, 0) + user2age(dst) + (if (relationship == "friend") 1 else 0)) + a.put(dst, a.getOrElse(dst, 0) + user2age(src)) + case _ => throw new GraphFramesUnreachableException() + } + a.toMap + } + // Compare to the true values. + aggMap.keys.foreach { case user => + assert(aggMap(user) === trueAgg(user), s"Failure on user $user") + } + // Perform the aggregation again, this time providing the messages as Strings instead. + val msgToSrc2 = "(dst['age'] + CASE WHEN (edge['relationship'] = 'friend') THEN 1 ELSE 0 END)" + val msgToDst2 = "src['age']" + val agg2 = g.aggregateMessages + .sendToSrc(msgToSrc2) + .sendToDst(msgToDst2) + .agg("sum(MSG) AS `summedAges`") + // Convert agg2 to a Map. + val agg2Map: Map[String, Long] = agg2 + .select("id", "summedAges") + .collect() + .map { + case Row(id: String, s: Long) => + id -> s + case _: Row => throw new GraphFramesUnreachableException() + } + .toMap + agg2.unpersist() + // Compare to the true values. + agg2Map.keys.foreach { case user => + assert(agg2Map(user) === trueAgg(user), s"Failure on user $user") + } + } + + test("aggregateMessages with multiple message and aggregation columns") { + val AM = AggregateMessages + val vertices = + sqlContext.createDataFrame(List((1, 30), (2, 40), (3, 50), (4, 60))).toDF("id", "att1") + val edges = + sqlContext.createDataFrame(List((1, 2, 4), (2, 3, 5), (1, 4, 6))).toDF("src", "dst", "att2") + val expectedValues = Map( + 1 -> Tuple2(100L, 5.0), + 2 -> Tuple2(80L, 4.5), + 3 -> Tuple2(40L, 5.0), + 4 -> Tuple2(30L, 6.0)) + + val g = GraphFrame(vertices, edges) + // aggregateMessages with column aliases + val agg = g.aggregateMessages + .sendToDst(AM.src("att1").as("att1"), AM.edge("att2").as("att2")) + .sendToSrc(AM.dst("att1").as("att1"), AM.edge("att2").as("att2")) + .agg(sum(AM.msg("att1")).as("sum_att1"), avg(AM.msg("att2")).as("avg_att2")) + + // aggregateMessages with columns and no aliases + val agg2 = g.aggregateMessages + .sendToDst(AM.src("att1"), AM.edge("att2")) + .sendToSrc(AM.dst("att1"), AM.edge("att2")) + .agg(sum(AM.msg("att1")).as("sum_att1"), avg(AM.msg("att2")).as("avg_att2")) + + // aggregateMessages with column expressions + val agg3 = g.aggregateMessages + .sendToDst("src['att1'] as att1", "edge['att2'] as att2") + .sendToSrc("dst['att1'] as att1", "edge['att2'] as att2") + .agg("sum(MSG['att1']) AS sum_att1", "avg(MSG['att2']) AS avg_att2") + + // validate schema + assert(agg.schema.size === 3) + TestUtils.checkColumnType(agg.schema, "id", IntegerType) + TestUtils.checkColumnType(agg.schema, "sum_att1", LongType) + TestUtils.checkColumnType(agg.schema, "avg_att2", DoubleType) + + assert(agg.schema === agg2.schema) + assert(agg.schema === agg3.schema) + + // validate content + val output1 = agg + .collect() + .map { + case Row(id: Int, sumAtt1: Long, avgAtt2: Double) => + id -> Tuple2(sumAtt1, avgAtt2) + case _ => throw new GraphFramesUnreachableException() + } + .toMap + val output2 = agg2 + .collect() + .map { + case Row(id: Int, sumAtt1: Long, avgAtt2: Double) => + id -> Tuple2(sumAtt1, avgAtt2) + case _ => throw new GraphFramesUnreachableException() + } + .toMap + val output3 = agg3 + .collect() + .map { + case Row(id: Int, sumAtt1: Long, avgAtt2: Double) => + id -> Tuple2(sumAtt1, avgAtt2) + case _ => throw new GraphFramesUnreachableException() + } + .toMap + assert(output1 === expectedValues) + assert(output2 === expectedValues) + assert(output3 === expectedValues) + agg.unpersist() + agg2.unpersist() + agg3.unpersist() + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateNeighborsSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateNeighborsSuite.scala new file mode 100644 index 0000000000000..ccaa45d96ef0b --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateNeighborsSuite.scala @@ -0,0 +1,463 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.functions._ +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite + +class AggregateNeighborsSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + test("find all paths between two vertices using AggregateNeighbors") { + // Create a simple graph: 1 -> 2 -> 3 -> 4, and also 1 -> 3 + val vertices = + spark.createDataFrame(Seq((1L, "A"), (2L, "B"), (3L, "C"), (4L, "D"))).toDF("id", "name") + + val edges = spark + .createDataFrame( + Seq((1L, 2L, "edge1"), (2L, 3L, "edge2"), (3L, 4L, "edge3"), (1L, 3L, "edge4"))) + .toDF("src", "dst", "edgeAttr") + + val graph = GraphFrame(vertices, edges) + + // We want to find all paths from vertex 1 to vertex 4 + val sourceId = 1L + val targetId = 4L + + // Use AggregateNeighbors to find paths + // We'll track paths as strings in an accumulator + val agg = graph.aggregateNeighbors + .setStartingVertices(col("id") === sourceId) + .setMaxHops(5) // Enough to reach vertex 4 + .setTargetCondition( + AggregateNeighbors.dstAttr("id") === lit(targetId) + ) // Using dst_id from the internal join + .addAccumulator( + "path", + // Initialize accumulator: start with source vertex ID + lit(sourceId.toString), + // Update accumulator: append current destination vertex ID + concat(col("path"), lit("->"), col("dst_id").cast("string"))) + .setRequiredVertexAttributes(Seq("id", "name")) + .setRequiredEdgeAttributes(Seq("edgeAttr")) + .run() + + // The result should contain paths from source to target + // Let's collect and analyze the results + val results = agg.collect() + + // Expected paths: + // Path 1: 1 -> 2 -> 3 -> 4 + // Path 2: 1 -> 3 -> 4 + val expectedPaths = Set(s"$sourceId->2->3->$targetId", s"$sourceId->3->$targetId") + + // Extract paths from results + val actualPaths = results.map { row => + // The path accumulator should be in the result + row.getAs[String]("path") + }.toSet + + // Verify we found the correct number of paths + assert(results.length === 2, s"Expected 2 paths but found ${results.length}") + + // Verify each expected path is present + expectedPaths.foreach { expectedPath => + assert( + actualPaths.contains(expectedPath), + s"Path $expectedPath not found in results: $actualPaths") + } + + // Also verify that each result has the correct target vertex ID + results.foreach { row => + val id = row.getAs[Long]("id") + assert(id === targetId, s"Result should have target vertex ID $targetId, but found $id") + } + } + + test("AggregateNeighbors with multiple accumulators") { + // Create a simple graph + val vertices = spark.createDataFrame(Seq((1L, 10), (2L, 20), (3L, 30))).toDF("id", "value") + + val edges = spark + .createDataFrame(Seq((1L, 2L, 5.0), (2L, 3L, 6.0), (1L, 3L, 7.0))) + .toDF("src", "dst", "weight") + + val graph = GraphFrame(vertices, edges) + + // Test with multiple accumulators: sum of values and product of weights + val sourceId = 1L + val targetId = 3L + + val agg = graph.aggregateNeighbors + .setStartingVertices(col("id") === sourceId) + .setMaxHops(3) + .setTargetCondition(AggregateNeighbors.dstAttr("id") === lit(targetId)) + .addAccumulator( + "sum_values", + lit(0L), + col("sum_values") + AggregateNeighbors.dstAttr("value")) + .addAccumulator( + "product_weights", + lit(1.0), + col("product_weights") * AggregateNeighbors.edgeAttr("weight")) + .setRequiredVertexAttributes(Seq("id", "value")) + .setRequiredEdgeAttributes(Seq("weight")) + .run() + + val results = agg.collect() + + // There should be 2 paths from 1 to 3 + assert(results.length === 2) + + // Sort results by product_weights for consistent checking + val sortedResults = results.sortBy(row => row.getAs[Double]("product_weights")) + + // Path 1: 1 -> 3 directly + val directPath = sortedResults(0) + assert(directPath.getAs[Long]("sum_values") === 30L) // Only vertex 3's value (30) + assert(math.abs(directPath.getAs[Double]("product_weights") - 7.0) < 0.001) + + // Path 2: 1 -> 2 -> 3 + val indirectPath = sortedResults(1) + assert(indirectPath.getAs[Long]("sum_values") === 50L) // 20 + 30 = 50 + assert( + math.abs(indirectPath.getAs[Double]("product_weights") - 30.0) < 0.001 + ) // 5.0 * 6.0 = 30.0 + } + + test("AggregateNeighbors with stopping condition") { + val vertices = spark.createDataFrame(Seq((1L, "A"), (2L, "B"), (3L, "C"))).toDF("id", "name") + + val edges = spark + .createDataFrame(Seq((1L, 2L), (2L, 3L), (3L, 1L), (1L, 3L), (2L, 1L))) + .toDF("src", "dst") + + val graph = GraphFrame(vertices, edges) + + // We want to stop when we revisit a vertex (to avoid infinite loops) + // We'll track visited vertices in an accumulator + val sourceId = 1L + + val agg = graph.aggregateNeighbors + .setStartingVertices(col("id") === sourceId) + .setMaxHops(10) + .setStoppingCondition( + // Stop when we've already visited the destination vertex + array_contains(col("visited_vertices"), AggregateNeighbors.dstAttr("id"))) + .addAccumulator( + "visited_vertices", + array(lit(sourceId)), + // Add the current destination vertex to the visited list + array_append(col("visited_vertices"), AggregateNeighbors.dstAttr("id"))) + .addAccumulator("path_length", lit(0L), col("path_length") + lit(1L)) + .setRequiredVertexAttributes(Seq("id")) + .run() + + val results = agg.collect() + + // We should get results for each vertex reachable without cycles + // Since we stop when revisiting a vertex, we should get: + // 1 - 3 - 1 + // 1 - 2 - 1 + // 1 - 2 - 3 - 1 + // Count results + assert(results.length === 3) + + // Check path lengths + assert(results.map(_.getAs[Long]("path_length")).toSet == Set(2, 3)) + assert(results.forall(r => r.getAs[Long]("path_length") == r.getAs[Int]("hop"))) + } + + test("AggregateNeighbors with edge filter") { + // Create a graph with different types of edges + val vertices = + spark.createDataFrame(Seq((1L, "A"), (2L, "B"), (3L, "C"), (4L, "D"))).toDF("id", "name") + + val edges = spark + .createDataFrame( + Seq( + (1L, 2L, "allowed"), + (2L, 3L, "allowed"), + (3L, 4L, "allowed"), + (1L, 3L, "blocked"), + (2L, 4L, "blocked"))) + .toDF("src", "dst", "type") + + val graph = GraphFrame(vertices, edges) + + // Only traverse edges with type "allowed" + val sourceId = 1L + val targetId = 4L + + val agg = graph.aggregateNeighbors + .setStartingVertices(col("id") === sourceId) + .setMaxHops(5) + .setTargetCondition(AggregateNeighbors.dstAttr("id") === lit(targetId)) + .setEdgeFilter(AggregateNeighbors.edgeAttr("type") === lit("allowed")) + .addAccumulator( + "path", + lit(sourceId.toString), + concat(col("path"), lit("->"), AggregateNeighbors.dstAttr("id").cast("string"))) + .setRequiredEdgeAttributes(Seq("type")) + .setRequiredVertexAttributes(Seq("id")) + .run() + + val results = agg.collect() + + // With the edge filter, only allowed edges should be traversed + // Allowed path: 1 -> 2 -> 3 -> 4 + // The direct path 1 -> 3 uses a "blocked" edge, so it shouldn't be found + // The path 1 -> 2 -> 4 uses a "blocked" edge for the last hop + + // So only one path should be found + assert(results.length === 1) + + val actualPath = results(0).getAs[String]("path") + val expectedPath = s"$sourceId->2->3->$targetId" + assert(actualPath === expectedPath) + } + + test("AggregateNeighbors with empty graph") { + // Create an empty graph with no vertices and no edges + val vertices = spark.createDataFrame(Seq.empty[(Long, String)]).toDF("id", "name") + val edges = spark.createDataFrame(Seq.empty[(Long, Long)]).toDF("src", "dst") + + val graph = GraphFrame(vertices, edges) + + // Run aggregate neighbors with arbitrary starting vertex + val agg = graph.aggregateNeighbors + .setStartingVertices(col("id") === 1L) + .setStoppingCondition(lit(false)) + .setMaxHops(3) + .addAccumulator("count", lit(0L), col("count") + lit(1L)) + .run() + + val results = agg.collect() + + // Should return empty results without errors + assert(results.length === 0) + } + + test("AggregateNeighbors with disconnected vertices") { + // Create a graph with disconnected components + // Component 1: 1 -> 2 + // Component 2: 3 -> 4 (disconnected from component 1) + val vertices = spark + .createDataFrame(Seq((1L, "A"), (2L, "B"), (3L, "C"), (4L, "D"))) + .toDF("id", "name") + + val edges = spark + .createDataFrame(Seq((1L, 2L, "edge1"), (3L, 4L, "edge2"))) + .toDF("src", "dst", "edgeAttr") + + val graph = GraphFrame(vertices, edges) + + // Start from vertex 1, try to reach vertex 4 (disconnected) + val sourceId = 1L + val targetId = 4L + + val agg = graph.aggregateNeighbors + .setStartingVertices(col("id") === sourceId) + .setMaxHops(5) + .setTargetCondition(AggregateNeighbors.dstAttr("id") === lit(targetId)) + .addAccumulator( + "path", + lit(sourceId.toString), + concat(col("path"), lit("->"), AggregateNeighbors.dstAttr("id").cast("string"))) + .setRequiredVertexAttributes(Seq("id")) + .setRequiredEdgeAttributes(Seq("edgeAttr")) + .run() + + val results = agg.collect() + + // No path should exist to disconnected vertex 4 + assert(results.length === 0) + + // Now test within connected component + val agg2 = graph.aggregateNeighbors + .setStartingVertices(col("id") === sourceId) + .setMaxHops(5) + .setTargetCondition(AggregateNeighbors.dstAttr("id") === lit(2L)) + .addAccumulator( + "path", + lit(sourceId.toString), + concat(col("path"), lit("->"), AggregateNeighbors.dstAttr("id").cast("string"))) + .setRequiredVertexAttributes(Seq("id")) + .setRequiredEdgeAttributes(Seq("edgeAttr")) + .run() + + val results2 = agg2.collect() + + // Path should exist within connected component + assert(results2.length === 1) + assert(results2(0).getAs[String]("path") === s"$sourceId->2") + } + + test("AggregateNeighbors with self-loops") { + // Create a graph with a self-loop + val vertices = + spark.createDataFrame(Seq((1L, "A"), (2L, "B"), (3L, "C"))).toDF("id", "name") + + val edges = spark + .createDataFrame( + Seq( + (1L, 2L, "edge1"), + (2L, 2L, "self-loop"), // Self-loop on vertex 2 + (2L, 3L, "edge2"))) + .toDF("src", "dst", "edgeAttr") + + val graph = GraphFrame(vertices, edges) + + val sourceId = 1L + val targetId = 3L + + // Test without stopping condition to see self-loop behavior + val agg = graph.aggregateNeighbors + .setStartingVertices(col("id") === sourceId) + .setMaxHops(3) + .setTargetCondition(AggregateNeighbors.dstAttr("id") === lit(targetId)) + .addAccumulator( + "path", + lit(sourceId.toString), + concat(col("path"), lit("->"), AggregateNeighbors.dstAttr("id").cast("string"))) + .setRequiredVertexAttributes(Seq("id")) + .setRequiredEdgeAttributes(Seq("edgeAttr")) + .run() + + val results = agg.collect() + + // Should find at least one path to target + assert(results.length >= 1) + + // Verify direct path exists + val paths = results.map(_.getAs[String]("path")).toSet + assert(paths.contains(s"$sourceId->2->3")) + } + + test("AggregateNeighbors with multiple edge types") { + // Create a graph with multiple edge types + val vertices = spark + .createDataFrame(Seq((1L, "A"), (2L, "B"), (3L, "C"), (4L, "D"))) + .toDF("id", "name") + + val edges = spark + .createDataFrame( + Seq( + (1L, 2L, "friend"), + (2L, 3L, "colleague"), + (3L, 4L, "friend"), + (1L, 3L, "colleague"), + (2L, 4L, "friend"))) + .toDF("src", "dst", "edgeType") + + val graph = GraphFrame(vertices, edges) + + val sourceId = 1L + val targetId = 4L + + // Test with edge type filter for "friend" edges only + val agg = graph.aggregateNeighbors + .setStartingVertices(col("id") === sourceId) + .setMaxHops(5) + .setTargetCondition(AggregateNeighbors.dstAttr("id") === lit(targetId)) + .setEdgeFilter(AggregateNeighbors.edgeAttr("edgeType") === lit("friend")) + .addAccumulator( + "path", + lit(sourceId.toString), + concat(col("path"), lit("->"), AggregateNeighbors.dstAttr("id").cast("string"))) + .setRequiredVertexAttributes(Seq("id")) + .setRequiredEdgeAttributes(Seq("edgeType")) + .run() + + val results = agg.collect() + + // With friend-only filter, should find paths using only friend edges + // Path: 1 -> 2 -> 4 (all friend edges) + assert(results.length === 1) + assert(results(0).getAs[String]("path") === s"$sourceId->2->$targetId") + + // Test without filter - should find all paths + val agg2 = graph.aggregateNeighbors + .setStartingVertices(col("id") === sourceId) + .setMaxHops(5) + .setTargetCondition(AggregateNeighbors.dstAttr("id") === lit(targetId)) + .addAccumulator( + "path", + lit(sourceId.toString), + concat(col("path"), lit("->"), AggregateNeighbors.dstAttr("id").cast("string"))) + .setRequiredVertexAttributes(Seq("id")) + .setRequiredEdgeAttributes(Seq("edgeType")) + .run() + + val results2 = agg2.collect() + + // Should find multiple paths without filtering + assert(results2.length > 1) + } + + test("AggregateNeighbors with large vertex degrees") { + // Create a star graph: center vertex (1) connected to many leaf vertices + val numLeaves = 100 + + // Center vertex + val centerVertex = (1L, "center") + + // Leaf vertices (2 to numLeaves+1) + val leafVertices = (2L to (numLeaves + 1).toLong).map(id => (id, s"leaf_$id")) + + val vertices = spark + .createDataFrame(centerVertex +: leafVertices) + .toDF("id", "name") + + // Edges from center to all leaves + val edges = leafVertices + .map { case (leafId, _) => + (1L, leafId, "connects") + } + + val edgesDF = spark + .createDataFrame(edges) + .toDF("src", "dst", "edgeType") + + val graph = GraphFrame(vertices, edgesDF) + + val sourceId = 1L + + // Test aggregation from high-degree vertex + val agg = graph.aggregateNeighbors + .setStartingVertices(col("id") === sourceId) + .setMaxHops(2) + .addAccumulator("count", lit(0L), col("count") + lit(1L)) + .setTargetCondition( + AggregateNeighbors.dstAttr("id").isInCollection((2L to (numLeaves + 1).toLong).toSeq)) + .setRequiredVertexAttributes(Seq("id")) + .setRequiredEdgeAttributes(Seq("edgeType")) + .run() + + val results = agg.collect() + + // Should find all leaf vertices without performance issues + assert(results.length === numLeaves) + + // Verify all accumulators are correctly computed + results.foreach { row => + assert(row.getAs[Long]("count") === 1L) + } + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AllPathsSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AllPathsSuite.scala new file mode 100644 index 0000000000000..63abcbfdad7cd --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AllPathsSuite.scala @@ -0,0 +1,151 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.types.ArrayType +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.types.LongType +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite + +class AllPathsSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + test("directed: enumerate all simple paths") { + val vertices = + spark.createDataFrame(Seq((1L, "A"), (2L, "B"), (3L, "C"), (4L, "D"))).toDF("id", "name") + val edges = + spark + .createDataFrame(Seq((1L, 2L, "x"), (2L, 4L, "x"), (1L, 3L, "x"), (3L, 4L, "x"))) + .toDF("src", "dst", "label") + val g = GraphFrame(vertices, edges) + + val result = + g.allPaths.fromExpr(col("id") === 1L).toExpr(col("id") === 4L).maxPathLength(3).run() + + val actual = result + .collect() + .map { row => + row.getAs[Seq[Long]]("path") -> row.getAs[Long]("len") + } + .toSet + + val expected = Set((Seq(1L, 2L, 4L), 2L), (Seq(1L, 3L, 4L), 2L)) + assert(actual === expected) + } + + test("undirected: traverse both edge directions") { + val vertices = + spark.createDataFrame(Seq((1L, "A"), (2L, "B"), (3L, "C"), (4L, "D"))).toDF("id", "name") + val edges = + spark + .createDataFrame(Seq((1L, 2L, "x"), (2L, 4L, "x"), (1L, 3L, "x"), (3L, 4L, "x"))) + .toDF("src", "dst", "label") + val g = GraphFrame(vertices, edges) + + val result = g.allPaths + .fromExpr(col("id") === 4L) + .toExpr(col("id") === 1L) + .setIsDirected(false) + .maxPathLength(3) + .run() + + val actual = result + .collect() + .map { row => + row.getAs[Seq[Long]]("path") -> row.getAs[Long]("len") + } + .toSet + + val expected = Set((Seq(4L, 2L, 1L), 2L), (Seq(4L, 3L, 1L), 2L)) + assert(actual === expected) + } + + test("edge filter excludes blocked edges") { + val vertices = + spark.createDataFrame(Seq((1L, "A"), (2L, "B"), (3L, "C"), (4L, "D"))).toDF("id", "name") + val edges = spark + .createDataFrame( + Seq((1L, 2L, "allowed"), (2L, 4L, "allowed"), (1L, 3L, "blocked"), (3L, 4L, "allowed"))) + .toDF("src", "dst", "label") + val g = GraphFrame(vertices, edges) + + val result = g.allPaths + .fromExpr("id = 1") + .toExpr("id = 4") + .edgeFilter(col("label") =!= "blocked") + .maxPathLength(3) + .run() + + val rows = result.collect() + assert(rows.length === 1) + assert(rows.head.getAs[Seq[Long]]("path") === Seq(1L, 2L, 4L)) + assert(rows.head.getAs[Long]("len") === 2L) + } + + test("cycle handling keeps paths simple") { + val vertices = spark.createDataFrame(Seq((1L, "A"), (2L, "B"), (3L, "C"))).toDF("id", "name") + val edges = spark + .createDataFrame(Seq((1L, 2L), (2L, 1L), (2L, 3L), (1L, 3L))) + .toDF("src", "dst") + val g = GraphFrame(vertices, edges) + + val result = g.allPaths.fromExpr("id = 1").toExpr("id = 3").maxPathLength(4).run() + + val actualPaths = + result.collect().map(_.getAs[scala.collection.Seq[Long]]("path").toList).toSet + assert(actualPaths === Set(List(1L, 3L), List(1L, 2L, 3L))) + assert(!actualPaths.contains(List(1L, 2L, 1L, 3L))) + } + + test("schema and required arguments") { + val vertices = spark.createDataFrame(Seq((1L, "A"), (2L, "B"))).toDF("id", "name") + val edges = spark.createDataFrame(Seq((1L, 2L))).toDF("src", "dst") + val g = GraphFrame(vertices, edges) + + val result = g.allPaths.fromExpr("id = 1").toExpr("id = 2").run() + + assert(result.columns.toSeq === Seq("path", "len")) + assert(result.schema("path").dataType.asInstanceOf[ArrayType].elementType === LongType) + assert(result.schema("len").dataType === IntegerType) + + intercept[IllegalArgumentException] { + g.allPaths.run() + } + intercept[IllegalArgumentException] { + g.allPaths.fromExpr("id = 1").run() + } + intercept[IllegalArgumentException] { + g.allPaths.toExpr("id = 2").run() + } + intercept[IllegalArgumentException] { + g.allPaths.fromExpr("id = 1").toExpr("id = 2").maxPathLength(0) + } + } + + test("no matching path returns empty dataframe") { + val vertices = spark.createDataFrame(Seq((1L, "A"), (2L, "B"), (3L, "C"))).toDF("id", "name") + val edges = spark.createDataFrame(Seq((1L, 2L))).toDF("src", "dst") + val g = GraphFrame(vertices, edges) + + val result = + g.allPaths.fromExpr(col("id") === 1L).toExpr(col("id") === 3L).maxPathLength(4).run() + assert(result.collect().isEmpty) + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/BFSSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/BFSSuite.scala new file mode 100644 index 0000000000000..c67204fffe70f --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/BFSSuite.scala @@ -0,0 +1,182 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions.col +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.graphframes.SparkFunSuite + +class BFSSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + // First graph uses String IDs + @transient var v: DataFrame = _ + @transient var e: DataFrame = _ + @transient var g: GraphFrame = _ + + // Second graph uses Int IDs + @transient var v2: DataFrame = _ + @transient var e2: DataFrame = _ + @transient var g2: GraphFrame = _ + + override def beforeAll(): Unit = { + super.beforeAll() + + /* + a -> b <-> c + ^ ^ + | | + d <- e --> f Also, self-edge for f + */ + v = spark + .createDataFrame( + List(("a", "f"), ("b", "f"), ("c", "m"), ("d", "f"), ("e", "m"), ("f", "m"))) + .toDF("id", "gender") + e = spark + .createDataFrame( + List( + ("a", "b", "friend"), + ("b", "c", "follow"), + ("c", "b", "follow"), + ("f", "c", "follow"), + ("e", "f", "follow"), + ("e", "d", "friend"), + ("d", "a", "friend"), + ("f", "f", "self"))) + .toDF("src", "dst", "relationship") + g = GraphFrame(v, e) + + v2 = + spark.createDataFrame(List((0L, "f"), (1L, "m"), (2L, "m"), (3L, "f"))).toDF("id", "gender") + e2 = spark.createDataFrame(List((0L, 1L), (1L, 2L), (2L, 3L), (2L, 0L))).toDF("src", "dst") + g2 = GraphFrame(v2, e2) + } + + override def afterAll(): Unit = { + v = null + e = null + g = null + v2 = null + e2 = null + g2 = null + super.afterAll() + } + + test("unmatched queries should return nothing") { + val badStart = g.bfs.fromExpr(col("id") === "howdy").toExpr(col("id") === "a").run() + assert(badStart.count() === 0) + val badEnd = g.bfs.fromExpr(col("id") === "a").toExpr(col("id") === "howdy").run() + assert(badEnd.count() === 0) + } + + test("0 hops, aka from=to") { + val paths = g.bfs.fromExpr(col("id") === "a").toExpr(col("id") === "a").run() + assert(paths.count() === 1) + assert(paths.columns === Seq("from", "to")) + assert(paths.select("from.id").head().getString(0) === "a") + assert(paths.select("to.id").head().getString(0) === "a") + } + + test("1 hop, aka single edge paths") { + val paths = g.bfs.fromExpr(col("id") === "a").toExpr(col("id") === "b").run() + assert(paths.count() === 1) + assert(paths.columns === Seq("from", "e0", "to")) + assert(paths.select("from.id", "to.id").head() === Row("a", "b")) + } + + test("ties") { + val paths = g.bfs.fromExpr(col("id") === "e").toExpr(col("id") === "b").run() + assert(paths.count() === 2) + assert(paths.columns === Seq("from", "e0", "v1", "e1", "v2", "e2", "to")) + paths.select("to.id").collect().foreach { + case Row(id: String) => + assert(id === "b") + case _ => throw new GraphFramesUnreachableException() + } + } + + test("maxPathLength: length 1") { + val paths = g.bfs.fromExpr(col("id") === "e").toExpr(col("id") === "f").maxPathLength(1).run() + assert(paths.count() === 1) + val paths0 = + g.bfs.fromExpr(col("id") === "e").toExpr(col("id") === "f").maxPathLength(0).run() + assert(paths0.count() === 0) + } + + test("maxPathLength: length > 1") { + val paths = g.bfs.fromExpr(col("id") === "e").toExpr(col("id") === "b").maxPathLength(3).run() + assert(paths.count() === 2) + val paths0 = + g.bfs.fromExpr(col("id") === "e").toExpr(col("id") === "b").maxPathLength(2).run() + assert(paths0.count() === 0) + } + + test("edge filter") { + val paths1 = g.bfs + .fromExpr(col("id") === "e") + .toExpr(col("id") === "b") + .edgeFilter(col("src") =!= "d") + .run() + assert(paths1.count() === 1) + paths1.select("e0.dst").collect().foreach { + case Row(id: String) => + assert(id === "f") + case _: Row => throw new GraphFramesUnreachableException() + } + val paths2 = g.bfs + .fromExpr(col("id") === "e") + .toExpr(col("id") === "b") + .edgeFilter(col("relationship") === "friend") + .run() + assert(paths2.count() === 1) + paths2.select("e0.dst").collect().foreach { + case Row(id: String) => + assert(id === "d") + case _: Row => throw new GraphFramesUnreachableException() + } + } + + test("string expressions") { + val paths1 = g.bfs + .fromExpr("id = 'e'") + .toExpr("id = 'b'") + .edgeFilter("src != 'd'") + .run() + assert(paths1.count() === 1) + paths1.select("e0.dst").collect().foreach { + case Row(id: String) => + assert(id === "f") + case _: Row => throw new GraphFramesUnreachableException() + } + } + + test("fromExpr and toExpr are required") { + intercept[IllegalArgumentException] { + g.bfs.run() + } + intercept[IllegalArgumentException] { + g.bfs.fromExpr("id = 'e'").run() + } + intercept[IllegalArgumentException] { + g.bfs.toExpr("id = 'b'").run() + } + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ConnectedComponentsSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ConnectedComponentsSuite.scala new file mode 100644 index 0000000000000..78f9c1602b8ea --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ConnectedComponentsSuite.scala @@ -0,0 +1,419 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.types.DataTypes +import org.apache.spark.sql.types.LongType +import org.apache.spark.storage.StorageLevel +import org.apache.spark.graphframes._ +import org.apache.spark.graphframes.GraphFrame._ +import org.apache.spark.graphframes.examples.Graphs + +import scala.reflect.ClassTag +import scala.reflect.runtime.universe.TypeTag + +class ConnectedComponentsSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + // vertices and edges for pruning node optimization tests. + var verticesOpt: DataFrame = _ + var edgesOpt: DataFrame = _ + + override def beforeAll(): Unit = { + super.beforeAll() + verticesOpt = spark.range(7L).toDF(ID) + edgesOpt = spark + .createDataFrame(Seq((0L, 1L), (0L, 2L), (0L, 3L), (0L, 4L), (1L, 2L), (1L, 5L))) + .toDF(SRC, DST) + } + + test("default params") { + val g = Graphs.empty[Int] + val cc = g.connectedComponents + // That is OK! It is just a name-change + assert(cc.getAlgorithm === "two_phase") + assert(cc.getBroadcastThreshold === 1000000) + assert(cc.getCheckpointInterval === 2) + assert(!cc.getUseLabelsAsComponents) + spark.conf.set("spark.graphframes.useLocalCheckpoints", "false") + } + + test("using labels as components") { + spark.conf.set("spark.graphframes.useLabelsAsComponents", "true") + val vertices = + spark.createDataFrame(Seq("a", "b", "c", "d", "e").map(Tuple1.apply)).toDF(ID) + val edges = spark.createDataFrame(Seq.empty[(String, String)]).toDF(SRC, DST) + val g = GraphFrame(vertices, edges) + val components = g.connectedComponents.run() + val expected = Seq("a", "b", "c", "d", "e").map(Set(_)).toSet + assertComponents(components, expected) + components.unpersist() + spark.conf.set("spark.graphframes.useLabelsAsComponents", "false") + } + + test("don't using labels as components") { + val vertices = + spark.createDataFrame(Seq("a", "b", "c", "d", "e").map(Tuple1.apply)).toDF(ID) + val edges = spark.createDataFrame(Seq.empty[(String, String)]).toDF(SRC, DST) + val g = GraphFrame(vertices, edges) + val components = g.connectedComponents.run() + assert(components.schema("component").dataType == LongType) + components.unpersist() + } + + test("friends graph with different broadcast thresholds") { + val friends = Graphs.friends + val expected = Set(Set("a", "b", "c", "d", "e", "f"), Set("g")) + for ((algorithm, broadcastThreshold) <- + Seq( + ("graphx", 1000000), + ("graphframes", 100000), + ("graphframes", 1), + ("graphframes", -1))) { + val components = friends.connectedComponents + .setAlgorithm(algorithm) + .setBroadcastThreshold(broadcastThreshold) + .run() + assertComponents(components, expected) + components.unpersist() + } + } + + Seq(true, false).foreach(useSkewedJoin => { + Seq(true, false).foreach(useLocalCheckpoint => { + val testPostfixName = s"${if (useLocalCheckpoint) " with local checkpoint" + else ""}${if (useSkewedJoin) ", skewed join" else ", AQE join"}" + val broadcastThreshold = if (useSkewedJoin) 1000000 else -1 + + test(s"non trivial graph #873$testPostfixName") { + val edges = spark + .createDataFrame( + Seq((1L, 2L), (2L, 3L), (3L, 4L), (4L, 5L), (5L, 1L), (6L, 7L), (7L, 8L), (8L, 6L))) + .toDF("src", "dst") + + val vertices = + edges.select("src").union(edges.select("dst")).distinct().withColumnRenamed("src", "id") + val g = GraphFrame(vertices, edges) + + val result = g.connectedComponents + .setBroadcastThreshold(broadcastThreshold) + .setUseLocalCheckpoints(useLocalCheckpoint) + .run() + val numComps = result.select("component").distinct().count() + assert(numComps == 2L) + result.unpersist() + } + + test(s"empty graph$testPostfixName") { + for (empty <- Seq(Graphs.empty[Int], Graphs.empty[Long], Graphs.empty[String])) { + val components = empty.connectedComponents + .setBroadcastThreshold(broadcastThreshold) + .setUseLocalCheckpoints(useLocalCheckpoint) + .run() + assert(components.count() === 0L) + components.unpersist() + } + } + + test(s"single vertex$testPostfixName") { + val v = spark.createDataFrame(List((0L, "a", "b"))).toDF("id", "vattr", "gender") + // Create an empty dataframe with the proper columns. + val e = spark + .createDataFrame(List((0L, 0L, 1L))) + .toDF("src", "dst", "test") + .filter("src > 10") + val g = GraphFrame(v, e) + val comps = g.connectedComponents + .setBroadcastThreshold(broadcastThreshold) + .setUseLocalCheckpoints(useLocalCheckpoint) + .run() + TestUtils.testSchemaInvariants(g, comps) + TestUtils.checkColumnType(comps.schema, "component", DataTypes.LongType) + assert(comps.count() === 1) + assert( + comps.select("id", "component", "vattr", "gender").collect() + === Seq(Row(0L, 0L, "a", "b"))) + comps.unpersist() + } + + test(s"disconnected vertices$testPostfixName") { + val n = 5L + val vertices = spark.range(n).toDF(ID) + val edges = spark.createDataFrame(Seq.empty[(Long, Long)]).toDF(SRC, DST) + val g = GraphFrame(vertices, edges) + val components = g.connectedComponents + .setUseLocalCheckpoints(useLocalCheckpoint) + .setBroadcastThreshold(broadcastThreshold) + .run() + val expected = (0L until n).map(Set(_)).toSet + assertComponents(components, expected) + components.unpersist() + } + + test(s"two connected vertices$testPostfixName") { + val v = + spark.createDataFrame(List((0L, "a0", "b0"), (1L, "a1", "b1"))).toDF("id", "A", "B") + val e = spark.createDataFrame(List((0L, 1L, "a01", "b01"))).toDF("src", "dst", "A", "B") + val g = GraphFrame(v, e) + val comps = g.connectedComponents + .setUseLocalCheckpoints(useLocalCheckpoint) + .setBroadcastThreshold(broadcastThreshold) + .run() + TestUtils.testSchemaInvariants(g, comps) + assert(comps.count() === 2) + val vxs = comps.sort("id").select("id", "component", "A", "B").collect() + assert(List(Row(0L, 0L, "a0", "b0"), Row(1L, 0L, "a1", "b1")) === vxs) + comps.unpersist() + } + + test(s"chain graph$testPostfixName") { + val n = 5L + val g = Graphs.chain(5L) + val components = g.connectedComponents + .setUseLocalCheckpoints(useLocalCheckpoint) + .setBroadcastThreshold(broadcastThreshold) + .run() + val expected = Set((0L until n).toSet) + assertComponents(components, expected) + components.unpersist() + } + + test(s"star graph$testPostfixName") { + val n = 5L + val g = Graphs.star(5L) + val components = g.connectedComponents + .setUseLocalCheckpoints(useLocalCheckpoint) + .setBroadcastThreshold(broadcastThreshold) + .run() + val expected = Set((0L to n).toSet) + assertComponents(components, expected) + components.unpersist() + } + + test(s"two blobs$testPostfixName") { + val n = 5L + val g = Graphs.twoBlobs(n.toInt) + val components = g.connectedComponents + .setUseLocalCheckpoints(useLocalCheckpoint) + .setBroadcastThreshold(broadcastThreshold) + .run() + val expected = Set((0L until 2 * n).toSet) + assertComponents(components, expected) + components.unpersist() + } + + test(s"two components$testPostfixName") { + val vertices = spark.range(6L).toDF(ID) + val edges = spark + .createDataFrame(Seq((0L, 1L), (1L, 2L), (2L, 0L), (3L, 4L), (4L, 5L), (5L, 3L))) + .toDF(SRC, DST) + val g = GraphFrame(vertices, edges) + val components = g.connectedComponents + .setUseLocalCheckpoints(useLocalCheckpoint) + .setBroadcastThreshold(broadcastThreshold) + .run() + val expected = Set(Set(0L, 1L, 2L), Set(3L, 4L, 5L)) + assertComponents(components, expected) + components.unpersist() + } + + test(s"one component, differing edge directions$testPostfixName") { + val vertices = spark.range(5L).toDF(ID) + val edges = spark + .createDataFrame( + Seq( + // 0 -> 4 -> 3 <- 2 -> 1 + (0L, 4L), + (4L, 3L), + (2L, 3L), + (2L, 1L))) + .toDF(SRC, DST) + val g = GraphFrame(vertices, edges) + val components = g.connectedComponents + .setUseLocalCheckpoints(useLocalCheckpoint) + .setBroadcastThreshold(broadcastThreshold) + .run() + val expected = Set((0L to 4L).toSet) + assertComponents(components, expected) + components.unpersist() + } + + test(s"two components and two dangling vertices$testPostfixName") { + val vertices = spark.range(8L).toDF(ID) + val edges = spark + .createDataFrame(Seq((0L, 1L), (1L, 2L), (2L, 0L), (3L, 4L), (4L, 5L), (5L, 3L))) + .toDF(SRC, DST) + val g = GraphFrame(vertices, edges) + val components = g.connectedComponents + .setUseLocalCheckpoints(useLocalCheckpoint) + .setBroadcastThreshold(broadcastThreshold) + .run() + val expected = Set(Set(0L, 1L, 2L), Set(3L, 4L, 5L), Set(6L), Set(7L)) + assertComponents(components, expected) + components.unpersist() + } + + test(s"really large long IDs$testPostfixName") { + val max = Long.MaxValue + val chain = examples.Graphs.chain(10L) + val vertices = chain.vertices.select((lit(max) - col(ID)).as(ID)) + val edges = + chain.edges.select((lit(max) - col(SRC)).as(SRC), (lit(max) - col(DST)).as(DST)) + val g = GraphFrame(vertices, edges) + val components = g.connectedComponents + .setUseLocalCheckpoints(useLocalCheckpoint) + .setBroadcastThreshold(broadcastThreshold) + .run() + assert(components.count() === 10L) + assert(components.groupBy("component").count().count() === 1L) + components.unpersist() + } + }) + }) + + test("set configuration from spark conf") { + spark.conf.set("spark.graphframes.connectedComponents.algorithm", "GRAPHX") + assert(Graphs.friends.connectedComponents.getAlgorithm == "graphx") + + spark.conf.set("spark.graphframes.connectedComponents.broadcastthreshold", "1000") + assert(Graphs.friends.connectedComponents.getBroadcastThreshold == 1000) + + spark.conf.set("spark.graphframes.connectedComponents.checkpointinterval", "5") + assert(Graphs.friends.connectedComponents.getCheckpointInterval == 5) + + spark.conf + .set("spark.graphframes.connectedComponents.intermediatestoragelevel", "memory_only") + assert( + Graphs.friends.connectedComponents.getIntermediateStorageLevel == StorageLevel.MEMORY_ONLY) + + spark.conf.unset("spark.graphframes.connectedComponents.algorithm") + spark.conf.unset("spark.graphframes.connectedComponents.broadcastthreshold") + spark.conf.unset("spark.graphframes.connectedComponents.checkpointinterval") + spark.conf.unset("spark.graphframes.connectedComponents.intermediatestoragelevel") + } + + Seq(StorageLevel.DISK_ONLY, StorageLevel.MEMORY_ONLY, StorageLevel.NONE).foreach( + storageLevel => { + test(s"intermediate storage level $storageLevel") { + val friends = Graphs.friends + val expected = Set(Set("a", "b", "c", "d", "e", "f"), Set("g")) + + val components = + friends.connectedComponents.setIntermediateStorageLevel(storageLevel).run() + assertComponents(components, expected) + components.unpersist() + () + } + }) + + Seq(StorageLevel.DISK_ONLY, StorageLevel.MEMORY_ONLY).foreach(storageLevel => { + test(s"intermediate storage level without skewedJoin $storageLevel") { + val friends = Graphs.friends + val expected = Set(Set("a", "b", "c", "d", "e", "f"), Set("g")) + + val components = + friends.connectedComponents + .setIntermediateStorageLevel(storageLevel) + .setBroadcastThreshold(-1) + .run() + assertComponents(components, expected) + components.unpersist() + () + } + }) + + test("not leaking cached data") { + val priorCachedDFsSize = spark.sparkContext.getPersistentRDDs.size + + val cc = Graphs.friends.connectedComponents + val components = cc.run() + + components.unpersist(blocking = true) + + assert(spark.sparkContext.getPersistentRDDs.size === priorCachedDFsSize) + } + + test("prune process for pruning nodes optimization") { + val intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK + val shrinkageThreshold = 2.0 + val Some(r1) = TwoPhase.pruneLeafNodes( + edgesOpt, + intermediateStorageLevel, + verticesOpt.count(), + shrinkageThreshold) + + val expectedV = Set(Row(0L), Row(1L), Row(2L)) + val expectedE = Set(Row(0L, 1L), Row(1L, 2L), Row(0L, 2L)) + + assert(r1._1.collect().toSet == expectedV) + assert(r1._2.select(SRC, DST).collect().toSet == expectedE) + assert(r1._3 == expectedV.size) + r1._1.unpersist() + r1._2.unpersist() + } + + test("shrinkage condition for pruning nodes optimization") { + val intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK + val shrinkageThreshold = 4.0 + // new_vv_cnt = 3, nodeNum = 7, shrinkageThreshold = 4 + // new_vv_cnt * shrinkageThreshold > nodeNum. Do not perform the optimization. + val r1 = TwoPhase.pruneLeafNodes( + edgesOpt, + intermediateStorageLevel, + verticesOpt.count(), + shrinkageThreshold) + assert(r1 == None) + } + + test("join back for pruning node optimization") { + val v1 = spark.range(3L).toDF(ID) + val e1 = spark.createDataFrame(Seq((0L, 1L), (0L, 2L))).toDF(SRC, DST) + val r = TwoPhase.joinBack(v1, e1, edgesOpt) + val expectedR = + Set(Row(0L, 0L), Row(0L, 1L), Row(0L, 2L), Row(0L, 3L), Row(0L, 4L), Row(0L, 5L)) + assert(r.collect().toSet == expectedR) + } + + private def assertComponents[T: ClassTag: TypeTag]( + actual: DataFrame, + expected: Set[Set[T]]): Unit = { + import actual.sparkSession.implicits._ + // note: not using agg + collect_list because collect_list is not available in 1.6.2 w/o hive + val actualComponents = actual + .select("component", "id") + .as[(T, T)] + .rdd + .groupByKey() + .values + .map(_.toSeq) + .collect() + .map { ids => + val idSet = ids.toSet + assert( + idSet.size === ids.size, + s"Found duplicated component assignment in [${ids.mkString(",")}].") + idSet + } + .toSet + assert(actualComponents === expected) + () + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/DetectingCyclesSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/DetectingCyclesSuite.scala new file mode 100644 index 0000000000000..1a72c9d3223e0 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/DetectingCyclesSuite.scala @@ -0,0 +1,80 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite + +import scala.annotation.nowarn +import scala.collection.mutable + +class DetectingCyclesSuite extends SparkFunSuite with GraphFrameTestSparkContext { + test("test detecting cycles") { + val graph = GraphFrame( + spark + .createDataFrame(Seq((1L, "a"), (2L, "b"), (3L, "c"), (4L, "d"), (5L, "e"))) + .toDF("id", "attr"), + spark + .createDataFrame(Seq((1L, 2L), (2L, 3L), (3L, 1L), (1L, 4L), (2L, 5L))) + .toDF("src", "dst")) + val res = graph.detectingCycles.setUseLocalCheckpoints(true).run() + assert(res.count() == 1) + @nowarn val collected = + res + .collect() + .map(r => r.getAs[mutable.WrappedArray[Long]](0)) + + assert(collected(0) == Seq(1, 2, 3, 1)) + res.unpersist() + } + + test("test no cycles") { + val graph = GraphFrame( + spark + .createDataFrame(Seq((1L, "a"), (2L, "b"), (3L, "c"), (4L, "d"), (5L, "e"))) + .toDF("id", "attr"), + spark + .createDataFrame(Seq((1L, 2L), (2L, 3L), (3L, 4L), (4L, 5L))) + .toDF("src", "dst")) + val res = graph.detectingCycles.setUseLocalCheckpoints(true).run() + assert(res.count() == 0) + res.unpersist() + } + + test("test multiple cycles from one source") { + val graph = GraphFrame( + spark + .createDataFrame(Seq((1L, "a"), (2L, "b"), (3L, "c"), (4L, "d"), (5L, "e"))) + .toDF("id", "attr"), + spark + .createDataFrame(Seq((1L, 2L), (2L, 1L), (1L, 3L), (3L, 1L), (2L, 5L), (5L, 1L))) + .toDF("src", "dst")) + val res = graph.detectingCycles.setUseLocalCheckpoints(true).run() + assert(res.count() == 3) + @nowarn val collected = + res + .sort(DetectingCycles.foundSeqCol) + .collect() + .map(r => r.getAs[mutable.WrappedArray[Long]](0)) + assert(collected(0) == Seq(1, 2, 1)) + assert(collected(1) == Seq(1, 2, 5, 1)) + assert(collected(2) == Seq(1, 3, 1)) + res.unpersist() + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/HyperANFSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/HyperANFSuite.scala new file mode 100644 index 0000000000000..6d6e5f4070580 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/HyperANFSuite.scala @@ -0,0 +1,152 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.expr +import org.apache.spark.sql.functions.hll_sketch_estimate +import org.apache.spark.sql.types.DataTypes +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.TestUtils + +class HyperANFSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + private def diamondCycleGraph: GraphFrame = { + val vertices = + spark.createDataFrame((1L to 5L).map(id => (id, s"v$id"))).toDF("id", "name") + val edges = spark + .createDataFrame(Seq((1L, 2L), (1L, 3L), (2L, 4L), (3L, 4L), (4L, 5L), (5L, 1L))) + .toDF("src", "dst") + GraphFrame(vertices, edges) + } + + private def estimateHopCounts(result: DataFrame, nHops: Int): Map[Long, Seq[Long]] = { + val estimateColumns = (0 to nHops).map { hop => + hll_sketch_estimate(col(s"hop_$hop")).alias(s"hop_${hop}_estimate") + } + + result + .select((Seq(col("id")) ++ estimateColumns): _*) + .collect() + .map { row => + row.getAs[Long]("id") -> (0 to nHops).map { hop => + row.getAs[Long](s"hop_${hop}_estimate") + } + } + .toMap + } + + test("HyperANF returns exact 0-hop through 2-hop reachable cardinalities") { + val graph = diamondCycleGraph + val result = new HyperANF(graph) + .setNHops(2) + .setLgNomEntries(12) + .run() + + TestUtils.checkColumnType(result.schema, "hop_0", DataTypes.BinaryType) + TestUtils.checkColumnType(result.schema, "hop_1", DataTypes.BinaryType) + TestUtils.checkColumnType(result.schema, "hop_2", DataTypes.BinaryType) + + val expected = Map( + 1L -> Seq(1L, 2L, 1L), + 2L -> Seq(1L, 1L, 1L), + 3L -> Seq(1L, 1L, 1L), + 4L -> Seq(1L, 1L, 1L), + 5L -> Seq(1L, 1L, 2L)) + + assert(estimateHopCounts(result, 2) === expected) + result.unpersist() + } + + test("HyperANF returns exact 0-hop through 3-hop reachable cardinalities") { + val graph = diamondCycleGraph + val result = new HyperANF(graph) + .setNHops(3) + .setLgNomEntries(12) + .run() + + TestUtils.checkColumnType(result.schema, "hop_0", DataTypes.BinaryType) + TestUtils.checkColumnType(result.schema, "hop_1", DataTypes.BinaryType) + TestUtils.checkColumnType(result.schema, "hop_2", DataTypes.BinaryType) + TestUtils.checkColumnType(result.schema, "hop_3", DataTypes.BinaryType) + + val expected = Map( + 1L -> Seq(1L, 2L, 1L, 1L), + 2L -> Seq(1L, 1L, 1L, 1L), + 3L -> Seq(1L, 1L, 1L, 1L), + 4L -> Seq(1L, 1L, 1L, 2L), + 5L -> Seq(1L, 1L, 2L, 1L)) + + assert(estimateHopCounts(result, 3) === expected) + result.unpersist() + } + + test( + "HyperANF starting vertices expression limits output to selected vertices with outgoing edges") { + val graph = diamondCycleGraph + val result = new HyperANF(graph) + .setEdgesFilterExpression(expr("src IN (1, 3, 42)")) + .setNHops(2) + .setLgNomEntries(12) + .run() + + val ids = result.select("id").collect().map(_.getAs[Long]("id")).toSet + + assert(ids === Set(1L, 3L)) + result.unpersist() + } + + test("HyperANF does not fail on dead-ends") { + // Graph: 1 -> 2 -> 3, vertex 4 is isolated (no edges at all) + // Dead-ends: vertex 3 (no outgoing edges), vertex 4 (no edges at all) + val vertices = spark + .createDataFrame((1L to 4L).map(id => (id, s"v$id"))) + .toDF("id", "name") + val edges = spark + .createDataFrame(Seq((1L, 2L), (2L, 3L))) + .toDF("src", "dst") + val graph = GraphFrame(vertices, edges) + + val nHops = 3 + val result = new HyperANF(graph) + .setNHops(nHops) + .setLgNomEntries(12) + .run() + + val ids = result.select("id").collect().map(_.getAs[Long]("id")).toSet + + // 1. Dead-end vertices are not present in the output + assert(!ids.contains(3L), "Dead-end vertex 3 (no outgoing edges) should not be in output") + assert(!ids.contains(4L), "Isolated vertex 4 (no edges at all) should not be in output") + assert(ids === Set(1L, 2L)) + + // 2. hop_2 for vertex 2 reaches dead-end 3, producing an empty sketch (estimate 0, not null) + val estimates = estimateHopCounts(result, nHops) + assert(estimates(2L)(2) === 0L, "hop_2 estimate for vertex 2 should be 0") + + val row2 = result.filter(col("id") === 2L).collect()(0) + assert( + row2.getAs[Array[Byte]]("hop_2") !== null, + "hop_2 sketch for vertex 2 should not be null") + + result.unpersist() + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/KCoreSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/KCoreSuite.scala new file mode 100644 index 0000000000000..71b4a7b631b44 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/KCoreSuite.scala @@ -0,0 +1,342 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.types.DataTypes +import org.apache.spark.graphframes._ +import org.apache.spark.graphframes.examples.Graphs + +class KCoreSuite extends SparkFunSuite with GraphFrameTestSparkContext { + test("empty graph") { + val empty = Graphs.empty[Int] + val result = empty.kCore.run() + assert(result.count() === 0L) + result.unpersist() + } + + test("single vertex") { + val v = spark.createDataFrame(Seq((0L, "a"))).toDF("id", "name") + // Create an empty dataframe with the proper columns. + val e = spark.createDataFrame(Seq.empty[(Long, Long)]).toDF("src", "dst") + val g = GraphFrame(v, e) + val result = g.kCore.run() + TestUtils.checkColumnType(result.schema, "kcore", DataTypes.IntegerType) + assert(result.count() === 1) + val rows = result.collect() + assert(rows.head.getAs[Int]("kcore") === 0) + result.unpersist() + } + + test("two connected vertices") { + val v = spark.createDataFrame(Seq((0L, "a"), (1L, "b"))).toDF("id", "name") + val e = spark.createDataFrame(Seq((0L, 1L))).toDF("src", "dst") + val g = GraphFrame(v, e) + val result = g.kCore.run() + TestUtils.checkColumnType(result.schema, "kcore", DataTypes.IntegerType) + assert(result.count() === 2) + val rows = result.collect() + // Both vertices should have k-core value of 1 + rows.foreach { row => + assert(row.getAs[Int]("kcore") === 1) + } + result.unpersist() + } + + test("triangle graph") { + val v = spark.createDataFrame(Seq((0L, "a"), (1L, "b"), (2L, "c"))).toDF("id", "name") + val e = spark.createDataFrame(Seq((0L, 1L), (1L, 2L), (2L, 0L))).toDF("src", "dst") + val g = GraphFrame(v, e) + val result = g.kCore.run() + TestUtils.checkColumnType(result.schema, "kcore", DataTypes.IntegerType) + assert(result.count() === 3) + val rows = result.collect() + // All vertices should have k-core value of 2 + rows.foreach { row => + assert(row.getAs[Int]("kcore") === 2) + } + result.unpersist() + } + + test("star graph") { + val v = spark + .createDataFrame(Seq((0L, "center"), (1L, "leaf1"), (2L, "leaf2"), (3L, "leaf3"))) + .toDF("id", "name") + val e = spark.createDataFrame(Seq((0L, 1L), (0L, 2L), (0L, 3L))).toDF("src", "dst") + val g = GraphFrame(v, e) + val result = g.kCore.run() + TestUtils.checkColumnType(result.schema, "kcore", DataTypes.IntegerType) + assert(result.count() === 4) + val rows = result.collect() + // All vertices have k-core = 1: despite the center having degree 3, each leaf has + // only one neighbor so no 2-core can form, pulling everything down to 1. + rows.foreach { row => + assert(row.getAs[Int]("kcore") === 1) + } + result.unpersist() + } + + test("chain graph") { + // Open chain: 0 - 1 - 2. + // No 2-core exists: endpoints have degree 1, so the whole graph is only a 1-core. + // All vertices get kcore = 1. + val v = spark.createDataFrame(Seq((0L, "a"), (1L, "b"), (2L, "c"))).toDF("id", "name") + val e = spark.createDataFrame(Seq((0L, 1L), (1L, 2L))).toDF("src", "dst") + val g = GraphFrame(v, e) + val result = g.kCore.run() + TestUtils.checkColumnType(result.schema, "kcore", DataTypes.IntegerType) + assert(result.count() === 3) + val rows = result.collect() + rows.foreach { row => + assert(row.getAs[Int]("kcore") === 1) + } + result.unpersist() + } + + test("disconnected vertices") { + val v = spark.createDataFrame(Seq((0L, "a"), (1L, "b"), (2L, "c"))).toDF("id", "name") + val e = spark.createDataFrame(Seq.empty[(Long, Long)]).toDF("src", "dst") + val g = GraphFrame(v, e) + val result = g.kCore.run() + TestUtils.checkColumnType(result.schema, "kcore", DataTypes.IntegerType) + assert(result.count() === 3) + val rows = result.collect() + // All vertices should have k-core value of 0 + rows.foreach { row => + assert(row.getAs[Int]("kcore") === 0) + } + result.unpersist() + } + + test("friends graph") { + val friends = Graphs.friends + val result = friends.kCore.run() + TestUtils.checkColumnType(result.schema, "kcore", DataTypes.IntegerType) + assert(result.count() === friends.vertices.count()) + // In the friends graph, all vertices except 'g' should have k-core >= 1 + // 'g' is isolated, so it should have k-core 0 + val rows = result.collect() + rows.foreach { row => + val id = row.getAs[String]("id") + val kcore = row.getAs[Int]("kcore") + if (id == "g") { + assert(kcore === 0) + } else { + assert(kcore >= 1) + } + } + result.unpersist() + } + + test("medium graph with varying k-core values") { + // Create a graph with 25 vertices and varying degrees to get different k-core values + val v = + spark.createDataFrame((0L until 25L).map(id => (id, s"vertex_$id"))).toDF("id", "name") + + // Create edges to form a graph with diverse connectivity + val edges = Seq( + // High degree cluster around vertex 0 (should have high k-core) + (0L, 1L), + (0L, 2L), + (0L, 3L), + (0L, 4L), + (0L, 5L), + (1L, 2L), + (1L, 3L), + (2L, 3L), + (2L, 4L), + (3L, 4L), + (1L, 6L), + (2L, 7L), + (3L, 8L), + (4L, 9L), + (5L, 10L), + + // Medium degree cluster around vertex 11 + (11L, 12L), + (11L, 13L), + (11L, 14L), + (12L, 13L), + (12L, 15L), + (13L, 14L), + (13L, 16L), + (14L, 17L), + + // Chain structure (lower k-core values) + (18L, 19L), + (19L, 20L), + (20L, 21L), + (21L, 22L), + + // Some additional connections to create more varied structure + (6L, 12L), + (7L, 13L), + (8L, 14L), + (9L, 15L), + (10L, 16L), + + // Isolated vertices or low-degree vertices + (23L, 24L)) + + val e = spark.createDataFrame(edges).toDF("src", "dst") + val g = GraphFrame(v, e) + val result = g.kCore.run() + TestUtils.checkColumnType(result.schema, "kcore", DataTypes.IntegerType) + assert(result.count() === 25) + + val rows = result.collect() + // Check that we have a range of k-core values + val kcoreValues = rows.map(_.getAs[Int]("kcore")).distinct.sorted + assert(kcoreValues.length > 2, "Should have at least 3 distinct k-core values") + + // Verify specific expected patterns + val kcoreMap = rows.map(row => row.getAs[Long]("id") -> row.getAs[Int]("kcore")).toMap + + // Vertices in the highly connected cluster should have higher k-core values + assert(kcoreMap(0L) >= 3, "Central vertex should have high k-core") + assert(kcoreMap(1L) >= 3, "Well-connected vertex should have high k-core") + + // Leaf nodes should have lower k-core values + assert(kcoreMap(18L) <= 2, "Leaf node should have low k-core") + assert(kcoreMap(23L) <= 1, "Low-degree node should have very low k-core") + + result.unpersist() + } + + test("graph with clear hierarchical k-core structure") { + // Create a graph designed to have clear k-core layers + val v = spark.createDataFrame((0L until 30L).map(id => (id, s"v$id"))).toDF("id", "name") + + // Build edges to create a hierarchical structure: + // Core (k=5): vertices 0-4 - fully connected + // Next layer (k=3): vertices 5-14 - each connects to multiple core vertices + // Outer layer (k=1): vertices 15-29 - sparse connections + val coreEdges = for { + i <- 0 until 5 + j <- (i + 1) until 5 + } yield (i.toLong, j.toLong) + + val midLayerEdges = Seq( + (5L, 0L), + (5L, 1L), + (5L, 2L), // Connect to core + (6L, 0L), + (6L, 1L), + (6L, 3L), + (7L, 1L), + (7L, 2L), + (7L, 4L), + (8L, 0L), + (8L, 3L), + (8L, 4L), + (9L, 1L), + (9L, 2L), + (9L, 3L), + (10L, 0L), + (10L, 4L), + (11L, 2L), + (11L, 3L), + (12L, 1L), + (12L, 4L), + (13L, 0L), + (13L, 2L), + (14L, 3L), + (14L, 4L)) + + val outerEdges = Seq( + (15L, 5L), + (16L, 6L), + (17L, 7L), + (18L, 8L), + (19L, 9L), + (20L, 10L), + (21L, 11L), + (22L, 12L), + (23L, 13L), + (24L, 14L), + (25L, 15L), + (26L, 16L), + (27L, 17L), + (28L, 18L), + (29L, 19L)) + + val allEdges = coreEdges ++ midLayerEdges ++ outerEdges + val e = spark.createDataFrame(allEdges).toDF("src", "dst") + val g = GraphFrame(v, e) + val result = g.kCore.run() + TestUtils.checkColumnType(result.schema, "kcore", DataTypes.IntegerType) + assert(result.count() === 30) + + val rows = result.collect() + val kcoreMap = rows.map(row => row.getAs[Long]("id") -> row.getAs[Int]("kcore")).toMap + + // Validate hierarchical structure + // Core vertices (0-4) should have highest k-core + (0L to 4L).foreach { id => + assert(kcoreMap(id) >= 4, s"Core vertex $id should have high k-core, got ${kcoreMap(id)}") + } + + // Mid-layer vertices (5-14) should have medium k-core + (5L to 14L).foreach { id => + assert( + kcoreMap(id) >= 2, + s"Mid-layer vertex $id should have medium k-core, got ${kcoreMap(id)}") + assert( + kcoreMap(id) <= 4, + s"Mid-layer vertex $id should not have too high k-core, got ${kcoreMap(id)}") + } + + // Outer vertices (15-29) should have low k-core + (15L to 29L).foreach { id => + assert(kcoreMap(id) <= 2, s"Outer vertex $id should have low k-core, got ${kcoreMap(id)}") + } + + result.unpersist() + } + + test("triangle with tail - exact kcore values") { + // This graph has vertices where degree != kcore, which is important to test correctness: + // it would catch a buggy implementation that converges too early (e.g. after one superstep), which + // would return kcore = degree for all vertices and pass simpler tests. + // + // Undirected graph: + // + // Triangle: 1 - 2 - 3 - 1 (kcore = 2, they form the 2-core) + // Tail: 1 - 4 - 5 (kcore = 1, pendant chain excluded from the 2-core) + // + // Degrees: 1→3, 2→2, 3→2, 4→2, 5→1 (degree != kcore for vertices 1 and 4) + val v = spark + .createDataFrame(Seq((1L, "a"), (2L, "b"), (3L, "c"), (4L, "d"), (5L, "e"))) + .toDF("id", "name") + val e = spark + .createDataFrame(Seq((1L, 2L), (2L, 3L), (3L, 1L), (1L, 4L), (4L, 5L))) + .toDF("src", "dst") + val g = GraphFrame(v, e) + val result = g.kCore.run() + TestUtils.checkColumnType(result.schema, "kcore", DataTypes.IntegerType) + assert(result.count() === 5) + val kcoreMap = result.collect().map(r => r.getAs[Long]("id") -> r.getAs[Int]("kcore")).toMap + // Triangle vertices form the 2-core + assert(kcoreMap(1L) === 2) + assert(kcoreMap(2L) === 2) + assert(kcoreMap(3L) === 2) + // Tail vertices are only in the 1-core + assert(kcoreMap(4L) === 1) + assert(kcoreMap(5L) === 1) + result.unpersist() + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/LabelPropagationSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/LabelPropagationSuite.scala new file mode 100644 index 0000000000000..c2036ee7b64dc --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/LabelPropagationSuite.scala @@ -0,0 +1,44 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.types.DataTypes +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.TestUtils +import org.apache.spark.graphframes.examples.Graphs + +class LabelPropagationSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + val n = 5 + + test("Toy example") { + val g = Graphs.twoBlobs(n) + val labels = g.labelPropagation.maxIter(4 * n).run() + TestUtils.testSchemaInvariants(g, labels) + TestUtils.checkColumnType(labels.schema, "label", DataTypes.LongType) + val clique1 = + labels.filter(s"id < $n").select("label").collect().toSeq.map(_.getLong(0)).toSet + assert(clique1.size === 1) + val clique2 = + labels.filter(s"id >= $n").select("label").collect().toSeq.map(_.getLong(0)).toSet + assert(clique2.size === 1) + assert(clique1 !== clique2) + labels.unpersist() + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/MaximalIndependentSetSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/MaximalIndependentSetSuite.scala new file mode 100644 index 0000000000000..335d114562b53 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/MaximalIndependentSetSuite.scala @@ -0,0 +1,138 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.col +import org.apache.spark.graphframes._ +import org.apache.spark.graphframes.examples.Graphs + +class MaximalIndependentSetSuite extends SparkFunSuite with GraphFrameTestSparkContext { + test("isolated vertices should be included in MIS") { + // Create a graph with isolated vertices + val vertices = + spark.createDataFrame(Seq((0L, "a"), (1L, "b"), (2L, "c"), (3L, "d"))).toDF("id", "name") + + // Only connect vertices 0 and 1 + val edges = spark.createDataFrame(Seq((0L, 1L, "edge1"))).toDF("src", "dst", "name") + + val graph = GraphFrame(vertices, edges) + val mis = graph.maximalIndependentSet.run(seed = 12345L) + + // Check that all vertices are in the MIS (since 2 and 3 are isolated) + val misIds = mis.select("id").collect().map(_.getLong(0)).toSet + assert(misIds.size == 3, "MIS should contain 2 isolated vertices and one of linked") + assert(misIds.contains(2L), "Isolated vertex 2 should be in MIS") + assert(misIds.contains(3L), "Isolated vertex 3 should be in MIS") + + mis.unpersist() + } + + def isIndependent(graph: GraphFrame, mis: DataFrame): Boolean = { + graph.edges + .join(mis, col(GraphFrame.SRC) === col(GraphFrame.ID)) + .select(col(GraphFrame.DST)) + .join(mis, col(GraphFrame.DST) === col(GraphFrame.ID)) + .count() == 0 + } + + def isMaximal(graph: GraphFrame, mis: DataFrame): Boolean = { + val undirectedG = graph.asUndirected() + val verticesNotInMIS = undirectedG.vertices.join(mis, Seq(GraphFrame.ID), "left_anti") + + val verticesWithEdgesToMIS = undirectedG.edges + .join(mis, col(GraphFrame.ID) === col(GraphFrame.DST)) + .select(GraphFrame.SRC) + .distinct() + + val countVerticesNotInMIS = verticesNotInMIS.count() + val countVerticesWithEdgesToMIS = verticesWithEdgesToMIS.count() + + countVerticesNotInMIS == countVerticesWithEdgesToMIS + } + + test("correct MIS, seed 12345") { + val graph = Graphs.friends + + val mis = graph.maximalIndependentSet.run(seed = 12345L) + + assert(isIndependent(graph, mis)) + assert(isMaximal(graph, mis)) + + mis.unpersist() + } + + test("correct MIS, seed 23456") { + val graph = Graphs.friends + + val mis = graph.maximalIndependentSet.run(seed = 23456L) + + assert(isIndependent(graph, mis)) + assert(isMaximal(graph, mis)) + + mis.unpersist() + } + + test("MIS on empty graph") { + val emptyGraph = Graphs.empty[Long] + val mis = emptyGraph.maximalIndependentSet.run(seed = 12345L) + assert(mis.count() == 0, "MIS of empty graph should be empty") + mis.unpersist() + } + + test("MIS on single vertex graph") { + val vertices = spark.createDataFrame(Seq((0L, "vertex"))).toDF("id", "name") + val edges = spark.createDataFrame(Seq.empty[(Long, Long)]).toDF("src", "dst") + val graph = GraphFrame(vertices, edges) + + val mis = graph.maximalIndependentSet.run(seed = 12345L) + assert(mis.count() == 1, "MIS of single vertex graph should contain one vertex") + + val misId = mis.select("id").collect()(0).getLong(0) + assert(misId == 0L, "MIS should contain vertex with id 0") + + mis.unpersist() + } + + test("MIS on disconnected vertices") { + val n = 5L + val vertices = spark.range(n).toDF("id") + val edges = spark.createDataFrame(Seq.empty[(Long, Long)]).toDF("src", "dst") + val graph = GraphFrame(vertices, edges) + + val mis = graph.maximalIndependentSet.run(seed = 12345L) + assert(mis.count() == n, s"MIS should contain all $n vertices for disconnected graph") + + mis.unpersist() + } + + test("MIS on complete graph of 5 vertices") { + val vertices = spark.range(5).toDF("id") + val edges = for { + i <- 0L until 5L + j <- (i + 1L) until 5L + } yield (i, j) + val edgeDF = spark.createDataFrame(edges).toDF("src", "dst") + val graph = GraphFrame(vertices, edgeDF) + + val mis = graph.maximalIndependentSet.run(seed = 12345L) + assert(mis.count() == 1, "MIS of complete graph should contain exactly one vertex") + + mis.unpersist() + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/PageRankSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/PageRankSuite.scala new file mode 100644 index 0000000000000..ba70519d1e172 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/PageRankSuite.scala @@ -0,0 +1,82 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.types.DataTypes +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.TestUtils +import org.apache.spark.graphframes.examples.Graphs + +class PageRankSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + val n = 100L + + test("Star example") { + val g = Graphs.star(n) + val resetProb = 0.15 + val errorTol = 1.0e-5 + val pr = g.pageRank + .resetProbability(resetProb) + .tol(errorTol) + .run() + TestUtils.testSchemaInvariants(g, pr) + TestUtils.checkColumnType(pr.vertices.schema, "pagerank", DataTypes.DoubleType) + TestUtils.checkColumnType(pr.edges.schema, "weight", DataTypes.DoubleType) + pr.unpersist() + } + + test("friends graph with personalized PageRank") { + val results = Graphs.friends.pageRank.resetProbability(0.15).maxIter(10).sourceId("a").run() + + val gRank = results.vertices.filter(col("id") === "g").select("pagerank").first().getDouble(0) + assert( + gRank === 0.0, + s"User g (Gabby) doesn't connect with a. So its pagerank should be 0 but we got $gRank.") + results.unpersist() + } + + test("graph with three disconnected components") { + import sqlImplicits._ + + val v = Seq((0L, "a"), (1L, "b"), (2L, "c"), (3L, "d"), (4L, "e"), (5L, "f"), (6L, "g")) + .toDF("id", "name") + + val e = Seq( + (0L, 1L, "friend"), // First component: a->b->c + (1L, 2L, "friend"), + (3L, 4L, "friend"), // Second component: d->e->f + (4L, 5L, "friend") + // Third component: isolated vertex g (6) + ).toDF("src", "dst", "relationship") + + val g = org.apache.spark.graphframes.GraphFrame(v, e) + val results = g.pageRank.resetProbability(0.15).maxIter(10).run() + + // Verify that all original vertices are present in the result + assert( + results.vertices.count() === v.count(), + "PageRank results should contain all original vertices") + + val originalIds = v.select("id").collect().map(_.getLong(0)).toSet + val resultIds = results.vertices.select("id").collect().map(_.getLong(0)).toSet + assert(originalIds === resultIds, "PageRank results should preserve all vertex IDs") + results.unpersist() + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRankSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRankSuite.scala new file mode 100644 index 0000000000000..0976ad97e9836 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRankSuite.scala @@ -0,0 +1,102 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib +import org.apache.spark.ml.linalg.SQLDataTypes +import org.apache.spark.ml.linalg.SparseVector +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.types.DataTypes +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.TestUtils +import org.apache.spark.graphframes.examples.Graphs + +class ParallelPersonalizedPageRankSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + val n = 100L + + test("Illegal function call argument setting") { + val g = Graphs.star(n) + val vertexIds: Array[Any] = Array(1L, 2L, 3L) + + // Not providing number of iterations + intercept[IllegalArgumentException] { + g.parallelPersonalizedPageRank.sourceIds(vertexIds).run() + } + + // Not providing sourceIds + intercept[IllegalArgumentException] { + g.parallelPersonalizedPageRank.maxIter(15).run() + } + + // Provided empty sourceIds + intercept[IllegalArgumentException] { + g.parallelPersonalizedPageRank.maxIter(15).sourceIds(Array()).run() + } + } + + test("Star example parallel personalized PageRank") { + val g = Graphs.star(n) + val resetProb = 0.15 + val maxIter = 10 + val vertexIds: Array[Any] = Array(1L, 2L, 3L) + + lazy val prc = g.parallelPersonalizedPageRank + .maxIter(maxIter) + .sourceIds(vertexIds) + .resetProbability(resetProb) + + val pr = prc.run() + TestUtils.testSchemaInvariants(g, pr) + TestUtils.checkColumnType(pr.vertices.schema, "pageranks", SQLDataTypes.VectorType) + TestUtils.checkColumnType(pr.edges.schema, "weight", DataTypes.DoubleType) + pr.unpersist() + } + + test("friends graph with parallel personalized PageRank") { + val g = Graphs.friends + val resetProb = 0.15 + val maxIter = 10 + val vertexIds: Array[Any] = Array("a") + lazy val prc = g.parallelPersonalizedPageRank + .maxIter(maxIter) + .sourceIds(vertexIds) + .resetProbability(resetProb) + + val pr = prc.run() + val prInvalid = pr.vertices + .select("pageranks") + .collect() + .filter { (row: Row) => + vertexIds.size != row.getAs[SparseVector](0).size + } + assert( + prInvalid.size === 0, + s"found ${prInvalid.size} entries with invalid number of returned personalized pagerank vector") + + val gRank = pr.vertices + .filter(col("id") === "g") + .select("pageranks") + .first() + .getAs[SparseVector](0) + assert( + gRank.numNonzeros === 0, + s"User g (Gabby) doesn't connect with a. So its pagerank should be 0 but we got ${gRank.numNonzeros}.") + pr.unpersist() + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/PregelSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/PregelSuite.scala new file mode 100644 index 0000000000000..e9b80471ea593 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/PregelSuite.scala @@ -0,0 +1,590 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.functions._ +import org.apache.spark.graphframes._ +import org.scalactic.Tolerance._ + +class PregelSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + import sqlImplicits._ + + Seq(true, false).foreach(useLocalCheckpoint => { + test(s"page rank${if (useLocalCheckpoint) " with local checkpoint" else ""}") { + spark.conf.set("spark.graphframes.useLocalCheckpoints", useLocalCheckpoint.toString) + val edges = Seq( + (0L, 1L), + (1L, 2L), + (2L, 4L), + (2L, 0L), + (3L, 4L), // 3 has no in-links + (4L, 0L), + (4L, 2L)).toDF("src", "dst").cache() + val vertices = GraphFrame.fromEdges(edges).outDegrees.cache() + val numVertices = vertices.count() + val graph = GraphFrame(vertices, edges) + + val alpha = 0.15 + // NOTE: This version doesn't handle nodes with no out-links. + val ranks = graph.pregel + .setMaxIter(5) + .withVertexColumn( + "rank", + lit(1.0 / numVertices), + coalesce(Pregel.msg, lit(0.0)) * (1.0 - alpha) + alpha / numVertices) + .sendMsgToDst(Pregel.src("rank") / Pregel.src("outDegree")) + .aggMsgs(sum(Pregel.msg)) + .run() + + val result = ranks + .sort(col("id")) + .select("rank") + .as[Double] + .collect() + assert(result.sum === 1.0 +- 1e-6) + val expected = Seq(0.245, 0.224, 0.303, 0.03, 0.197) + result.zip(expected).foreach { case (r, e) => + assert(r === e +- 1e-3) + } + spark.conf.set("spark.graphframes.useLocalCheckpoints", "false") + } + + test(s"chain propagation${if (useLocalCheckpoint) " with local checkpoint" else ""}") { + spark.conf.set("spark.graphframes.useLocalCheckpoints", useLocalCheckpoint.toString) + val n = 5 + val verDF = (1 to n).toDF("id").repartition(3) + val edgeDF = (1 until n) + .map(x => (x, x + 1)) + .toDF("src", "dst") + .repartition(3) + + val graph = GraphFrame(verDF, edgeDF) + + val resultDF = graph.pregel + .setMaxIter(n - 1) + .withVertexColumn( + "value", + when(col("id") === lit(1), lit(1)).otherwise(lit(0)), + when(Pregel.msg > col("value"), Pregel.msg).otherwise(col("value"))) + .sendMsgToDst(when(Pregel.dst("value") =!= Pregel.src("value"), Pregel.src("value"))) + .aggMsgs(max(Pregel.msg)) + .run() + + assert(resultDF.sort("id").select("value").as[Int].collect() === Array.fill(n)(1)) + spark.conf.set("spark.graphframes.useLocalCheckpoints", "false") + } + + test( + s"reverse chain propagation${if (useLocalCheckpoint) " with local checkpoint" else ""}") { + spark.conf.set("spark.graphframes.useLocalCheckpoints", useLocalCheckpoint.toString) + val n = 5 + val verDF = (1 to n).toDF("id").repartition(3) + val edgeDF = (1 until n) + .map(x => (x + 1, x)) + .toDF("src", "dst") + .repartition(3) + + val graph = GraphFrame(verDF, edgeDF) + + val resultDF = graph.pregel + .setMaxIter(n - 1) + .withVertexColumn( + "value", + when(col("id") === lit(1), lit(1)).otherwise(lit(0)), + when(Pregel.msg > col("value"), Pregel.msg).otherwise(col("value"))) + .sendMsgToSrc(when(Pregel.dst("value") =!= Pregel.src("value"), Pregel.dst("value"))) + .aggMsgs(max(Pregel.msg)) + .run() + + assert(resultDF.sort("id").select("value").as[Int].collect() === Array.fill(n)(1)) + spark.conf.set("spark.graphframes.useLocalCheckpoints", "false") + } + + test(s"chain propagation with termination${if (useLocalCheckpoint) " with local checkpoint" + else ""}") { + spark.conf.set("spark.graphframes.useLocalCheckpoints", useLocalCheckpoint.toString) + val n = 5 + val verDF = (1 to n).toDF("id").repartition(3) + val edgeDF = (1 until n) + .map(x => (x, x + 1)) + .toDF("src", "dst") + .repartition(3) + + val graph = GraphFrame(verDF, edgeDF) + + val resultDF = graph.pregel + .setMaxIter(1000) + .setEarlyStopping(true) + .withVertexColumn( + "value", + when(col("id") === lit(1), lit(1)).otherwise(lit(0)), + when(Pregel.msg > col("value"), Pregel.msg).otherwise(col("value"))) + .sendMsgToDst(when(Pregel.dst("value") =!= Pregel.src("value"), Pregel.src("value"))) + .aggMsgs(max(Pregel.msg)) + .run() + + assert(resultDF.sort("id").select("value").as[Int].collect() === Array.fill(n)(1)) + spark.conf.set("spark.graphframes.useLocalCheckpoints", "false") + } + }) + + test("new vertex column is based on the nullable column") { + val verDF = Seq(1L, 2L, 3L, 4L) + .toDF("id") + .withColumn( + "nullableColumn", + when(col("id") % lit(2) === lit(0), lit(null)).otherwise(lit(1))) + val edgeDF = Seq((1L, 2L), (2L, 3L), (3L, 4L), (4L, 1L)).toDF("src", "dst") + val graph = GraphFrame(verDF, edgeDF) + val pregel = graph.pregel + .withVertexColumn( + "newColumn", + when(col("nullableColumn").isNull, lit(0)).otherwise(lit(1)), + col("newColumn") + Pregel.msg) + .sendMsgToDst(lit(1)) + .aggMsgs(last(Pregel.msg)) + .setCheckpointInterval(0) + .setMaxIter(1) + + val resultDF = pregel.run() + assert( + resultDF + .select("id", "newColumn") + .collect() + .map(r => r.getAs[Long]("id") -> r.getAs[Int]("newColumn")) + .toMap === Map(1L -> 2, 2L -> 1, 3L -> 2, 4L -> 1)) + } + + test("requiredSrcColumns - only specified columns are used in triplets") { + // Test that requiredSrcColumns correctly limits the columns in triplets + // This is a memory optimization test - we verify the result is correct + // with only required source columns + + val edges = Seq((0L, 1L), (1L, 2L), (2L, 4L), (2L, 0L), (3L, 4L), (4L, 0L), (4L, 2L)) + .toDF("src", "dst") + .cache() + val vertices = GraphFrame.fromEdges(edges).outDegrees.cache() + val numVertices = vertices.count() + val graph = GraphFrame(vertices, edges) + + val alpha = 0.15 + // PageRank only needs "rank" and "outDegree" from source vertex + val ranks = graph.pregel + .setMaxIter(5) + .withVertexColumn( + "rank", + lit(1.0 / numVertices), + coalesce(Pregel.msg, lit(0.0)) * (1.0 - alpha) + alpha / numVertices) + .sendMsgToDst(Pregel.src("rank") / Pregel.src("outDegree")) + .aggMsgs(sum(Pregel.msg)) + .requiredSrcColumns("rank", "outDegree") + .run() + + val result = ranks + .sort(col("id")) + .select("rank") + .as[Double] + .collect() + assert(result.sum === 1.0 +- 1e-6) + val expected = Seq(0.245, 0.224, 0.303, 0.03, 0.197) + result.zip(expected).foreach { case (r, e) => + assert(r === e +- 1e-3) + } + } + + test("requiredDstColumns - only specified columns are used in triplets") { + // Test that requiredDstColumns correctly limits the columns in triplets + // Reverse chain propagation where we only need dst("value") from destination + + val n = 5 + val verDF = (1 to n).toDF("id").repartition(3) + val edgeDF = (1 until n) + .map(x => (x + 1, x)) + .toDF("src", "dst") + .repartition(3) + + val graph = GraphFrame(verDF, edgeDF) + + val resultDF = graph.pregel + .setMaxIter(n - 1) + .withVertexColumn( + "value", + when(col("id") === lit(1), lit(1)).otherwise(lit(0)), + when(Pregel.msg > col("value"), Pregel.msg).otherwise(col("value"))) + .sendMsgToSrc(when(Pregel.dst("value") =!= Pregel.src("value"), Pregel.dst("value"))) + .aggMsgs(max(Pregel.msg)) + .requiredDstColumns("value") // Only need "value" from destination + .run() + + assert(resultDF.sort("id").select("value").as[Int].collect() === Array.fill(n)(1)) + } + + test("requiredSrcColumns and requiredDstColumns together") { + // Test using both requiredSrcColumns and requiredDstColumns + // Chain propagation where we need "value" from both src and dst + + val n = 5 + val verDF = (1 to n).toDF("id").repartition(3) + val edgeDF = (1 until n) + .map(x => (x, x + 1)) + .toDF("src", "dst") + .repartition(3) + + val graph = GraphFrame(verDF, edgeDF) + + val resultDF = graph.pregel + .setMaxIter(n - 1) + .withVertexColumn( + "value", + when(col("id") === lit(1), lit(1)).otherwise(lit(0)), + when(Pregel.msg > col("value"), Pregel.msg).otherwise(col("value"))) + .sendMsgToDst(when(Pregel.dst("value") =!= Pregel.src("value"), Pregel.src("value"))) + .aggMsgs(max(Pregel.msg)) + .requiredSrcColumns("value") // Only need "value" from source + .requiredDstColumns("value") // Only need "value" from destination + .run() + + assert(resultDF.sort("id").select("value").as[Int].collect() === Array.fill(n)(1)) + } + + test("requiredSrcColumns with empty list uses all columns (default behavior)") { + // Verify that not calling requiredSrcColumns means all columns are used + // This is the same as the original page rank test + + val edges = Seq((0L, 1L), (1L, 2L), (2L, 4L), (2L, 0L), (3L, 4L), (4L, 0L), (4L, 2L)) + .toDF("src", "dst") + .cache() + val vertices = GraphFrame.fromEdges(edges).outDegrees.cache() + val numVertices = vertices.count() + val graph = GraphFrame(vertices, edges) + + val alpha = 0.15 + val ranks = graph.pregel + .setMaxIter(5) + .withVertexColumn( + "rank", + lit(1.0 / numVertices), + coalesce(Pregel.msg, lit(0.0)) * (1.0 - alpha) + alpha / numVertices) + .sendMsgToDst(Pregel.src("rank") / Pregel.src("outDegree")) + .aggMsgs(sum(Pregel.msg)) + // No requiredSrcColumns or requiredDstColumns - should use all columns + .run() + + val result = ranks + .sort(col("id")) + .select("rank") + .as[Double] + .collect() + assert(result.sum === 1.0 +- 1e-6) + } + + test("automatic dst join skipping - PageRank only uses src columns") { + // PageRank only references Pregel.src("rank") and Pregel.src("outDegree"), + // so the second join (for dst vertex state) should be automatically skipped. + // This test verifies the optimization produces correct results. + + val edges = Seq((0L, 1L), (1L, 2L), (2L, 4L), (2L, 0L), (3L, 4L), (4L, 0L), (4L, 2L)) + .toDF("src", "dst") + .cache() + val vertices = GraphFrame.fromEdges(edges).outDegrees.cache() + val numVertices = vertices.count() + val graph = GraphFrame(vertices, edges) + + val alpha = 0.15 + // PageRank only uses Pregel.src(...) - dst state should be automatically skipped + val ranks = graph.pregel + .setMaxIter(5) + .withVertexColumn( + "rank", + lit(1.0 / numVertices), + coalesce(Pregel.msg, lit(0.0)) * (1.0 - alpha) + alpha / numVertices) + .sendMsgToDst(Pregel.src("rank") / Pregel.src("outDegree")) + .aggMsgs(sum(Pregel.msg)) + .run() + + val result = ranks + .sort(col("id")) + .select("rank") + .as[Double] + .collect() + assert(result.sum === 1.0 +- 1e-6) + val expected = Seq(0.245, 0.224, 0.303, 0.03, 0.197) + result.zip(expected).foreach { case (r, e) => + assert(r === e +- 1e-3) + } + } + + test("automatic dst join NOT skipped when dst columns are referenced") { + // This test uses Pregel.dst("value") in the message expression, + // so the second join must NOT be skipped. + + val n = 5 + val verDF = (1 to n).toDF("id").repartition(3) + val edgeDF = (1 until n) + .map(x => (x, x + 1)) + .toDF("src", "dst") + .repartition(3) + + val graph = GraphFrame(verDF, edgeDF) + + val resultDF = graph.pregel + .setMaxIter(n - 1) + .withVertexColumn( + "value", + when(col("id") === lit(1), lit(1)).otherwise(lit(0)), + when(Pregel.msg > col("value"), Pregel.msg).otherwise(col("value"))) + // This references BOTH src and dst - dst join should NOT be skipped + .sendMsgToDst(when(Pregel.dst("value") =!= Pregel.src("value"), Pregel.src("value"))) + .aggMsgs(max(Pregel.msg)) + .run() + + assert(resultDF.sort("id").select("value").as[Int].collect() === Array.fill(n)(1)) + } + + test("automatic dst join skipping with skipMessagesFromNonActiveVertices enabled") { + // When skipMessagesFromNonActiveVertices is true but message expressions don't + // reference dst columns, the dst join is still skipped. Active-vertex filtering + // is pushed before the src-edge join to reduce data volume. + + val n = 5 + val verDF = (1 to n).toDF("id").repartition(3) + val edgeDF = (1 until n) + .map(x => (x, x + 1)) + .toDF("src", "dst") + .repartition(3) + + val graph = GraphFrame(verDF, edgeDF) + + // This only uses Pregel.src("value") - dst join should be skipped, + // and active-vertex filtering is applied before the src-edge join. + val resultDF = graph.pregel + .setMaxIter(n - 1) + .setSkipMessagesFromNonActiveVertices(true) + .setUpdateActiveVertexExpression(Pregel.msg.isNotNull) + .withVertexColumn( + "value", + when(col("id") === lit(1), lit(1)).otherwise(lit(0)), + when(Pregel.msg > col("value"), Pregel.msg).otherwise(col("value"))) + .sendMsgToDst(Pregel.src("value")) + .aggMsgs(max(Pregel.msg)) + .run() + + assert(resultDF.sort("id").select("value").as[Int].collect() === Array.fill(n)(1)) + } + + test("automatic dst join skipping - sendMsgToSrc with only edge columns") { + // When sending messages to src using only edge columns, dst join should be skipped + + val edges = Seq((1L, 0L, 10L), (2L, 1L, 20L), (3L, 2L, 30L), (4L, 3L, 40L)) + .toDF("src", "dst", "weight") + .cache() + val vertices = (0L to 4L).toDF("id").cache() + + val graph = GraphFrame(vertices, edges) + + // Only uses Pregel.edge("weight") - dst join should be skipped + val resultDF = graph.pregel + .requiredEdgeColumns("weight") + .setMaxIter(1) + .withVertexColumn("received", lit(0L), coalesce(Pregel.msg, col("received"))) + .sendMsgToSrc(Pregel.edge("weight")) + .aggMsgs(sum(Pregel.msg)) + .run() + + // Each src vertex receives the weight from its outgoing edge + val received = resultDF.sort("id").select("received").as[Long].collect() + assert(received(0) === 0L) // vertex 0: no outgoing edges + assert(received(1) === 10L) // vertex 1: edge 1->0 with weight 10 + assert(received(2) === 20L) // vertex 2: edge 2->1 with weight 20 + assert(received(3) === 30L) // vertex 3: edge 3->2 with weight 30 + assert(received(4) === 40L) // vertex 4: edge 4->3 with weight 40 + } + + test("automatic dst join skipping - edge columns only") { + // When message expressions only reference edge columns, dst join should be skipped + + val edges = + Seq((0L, 1L, 1.0), (1L, 2L, 2.0), (2L, 3L, 3.0)).toDF("src", "dst", "weight").cache() + val vertices = Seq(0L, 1L, 2L, 3L).toDF("id").cache() + val graph = GraphFrame(vertices, edges) + + // Only uses Pregel.edge("weight") - dst join should be skipped + val result = graph.pregel + .requiredEdgeColumns("weight") + .setMaxIter(1) // Single iteration to simplify testing + .withVertexColumn("total", lit(0.0), coalesce(Pregel.msg, col("total"))) + .sendMsgToDst(Pregel.edge("weight")) + .aggMsgs(sum(Pregel.msg)) + .run() + + // Verify results: vertex 1 gets weight 1.0, vertex 2 gets 2.0, vertex 3 gets 3.0 + val totals = result.sort("id").select("total").as[Double].collect() + assert(totals(0) === 0.0 +- 1e-6) // vertex 0: no incoming edges + assert(totals(1) === 1.0 +- 1e-6) // vertex 1: edge 0->1 with weight 1.0 + assert(totals(2) === 2.0 +- 1e-6) // vertex 2: edge 1->2 with weight 2.0 + assert(totals(3) === 3.0 +- 1e-6) // vertex 3: edge 2->3 with weight 3.0 + } + + test("automatic dst join skipping - sendMsgToDst with only src columns in message") { + // When sendMsgToDst is used but the message expression only references src columns, + // the second join should be skipped. The dst.id needed for message routing is + // obtained from the edge's dst column, not from a vertex join. + + val edges = Seq((0L, 1L), (1L, 2L), (2L, 3L)).toDF("src", "dst").cache() + val vertices = (0L to 3L).toDF("id").cache() + val graph = GraphFrame(vertices, edges) + + // sendMsgToDst but message only uses Pregel.src("id") - dst join should be skipped + val result = graph.pregel + .setMaxIter(1) + .withVertexColumn("received", lit(0L), coalesce(Pregel.msg, col("received"))) + .sendMsgToDst(Pregel.src("id")) // Message only uses src.id + .aggMsgs(sum(Pregel.msg)) + .run() + + // Each dst vertex receives the src.id from incoming edges + val received = result.sort("id").select("received").as[Long].collect() + assert(received(0) === 0L) // vertex 0: no incoming edges + assert(received(1) === 0L) // vertex 1: edge 0->1, receives src.id = 0 + assert(received(2) === 1L) // vertex 2: edge 1->2, receives src.id = 1 + assert(received(3) === 2L) // vertex 3: edge 2->3, receives src.id = 2 + } + + test("automatic dst join skipping - message references only dst.id") { + // When message expressions only reference dst.id (not other dst fields), + // the join should still be skipped since dst.id is available from the edge. + + val edges = Seq((0L, 1L), (1L, 2L), (2L, 3L)).toDF("src", "dst").cache() + val vertices = (0L to 3L).toDF("id").cache() + val graph = GraphFrame(vertices, edges) + + // Message uses Pregel.dst("id") - but since only id is used, dst join should be skipped + val result = graph.pregel + .setMaxIter(1) + .withVertexColumn("received", lit(0L), coalesce(Pregel.msg, col("received"))) + .sendMsgToDst(Pregel.src("id") + Pregel.dst("id")) // Uses dst.id only + .aggMsgs(sum(Pregel.msg)) + .run() + + // Each dst vertex receives src.id + dst.id from incoming edges + val received = result.sort("id").select("received").as[Long].collect() + assert(received(0) === 0L) // vertex 0: no incoming edges + assert(received(1) === 1L) // vertex 1: edge 0->1, receives 0 + 1 = 1 + assert(received(2) === 3L) // vertex 2: edge 1->2, receives 1 + 2 = 3 + assert(received(3) === 5L) // vertex 3: edge 2->3, receives 2 + 3 = 5 + } + + // ============================================================================ + // Integration tests for complex expression patterns + // These verify that dst join is correctly performed when dst columns are used + // in non-trivial ways (map keys, array indices, conditionals, nested structs) + // ============================================================================ + + test("dst join required when dst column used in conditional") { + // when(Pregel.dst("flag"), Pregel.src("value")) - dst.flag requires the join + val vertices = Seq((0L, true, 10L), (1L, false, 20L), (2L, true, 30L)) + .toDF("id", "flag", "value") + val edges = Seq((0L, 1L), (1L, 2L)).toDF("src", "dst") + val graph = GraphFrame(vertices, edges) + + val result = graph.pregel + .setMaxIter(1) + .withVertexColumn("received", lit(0L), coalesce(Pregel.msg, col("received"))) + .sendMsgToDst(when(Pregel.dst("flag"), Pregel.src("value")).otherwise(lit(null))) + .aggMsgs(sum(Pregel.msg)) + .run() + + // Verify correct behavior: message only sent when dst.flag is true + val received = result.sort("id").select("received").as[Long].collect() + assert(received(0) === 0L) // vertex 0: no incoming + assert(received(1) === 0L) // vertex 1: dst.flag=false, so null message (filtered) + assert(received(2) === 20L) // vertex 2: dst.flag=true, receives src.value=20 + } + + test("dst join required when dst column used as map key") { + // Create edges with a map column, use dst vertex attribute as key + val vertices = Seq((0L, "a"), (1L, "b"), (2L, "a")).toDF("id", "key") + val edges = Seq((0L, 1L, Map("a" -> 10L, "b" -> 20L)), (1L, 2L, Map("a" -> 30L, "b" -> 40L))) + .toDF("src", "dst", "weights") + val graph = GraphFrame(vertices, edges) + + val result = graph.pregel + .requiredEdgeColumns("weights") + .setMaxIter(1) + .withVertexColumn("received", lit(0L), coalesce(Pregel.msg, col("received"))) + // Use dst.key to look up value in edge.weights map + .sendMsgToDst(element_at(Pregel.edge("weights"), Pregel.dst("key"))) + .aggMsgs(sum(Pregel.msg)) + .run() + + // Verify: dst.key is used to index into map + val received = result.sort("id").select("received").as[Long].collect() + assert(received(0) === 0L) // vertex 0: no incoming + assert(received(1) === 20L) // vertex 1: key="b", edge weights has b->20 + assert(received(2) === 30L) // vertex 2: key="a", edge weights has a->30 + } + + test("dst join required when dst column used as array index") { + // Create edges with array column, use dst vertex attribute as index + val vertices = Seq((0L, 1), (1L, 2), (2L, 1)).toDF("id", "idx") + val edges = Seq((0L, 1L, Array(100L, 200L)), (1L, 2L, Array(300L, 400L))) + .toDF("src", "dst", "values") + val graph = GraphFrame(vertices, edges) + + val result = graph.pregel + .requiredEdgeColumns("values") + .setMaxIter(1) + .withVertexColumn("received", lit(0L), coalesce(Pregel.msg, col("received"))) + // Use dst.idx to index into edge.values array (element_at is 1-based) + .sendMsgToDst(element_at(Pregel.edge("values"), Pregel.dst("idx"))) + .aggMsgs(sum(Pregel.msg)) + .run() + + // Verify: dst.idx is used to index into array + val received = result.sort("id").select("received").as[Long].collect() + assert(received(0) === 0L) // vertex 0: no incoming + assert(received(1) === 200L) // vertex 1: idx=2, array element 2 = 200 + assert(received(2) === 300L) // vertex 2: idx=1, array element 1 = 300 + } + + test("dst join required for nested struct field access") { + // Create vertices with nested struct + val vertices = spark + .createDataFrame(Seq((0L, 1.0, 2.0), (1L, 3.0, 4.0), (2L, 5.0, 6.0))) + .toDF("id", "x", "y") + .selectExpr("id", "named_struct('x', x, 'y', y) as location") + + val edges = Seq((0L, 1L), (1L, 2L)).toDF("src", "dst") + val graph = GraphFrame(vertices, edges) + + val result = graph.pregel + .setMaxIter(1) + .withVertexColumn("received", lit(0.0), coalesce(Pregel.msg, col("received"))) + // Access nested field dst.location.x and src.location.y + .sendMsgToDst(Pregel.dst("location")("x") + Pregel.src("location")("y")) + .aggMsgs(sum(Pregel.msg)) + .run() + + // Verify: nested struct fields are accessed correctly + val received = result.sort("id").select("received").as[Double].collect() + assert(received(0) === 0.0 +- 1e-6) // vertex 0: no incoming + assert(received(1) === 5.0 +- 1e-6) // vertex 1: dst.location.x=3.0 + src.location.y=2.0 + assert(received(2) === 9.0 +- 1e-6) // vertex 2: dst.location.x=5.0 + src.location.y=4.0 + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/RandomizedContractionSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/RandomizedContractionSuite.scala new file mode 100644 index 0000000000000..3a51cc815da50 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/RandomizedContractionSuite.scala @@ -0,0 +1,294 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.FunctionIdentifier +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.lit +import org.apache.spark.storage.StorageLevel +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.examples.Graphs + +class RandomizedContractionSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + test("RandomizedContraction: empty graph") { + val graph = Graphs.empty[Long] + val components = RandomizedContraction.run( + inputGraph = graph, + useLabelsAsComponents = false, + intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK, + useLocalCheckpoints = true, + checkpointInterval = 1, + isGraphPrepared = false) + assert(components.count() === 0L) + assertFunctionRegistryClean() + } + + test("RandomizedContraction: single isolated vertex") { + val vertices = spark.createDataFrame(List((0L, "a", "b"))).toDF("id", "vattr", "gender") + val e = + spark.createDataFrame(List((0L, 0L, 1L))).toDF("src", "dst", "test").filter("src > 10") + val graph = GraphFrame(vertices, e) + val components = RandomizedContraction.run( + inputGraph = graph, + useLabelsAsComponents = false, + intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK, + useLocalCheckpoints = true, + checkpointInterval = 1, + isGraphPrepared = false) + assert(components.count() === 1L) + assert(components.select("id", "component").collect().toSet === Set(Row(0L, 0L))) + assertFunctionRegistryClean() + } + + test("RandomizedContraction: two connected vertices") { + val vertices = + spark.createDataFrame(List((0L, "a0", "b0"), (1L, "a1", "b1"))).toDF("id", "A", "B") + val edges = spark.createDataFrame(List((0L, 1L, "a01", "b01"))).toDF("src", "dst", "A", "B") + val graph = GraphFrame(vertices, edges) + val components = RandomizedContraction.run( + inputGraph = graph, + useLabelsAsComponents = false, + intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK, + useLocalCheckpoints = true, + checkpointInterval = 1, + isGraphPrepared = false) + assert(components.count() === 2L) + val compValues = components.select("id", "component").collect() + assert(compValues.map(_.getLong(1)).toSet.size === 1) + assert(compValues.map(_.getLong(0)) === Array(0L, 1L)) + assertFunctionRegistryClean() + } + + test("RandomizedContraction: chain graph") { + val n = 5L + val graph = Graphs.chain(n) + val components = RandomizedContraction.run( + inputGraph = graph, + useLabelsAsComponents = false, + intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK, + useLocalCheckpoints = true, + checkpointInterval = 1, + isGraphPrepared = false) + assert(components.count() === n) + assert(components.select("component").distinct().count() === 1L) + assertFunctionRegistryClean() + } + + test("RandomizedContraction: disconnected vertices") { + val n = 5L + val vertices = spark.range(n).toDF(GraphFrame.ID) + val edges = + spark.createDataFrame(Seq.empty[(Long, Long)]).toDF(GraphFrame.SRC, GraphFrame.DST) + val graph = GraphFrame(vertices, edges) + val components = RandomizedContraction.run( + inputGraph = graph, + useLabelsAsComponents = false, + intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK, + useLocalCheckpoints = true, + checkpointInterval = 1, + isGraphPrepared = false) + assert(components.count() === n) + assert(components.select("component").distinct().count() === n) + assertFunctionRegistryClean() + } + + test("RandomizedContraction: two separate components") { + val vertices = spark.range(6L).toDF(GraphFrame.ID) + val edges = spark + .createDataFrame(Seq((0L, 1L), (1L, 2L), (2L, 0L), (3L, 4L), (4L, 5L), (5L, 3L))) + .toDF(GraphFrame.SRC, GraphFrame.DST) + val graph = GraphFrame(vertices, edges) + val components = RandomizedContraction.run( + inputGraph = graph, + useLabelsAsComponents = false, + intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK, + useLocalCheckpoints = true, + checkpointInterval = 1, + isGraphPrepared = false) + assert(components.count() === 6L) + val compGroups = + components.groupBy("component").count().collect().map(r => r.getLong(1)).toSet + assert(compGroups === Set(3L, 3L)) + assertFunctionRegistryClean() + } + + test("RandomizedContraction: with dangling vertices") { + val vertices = spark.range(8L).toDF(GraphFrame.ID) + val edges = spark + .createDataFrame(Seq((0L, 1L), (1L, 2L), (2L, 0L), (3L, 4L), (4L, 5L), (5L, 3L))) + .toDF(GraphFrame.SRC, GraphFrame.DST) + val graph = GraphFrame(vertices, edges) + val components = RandomizedContraction.run( + inputGraph = graph, + useLabelsAsComponents = false, + intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK, + useLocalCheckpoints = true, + checkpointInterval = 1, + isGraphPrepared = false) + assert(components.count() === 8L) + val compCounts = + components.groupBy("component").count().collect().map(r => r.getLong(1)).toSet + assert(compCounts === Set(1L, 1L, 3L, 3L)) + assertFunctionRegistryClean() + } + + test("RandomizedContraction: useLabelsAsComponents with string IDs") { + val vertices = + spark.createDataFrame(Seq("a", "b", "c", "d").map(Tuple1.apply)).toDF(GraphFrame.ID) + val edges = + spark.createDataFrame(Seq(("a", "b"), ("b", "c"))).toDF(GraphFrame.SRC, GraphFrame.DST) + val graph = GraphFrame(vertices, edges) + val components = RandomizedContraction.run( + inputGraph = graph, + useLabelsAsComponents = true, + intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK, + useLocalCheckpoints = true, + checkpointInterval = 1, + isGraphPrepared = false) + assert(components.count() === 4L) + val compIds = components.select("component").collect().map(_.getString(0)).toSet + assert(compIds.size === 2) + assert(compIds == Set("a", "d")) + assertFunctionRegistryClean() + } + + test("RandomizedContraction: useLabelsAsComponents with long IDs") { + val vertices = + spark.createDataFrame(Seq(1L, 2L, 3L, 4L).map(Tuple1.apply)).toDF(GraphFrame.ID) + val edges = + spark.createDataFrame(Seq((1L, 2L), (2L, 3L))).toDF(GraphFrame.SRC, GraphFrame.DST) + val graph = GraphFrame(vertices, edges) + val components = RandomizedContraction.run( + inputGraph = graph, + useLabelsAsComponents = true, + intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK, + useLocalCheckpoints = true, + checkpointInterval = 1, + isGraphPrepared = false) + assert(components.count() === 4L) + val compIds = components.select("component").collect().map(_.getLong(0)).toSet + assert(compIds.size === 2) + assert(compIds == Set(1L, 4L)) + assertFunctionRegistryClean() + } + + test("RandomizedContraction: no parquet file leaks") { + val graph = Graphs.chain(3L) + val initialParquetFiles = listParquetFiles() + + val components = RandomizedContraction.run( + inputGraph = graph, + useLabelsAsComponents = false, + intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK, + useLocalCheckpoints = true, + checkpointInterval = 1, + isGraphPrepared = false) + components.count() + + val finalParquetFiles = listParquetFiles() + assert(finalParquetFiles === initialParquetFiles) + assertFunctionRegistryClean() + } + + test("RandomizedContraction: no memory leaks") { + val priorCached = spark.sparkContext.getPersistentRDDs + + val graph = Graphs.chain(10L) + val components = RandomizedContraction.run( + inputGraph = graph, + useLabelsAsComponents = false, + intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK, + useLocalCheckpoints = false, + checkpointInterval = 1, + isGraphPrepared = false) + components.count() + components.unpersist() + + val postCached = spark.sparkContext.getPersistentRDDs + assert(postCached.size === priorCached.size) + assertFunctionRegistryClean() + } + + test("RandomizedContraction: large long IDs") { + val max = Long.MaxValue + val chain = Graphs.chain(10L) + val vertices = chain.vertices.select((col(GraphFrame.ID) - lit(max)).as(GraphFrame.ID)) + val edges = chain.edges.select( + (col(GraphFrame.SRC) - lit(max)).as(GraphFrame.SRC), + (col(GraphFrame.DST) - lit(max)).as(GraphFrame.DST)) + val graph = GraphFrame(vertices, edges) + val components = RandomizedContraction.run( + inputGraph = graph, + useLabelsAsComponents = false, + intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK, + useLocalCheckpoints = true, + checkpointInterval = 1, + isGraphPrepared = false) + assert(components.count() === 10L) + assert(components.select("component").distinct().count() === 1L) + assertFunctionRegistryClean() + } + + test("RandomizedContraction: directed edges still produce connected components") { + val vertices = spark.range(5L).toDF(GraphFrame.ID) + val edges = spark + .createDataFrame(Seq((0L, 4L), (4L, 3L), (2L, 3L), (2L, 1L))) + .toDF(GraphFrame.SRC, GraphFrame.DST) + val graph = GraphFrame(vertices, edges) + val components = RandomizedContraction.run( + inputGraph = graph, + useLabelsAsComponents = false, + intermediateStorageLevel = StorageLevel.MEMORY_AND_DISK, + useLocalCheckpoints = true, + checkpointInterval = 1, + isGraphPrepared = false) + assert(components.count() === 5L) + assert(components.select("component").distinct().count() === 1L) + assertFunctionRegistryClean() + } + + private def assertFunctionRegistryClean(): Unit = { + val functionRegistry = spark.sessionState.functionRegistry + val identifier = new FunctionIdentifier("_axpb", Some("builtin"), Some("system")) + val _ = assert(!functionRegistry.functionExists(identifier)) + } + + private def listParquetFiles(): Set[String] = { + val hadoopConf = spark.sparkContext.hadoopConfiguration + val fs = org.apache.hadoop.fs.FileSystem.get(hadoopConf) + val rootPath = new org.apache.hadoop.fs.Path(spark.conf.get("spark.sql.warehouse.dir")) + + def listFiles(path: org.apache.hadoop.fs.Path): Set[String] = { + if (fs.exists(path)) { + fs.listStatus(path) + .flatMap { status => + if (status.isDirectory) listFiles(status.getPath) + else if (status.getPath.getName.endsWith(".parquet")) Set(status.getPath.toString) + else Set.empty[String] + } + .toSet + } else Set.empty[String] + } + + listFiles(rootPath) + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/SVDPlusPlusSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/SVDPlusPlusSuite.scala new file mode 100644 index 0000000000000..c4b4957548ffb --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/SVDPlusPlusSuite.scala @@ -0,0 +1,104 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.types.DataTypes +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.TestUtils +import org.apache.spark.graphframes.examples.Graphs + +class SVDPlusPlusSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + test("Test SVD++ with mean square error on training set") { + val svdppErr = 8.0 + val g = Graphs.ALSSyntheticData() + + val v2 = g.svdPlusPlus.maxIter(2).run() + TestUtils.testSchemaInvariants(g, v2) + Seq(SVDPlusPlus.COLUMN1, SVDPlusPlus.COLUMN2).foreach { case c => + TestUtils.checkColumnType( + v2.schema, + c, + DataTypes.createArrayType(DataTypes.DoubleType, false)) + } + Seq(SVDPlusPlus.COLUMN3, SVDPlusPlus.COLUMN4).foreach { case c => + TestUtils.checkColumnType(v2.schema, c, DataTypes.DoubleType) + } + val err = v2 + .select(GraphFrame.ID, SVDPlusPlus.COLUMN4) + .rdd + .map { + case Row(vid: Long, vd: Double) => + if (vid % 2 == 1) vd else 0.0 + case _ => throw new GraphFramesUnreachableException() + } + .reduce(_ + _) / g.edges.count() + assert(err <= svdppErr) + v2.unpersist() + } + + Seq( + ("int", "float"), + ("short", "double"), + ("long", "float"), + ("byte", "double"), + ("string", "float")).foreach(types => + test(s"Test SVD++ with mean square error on training set, ${types._1}/${types._2} types") { + val svdppErr = 8.0 + val g = { + val gg = Graphs.ALSSyntheticData() + GraphFrame( + gg.vertices.select(col(GraphFrame.ID).cast(types._1)), + gg.edges.select( + col(GraphFrame.SRC).cast(types._1), + col(GraphFrame.DST).cast(types._1), + col("weight").cast(types._2))) + } + + val v2 = g.svdPlusPlus.maxIter(2).run() + TestUtils.testSchemaInvariants(g, v2) + Seq(SVDPlusPlus.COLUMN1, SVDPlusPlus.COLUMN2).foreach { case c => + TestUtils.checkColumnType( + v2.schema, + c, + DataTypes.createArrayType(DataTypes.DoubleType, false)) + } + Seq(SVDPlusPlus.COLUMN3, SVDPlusPlus.COLUMN4).foreach { case c => + TestUtils.checkColumnType(v2.schema, c, DataTypes.DoubleType) + } + val err = v2 + .select(GraphFrame.ID, SVDPlusPlus.COLUMN4) + .rdd + .map { row => + { + val vid = if (types._1 == "string") { row.getAs[String](0).toLong } + else { row.getAs[Number](0).longValue() } + val vd = row.getAs[Number](1).doubleValue() + if (vid % 2 == 1) vd else 0.0 + } + } + .reduce(_ + _) / g.edges.count() + assert(err <= svdppErr) + v2.unpersist() + }) +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ShortestPathsSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ShortestPathsSuite.scala new file mode 100644 index 0000000000000..7c5a708a6ed7d --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ShortestPathsSuite.scala @@ -0,0 +1,180 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.types.DataTypes +import org.apache.spark.graphframes._ +import org.apache.spark.graphframes.GraphFrame.quote + +class ShortestPathsSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + test("Simple test") { + doTest() + } + + def doTest(vertices: Option[DataFrame] = Option.empty): Unit = { + val spark = this.spark + import spark.implicits._ + + val edgeSeq = Seq((1, 2), (1, 5), (2, 3), (2, 5), (3, 4), (4, 5), (4, 6)) + .flatMap { case e => + Seq(e, e.swap) + } + .map { case (src, dst) => (src.toLong, dst.toLong) } + val edges = spark.createDataFrame(edgeSeq).toDF("src", "dst") + val graph = vertices.map(GraphFrame(_, edges)).getOrElse(GraphFrame.fromEdges(edges)) + + // Ground truth + val shortestPaths = Seq( + (1, Map(1 -> 0, 4 -> 2)), + (2, Map(1 -> 1, 4 -> 2)), + (3, Map(1 -> 2, 4 -> 1)), + (4, Map(1 -> 2, 4 -> 0)), + (5, Map(1 -> 1, 4 -> 1)), + (6, Map(1 -> 3, 4 -> 1))).toDF("id", "distances") + val expectedCols = vertices.map(_.columns.toSeq).getOrElse(Seq("id")) :+ "distances" + val expected = vertices + .foldLeft(shortestPaths) { case (shortestPaths, vertices) => + shortestPaths.join(vertices, "id") + } + .select(expectedCols.map(quote).map(col): _*) + .collect() + .toSet + + val landmarks = Seq(1, 4).map(_.toLong) + val v2 = graph.shortestPaths.landmarks(landmarks).run() + + TestUtils.testSchemaInvariants(graph, v2) + TestUtils.checkColumnType( + v2.schema, + "distances", + DataTypes.createMapType(v2.schema("id").dataType, DataTypes.IntegerType, false)) + val results = v2.collect().toSet + assert(results === expected) + v2.unpersist() + () + } + + test("Simple test with GraphFrames") { + val edgeSeq = Seq((1, 2), (1, 5), (2, 3), (2, 5), (3, 4), (4, 5), (4, 6)) + .flatMap { case e => + Seq(e, e.swap) + } + .map { case (src, dst) => (src.toLong, dst.toLong) } + val edges = spark.createDataFrame(edgeSeq).toDF("src", "dst") + val graph = GraphFrame.fromEdges(edges) + + // Ground truth + val shortestPaths = Set( + (1, Map(1 -> 0, 4 -> 2)), + (2, Map(1 -> 1, 4 -> 2)), + (3, Map(1 -> 2, 4 -> 1)), + (4, Map(1 -> 2, 4 -> 0)), + (5, Map(1 -> 1, 4 -> 1)), + (6, Map(1 -> 3, 4 -> 1))) + + val landmarks = Seq(1, 4).map(_.toLong) + val v2 = graph.shortestPaths.landmarks(landmarks).setAlgorithm("graphframes").run() + + TestUtils.testSchemaInvariants(graph, v2) + TestUtils.checkColumnType( + v2.schema, + "distances", + DataTypes.createMapType(v2.schema("id").dataType, DataTypes.IntegerType, true)) + val newVs = v2.select("id", "distances").collect().toSeq + val results = newVs.map { + case Row(id: Long, spMap: Map[Long, Int] @unchecked) => + (id, spMap) + case _ => throw new GraphFramesUnreachableException() + } + assert(results.toSet === shortestPaths) + v2.unpersist() + } + + test("friends graph") { + val friends = examples.Graphs.friends + val v = friends.shortestPaths.landmarks(Seq("a", "d")).run() + val expected = Set[(String, Map[String, Int])]( + ("a", Map("a" -> 0, "d" -> 2)), + ("b", Map.empty), + ("c", Map.empty), + ("d", Map("a" -> 1, "d" -> 0)), + ("e", Map("a" -> 2, "d" -> 1)), + ("f", Map.empty), + ("g", Map.empty)) + val results = v + .select("id", "distances") + .collect() + .map { + case Row(id: String, spMap: Map[String, Int] @unchecked) => + (id, spMap) + case _ => throw new GraphFramesUnreachableException() + } + .toSet + assert(results === expected) + v.unpersist() + } + + test("friends graph with GraphFrames") { + val friends = examples.Graphs.friends + val v = friends.shortestPaths.landmarks(Seq("a", "d")).setAlgorithm("graphframes").run() + val expected = Set[(String, Map[String, Int])]( + ("a", Map("a" -> 0, "d" -> 2)), + ("b", Map.empty), + ("c", Map.empty), + ("d", Map("a" -> 1, "d" -> 0)), + ("e", Map("a" -> 2, "d" -> 1)), + ("f", Map.empty), + ("g", Map.empty)) + val results = v + .select("id", "distances") + .collect() + .map { + case Row(id: String, spMap: Map[String, Int] @unchecked) => + (id, spMap) + case _ => throw new GraphFramesUnreachableException() + } + .toSet + assert(results === expected) + v.unpersist() + } + test("Test vertices with column name") { + val verticeSeq = + Seq((1L, "one"), (2L, "two"), (3L, "three"), (4L, "four"), (5L, "five"), (6L, "six")) + val vertices = sqlContext.createDataFrame(verticeSeq).toDF("id", "name") + doTest(Some(vertices)) + } + + test("Test vertices with dot column name") { + val verticeSeq = + Seq((1L, "one"), (2L, "two"), (3L, "three"), (4L, "four"), (5L, "five"), (6L, "six")) + val vertices = sqlContext.createDataFrame(verticeSeq).toDF("id", "a.name") + doTest(Some(vertices)) + } + + test("Test vertices with backquote in column name") { + val verticeSeq = + Seq((1L, "one"), (2L, "two"), (3L, "three"), (4L, "four"), (5L, "five"), (6L, "six")) + val vertices = sqlContext.createDataFrame(verticeSeq).toDF("id", "a `name`") + doTest(Some(vertices)) + } + +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponentsSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponentsSuite.scala new file mode 100644 index 0000000000000..b7c9811754c09 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponentsSuite.scala @@ -0,0 +1,42 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.DataTypes +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.TestUtils + +class StronglyConnectedComponentsSuite extends SparkFunSuite with GraphFrameTestSparkContext { + test("Island Strongly Connected Components") { + val vertices = spark + .createDataFrame(Seq((1L, "a"), (2L, "b"), (3L, "c"), (4L, "d"), (5L, "e"))) + .toDF("id", "value") + val edges = spark.createDataFrame(Seq.empty[(Long, Long)]).toDF("src", "dst") + val graph = GraphFrame(vertices, edges) + val c = graph.stronglyConnectedComponents.maxIter(5).run() + TestUtils.testSchemaInvariants(graph, c) + TestUtils.checkColumnType(c.schema, "component", DataTypes.LongType) + for (Row(id: Long, component: Long, _) <- c.select("id", "component", "value").collect()) { + assert(id === component) + } + c.unpersist() + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala new file mode 100644 index 0000000000000..6c12ff4be9e3b --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala @@ -0,0 +1,360 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.types.DataTypes +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.TestUtils + +class StructureAwareLabelPropagationSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + test("basic flow: one iteration propagates strongest incoming label") { + assume(TestUtils.requireSparkVersionGE(4, 1, spark.version)) + + val vertices = spark.createDataFrame(Seq(1L, 2L, 3L).map(Tuple1(_))).toDF("id") + val edges = spark.createDataFrame(Seq((1L, 2L), (2L, 3L), (3L, 1L))).toDF("src", "dst") + val g = GraphFrame(vertices, edges) + + val result = new StructureAwareLabelPropagation(g) + .maxIter(1) + .setIgnoreDirectLinks(false) + .setStructuralSimilarityMultiplier(0.5) + .run() + + TestUtils.testSchemaInvariants(g, result) + TestUtils.checkColumnType(result.schema, "label", DataTypes.LongType) + + val labels = result + .select("id", "label") + .collect() + .map { row => + row.getLong(0) -> row.getLong(1) + } + .toMap + + assert(labels === Map(1L -> 3L, 2L -> 1L, 3L -> 2L)) + + result.unpersist() + } + + test( + "different structuralSimilarityMultiplier values can change winner between direct-link mass and common-neighbor overlap") { + assume(TestUtils.requireSparkVersionGE(4, 1, spark.version)) + + val vertices = spark + .createDataFrame(Seq((1L, "A"), (2L, "B"), (3L, "B"), (4L, "T"), (7L, "X"), (8L, "Y"))) + .toDF("id", "initLabel") + + // For destination 4 when direct links are included: + // - src 1 has overlap 2 with dst 4 via neighbors {7, 8} + // - src 2 and src 3 each have overlap 0 with dst 4 + // Messages are aggregated by label: + // label A total = 1 + 2m + // label B total = 1 + 1 = 2 + // where m = structuralSimilarityMultiplier. + // So smaller m favors B; larger m favors A. + val edges = spark + .createDataFrame(Seq((1L, 4L), (2L, 4L), (3L, 4L), (4L, 7L), (4L, 8L), (1L, 7L), (1L, 8L))) + .toDF("src", "dst") + + val g = GraphFrame(vertices, edges) + + val lowMultiplier = new StructureAwareLabelPropagation(g) + .maxIter(1) + .setInitialLabelCol("initLabel") + .setIgnoreDirectLinks(false) + .setStructuralSimilarityMultiplier(0.1) + .run() + + val highMultiplier = new StructureAwareLabelPropagation(g) + .maxIter(1) + .setInitialLabelCol("initLabel") + .setIgnoreDirectLinks(false) + .setStructuralSimilarityMultiplier(1.0) + .run() + + TestUtils.checkColumnType(lowMultiplier.schema, "label", DataTypes.StringType) + TestUtils.checkColumnType(highMultiplier.schema, "label", DataTypes.StringType) + + val labelWithLowMultiplier = + lowMultiplier.filter("id = 4").select("label").collect().head.getString(0) + val labelWithHighMultiplier = + highMultiplier.filter("id = 4").select("label").collect().head.getString(0) + + assert(labelWithLowMultiplier === "B") + assert(labelWithHighMultiplier === "A") + + lowMultiplier.unpersist() + highMultiplier.unpersist() + } + + test("isolated vertex keeps its own ID label") { + assume(TestUtils.requireSparkVersionGE(4, 1, spark.version)) + + val vertices = spark.createDataFrame(Seq(1L, 2L, 99L).map(Tuple1(_))).toDF("id") + val edges = spark.createDataFrame(Seq((1L, 2L))).toDF("src", "dst") + val g = GraphFrame(vertices, edges) + + val result = new StructureAwareLabelPropagation(g) + .maxIter(3) + .setIgnoreDirectLinks(false) + .setStructuralSimilarityMultiplier(0.5) + .run() + + assert(result.count() == 3) + val isolatedLabel = + result.filter(col("id") === lit(99L)).select("label").collect().head.getLong(0) + assert(isolatedLabel === 99L) + + result.unpersist() + } + + test("disconnected graph propagates labels independently per component") { + assume(TestUtils.requireSparkVersionGE(4, 1, spark.version)) + + val vertices = spark.createDataFrame(Seq(1L, 2L, 3L, 10L, 11L, 12L).map(Tuple1(_))).toDF("id") + val edges = spark + .createDataFrame(Seq((1L, 2L), (2L, 3L), (10L, 11L), (11L, 12L))) + .toDF("src", "dst") + val g = GraphFrame(vertices, edges) + + val result = new StructureAwareLabelPropagation(g) + .maxIter(1) + .setIgnoreDirectLinks(false) + .setStructuralSimilarityMultiplier(0.5) + .run() + + val labels = result + .select("id", "label") + .collect() + .map { row => + row.getLong(0) -> row.getLong(1) + } + .toMap + + assert(labels === Map(1L -> 1L, 2L -> 1L, 3L -> 2L, 10L -> 10L, 11L -> 10L, 12L -> 11L)) + + result.unpersist() + } + + test("changing only structuralSimilarityMultiplier can flip the winning label") { + assume(TestUtils.requireSparkVersionGE(4, 1, spark.version)) + + val vertices = spark + .createDataFrame(Seq((1L, "A"), (2L, "B"), (3L, "B"), (4L, "T"), (7L, "X"), (8L, "Y"))) + .toDF("id", "initLabel") + + // For destination 4 when direct links are included: + // label A score = 1 + 2m (from src 1) + // label B score = 2 (from src 2 and src 3) + // Keep direct-link handling fixed and vary only m. + val edges = spark + .createDataFrame(Seq((1L, 4L), (2L, 4L), (3L, 4L), (4L, 7L), (4L, 8L), (1L, 7L), (1L, 8L))) + .toDF("src", "dst") + + val g = GraphFrame(vertices, edges) + + val lowMultiplier = new StructureAwareLabelPropagation(g) + .maxIter(1) + .setInitialLabelCol("initLabel") + .setIgnoreDirectLinks(false) + .setStructuralSimilarityMultiplier(0.1) + .run() + + val highMultiplier = new StructureAwareLabelPropagation(g) + .maxIter(1) + .setInitialLabelCol("initLabel") + .setIgnoreDirectLinks(false) + .setStructuralSimilarityMultiplier(1.0) + .run() + + val labelWithLowMultiplier = + lowMultiplier.filter("id = 4").select("label").collect().head.getString(0) + val labelWithHighMultiplier = + highMultiplier.filter("id = 4").select("label").collect().head.getString(0) + + assert(labelWithLowMultiplier === "B") + assert(labelWithHighMultiplier === "A") + + lowMultiplier.unpersist() + highMultiplier.unpersist() + } + + test("changing ignoreDirectLinks can flip the winning label") { + assume(TestUtils.requireSparkVersionGE(4, 1, spark.version)) + + val vertices = spark + .createDataFrame(Seq((1L, "A"), (2L, "B"), (3L, "B"), (4L, "T"), (7L, "X"), (8L, "Y"))) + .toDF("id", "initLabel") + + // For destination 4 with m = 0.25: + // ignoreDirectLinks = false: + // label A score = 1 + 2m = 1.5 + // label B score = 2 + // ignoreDirectLinks = true: + // label A score = 2m = 0.5 + // label B score = 0 + val edges = spark + .createDataFrame(Seq((1L, 4L), (2L, 4L), (3L, 4L), (4L, 7L), (4L, 8L), (1L, 7L), (1L, 8L))) + .toDF("src", "dst") + + val g = GraphFrame(vertices, edges) + + val includeDirectLinks = new StructureAwareLabelPropagation(g) + .maxIter(1) + .setInitialLabelCol("initLabel") + .setIgnoreDirectLinks(false) + .setStructuralSimilarityMultiplier(0.25) + .run() + + val ignoreDirectLinks = new StructureAwareLabelPropagation(g) + .maxIter(1) + .setInitialLabelCol("initLabel") + .setIgnoreDirectLinks(true) + .setStructuralSimilarityMultiplier(0.25) + .run() + + val labelWithDirectLinks = + includeDirectLinks.filter("id = 4").select("label").collect().head.getString(0) + val labelWithoutDirectLinks = + ignoreDirectLinks.filter("id = 4").select("label").collect().head.getString(0) + + assert(labelWithDirectLinks === "B") + assert(labelWithoutDirectLinks === "A") + + includeDirectLinks.unpersist() + ignoreDirectLinks.unpersist() + } + + test("setStructuralSimilarityMultiplier allows zero but rejects negative values") { + assume(TestUtils.requireSparkVersionGE(4, 1, spark.version)) + + val vertices = spark.createDataFrame(Seq(1L, 2L).map(Tuple1(_))).toDF("id") + val edges = spark.createDataFrame(Seq((1L, 2L))).toDF("src", "dst") + val g = GraphFrame(vertices, edges) + + new StructureAwareLabelPropagation(g).setStructuralSimilarityMultiplier(0.0) + + intercept[IllegalArgumentException] { + new StructureAwareLabelPropagation(g).setStructuralSimilarityMultiplier(-0.1) + } + } + + test("zero structuralSimilarityMultiplier is invalid when ignoreDirectLinks is true") { + assume(TestUtils.requireSparkVersionGE(4, 1, spark.version)) + + val vertices = spark.createDataFrame(Seq(1L, 2L).map(Tuple1(_))).toDF("id") + val edges = spark.createDataFrame(Seq((1L, 2L))).toDF("src", "dst") + val g = GraphFrame(vertices, edges) + + intercept[IllegalArgumentException] { + new StructureAwareLabelPropagation(g) + .maxIter(1) + .setIgnoreDirectLinks(true) + .setStructuralSimilarityMultiplier(0.0) + .run() + } + } + + test("setIsDirected(false) changes propagation by adding reverse links") { + assume(TestUtils.requireSparkVersionGE(4, 1, spark.version)) + + val vertices = spark.createDataFrame(Seq(1L, 2L).map(Tuple1(_))).toDF("id") + val edges = spark.createDataFrame(Seq((1L, 2L))).toDF("src", "dst") + val g = GraphFrame(vertices, edges) + + val directed = new StructureAwareLabelPropagation(g) + .maxIter(1) + .setIgnoreDirectLinks(false) + .setStructuralSimilarityMultiplier(0.5) + .setIsDirected(true) + .run() + + val undirected = new StructureAwareLabelPropagation(g) + .maxIter(1) + .setIgnoreDirectLinks(false) + .setStructuralSimilarityMultiplier(0.5) + .setIsDirected(false) + .run() + + val directedLabels = directed + .select("id", "label") + .collect() + .map(r => r.getLong(0) -> r.getLong(1)) + .toMap + val undirectedLabels = undirected + .select("id", "label") + .collect() + .map(r => r.getLong(0) -> r.getLong(1)) + .toMap + + assert(directedLabels === Map(1L -> 1L, 2L -> 1L)) + assert(undirectedLabels === Map(1L -> 2L, 2L -> 1L)) + + directed.unpersist() + undirected.unpersist() + } + + test("undirected mode matches explicitly symmetrized directed edge set") { + assume(TestUtils.requireSparkVersionGE(4, 1, spark.version)) + + val vertices = spark.createDataFrame(Seq(1L, 2L, 3L, 4L).map(Tuple1(_))).toDF("id") + val directedEdges = + spark.createDataFrame(Seq((1L, 2L), (2L, 3L), (2L, 4L))).toDF("src", "dst") + val symmetrizedEdges = spark + .createDataFrame(Seq((1L, 2L), (2L, 1L), (2L, 3L), (3L, 2L), (2L, 4L), (4L, 2L))) + .toDF("src", "dst") + + val g = GraphFrame(vertices, directedEdges) + val gSym = GraphFrame(vertices, symmetrizedEdges) + + val internalUndirected = new StructureAwareLabelPropagation(g) + .maxIter(1) + .setIgnoreDirectLinks(false) + .setStructuralSimilarityMultiplier(0.5) + .setIsDirected(false) + .run() + + val explicitSymDirected = new StructureAwareLabelPropagation(gSym) + .maxIter(1) + .setIgnoreDirectLinks(false) + .setStructuralSimilarityMultiplier(0.5) + .setIsDirected(true) + .run() + + val internalMap = internalUndirected + .select("id", "label") + .collect() + .map(r => r.getLong(0) -> r.getLong(1)) + .toMap + val explicitMap = explicitSymDirected + .select("id", "label") + .collect() + .map(r => r.getLong(0) -> r.getLong(1)) + .toMap + + assert(internalMap === explicitMap) + + internalUndirected.unpersist() + explicitSymDirected.unpersist() + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/TriangleCountSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/TriangleCountSuite.scala new file mode 100644 index 0000000000000..5147e9844d8d2 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/TriangleCountSuite.scala @@ -0,0 +1,218 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.lib + +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.DataTypes +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrame.quote +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.TestUtils + +class TriangleCountSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + test("Count a single triangle") { + val edges = spark.createDataFrame(Seq(0L -> 1L, 1L -> 2L, 2L -> 0L)).toDF("src", "dst") + val vertices = spark + .createDataFrame(Seq((0L, "a"), (1L, "b"), (2L, "c"))) + .toDF("id", "a") + val g = GraphFrame(vertices, edges) + val v2 = g.triangleCount.run() + TestUtils.testSchemaInvariants(g, v2) + TestUtils.checkColumnType(v2.schema, "count", DataTypes.LongType) + v2.select("id", "count", "a") + .collect() + .foreach { + case Row(_: Long, count: Long, _) => assert(count === 1) + case _: Row => throw new GraphFramesUnreachableException() + } + v2.unpersist() + } + + test("Count two triangles") { + val edges = spark + .createDataFrame( + Seq(0L -> 1L, 1L -> 2L, 2L -> 0L) ++ + Seq(0L -> -1L, -1L -> -2L, -2L -> 0L)) + .toDF("src", "dst") + val g = GraphFrame.fromEdges(edges) + val v2 = g.triangleCount.run() + v2.select("id", "count").collect().foreach { + case Row(id: Long, count: Long) => + if (id == 0) { + assert(count === 2) + } else { + assert(count === 1) + } + case _: Row => throw new GraphFramesUnreachableException() + } + v2.unpersist() + } + + test("Count one triangles with bi-directed edges") { + // Note: This is different from GraphX, which double-counts triangles with bidirected edges. + val triangles = Seq(0L -> 1L, 1L -> 2L, 2L -> 0L) ++ Seq(0L -> -1L, -1L -> -2L, -2L -> 0L) + val revTriangles = triangles.map { case (a, b) => (b, a) } + val edges = spark.createDataFrame(triangles ++ revTriangles).toDF("src", "dst") + val g = GraphFrame.fromEdges(edges) + val v2 = g.triangleCount.run() + v2.select("id", "count").collect().foreach { + case Row(id: Long, count: Long) => + if (id == 0) { + assert(count === 2) + } else { + assert(count === 1) + } + case _: Row => throw new GraphFramesUnreachableException() + } + v2.unpersist() + } + + test("Count a single triangle with duplicate edges") { + val edges = spark + .createDataFrame( + Seq(0L -> 1L, 1L -> 2L, 2L -> 0L) ++ + Seq(0L -> 1L, 1L -> 2L, 2L -> 0L)) + .toDF("src", "dst") + val g = GraphFrame.fromEdges(edges) + val v2 = g.triangleCount.run() + v2.select("id", "count").collect().foreach { + case Row(_: Long, count: Long) => + assert(count === 1) + case _: Row => throw new GraphFramesUnreachableException() + } + v2.unpersist() + } + + test("Count with dot column name") { + val edges = sqlContext.createDataFrame(Seq(0L -> 1L, 1L -> 2L, 2L -> 0L)).toDF("src", "dst") + val vertices = sqlContext + .createDataFrame(Seq((0L, "a"), (1L, "b"), (2L, "c"))) + .toDF("id", "a.column") + val g = GraphFrame(vertices, edges) + val v2 = g.triangleCount.run() + TestUtils.testSchemaInvariants(g, v2) + TestUtils.checkColumnType(v2.schema, "count", DataTypes.LongType) + v2.select("id", "count", quote("a.column")) + .collect() + .foreach { + case Row(_: Long, count: Long, _) => assert(count === 1) + case _: Row => throw new GraphFramesUnreachableException() + } + v2.unpersist() + } + + test("Count with backquote in column name") { + val edges = sqlContext.createDataFrame(Seq(0L -> 1L, 1L -> 2L, 2L -> 0L)).toDF("src", "dst") + val vertices = sqlContext + .createDataFrame(Seq((0L, "a"), (1L, "b"), (2L, "c"))) + .toDF("id", "a `column`") + val g = GraphFrame(vertices, edges) + val v2 = g.triangleCount.run() + TestUtils.testSchemaInvariants(g, v2) + TestUtils.checkColumnType(v2.schema, "count", DataTypes.LongType) + v2.select("id", "count", quote("a `column`")) + .collect() + .foreach { + case Row(_: Long, count: Long, _) => assert(count === 1) + case _: Row => throw new GraphFramesUnreachableException() + } + v2.unpersist() + } + + test("no triangle") { + val edges = spark.createDataFrame(Seq(0L -> 1L, 1L -> 2L)).toDF("src", "dst") + val g = GraphFrame.fromEdges(edges) + val v2 = g.triangleCount.run() + v2.select("count").collect().foreach { + case Row(count: Long) => + assert(count === 0) + case _: Row => throw new GraphFramesUnreachableException() + } + v2.unpersist() + } + + test("Approximate triangle count") { + assume(TestUtils.requireSparkVersionGE(4, 1, spark.version)) + + val edges = spark + .createDataFrame( + Seq(0L -> 1L, 1L -> 2L, 2L -> 0L) ++ + Seq(0L -> -1L, -1L -> -2L, -2L -> 0L)) + .toDF("src", "dst") + val g = GraphFrame.fromEdges(edges) + val v2 = g.triangleCount.setAlgorithm("approx").run() + + v2.select("id", "count").collect().foreach { + case Row(id: Long, count: Long) => + if (id == 0L) { + // Approx might have variation but for this small graph should be exact + assert(count >= 1) + } else { + assert(count >= 0) + } + case _ => throw new GraphFramesUnreachableException() + } + v2.unpersist() + } + + test("Approximate triangle count - no triangles") { + assume(TestUtils.requireSparkVersionGE(4, 1, spark.version)) + + val edges = spark.createDataFrame(Seq(0L -> 1L, 1L -> 2L, 3L -> 4L)).toDF("src", "dst") + val g = GraphFrame.fromEdges(edges) + val v2 = g.triangleCount.setAlgorithm("approx").run() + v2.select("count").collect().foreach { + case Row(count: Long) => assert(count === 0) + case _ => throw new GraphFramesUnreachableException() + } + v2.unpersist() + } + + test("Approximate triangle count - bipartite graph") { + assume(TestUtils.requireSparkVersionGE(4, 1, spark.version)) + + val edges = + spark.createDataFrame(Seq(0L -> 2L, 0L -> 3L, 1L -> 2L, 1L -> 3L)).toDF("src", "dst") + val g = GraphFrame.fromEdges(edges) + val v2 = g.triangleCount.setAlgorithm("approx").run() + v2.select("count").collect().foreach { + case Row(count: Long) => assert(count === 0) + case _ => throw new GraphFramesUnreachableException() + } + v2.unpersist() + } + + test("Approximate triangle count - large lgNomEntries") { + assume(TestUtils.requireSparkVersionGE(4, 1, spark.version)) + + val edges = spark.createDataFrame(Seq(0L -> 1L, 1L -> 2L, 2L -> 0L)).toDF("src", "dst") + val g = GraphFrame.fromEdges(edges) + val v2 = g.triangleCount + .setAlgorithm("approx") + .setLgNomEntries(16) + .run() + v2.select("count").collect().foreach { + case Row(count: Long) => assert(count === 1) + case _ => throw new GraphFramesUnreachableException() + } + v2.unpersist() + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/pattern/PatternSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/pattern/PatternSuite.scala new file mode 100644 index 0000000000000..7c1c9b200c57e --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/pattern/PatternSuite.scala @@ -0,0 +1,282 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.pattern + +import org.apache.spark.graphframes.InvalidParseException +import org.apache.spark.graphframes.SparkFunSuite + +class PatternSuite extends SparkFunSuite { + + test("good parses") { + assert(Pattern.parse("(abc)") === Seq(NamedVertex("abc"))) + + assert( + Pattern.parse("(u)-[e]->(v)") === + Seq(NamedEdge("e", NamedVertex("u"), NamedVertex("v")))) + + assert( + Pattern.parse("(u)-[e*1]->(v)") === + Seq(NamedEdge("_e1", NamedVertex("u"), NamedVertex("v")))) + + assert( + Pattern.parse("(u)-[e*3]->(v)") === + Seq( + NamedEdge("_e1", NamedVertex("u"), NamedVertex("_uv1")), + NamedEdge("_e2", NamedVertex("_uv1"), NamedVertex("_uv2")), + NamedEdge("_e3", NamedVertex("_uv2"), NamedVertex("v")))) + + assert( + Pattern.parse("(u)-[e*3]->(v);(v)-[l*2]->(w);(w)-[k*1]->(p)") === + Seq( + NamedEdge("_e1", NamedVertex("u"), NamedVertex("_uv1")), + NamedEdge("_e2", NamedVertex("_uv1"), NamedVertex("_uv2")), + NamedEdge("_e3", NamedVertex("_uv2"), NamedVertex("v")), + NamedEdge("_l1", NamedVertex("v"), NamedVertex("_vw1")), + NamedEdge("_l2", NamedVertex("_vw1"), NamedVertex("w")), + NamedEdge("_k1", NamedVertex("w"), NamedVertex("p")))) + + assert( + Pattern.parse("()-[]->(v)") === + Seq(AnonymousEdge(AnonymousVertex, NamedVertex("v")))) + + assert( + Pattern.parse("()-[e]->()") === + Seq(NamedEdge("e", AnonymousVertex, AnonymousVertex))) + + assert( + Pattern.parse("(u)-[e]->(u)") === + Seq(NamedEdge("e", NamedVertex("u"), NamedVertex("u")))) + + assert( + Pattern.parse("(u); ()-[]->(v)") === + Seq(NamedVertex("u"), AnonymousEdge(AnonymousVertex, NamedVertex("v")))) + + assert( + Pattern.parse("(u)-[]->(v); (v)-[]->(w); !(u)-[]->(w)") === + Seq( + AnonymousEdge(NamedVertex("u"), NamedVertex("v")), + AnonymousEdge(NamedVertex("v"), NamedVertex("w")), + Negation(AnonymousEdge(NamedVertex("u"), NamedVertex("w"))))) + + assert( + Pattern.parse("(u)-[*3]->(v)") === + Seq( + AnonymousEdge(NamedVertex("u"), NamedVertex("_uv1")), + AnonymousEdge(NamedVertex("_uv1"), NamedVertex("_uv2")), + AnonymousEdge(NamedVertex("_uv2"), NamedVertex("v")))) + + assert( + Pattern.parse("(u)-[*5]->(v)") === + Seq( + AnonymousEdge(NamedVertex("u"), NamedVertex("_uv1")), + AnonymousEdge(NamedVertex("_uv1"), NamedVertex("_uv2")), + AnonymousEdge(NamedVertex("_uv2"), NamedVertex("_uv3")), + AnonymousEdge(NamedVertex("_uv3"), NamedVertex("_uv4")), + AnonymousEdge(NamedVertex("_uv4"), NamedVertex("v")))) + + assert( + Pattern.parse("(u)-[*10]->(v)") === + Seq( + AnonymousEdge(NamedVertex("u"), NamedVertex("_uv1")), + AnonymousEdge(NamedVertex("_uv1"), NamedVertex("_uv2")), + AnonymousEdge(NamedVertex("_uv2"), NamedVertex("_uv3")), + AnonymousEdge(NamedVertex("_uv3"), NamedVertex("_uv4")), + AnonymousEdge(NamedVertex("_uv4"), NamedVertex("_uv5")), + AnonymousEdge(NamedVertex("_uv5"), NamedVertex("_uv6")), + AnonymousEdge(NamedVertex("_uv6"), NamedVertex("_uv7")), + AnonymousEdge(NamedVertex("_uv7"), NamedVertex("_uv8")), + AnonymousEdge(NamedVertex("_uv8"), NamedVertex("_uv9")), + AnonymousEdge(NamedVertex("_uv9"), NamedVertex("v")))) + } + + test("good parses - undirected pattern") { + assert( + Pattern.parse("(u)-[e]-(v)") === + Seq(UndirectedEdge(NamedEdge("e", NamedVertex("u"), NamedVertex("v"))))) + + assert( + Pattern.parse("(u)-[e]-(v);(v)-[]-(k)") === + Seq( + UndirectedEdge(NamedEdge("e", NamedVertex("u"), NamedVertex("v"))), + UndirectedEdge(AnonymousEdge(NamedVertex("v"), NamedVertex("k"))))) + } + + test("rewrite incoming edges") { + assert(Pattern.rewriteIncomingEdges("(u)<-[e]-(v);") === "(v)-[e]->(u)") + assert(Pattern.rewriteIncomingEdges("!(u)<-[e]-(v);") === "!(v)-[e]->(u)") + assert( + Pattern.rewriteIncomingEdges("(u)<-[]-(v);(u)-[e]->(v)") === "(v)-[]->(u);(u)-[e]->(v)") + assert(Pattern.rewriteIncomingEdges("(u)<-[]->(v)") === "(u)-[]->(v);(v)-[]->(u)") + assert(Pattern.rewriteIncomingEdges("(u)<-[e]->(v)") === "(u)-[e1]->(v);(v)-[e2]->(u)") + assert(Pattern.rewriteIncomingEdges("(u)<-[*5]-(v)") === "(v)-[*5]->(u)") + assert(Pattern.rewriteIncomingEdges("(u)<-[*5]->(v)") === "(u)-[*5]->(v);(v)-[*5]->(u)") + assert( + Pattern.rewriteIncomingEdges( + "(v1)<-[e*1..2]->(v2)") === "(v1)-[e*1..2]->(v2);(v2)-[e*1..2]->(v1)") + } + + test("rewrite incoming edges and parse") { + Pattern.parse("(v)<-[e]-(u)") === Pattern.parse("(u)-[e]->(v)") + Pattern.parse("(v)<-[]-(u)") === Pattern.parse("(u)-[]->(v)") + Pattern.parse("!(v)<-[]-(u)") === Pattern.parse("!(u)-[]->(v)") + Pattern.parse("()<-[e]-()") === Pattern.parse("()-[e]->()") + Pattern.parse("(u)-[]->(v); (w)<-[]-(v); !(w)<-[]-(u)") === Pattern.parse( + "(u)-[]->(v); (v)-[]->(w); !(u)-[]->(w)") + Pattern.parse("(v)<-[*5]-(u)") === Pattern.parse("(u)-[*5]->(v)") + } + + test("bad parses") { + withClue("Failed to catch parse error with lone anonymous vertex") { + intercept[InvalidParseException] { + Pattern.parse("()") + } + } + withClue("Failed to catch parse error with lone anonymous vertex") { + intercept[InvalidParseException] { + Pattern.parse("(a)-[]->(b); ()") + } + } + withClue("Failed to catch parse error") { + intercept[InvalidParseException] { + Pattern.parse("(") + } + } + withClue("Failed to catch parse error") { + intercept[InvalidParseException] { + Pattern.parse("->(a)") + } + } + withClue("Failed to catch parse error with negated vertex") { + intercept[InvalidParseException] { + Pattern.parse("!(a)") + } + } + withClue("Failed to catch parse error with negated named edge") { + val msg = intercept[InvalidParseException] { + Pattern.parse("!(a)-[ab]->(b)") + } + assert(msg.getMessage.contains("does not support negated named edges")) + } + withClue("Failed to catch parse error with negated named edge") { + val msg = intercept[InvalidParseException] { + Pattern.parse("!()-[ab]->()") + } + assert(msg.getMessage.contains("does not support negated named edges")) + } + withClue("Failed to catch parse error with double negative") { + intercept[InvalidParseException] { + Pattern.parse("!!(a)-[]->(b)") + } + } + withClue("Failed to catch parse error with completely anonymous edge ()-[]->()") { + intercept[InvalidParseException] { + Pattern.parse("()-[]->()") + } + } + withClue("Failed to catch parse error with completely anonymous negated edge !()-[]->()") { + intercept[InvalidParseException] { + Pattern.parse("!()-[]->()") + } + } + withClue("Failed to catch parse error with completely anonymous undirected edge ()-[]-()") { + intercept[InvalidParseException] { + Pattern.parse("()-[]-()") + } + } + withClue( + "Failed to catch parse error with completely anonymous negated and undirected edge !()-[]-()") { + intercept[InvalidParseException] { + Pattern.parse("!()-[]-()") + } + } + withClue("Failed to catch parse error with reused element name") { + intercept[InvalidParseException] { + Pattern.parse("(a)-[]->(b); ()-[a]->()") + } + } + withClue("Failed to catch parse error with reused element name") { + intercept[InvalidParseException] { + Pattern.parse("(a)-[a]->(b)") + } + } + withClue("Failed to catch parse error with reused edge name") { + intercept[InvalidParseException] { + Pattern.parse("(a)-[e]->(b); ()-[e]->()") + } + } + withClue("Failed to catch parse error with not support negated bidirectional edge") { + intercept[InvalidParseException] { + Pattern.parse("!(u)<-[]->(v)") + } + } + } + + test("unsupported parse on the fixed length patterns") { + withClue("Failed to catch parse error with graph frame unreachable") { + intercept[InvalidParseException] { + Pattern.parse("(u)-[*0]->(v)") + } + } + + withClue("Failed to catch parse error with bad motif string") { + intercept[InvalidParseException] { + Pattern.parse("(u)-[*]->(v)") + } + } + } + + test("empty pattern should be parsable") { + Pattern.parse("") + } + + def testFindNamedVerticesOnlyInNegatedTerms(pattern: String, expected: Seq[String]): Unit = { + test(s"findNamedVerticesOnlyInNegatedTerms: $pattern") { + val patterns = Pattern.parse(pattern) + val result = Pattern.findNamedVerticesOnlyInNegatedTerms(patterns) + assert(result === expected) + } + } + + testFindNamedVerticesOnlyInNegatedTerms( + "(u)-[]->(v); (v)-[]->(w); !(u)-[]->(w)", + Seq.empty[String]) + + testFindNamedVerticesOnlyInNegatedTerms("(u)-[]->(v); (v)-[]->(w)", Seq.empty[String]) + + testFindNamedVerticesOnlyInNegatedTerms("!(u)-[]->(v)", Seq("u", "v")) + + testFindNamedVerticesOnlyInNegatedTerms( + "(u)-[]->(v); (v)-[]->(w); !(a)-[]->(b); !(v)-[]->(c)", + Seq("a", "b", "c")) + + def testFindNamedElementsInOrder(pattern: String, expected: Seq[String]): Unit = { + test(s"testFindNamedElementsInOrder: $pattern") { + val patterns = Pattern.parse(pattern) + val result = Pattern.findNamedElementsInOrder(patterns, includeEdges = true) + assert(result === expected) + } + } + + testFindNamedElementsInOrder("(u)-[]->(v); (v)-[]->(w); !(u)-[]->(w)", Seq("u", "v", "w")) + + testFindNamedElementsInOrder("(u)-[]->(v); ()-[vw]->()", Seq("u", "v", "vw")) + + testFindNamedElementsInOrder( + "(u)-[uv]->(v); (v)-[vw]->(w); !(u)-[]->(w); (x)", + Seq("u", "uv", "v", "vw", "w", "x")) +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestartSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestartSuite.scala new file mode 100644 index 0000000000000..f0bbe33860ecf --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestartSuite.scala @@ -0,0 +1,178 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.rw + +import org.apache.spark.sql.functions.array_size +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.types.ArrayType +import org.apache.spark.sql.types.StringType +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.examples.Graphs + +class RandomWalkWithRestartSuite extends SparkFunSuite with GraphFrameTestSparkContext { + test("test RW base") { + val g = Graphs.friends + + val numBatches = 5 + val rwRunner = new RandomWalkWithRestart() + .onGraph(g) + .setRestartProbability(0.2) + .setGlobalSeed(42L) + .setBatchSize(5) + .setNumBatches(numBatches) + .setNumWalksPerNode(10) + .setTemporaryPrefix("/tmp") + + val runId = rwRunner.getRunId() + try { + val walks = rwRunner.run() + + assert(walks.schema.fields.length === 2) + // friends has string as ID type + assert(walks.schema(RandomWalkBase.rwColName).dataType === ArrayType(StringType)) + + // num rows should be: + // - 10 walks for each vertex that has edge + // - vertex "g" is isolated + // - total vertices 7 + // - total walks (7 - 1) * 10 = 60 + assert(walks.count() === 60) + + // each walk should have length numBatches * batchSize = 25 + assert(walks.filter(array_size(col(RandomWalkBase.rwColName)) =!= lit(25)).count() === 0) + + // all walk IDs are unique + assert(walks.select(col(RandomWalkBase.walkIdCol)).distinct().count() === 60) + } finally { + // Clean up temporary files after the test + rwRunner.cleanUp() + + // Verify that all temporary files have been deleted + val hadoopConf = spark.sparkContext.hadoopConfiguration + val fs = org.apache.hadoop.fs.FileSystem.get(hadoopConf) + val basePath = "/tmp" + val runPath = s"$basePath/${runId}_batch_" + (1 to numBatches).foreach { i => + val path = new org.apache.hadoop.fs.Path(s"${runPath}${i}") + assert(!fs.exists(path), s"Temporary file not deleted: $path") + } + } + } + + test("test RW with restart from middle iteration") { + val g = Graphs.friends + + val numBatches = 6 + val batchSize = 5 + val totalSteps = numBatches * batchSize + + // First run: generate all 6 batches + val rwRunner1 = new RandomWalkWithRestart() + .onGraph(g) + .setRestartProbability(0.2) + .setGlobalSeed(42L) + .setBatchSize(batchSize) + .setNumBatches(numBatches) + .setNumWalksPerNode(10) + .setTemporaryPrefix("/tmp") + + val runId = rwRunner1.getRunId() + println(s"Using runId: $runId") + + // Run and persist the result + val walks1 = rwRunner1.run() + // Materialize and persist + walks1.persist() + walks1.count() // Force materialization + + // Verify walks1 properties + assert(walks1.schema.fields.length === 2) + assert(walks1.schema(RandomWalkBase.rwColName).dataType === ArrayType(StringType)) + assert(walks1.count() === 60) // (7-1)*10 = 60 walks + assert( + walks1.filter(array_size(col(RandomWalkBase.rwColName)) =!= lit(totalSteps)).count() === 0) + + // Second run: same runID, start from iteration 3 + val rwRunner2 = new RandomWalkWithRestart() + .onGraph(g) + .setRestartProbability(0.2) + .setGlobalSeed(42L) + .setBatchSize(batchSize) + .setNumBatches(numBatches) + .setNumWalksPerNode(10) + .setTemporaryPrefix("/tmp") + .setRunId(runId) + .setStartingFromBatch(3) + + val walks2 = rwRunner2.run() + walks2.persist() + walks2.count() // Force materialization + + // Verify walks2 properties + assert(walks2.schema.fields.length === 2) + assert(walks2.schema(RandomWalkBase.rwColName).dataType === ArrayType(StringType)) + assert(walks2.count() === 60) + assert( + walks2.filter(array_size(col(RandomWalkBase.rwColName)) =!= lit(totalSteps)).count() === 0) + + // Compare results: they should be identical + // Since walk IDs are UUIDs generated during the first batch, they will be different + // between runs. So we need to compare the walks without considering walk IDs. + // We'll sort both DataFrames by the walk array and compare the arrays directly. + + // Extract just the walk arrays and sort them + val walks1Sorted = walks1 + .select(col(RandomWalkBase.rwColName)) + .orderBy(col(RandomWalkBase.rwColName).asc) + .collect() + .map(_.getSeq[String](0)) + + val walks2Sorted = walks2 + .select(col(RandomWalkBase.rwColName)) + .orderBy(col(RandomWalkBase.rwColName).asc) + .collect() + .map(_.getSeq[String](0)) + + // Both should have the same number of walks + assert(walks1Sorted.length === walks2Sorted.length) + + // Compare each walk array + walks1Sorted.zip(walks2Sorted).foreach { case (walk1, walk2) => + assert(walk1 === walk2, "Walk sequences should be identical") + } + + // Clean up temporary files + rwRunner2.cleanUp() + + // Verify cleanup + val hadoopConf = spark.sparkContext.hadoopConfiguration + val fs = org.apache.hadoop.fs.FileSystem.get(hadoopConf) + val basePath = "/tmp" + val runPath = s"$basePath/${runId}_batch_" + (1 to numBatches).foreach { i => + val path = new org.apache.hadoop.fs.Path(s"${runPath}${i}") + assert(!fs.exists(path), s"Temporary file not deleted: $path") + } + + // Unpersist + walks1.unpersist() + walks2.unpersist() + } +} diff --git a/python/pyspark/graphframes/classic/__init__.py b/python/pyspark/graphframes/classic/__init__.py new file mode 100644 index 0000000000000..cce3acad34a49 --- /dev/null +++ b/python/pyspark/graphframes/classic/__init__.py @@ -0,0 +1,16 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# diff --git a/python/pyspark/graphframes/classic/graphframe.py b/python/pyspark/graphframes/classic/graphframe.py new file mode 100644 index 0000000000000..edc5fcc737b5a --- /dev/null +++ b/python/pyspark/graphframes/classic/graphframe.py @@ -0,0 +1,603 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from __future__ import annotations + +from typing import cast, final + +from py4j.java_gateway import JavaObject + +from pyspark import SparkContext +from pyspark.graphframes.classic.utils import storage_level_to_jvm +from pyspark.graphframes.internal.utils import _RandomWalksEmbeddingsParameters +from pyspark.graphframes.lib import Pregel +from pyspark.sql import SparkSession +from pyspark.sql import functions as F +from pyspark.sql.classic.column import Column, _to_seq +from pyspark.sql.classic.dataframe import DataFrame +from pyspark.storagelevel import StorageLevel + + +def _from_java_gf(jgf: JavaObject, spark: SparkSession) -> "GraphFrame": + """ + (internal) creates a python GraphFrame wrapper from a java GraphFrame. + + :param jgf: + """ + pv = DataFrame(jgf.vertices(), spark) + pe = DataFrame(jgf.edges(), spark) + return GraphFrame(pv, pe) + + +def _java_api(jsc: SparkContext) -> JavaObject: + javaClassName = "org.apache.spark.graphframes.GraphFramePythonAPI" + if jsc._jvm is None: + raise RuntimeError( + "Spark Driver's JVM is dead or did not start properly. See driver logs for details." + ) + return ( + jsc._jvm.Thread.currentThread() + .getContextClassLoader() + .loadClass(javaClassName) + .newInstance() + ) + + +@final +class GraphFrame: + def __init__(self, v: DataFrame, e: DataFrame) -> None: + self._vertices = v + self._edges = e + self._spark = v.sparkSession + self._sc = self._spark._sc + self._jvm_gf_api = _java_api(self._sc) + self._jvm = self._spark._jvm + + self._ATTR: str = self._jvm_gf_api.ATTR() + + self._jvm_graph = self._jvm_gf_api.createGraph(v._jdf, e._jdf) + + @property + def triplets(self) -> DataFrame: + jdf = self._jvm_graph.triplets() + return DataFrame(jdf, self._spark) + + @property + def pregel(self) -> Pregel: + return Pregel(self) + + def find(self, pattern: str) -> DataFrame: + jdf = self._jvm_graph.find(pattern) + return DataFrame(jdf, self._spark) + + def filterVertices(self, condition: str | Column) -> "GraphFrame": + if isinstance(condition, str): + jdf = self._jvm_graph.filterVertices(condition) + else: + jdf = self._jvm_graph.filterVertices(condition._jc) + + return _from_java_gf(jdf, self._spark) + + def filterEdges(self, condition: str | Column) -> "GraphFrame": + if isinstance(condition, str): + jdf = self._jvm_graph.filterEdges(condition) + else: + jdf = self._jvm_graph.filterEdges(condition._jc) + + return _from_java_gf(jdf, self._spark) + + def detectingCycles( + self, + checkpoint_interval: int = 2, + use_local_checkpoints: bool = False, + intermediate_storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> DataFrame: + builder = self._jvm_graph.detectingCycles() + builder.setUseLocalCheckpoints(use_local_checkpoints) + builder.setCheckpointInterval(checkpoint_interval) + builder.setIntermediateStorageLevel( + storage_level_to_jvm(intermediate_storage_level, self._spark) + ) + jdf = builder.run() + + return DataFrame(jdf, self._spark) + + def dropIsolatedVertices(self) -> "GraphFrame": + jdf = self._jvm_graph.dropIsolatedVertices() + return _from_java_gf(jdf, self._spark) + + def bfs( + self, + fromExpr: str, + toExpr: str, + edgeFilter: str | None = None, + maxPathLength: int = 10, + ) -> DataFrame: + builder = ( + self._jvm_graph.bfs().fromExpr(fromExpr).toExpr(toExpr).maxPathLength(maxPathLength) + ) + if edgeFilter is not None: + builder.edgeFilter(edgeFilter) + jdf = builder.run() + return DataFrame(jdf, self._spark) + + def all_paths( + self, + from_expr: Column | str, + to_expr: Column | str, + edge_filter: Column | str | None = None, + max_path_length: int = 5, + is_directed: bool = True, + checkpoint_interval: int = 2, + use_local_checkpoints: bool = False, + storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> DataFrame: + builder = self._jvm_graph.allPaths() + if isinstance(from_expr, Column): + builder.fromExpr(from_expr._jc) + else: + builder.fromExpr(from_expr) + if isinstance(to_expr, Column): + builder.toExpr(to_expr._jc) + else: + builder.toExpr(to_expr) + builder.maxPathLength(max_path_length).setIsDirected(is_directed) + if edge_filter is not None: + if isinstance(edge_filter, Column): + builder.edgeFilter(edge_filter._jc) + else: + builder.edgeFilter(edge_filter) + + if checkpoint_interval > 0: + builder.setCheckpointInterval(checkpoint_interval) + + builder.setUseLocalCheckpoints(use_local_checkpoints) + builder.setIntermediateStorageLevel(storage_level_to_jvm(storage_level, self._spark)) + + jdf = builder.run() + return DataFrame(jdf, self._spark) + + def aggregateMessages( + self, + aggCol: list[Column | str], + sendToSrc: list[Column | str], + sendToDst: list[Column | str], + intermediate_storage_level: StorageLevel, + ) -> DataFrame: + builder = self._jvm_graph.aggregateMessages() + builder.setIntermediateStorageLevel( + storage_level_to_jvm(intermediate_storage_level, self._spark) + ) + if len(sendToSrc) == 1: + if isinstance(sendToSrc[0], Column): + builder.sendToSrc(sendToSrc[0]._jc) + elif isinstance(sendToSrc[0], str): + builder.sendToSrc(sendToSrc[0]) + else: + raise TypeError("Provide message either as `Column` or `str`") + elif len(sendToSrc) > 1: + if all(isinstance(x, Column) for x in sendToSrc): + send2src = [x._jc for x in cast(list[Column], sendToSrc)] + builder.sendToSrc(send2src[0], _to_seq(self._sc, send2src[1:])) + elif all(isinstance(x, str) for x in sendToSrc): + builder.sendToSrc(sendToSrc[0], _to_seq(self._sc, sendToSrc[1:])) + else: + raise TypeError( + "Multiple messages should all be `Column` or `str`, not a mix of them." + ) + + if len(sendToDst) == 1: + if isinstance(sendToDst[0], Column): + builder.sendToDst(sendToDst[0]._jc) + elif isinstance(sendToDst[0], str): + builder.sendToDst(sendToDst[0]) + else: + raise TypeError("Provide message either as `Column` or `str`") + elif len(sendToDst) > 1: + if all(isinstance(x, Column) for x in sendToDst): + send2dst = [x._jc for x in cast(list[Column], sendToDst)] + builder.sendToDst(send2dst[0], _to_seq(self._sc, send2dst[1:])) + elif all(isinstance(x, str) for x in sendToDst): + builder.sendToDst(sendToDst[0], _to_seq(self._sc, sendToDst[1:])) + else: + raise TypeError( + "Multiple messages should all be `Column` or `str`, not a mix of them." + ) + + if len(aggCol) == 1: + if isinstance(aggCol[0], Column): + jdf = builder.agg(aggCol[0]._jc) + elif isinstance(aggCol[0], str): + jdf = builder.agg(aggCol[0]) + elif len(aggCol) > 1: + if all(isinstance(x, Column) for x in aggCol): + jdf = builder.agg( + cast(Column, aggCol[0])._jc, + _to_seq(self._sc, [x._jc for x in cast(list[Column], aggCol)]), + ) + elif all(isinstance(x, str) for x in aggCol): + jdf = builder.agg(aggCol[0], _to_seq(self._sc, aggCol[1:])) + else: + raise TypeError( + "Multiple agg cols should all be `Column` or `str`, not a mix of them." + ) + return DataFrame(jdf, self._spark) + + def connectedComponents( + self, + algorithm: str, + checkpointInterval: int, + broadcastThreshold: int, + useLabelsAsComponents: bool, + use_local_checkpoints: bool, + max_iter: int, + storage_level: StorageLevel, + ) -> DataFrame: + java_cc = self._jvm_graph.connectedComponents() + java_cc.setAlgorithm(algorithm) + java_cc.setCheckpointInterval(checkpointInterval) + java_cc.setBroadcastThreshold(broadcastThreshold) + java_cc.setUseLabelsAsComponents(useLabelsAsComponents) + java_cc.setUseLocalCheckpoints(use_local_checkpoints) + java_cc.maxIter(max_iter) + java_cc.setIntermediateStorageLevel(storage_level_to_jvm(storage_level, self._spark)) + jdf = java_cc.run() + + return DataFrame(jdf, self._spark) + + def labelPropagation( + self, + maxIter: int, + algorithm: str, + use_local_checkpoints: bool, + checkpoint_interval: int, + storage_level: StorageLevel, + ) -> DataFrame: + java_cdlp = self._jvm_graph.labelPropagation() + java_cdlp.maxIter(maxIter) + java_cdlp.setAlgorithm(algorithm) + java_cdlp.setUseLocalCheckpoints(use_local_checkpoints) + java_cdlp.setCheckpointInterval(checkpoint_interval) + java_cdlp.setIntermediateStorageLevel(storage_level_to_jvm(storage_level, self._spark)) + jdf = java_cdlp.run() + + return DataFrame(jdf, self._spark) + + def neighborhood_aware_cdlp( + self, + max_iter: int, + structural_similarity_multiplier: float, + ignore_direct_links: bool, + use_local_checkpoints: bool, + checkpoint_interval: int, + storage_level: StorageLevel, + is_directed: bool, + lg_nom_entries: int, + initial_label_col: str | None, + ) -> DataFrame: + java_nacdlp = self._jvm_graph.structureAwareLabelPropagation() + java_nacdlp.maxIter(max_iter) + java_nacdlp.setStructuralSimilarityMultiplier(structural_similarity_multiplier) + java_nacdlp.setIgnoreDirectLinks(ignore_direct_links) + java_nacdlp.setUseLocalCheckpoints(use_local_checkpoints) + java_nacdlp.setCheckpointInterval(checkpoint_interval) + java_nacdlp.setIntermediateStorageLevel(storage_level_to_jvm(storage_level, self._spark)) + java_nacdlp.setIsDirected(is_directed) + java_nacdlp.setLgNomEntries(lg_nom_entries) + + if initial_label_col is not None: + java_nacdlp.setInitialLabelCol(initial_label_col) + + jdf = java_nacdlp.run() + return DataFrame(jdf, self._spark) + + def pageRank( + self, + resetProbability: float = 0.15, + sourceId: str | int | None = None, + maxIter: int | None = None, + tol: float | None = None, + ) -> "GraphFrame": + builder = self._jvm_graph.pageRank().resetProbability(resetProbability) + if sourceId is not None: + builder.sourceId(sourceId) + if maxIter is not None: + builder.maxIter(maxIter) + assert tol is None, "Exactly one of maxIter or tol should be set." + else: + assert tol is not None, "Exactly one of maxIter or tol should be set." + builder.tol(tol) + jgf = builder.run() + return _from_java_gf(jgf, self._spark) + + def parallelPersonalizedPageRank( + self, + resetProbability: float = 0.15, + sourceIds: list[str | int] | None = None, + maxIter: int | None = None, + ) -> "GraphFrame": + assert sourceIds is not None and len(sourceIds) > 0, ( + "Source vertices Ids sourceIds must be provided" + ) + assert maxIter is not None, "Max number of iterations maxIter must be provided" + jvm = self._sc._jvm + assert jvm is not None + sourceIds = jvm.PythonUtils.toArray(sourceIds) + builder = self._jvm_graph.parallelPersonalizedPageRank() + builder.resetProbability(resetProbability) + builder.sourceIds(sourceIds) + builder.maxIter(maxIter) + jgf = builder.run() + return _from_java_gf(jgf, self._spark) + + def shortestPaths( + self, + landmarks: list[str | int], + algorithm: str, + use_local_checkpoints: bool, + checkpoint_interval: int, + storage_level: StorageLevel, + is_directed: bool, + ) -> DataFrame: + java_sp = self._jvm_graph.shortestPaths() + java_sp.landmarks(landmarks) + java_sp.setAlgorithm(algorithm) + java_sp.setUseLocalCheckpoints(use_local_checkpoints) + java_sp.setCheckpointInterval(checkpoint_interval) + java_sp.setIntermediateStorageLevel(storage_level_to_jvm(storage_level, self._spark)) + java_sp.setIsDirected(is_directed) + jdf = java_sp.run() + + return DataFrame(jdf, self._spark) + + def stronglyConnectedComponents(self, maxIter: int) -> DataFrame: + builder = self._jvm_graph.stronglyConnectedComponents() + builder.maxIter(maxIter) + jdf = builder.run() + return DataFrame(jdf, self._spark) + + def svdPlusPlus( + self, + rank: int = 10, + maxIter: int = 2, + minValue: float = 0.0, + maxValue: float = 5.0, + gamma1: float = 0.007, + gamma2: float = 0.007, + gamma6: float = 0.005, + gamma7: float = 0.015, + ) -> tuple[DataFrame, float]: + # This call is actually useless, because one needs to build the configuration first... + builder = self._jvm_graph.svdPlusPlus() + builder.rank(rank).maxIter(maxIter).minValue(minValue).maxValue(maxValue) + builder.gamma1(gamma1).gamma2(gamma2).gamma6(gamma6).gamma7(gamma7) + jdf = builder.run() + loss = builder.loss() + v = DataFrame(jdf, self._spark) + return (v, loss) + + def triangleCount( + self, storage_level: StorageLevel, algorithm: str, log_nom_entries: int + ) -> DataFrame: + builder = self._jvm_graph.triangleCount() + builder.setIntermediateStorageLevel(storage_level_to_jvm(storage_level, self._spark)) + builder.setAlgorithm(algorithm) + builder.setLgNomEntries(log_nom_entries) + jdf = builder.run() + return DataFrame(jdf, self._spark) + + def powerIterationClustering( + self, k: int, maxIter: int, weightCol: str | None = None + ) -> DataFrame: + jvm = self._spark._jvm + assert jvm is not None + if weightCol: + weightCol = jvm.scala.Option.apply(weightCol) + else: + weightCol = jvm.scala.Option.empty() + jdf = self._jvm_graph.powerIterationClustering(k, maxIter, weightCol) + return DataFrame(jdf, self._spark) + + def maximal_independent_set( + self, + checkpoint_interval: int, + storage_level: StorageLevel, + use_local_checkpoints: bool, + seed: int, + ) -> DataFrame: + builder = self._jvm_graph.maximalIndependentSet() + builder.setCheckpointInterval(checkpoint_interval) + builder.setIntermediateStorageLevel(storage_level_to_jvm(storage_level, self._spark)) + builder.setUseLocalCheckpoints(use_local_checkpoints) + + jdf = builder.run(seed) + return DataFrame(jdf, self._spark) + + def k_core( + self, + checkpoint_interval: int, + use_local_checkpoints: bool, + storage_level: StorageLevel, + ) -> DataFrame: + java_kcore = self._jvm_graph.kCore() + java_kcore.setUseLocalCheckpoints(use_local_checkpoints) + java_kcore.setCheckpointInterval(checkpoint_interval) + java_kcore.setIntermediateStorageLevel(storage_level_to_jvm(storage_level, self._spark)) + jdf = java_kcore.run() + + return DataFrame(jdf, self._spark) + + def hyper_anf( + self, + n_hops: int, + lg_nom_entries: int, + edge_filter: Column | str | None, + checkpoint_interval: int, + use_local_checkpoints: bool, + storage_level: StorageLevel, + ) -> DataFrame: + builder = self._jvm_graph.hyperANF() + builder.setNHops(n_hops) + builder.setLgNomEntries(lg_nom_entries) + if edge_filter is not None: + if isinstance(edge_filter, Column): + builder.setEdgesFilterExpression(edge_filter._jc) + else: + builder.setEdgesFilterExpression(edge_filter) + builder.setCheckpointInterval(checkpoint_interval) + builder.setUseLocalCheckpoints(use_local_checkpoints) + builder.setIntermediateStorageLevel(storage_level_to_jvm(storage_level, self._spark)) + jdf = builder.run() + + return DataFrame(jdf, self._spark) + + def aggregate_neighbors( + self, + starting_vertices: Column | str, + max_hops: int, + accumulator_names: list[str], + accumulator_inits: list[Column | str], + accumulator_updates: list[Column | str], + stopping_condition: Column | str | None = None, + target_condition: Column | str | None = None, + required_vertex_attributes: list[str] | None = None, + required_edge_attributes: list[str] | None = None, + edge_filter: Column | str | None = None, + remove_loops: bool = False, + checkpoint_interval: int = 0, + use_local_checkpoints: bool = False, + storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> DataFrame: + builder = self._jvm_graph.aggregateNeighbors() + + # Set required parameters + if isinstance(starting_vertices, Column): + builder.setStartingVertices(starting_vertices._jc) + else: + builder.setStartingVertices(starting_vertices) + + builder.setMaxHops(max_hops) + + jvm = self._sc._jvm + assert jvm is not None + # Handle accumulators with proper py4j conversion + if len(accumulator_names) > 0: + names_seq = jvm.scala.collection.JavaConverters.asScalaBuffer(accumulator_names).toSeq() + + inits_list = [] + for init in accumulator_inits: + if isinstance(init, Column): + inits_list.append(init._jc) + else: + inits_list.append(F.expr(init)._jc) + inits_seq = jvm.scala.collection.JavaConverters.asScalaBuffer(inits_list).toSeq() + + updates_list = [] + for update in accumulator_updates: + if isinstance(update, Column): + updates_list.append(update._jc) + else: + updates_list.append(F.expr(update)._jc) + updates_seq = jvm.scala.collection.JavaConverters.asScalaBuffer(updates_list).toSeq() + + builder.setAccumulators(names_seq, inits_seq, updates_seq) + + # Set optional parameters + if stopping_condition is not None: + if isinstance(stopping_condition, Column): + builder.setStoppingCondition(stopping_condition._jc) + else: + builder.setStoppingCondition(stopping_condition) + + if target_condition is not None: + if isinstance(target_condition, Column): + builder.setTargetCondition(target_condition._jc) + else: + builder.setTargetCondition(target_condition) + + if required_vertex_attributes is not None and len(required_vertex_attributes) > 0: + attrs_seq = jvm.scala.collection.JavaConverters.asScalaBuffer( + required_vertex_attributes + ).toSeq() + builder.setRequiredVertexAttributes(attrs_seq) + + if required_edge_attributes is not None and len(required_edge_attributes) > 0: + attrs_seq = jvm.scala.collection.JavaConverters.asScalaBuffer( + required_edge_attributes + ).toSeq() + builder.setRequiredEdgeAttributes(attrs_seq) + + if edge_filter is not None: + if isinstance(edge_filter, Column): + builder.setEdgeFilter(edge_filter._jc) + else: + builder.setEdgeFilter(edge_filter) + + builder.setRemoveLoops(remove_loops) + + if checkpoint_interval > 0: + builder.setCheckpointInterval(checkpoint_interval) + + builder.setUseLocalCheckpoints(use_local_checkpoints) + builder.setIntermediateStorageLevel(storage_level_to_jvm(storage_level, self._spark)) + + jdf = builder.run() + assert jdf is not None + + return DataFrame(jdf, self._spark) + + def rw_embeddings(self, params: _RandomWalksEmbeddingsParameters) -> DataFrame: + assert self._jvm is not None + j_rw_embeddings = self._jvm.org.apache.spark.graphframes.embeddings.RandomWalkEmbeddings + assert j_rw_embeddings is not None + jdf: JavaObject = j_rw_embeddings.pythonAPI( + self._jvm_graph, + params.use_edge_direction, + params.rw_model, + params.rw_max_nbrs, + params.rw_num_walks_per_node, + params.rw_batch_size, + params.rw_num_batches, + params.rw_seed, + params.rw_restart_probability, + params.rw_temporary_prefix, + params.rw_cached_walks, + params.sequence_model, + params.hash2vec_context_size, + params.hash2vec_num_partitions, + params.hash2vec_embeddings_dim, + params.hash2vec_decay_function, + params.hash2vec_gaussian_sigma, + params.hash2vec_hashing_seed, + params.hash2vec_sign_seed, + params.hash2vec_do_l2_norm, + params.hash2vec_safe_l2, + params.word2vec_max_iter, + params.word2vec_embeddings_dim, + params.word2vec_window_size, + params.word2vec_num_partitions, + params.word2vec_min_count, + params.word2vec_max_sentence_length, + params.word2vec_seed, + params.word2vec_step_size, + params.aggregate_neighbors, + params.aggregate_neighbors_max_nbrs, + params.aggregate_neighbors_seed, + params.clean_up_after_run, + ) + assert jdf is not None + + return DataFrame(jdf, self._spark) diff --git a/python/pyspark/graphframes/classic/utils.py b/python/pyspark/graphframes/classic/utils.py new file mode 100644 index 0000000000000..d1b7a8e1b9481 --- /dev/null +++ b/python/pyspark/graphframes/classic/utils.py @@ -0,0 +1,31 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from py4j.java_gateway import JavaObject + +from pyspark.sql import SparkSession +from pyspark.storagelevel import StorageLevel + + +def storage_level_to_jvm(storage_level: StorageLevel, spark: SparkSession) -> JavaObject: + assert spark._jvm is not None + return spark._jvm.org.apache.spark.storage.StorageLevel.apply( + storage_level.useDisk, + storage_level.useMemory, + storage_level.useOffHeap, + storage_level.deserialized, + storage_level.replication, + ) diff --git a/python/pyspark/graphframes/connect/__init__.py b/python/pyspark/graphframes/connect/__init__.py new file mode 100644 index 0000000000000..cce3acad34a49 --- /dev/null +++ b/python/pyspark/graphframes/connect/__init__.py @@ -0,0 +1,16 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# diff --git a/python/pyspark/graphframes/connect/graphframes_client.py b/python/pyspark/graphframes/connect/graphframes_client.py new file mode 100644 index 0000000000000..4bf0121f79c57 --- /dev/null +++ b/python/pyspark/graphframes/connect/graphframes_client.py @@ -0,0 +1,1642 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from __future__ import annotations + +from typing import cast, final + +from typing_extensions import Self, override + +from pyspark.graphframes.internal.utils import _RandomWalksEmbeddingsParameters +from pyspark.sql.connect import functions as F +from pyspark.sql.connect import proto +from pyspark.sql.connect.client import SparkConnectClient +from pyspark.sql.connect.column import Column +from pyspark.sql.connect.dataframe import DataFrame +from pyspark.sql.connect.plan import LogicalPlan +from pyspark.sql.connect.proto import graphframes_pb2 as pb +from pyspark.sql.connect.session import SparkSession +from pyspark.storagelevel import StorageLevel + +from .utils import ( + dataframe_to_proto, + make_column_or_expr, + make_str_or_long_id, + storage_level_to_proto, +) + + +def _dataframe_from_plan(plan: LogicalPlan, session: SparkSession) -> DataFrame: + return DataFrame(plan, session) + + +@final +class PregelConnect: + def __init__(self, graph: "GraphFrameConnect") -> None: + self.graph = graph + self._max_iter = 10 + self._checkpoint_interval = 2 + self._col_name: str | None = None + self._initial_expr: Column | str | None = None + self._update_after_agg_msgs_expr: Column | str | None = None + self._send_msg_to_src: list[Column | str] = [] + self._send_msg_to_dst: list[Column | str] = [] + self._agg_msg: Column | None = None + self._early_stopping = False + self._use_local_checkpoints = False + self._storage_level = StorageLevel.MEMORY_AND_DISK_DESER + self._initial_active_expr: Column | str | None = None + self._update_active_expr: Column | str | None = None + self._stop_if_all_non_active = False + self._skip_messages_from_non_active = False + self._required_src_columns: list[str] = [] + self._required_dst_columns: list[str] = [] + self._required_edge_columns: list[str] = [] + + def setMaxIter(self, value: int) -> Self: + self._max_iter = value + return self + + def setCheckpointInterval(self, value: int) -> Self: + self._checkpoint_interval = value + return self + + def setEarlyStopping(self, value: bool) -> Self: + self._early_stopping = value + return self + + def withVertexColumn( + self, + colName: str, + initialExpr: Column | str, + updateAfterAggMsgsExpr: Column | str, + ) -> Self: + self._col_name = colName + self._initial_expr = initialExpr + self._update_after_agg_msgs_expr = updateAfterAggMsgsExpr + return self + + def sendMsgToSrc(self, msgExpr: Column | str) -> Self: + self._send_msg_to_src.append(msgExpr) + return self + + def sendMsgToDst(self, msgExpr: Column | str) -> Self: + self._send_msg_to_dst.append(msgExpr) + return self + + def aggMsgs(self, aggExpr: Column) -> Self: + self._agg_msg = aggExpr + return self + + def setStopIfAllNonActiveVertices(self, value: bool) -> Self: + self._stop_if_all_non_active = value + return self + + def setInitialActiveVertexExpression(self, value: Column | str) -> Self: + self._initial_active_expr = value + return self + + def setUpdateActiveVertexExpression(self, value: Column | str) -> Self: + self._update_active_expr = value + return self + + def setSkipMessagesFromNonActiveVertices(self, value: bool) -> Self: + self._skip_messages_from_non_active = value + return self + + def setUseLocalCheckpoints(self, value: bool) -> Self: + self._use_local_checkpoints = value + return self + + def setIntermediateStorageLevel(self, storage_level: StorageLevel) -> Self: + self._storage_level = storage_level + return self + + def required_src_columns(self, col_name: str, *col_names: str) -> Self: + """Specifies which source vertex columns are required when constructing triplets. + + By default, all source vertex columns are included in triplets, which can create large + intermediate datasets for algorithms with significant state. Use this method to reduce + memory usage by specifying only the columns that are actually needed. + + :param col_name: the first required source vertex column name + :param col_names: additional required source vertex column names + """ + self._required_src_columns = [col_name] + list(col_names) + return self + + def required_dst_columns(self, col_name: str, *col_names: str) -> Self: + """Specifies which destination vertex columns are required when constructing triplets. + + By default, all destination vertex columns are included in triplets, which can create large + intermediate datasets for algorithms with significant state. Use this method to reduce + memory usage by specifying only the columns that are actually needed. + + :param col_name: the first required destination vertex column name + :param col_names: additional required destination vertex column names + """ + self._required_dst_columns = [col_name] + list(col_names) + return self + + def required_edge_columns(self, col_name: str, *col_names: str) -> Self: + """Specifies which edge columns are required when constructing triplets. + + By default, only src and dst columns are included. Use this method to specify + additional edge columns that are needed by the sendMsgToSrc and sendMsgToDst + expressions. + + :param col_name: the first required edge column name + :param col_names: additional required edge column names + """ + self._required_edge_columns = [col_name] + list(col_names) + return self + + def run(self) -> DataFrame: + @final + class Pregel(LogicalPlan): + def __init__( + self, + max_iter: int, + checkpoint_interval: int, + early_stopping: bool, + vertex_col_name: str, + agg_msg: Column | str, + send2dst: list[Column | str], + send2src: list[Column | str], + vertex_col_init: Column | str, + vertex_col_upd: Column | str, + use_local_checkpoints: bool, + storage_level: StorageLevel, + initial_active_col: Column | str | None, + update_active_col: Column | str | None, + stop_if_all_non_active: bool, + skip_message_from_non_active: bool, + required_src_columns: list[str], + required_dst_columns: list[str], + required_edge_columns: list[str], + vertices: DataFrame, + edges: DataFrame, + ) -> None: + super().__init__(None) + self.max_iter = max_iter + self.checkpoint_interval = checkpoint_interval + self.early_stopping = early_stopping + self.vertex_col_name = vertex_col_name + self.agg_msg = agg_msg + self.send2dst = send2dst + self.send2src = send2src + self.vertex_col_init = vertex_col_init + self.vertex_col_upd = vertex_col_upd + self.use_local_checkpoints = use_local_checkpoints + self.storage_level = storage_level + self.initial_active_expr = initial_active_col + self.update_active_expr = update_active_col + self.stop_if_all_non_active = stop_if_all_non_active + self.skip_message_from_non_active = skip_message_from_non_active + self.required_src_columns = required_src_columns + self.required_dst_columns = required_dst_columns + self.required_edge_columns = required_edge_columns + self.vertices = vertices + self.edges = edges + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + pregel = pb.Pregel( + agg_msgs=make_column_or_expr(self.agg_msg, session), + send_msg_to_dst=[ + make_column_or_expr(c_or_e, session) for c_or_e in self.send2dst + ], + send_msg_to_src=[ + make_column_or_expr(c_or_e, session) for c_or_e in self.send2src + ], + checkpoint_interval=self.checkpoint_interval, + max_iter=self.max_iter, + additional_col_name=self.vertex_col_name, + additional_col_initial=make_column_or_expr(self.vertex_col_init, session), + additional_col_upd=make_column_or_expr(self.vertex_col_upd, session), + early_stopping=self.early_stopping, + use_local_checkpoints=self.use_local_checkpoints, + storage_level=storage_level_to_proto(self.storage_level), + stop_if_all_non_active=self.stop_if_all_non_active, + skip_messages_from_non_active=self.skip_message_from_non_active, + initial_active_expr=make_column_or_expr(self.initial_active_expr, session) + if self.initial_active_expr is not None + else None, + update_active_expr=make_column_or_expr(self.update_active_expr, session) + if self.update_active_expr is not None + else None, + required_src_columns=",".join(self.required_src_columns) + if self.required_src_columns + else None, + required_dst_columns=",".join(self.required_dst_columns) + if self.required_dst_columns + else None, + required_edge_columns=",".join(self.required_edge_columns) + if self.required_edge_columns + else None, + ) + pb_message = pb.GraphFramesAPI( + vertices=dataframe_to_proto(self.vertices, session), + edges=dataframe_to_proto(self.edges, session), + ) + pb_message.pregel.CopyFrom(pregel) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(pb_message) + return plan + + if ( + (self._col_name is None) + or (self._initial_expr is None) + or (self._update_after_agg_msgs_expr is None) + ): + raise ValueError("Initial vertex column is not initialized!") + + if self._agg_msg is None: + raise ValueError("AggMsg is not initialized!") + + return _dataframe_from_plan( + Pregel( + max_iter=self._max_iter, + checkpoint_interval=self._checkpoint_interval, + vertex_col_name=self._col_name, + vertex_col_init=self._initial_expr, + vertex_col_upd=self._update_after_agg_msgs_expr, + agg_msg=self._agg_msg, + send2dst=self._send_msg_to_dst, + send2src=self._send_msg_to_src, + early_stopping=self._early_stopping, + use_local_checkpoints=self._use_local_checkpoints, + initial_active_col=self._initial_active_expr, + update_active_col=self._update_active_expr, + stop_if_all_non_active=self._stop_if_all_non_active, + skip_message_from_non_active=self._skip_messages_from_non_active, + required_src_columns=self._required_src_columns, + required_dst_columns=self._required_dst_columns, + required_edge_columns=self._required_edge_columns, + storage_level=self._storage_level, + vertices=self.graph._vertices, + edges=self.graph._edges, + ), + session=self.graph._spark, + ) + + @staticmethod + def msg() -> Column: + return cast(Column, F.col("_pregel_msg_")) + + @staticmethod + def src(colName: str) -> Column: + return cast(Column, F.col("src." + colName)) + + @staticmethod + def dst(colName: str) -> Column: + return cast(Column, F.col("dst." + colName)) + + @staticmethod + def edge(colName: str) -> Column: + return cast(Column, F.col("edge." + colName)) + + +@final +class GraphFrameConnect: + _ID: str = "id" + _SRC: str = "src" + _DST: str = "dst" + _EDGE: str = "edge" + + def __init__(self, v: DataFrame, e: DataFrame) -> None: + self._vertices = v + self._edges = e + self._spark = v.sparkSession + + @staticmethod + def _get_pb_api_message( + vertices: DataFrame, edges: DataFrame, client: SparkConnectClient + ) -> pb.GraphFramesAPI: + return pb.GraphFramesAPI( + vertices=dataframe_to_proto(vertices, client), + edges=dataframe_to_proto(edges, client), + ) + + @property + def triplets(self) -> DataFrame: + @final + class Triplets(LogicalPlan): + def __init__(self, v: DataFrame, e: DataFrame) -> None: + super().__init__(None) + self.v = v + self.e = e + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.triplets.CopyFrom(pb.Triplets()) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + return _dataframe_from_plan(Triplets(self._vertices, self._edges), self._spark) + + @property + def pregel(self) -> PregelConnect: + return PregelConnect(self) + + def find(self, pattern: str) -> DataFrame: + @final + class Find(LogicalPlan): + def __init__(self, v: DataFrame, e: DataFrame, pattern: str) -> None: + super().__init__(None) + self.v = v + self.e = e + self.p = pattern + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.find.CopyFrom(pb.Find(pattern=self.p)) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + return _dataframe_from_plan(Find(self._vertices, self._edges, pattern), self._spark) + + def filterVertices(self, condition: str | Column) -> "GraphFrameConnect": + @final + class FilterVertices(LogicalPlan): + def __init__(self, v: DataFrame, e: DataFrame, condition: str | Column) -> None: + super().__init__(None) + self.v = v + self.e = e + self.c = condition + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + col_or_expr = make_column_or_expr(self.c, session) + graphframes_api_call.filter_vertices.CopyFrom( + pb.FilterVertices(condition=col_or_expr) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + new_vertices = _dataframe_from_plan( + FilterVertices(self._vertices, self._edges, condition), self._spark + ) + # Exactly like in the scala-core + new_edges = self._edges.join( + new_vertices.withColumn(self._SRC, F.col(self._ID)), + on=[self._SRC], + how="left_semi", + ).join( + new_vertices.withColumn(self._DST, F.col(self._ID)), + on=[self._DST], + how="left_semi", + ) + return GraphFrameConnect(new_vertices, cast(DataFrame, new_edges)) + + def filterEdges(self, condition: str | Column) -> "GraphFrameConnect": + @final + class FilterEdges(LogicalPlan): + def __init__(self, v: DataFrame, e: DataFrame, condition: str | Column) -> None: + super().__init__(None) + self.v = v + self.e = e + self.c = condition + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + col_or_expr = make_column_or_expr(self.c, session) + graphframes_api_call.filter_edges.CopyFrom(pb.FilterEdges(condition=col_or_expr)) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + new_edges = _dataframe_from_plan( + FilterEdges(self._vertices, self._edges, condition), self._spark + ) + return GraphFrameConnect(self._vertices, new_edges) + + def detectingCycles( + self, + checkpoint_interval: int, + use_local_checkpoints: bool, + intermediate_storage_level: StorageLevel, + ) -> DataFrame: + @final + class DetectingCycles(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + checkpoint_interval: int, + use_local_checkpoints: bool, + storage_level: StorageLevel, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.checkpoint_interval = checkpoint_interval + self.use_local_checkpoints = use_local_checkpoints + self.storage_level = storage_level + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.detecting_cycles.CopyFrom( + pb.DetectingCycles( + use_local_checkpoints=self.use_local_checkpoints, + checkpoint_interval=self.checkpoint_interval, + storage_level=storage_level_to_proto(self.storage_level), + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + return _dataframe_from_plan( + DetectingCycles( + self._vertices, + self._edges, + checkpoint_interval, + use_local_checkpoints, + intermediate_storage_level, + ), + self._spark, + ) + + def dropIsolatedVertices(self) -> "GraphFrameConnect": + @final + class DropIsolatedVertices(LogicalPlan): + def __init__(self, v: DataFrame, e: DataFrame) -> None: + super().__init__(None) + self.v = v + self.e = e + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.drop_isolated_vertices.CopyFrom(pb.DropIsolatedVertices()) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + new_vertices = _dataframe_from_plan( + DropIsolatedVertices(self._vertices, self._edges), self._spark + ) + return GraphFrameConnect(new_vertices, self._edges) + + def bfs( + self, + fromExpr: Column | str, + toExpr: Column | str, + edgeFilter: Column | str | None = None, + maxPathLength: int = 10, + ) -> DataFrame: + @final + class BFS(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + from_expr: Column | str, + to_expr: Column | str, + edge_filter: Column | str, + max_path_len: int, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.from_expr = from_expr + self.to_expr = to_expr + self.edge_filter = edge_filter + self.max_path_len = max_path_len + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.bfs.CopyFrom( + pb.BFS( + from_expr=make_column_or_expr(self.from_expr, session), + to_expr=make_column_or_expr(self.to_expr, session), + edge_filter=make_column_or_expr(self.edge_filter, session), + max_path_length=self.max_path_len, + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + actual_edge_filter: Column | str = ( + cast(Column, F.lit(True)) if edgeFilter is None else edgeFilter + ) + + return _dataframe_from_plan( + BFS( + v=self._vertices, + e=self._edges, + from_expr=fromExpr, + to_expr=toExpr, + edge_filter=actual_edge_filter, + max_path_len=maxPathLength, + ), + self._spark, + ) + + def all_paths( + self, + from_expr: Column | str, + to_expr: Column | str, + edge_filter: Column | str | None = None, + max_path_length: int = 5, + is_directed: bool = True, + checkpoint_interval: int = 2, + use_local_checkpoints: bool = False, + storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> DataFrame: + @final + class AllPaths(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + from_expr: Column | str, + to_expr: Column | str, + edge_filter: Column | str, + max_path_length: int, + is_directed: bool, + checkpoint_interval: int, + use_local_checkpoints: bool, + storage_level: StorageLevel, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.from_expr = from_expr + self.to_expr = to_expr + self.edge_filter = edge_filter + self.max_path_length = max_path_length + self.is_directed = is_directed + self.checkpoint_interval = checkpoint_interval + self.use_local_checkpoints = use_local_checkpoints + self.storage_level = storage_level + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.all_paths.CopyFrom( + pb.AllPaths( + from_expr=make_column_or_expr(self.from_expr, session), + to_expr=make_column_or_expr(self.to_expr, session), + edge_filter=make_column_or_expr(self.edge_filter, session), + max_path_length=self.max_path_length, + is_directed=self.is_directed, + checkpoint_interval=self.checkpoint_interval, + use_local_checkpoints=self.use_local_checkpoints, + storage_level=storage_level_to_proto(self.storage_level), + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + actual_edge_filter: Column | str = ( + cast(Column, F.lit(True)) if edge_filter is None else edge_filter + ) + + return _dataframe_from_plan( + AllPaths( + v=self._vertices, + e=self._edges, + from_expr=from_expr, + to_expr=to_expr, + edge_filter=actual_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, + ) + + def aggregateMessages( + self, + aggCol: list[Column | str], + sendToSrc: list[Column | str], + sendToDst: list[Column | str], + intermediate_storage_level: StorageLevel, + ) -> DataFrame: + @final + class AggregateMessages(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + agg_col: list[Column | str], + send2src: list[Column | str], + send2dst: list[Column | str], + storage_level: StorageLevel, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.agg_col = agg_col + self.send2src = send2src + self.send2dst = send2dst + self.storage_level = storage_level + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.aggregate_messages.CopyFrom( + pb.AggregateMessages( + agg_col=[make_column_or_expr(x, session) for x in self.agg_col], + send_to_src=[make_column_or_expr(x, session) for x in self.send2src], + send_to_dst=[make_column_or_expr(x, session) for x in self.send2dst], + storage_level=storage_level_to_proto(self.storage_level), + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + if (len(sendToSrc) == 0) and (len(sendToDst) == 0): + raise ValueError("Either `sendToSrc`, `sendToDst`, or both have to be provided") + + return _dataframe_from_plan( + AggregateMessages( + self._vertices, + self._edges, + aggCol, + sendToSrc, + sendToDst, + intermediate_storage_level, + ), + self._spark, + ) + + def connectedComponents( + self, + algorithm: str, + checkpointInterval: int, + broadcastThreshold: int, + useLabelsAsComponents: bool, + use_local_checkpoints: bool, + max_iter: int, + storage_level: StorageLevel, + ) -> DataFrame: + @final + class ConnectedComponents(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + algorithm: str, + checkpoint_interval: int, + broadcast_threshold: int, + use_labels_as_components: bool, + use_local_checkpoints: bool, + max_iter: int, + storage_level: StorageLevel, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.algorithm = algorithm + self.checkpoint_interval = checkpoint_interval + self.broadcast_threshold = broadcast_threshold + self.use_labels_as_components = use_labels_as_components + self.use_local_checkpoints = use_local_checkpoints + self.max_iter = max_iter + self.storage_level = storage_level + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.connected_components.CopyFrom( + pb.ConnectedComponents( + algorithm=self.algorithm, + checkpoint_interval=self.checkpoint_interval, + broadcast_threshold=self.broadcast_threshold, + use_labels_as_components=self.use_labels_as_components, + use_local_checkpoints=self.use_local_checkpoints, + max_iter=self.max_iter, + storage_level=storage_level_to_proto(self.storage_level), + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + return _dataframe_from_plan( + ConnectedComponents( + self._vertices, + self._edges, + algorithm, + checkpointInterval, + broadcastThreshold, + useLabelsAsComponents, + use_local_checkpoints, + max_iter, + storage_level, + ), + self._spark, + ) + + def labelPropagation( + self, + maxIter: int, + algorithm: str, + use_local_checkpoints: bool, + checkpoint_interval: int, + storage_level: StorageLevel, + ) -> DataFrame: + @final + class LabelPropagation(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + max_iter: int, + algorithm: str, + use_local_checkpoints: bool, + checkpoint_interval: int, + storage_level: StorageLevel, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.max_iter = max_iter + self.algorithm = algorithm + self.use_local_checkpoints = use_local_checkpoints + self.checkpoint_interval = checkpoint_interval + self.storage_level = storage_level + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.label_propagation.CopyFrom( + pb.LabelPropagation( + algorithm=self.algorithm, + max_iter=self.max_iter, + use_local_checkpoints=self.use_local_checkpoints, + checkpoint_interval=self.checkpoint_interval, + storage_level=storage_level_to_proto(self.storage_level), + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + return _dataframe_from_plan( + LabelPropagation( + self._vertices, + self._edges, + maxIter, + algorithm, + use_local_checkpoints, + checkpoint_interval, + storage_level, + ), + self._spark, + ) + + def neighborhood_aware_cdlp( + self, + max_iter: int, + structural_similarity_multiplier: float, + ignore_direct_links: bool, + use_local_checkpoints: bool, + checkpoint_interval: int, + storage_level: StorageLevel, + is_directed: bool, + lg_nom_entries: int, + initial_label_col: str | None, + ) -> DataFrame: + @final + class NeighborhoodAwareCDLP(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + max_iter: int, + structural_similarity_multiplier: float, + ignore_direct_links: bool, + use_local_checkpoints: bool, + checkpoint_interval: int, + storage_level: StorageLevel, + is_directed: bool, + lg_nom_entries: int, + initial_label_col: str | None, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.max_iter = max_iter + self.structural_similarity_multiplier = structural_similarity_multiplier + self.ignore_direct_links = ignore_direct_links + self.use_local_checkpoints = use_local_checkpoints + self.checkpoint_interval = checkpoint_interval + self.storage_level = storage_level + self.is_directed = is_directed + self.lg_nom_entries = lg_nom_entries + self.initial_label_col = initial_label_col + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.neighborhood_aware_cdlp.CopyFrom( + pb.NeighborhoodAwareCDLP( + max_iter=self.max_iter, + structural_similarity_multiplier=self.structural_similarity_multiplier, + ignore_direct_links=self.ignore_direct_links, + use_local_checkpoints=self.use_local_checkpoints, + checkpoint_interval=self.checkpoint_interval, + storage_level=storage_level_to_proto(self.storage_level), + is_directed=self.is_directed, + lg_nom_entries=self.lg_nom_entries, + initial_label_col=self.initial_label_col, + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + return _dataframe_from_plan( + NeighborhoodAwareCDLP( + self._vertices, + self._edges, + max_iter, + structural_similarity_multiplier, + ignore_direct_links, + use_local_checkpoints, + checkpoint_interval, + storage_level, + is_directed, + lg_nom_entries, + initial_label_col, + ), + self._spark, + ) + + def _update_page_rank_edge_weights(self, new_vertices: DataFrame) -> "GraphFrameConnect": + cols2select = [col for col in self._edges.columns if col != "weight"] + ["weight"] + out_degrees = self._edges.groupBy(self._SRC).agg(F.count("*").alias("outDegree")) + new_edges = ( + self._edges.join(out_degrees, on=[self._SRC], how="inner") + .withColumn("weight", F.lit(1.0) / F.col("outDegree")) + .select(*cols2select) + ) + return GraphFrameConnect(new_vertices, cast(DataFrame, new_edges)) + + def pageRank( + self, + resetProbability: float = 0.15, + sourceId: str | int | None = None, + maxIter: int | None = None, + tol: float | None = None, + ) -> "GraphFrameConnect": + @final + class PageRank(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + reset_prob: float, + source_id: str | int | None, + max_iter: int | None, + tol: float | None, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.reset_prob = reset_prob + self.source_id = source_id + self.max_iter = max_iter + self.tol = tol + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.page_rank.CopyFrom( + pb.PageRank( + reset_probability=self.reset_prob, + source_id=( + None if self.source_id is None else make_str_or_long_id(self.source_id) + ), + max_iter=self.max_iter, + tol=self.tol, + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + if (maxIter is None) == (tol is None): + # TODO: in classic it is not an axception but assert; + # at the same time I think it should be an exception. + raise ValueError("Exactly one of maxIter or tol should be set.") + + new_vertices = _dataframe_from_plan( + PageRank( + self._vertices, + self._edges, + reset_prob=resetProbability, + source_id=sourceId, + max_iter=maxIter, + tol=tol, + ), + self._spark, + ) + # TODO: should this part to be optional? Like 'compute_edge_weights'? + return self._update_page_rank_edge_weights(new_vertices) + + def parallelPersonalizedPageRank( + self, + resetProbability: float = 0.15, + sourceIds: list[str | int] | None = None, + maxIter: int | None = None, + ) -> "GraphFrameConnect": + @final + class ParallelPersonalizedPageRank(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + reset_prob: float, + source_ids: list[str | int], + max_iter: int, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.reset_prob = reset_prob + self.source_ids = source_ids + self.max_iter = max_iter + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.parallel_personalized_page_rank.CopyFrom( + pb.ParallelPersonalizedPageRank( + reset_probability=self.reset_prob, + source_ids=[make_str_or_long_id(raw_id) for raw_id in self.source_ids], + max_iter=self.max_iter, + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + assert sourceIds is not None and len(sourceIds) > 0, ( + "Source vertices Ids sourceIds must be provided" + ) + assert maxIter is not None, "Max number of iterations maxIter must be provided" + + new_vertices = _dataframe_from_plan( + ParallelPersonalizedPageRank( + self._vertices, + self._edges, + reset_prob=resetProbability, + source_ids=sourceIds, + max_iter=maxIter, + ), + self._spark, + ) + return self._update_page_rank_edge_weights(new_vertices) + + def powerIterationClustering( + self, k: int, maxIter: int, weightCol: str | None = None + ) -> DataFrame: + @final + class PowerIterationClustering(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + k: int, + max_iter: int, + weight_col: str | None, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.k = k + self.max_iter = max_iter + self.weight_col = weight_col + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.power_iteration_clustering.CopyFrom( + pb.PowerIterationClustering( + k=self.k, + max_iter=self.max_iter, + weight_col=self.weight_col, + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + return _dataframe_from_plan( + PowerIterationClustering(self._vertices, self._edges, k, maxIter, weightCol), + self._spark, + ) + + def shortestPaths( + self, + landmarks: list[str | int], + algorithm: str, + use_local_checkpoints: bool, + checkpoint_interval: int, + storage_level: StorageLevel, + is_directed: bool, + ) -> DataFrame: + @final + class ShortestPaths(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + landmarks: list[str | int], + algorithm: str, + use_local_checkpoints: bool, + checkpoint_interval: int, + storage_level: StorageLevel, + is_directed: bool, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.landmarks = landmarks + self.algorithm = algorithm + self.use_local_checkpoints = use_local_checkpoints + self.checkpoint_interval = checkpoint_interval + self.storage_level = storage_level + self.is_directed = is_directed + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.shortest_paths.CopyFrom( + pb.ShortestPaths( + landmarks=[make_str_or_long_id(raw_id) for raw_id in self.landmarks], + algorithm=self.algorithm, + use_local_checkpoints=self.use_local_checkpoints, + checkpoint_interval=self.checkpoint_interval, + storage_level=storage_level_to_proto(self.storage_level), + is_directed=self.is_directed, + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + return _dataframe_from_plan( + ShortestPaths( + self._vertices, + self._edges, + landmarks, + algorithm, + use_local_checkpoints, + checkpoint_interval, + storage_level, + is_directed, + ), + self._spark, + ) + + def stronglyConnectedComponents(self, maxIter: int) -> DataFrame: + @final + class StronglyConnectedComponents(LogicalPlan): + def __init__(self, v: DataFrame, e: DataFrame, max_iter: int) -> None: + super().__init__(None) + self.v = v + self.e = e + self.max_iter = max_iter + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.strongly_connected_components.CopyFrom( + pb.StronglyConnectedComponents(max_iter=self.max_iter) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + return _dataframe_from_plan( + StronglyConnectedComponents(self._vertices, self._edges, maxIter), + self._spark, + ) + + def svdPlusPlus( + self, + rank: int = 10, + maxIter: int = 2, + minValue: float = 0.0, + maxValue: float = 5.0, + gamma1: float = 0.007, + gamma2: float = 0.007, + gamma6: float = 0.005, + gamma7: float = 0.015, + ) -> tuple[DataFrame, float]: + @final + class SVDPlusPlus(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + rank: int, + max_iter: int, + min_value: float, + max_value: float, + gamma1: float, + gamma2: float, + gamma6: float, + gamma7: float, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.rank = rank + self.max_iter = max_iter + self.min_value = min_value + self.max_value = max_value + self.gamma1 = gamma1 + self.gamma2 = gamma2 + self.gamma6 = gamma6 + self.gamma7 = gamma7 + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.svd_plus_plus.CopyFrom( + pb.SVDPlusPlus( + rank=self.rank, + max_iter=self.max_iter, + min_value=self.min_value, + max_value=self.max_value, + gamma1=self.gamma1, + gamma2=self.gamma2, + gamma6=self.gamma6, + gamma7=self.gamma7, + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + output = _dataframe_from_plan( + SVDPlusPlus( + self._vertices, + self._edges, + rank=rank, + max_iter=maxIter, + min_value=minValue, + max_value=maxValue, + gamma1=gamma1, + gamma2=gamma2, + gamma6=gamma6, + gamma7=gamma7, + ), + self._spark, + ) + + loss_row = output.select("loss").first() + assert loss_row is not None + return (cast(DataFrame, output.drop("loss")), float(loss_row["loss"])) + + def triangleCount( + self, storage_level: StorageLevel, algorithm: str, log_nom_entries: int + ) -> DataFrame: + @final + class TriangleCount(LogicalPlan): + def __init__(self, v: DataFrame, e: DataFrame, storage_level: StorageLevel) -> None: + super().__init__(None) + self.v = v + self.e = e + self.storage_level = storage_level + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.triangle_count.CopyFrom( + pb.TriangleCount( + storage_level=storage_level_to_proto(self.storage_level), + algorithm=algorithm, + lg_nom_entries=log_nom_entries, + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + return _dataframe_from_plan( + TriangleCount(self._vertices, self._edges, storage_level), self._spark + ) + + def maximal_independent_set( + self, + checkpoint_interval: int, + storage_level: StorageLevel, + use_local_checkpoints: bool, + seed: int, + ) -> DataFrame: + @final + class MaximalIndependentSet(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + checkpoint_interval: int, + storage_level: StorageLevel, + use_local_checkpoints: bool, + seed: int, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.checkpoint_interval = checkpoint_interval + self.storage_level = storage_level + self.use_local_checkpoints = use_local_checkpoints + self.seed = seed + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.mis.CopyFrom( + pb.MaximalIndependentSet( + checkpoint_interval=self.checkpoint_interval, + storage_level=storage_level_to_proto(self.storage_level), + use_local_checkpoints=self.use_local_checkpoints, + seed=self.seed, + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + return _dataframe_from_plan( + MaximalIndependentSet( + self._vertices, + self._edges, + checkpoint_interval, + storage_level, + use_local_checkpoints, + seed, + ), + self._spark, + ) + + def k_core( + self, + checkpoint_interval: int, + use_local_checkpoints: bool, + storage_level: StorageLevel, + ) -> DataFrame: + @final + class KCore(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + checkpoint_interval: int, + use_local_checkpoints: bool, + storage_level: StorageLevel, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.checkpoint_interval = checkpoint_interval + self.use_local_checkpoints = use_local_checkpoints + self.storage_level = storage_level + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.kcore.CopyFrom( + pb.KCore( + checkpoint_interval=self.checkpoint_interval, + use_local_checkpoints=self.use_local_checkpoints, + storage_level=storage_level_to_proto(self.storage_level), + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + return _dataframe_from_plan( + KCore( + self._vertices, + self._edges, + checkpoint_interval, + use_local_checkpoints, + storage_level, + ), + self._spark, + ) + + def hyper_anf( + self, + n_hops: int, + lg_nom_entries: int, + edge_filter: Column | str | None, + checkpoint_interval: int, + use_local_checkpoints: bool, + storage_level: StorageLevel, + ) -> DataFrame: + @final + class HyperANF(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + n_hops: int, + lg_nom_entries: int, + edge_filter: Column | str | None, + checkpoint_interval: int, + use_local_checkpoints: bool, + storage_level: StorageLevel, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.n_hops = n_hops + self.lg_nom_entries = lg_nom_entries + self.edge_filter = edge_filter + self.checkpoint_interval = checkpoint_interval + self.use_local_checkpoints = use_local_checkpoints + self.storage_level = storage_level + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + ha_message = pb.HyperANF( + n_hops=self.n_hops, + lg_nom_entries=self.lg_nom_entries, + checkpoint_interval=self.checkpoint_interval, + use_local_checkpoints=self.use_local_checkpoints, + storage_level=storage_level_to_proto(self.storage_level), + ) + if self.edge_filter is not None: + ha_message.edges_filter_expression.CopyFrom( + make_column_or_expr(self.edge_filter, session) + ) + graphframes_api_call.hyper_anf.CopyFrom(ha_message) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + return _dataframe_from_plan( + HyperANF( + self._vertices, + self._edges, + n_hops, + lg_nom_entries, + edge_filter, + checkpoint_interval, + use_local_checkpoints, + storage_level, + ), + self._spark, + ) + + def aggregate_neighbors( + self, + starting_vertices: Column | str, + max_hops: int, + accumulator_names: list[str], + accumulator_inits: list[Column | str], + accumulator_updates: list[Column | str], + stopping_condition: Column | str | None = None, + target_condition: Column | str | None = None, + required_vertex_attributes: list[str] | None = None, + required_edge_attributes: list[str] | None = None, + edge_filter: Column | str | None = None, + remove_loops: bool = False, + checkpoint_interval: int = 0, + use_local_checkpoints: bool = False, + storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> DataFrame: + @final + class AggregateNeighbors(LogicalPlan): + def __init__( + self, + v: DataFrame, + e: DataFrame, + starting_vertices: Column | str, + max_hops: int, + accumulator_names: list[str], + accumulator_inits: list[Column | str], + accumulator_updates: list[Column | str], + stopping_condition: Column | str | None, + target_condition: Column | str | None, + required_vertex_attributes: list[str] | None, + required_edge_attributes: list[str] | None, + edge_filter: Column | str | None, + remove_loops: bool, + checkpoint_interval: int, + use_local_checkpoints: bool, + storage_level: StorageLevel, + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.starting_vertices = starting_vertices + self.max_hops = max_hops + self.accumulator_names = accumulator_names + self.accumulator_inits = accumulator_inits + self.accumulator_updates = accumulator_updates + self.stopping_condition = stopping_condition + self.target_condition = target_condition + self.required_vertex_attributes = required_vertex_attributes or [] + self.required_edge_attributes = required_edge_attributes or [] + self.edge_filter = edge_filter + self.remove_loops = remove_loops + self.checkpoint_interval = checkpoint_interval + self.use_local_checkpoints = use_local_checkpoints + self.storage_level = storage_level + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + # Build the protobuf message + an_message = pb.AggregateNeighbors( + starting_vertices=make_column_or_expr(self.starting_vertices, session), + max_hops=self.max_hops, + accumulator_names=self.accumulator_names, + accumulator_inits=[ + make_column_or_expr(init, session) for init in self.accumulator_inits + ], + accumulator_updates=[ + make_column_or_expr(update, session) for update in self.accumulator_updates + ], + required_vertex_attributes=self.required_vertex_attributes, + required_edge_attributes=self.required_edge_attributes, + remove_loops=self.remove_loops, + checkpoint_interval=self.checkpoint_interval, + use_local_checkpoints=self.use_local_checkpoints, + storage_level=storage_level_to_proto(self.storage_level), + ) + + # Add optional fields if present + if self.stopping_condition is not None: + an_message.stopping_condition.CopyFrom( + make_column_or_expr(self.stopping_condition, session) + ) + + if self.target_condition is not None: + an_message.target_condition.CopyFrom( + make_column_or_expr(self.target_condition, session) + ) + + if self.edge_filter is not None: + an_message.edge_filter.CopyFrom(make_column_or_expr(self.edge_filter, session)) + + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.aggregate_neighbors.CopyFrom(an_message) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + return _dataframe_from_plan( + AggregateNeighbors( + self._vertices, + self._edges, + starting_vertices, + max_hops, + accumulator_names, + accumulator_inits, + accumulator_updates, + stopping_condition, + target_condition, + required_vertex_attributes, + required_edge_attributes, + edge_filter, + remove_loops, + checkpoint_interval, + use_local_checkpoints, + storage_level, + ), + self._spark, + ) + + def rw_embeddings(self, params: _RandomWalksEmbeddingsParameters) -> DataFrame: + @final + class RWEmbeddings(LogicalPlan): + def __init__( + self, v: DataFrame, e: DataFrame, params: _RandomWalksEmbeddingsParameters + ) -> None: + super().__init__(None) + self.v = v + self.e = e + self.params = params + + @override + def plan(self, session: SparkConnectClient) -> proto.Relation: + graphframes_api_call = GraphFrameConnect._get_pb_api_message( + self.v, self.e, session + ) + graphframes_api_call.rw_embeddings.CopyFrom( + pb.RandomWalkEmbeddings( + use_edge_direction=self.params.use_edge_direction, + rw_model=self.params.rw_model, + rw_max_nbrs=self.params.rw_max_nbrs, + rw_num_walks_per_node=self.params.rw_num_walks_per_node, + rw_batch_size=self.params.rw_batch_size, + rw_num_batches=self.params.rw_num_batches, + rw_seed=self.params.rw_seed, + rw_restart_probability=self.params.rw_restart_probability, + rw_temporary_prefix=self.params.rw_temporary_prefix, + rw_cached_walks=self.params.rw_cached_walks, + sequence_model=self.params.sequence_model, + hash2vec_context_size=self.params.hash2vec_context_size, + hash2vec_num_partitions=self.params.hash2vec_num_partitions, + hash2vec_embeddings_dim=self.params.hash2vec_embeddings_dim, + hash2vec_decay_function=self.params.hash2vec_decay_function, + hash2vec_gaussian_sigma=self.params.hash2vec_gaussian_sigma, + hash2vec_hashing_seed=self.params.hash2vec_hashing_seed, + hash2vec_sign_seed=self.params.hash2vec_sign_seed, + hash2vec_do_l2_norm=self.params.hash2vec_do_l2_norm, + hash2vec_safe_l2=self.params.hash2vec_safe_l2, + word2vec_max_iter=self.params.word2vec_max_iter, + word2vec_embeddings_dim=self.params.word2vec_embeddings_dim, + word2vec_window_size=self.params.word2vec_window_size, + word2vec_num_partitions=self.params.word2vec_num_partitions, + word2vec_min_count=self.params.word2vec_min_count, + word2vec_max_sentence_length=self.params.word2vec_max_sentence_length, + word2vec_seed=self.params.word2vec_seed, + word2vec_step_size=self.params.word2vec_step_size, + aggregate_neighbors=self.params.aggregate_neighbors, + aggregate_neighbors_max_nbrs=self.params.aggregate_neighbors_max_nbrs, + aggregate_neighbors_seed=self.params.aggregate_neighbors_seed, + clean_up_after_run=self.params.clean_up_after_run, + ) + ) + plan = self._create_proto_relation() + plan.graph_frames.CopyFrom(graphframes_api_call) + return plan + + return _dataframe_from_plan(RWEmbeddings(self._vertices, self._edges, params), self._spark) diff --git a/python/pyspark/graphframes/connect/utils.py b/python/pyspark/graphframes/connect/utils.py new file mode 100644 index 0000000000000..19804ac61b168 --- /dev/null +++ b/python/pyspark/graphframes/connect/utils.py @@ -0,0 +1,78 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from __future__ import annotations + +from pyspark.sql.connect.client import SparkConnectClient +from pyspark.sql.connect.column import Column +from pyspark.sql.connect.dataframe import DataFrame +from pyspark.sql.connect.expressions import Expression +from pyspark.sql.connect.plan import LogicalPlan +from pyspark.sql.connect.proto.graphframes_pb2 import ( + ColumnOrExpression, + StringOrLongID, +) +from pyspark.sql.connect.proto.graphframes_pb2 import StorageLevel as StorageLevelProto +from pyspark.storagelevel import StorageLevel + + +def dataframe_to_proto(df: DataFrame, client: SparkConnectClient) -> bytes: + plan = df._plan + assert plan is not None + assert isinstance(plan, LogicalPlan) + return plan.to_proto(client).SerializeToString() + + +def column_to_proto(col: Column, client: SparkConnectClient) -> bytes: + expr = col._expr + assert expr is not None + assert isinstance(expr, Expression) + return expr.to_plan(client).SerializeToString() + + +def make_column_or_expr(col: Column | str, client: SparkConnectClient) -> ColumnOrExpression: + if isinstance(col, Column): + return ColumnOrExpression(col=column_to_proto(col, client)) + else: + return ColumnOrExpression(expr=col) + + +def make_str_or_long_id(str_or_long: str | int) -> StringOrLongID: + if isinstance(str_or_long, str): + return StringOrLongID(string_id=str_or_long) + else: + return StringOrLongID(long_id=str_or_long) + + +def storage_level_to_proto(storage_level: StorageLevel) -> StorageLevelProto: + if storage_level == StorageLevel.DISK_ONLY: + return StorageLevelProto(disk_only=True) + elif storage_level == StorageLevel.DISK_ONLY_2: + return StorageLevelProto(disk_only_2=True) + elif storage_level == StorageLevel.DISK_ONLY_3: + return StorageLevelProto(disk_only_3=True) + elif storage_level == StorageLevel.MEMORY_AND_DISK: + return StorageLevelProto(memory_and_disk=True) + elif storage_level == StorageLevel.MEMORY_AND_DISK_2: + return StorageLevelProto(memory_and_disk_2=True) + elif storage_level == StorageLevel.MEMORY_ONLY: + return StorageLevelProto(memory_only=True) + elif storage_level == StorageLevel.MEMORY_ONLY_2: + return StorageLevelProto(memory_only_2=True) + elif storage_level == StorageLevel.MEMORY_AND_DISK_DESER: + return StorageLevelProto(memory_and_disk_deser=True) + else: + raise ValueError(f"Unknown storage level: {storage_level}") diff --git a/python/pyspark/graphframes/graphframe.py b/python/pyspark/graphframes/graphframe.py index edbe2de924910..ee785419ad7b6 100644 --- a/python/pyspark/graphframes/graphframe.py +++ b/python/pyspark/graphframes/graphframe.py @@ -17,221 +17,1456 @@ from __future__ import annotations -from typing import Union +import warnings +from typing import TYPE_CHECKING, Any, cast, final + +from typing_extensions import override -from pyspark.sql import Column, DataFrame from pyspark.sql import functions as F from pyspark.storagelevel import StorageLevel +class AggregateNeighbors: + """Helper class for referencing attributes in AggregateNeighbors expressions. + + Use these static methods in accumulator update expressions, stopping conditions, + and target conditions to reference vertex and edge attributes. + + **Example:** + + >>> result = g.aggregate_neighbors( + ... starting_vertices=F.col("id") == 1, + ... max_hops=5, + ... accumulator_names=["sum_values"], + ... accumulator_inits=[F.lit(0)], + ... accumulator_updates=[ + ... F.col("sum_values") + AggregateNeighbors.dst_attr("value") + ... ], + ... target_condition=AggregateNeighbors.dst_attr("id") == 10 + ... ) + """ + + @staticmethod + def src_attr(colName: str) -> Column: + """Reference a source vertex attribute. + + :param colName: Name of the source vertex attribute + :return: Column expression referencing the attribute + """ + return F.col("src_attributes").getField(colName) + + @staticmethod + def dst_attr(colName: str) -> Column: + """Reference a destination vertex attribute. + + :param colName: Name of the destination vertex attribute + :return: Column expression referencing the attribute + """ + return F.col("dst_attributes").getField(colName) + + @staticmethod + def edge_attr(colName: str) -> Column: + """Reference an edge attribute. + + :param colName: Name of the edge attribute + :return: Column expression referencing the attribute + """ + return F.col("edge_attributes").getField(colName) + + +from pyspark.graphframes.internal.utils import ( + _HASH2VEC_DECAY_FUNCTIONS, + _RandomWalksEmbeddingsParameters, +) +from pyspark.sql.utils import is_remote + +if TYPE_CHECKING: + from pyspark.graphframes.lib import Pregel + from pyspark.sql import Column, DataFrame + +"""Constant for the vertices ID column name.""" +ID = "id" + +"""Constant for the edge src column name.""" +SRC = "src" + +"""Constant for the edge dst column name.""" +DST = "dst" + +"""Constant for the edge column name.""" +EDGE = "edge" + +"""Constant for the weight column name.""" +WEIGHT = "weight" + + class GraphFrame: - """A graph whose vertices and edges are represented by Spark DataFrames. + """ + Represents a graph with vertices and edges stored as DataFrames. - The vertex DataFrame must contain a unique ``id`` column. The edge DataFrame must contain - ``src`` and ``dst`` columns identifying its source and destination vertices. All additional - columns are retained as graph attributes. + :param v: :class:`DataFrame` holding vertex information. + Must contain a column named "id" that stores unique + vertex IDs. + :param e: :class:`DataFrame` holding edge information. + Must contain two columns "src" and "dst" storing source + vertex IDs and destination vertex IDs of edges, respectively. - The initial in-tree API consists entirely of DataFrame operations and therefore supports both - classic Spark and Spark Connect. + >>> 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"]) + >>> g = GraphFrame(v, e) """ - ID = "id" - SRC = "src" - DST = "dst" - EDGE = "edge" + ID: str = ID + SRC: str = SRC + DST: str = DST + EDGE: str = EDGE + WEIGHT: str = WEIGHT + + @staticmethod + def _from_impl(impl: Any) -> "GraphFrame": + return GraphFrame(impl._vertices, impl._edges) + + def __init__(self, v: DataFrame, e: DataFrame) -> None: + """ + Initialize a GraphFrame from vertex DataFrame and edges DataFrame. + + :param v: :class:`DataFrame` holding vertex information. + Must contain a column named "id" that stores unique + vertex IDs. + :param e: :class:`DataFrame` holding edge information. + Must contain two columns "src" and "dst" storing source + vertex IDs and destination vertex IDs of edges, respectively. + """ + self._impl: Any + if self.ID not in v.columns: + raise ValueError( + "Vertex ID column {} missing from vertex DataFrame, which has columns: {}".format( + self.ID, ",".join(v.columns) + ) + ) + if self.SRC not in e.columns: + raise ValueError( + "Source vertex ID column {} missing from edge DataFrame, which has columns: {}".format( + self.SRC, ",".join(e.columns) + ) + ) + if self.DST not in e.columns: + raise ValueError( + "Destination vertex ID column {} missing from edge DataFrame, which has columns: {}".format( + self.DST, ",".join(e.columns) + ) + ) + if is_remote(): + from pyspark.graphframes.connect.graphframes_client import GraphFrameConnect + + self._impl = GraphFrameConnect(cast(Any, v), cast(Any, e)) # ty: ignore[invalid-argument-type] + else: + from pyspark.graphframes.classic.graphframe import GraphFrame as GraphFrameClassic - def __init__(self, vertices: DataFrame, edges: DataFrame) -> None: - self._require_column(vertices, self.ID, "Vertex ID") - self._require_column(edges, self.SRC, "Source vertex ID") - self._require_column(edges, self.DST, "Destination vertex ID") - self._vertices = vertices - self._edges = edges + self._impl = GraphFrameClassic(cast(Any, v), cast(Any, e)) # ty: ignore[invalid-argument-type] @property def vertices(self) -> DataFrame: - """The graph's vertex DataFrame.""" - return self._vertices + """ + :class:`DataFrame` holding vertex information, with unique column "id" + for vertex IDs. + """ + return self._impl._vertices + + @property + def edges(self) -> DataFrame: + """ + :class:`DataFrame` holding edge information, with unique columns "src" and + "dst" storing source vertex IDs and destination vertex IDs of edges, + respectively. + """ + return self._impl._edges @property def nodes(self) -> DataFrame: - """An alias for :attr:`vertices`.""" + """Alias to vertices.""" return self.vertices - @property - def edges(self) -> DataFrame: - """The graph's edge DataFrame.""" - return self._edges + @override + def __repr__(self) -> str: + # Exactly like in the scala core + v_cols = [self.ID] + [col for col in self._impl._vertices.columns if col != self.ID] + e_cols = [self.SRC, self.DST] + [ + col for col in self._impl._edges.columns if col not in {self.SRC, self.DST} + ] + v = self._impl._vertices.select(*v_cols).__repr__() + e = self._impl._edges.select(*e_cols).__repr__() - @property - def triplets(self) -> DataFrame: - """Return ``(source vertex)-[edge]->(destination vertex)`` triplets.""" - source_vertices = self.vertices.select( - self.vertices[self.ID].alias("__graphframes_src_id"), - self._nested(self.vertices, self.SRC), - ) - graph_edges = self.edges.select( - self.edges[self.SRC].alias("__graphframes_edge_src"), - self.edges[self.DST].alias("__graphframes_edge_dst"), - self._nested(self.edges, self.EDGE), - ) - destination_vertices = self.vertices.select( - self.vertices[self.ID].alias("__graphframes_dst_id"), - self._nested(self.vertices, self.DST), - ) - return ( - source_vertices.join( - graph_edges, - F.col("__graphframes_src_id") == F.col("__graphframes_edge_src"), - ) - .join( - destination_vertices, - F.col("__graphframes_dst_id") == F.col("__graphframes_edge_dst"), - ) - .select(self.SRC, self.EDGE, self.DST) - ) + return f"GraphFrame(v:{v}, e:{e})" + + def cache(self) -> "GraphFrame": + """Persist the dataframe representation of vertices and edges of the graph with the default + storage level. + """ + new_vertices = self._impl._vertices.cache() + new_edges = self._impl._edges.cache() + return GraphFrame(new_vertices, new_edges) + + def persist(self, storageLevel: StorageLevel = StorageLevel.MEMORY_ONLY) -> "GraphFrame": + """Persist the dataframe representation of vertices and edges of the graph with the given + storage level. + """ + new_vertices = self._impl._vertices.persist(storageLevel=storageLevel) + new_edges = self._impl._edges.persist(storageLevel=storageLevel) + return GraphFrame(new_vertices, new_edges) + + def unpersist(self, blocking: bool = False) -> "GraphFrame": + """Mark the dataframe representation of vertices and edges of the graph as non-persistent, + and remove all blocks for it from memory and disk. + """ + new_vertices = self._impl._vertices.unpersist(blocking=blocking) + new_edges = self._impl._edges.unpersist(blocking=blocking) + return GraphFrame(new_vertices, new_edges) @property def outDegrees(self) -> DataFrame: - """Return the out-degree of vertices having at least one outgoing edge.""" - return self.edges.groupBy(self.edges[self.SRC].alias(self.ID)).agg( - F.count("*").cast("int").alias("outDegree") + """ + The out-degree of each vertex in the graph, returned as a DataFrame with two columns: + - "id": the ID of the vertex + - "outDegree" (integer) storing the out-degree of the vertex + + Note that vertices with 0 out-edges are not returned in the result. + + :return: DataFrame with new vertices column "outDegree" + """ + return self._impl._edges.groupBy(F.col(self.SRC).alias(self.ID)).agg( + F.count("*").alias("outDegree") ) @property def inDegrees(self) -> DataFrame: - """Return the in-degree of vertices having at least one incoming edge.""" - return self.edges.groupBy(self.edges[self.DST].alias(self.ID)).agg( - F.count("*").cast("int").alias("inDegree") + """ + The in-degree of each vertex in the graph, returned as a DataFame with two columns: + - "id": the ID of the vertex + - "inDegree" (int) storing the in-degree of the vertex + + Note that vertices with 0 in-edges are not returned in the result. + + :return: DataFrame with new vertices column "inDegree" + """ + return self._impl._edges.groupBy(F.col(self.DST).alias(self.ID)).agg( + F.count("*").alias("inDegree") ) @property def degrees(self) -> DataFrame: - """Return the total degree of vertices incident to at least one edge.""" + """ + The degree of each vertex in the graph, returned as a DataFrame with two columns: + - "id": the ID of the vertex + - 'degree' (integer) the degree of the vertex + + Note that vertices with 0 edges are not returned in the result. + + :return: DataFrame with new vertices column "degree" + """ return ( - self.edges.select( - F.explode(F.array(self.edges[self.SRC], self.edges[self.DST])).alias(self.ID) + self._impl._edges.select( + F.explode(F.array(F.col(self.SRC), F.col(self.DST))).alias(self.ID) ) .groupBy(self.ID) - .agg(F.count("*").cast("int").alias("degree")) + .agg(F.count("*").alias("degree")) ) - def cache(self) -> "GraphFrame": - """Persist the vertex and edge DataFrames with their default storage level.""" - self.vertices.cache() - self.edges.cache() - return self + def type_out_degree(self, edge_type_col: str, edge_types: list[Any] | None = None) -> DataFrame: + """ + The out-degree of each vertex per edge type, returned as a DataFrame with two columns: + - "id": the ID of the vertex + - "outDegrees": a struct with a field for each edge type, storing the out-degree count - def persist( - self, storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER - ) -> "GraphFrame": - """Persist the vertex and edge DataFrames with ``storage_level``.""" - self.vertices.persist(storage_level) - self.edges.persist(storage_level) - return self + :param edge_type_col: Name of the column in edges DataFrame that contains edge types + :param edge_types: Optional list of edge type values. If None, edge types will be discovered automatically. + :return: DataFrame with columns "id" and "outDegrees" (struct type) + """ + if edge_types is not None: + pivot_df = self._impl._edges.groupBy(F.col(self.SRC).alias(self.ID)).pivot( + edge_type_col, edge_types + ) + else: + pivot_df = self._impl._edges.groupBy(F.col(self.SRC).alias(self.ID)).pivot( + edge_type_col + ) - def unpersist(self, blocking: bool = False) -> "GraphFrame": - """Remove the vertex and edge DataFrames from the cache.""" - self.vertices.unpersist(blocking) - self.edges.unpersist(blocking) - return self - - def filterVertices(self, condition: Union[Column, str]) -> "GraphFrame": - """Filter vertices and remove edges incident to any removed vertex.""" - filtered_vertices = self.vertices.filter(condition) - vertex_ids = filtered_vertices.select(filtered_vertices[self.ID]) - filtered_edges = self.edges.join( - vertex_ids, - self.edges[self.SRC] == vertex_ids[self.ID], - "left_semi", - ).join( - vertex_ids, - self.edges[self.DST] == vertex_ids[self.ID], - "left_semi", - ) - return GraphFrame(filtered_vertices, filtered_edges) - - def filterEdges(self, condition: Union[Column, str]) -> "GraphFrame": - """Filter edges while keeping all vertices.""" - return GraphFrame(self.vertices, self.edges.filter(condition)) + count_df = pivot_df.agg(F.count(F.lit(1))).na.fill(0) + struct_cols = [ + F.col(col_name).cast("int").alias(col_name) + for col_name in count_df.columns + if col_name != self.ID + ] - def dropIsolatedVertices(self) -> "GraphFrame": - """Return a graph without vertices that are not incident to an edge.""" - incident_ids = self.edges.select( - F.explode(F.array(self.edges[self.SRC], self.edges[self.DST])).alias(self.ID) + return count_df.select(F.col(self.ID), F.struct(*struct_cols).alias("outDegrees")) + + def type_in_degree(self, edge_type_col: str, edge_types: list[Any] | None = None) -> DataFrame: + """ + The in-degree of each vertex per edge type, returned as a DataFrame with two columns: + - "id": the ID of the vertex + - "inDegrees": a struct with a field for each edge type, storing the in-degree count + + :param edge_type_col: Name of the column in edges DataFrame that contains edge types + :param edge_types: Optional list of edge type values. If None, edge types will be discovered automatically. + :return: DataFrame with columns "id" and "inDegrees" (struct type) + """ + if edge_types is not None: + pivot_df = self._impl._edges.groupBy(F.col(self.DST).alias(self.ID)).pivot( + edge_type_col, edge_types + ) + else: + pivot_df = self._impl._edges.groupBy(F.col(self.DST).alias(self.ID)).pivot( + edge_type_col + ) + + count_df = pivot_df.agg(F.count(F.lit(1))).na.fill(0) + struct_cols = [ + F.col(col_name).cast("int").alias(col_name) + for col_name in count_df.columns + if col_name != self.ID + ] + return count_df.select(F.col(self.ID), F.struct(*struct_cols).alias("inDegrees")) + + def type_degree(self, edge_type_col: str, edge_types: list[Any] | None = None) -> DataFrame: + """ + The total degree of each vertex per edge type (both in and out), returned as a DataFrame + with two columns: + + - "id": the ID of the vertex + - "degrees": a struct with a field for each edge type, storing the total degree count + + :param edge_type_col: Name of the column in edges DataFrame that contains edge types + :param edge_types: Optional list of edge type values. If None, edge types will be discovered automatically. + :return: DataFrame with columns "id" and "degrees" (struct type) + """ + exploded_edges = self._impl._edges.select( + F.explode(F.array(F.col(self.SRC), F.col(self.DST))).alias(self.ID), + F.col(edge_type_col), ) - return GraphFrame(self.vertices.join(incident_ids, self.ID, "left_semi"), self.edges) - def as_reversed(self) -> "GraphFrame": - """Return a graph with the direction of every edge reversed.""" - attributes = [ - self.edges[name] for name in self.edges.columns if name not in {self.SRC, self.DST} + if edge_types is not None: + pivot_df = exploded_edges.groupBy(self.ID).pivot(edge_type_col, edge_types) + else: + pivot_df = exploded_edges.groupBy(self.ID).pivot(edge_type_col) + + count_df = pivot_df.agg(F.count(F.lit(1))).na.fill(0) + struct_cols = [ + F.col(col_name).cast("int").alias(col_name) + for col_name in count_df.columns + if col_name != self.ID ] - reversed_edges = self.edges.select( - self.edges[self.DST].alias(self.SRC), - self.edges[self.SRC].alias(self.DST), - *attributes, + + return count_df.select(F.col(self.ID), F.struct(*struct_cols).alias("degrees")) + + @property + def triplets(self) -> DataFrame: + """ + The triplets (source vertex)-[edge]->(destination vertex) for all edges in the graph. + + Returned as a :class:`DataFrame` with three columns: + - "src": source vertex with schema matching 'vertices' + - "edge": edge with schema matching 'edges' + - 'dst': destination vertex with schema matching 'vertices' + + :return: DataFrame with columns 'src', 'edge', and 'dst' + """ + return self._impl.triplets + + @property + def pregel(self) -> Pregel: + """ + Get the :class:`graphframes.classic.pregel.Pregel` + or :class`graphframes.connect.graphframes_client.Pregel` + object for running pregel. + + See :class:`graphframes.lib.Pregel` for more details. + """ + return self._impl.pregel + + def find(self, pattern: str) -> DataFrame: + """ + Motif finding: searching the graph for structural patterns. + + Motif finding uses a simple Domain-Specific Language (DSL) for expressing structural + queries. For example, ``graph.find("(a)-[e1]->(b); (b)-[e2]->(a)")`` will search for + pairs of vertices ``a``, ``b`` connected by edges in both directions. It returns a + :class:`DataFrame` of all such structures, with columns for each named element (vertex + or edge) in the motif. + + **Performance tip:** Motif finding translates patterns into a series of joins. Enabling + Spark's Cost-Based Optimizer (CBO) and join reordering can significantly improve + performance:: + + spark.conf.set("spark.sql.cbo.enabled", "true") + spark.conf.set("spark.sql.cbo.joinReorder.enabled", "true") + + The join reorder algorithm is bounded by ``spark.sql.cbo.joinReorder.dp.threshold`` + (default: ``12``). If the estimated number of joins in your motif exceeds this threshold, + increase it accordingly:: + + spark.conf.set("spark.sql.cbo.joinReorder.dp.threshold", "20") + + CBO relies on table statistics, so run ``ANALYZE TABLE COMPUTE STATISTICS`` on + the vertices and edges tables (or temp views) to ensure accurate statistics are available. + + :param pattern: String describing the motif to search for. + :return: DataFrame with one Row for each instance of the motif found. + """ + return self._impl.find(pattern=pattern) + + def filterVertices(self, condition: str | Column) -> "GraphFrame": + """ + Filters the vertices based on expression, remove edges containing any dropped vertices. + + :param condition: String or Column describing the condition expression for filtering. + :return: GraphFrame with filtered vertices and edges. + """ + return GraphFrame._from_impl(self._impl.filterVertices(condition=condition)) + + def filterEdges(self, condition: str | Column) -> "GraphFrame": + """ + Filters the edges based on expression, keep all vertices. + + :param condition: String or Column describing the condition expression for filtering. + :return: GraphFrame with filtered edges. + """ + + return GraphFrame._from_impl(self._impl.filterEdges(condition=condition)) + + def dropIsolatedVertices(self) -> "GraphFrame": + """ + Drops isolated vertices, vertices are not contained in any edges. + + :return: GraphFrame with filtered vertices. + """ + return GraphFrame._from_impl(self._impl.dropIsolatedVertices()) + + def detectingCycles( + self, + checkpoint_interval: int = 2, + use_local_checkpoints: bool = False, + storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> DataFrame: + """Find all cycles in the graph. + + An implementation of the Rocha-Thatte cycle detection algorithm. + Rocha, Rodrigo Caetano, and Bhalchandra D. Thatte. "Distributed cycle detection in + large-scale sparse graphs." Proceedings of Simpósio Brasileiro de Pesquisa Operacional + (SBPO'15) (2015): 1-11. + + Returns a DataFrame with unique cycles. + + :param checkpoint_interval: Pregel checkpoint interval, default is 2 + :param use_local_checkpoints: should local checkpoints be used instead of checkpointDir + :storage_level: the level of storage for both intermediate results and an output DataFrame + + :return: Persisted DataFrame with all the cycles + """ + return self._impl.detectingCycles(checkpoint_interval, use_local_checkpoints, storage_level) + + def bfs( + self, + fromExpr: str, + toExpr: str, + edgeFilter: str | None = None, + maxPathLength: int = 10, + ) -> DataFrame: + """ + Breadth-first search (BFS). + + See Scala documentation for more details. + + :return: DataFrame with one Row for each shortest path between matching vertices. + """ + return self._impl.bfs( + fromExpr=fromExpr, + toExpr=toExpr, + edgeFilter=edgeFilter, + maxPathLength=maxPathLength, ) - return GraphFrame(self.vertices, reversed_edges) - def as_undirected(self) -> "GraphFrame": - """Return an undirected graph by adding a reversed copy of every edge.""" - return GraphFrame(self.vertices, self.edges.unionByName(self.as_reversed().edges)) - - def validate(self) -> None: - """Run jobs that validate vertex uniqueness and edge endpoint integrity.""" - vertex_count = self.vertices.count() - distinct_vertex_count = self.vertices.select(self.ID).distinct().count() - if vertex_count != distinct_vertex_count: - raise ValueError( - f"Graph contains {vertex_count - distinct_vertex_count} duplicate vertices" + def all_paths( + self, + from_expr: Column | str, + to_expr: Column | str, + edge_filter: Column | str | None = None, + max_path_length: int = 5, + is_directed: bool = True, + checkpoint_interval: int = 2, + use_local_checkpoints: bool = False, + storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> DataFrame: + """ + Computes all simple paths between source and destination vertices. + + This algorithm enumerates paths up to ``max_path_length`` hops. It supports directed + and undirected traversal as well as optional edge filtering. It returns all simple + paths between source and destination vertices. Here the term "simple" means no + repeated vertices. For example, if there are paths A-B-C, A-D-C and the edge B-A, + and the user asked to find all the paths between "A" and "C", only A-B-C and A-D-C + will be returned, but not A-B-A-D-C. + + The default value of ``max_path_length`` is 5. Keep in mind that requesting + ``max_path_length`` on the scale of the graph diameter may cause the algorithm to + try to return (almost) all simple paths in the graph, which can create huge + performance degradation or even OOM-like errors. + + **Returned DataFrame schema:** + + - ``path``: array of vertex ids in traversal order + - ``len``: number of edges in the path (Long) + + .. note:: + In the case of an undirected graph, the algorithm runs on an internal graph + made by union of edges and reversed edges. It is assumed that the graph does + not have multi-edges. Results may be unstable and unpredictable for graphs + with multi-edges. + + **Example:** + + >>> paths = g.all_paths( + ... from_expr="name = 'A'", + ... to_expr="name = 'C'", + ... max_path_length=3, + ... ) + >>> paths.show() + + :param from_expr: Column expression or SQL expression string identifying the + source (starting) vertices. + :param to_expr: Column expression or SQL expression string identifying the + destination (target) vertices. + :param edge_filter: Optional Column expression or SQL expression string applied + to edges during traversal. Only edges satisfying this condition are considered. + If not provided, all edges are considered. + :param max_path_length: Maximum number of edges in a path; must be greater than 0. + Default is 5. Setting a large value (e.g., on the scale of the graph diameter) + may cause severe performance degradation or out-of-memory errors. + :param is_directed: Whether to use directed traversal. If False, the graph is + treated as undirected by internally unioning edges with reversed edges. + Default is True. + :param checkpoint_interval: Checkpoint every N iterations, 0 = disabled (default: 0) + :param use_local_checkpoints: Use local checkpoints (faster but less reliable) + :param storage_level: Storage level for intermediate results + + :return: DataFrame with columns ``path`` (array of vertex ids) and ``len`` + (number of edges in the path). + """ + return self._impl.all_paths( + from_expr=from_expr, + to_expr=to_expr, + 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, + ) + + def aggregateMessages( + self, + aggCol: list[Column | str] | Column, + sendToSrc: list[Column | str] | Column | str | None = None, + sendToDst: list[Column | str] | Column | str | None = None, + intermediate_storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> DataFrame: + """ + Aggregates messages from the neighbours. + + When specifying the messages and aggregation function, the user may reference columns using + the static methods in :class:`graphframes.lib.AggregateMessages`. + + See Scala documentation for more details. + + Warning! The result of this method is persisted DataFrame object! Users should handle unpersist + to avoid possible memory leaks! + + :param aggCol: the requested aggregation output either as a collection of + :class:`pyspark.sql.Column` or SQL expression string + :param sendToSrc: message sent to the source vertex of each triplet either as + a collection of :class:`pyspark.sql.Column` or SQL expression string (default: None) + :param sendToDst: message sent to the destination vertex of each triplet either as + collection of :class:`pyspark.sql.Column` or SQL expression string (default: None) + :param intermediate_storage_level: the level of intermediate storage that will be used + for both intermediate result and the output. + + :return: Persisted DataFrame with columns for the vertex ID and the resulting aggregated message. + The name of the resulted message column is based on the alias of the provided aggCol! + """ + + if sendToDst is None: + sendToDst = [] + if sendToSrc is None: + sendToSrc = [] + + # Back-compatibility workaround + if not isinstance(aggCol, list): + warnings.warn( + "Passing single column to aggCol is deprecated, use list", + DeprecationWarning, + ) + return self.aggregateMessages( + [aggCol], sendToSrc, sendToDst, intermediate_storage_level + ) + if not isinstance(sendToSrc, list): + warnings.warn( + "Passing single column to sendToSrc is deprecated, use list", + DeprecationWarning, + ) + return self.aggregateMessages( + aggCol, [sendToSrc], sendToDst, intermediate_storage_level + ) + if not isinstance(sendToDst, list): + warnings.warn( + "Passing single column to sendToDst is deprecated, use list", + DeprecationWarning, + ) + return self.aggregateMessages( + aggCol, sendToSrc, [sendToDst], intermediate_storage_level ) - endpoints = ( - self.edges.select(self.edges[self.SRC].alias(self.ID)) - .union(self.edges.select(self.edges[self.DST].alias(self.ID))) - .distinct() + if len(aggCol) == 0: + raise TypeError("At least one aggregation column should be provided!") + + if (len(sendToSrc) == 0) and (len(sendToDst) == 0): + raise ValueError("Either `sendToSrc`, `sendToDst`, or both have to be provided") + return self._impl.aggregateMessages( + aggCol=aggCol, + sendToSrc=sendToSrc, + sendToDst=sendToDst, + intermediate_storage_level=intermediate_storage_level, + ) + + # Standard algorithms + + def connectedComponents( + self, + algorithm: str = "graphframes", + checkpointInterval: int = 2, + broadcastThreshold: int = 1000000, + useLabelsAsComponents: bool = False, + use_local_checkpoints: bool = False, + max_iter: int = (1 << 31) - 2, + storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> DataFrame: + """ + Computes the connected components of the graph. + + See Scala documentation for more details. + + :param algorithm: connected components algorithm to use (default: "graphframes") + Supported algorithms are "two_phase", "randomized_contraction", + "graphframes" (deprecated alias for "two_phase") and "graphx". + :param checkpointInterval: checkpoint interval in terms of number of iterations (default: 2) + :param broadcastThreshold: broadcast threshold in propagating component assignments + (default: 1000000). Passing -1 disable manual broadcasting and + allows AQE to handle skewed joins. This mode is much faster + and is recommended to use. Default value may be changed to -1 + in the future versions of GraphFrames. + :param useLabelsAsComponents: if True, uses the vertex labels as components, otherwise will + use longs + :param use_local_checkpoints: should local checkpoints be used, default false; + local checkpoints are faster and does not require to set + a persistent checkpointDir; from the other side, local + checkpoints are less reliable and require executors to have + big enough local disks. + :param storage_level: storage level for both intermediate and final dataframes. + + :return: DataFrame with new vertices column "component" + """ + return self._impl.connectedComponents( + algorithm=algorithm, + checkpointInterval=checkpointInterval, + broadcastThreshold=broadcastThreshold, + useLabelsAsComponents=useLabelsAsComponents, + use_local_checkpoints=use_local_checkpoints, + max_iter=max_iter, + storage_level=storage_level, + ) + + def maximal_independent_set( + self, + seed: int = 42, + checkpoint_interval: int = 2, + use_local_checkpoints: bool = False, + storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> DataFrame: + """ + This method implements a distributed algorithm for finding a Maximal Independent Set (MIS) + in a graph. + + An MIS is a set of vertices such that no two vertices in the set are adjacent (i.e., there + is no edge between any two vertices in the set), and the set is maximal, meaning that adding + any other vertex to the set would violate the independence property. Note that this + implementation finds a maximal (but not necessarily maximum) independent set; that is, it + ensures no more vertices can be added to the set, but does not guarantee that the set has + the largest possible number of vertices among all possible independent sets in the graph. + + The algorithm implemented here is based on the paper: Ghaffari, Mohsen. "An improved + distributed algorithm for maximal independent set." Proceedings of the twenty-seventh annual + ACM-SIAM symposium on Discrete algorithms. Society for Industrial and Applied Mathematics, + 2016. + + Note: This is a randomized, non-deterministic algorithm. The result may vary between runs + even if a fixed random seed is provided because of how Apache Spark works. + + :param seed: random seed used for tie-breaking in the algorithm (default: 42) + :param checkpoint_interval: checkpoint interval in terms of number of iterations (default: 2) + :param use_local_checkpoints: whether to use local checkpoints (default: False); + local checkpoints are faster and do not require setting + a persistent checkpoint directory; however, they are less + reliable and require executors to have sufficient local disk space. + :param storage_level: storage level for both intermediate and final DataFrames + (default: MEMORY_AND_DISK_DESER) + + :return: DataFrame containing the ``id`` of each vertex in the Maximal Independent Set + """ + return self._impl.maximal_independent_set( + checkpoint_interval, storage_level, use_local_checkpoints, seed + ) + + def k_core( + self, + checkpoint_interval: int = 2, + use_local_checkpoints: bool = False, + storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> DataFrame: + """ + The k-core is the maximal subgraph such that every vertex has at least degree k. + The k-core metric is a measure of the centrality of a node in a network, based on its + degree and the degrees of its neighbors. Nodes with higher k-core values are considered + to be more central and influential within the network. + + This implementation is based on the algorithm described in: + Mandal, Aritra, and Mohammad Al Hasan. "A distributed k-core decomposition algorithm + on spark." 2017 IEEE International Conference on Big Data (Big Data). IEEE, 2017. + + :param checkpoint_interval: Pregel checkpoint interval, default is 2 + :param use_local_checkpoints: should local checkpoints be used instead of checkpointDir + :param storage_level: the level of storage for both intermediate results and an output DataFrame + + :return: Persisted DataFrame with ID and k-core values (column "kcore") + """ + return self._impl.k_core(checkpoint_interval, use_local_checkpoints, storage_level) + + def hyper_anf( + self, + n_hops: int = 3, + lg_nom_entries: int = 12, + edge_filter: Column | str | None = None, + checkpoint_interval: int = 2, + use_local_checkpoints: bool = False, + storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> DataFrame: + """HyperANF-style approximation of the neighbourhood function. + + This implementation is inspired by + `HyperANF: Approximating the Neighbourhood Function of Very Large Graphs on a Budget + `_ + (Vigna, Boldi, Rosa; 2010). + + The input graph is treated as directed: for each vertex, reachability is computed by + following outgoing edges from ``src`` to ``dst``. + + Compared with the cumulative neighbourhood-function presentation in the paper, this + implementation returns one column per hop: ``hop_0``, ``hop_1``, ``hop_2``, …, ``hop_N``. + The ``hop_0`` column contains a HyperLogLog sketch of the source vertex itself, and each + ``hop_k`` column for ``k >= 1`` contains a HyperLogLog sketch of the set of vertices + reachable in exactly ``k`` hops. To derive the cumulative approximate neighbourhood + function for distances up to some hop ``k``, combine ``hop_0`` through ``hop_k`` with + ``hll_union`` and then apply ``hll_sketch_estimate`` to the merged sketch. + + The computation can also be restricted to a subgraph by supplying an edge filter + expression via ``edge_filter``. A common use case is to filter on ``src``, for example + ``"src IN (1, 2, 3)"``, to obtain sketches only for a selected set of starting vertices. + + **Example:** + + >>> result = g.hyper_anf(n_hops=2) + >>> result.columns + ['id', 'hop_0', 'hop_1', 'hop_2'] + + :param n_hops: Maximum hop distance to compute (positive integer). The result will + contain columns ``hop_0`` through ``hop_N`` where ``N = n_hops``. + :param lg_nom_entries: Log2 of nominal entries used by HLL sketch aggregations. + Must be between 4 and 21 (inclusive). Higher values increase accuracy at the + cost of memory. Default is 12. + :param edge_filter: Optional column expression or SQL expression string applied to + edges before computation. Only edges satisfying this predicate participate in + the directed reachability expansion. + :param checkpoint_interval: Checkpoint interval in terms of number of iterations + (default: 2). Use 0 to disable checkpointing. + :param use_local_checkpoints: Whether to use local checkpoints instead of a + persistent checkpoint directory. Local checkpoints are faster but less reliable. + :param storage_level: Storage level for intermediate and final DataFrames. + + :return: Persisted DataFrame with vertex ``id`` and one sketch column per hop + (``hop_0``, ``hop_1``, …, ``hop_N``). + """ + if n_hops <= 0: + raise ValueError("n_hops must be a positive integer") + if not (4 <= lg_nom_entries <= 21): + raise ValueError("lg_nom_entries must be between 4 and 21 (inclusive)") + + return self._impl.hyper_anf( + n_hops=n_hops, + lg_nom_entries=lg_nom_entries, + edge_filter=edge_filter, + checkpoint_interval=checkpoint_interval, + use_local_checkpoints=use_local_checkpoints, + storage_level=storage_level, ) - missing_endpoint_count = endpoints.join(self.vertices, self.ID, "left_anti").count() - if missing_endpoint_count: + + def labelPropagation( + self, + maxIter: int, + algorithm: str = "graphx", + use_local_checkpoints: bool = False, + checkpoint_interval: int = 2, + storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> DataFrame: + """ + Runs static label propagation for detecting communities in networks. + + See Scala documentation for more details. + + :param maxIter: the number of iterations to be performed + :param algorithm: implementation to use, possible values are "graphframes" and "graphx"; + "graphx" is faster for small-medium sized graphs, + "graphframes" requires less amount of memory + :param use_local_checkpoints: should local checkpoints be used, default false; + local checkpoints are faster and does not require to set + a persistent checkpointDir; from the other side, local + checkpoints are less reliable and require executors to have + big enough local disks. + :checkpoint_interval: How often should the intermediate result be checkpointed; + Using big value here may tend to huge logical plan growth due + to the iterative nature of the algorithm. + :param storage_level: storage level for both intermediate and final dataframes. + + :return: Persisted DataFrame with new vertices column "label" + """ + return self._impl.labelPropagation( + maxIter=maxIter, + algorithm=algorithm, + use_local_checkpoints=use_local_checkpoints, + checkpoint_interval=checkpoint_interval, + storage_level=storage_level, + ) + + def neighborhood_aware_cdlp( + self, + max_iter: int, + structural_similarity_multiplier: float = 0.5, + ignore_direct_links: bool = False, + initial_label_col: str | None = None, + is_directed: bool = True, + lg_nom_entries: int = 12, + use_local_checkpoints: bool = False, + checkpoint_interval: int = 2, + storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> DataFrame: + """ + Neighborhood-aware community detection via weighted label propagation. + + This algorithm is a Label Propagation variant where each incoming label vote is weighted + by a combination of: + - optional direct-link baseline strength (enabled unless + ``ignore_direct_links = True``), and + - neighborhood-overlap strength + (``structural_similarity_multiplier * common_neighbors``). + + Intuitively, labels from neighbors that are structurally similar to the destination + (many common neighbors) can be amplified instead of treating all edges equally. + + At each iteration, every vertex aggregates weighted incoming votes by label and picks + the label with maximum total weight. + + Edge-weight regimes: + - ``ignore_direct_links = False``: + ``edge_weight = 1 + structural_similarity_multiplier * common_neighbors`` + - ``ignore_direct_links = True``: + ``edge_weight = structural_similarity_multiplier * common_neighbors`` + + :param max_iter: maximum number of propagation rounds. + :param structural_similarity_multiplier: scales neighborhood-overlap contribution. + Must be non-negative. + :param ignore_direct_links: whether to drop direct-link baseline vote mass. + :param initial_label_col: optional vertex column used to initialize labels. + :param is_directed: whether to treat edges as directed. + :param lg_nom_entries: log2 nominal entries used by Theta sketch aggregations. + :param use_local_checkpoints: whether to use local checkpoints. + :param checkpoint_interval: checkpoint interval in iterations. + :param storage_level: storage level for intermediate datasets. + + :return: Persisted DataFrame with new vertex column ``label``. + """ + if structural_similarity_multiplier < 0.0: + raise ValueError("structural_similarity_multiplier must be >= 0") + + if ignore_direct_links and structural_similarity_multiplier == 0.0: raise ValueError( - f"Graph contains {missing_endpoint_count} edge endpoints without matching vertices" + "structural_similarity_multiplier must be > 0 when ignore_direct_links is True" ) - @classmethod - def from_edges( - cls, - edges: DataFrame, - storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + return self._impl.neighborhood_aware_cdlp( + max_iter=max_iter, + structural_similarity_multiplier=structural_similarity_multiplier, + ignore_direct_links=ignore_direct_links, + use_local_checkpoints=use_local_checkpoints, + checkpoint_interval=checkpoint_interval, + storage_level=storage_level, + is_directed=is_directed, + lg_nom_entries=lg_nom_entries, + initial_label_col=initial_label_col, + ) + + def pageRank( + self, + resetProbability: float = 0.15, + sourceId: Any | None = None, + maxIter: int | None = None, + tol: float | None = None, ) -> "GraphFrame": - """Create a graph by deriving and persisting distinct vertices from ``edges``.""" - cls._require_column(edges, cls.SRC, "Source vertex ID") - cls._require_column(edges, cls.DST, "Destination vertex ID") - vertices = ( - edges.select(edges[cls.SRC].alias(cls.ID)) - .union(edges.select(edges[cls.DST].alias(cls.ID))) - .distinct() - .persist(storage_level) + """ + Runs the PageRank algorithm on the graph. + Note: Exactly one of fixed_num_iter or tolerance must be set. + + See Scala documentation for more details. + + :param resetProbability: Probability of resetting to a random vertex. + :param sourceId: (optional) the source vertex for a personalized PageRank. + :param maxIter: If set, the algorithm is run for a fixed number + of iterations. This may not be set if the `tol` parameter is set. + :param tol: If set, the algorithm is run until the given tolerance. + This may not be set if the `numIter` parameter is set. + :return: GraphFrame with new vertices column "pagerank" and new edges column "weight" + """ + return GraphFrame._from_impl( + self._impl.pageRank( + resetProbability=resetProbability, + sourceId=sourceId, + maxIter=maxIter, + tol=tol, + ) ) - return cls(vertices, edges) - def __repr__(self) -> str: - vertex_columns = [self.ID] + [name for name in self.vertices.columns if name != self.ID] - edge_columns = [self.SRC, self.DST] + [ - name for name in self.edges.columns if name not in {self.SRC, self.DST} - ] - return ( - f"GraphFrame(v:{self.vertices.select(*vertex_columns)!r}, " - f"e:{self.edges.select(*edge_columns)!r})" + def parallelPersonalizedPageRank( + self, + resetProbability: float = 0.15, + sourceIds: list[Any] | None = None, + maxIter: int | None = None, + ) -> "GraphFrame": + """ + Run the personalized PageRank algorithm on the graph, + from the provided list of sources in parallel for a fixed number of iterations. + + See Scala documentation for more details. + + :param resetProbability: Probability of resetting to a random vertex + :param sourceIds: the source vertices for a personalized PageRank + :param maxIter: the fixed number of iterations this algorithm runs + :return: GraphFrame with new vertices column "pageranks" and new edges column "weight" + """ + return GraphFrame._from_impl( + self._impl.parallelPersonalizedPageRank( + resetProbability=resetProbability, sourceIds=sourceIds, maxIter=maxIter + ) ) - @staticmethod - def _nested(dataframe: DataFrame, name: str) -> Column: - return F.struct(*[dataframe[column] for column in dataframe.columns]).alias(name) + def shortestPaths( + self, + landmarks: list[str | int], + algorithm: str = "graphx", + use_local_checkpoints: bool = False, + checkpoint_interval: int = 2, + storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + is_directed: bool = True, + ) -> DataFrame: + """ + Runs the shortest path algorithm from a set of landmark vertices in the graph. - @staticmethod - def _require_column(dataframe: DataFrame, column: str, label: str) -> None: - if column not in dataframe.columns: - available = ", ".join(dataframe.columns) - raise ValueError( - f"{label} column '{column}' is missing; available columns: {available}" + See Scala documentation for more details. + + :param landmarks: a set of one or more landmarks + :param algorithm: implementation to use, possible values are "graphframes" and "graphx"; + "graphx" is faster for small-medium sized graphs, + "graphframes" requires less amount of memory + :param use_local_checkpoints: should local checkpoints be used, default false; + local checkpoints are faster and does not require to set + a persistent checkpointDir; from the other side, local + checkpoints are less reliable and require executors to have + big enough local disks. + :param checkpoint_interval: How often should the intermediate result be checkpointed; + Using big value here may tend to huge logical plan growth due + to the iterative nature of the algorithm. + :param storage_level: storage level for both intermediate and final dataframes. + :param is_directed: should algorithm find directed paths or any paths. + + :return: persistent DataFrame with new vertices column "distances" + """ + return self._impl.shortestPaths( + landmarks=landmarks, + algorithm=algorithm, + use_local_checkpoints=use_local_checkpoints, + checkpoint_interval=checkpoint_interval, + storage_level=storage_level, + is_directed=is_directed, + ) + + def stronglyConnectedComponents(self, maxIter: int) -> DataFrame: + """ + Runs the strongly connected components algorithm on this graph. + + See Scala documentation for more details. + + :param maxIter: the number of iterations to run + :return: DataFrame with new vertex column "component" + """ + return self._impl.stronglyConnectedComponents(maxIter=maxIter) + + def svdPlusPlus( + self, + rank: int = 10, + maxIter: int = 2, + minValue: float = 0.0, + maxValue: float = 5.0, + gamma1: float = 0.007, + gamma2: float = 0.007, + gamma6: float = 0.005, + gamma7: float = 0.015, + ) -> tuple[DataFrame, float]: + """Runs the SVD++ algorithm for Collaborative Filtering. + + Based on the paper "Factorization Meets the Neighborhood: a Multifaceted Collaborative + Filtering Model" by Yehuda Koren (2008). + + **Algorithm Description** + SVD++ improves upon standard Matrix Factorization by incorporating implicit feedback + (the history of items a user has interacted with) alongside explicit ratings. + The prediction rule is: + ``r_ui = µ + b_u + b_i + q_i^T * (p_u + |N(u)|^-0.5 * sum(y_j for j in N(u)))`` + + **Input Requirements** + The input graph must be a **Directed Bipartite Graph**: + - **Vertices**: A mix of Users and Items. + - **Edges**: Directed strictly from **User (src) -> Item (dst)**. + - **Edge Attribute**: Represents the rating (weight). + + :param rank: The number of latent factors (embedding size). + :param maxIter: The maximum number of iterations. + :param minValue: The minimum possible rating value (used for clipping predictions). + :param maxValue: The maximum possible rating value (used for clipping predictions). + :param gamma1: Learning rate for bias parameters (`b_u`, `b_i`). + :param gamma2: Learning rate for factor parameters (`p_u`, `q_i`, `y_j`). + :param gamma6: Regularization coefficient for bias parameters. + :param gamma7: Regularization coefficient for factor parameters. + :return: A tuple ``(v, loss)`` where: + - ``v`` is a DataFrame of vertices containing the trained model parameters (embeddings). + - ``loss`` is the final training loss (double). + + **Output DataFrame Columns** + The returned DataFrame ``v`` contains the following new columns containing the model parameters: + + - **column1** (Array[Double]): Primary Latent Factors (Explicit Embedding). + - For Users: Preferences vector (`p_u`). + - For Items: Characteristics vector (`q_i`). + - **column2** (Array[Double]): Implicit Factors (Implicit Embedding). + - For Items: Influence vector (`y_i`). + - For Users: Unused/Zero (users aggregate `y` from neighbors). + - **column3** (Double): Bias term. + - For Users: User bias (`b_u`). + - For Items: Item bias (`b_i`). + - **column4** (Double): Implicit Normalization term. + - For Users: Precomputed ``|N(u)|^-0.5``. + - For Items: Unused. + """ + return self._impl.svdPlusPlus( + rank=rank, + maxIter=maxIter, + minValue=minValue, + maxValue=maxValue, + gamma1=gamma1, + gamma2=gamma2, + gamma6=gamma6, + gamma7=gamma7, + ) + + def triangleCount( + self, storage_level: StorageLevel, algorithm: str = "exact", lg_nom_entries: int = 12 + ) -> DataFrame: + """ + Computes the number of triangles passing through each vertex. + This algorithm identifies sets of three vertices where each pair is connected by an edge. + + The implementation provides two algorithms: + + - "exact": Computes the exact triangle count using set intersection of neighbor lists. + + Note: This method can fail or encounter OOM errors on power-law graphs or graphs with + very high-degree nodes, as it requires collecting and intersecting the full neighbor + lists for the source and destination vertices of every edge. + + - "approx": Uses DataSketches (Theta sketches) to estimate the triangle count. This trades off perfect accuracy for significantly improved performance and lower memory overhead, making it suitable for large-scale or dense graphs. + + :param storage_level: Storage level for caching intermediate DataFrames. + :param algorithm: The triangle counting algorithm to use, "exact" or "approx" (default: "exact"). + :param lg_nom_entries: The log2 of the nominal entries for the Theta sketch (only used + if algorithm="approx"). Higher values increase accuracy at the + cost of memory. (default: 12). + :return: A DataFrame containing the vertex "id" and the triangle "count". + """ + spark_version = self._impl._spark.version + if (spark_version[:3] < "4.1") and (algorithm == "approx"): + err_msg = "approximate algorithm requires Spark 4.1+" + err_msg += f" version {spark_version[:3]} is not supported" + raise ValueError(err_msg) + return self._impl.triangleCount( + storage_level=storage_level, algorithm=algorithm, log_nom_entries=lg_nom_entries + ) + + def powerIterationClustering( + self, k: int, maxIter: int, weightCol: str | None = None + ) -> DataFrame: + """ + Power Iteration Clustering (PIC), a scalable graph clustering algorithm developed by Lin and Cohen. From the abstract: PIC finds a very low-dimensional embedding of a dataset using truncated power iteration on a normalized pair-wise similarity matrix of the data. + + :param k: the numbers of clusters to create + :param maxIter: param for maximum number of iterations (>= 0) + :param weightCol: optional name of weight column, 1.0 is used if not provided + + :return: DataFrame with new column "cluster" + """ + return self._impl.powerIterationClustering(k, maxIter, weightCol) + + def validate( + self, + check_vertices: bool = True, + intermediate_storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> None: + """ + Validates the consistency and integrity of a graph by performing checks on the vertices and + edges. + + :param check_vertices: a flag to indicate whether additional vertex consistency checks + should be performed. If true, the method will verify that all vertices in the vertex + DataFrame are represented in the edge DataFrame and vice versa. It is slow on big graphs. + :param intermediate_storage_level: the storage level to be used when persisting + intermediate DataFrame computations during the validation process. + :return: Unit, as the method performs validation checks and throws an exception if + validation fails. + :raises ValueError: if there are any inconsistencies in the graph, such as duplicate + vertices, mismatched vertices between edges and vertex DataFrames or missing + connections. + """ + persisted_vertices = self.vertices.persist(intermediate_storage_level) + row = persisted_vertices.select(F.count_distinct(F.col(ID))).first() + assert row is not None # for type checker + count_distinct_vertices = row[0] + assert isinstance(count_distinct_vertices, int) # for type checker + total_count_vertices = persisted_vertices.count() + if count_distinct_vertices != total_count_vertices: + _msg = "Graph contains ({}) duplicate vertices." + + raise ValueError(_msg.format(total_count_vertices - count_distinct_vertices)) + if check_vertices: + vertices_set_from_edges = ( + self.edges.select(F.col(SRC).alias(ID)) + .union(self.edges.select(F.col(DST).alias(ID))) + .distinct() + .persist(intermediate_storage_level) + ) + count_vertices_from_edges = vertices_set_from_edges.count() + if count_vertices_from_edges > count_distinct_vertices: + _msg = "Graph is inconsistent: edges has {} " + _msg += "vertices, but vertices has {} vertices." + raise ValueError(_msg.format(count_vertices_from_edges, count_distinct_vertices)) + + combined = vertices_set_from_edges.join(self.vertices, ID, "left_anti") + count_of_bad_vertices = combined.count() + if count_of_bad_vertices > 0: + _msg = "Vertices DataFrame does not contain all edges src/dst. " + _msg += "Found {} edges src/dst that are not in the vertices DataFrame." + raise ValueError(_msg.format(count_of_bad_vertices)) + _ = vertices_set_from_edges.unpersist() + _ = persisted_vertices.unpersist() + + def as_undirected(self) -> "GraphFrame": + """ + Converts the directed graph into an undirected graph by ensuring that all directed edges are + bidirectional. For every directed edge (src, dst), a corresponding edge (dst, src) is added. + + :return: A new GraphFrame representing the undirected graph. + """ + + edge_attr_columns = [c for c in self.edges.columns if c not in [SRC, DST]] + + # Create the undirected edges by duplicating each edge in both directions + + # 3.5.x problem: selecting empty struct fails on spark connect + # TODO: remove after removing 3.5.x + + if edge_attr_columns: + forward_edges = self.edges.select( + F.col(SRC), F.col(DST), F.struct(*edge_attr_columns).alias(EDGE) + ) + backward_edges = self.edges.select( + F.col(DST).alias(SRC), + F.col(SRC).alias(DST), + F.struct(*edge_attr_columns).alias(EDGE), + ) + new_edges = forward_edges.union(backward_edges).select(SRC, DST, EDGE) + else: + forward_edges = self.edges.select(F.col(SRC), F.col(DST)) + backward_edges = self.edges.select(F.col(DST).alias(SRC), F.col(SRC).alias(DST)) + new_edges = forward_edges.union(backward_edges).select(SRC, DST) + + # Preserve additional edge attributes + edge_columns = [F.col(EDGE).getField(c).alias(c) for c in edge_attr_columns] + + # Select all columns including the new edge attributes + selected_columns = [F.col(SRC), F.col(DST)] + edge_columns + new_edges = new_edges.select(*selected_columns) + + return GraphFrame(self.vertices, new_edges) + + def as_reversed(self) -> "GraphFrame": + """ + Reverses the direction of all edges in the graph. For every directed edge (src, dst), + the resulting graph will contain an edge (dst, src) with the same attributes. + + :return: A new GraphFrame with all edge directions reversed. + """ + edge_attr_columns = [c for c in self.edges.columns if c not in [SRC, DST]] + + if edge_attr_columns: + reversed_edges = self.edges.select( + F.col(DST).alias(SRC), + F.col(SRC).alias(DST), + F.struct(*edge_attr_columns).alias(EDGE), ) + reversed_edges = reversed_edges.select(SRC, DST, EDGE) + else: + reversed_edges = self.edges.select(F.col(DST).alias(SRC), F.col(SRC).alias(DST)) + + edge_columns = [F.col(EDGE).getField(c).alias(c) for c in edge_attr_columns] + selected_columns = [F.col(SRC), F.col(DST)] + edge_columns + reversed_edges = reversed_edges.select(*selected_columns) + + return GraphFrame(self.vertices, reversed_edges) + + def aggregate_neighbors( + self, + starting_vertices: Column | str, + accumulator_names: list[str], + accumulator_inits: list[Column | str], + accumulator_updates: list[Column | str], + max_hops: int = 3, + stopping_condition: Column | str | None = None, + target_condition: Column | str | None = None, + required_vertex_attributes: list[str] | None = None, + required_edge_attributes: list[str] | None = None, + edge_filter: Column | str | None = None, + remove_loops: bool = False, + checkpoint_interval: int = 2, + use_local_checkpoints: bool = False, + storage_level: StorageLevel = StorageLevel.MEMORY_AND_DISK_DESER, + ) -> DataFrame: + """Multi-hop neighbor aggregation with customizable accumulators. + + AggregateNeighbors allows you to explore the graph up to a specified number of hops, + accumulating values along paths using customizable accumulator expressions. It supports + both stopping conditions (when to stop exploring) and target conditions (when to + collect a result). + + **Basic Example:** + + >>> from pyspark.sql import functions as F + >>> g = GraphFrame(vertices, edges) + >>> result = g.aggregate_neighbors( + ... starting_vertices=F.col("id") == 1, + ... max_hops=3, + ... accumulator_names=["path_length"], + ... accumulator_inits=[F.lit(0)], + ... accumulator_updates=[F.col("path_length") + 1], + ... target_condition=AggregateNeighbors.dst_attr("id") == 4 + ... ) + + **Using Accumulators:** + + Accumulators track values as the algorithm traverses the graph. Each accumulator has: + - A name (becomes a column in the result) + - An initial value expression (evaluated on starting vertices) + - An update expression (evaluated when traversing each edge) + + In update expressions, you can reference: + - Source vertex attributes via ``src_attr("attrName")`` + - Destination vertex attributes via ``dst_attr("attrName")`` + - Edge attributes via ``edge_attr("attrName")`` + + **Example with Multiple Accumulators:** + + >>> result = g.aggregate_neighbors( + ... starting_vertices=F.col("id") == 1, + ... max_hops=5, + ... accumulator_names=["sum_values", "product_weights"], + ... accumulator_inits=[F.lit(0), F.lit(1.0)], + ... accumulator_updates=[ + ... F.col("sum_values") + AggregateNeighbors.dst_attr("value"), + ... F.col("product_weights") * AggregateNeighbors.edge_attr("weight") + ... ], + ... target_condition=F.col("dst.id") == 10 + ... ) + + **Stopping vs Target Conditions:** + + - ``stopping_condition``: Stops traversal along a path when true. Useful for avoiding + cycles or limiting search depth based on custom criteria. + - ``target_condition``: Marks vertices as results when true. Only accumulators reaching + target vertices are returned. + - If both are provided, only accumulators that reach ``target_condition`` are saved. + - At least one must be provided. + + **Performance Considerations:** + + - Use ``required_vertex_attributes`` and ``required_edge_attributes`` to limit the + columns carried through traversal, reducing memory usage. + - Use ``edge_filter`` to limit traversable edges. + - Set ``remove_loops=True`` to exclude self-loop edges. + - Use ``checkpoint_interval`` to prevent logical plan growth on deep traversals. + - Be cautious with high ``max_hops`` on dense graphs (risk of OOM). + + **Warning:** The result is a persisted DataFrame. Call ``.unpersist()`` when done + to release memory. + + :param starting_vertices: Column expression selecting seed vertices (e.g., ``F.col("id") == 1``) + :param max_hops: Maximum number of hops to explore (positive integer) + :param accumulator_names: List of names for accumulators (become result columns) + :param accumulator_inits: List of initial value expressions for accumulators + :param accumulator_updates: List of update expressions for accumulators + :param stopping_condition: Optional condition to stop traversal along a path + :param target_condition: Optional condition to mark vertices as results + :param required_vertex_attributes: Vertex columns to carry (None = all columns) + :param required_edge_attributes: Edge columns to carry (None = all columns) + :param edge_filter: Optional condition to filter traversable edges + :param remove_loops: If True, exclude self-loop edges (default: False) + :param checkpoint_interval: Checkpoint every N iterations, 0 = disabled (default: 0) + :param use_local_checkpoints: Use local checkpoints (faster but less reliable) + :param storage_level: Storage level for intermediate results + :return: DataFrame with columns: ``id`` (target vertex), ``hop`` (path length), and + one column per accumulator with its final value + """ + return self._impl.aggregate_neighbors( + starting_vertices=starting_vertices, + max_hops=max_hops, + accumulator_names=accumulator_names, + accumulator_inits=accumulator_inits, + accumulator_updates=accumulator_updates, + stopping_condition=stopping_condition, + target_condition=target_condition, + required_vertex_attributes=required_vertex_attributes, + required_edge_attributes=required_edge_attributes, + edge_filter=edge_filter, + remove_loops=remove_loops, + checkpoint_interval=checkpoint_interval, + use_local_checkpoints=use_local_checkpoints, + storage_level=storage_level, + ) + + +@final +class RandomWalkEmbeddings: + def __init__(self, graph: GraphFrame) -> None: + self._graph: GraphFrame = graph + self._params: _RandomWalksEmbeddingsParameters = _RandomWalksEmbeddingsParameters() + + def use_cached_random_walks(self, cached_walks_path: str) -> None: + if cached_walks_path == "": + raise ValueError("cached walks path cannot be empty") + self._params.rw_cached_walks = cached_walks_path + + def set_rw_model( + self, + temporary_prefix: str, + use_edge_direction: bool = False, + max_neighbors_per_vertex: int = 50, + num_walks_per_node: int = 5, + num_batches: int = 5, + walks_per_batch: int = 10, + restart_probability: float = 0.1, + seed: int = 42, + ) -> None: + self._params.rw_model = "rw_with_restart" + self._params.rw_temporary_prefix = temporary_prefix + self._params.use_edge_direction = use_edge_direction + self._params.rw_max_nbrs = max_neighbors_per_vertex + self._params.rw_num_walks_per_node = num_walks_per_node + self._params.rw_num_batches = num_batches + self._params.rw_batch_size = walks_per_batch + self._params.rw_restart_probability = restart_probability + self._params.rw_seed = seed + + def set_hash2vec( + self, + context_size: int = 5, + num_partitions: int = 5, + embeddings_dim: int = 512, + decay_function: str = "gaussian", + gaussian_sigma: float = 1.0, + hashing_seed: int = 42, + sign_seed: int = 18, + l2_norm: bool = True, + save_norm: bool = True, + ) -> None: + if decay_function not in _HASH2VEC_DECAY_FUNCTIONS: + raise ValueError(f"supported decay functions are {str(_HASH2VEC_DECAY_FUNCTIONS)}") + + self._params.sequence_model = "hash2vec" + self._params.hash2vec_context_size = context_size + self._params.hash2vec_num_partitions = num_partitions + self._params.hash2vec_embeddings_dim = embeddings_dim + self._params.hash2vec_decay_function = decay_function + self._params.hash2vec_gaussian_sigma = gaussian_sigma + self._params.hash2vec_hashing_seed = hashing_seed + self._params.hash2vec_sign_seed = sign_seed + self._params.hash2vec_do_l2_norm = l2_norm + self._params.hash2vec_safe_l2 = save_norm + + def set_word2vec( + self, + max_iter: int = 1, + embeddings_dim: int = 100, + window_size: int = 5, + num_partitions: int = 1, + min_count: int = 5, + max_sentence_length: int = 1000, + seed: int = 42, + step_size: float = 0.025, + ) -> None: + self._params.sequence_model = "word2vec" + self._params.word2vec_max_iter = max_iter + self._params.word2vec_embeddings_dim = embeddings_dim + self._params.word2vec_window_size = window_size + self._params.word2vec_num_partitions = num_partitions + self._params.word2vec_min_count = min_count + self._params.word2vec_max_sentence_length = max_sentence_length + self._params.word2vec_seed = seed + self._params.word2vec_step_size = step_size + + def unset_neighbors_aggregation(self) -> None: + self._params.aggregate_neighbors = False + + def set_neighbors_aggregation(self, max_neighbors: int = 50, seed: int = 42) -> None: + self._params.aggregate_neighbors = True + self._params.aggregate_neighbors_max_nbrs = max_neighbors + self._params.aggregate_neighbors_seed = seed + + def set_clean_up_after_run(self, clean_up: bool = True) -> None: + self._params.clean_up_after_run = clean_up + + def run(self) -> DataFrame: + if self._params.rw_temporary_prefix == "": + if self._params.rw_cached_walks == "": + raise ValueError("TMP path or cached walks path should be provided!") + return self._graph._impl.rw_embeddings(self._params) diff --git a/python/pyspark/graphframes/internal/__init__.py b/python/pyspark/graphframes/internal/__init__.py new file mode 100644 index 0000000000000..cce3acad34a49 --- /dev/null +++ b/python/pyspark/graphframes/internal/__init__.py @@ -0,0 +1,16 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# diff --git a/python/pyspark/graphframes/internal/utils.py b/python/pyspark/graphframes/internal/utils.py new file mode 100644 index 0000000000000..076cb71e438dc --- /dev/null +++ b/python/pyspark/graphframes/internal/utils.py @@ -0,0 +1,74 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from __future__ import annotations + +from dataclasses import dataclass + +"""Internal constants.""" +_SEQUENCE_MODELS: tuple[str, ...] = ( + "hash2vec", + "word2vec", +) +_RW_MODELS: tuple[str, ...] = ("rw_with_restart",) +_HASH2VEC_DECAY_FUNCTIONS: tuple[str, ...] = ( + "gaussian", + "constant", +) + + +@dataclass +class _RandomWalksEmbeddingsParameters: + """Internal class represents all possible parameters with defaults.""" + + use_edge_direction: bool = False + rw_model: str = "rw_with_restart" + rw_max_nbrs: int = 50 + rw_num_walks_per_node: int = 5 + rw_batch_size: int = 10 + rw_num_batches: int = 5 + rw_seed: int = 42 + rw_restart_probability: float = 0.1 + rw_temporary_prefix: str = "" + rw_cached_walks: str = "" + sequence_model: str = "hash2vec" + hash2vec_context_size: int = 5 + hash2vec_num_partitions: int = 5 + hash2vec_embeddings_dim: int = 512 + hash2vec_decay_function: str = "gaussian" + hash2vec_gaussian_sigma: float = 1.0 + hash2vec_hashing_seed: int = 42 + hash2vec_sign_seed: int = 18 + hash2vec_do_l2_norm: bool = True + hash2vec_safe_l2: bool = True + word2vec_max_iter: int = 1 + word2vec_embeddings_dim: int = 100 + word2vec_window_size: int = 5 + word2vec_num_partitions: int = 1 + word2vec_min_count: int = 5 + word2vec_max_sentence_length: int = 1000 + word2vec_seed: int = 42 + word2vec_step_size: float = 0.025 + aggregate_neighbors: bool = True + aggregate_neighbors_max_nbrs: int = 50 + aggregate_neighbors_seed: int = 42 + clean_up_after_run: bool = False + + def validate(self) -> None: + if self.sequence_model not in _SEQUENCE_MODELS: + raise ValueError(f"supported seq2vec models are {str(_SEQUENCE_MODELS)}") + if self.rw_model not in _RW_MODELS: + raise ValueError(f"supported RW models are {str(_RW_MODELS)}") diff --git a/python/pyspark/graphframes/lib/__init__.py b/python/pyspark/graphframes/lib/__init__.py new file mode 100644 index 0000000000000..8206984bc07f0 --- /dev/null +++ b/python/pyspark/graphframes/lib/__init__.py @@ -0,0 +1,20 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from .aggregate_messages import AggregateMessages +from .pregel import Pregel + +__all__ = ["AggregateMessages", "Pregel"] diff --git a/python/pyspark/graphframes/lib/aggregate_messages.py b/python/pyspark/graphframes/lib/aggregate_messages.py new file mode 100644 index 0000000000000..932856636ebbc --- /dev/null +++ b/python/pyspark/graphframes/lib/aggregate_messages.py @@ -0,0 +1,60 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + + +from typing import Any + +from pyspark.sql import Column +from pyspark.sql import functions as F + + +class _ClassProperty: + """Custom read-only class property descriptor. + + The underlying method should take the class as the sole argument. + """ + + def __init__(self, f: Any) -> None: + self.f = f + self.__doc__ = f.__doc__ + + def __get__(self, instance: Any, owner: type) -> Any: + return self.f(owner) + + +class AggregateMessages: + """Collection of utilities usable with :meth:`graphframes.GraphFrame.aggregateMessages()`.""" + + @_ClassProperty + def src(cls) -> Column: + """Reference for source column, used for specifying messages.""" + return F.col("src") + + @_ClassProperty + def dst(cls) -> Column: + """Reference for destination column, used for specifying messages.""" + return F.col("dst") + + @_ClassProperty + def edge(cls) -> Column: + """Reference for edge column, used for specifying messages.""" + return F.col("edge") + + @_ClassProperty + def msg(cls) -> Column: + """Reference for message column, used for specifying aggregation function.""" + return F.col("MSG") diff --git a/python/pyspark/graphframes/lib/pregel.py b/python/pyspark/graphframes/lib/pregel.py new file mode 100644 index 0000000000000..1863bb9bb9849 --- /dev/null +++ b/python/pyspark/graphframes/lib/pregel.py @@ -0,0 +1,342 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from typing import TYPE_CHECKING, Any, final + +from typing_extensions import Self + +from pyspark.graphframes.classic.utils import storage_level_to_jvm +from pyspark.ml.wrapper import JavaWrapper +from pyspark.sql import Column, DataFrame, SparkSession +from pyspark.sql.classic.column import _to_seq +from pyspark.sql.functions import col +from pyspark.storagelevel import StorageLevel + +if TYPE_CHECKING: + from pyspark.graphframes.classic.graphframe import GraphFrame + + +@final +class Pregel(JavaWrapper): + """Implements a Pregel-like bulk-synchronous message-passing API based on DataFrame operations. + + See `Malewicz et al., Pregel: a system for large-scale graph processing `_ + for a detailed description of the Pregel algorithm. + + You can construct a Pregel instance using either this constructor or :attr:`graphframes.GraphFrame.pregel`, + then use builder pattern to describe the operations, and then call :func:`run` to start a run. + It returns a DataFrame of vertices from the last iteration. + + When a run starts, it expands the vertices DataFrame using column expressions defined by :func:`withVertexColumn`. + Those additional vertex properties can be changed during Pregel iterations. + In each Pregel iteration, there are three phases: + + - Given each edge triplet, generate messages and specify target vertices to send, described by :func:`sendMsgToDst` and :func:`sendMsgToSrc`. + - Aggregate messages by target vertex IDs, described by :func:`aggMsgs`. + - Update additional vertex properties based on aggregated messages and states from previous iteration, described by :func:`withVertexColumn`. + + Please find what columns you can reference at each phase in the method API docs. + + You can control the number of iterations by :func:`setMaxIter` and check API docs for advanced controls. + + :param graph: a :class:`graphframes.GraphFrame` object holding a graph with vertices and edges stored as DataFrames. + """ + + def __init__(self, graph: "GraphFrame") -> None: + super(Pregel, self).__init__() + + self.graph = graph + self._java_obj: Any = self._new_java_obj( + "org.apache.spark.graphframes.lib.Pregel", graph._jvm_graph + ) + + def setMaxIter(self, value: int) -> "Pregel": + """Sets the max number of iterations (default: 10). + + :param value: the number of Pregel iterations + """ + self._java_obj.setMaxIter(int(value)) + return self + + def setCheckpointInterval(self, value: int) -> "Pregel": + """Sets the number of iterations between two checkpoints (default: 2). + + This is an advanced control to balance query plan optimization and checkpoint data I/O cost. + In most cases, you should keep the default value. + + Checkpoint is disabled if this is set to 0. + """ + self._java_obj.setCheckpointInterval(int(value)) + return self + + def setEarlyStopping(self, value: bool) -> "Pregel": + """Set should Pregel stop earlier in case of no new messages to send or not. + + Early stopping allows to terminate Pregel before reaching maxIter by checking if there are any non-null messages. + While in some cases it may gain significant performance boost, in other cases it can lead to performance degradation, + because checking if the messages DataFrame is empty or not is an action and requires materialization of the Spark Plan + with some additional computations. + + In the case when the user can assume a good value of maxIter, it is recommended to leave this value to the default "false". + In the case when it is hard to estimate the number of iterations required for convergence, + it is recommended to set this value to "false" to avoid iterating over convergence until reaching maxIter. + When this value is "true", maxIter can be set to a bigger value without risks. + """ + self._java_obj.setEarlyStopping(bool(value)) + return self + + def withVertexColumn( + self, colName: str, initialExpr: Column, updateAfterAggMsgsExpr: Column + ) -> "Pregel": + """Defines an additional vertex column at the start of run and how to update it in each iteration. + + You can call it multiple times to add more than one additional vertex columns. + + :param colName: the name of the additional vertex column. + It cannot be an existing vertex column in the graph. + :param initialExpr: the expression to initialize the additional vertex column. + You can reference all original vertex columns in this expression. + :param updateAfterAggMsgsExpr: the expression to update the additional vertex column after messages aggregation. + You can reference all original vertex columns, additional vertex columns, and the + aggregated message column using :func:`msg`. + If the vertex received no messages, the message column would be null. + """ + self._java_obj.withVertexColumn(colName, initialExpr._jc, updateAfterAggMsgsExpr._jc) + return self + + def sendMsgToSrc(self, msgExpr: Column) -> "Pregel": + """Defines a message to send to the source vertex of each edge triplet. + + You can call it multiple times to send more than one messages. + + See method :func:`sendMsgToDst`. + + :param msgExpr: the expression of the message to send to the source vertex given a (src, edge, dst) triplet. + Source/destination vertex properties and edge properties are nested under columns `src`, `dst`, + and `edge`, respectively. + You can reference them using :func:`src`, :func:`dst`, and :func:`edge`. + Null messages are not included in message aggregation. + """ + self._java_obj.sendMsgToSrc(msgExpr._jc) + return self + + def sendMsgToDst(self, msgExpr: Column) -> "Pregel": + """Defines a message to send to the destination vertex of each edge triplet. + + You can call it multiple times to send more than one messages. + + See method :func:`sendMsgToSrc`. + + :param msgExpr: the message expression to send to the destination vertex given a (`src`, `edge`, `dst`) triplet. + Source/destination vertex properties and edge properties are nested under columns `src`, `dst`, + and `edge`, respectively. + You can reference them using :func:`src`, :func:`dst`, and :func:`edge`. + Null messages are not included in message aggregation. + """ + self._java_obj.sendMsgToDst(msgExpr._jc) + return self + + def aggMsgs(self, aggExpr: Column) -> "Pregel": + """Defines how messages are aggregated after grouped by target vertex IDs. + + :param aggExpr: the message aggregation expression, such as `sum(Pregel.msg())`. + You can reference the message column by :func:`msg` and the vertex ID by `col("id")`, + while the latter is usually not used. + """ + self._java_obj.aggMsgs(aggExpr._jc) + return self + + def setStopIfAllNonActiveVertices(self, value: bool) -> Self: + """Set should Pregel stop if all the vertices voted to halt. + + Activity (or vote) is determined based on the activity_col. + See methods :func:`setInitialActiveVertexExpression` and :func:`setUpdateActiveVertexExpression` for details + how to set and update activity_col. + + Be aware that checking of the vote is not free but a Spark Action. In case the + condition is not realistically reachable but set, it will just slow down the algorithm. + + :param value: the boolean value. + """ + self._java_obj.setStopIfAllNonActiveVertices(value) + return self + + def setInitialActiveVertexExpression(self, value: Column) -> Self: + """Sets the initial expression for the active vertex column. + + The active vertex column is used to determine if a vertices voting result on each iteration of Pregel. + This expression is evaluated on the initial vertices DataFrame to set the initial state of the activity column. + + :param value: expression to compute the initial active state of vertices. + You can reference all original vertex columns in this expression. + """ + self._java_obj.setInitialActiveVertexExpression(value._jc) + return self + + def setUpdateActiveVertexExpression(self, value: Column) -> Self: + """Sets the expression to update the active vertex column. + + The active vertex column is used to determine if a vertices voting result on each iteration of Pregel. + This expression is evaluated on the updated vertices DataFrame to set the new state of the activity column. + + :param value: expression to compute the new active state of vertices. + You can reference all original vertex columns and additional vertex columns in this expression. + """ + self._java_obj.setUpdateActiveVertexExpression(value._jc) + return self + + def setSkipMessagesFromNonActiveVertices(self, value: bool) -> Self: + """Set should Pregel skip sending messages from non-active vertices. + + When this option is enabled, messages will not be sent from vertices that are marked as inactive. + This can help optimize performance by avoiding unnecessary message propagation from inactive vertices. + + :param value: boolean value. + """ + self._java_obj.setSkipMessagesFromNonActiveVertices(value) + return self + + def setUseLocalCheckpoints(self, value: bool) -> Self: + """Set should Pregel use local checkpoints. + + Local checkpoints are faster and do not require configuring a persistent storage. + At the same time, local checkpoints are less reliable and may create a big load on local disks of executors. + + :param value: boolean value. + """ + self._java_obj.setUseLocalCheckpoints(value) + return self + + def setIntermediateStorageLevel(self, storage_level: StorageLevel) -> Self: + """Set the intermediate storage level. + On each iteration, Pregel cache results with a requested storage level. + + For very big graphs it is recommended to use DISK_ONLY. + + :param storage_level: storage level to use. + """ + self._java_obj.setIntermediateStorageLevel( + storage_level_to_jvm(storage_level, self.graph._spark) + ) + return self + + def required_src_columns(self, col_name: str, *col_names: str) -> Self: + """Specifies which source vertex columns are required when constructing triplets. + + By default, all source vertex columns are included in triplets, which can create large + intermediate datasets for algorithms with significant state (e.g., cycle detection, + random walks). Use this method to reduce memory usage by specifying only the columns + that are actually needed by the sendMsgToSrc and sendMsgToDst expressions. + + The ID column and the active flag column (if used) are always included automatically. + + :param col_name: the first required source vertex column name + :param col_names: additional required source vertex column names + + See also :func:`required_dst_columns` + """ + self._java_obj.requiredSrcColumns( + col_name, _to_seq(self.graph._spark.sparkContext, col_names) + ) + return self + + def required_dst_columns(self, col_name: str, *col_names: str) -> Self: + """Specifies which destination vertex columns are required when constructing triplets. + + By default, all destination vertex columns are included in triplets, which can create large + intermediate datasets for algorithms with significant state (e.g., cycle detection, + random walks). Use this method to reduce memory usage by specifying only the columns + that are actually needed by the sendMsgToSrc and sendMsgToDst expressions. + + The ID column and the active flag column (if used) are always included automatically. + + :param col_name: the first required destination vertex column name + :param col_names: additional required destination vertex column names + + See also :func:`required_src_columns` + """ + self._java_obj.requiredDstColumns( + col_name, _to_seq(self.graph._spark.sparkContext, col_names) + ) + return self + + def required_edge_columns(self, col_name: str, *col_names: str) -> Self: + """Specifies which edge columns are required when constructing triplets. + + By default, only src and dst columns are included from edges. Use this method to + specify additional edge columns that are needed by the sendMsgToSrc and sendMsgToDst + expressions. + + :param col_name: the first required edge column name + :param col_names: additional required edge column names + + See also :func:`required_src_columns` and :func:`required_dst_columns` + """ + self._java_obj.requiredEdgeColumns( + col_name, _to_seq(self.graph._spark.sparkContext, col_names) + ) + return self + + def run(self) -> DataFrame: + """Runs the defined Pregel algorithm. + + :return: the result vertex DataFrame from the final iteration including both original and additional columns. + """ + spark = SparkSession.getActiveSession() + if spark is None: + raise ValueError("SparkSession is dead or did not started.") + return DataFrame(self._java_obj.run(), spark) + + @staticmethod + def msg() -> Column: + """References the message column in aggregating messages and updating additional vertex columns. + + See :func:`aggMsgs` and :func:`withVertexColumn` + """ + return col("_pregel_msg_") + + @staticmethod + def src(colName: str) -> Column: + """References a source vertex column in generating messages to send. + + See :func:`sendMsgToSrc` and :func:`sendMsgToDst` + + :param colName: the vertex column name. + """ + return col("src." + colName) + + @staticmethod + def dst(colName: str) -> Column: + """ + References a destination vertex column in generating messages to send. + + See :func:`sendMsgToSrc` and :func:`sendMsgToDst` + + :param colName: the vertex column name. + """ + return col("dst." + colName) + + @staticmethod + def edge(colName: str) -> Column: + """ + References an edge column in generating messages to send. + + See :func:`sendMsgToSrc` and :func:`sendMsgToDst` + + :param colName: the edge column name. + """ + return col("edge." + colName) diff --git a/python/pyspark/graphframes/tests/connect/test_all_algorithms.py b/python/pyspark/graphframes/tests/connect/test_all_algorithms.py new file mode 100644 index 0000000000000..dae483040f891 --- /dev/null +++ b/python/pyspark/graphframes/tests/connect/test_all_algorithms.py @@ -0,0 +1,221 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from pyspark.graphframes import GraphFrame +from pyspark.graphframes.graphframe import AggregateNeighbors, RandomWalkEmbeddings +from pyspark.graphframes.lib import AggregateMessages +from pyspark.sql import functions as F +from pyspark.storagelevel import StorageLevel +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class GraphFrameConnectAlgorithmTests(ReusedConnectTestCase): + def setUp(self) -> None: + super().setUp() + vertices = self.spark.createDataFrame( + [(1, "a", 10), (2, "b", 20), (3, "c", 30), (4, "isolated", 40)], + ["id", "name", "age"], + ) + edges = self.spark.createDataFrame( + [ + (1, 2, "friend", 1.0), + (2, 3, "follow", 2.0), + (3, 1, "friend", 3.0), + (2, 1, "friend", 4.0), + ], + ["src", "dst", "relationship", "weight"], + ) + self.graph = GraphFrame(vertices, edges) + + def test_traversal_algorithms(self) -> None: + self.assertEqual(self.graph.find("(a)-[e]->(b)").count(), 4) + self.assertEqual(self.graph.bfs("id = 1", "id = 3").count(), 1) + + paths = self.graph.all_paths( + "id = 1", + "id = 3", + max_path_length=3, + use_local_checkpoints=True, + ) + self.assertEqual(paths.count(), 1) + self.assertEqual(paths.first()["len"], 2) + + cycles = self.graph.detectingCycles(use_local_checkpoints=True) + self.assertGreater(cycles.count(), 0) + + def test_message_aggregation_and_pregel(self) -> None: + messages = self.graph.aggregateMessages( + [F.sum(AggregateMessages.msg).alias("ageSum")], + sendToDst=[AggregateMessages.src["age"]], + ) + self.assertEqual(messages.count(), 3) + messages.unpersist() + + pregel = self.graph.pregel + result = ( + pregel.setMaxIter(2) + .setUseLocalCheckpoints(True) + .withVertexColumn( + "value", + F.lit(0), + F.coalesce(pregel.msg(), F.lit(0)), + ) + .sendMsgToDst(F.lit(1)) + .aggMsgs(F.sum(pregel.msg())) + .run() + ) + self.assertEqual(result.count(), 4) + self.assertIn("value", result.columns) + result.unpersist() + + def test_component_and_community_algorithms(self) -> None: + components = self.graph.connectedComponents( + algorithm="two_phase", + use_local_checkpoints=True, + max_iter=10, + ) + self.assertEqual(components.count(), 4) + self.assertIn("component", components.columns) + components.unpersist() + + labels = self.graph.labelPropagation( + maxIter=2, + algorithm="graphframes", + use_local_checkpoints=True, + ) + self.assertEqual(labels.count(), 4) + self.assertIn("label", labels.columns) + labels.unpersist() + + structure_labels = self.graph.neighborhood_aware_cdlp( + max_iter=2, + use_local_checkpoints=True, + ) + self.assertEqual(structure_labels.count(), 4) + self.assertIn("label", structure_labels.columns) + structure_labels.unpersist() + + def test_ranking_and_path_algorithms(self) -> None: + ranked = self.graph.pageRank(maxIter=2) + self.assertEqual(ranked.vertices.count(), 4) + self.assertEqual(ranked.edges.count(), 4) + self.assertIn("pagerank", ranked.vertices.columns) + self.assertIn("weight", ranked.edges.columns) + + personalized = self.graph.parallelPersonalizedPageRank(sourceIds=[1, 2], maxIter=2) + self.assertEqual(personalized.vertices.count(), 4) + self.assertIn("pageranks", personalized.vertices.columns) + + shortest = self.graph.shortestPaths( + landmarks=[1, 3], + algorithm="graphframes", + use_local_checkpoints=True, + ) + self.assertEqual(shortest.count(), 4) + self.assertIn("distances", shortest.columns) + shortest.unpersist() + + strongly_connected = self.graph.stronglyConnectedComponents(maxIter=5) + self.assertEqual(strongly_connected.count(), 4) + self.assertIn("component", strongly_connected.columns) + + def test_structural_algorithms(self) -> None: + triangles = self.graph.triangleCount(StorageLevel.MEMORY_AND_DISK_DESER) + self.assertEqual(triangles.count(), 4) + self.assertIn("count", triangles.columns) + triangles.unpersist() + + independent_set = self.graph.maximal_independent_set( + use_local_checkpoints=True, + seed=7, + ) + self.assertGreater(independent_set.count(), 0) + self.assertEqual(independent_set.columns, ["id"]) + independent_set.unpersist() + + cores = self.graph.k_core(use_local_checkpoints=True) + self.assertEqual(cores.count(), 4) + self.assertIn("kcore", cores.columns) + cores.unpersist() + + neighborhood = self.graph.hyper_anf(n_hops=2, use_local_checkpoints=True) + self.assertEqual(neighborhood.count(), 3) + self.assertEqual( + neighborhood.columns, + ["id", "hop_0", "hop_1", "hop_2"], + ) + neighborhood.unpersist() + + def test_power_iteration_clustering(self) -> None: + clusters = self.graph.powerIterationClustering( + k=2, + maxIter=5, + weightCol="weight", + ) + self.assertEqual(clusters.count(), 3) + self.assertIn("cluster", clusters.columns) + + def test_svd_plus_plus(self) -> None: + vertices = self.spark.createDataFrame([(1,), (2,), (3,), (4,)], ["id"]) + ratings = self.spark.createDataFrame( + [(1, 3, 4.0), (1, 4, 3.0), (2, 3, 2.0), (2, 4, 5.0)], + ["src", "dst", "weight"], + ) + model, loss = GraphFrame(vertices, ratings).svdPlusPlus(rank=2, maxIter=1) + self.assertEqual(model.count(), 4) + self.assertGreaterEqual(loss, 0.0) + + def test_aggregate_neighbors(self) -> None: + result = self.graph.aggregate_neighbors( + starting_vertices=F.col("id") == 1, + accumulator_names=["path_length"], + accumulator_inits=[F.lit(0)], + accumulator_updates=[F.col("path_length") + 1], + max_hops=3, + target_condition=AggregateNeighbors.dst_attr("id") == 3, + required_vertex_attributes=["id"], + use_local_checkpoints=True, + ) + rows = result.collect() + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0]["path_length"], 2) + result.unpersist() + + def test_random_walk_embeddings(self) -> None: + embeddings = RandomWalkEmbeddings(self.graph) + embeddings.set_rw_model( + "/tmp/spark-graphframes-connect-rw-test", + num_walks_per_node=1, + num_batches=1, + walks_per_batch=1, + ) + embeddings.set_hash2vec( + context_size=2, + num_partitions=1, + embeddings_dim=8, + ) + embeddings.unset_neighbors_aggregation() + embeddings.set_clean_up_after_run() + result = embeddings.run() + self.assertEqual(result.count(), 3) + self.assertIn("embedding", result.columns) + + +if __name__ == "__main__": + from pyspark.testing import main + + main() diff --git a/python/pyspark/sql/connect/proto/graphframes_pb2.py b/python/pyspark/sql/connect/proto/graphframes_pb2.py new file mode 100644 index 0000000000000..469322dbffc5d --- /dev/null +++ b/python/pyspark/sql/connect/proto/graphframes_pb2.py @@ -0,0 +1,110 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: spark/connect/graphframes.proto +# Protobuf Python Version: 6.33.5 +"""Generated protocol buffer code.""" + +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder + +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, 6, 33, 5, "", "spark/connect/graphframes.proto" +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( + b'\n\x1fspark/connect/graphframes.proto\x12\x19spark.connect.graphframes"\x98\x11\n\x0eGraphFramesAPI\x12\x1a\n\x08vertices\x18\x01 \x01(\x0cR\x08vertices\x12\x14\n\x05\x65\x64ges\x18\x02 \x01(\x0cR\x05\x65\x64ges\x12]\n\x12\x61ggregate_messages\x18\x03 \x01(\x0b\x32,.spark.connect.graphframes.AggregateMessagesH\x00R\x11\x61ggregateMessages\x12\x32\n\x03\x62\x66s\x18\x04 \x01(\x0b\x32\x1e.spark.connect.graphframes.BFSH\x00R\x03\x62\x66s\x12\x63\n\x14\x63onnected_components\x18\x05 \x01(\x0b\x32..spark.connect.graphframes.ConnectedComponentsH\x00R\x13\x63onnectedComponents\x12g\n\x16\x64rop_isolated_vertices\x18\x06 \x01(\x0b\x32/.spark.connect.graphframes.DropIsolatedVerticesH\x00R\x14\x64ropIsolatedVertices\x12W\n\x10\x64\x65tecting_cycles\x18\x07 \x01(\x0b\x32*.spark.connect.graphframes.DetectingCyclesH\x00R\x0f\x64\x65tectingCycles\x12K\n\x0c\x66ilter_edges\x18\x08 \x01(\x0b\x32&.spark.connect.graphframes.FilterEdgesH\x00R\x0b\x66ilterEdges\x12T\n\x0f\x66ilter_vertices\x18\t \x01(\x0b\x32).spark.connect.graphframes.FilterVerticesH\x00R\x0e\x66ilterVertices\x12\x35\n\x04\x66ind\x18\n \x01(\x0b\x32\x1f.spark.connect.graphframes.FindH\x00R\x04\x66ind\x12Z\n\x11label_propagation\x18\x0b \x01(\x0b\x32+.spark.connect.graphframes.LabelPropagationH\x00R\x10labelPropagation\x12\x42\n\tpage_rank\x18\x0c \x01(\x0b\x32#.spark.connect.graphframes.PageRankH\x00R\x08pageRank\x12\x80\x01\n\x1fparallel_personalized_page_rank\x18\r \x01(\x0b\x32\x37.spark.connect.graphframes.ParallelPersonalizedPageRankH\x00R\x1cparallelPersonalizedPageRank\x12s\n\x1apower_iteration_clustering\x18\x0e \x01(\x0b\x32\x33.spark.connect.graphframes.PowerIterationClusteringH\x00R\x18powerIterationClustering\x12;\n\x06pregel\x18\x0f \x01(\x0b\x32!.spark.connect.graphframes.PregelH\x00R\x06pregel\x12Q\n\x0eshortest_paths\x18\x10 \x01(\x0b\x32(.spark.connect.graphframes.ShortestPathsH\x00R\rshortestPaths\x12|\n\x1dstrongly_connected_components\x18\x11 \x01(\x0b\x32\x36.spark.connect.graphframes.StronglyConnectedComponentsH\x00R\x1bstronglyConnectedComponents\x12L\n\rsvd_plus_plus\x18\x12 \x01(\x0b\x32&.spark.connect.graphframes.SVDPlusPlusH\x00R\x0bsvdPlusPlus\x12Q\n\x0etriangle_count\x18\x13 \x01(\x0b\x32(.spark.connect.graphframes.TriangleCountH\x00R\rtriangleCount\x12\x41\n\x08triplets\x18\x14 \x01(\x0b\x32#.spark.connect.graphframes.TripletsH\x00R\x08triplets\x12\x38\n\x05kcore\x18\x15 \x01(\x0b\x32 .spark.connect.graphframes.KCoreH\x00R\x05kcore\x12\x44\n\x03mis\x18\x16 \x01(\x0b\x32\x30.spark.connect.graphframes.MaximalIndependentSetH\x00R\x03mis\x12V\n\rrw_embeddings\x18\x17 \x01(\x0b\x32/.spark.connect.graphframes.RandomWalkEmbeddingsH\x00R\x0crwEmbeddings\x12`\n\x13\x61ggregate_neighbors\x18\x18 \x01(\x0b\x32-.spark.connect.graphframes.AggregateNeighborsH\x00R\x12\x61ggregateNeighbors\x12j\n\x17neighborhood_aware_cdlp\x18\x19 \x01(\x0b\x32\x30.spark.connect.graphframes.NeighborhoodAwareCDLPH\x00R\x15neighborhoodAwareCdlp\x12\x42\n\tall_paths\x18\x1a \x01(\x0b\x32#.spark.connect.graphframes.AllPathsH\x00R\x08\x61llPaths\x12\x42\n\thyper_anf\x18\x1b \x01(\x0b\x32#.spark.connect.graphframes.HyperANFH\x00R\x08hyperAnfB\x08\n\x06method"\xd7\x02\n\x0cStorageLevel\x12\x1d\n\tdisk_only\x18\x01 \x01(\x08H\x00R\x08\x64iskOnly\x12 \n\x0b\x64isk_only_2\x18\x02 \x01(\x08H\x00R\tdiskOnly2\x12 \n\x0b\x64isk_only_3\x18\x03 \x01(\x08H\x00R\tdiskOnly3\x12(\n\x0fmemory_and_disk\x18\x04 \x01(\x08H\x00R\rmemoryAndDisk\x12+\n\x11memory_and_disk_2\x18\x05 \x01(\x08H\x00R\x0ememoryAndDisk2\x12\x33\n\x15memory_and_disk_deser\x18\x06 \x01(\x08H\x00R\x12memoryAndDiskDeser\x12!\n\x0bmemory_only\x18\x07 \x01(\x08H\x00R\nmemoryOnly\x12$\n\rmemory_only_2\x18\x08 \x01(\x08H\x00R\x0bmemoryOnly2B\x0f\n\rstorage_level"M\n\x12\x43olumnOrExpression\x12\x12\n\x03\x63ol\x18\x01 \x01(\x0cH\x00R\x03\x63ol\x12\x14\n\x04\x65xpr\x18\x02 \x01(\tH\x00R\x04\x65xprB\r\n\x0b\x63ol_or_expr"P\n\x0eStringOrLongID\x12\x19\n\x07long_id\x18\x01 \x01(\x03H\x00R\x06longId\x12\x1d\n\tstring_id\x18\x02 \x01(\tH\x00R\x08stringIdB\x04\n\x02id"\xde\x02\n\x11\x41ggregateMessages\x12\x46\n\x07\x61gg_col\x18\x01 \x03(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\x06\x61ggCol\x12M\n\x0bsend_to_src\x18\x02 \x03(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\tsendToSrc\x12M\n\x0bsend_to_dst\x18\x03 \x03(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\tsendToDst\x12Q\n\rstorage_level\x18\x04 \x01(\x0b\x32\'.spark.connect.graphframes.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\x91\x02\n\x03\x42\x46S\x12J\n\tfrom_expr\x18\x01 \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\x08\x66romExpr\x12\x46\n\x07to_expr\x18\x02 \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\x06toExpr\x12N\n\x0b\x65\x64ge_filter\x18\x03 \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\nedgeFilter\x12&\n\x0fmax_path_length\x18\x04 \x01(\x05R\rmaxPathLength"\x81\x04\n\x08\x41llPaths\x12J\n\tfrom_expr\x18\x01 \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\x08\x66romExpr\x12\x46\n\x07to_expr\x18\x02 \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\x06toExpr\x12N\n\x0b\x65\x64ge_filter\x18\x03 \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\nedgeFilter\x12&\n\x0fmax_path_length\x18\x04 \x01(\x05R\rmaxPathLength\x12\x1f\n\x0bis_directed\x18\x05 \x01(\x08R\nisDirected\x12/\n\x13\x63heckpoint_interval\x18\x06 \x01(\x05R\x12\x63heckpointInterval\x12\x32\n\x15use_local_checkpoints\x18\x07 \x01(\x08R\x13useLocalCheckpoints\x12Q\n\rstorage_level\x18\x08 \x01(\x0b\x32\'.spark.connect.graphframes.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\x99\x03\n\x08HyperANF\x12\x15\n\x06n_hops\x18\x01 \x01(\x05R\x05nHops\x12$\n\x0elg_nom_entries\x18\x02 \x01(\x05R\x0clgNomEntries\x12j\n\x17\x65\x64ges_filter_expression\x18\x03 \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionH\x00R\x15\x65\x64gesFilterExpression\x88\x01\x01\x12/\n\x13\x63heckpoint_interval\x18\x04 \x01(\x05R\x12\x63heckpointInterval\x12\x32\n\x15use_local_checkpoints\x18\x05 \x01(\x08R\x13useLocalCheckpoints\x12Q\n\rstorage_level\x18\x06 \x01(\x0b\x32\'.spark.connect.graphframes.StorageLevelH\x01R\x0cstorageLevel\x88\x01\x01\x42\x1a\n\x18_edges_filter_expressionB\x10\n\x0e_storage_level"\x82\x03\n\x13\x43onnectedComponents\x12\x1c\n\talgorithm\x18\x01 \x01(\tR\talgorithm\x12/\n\x13\x63heckpoint_interval\x18\x02 \x01(\x05R\x12\x63heckpointInterval\x12/\n\x13\x62roadcast_threshold\x18\x03 \x01(\x05R\x12\x62roadcastThreshold\x12\x37\n\x18use_labels_as_components\x18\x04 \x01(\x08R\x15useLabelsAsComponents\x12\x32\n\x15use_local_checkpoints\x18\x05 \x01(\x08R\x13useLocalCheckpoints\x12\x19\n\x08max_iter\x18\x06 \x01(\x05R\x07maxIter\x12Q\n\rstorage_level\x18\x07 \x01(\x0b\x32\'.spark.connect.graphframes.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\xdb\x01\n\x0f\x44\x65tectingCycles\x12\x32\n\x15use_local_checkpoints\x18\x01 \x01(\x08R\x13useLocalCheckpoints\x12/\n\x13\x63heckpoint_interval\x18\x02 \x01(\x05R\x12\x63heckpointInterval\x12Q\n\rstorage_level\x18\x03 \x01(\x0b\x32\'.spark.connect.graphframes.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\x16\n\x14\x44ropIsolatedVertices"Z\n\x0b\x46ilterEdges\x12K\n\tcondition\x18\x01 \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\tcondition"]\n\x0e\x46ilterVertices\x12K\n\tcondition\x18\x02 \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\tcondition" \n\x04\x46ind\x12\x18\n\x07pattern\x18\x01 \x01(\tR\x07pattern"\x95\x02\n\x10LabelPropagation\x12\x1c\n\talgorithm\x18\x01 \x01(\tR\talgorithm\x12\x19\n\x08max_iter\x18\x02 \x01(\x05R\x07maxIter\x12\x32\n\x15use_local_checkpoints\x18\x03 \x01(\x08R\x13useLocalCheckpoints\x12/\n\x13\x63heckpoint_interval\x18\x04 \x01(\x05R\x12\x63heckpointInterval\x12Q\n\rstorage_level\x18\x05 \x01(\x0b\x32\'.spark.connect.graphframes.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\x84\x04\n\x15NeighborhoodAwareCDLP\x12\x19\n\x08max_iter\x18\x01 \x01(\x05R\x07maxIter\x12.\n\x13ignore_direct_links\x18\x02 \x01(\x08R\x11ignoreDirectLinks\x12H\n structural_similarity_multiplier\x18\x03 \x01(\x01R\x1estructuralSimilarityMultiplier\x12\x32\n\x15use_local_checkpoints\x18\x04 \x01(\x08R\x13useLocalCheckpoints\x12/\n\x13\x63heckpoint_interval\x18\x05 \x01(\x05R\x12\x63heckpointInterval\x12Q\n\rstorage_level\x18\x06 \x01(\x0b\x32\'.spark.connect.graphframes.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x12\x1f\n\x0bis_directed\x18\x07 \x01(\x08R\nisDirected\x12$\n\x0elg_nom_entries\x18\x08 \x01(\x05R\x0clgNomEntries\x12/\n\x11initial_label_col\x18\t \x01(\tH\x01R\x0finitialLabelCol\x88\x01\x01\x42\x10\n\x0e_storage_levelB\x14\n\x12_initial_label_col"\xde\x01\n\x08PageRank\x12+\n\x11reset_probability\x18\x01 \x01(\x01R\x10resetProbability\x12K\n\tsource_id\x18\x02 \x01(\x0b\x32).spark.connect.graphframes.StringOrLongIDH\x00R\x08sourceId\x88\x01\x01\x12\x1e\n\x08max_iter\x18\x03 \x01(\x05H\x01R\x07maxIter\x88\x01\x01\x12\x15\n\x03tol\x18\x04 \x01(\x01H\x02R\x03tol\x88\x01\x01\x42\x0c\n\n_source_idB\x0b\n\t_max_iterB\x06\n\x04_tol"\xb0\x01\n\x1cParallelPersonalizedPageRank\x12+\n\x11reset_probability\x18\x01 \x01(\x01R\x10resetProbability\x12H\n\nsource_ids\x18\x02 \x03(\x0b\x32).spark.connect.graphframes.StringOrLongIDR\tsourceIds\x12\x19\n\x08max_iter\x18\x03 \x01(\x05R\x07maxIter"v\n\x18PowerIterationClustering\x12\x0c\n\x01k\x18\x01 \x01(\x05R\x01k\x12\x19\n\x08max_iter\x18\x02 \x01(\x05R\x07maxIter\x12"\n\nweight_col\x18\x03 \x01(\tH\x00R\tweightCol\x88\x01\x01\x42\r\n\x0b_weight_col"\xb9\x0b\n\x06Pregel\x12H\n\x08\x61gg_msgs\x18\x01 \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\x07\x61ggMsgs\x12T\n\x0fsend_msg_to_dst\x18\x02 \x03(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\x0csendMsgToDst\x12T\n\x0fsend_msg_to_src\x18\x03 \x03(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\x0csendMsgToSrc\x12/\n\x13\x63heckpoint_interval\x18\x04 \x01(\x05R\x12\x63heckpointInterval\x12\x19\n\x08max_iter\x18\x05 \x01(\x05R\x07maxIter\x12.\n\x13\x61\x64\x64itional_col_name\x18\x06 \x01(\tR\x11\x61\x64\x64itionalColName\x12\x63\n\x16\x61\x64\x64itional_col_initial\x18\x07 \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\x14\x61\x64\x64itionalColInitial\x12[\n\x12\x61\x64\x64itional_col_upd\x18\x08 \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\x10\x61\x64\x64itionalColUpd\x12*\n\x0e\x65\x61rly_stopping\x18\t \x01(\x08H\x00R\rearlyStopping\x88\x01\x01\x12\x32\n\x15use_local_checkpoints\x18\n \x01(\x08R\x13useLocalCheckpoints\x12Q\n\rstorage_level\x18\x0b \x01(\x0b\x32\'.spark.connect.graphframes.StorageLevelH\x01R\x0cstorageLevel\x88\x01\x01\x12\x37\n\x16stop_if_all_non_active\x18\x0c \x01(\x08H\x02R\x12stopIfAllNonActive\x88\x01\x01\x12\x62\n\x13initial_active_expr\x18\r \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionH\x03R\x11initialActiveExpr\x88\x01\x01\x12`\n\x12update_active_expr\x18\x0e \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionH\x04R\x10updateActiveExpr\x88\x01\x01\x12\x45\n\x1dskip_messages_from_non_active\x18\x0f \x01(\x08H\x05R\x19skipMessagesFromNonActive\x88\x01\x01\x12\x35\n\x14required_src_columns\x18\x10 \x01(\tH\x06R\x12requiredSrcColumns\x88\x01\x01\x12\x35\n\x14required_dst_columns\x18\x11 \x01(\tH\x07R\x12requiredDstColumns\x88\x01\x01\x12\x37\n\x15required_edge_columns\x18\x12 \x01(\tH\x08R\x13requiredEdgeColumns\x88\x01\x01\x42\x11\n\x0f_early_stoppingB\x10\n\x0e_storage_levelB\x19\n\x17_stop_if_all_non_activeB\x16\n\x14_initial_active_exprB\x15\n\x13_update_active_exprB \n\x1e_skip_messages_from_non_activeB\x17\n\x15_required_src_columnsB\x17\n\x15_required_dst_columnsB\x18\n\x16_required_edge_columns"\xf6\x02\n\rShortestPaths\x12G\n\tlandmarks\x18\x01 \x03(\x0b\x32).spark.connect.graphframes.StringOrLongIDR\tlandmarks\x12\x1c\n\talgorithm\x18\x02 \x01(\tR\talgorithm\x12\x32\n\x15use_local_checkpoints\x18\x03 \x01(\x08R\x13useLocalCheckpoints\x12/\n\x13\x63heckpoint_interval\x18\x04 \x01(\x05R\x12\x63heckpointInterval\x12Q\n\rstorage_level\x18\x05 \x01(\x0b\x32\'.spark.connect.graphframes.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x12$\n\x0bis_directed\x18\x06 \x01(\x08H\x01R\nisDirected\x88\x01\x01\x42\x10\n\x0e_storage_levelB\x0e\n\x0c_is_directed"8\n\x1bStronglyConnectedComponents\x12\x19\n\x08max_iter\x18\x01 \x01(\x05R\x07maxIter"\xd6\x01\n\x0bSVDPlusPlus\x12\x12\n\x04rank\x18\x01 \x01(\x05R\x04rank\x12\x19\n\x08max_iter\x18\x02 \x01(\x05R\x07maxIter\x12\x1b\n\tmin_value\x18\x03 \x01(\x01R\x08minValue\x12\x1b\n\tmax_value\x18\x04 \x01(\x01R\x08maxValue\x12\x16\n\x06gamma1\x18\x05 \x01(\x01R\x06gamma1\x12\x16\n\x06gamma2\x18\x06 \x01(\x01R\x06gamma2\x12\x16\n\x06gamma6\x18\x07 \x01(\x01R\x06gamma6\x12\x16\n\x06gamma7\x18\x08 \x01(\x01R\x06gamma7"\xe3\x01\n\rTriangleCount\x12Q\n\rstorage_level\x18\x01 \x01(\x0b\x32\'.spark.connect.graphframes.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x12!\n\talgorithm\x18\x02 \x01(\tH\x01R\talgorithm\x88\x01\x01\x12)\n\x0elg_nom_entries\x18\x03 \x01(\x05H\x02R\x0clgNomEntries\x88\x01\x01\x42\x10\n\x0e_storage_levelB\x0c\n\n_algorithmB\x11\n\x0f_lg_nom_entries"\n\n\x08Triplets"\xf5\x01\n\x15MaximalIndependentSet\x12/\n\x13\x63heckpoint_interval\x18\x01 \x01(\x05R\x12\x63heckpointInterval\x12Q\n\rstorage_level\x18\x02 \x01(\x0b\x32\'.spark.connect.graphframes.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x12\x32\n\x15use_local_checkpoints\x18\x03 \x01(\x08R\x13useLocalCheckpoints\x12\x12\n\x04seed\x18\x04 \x01(\x03R\x04seedB\x10\n\x0e_storage_level"\xd1\x01\n\x05KCore\x12\x32\n\x15use_local_checkpoints\x18\x01 \x01(\x08R\x13useLocalCheckpoints\x12/\n\x13\x63heckpoint_interval\x18\x02 \x01(\x05R\x12\x63heckpointInterval\x12Q\n\rstorage_level\x18\x03 \x01(\x0b\x32\'.spark.connect.graphframes.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\xac\x08\n\x12\x41ggregateNeighbors\x12Z\n\x11starting_vertices\x18\x01 \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\x10startingVertices\x12\x19\n\x08max_hops\x18\x02 \x01(\x05R\x07maxHops\x12+\n\x11\x61\x63\x63umulator_names\x18\x03 \x03(\tR\x10\x61\x63\x63umulatorNames\x12Z\n\x11\x61\x63\x63umulator_inits\x18\x04 \x03(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\x10\x61\x63\x63umulatorInits\x12^\n\x13\x61\x63\x63umulator_updates\x18\x05 \x03(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionR\x12\x61\x63\x63umulatorUpdates\x12\x61\n\x12stopping_condition\x18\x06 \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionH\x00R\x11stoppingCondition\x88\x01\x01\x12]\n\x10target_condition\x18\x07 \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionH\x01R\x0ftargetCondition\x88\x01\x01\x12<\n\x1arequired_vertex_attributes\x18\x08 \x03(\tR\x18requiredVertexAttributes\x12\x38\n\x18required_edge_attributes\x18\t \x03(\tR\x16requiredEdgeAttributes\x12S\n\x0b\x65\x64ge_filter\x18\n \x01(\x0b\x32-.spark.connect.graphframes.ColumnOrExpressionH\x02R\nedgeFilter\x88\x01\x01\x12!\n\x0cremove_loops\x18\x0b \x01(\x08R\x0bremoveLoops\x12/\n\x13\x63heckpoint_interval\x18\x0c \x01(\x05R\x12\x63heckpointInterval\x12\x32\n\x15use_local_checkpoints\x18\r \x01(\x08R\x13useLocalCheckpoints\x12Q\n\rstorage_level\x18\x0e \x01(\x0b\x32\'.spark.connect.graphframes.StorageLevelH\x03R\x0cstorageLevel\x88\x01\x01\x42\x15\n\x13_stopping_conditionB\x13\n\x11_target_conditionB\x0e\n\x0c_edge_filterB\x10\n\x0e_storage_level"\x81\x0c\n\x14RandomWalkEmbeddings\x12,\n\x12use_edge_direction\x18\x01 \x01(\x08R\x10useEdgeDirection\x12\x19\n\x08rw_model\x18\x02 \x01(\tR\x07rwModel\x12\x1e\n\x0brw_max_nbrs\x18\x03 \x01(\x05R\trwMaxNbrs\x12\x30\n\x15rw_num_walks_per_node\x18\x04 \x01(\x05R\x11rwNumWalksPerNode\x12"\n\rrw_batch_size\x18\x05 \x01(\x05R\x0brwBatchSize\x12$\n\x0erw_num_batches\x18\x06 \x01(\x05R\x0crwNumBatches\x12\x17\n\x07rw_seed\x18\x07 \x01(\x03R\x06rwSeed\x12\x34\n\x16rw_restart_probability\x18\x08 \x01(\x01R\x14rwRestartProbability\x12.\n\x13rw_temporary_prefix\x18\t \x01(\tR\x11rwTemporaryPrefix\x12&\n\x0frw_cached_walks\x18\n \x01(\tR\rrwCachedWalks\x12%\n\x0esequence_model\x18\x0b \x01(\tR\rsequenceModel\x12\x32\n\x15hash2vec_context_size\x18\x0c \x01(\x05R\x13hash2vecContextSize\x12\x36\n\x17hash2vec_num_partitions\x18\r \x01(\x05R\x15hash2vecNumPartitions\x12\x36\n\x17hash2vec_embeddings_dim\x18\x0e \x01(\x05R\x15hash2vecEmbeddingsDim\x12\x36\n\x17hash2vec_decay_function\x18\x0f \x01(\tR\x15hash2vecDecayFunction\x12\x36\n\x17hash2vec_gaussian_sigma\x18\x10 \x01(\x01R\x15hash2vecGaussianSigma\x12\x32\n\x15hash2vec_hashing_seed\x18\x11 \x01(\x05R\x13hash2vecHashingSeed\x12,\n\x12hash2vec_sign_seed\x18\x12 \x01(\x05R\x10hash2vecSignSeed\x12-\n\x13hash2vec_do_l2_norm\x18\x13 \x01(\x08R\x10hash2vecDoL2Norm\x12(\n\x10hash2vec_safe_l2\x18\x14 \x01(\x08R\x0ehash2vecSafeL2\x12*\n\x11word2vec_max_iter\x18\x15 \x01(\x05R\x0fword2vecMaxIter\x12\x36\n\x17word2vec_embeddings_dim\x18\x16 \x01(\x05R\x15word2vecEmbeddingsDim\x12\x30\n\x14word2vec_window_size\x18\x17 \x01(\x05R\x12word2vecWindowSize\x12\x36\n\x17word2vec_num_partitions\x18\x18 \x01(\x05R\x15word2vecNumPartitions\x12,\n\x12word2vec_min_count\x18\x19 \x01(\x05R\x10word2vecMinCount\x12?\n\x1cword2vec_max_sentence_length\x18\x1a \x01(\x05R\x19word2vecMaxSentenceLength\x12#\n\rword2vec_seed\x18\x1b \x01(\x03R\x0cword2vecSeed\x12,\n\x12word2vec_step_size\x18\x1c \x01(\x01R\x10word2vecStepSize\x12/\n\x13\x61ggregate_neighbors\x18\x1d \x01(\x08R\x12\x61ggregateNeighbors\x12?\n\x1c\x61ggregate_neighbors_max_nbrs\x18\x1e \x01(\x05R\x19\x61ggregateNeighborsMaxNbrs\x12\x38\n\x18\x61ggregate_neighbors_seed\x18\x1f \x01(\x03R\x16\x61ggregateNeighborsSeed\x12+\n\x12\x63lean_up_after_run\x18 \x01(\x08R\x0f\x63leanUpAfterRunB3\n*org.apache.spark.connect.proto.graphframesH\x01P\x01\xa0\x01\x01\x62\x06proto3' +) + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages( + DESCRIPTOR, "pyspark.sql.connect.proto.graphframes_pb2", _globals +) +if not _descriptor._USE_C_DESCRIPTORS: + _globals["DESCRIPTOR"]._loaded_options = None + _globals[ + "DESCRIPTOR" + ]._serialized_options = b"\n*org.apache.spark.connect.proto.graphframesH\001P\001\240\001\001" + _globals["_GRAPHFRAMESAPI"]._serialized_start = 63 + _globals["_GRAPHFRAMESAPI"]._serialized_end = 2263 + _globals["_STORAGELEVEL"]._serialized_start = 2266 + _globals["_STORAGELEVEL"]._serialized_end = 2609 + _globals["_COLUMNOREXPRESSION"]._serialized_start = 2611 + _globals["_COLUMNOREXPRESSION"]._serialized_end = 2688 + _globals["_STRINGORLONGID"]._serialized_start = 2690 + _globals["_STRINGORLONGID"]._serialized_end = 2770 + _globals["_AGGREGATEMESSAGES"]._serialized_start = 2773 + _globals["_AGGREGATEMESSAGES"]._serialized_end = 3123 + _globals["_BFS"]._serialized_start = 3126 + _globals["_BFS"]._serialized_end = 3399 + _globals["_ALLPATHS"]._serialized_start = 3402 + _globals["_ALLPATHS"]._serialized_end = 3915 + _globals["_HYPERANF"]._serialized_start = 3918 + _globals["_HYPERANF"]._serialized_end = 4327 + _globals["_CONNECTEDCOMPONENTS"]._serialized_start = 4330 + _globals["_CONNECTEDCOMPONENTS"]._serialized_end = 4716 + _globals["_DETECTINGCYCLES"]._serialized_start = 4719 + _globals["_DETECTINGCYCLES"]._serialized_end = 4938 + _globals["_DROPISOLATEDVERTICES"]._serialized_start = 4940 + _globals["_DROPISOLATEDVERTICES"]._serialized_end = 4962 + _globals["_FILTEREDGES"]._serialized_start = 4964 + _globals["_FILTEREDGES"]._serialized_end = 5054 + _globals["_FILTERVERTICES"]._serialized_start = 5056 + _globals["_FILTERVERTICES"]._serialized_end = 5149 + _globals["_FIND"]._serialized_start = 5151 + _globals["_FIND"]._serialized_end = 5183 + _globals["_LABELPROPAGATION"]._serialized_start = 5186 + _globals["_LABELPROPAGATION"]._serialized_end = 5463 + _globals["_NEIGHBORHOODAWARECDLP"]._serialized_start = 5466 + _globals["_NEIGHBORHOODAWARECDLP"]._serialized_end = 5982 + _globals["_PAGERANK"]._serialized_start = 5985 + _globals["_PAGERANK"]._serialized_end = 6207 + _globals["_PARALLELPERSONALIZEDPAGERANK"]._serialized_start = 6210 + _globals["_PARALLELPERSONALIZEDPAGERANK"]._serialized_end = 6386 + _globals["_POWERITERATIONCLUSTERING"]._serialized_start = 6388 + _globals["_POWERITERATIONCLUSTERING"]._serialized_end = 6506 + _globals["_PREGEL"]._serialized_start = 6509 + _globals["_PREGEL"]._serialized_end = 7974 + _globals["_SHORTESTPATHS"]._serialized_start = 7977 + _globals["_SHORTESTPATHS"]._serialized_end = 8351 + _globals["_STRONGLYCONNECTEDCOMPONENTS"]._serialized_start = 8353 + _globals["_STRONGLYCONNECTEDCOMPONENTS"]._serialized_end = 8409 + _globals["_SVDPLUSPLUS"]._serialized_start = 8412 + _globals["_SVDPLUSPLUS"]._serialized_end = 8626 + _globals["_TRIANGLECOUNT"]._serialized_start = 8629 + _globals["_TRIANGLECOUNT"]._serialized_end = 8856 + _globals["_TRIPLETS"]._serialized_start = 8858 + _globals["_TRIPLETS"]._serialized_end = 8868 + _globals["_MAXIMALINDEPENDENTSET"]._serialized_start = 8871 + _globals["_MAXIMALINDEPENDENTSET"]._serialized_end = 9116 + _globals["_KCORE"]._serialized_start = 9119 + _globals["_KCORE"]._serialized_end = 9328 + _globals["_AGGREGATENEIGHBORS"]._serialized_start = 9331 + _globals["_AGGREGATENEIGHBORS"]._serialized_end = 10399 + _globals["_RANDOMWALKEMBEDDINGS"]._serialized_start = 10402 + _globals["_RANDOMWALKEMBEDDINGS"]._serialized_end = 11939 +# @@protoc_insertion_point(module_scope) diff --git a/python/pyspark/sql/connect/proto/graphframes_pb2.pyi b/python/pyspark/sql/connect/proto/graphframes_pb2.pyi new file mode 100644 index 0000000000000..57b5f78487ecf --- /dev/null +++ b/python/pyspark/sql/connect/proto/graphframes_pb2.pyi @@ -0,0 +1,2090 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +""" +@generated by mypy-protobuf. Do not edit manually! +isort:skip_file + +Licensed to the Apache Software Foundation (ASF) under one or more +contributor license agreements. See the NOTICE file distributed with +this work for additional information regarding copyright ownership. +The ASF licenses this file to You under the Apache License, Version 2.0 +(the "License"); you may not use this file except in compliance with +the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import builtins +import collections.abc +import google.protobuf.descriptor +import google.protobuf.internal.containers +import google.protobuf.message +import sys +import typing + +if sys.version_info >= (3, 8): + import typing as typing_extensions +else: + import typing_extensions + +DESCRIPTOR: google.protobuf.descriptor.FileDescriptor + +class GraphFramesAPI(google.protobuf.message.Message): + """GraphFramesAPI represents the core message type for GraphFrames operations + containing graph data and the specific graph algorithm to be executed + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + VERTICES_FIELD_NUMBER: builtins.int + EDGES_FIELD_NUMBER: builtins.int + AGGREGATE_MESSAGES_FIELD_NUMBER: builtins.int + BFS_FIELD_NUMBER: builtins.int + CONNECTED_COMPONENTS_FIELD_NUMBER: builtins.int + DROP_ISOLATED_VERTICES_FIELD_NUMBER: builtins.int + DETECTING_CYCLES_FIELD_NUMBER: builtins.int + FILTER_EDGES_FIELD_NUMBER: builtins.int + FILTER_VERTICES_FIELD_NUMBER: builtins.int + FIND_FIELD_NUMBER: builtins.int + LABEL_PROPAGATION_FIELD_NUMBER: builtins.int + PAGE_RANK_FIELD_NUMBER: builtins.int + PARALLEL_PERSONALIZED_PAGE_RANK_FIELD_NUMBER: builtins.int + POWER_ITERATION_CLUSTERING_FIELD_NUMBER: builtins.int + PREGEL_FIELD_NUMBER: builtins.int + SHORTEST_PATHS_FIELD_NUMBER: builtins.int + STRONGLY_CONNECTED_COMPONENTS_FIELD_NUMBER: builtins.int + SVD_PLUS_PLUS_FIELD_NUMBER: builtins.int + TRIANGLE_COUNT_FIELD_NUMBER: builtins.int + TRIPLETS_FIELD_NUMBER: builtins.int + KCORE_FIELD_NUMBER: builtins.int + MIS_FIELD_NUMBER: builtins.int + RW_EMBEDDINGS_FIELD_NUMBER: builtins.int + AGGREGATE_NEIGHBORS_FIELD_NUMBER: builtins.int + NEIGHBORHOOD_AWARE_CDLP_FIELD_NUMBER: builtins.int + ALL_PATHS_FIELD_NUMBER: builtins.int + HYPER_ANF_FIELD_NUMBER: builtins.int + vertices: builtins.bytes + """Serialized vertex DataFrame containing node information""" + edges: builtins.bytes + """Serialized edge DataFrame containing relationship information""" + @property + def aggregate_messages(self) -> global___AggregateMessages: ... + @property + def bfs(self) -> global___BFS: ... + @property + def connected_components(self) -> global___ConnectedComponents: ... + @property + def drop_isolated_vertices(self) -> global___DropIsolatedVertices: ... + @property + def detecting_cycles(self) -> global___DetectingCycles: ... + @property + def filter_edges(self) -> global___FilterEdges: ... + @property + def filter_vertices(self) -> global___FilterVertices: ... + @property + def find(self) -> global___Find: ... + @property + def label_propagation(self) -> global___LabelPropagation: ... + @property + def page_rank(self) -> global___PageRank: ... + @property + def parallel_personalized_page_rank(self) -> global___ParallelPersonalizedPageRank: ... + @property + def power_iteration_clustering(self) -> global___PowerIterationClustering: ... + @property + def pregel(self) -> global___Pregel: ... + @property + def shortest_paths(self) -> global___ShortestPaths: ... + @property + def strongly_connected_components(self) -> global___StronglyConnectedComponents: ... + @property + def svd_plus_plus(self) -> global___SVDPlusPlus: ... + @property + def triangle_count(self) -> global___TriangleCount: ... + @property + def triplets(self) -> global___Triplets: ... + @property + def kcore(self) -> global___KCore: ... + @property + def mis(self) -> global___MaximalIndependentSet: ... + @property + def rw_embeddings(self) -> global___RandomWalkEmbeddings: ... + @property + def aggregate_neighbors(self) -> global___AggregateNeighbors: ... + @property + def neighborhood_aware_cdlp(self) -> global___NeighborhoodAwareCDLP: ... + @property + def all_paths(self) -> global___AllPaths: ... + @property + def hyper_anf(self) -> global___HyperANF: ... + def __init__( + self, + *, + vertices: builtins.bytes = ..., + edges: builtins.bytes = ..., + aggregate_messages: global___AggregateMessages | None = ..., + bfs: global___BFS | None = ..., + connected_components: global___ConnectedComponents | None = ..., + drop_isolated_vertices: global___DropIsolatedVertices | None = ..., + detecting_cycles: global___DetectingCycles | None = ..., + filter_edges: global___FilterEdges | None = ..., + filter_vertices: global___FilterVertices | None = ..., + find: global___Find | None = ..., + label_propagation: global___LabelPropagation | None = ..., + page_rank: global___PageRank | None = ..., + parallel_personalized_page_rank: global___ParallelPersonalizedPageRank | None = ..., + power_iteration_clustering: global___PowerIterationClustering | None = ..., + pregel: global___Pregel | None = ..., + shortest_paths: global___ShortestPaths | None = ..., + strongly_connected_components: global___StronglyConnectedComponents | None = ..., + svd_plus_plus: global___SVDPlusPlus | None = ..., + triangle_count: global___TriangleCount | None = ..., + triplets: global___Triplets | None = ..., + kcore: global___KCore | None = ..., + mis: global___MaximalIndependentSet | None = ..., + rw_embeddings: global___RandomWalkEmbeddings | None = ..., + aggregate_neighbors: global___AggregateNeighbors | None = ..., + neighborhood_aware_cdlp: global___NeighborhoodAwareCDLP | None = ..., + all_paths: global___AllPaths | None = ..., + hyper_anf: global___HyperANF | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "aggregate_messages", + b"aggregate_messages", + "aggregate_neighbors", + b"aggregate_neighbors", + "all_paths", + b"all_paths", + "bfs", + b"bfs", + "connected_components", + b"connected_components", + "detecting_cycles", + b"detecting_cycles", + "drop_isolated_vertices", + b"drop_isolated_vertices", + "filter_edges", + b"filter_edges", + "filter_vertices", + b"filter_vertices", + "find", + b"find", + "hyper_anf", + b"hyper_anf", + "kcore", + b"kcore", + "label_propagation", + b"label_propagation", + "method", + b"method", + "mis", + b"mis", + "neighborhood_aware_cdlp", + b"neighborhood_aware_cdlp", + "page_rank", + b"page_rank", + "parallel_personalized_page_rank", + b"parallel_personalized_page_rank", + "power_iteration_clustering", + b"power_iteration_clustering", + "pregel", + b"pregel", + "rw_embeddings", + b"rw_embeddings", + "shortest_paths", + b"shortest_paths", + "strongly_connected_components", + b"strongly_connected_components", + "svd_plus_plus", + b"svd_plus_plus", + "triangle_count", + b"triangle_count", + "triplets", + b"triplets", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "aggregate_messages", + b"aggregate_messages", + "aggregate_neighbors", + b"aggregate_neighbors", + "all_paths", + b"all_paths", + "bfs", + b"bfs", + "connected_components", + b"connected_components", + "detecting_cycles", + b"detecting_cycles", + "drop_isolated_vertices", + b"drop_isolated_vertices", + "edges", + b"edges", + "filter_edges", + b"filter_edges", + "filter_vertices", + b"filter_vertices", + "find", + b"find", + "hyper_anf", + b"hyper_anf", + "kcore", + b"kcore", + "label_propagation", + b"label_propagation", + "method", + b"method", + "mis", + b"mis", + "neighborhood_aware_cdlp", + b"neighborhood_aware_cdlp", + "page_rank", + b"page_rank", + "parallel_personalized_page_rank", + b"parallel_personalized_page_rank", + "power_iteration_clustering", + b"power_iteration_clustering", + "pregel", + b"pregel", + "rw_embeddings", + b"rw_embeddings", + "shortest_paths", + b"shortest_paths", + "strongly_connected_components", + b"strongly_connected_components", + "svd_plus_plus", + b"svd_plus_plus", + "triangle_count", + b"triangle_count", + "triplets", + b"triplets", + "vertices", + b"vertices", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["method", b"method"] + ) -> ( + typing_extensions.Literal[ + "aggregate_messages", + "bfs", + "connected_components", + "drop_isolated_vertices", + "detecting_cycles", + "filter_edges", + "filter_vertices", + "find", + "label_propagation", + "page_rank", + "parallel_personalized_page_rank", + "power_iteration_clustering", + "pregel", + "shortest_paths", + "strongly_connected_components", + "svd_plus_plus", + "triangle_count", + "triplets", + "kcore", + "mis", + "rw_embeddings", + "aggregate_neighbors", + "neighborhood_aware_cdlp", + "all_paths", + "hyper_anf", + ] + | None + ): ... + +global___GraphFramesAPI = GraphFramesAPI + +class StorageLevel(google.protobuf.message.Message): + """Mapping follows PySpark Storage Levels! + (not Scala-Spark Storage Levels) + """ + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + DISK_ONLY_FIELD_NUMBER: builtins.int + DISK_ONLY_2_FIELD_NUMBER: builtins.int + DISK_ONLY_3_FIELD_NUMBER: builtins.int + MEMORY_AND_DISK_FIELD_NUMBER: builtins.int + MEMORY_AND_DISK_2_FIELD_NUMBER: builtins.int + MEMORY_AND_DISK_DESER_FIELD_NUMBER: builtins.int + MEMORY_ONLY_FIELD_NUMBER: builtins.int + MEMORY_ONLY_2_FIELD_NUMBER: builtins.int + disk_only: builtins.bool + disk_only_2: builtins.bool + disk_only_3: builtins.bool + memory_and_disk: builtins.bool + memory_and_disk_2: builtins.bool + memory_and_disk_deser: builtins.bool + memory_only: builtins.bool + memory_only_2: builtins.bool + def __init__( + self, + *, + disk_only: builtins.bool = ..., + disk_only_2: builtins.bool = ..., + disk_only_3: builtins.bool = ..., + memory_and_disk: builtins.bool = ..., + memory_and_disk_2: builtins.bool = ..., + memory_and_disk_deser: builtins.bool = ..., + memory_only: builtins.bool = ..., + memory_only_2: builtins.bool = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "disk_only", + b"disk_only", + "disk_only_2", + b"disk_only_2", + "disk_only_3", + b"disk_only_3", + "memory_and_disk", + b"memory_and_disk", + "memory_and_disk_2", + b"memory_and_disk_2", + "memory_and_disk_deser", + b"memory_and_disk_deser", + "memory_only", + b"memory_only", + "memory_only_2", + b"memory_only_2", + "storage_level", + b"storage_level", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "disk_only", + b"disk_only", + "disk_only_2", + b"disk_only_2", + "disk_only_3", + b"disk_only_3", + "memory_and_disk", + b"memory_and_disk", + "memory_and_disk_2", + b"memory_and_disk_2", + "memory_and_disk_deser", + b"memory_and_disk_deser", + "memory_only", + b"memory_only", + "memory_only_2", + b"memory_only_2", + "storage_level", + b"storage_level", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["storage_level", b"storage_level"] + ) -> ( + typing_extensions.Literal[ + "disk_only", + "disk_only_2", + "disk_only_3", + "memory_and_disk", + "memory_and_disk_2", + "memory_and_disk_deser", + "memory_only", + "memory_only_2", + ] + | None + ): ... + +global___StorageLevel = StorageLevel + +class ColumnOrExpression(google.protobuf.message.Message): + """String expression or serialized column""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + COL_FIELD_NUMBER: builtins.int + EXPR_FIELD_NUMBER: builtins.int + col: builtins.bytes + expr: builtins.str + def __init__( + self, + *, + col: builtins.bytes = ..., + expr: builtins.str = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "col", b"col", "col_or_expr", b"col_or_expr", "expr", b"expr" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "col", b"col", "col_or_expr", b"col_or_expr", "expr", b"expr" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["col_or_expr", b"col_or_expr"] + ) -> typing_extensions.Literal["col", "expr"] | None: ... + +global___ColumnOrExpression = ColumnOrExpression + +class StringOrLongID(google.protobuf.message.Message): + """Connect supports only string or long-like IDs""" + + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + LONG_ID_FIELD_NUMBER: builtins.int + STRING_ID_FIELD_NUMBER: builtins.int + long_id: builtins.int + string_id: builtins.str + def __init__( + self, + *, + long_id: builtins.int = ..., + string_id: builtins.str = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "id", b"id", "long_id", b"long_id", "string_id", b"string_id" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "id", b"id", "long_id", b"long_id", "string_id", b"string_id" + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["id", b"id"] + ) -> typing_extensions.Literal["long_id", "string_id"] | None: ... + +global___StringOrLongID = StringOrLongID + +class AggregateMessages(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + AGG_COL_FIELD_NUMBER: builtins.int + SEND_TO_SRC_FIELD_NUMBER: builtins.int + SEND_TO_DST_FIELD_NUMBER: builtins.int + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + @property + def agg_col( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___ColumnOrExpression + ]: ... + @property + def send_to_src( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___ColumnOrExpression + ]: ... + @property + def send_to_dst( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___ColumnOrExpression + ]: ... + @property + def storage_level(self) -> global___StorageLevel: ... + def __init__( + self, + *, + agg_col: collections.abc.Iterable[global___ColumnOrExpression] | None = ..., + send_to_src: collections.abc.Iterable[global___ColumnOrExpression] | None = ..., + send_to_dst: collections.abc.Iterable[global___ColumnOrExpression] | None = ..., + storage_level: global___StorageLevel | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", b"_storage_level", "storage_level", b"storage_level" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", + b"_storage_level", + "agg_col", + b"agg_col", + "send_to_dst", + b"send_to_dst", + "send_to_src", + b"send_to_src", + "storage_level", + b"storage_level", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_storage_level", b"_storage_level"] + ) -> typing_extensions.Literal["storage_level"] | None: ... + +global___AggregateMessages = AggregateMessages + +class BFS(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + FROM_EXPR_FIELD_NUMBER: builtins.int + TO_EXPR_FIELD_NUMBER: builtins.int + EDGE_FILTER_FIELD_NUMBER: builtins.int + MAX_PATH_LENGTH_FIELD_NUMBER: builtins.int + @property + def from_expr(self) -> global___ColumnOrExpression: ... + @property + def to_expr(self) -> global___ColumnOrExpression: ... + @property + def edge_filter(self) -> global___ColumnOrExpression: ... + max_path_length: builtins.int + def __init__( + self, + *, + from_expr: global___ColumnOrExpression | None = ..., + to_expr: global___ColumnOrExpression | None = ..., + edge_filter: global___ColumnOrExpression | None = ..., + max_path_length: builtins.int = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "edge_filter", b"edge_filter", "from_expr", b"from_expr", "to_expr", b"to_expr" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "edge_filter", + b"edge_filter", + "from_expr", + b"from_expr", + "max_path_length", + b"max_path_length", + "to_expr", + b"to_expr", + ], + ) -> None: ... + +global___BFS = BFS + +class AllPaths(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + FROM_EXPR_FIELD_NUMBER: builtins.int + TO_EXPR_FIELD_NUMBER: builtins.int + EDGE_FILTER_FIELD_NUMBER: builtins.int + MAX_PATH_LENGTH_FIELD_NUMBER: builtins.int + IS_DIRECTED_FIELD_NUMBER: builtins.int + CHECKPOINT_INTERVAL_FIELD_NUMBER: builtins.int + USE_LOCAL_CHECKPOINTS_FIELD_NUMBER: builtins.int + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + @property + def from_expr(self) -> global___ColumnOrExpression: ... + @property + def to_expr(self) -> global___ColumnOrExpression: ... + @property + def edge_filter(self) -> global___ColumnOrExpression: ... + max_path_length: builtins.int + is_directed: builtins.bool + checkpoint_interval: builtins.int + use_local_checkpoints: builtins.bool + @property + def storage_level(self) -> global___StorageLevel: ... + def __init__( + self, + *, + from_expr: global___ColumnOrExpression | None = ..., + to_expr: global___ColumnOrExpression | None = ..., + edge_filter: global___ColumnOrExpression | None = ..., + max_path_length: builtins.int = ..., + is_directed: builtins.bool = ..., + checkpoint_interval: builtins.int = ..., + use_local_checkpoints: builtins.bool = ..., + storage_level: global___StorageLevel | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", + b"_storage_level", + "edge_filter", + b"edge_filter", + "from_expr", + b"from_expr", + "storage_level", + b"storage_level", + "to_expr", + b"to_expr", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", + b"_storage_level", + "checkpoint_interval", + b"checkpoint_interval", + "edge_filter", + b"edge_filter", + "from_expr", + b"from_expr", + "is_directed", + b"is_directed", + "max_path_length", + b"max_path_length", + "storage_level", + b"storage_level", + "to_expr", + b"to_expr", + "use_local_checkpoints", + b"use_local_checkpoints", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_storage_level", b"_storage_level"] + ) -> typing_extensions.Literal["storage_level"] | None: ... + +global___AllPaths = AllPaths + +class HyperANF(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + N_HOPS_FIELD_NUMBER: builtins.int + LG_NOM_ENTRIES_FIELD_NUMBER: builtins.int + EDGES_FILTER_EXPRESSION_FIELD_NUMBER: builtins.int + CHECKPOINT_INTERVAL_FIELD_NUMBER: builtins.int + USE_LOCAL_CHECKPOINTS_FIELD_NUMBER: builtins.int + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + n_hops: builtins.int + lg_nom_entries: builtins.int + @property + def edges_filter_expression(self) -> global___ColumnOrExpression: ... + checkpoint_interval: builtins.int + use_local_checkpoints: builtins.bool + @property + def storage_level(self) -> global___StorageLevel: ... + def __init__( + self, + *, + n_hops: builtins.int = ..., + lg_nom_entries: builtins.int = ..., + edges_filter_expression: global___ColumnOrExpression | None = ..., + checkpoint_interval: builtins.int = ..., + use_local_checkpoints: builtins.bool = ..., + storage_level: global___StorageLevel | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_edges_filter_expression", + b"_edges_filter_expression", + "_storage_level", + b"_storage_level", + "edges_filter_expression", + b"edges_filter_expression", + "storage_level", + b"storage_level", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_edges_filter_expression", + b"_edges_filter_expression", + "_storage_level", + b"_storage_level", + "checkpoint_interval", + b"checkpoint_interval", + "edges_filter_expression", + b"edges_filter_expression", + "lg_nom_entries", + b"lg_nom_entries", + "n_hops", + b"n_hops", + "storage_level", + b"storage_level", + "use_local_checkpoints", + b"use_local_checkpoints", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, + oneof_group: typing_extensions.Literal[ + "_edges_filter_expression", b"_edges_filter_expression" + ], + ) -> typing_extensions.Literal["edges_filter_expression"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_storage_level", b"_storage_level"] + ) -> typing_extensions.Literal["storage_level"] | None: ... + +global___HyperANF = HyperANF + +class ConnectedComponents(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + ALGORITHM_FIELD_NUMBER: builtins.int + CHECKPOINT_INTERVAL_FIELD_NUMBER: builtins.int + BROADCAST_THRESHOLD_FIELD_NUMBER: builtins.int + USE_LABELS_AS_COMPONENTS_FIELD_NUMBER: builtins.int + USE_LOCAL_CHECKPOINTS_FIELD_NUMBER: builtins.int + MAX_ITER_FIELD_NUMBER: builtins.int + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + algorithm: builtins.str + checkpoint_interval: builtins.int + broadcast_threshold: builtins.int + use_labels_as_components: builtins.bool + use_local_checkpoints: builtins.bool + max_iter: builtins.int + @property + def storage_level(self) -> global___StorageLevel: ... + def __init__( + self, + *, + algorithm: builtins.str = ..., + checkpoint_interval: builtins.int = ..., + broadcast_threshold: builtins.int = ..., + use_labels_as_components: builtins.bool = ..., + use_local_checkpoints: builtins.bool = ..., + max_iter: builtins.int = ..., + storage_level: global___StorageLevel | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", b"_storage_level", "storage_level", b"storage_level" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", + b"_storage_level", + "algorithm", + b"algorithm", + "broadcast_threshold", + b"broadcast_threshold", + "checkpoint_interval", + b"checkpoint_interval", + "max_iter", + b"max_iter", + "storage_level", + b"storage_level", + "use_labels_as_components", + b"use_labels_as_components", + "use_local_checkpoints", + b"use_local_checkpoints", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_storage_level", b"_storage_level"] + ) -> typing_extensions.Literal["storage_level"] | None: ... + +global___ConnectedComponents = ConnectedComponents + +class DetectingCycles(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + USE_LOCAL_CHECKPOINTS_FIELD_NUMBER: builtins.int + CHECKPOINT_INTERVAL_FIELD_NUMBER: builtins.int + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + use_local_checkpoints: builtins.bool + checkpoint_interval: builtins.int + @property + def storage_level(self) -> global___StorageLevel: ... + def __init__( + self, + *, + use_local_checkpoints: builtins.bool = ..., + checkpoint_interval: builtins.int = ..., + storage_level: global___StorageLevel | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", b"_storage_level", "storage_level", b"storage_level" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", + b"_storage_level", + "checkpoint_interval", + b"checkpoint_interval", + "storage_level", + b"storage_level", + "use_local_checkpoints", + b"use_local_checkpoints", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_storage_level", b"_storage_level"] + ) -> typing_extensions.Literal["storage_level"] | None: ... + +global___DetectingCycles = DetectingCycles + +class DropIsolatedVertices(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + def __init__( + self, + ) -> None: ... + +global___DropIsolatedVertices = DropIsolatedVertices + +class FilterEdges(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + CONDITION_FIELD_NUMBER: builtins.int + @property + def condition(self) -> global___ColumnOrExpression: ... + def __init__( + self, + *, + condition: global___ColumnOrExpression | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["condition", b"condition"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["condition", b"condition"] + ) -> None: ... + +global___FilterEdges = FilterEdges + +class FilterVertices(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + CONDITION_FIELD_NUMBER: builtins.int + @property + def condition(self) -> global___ColumnOrExpression: ... + def __init__( + self, + *, + condition: global___ColumnOrExpression | None = ..., + ) -> None: ... + def HasField( + self, field_name: typing_extensions.Literal["condition", b"condition"] + ) -> builtins.bool: ... + def ClearField( + self, field_name: typing_extensions.Literal["condition", b"condition"] + ) -> None: ... + +global___FilterVertices = FilterVertices + +class Find(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + PATTERN_FIELD_NUMBER: builtins.int + pattern: builtins.str + def __init__( + self, + *, + pattern: builtins.str = ..., + ) -> None: ... + def ClearField(self, field_name: typing_extensions.Literal["pattern", b"pattern"]) -> None: ... + +global___Find = Find + +class LabelPropagation(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + ALGORITHM_FIELD_NUMBER: builtins.int + MAX_ITER_FIELD_NUMBER: builtins.int + USE_LOCAL_CHECKPOINTS_FIELD_NUMBER: builtins.int + CHECKPOINT_INTERVAL_FIELD_NUMBER: builtins.int + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + algorithm: builtins.str + max_iter: builtins.int + use_local_checkpoints: builtins.bool + checkpoint_interval: builtins.int + @property + def storage_level(self) -> global___StorageLevel: ... + def __init__( + self, + *, + algorithm: builtins.str = ..., + max_iter: builtins.int = ..., + use_local_checkpoints: builtins.bool = ..., + checkpoint_interval: builtins.int = ..., + storage_level: global___StorageLevel | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", b"_storage_level", "storage_level", b"storage_level" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", + b"_storage_level", + "algorithm", + b"algorithm", + "checkpoint_interval", + b"checkpoint_interval", + "max_iter", + b"max_iter", + "storage_level", + b"storage_level", + "use_local_checkpoints", + b"use_local_checkpoints", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_storage_level", b"_storage_level"] + ) -> typing_extensions.Literal["storage_level"] | None: ... + +global___LabelPropagation = LabelPropagation + +class NeighborhoodAwareCDLP(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + MAX_ITER_FIELD_NUMBER: builtins.int + IGNORE_DIRECT_LINKS_FIELD_NUMBER: builtins.int + STRUCTURAL_SIMILARITY_MULTIPLIER_FIELD_NUMBER: builtins.int + USE_LOCAL_CHECKPOINTS_FIELD_NUMBER: builtins.int + CHECKPOINT_INTERVAL_FIELD_NUMBER: builtins.int + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + IS_DIRECTED_FIELD_NUMBER: builtins.int + LG_NOM_ENTRIES_FIELD_NUMBER: builtins.int + INITIAL_LABEL_COL_FIELD_NUMBER: builtins.int + max_iter: builtins.int + ignore_direct_links: builtins.bool + structural_similarity_multiplier: builtins.float + use_local_checkpoints: builtins.bool + checkpoint_interval: builtins.int + @property + def storage_level(self) -> global___StorageLevel: ... + is_directed: builtins.bool + lg_nom_entries: builtins.int + initial_label_col: builtins.str + def __init__( + self, + *, + max_iter: builtins.int = ..., + ignore_direct_links: builtins.bool = ..., + structural_similarity_multiplier: builtins.float = ..., + use_local_checkpoints: builtins.bool = ..., + checkpoint_interval: builtins.int = ..., + storage_level: global___StorageLevel | None = ..., + is_directed: builtins.bool = ..., + lg_nom_entries: builtins.int = ..., + initial_label_col: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_initial_label_col", + b"_initial_label_col", + "_storage_level", + b"_storage_level", + "initial_label_col", + b"initial_label_col", + "storage_level", + b"storage_level", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_initial_label_col", + b"_initial_label_col", + "_storage_level", + b"_storage_level", + "checkpoint_interval", + b"checkpoint_interval", + "ignore_direct_links", + b"ignore_direct_links", + "initial_label_col", + b"initial_label_col", + "is_directed", + b"is_directed", + "lg_nom_entries", + b"lg_nom_entries", + "max_iter", + b"max_iter", + "storage_level", + b"storage_level", + "structural_similarity_multiplier", + b"structural_similarity_multiplier", + "use_local_checkpoints", + b"use_local_checkpoints", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_initial_label_col", b"_initial_label_col"] + ) -> typing_extensions.Literal["initial_label_col"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_storage_level", b"_storage_level"] + ) -> typing_extensions.Literal["storage_level"] | None: ... + +global___NeighborhoodAwareCDLP = NeighborhoodAwareCDLP + +class PageRank(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + RESET_PROBABILITY_FIELD_NUMBER: builtins.int + SOURCE_ID_FIELD_NUMBER: builtins.int + MAX_ITER_FIELD_NUMBER: builtins.int + TOL_FIELD_NUMBER: builtins.int + reset_probability: builtins.float + @property + def source_id(self) -> global___StringOrLongID: ... + max_iter: builtins.int + tol: builtins.float + def __init__( + self, + *, + reset_probability: builtins.float = ..., + source_id: global___StringOrLongID | None = ..., + max_iter: builtins.int | None = ..., + tol: builtins.float | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_max_iter", + b"_max_iter", + "_source_id", + b"_source_id", + "_tol", + b"_tol", + "max_iter", + b"max_iter", + "source_id", + b"source_id", + "tol", + b"tol", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_max_iter", + b"_max_iter", + "_source_id", + b"_source_id", + "_tol", + b"_tol", + "max_iter", + b"max_iter", + "reset_probability", + b"reset_probability", + "source_id", + b"source_id", + "tol", + b"tol", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_max_iter", b"_max_iter"] + ) -> typing_extensions.Literal["max_iter"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_source_id", b"_source_id"] + ) -> typing_extensions.Literal["source_id"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_tol", b"_tol"] + ) -> typing_extensions.Literal["tol"] | None: ... + +global___PageRank = PageRank + +class ParallelPersonalizedPageRank(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + RESET_PROBABILITY_FIELD_NUMBER: builtins.int + SOURCE_IDS_FIELD_NUMBER: builtins.int + MAX_ITER_FIELD_NUMBER: builtins.int + reset_probability: builtins.float + @property + def source_ids( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___StringOrLongID + ]: ... + max_iter: builtins.int + def __init__( + self, + *, + reset_probability: builtins.float = ..., + source_ids: collections.abc.Iterable[global___StringOrLongID] | None = ..., + max_iter: builtins.int = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "max_iter", + b"max_iter", + "reset_probability", + b"reset_probability", + "source_ids", + b"source_ids", + ], + ) -> None: ... + +global___ParallelPersonalizedPageRank = ParallelPersonalizedPageRank + +class PowerIterationClustering(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + K_FIELD_NUMBER: builtins.int + MAX_ITER_FIELD_NUMBER: builtins.int + WEIGHT_COL_FIELD_NUMBER: builtins.int + k: builtins.int + max_iter: builtins.int + weight_col: builtins.str + def __init__( + self, + *, + k: builtins.int = ..., + max_iter: builtins.int = ..., + weight_col: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_weight_col", b"_weight_col", "weight_col", b"weight_col" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_weight_col", + b"_weight_col", + "k", + b"k", + "max_iter", + b"max_iter", + "weight_col", + b"weight_col", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_weight_col", b"_weight_col"] + ) -> typing_extensions.Literal["weight_col"] | None: ... + +global___PowerIterationClustering = PowerIterationClustering + +class Pregel(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + AGG_MSGS_FIELD_NUMBER: builtins.int + SEND_MSG_TO_DST_FIELD_NUMBER: builtins.int + SEND_MSG_TO_SRC_FIELD_NUMBER: builtins.int + CHECKPOINT_INTERVAL_FIELD_NUMBER: builtins.int + MAX_ITER_FIELD_NUMBER: builtins.int + ADDITIONAL_COL_NAME_FIELD_NUMBER: builtins.int + ADDITIONAL_COL_INITIAL_FIELD_NUMBER: builtins.int + ADDITIONAL_COL_UPD_FIELD_NUMBER: builtins.int + EARLY_STOPPING_FIELD_NUMBER: builtins.int + USE_LOCAL_CHECKPOINTS_FIELD_NUMBER: builtins.int + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + STOP_IF_ALL_NON_ACTIVE_FIELD_NUMBER: builtins.int + INITIAL_ACTIVE_EXPR_FIELD_NUMBER: builtins.int + UPDATE_ACTIVE_EXPR_FIELD_NUMBER: builtins.int + SKIP_MESSAGES_FROM_NON_ACTIVE_FIELD_NUMBER: builtins.int + REQUIRED_SRC_COLUMNS_FIELD_NUMBER: builtins.int + REQUIRED_DST_COLUMNS_FIELD_NUMBER: builtins.int + REQUIRED_EDGE_COLUMNS_FIELD_NUMBER: builtins.int + @property + def agg_msgs(self) -> global___ColumnOrExpression: ... + @property + def send_msg_to_dst( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___ColumnOrExpression + ]: ... + @property + def send_msg_to_src( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___ColumnOrExpression + ]: ... + checkpoint_interval: builtins.int + max_iter: builtins.int + additional_col_name: builtins.str + @property + def additional_col_initial(self) -> global___ColumnOrExpression: ... + @property + def additional_col_upd(self) -> global___ColumnOrExpression: ... + early_stopping: builtins.bool + use_local_checkpoints: builtins.bool + @property + def storage_level(self) -> global___StorageLevel: ... + stop_if_all_non_active: builtins.bool + @property + def initial_active_expr(self) -> global___ColumnOrExpression: ... + @property + def update_active_expr(self) -> global___ColumnOrExpression: ... + skip_messages_from_non_active: builtins.bool + required_src_columns: builtins.str + """Required columns for triplet construction (memory optimization) + Column names separated by comma + """ + required_dst_columns: builtins.str + required_edge_columns: builtins.str + def __init__( + self, + *, + agg_msgs: global___ColumnOrExpression | None = ..., + send_msg_to_dst: collections.abc.Iterable[global___ColumnOrExpression] | None = ..., + send_msg_to_src: collections.abc.Iterable[global___ColumnOrExpression] | None = ..., + checkpoint_interval: builtins.int = ..., + max_iter: builtins.int = ..., + additional_col_name: builtins.str = ..., + additional_col_initial: global___ColumnOrExpression | None = ..., + additional_col_upd: global___ColumnOrExpression | None = ..., + early_stopping: builtins.bool | None = ..., + use_local_checkpoints: builtins.bool = ..., + storage_level: global___StorageLevel | None = ..., + stop_if_all_non_active: builtins.bool | None = ..., + initial_active_expr: global___ColumnOrExpression | None = ..., + update_active_expr: global___ColumnOrExpression | None = ..., + skip_messages_from_non_active: builtins.bool | None = ..., + required_src_columns: builtins.str | None = ..., + required_dst_columns: builtins.str | None = ..., + required_edge_columns: builtins.str | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_early_stopping", + b"_early_stopping", + "_initial_active_expr", + b"_initial_active_expr", + "_required_dst_columns", + b"_required_dst_columns", + "_required_edge_columns", + b"_required_edge_columns", + "_required_src_columns", + b"_required_src_columns", + "_skip_messages_from_non_active", + b"_skip_messages_from_non_active", + "_stop_if_all_non_active", + b"_stop_if_all_non_active", + "_storage_level", + b"_storage_level", + "_update_active_expr", + b"_update_active_expr", + "additional_col_initial", + b"additional_col_initial", + "additional_col_upd", + b"additional_col_upd", + "agg_msgs", + b"agg_msgs", + "early_stopping", + b"early_stopping", + "initial_active_expr", + b"initial_active_expr", + "required_dst_columns", + b"required_dst_columns", + "required_edge_columns", + b"required_edge_columns", + "required_src_columns", + b"required_src_columns", + "skip_messages_from_non_active", + b"skip_messages_from_non_active", + "stop_if_all_non_active", + b"stop_if_all_non_active", + "storage_level", + b"storage_level", + "update_active_expr", + b"update_active_expr", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_early_stopping", + b"_early_stopping", + "_initial_active_expr", + b"_initial_active_expr", + "_required_dst_columns", + b"_required_dst_columns", + "_required_edge_columns", + b"_required_edge_columns", + "_required_src_columns", + b"_required_src_columns", + "_skip_messages_from_non_active", + b"_skip_messages_from_non_active", + "_stop_if_all_non_active", + b"_stop_if_all_non_active", + "_storage_level", + b"_storage_level", + "_update_active_expr", + b"_update_active_expr", + "additional_col_initial", + b"additional_col_initial", + "additional_col_name", + b"additional_col_name", + "additional_col_upd", + b"additional_col_upd", + "agg_msgs", + b"agg_msgs", + "checkpoint_interval", + b"checkpoint_interval", + "early_stopping", + b"early_stopping", + "initial_active_expr", + b"initial_active_expr", + "max_iter", + b"max_iter", + "required_dst_columns", + b"required_dst_columns", + "required_edge_columns", + b"required_edge_columns", + "required_src_columns", + b"required_src_columns", + "send_msg_to_dst", + b"send_msg_to_dst", + "send_msg_to_src", + b"send_msg_to_src", + "skip_messages_from_non_active", + b"skip_messages_from_non_active", + "stop_if_all_non_active", + b"stop_if_all_non_active", + "storage_level", + b"storage_level", + "update_active_expr", + b"update_active_expr", + "use_local_checkpoints", + b"use_local_checkpoints", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_early_stopping", b"_early_stopping"] + ) -> typing_extensions.Literal["early_stopping"] | None: ... + @typing.overload + def WhichOneof( + self, + oneof_group: typing_extensions.Literal["_initial_active_expr", b"_initial_active_expr"], + ) -> typing_extensions.Literal["initial_active_expr"] | None: ... + @typing.overload + def WhichOneof( + self, + oneof_group: typing_extensions.Literal["_required_dst_columns", b"_required_dst_columns"], + ) -> typing_extensions.Literal["required_dst_columns"] | None: ... + @typing.overload + def WhichOneof( + self, + oneof_group: typing_extensions.Literal["_required_edge_columns", b"_required_edge_columns"], + ) -> typing_extensions.Literal["required_edge_columns"] | None: ... + @typing.overload + def WhichOneof( + self, + oneof_group: typing_extensions.Literal["_required_src_columns", b"_required_src_columns"], + ) -> typing_extensions.Literal["required_src_columns"] | None: ... + @typing.overload + def WhichOneof( + self, + oneof_group: typing_extensions.Literal[ + "_skip_messages_from_non_active", b"_skip_messages_from_non_active" + ], + ) -> typing_extensions.Literal["skip_messages_from_non_active"] | None: ... + @typing.overload + def WhichOneof( + self, + oneof_group: typing_extensions.Literal[ + "_stop_if_all_non_active", b"_stop_if_all_non_active" + ], + ) -> typing_extensions.Literal["stop_if_all_non_active"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_storage_level", b"_storage_level"] + ) -> typing_extensions.Literal["storage_level"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_update_active_expr", b"_update_active_expr"] + ) -> typing_extensions.Literal["update_active_expr"] | None: ... + +global___Pregel = Pregel + +class ShortestPaths(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + LANDMARKS_FIELD_NUMBER: builtins.int + ALGORITHM_FIELD_NUMBER: builtins.int + USE_LOCAL_CHECKPOINTS_FIELD_NUMBER: builtins.int + CHECKPOINT_INTERVAL_FIELD_NUMBER: builtins.int + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + IS_DIRECTED_FIELD_NUMBER: builtins.int + @property + def landmarks( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___StringOrLongID + ]: ... + algorithm: builtins.str + use_local_checkpoints: builtins.bool + checkpoint_interval: builtins.int + @property + def storage_level(self) -> global___StorageLevel: ... + is_directed: builtins.bool + def __init__( + self, + *, + landmarks: collections.abc.Iterable[global___StringOrLongID] | None = ..., + algorithm: builtins.str = ..., + use_local_checkpoints: builtins.bool = ..., + checkpoint_interval: builtins.int = ..., + storage_level: global___StorageLevel | None = ..., + is_directed: builtins.bool | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_is_directed", + b"_is_directed", + "_storage_level", + b"_storage_level", + "is_directed", + b"is_directed", + "storage_level", + b"storage_level", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_is_directed", + b"_is_directed", + "_storage_level", + b"_storage_level", + "algorithm", + b"algorithm", + "checkpoint_interval", + b"checkpoint_interval", + "is_directed", + b"is_directed", + "landmarks", + b"landmarks", + "storage_level", + b"storage_level", + "use_local_checkpoints", + b"use_local_checkpoints", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_is_directed", b"_is_directed"] + ) -> typing_extensions.Literal["is_directed"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_storage_level", b"_storage_level"] + ) -> typing_extensions.Literal["storage_level"] | None: ... + +global___ShortestPaths = ShortestPaths + +class StronglyConnectedComponents(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + MAX_ITER_FIELD_NUMBER: builtins.int + max_iter: builtins.int + def __init__( + self, + *, + max_iter: builtins.int = ..., + ) -> None: ... + def ClearField( + self, field_name: typing_extensions.Literal["max_iter", b"max_iter"] + ) -> None: ... + +global___StronglyConnectedComponents = StronglyConnectedComponents + +class SVDPlusPlus(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + RANK_FIELD_NUMBER: builtins.int + MAX_ITER_FIELD_NUMBER: builtins.int + MIN_VALUE_FIELD_NUMBER: builtins.int + MAX_VALUE_FIELD_NUMBER: builtins.int + GAMMA1_FIELD_NUMBER: builtins.int + GAMMA2_FIELD_NUMBER: builtins.int + GAMMA6_FIELD_NUMBER: builtins.int + GAMMA7_FIELD_NUMBER: builtins.int + rank: builtins.int + max_iter: builtins.int + min_value: builtins.float + max_value: builtins.float + gamma1: builtins.float + gamma2: builtins.float + gamma6: builtins.float + gamma7: builtins.float + def __init__( + self, + *, + rank: builtins.int = ..., + max_iter: builtins.int = ..., + min_value: builtins.float = ..., + max_value: builtins.float = ..., + gamma1: builtins.float = ..., + gamma2: builtins.float = ..., + gamma6: builtins.float = ..., + gamma7: builtins.float = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "gamma1", + b"gamma1", + "gamma2", + b"gamma2", + "gamma6", + b"gamma6", + "gamma7", + b"gamma7", + "max_iter", + b"max_iter", + "max_value", + b"max_value", + "min_value", + b"min_value", + "rank", + b"rank", + ], + ) -> None: ... + +global___SVDPlusPlus = SVDPlusPlus + +class TriangleCount(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + ALGORITHM_FIELD_NUMBER: builtins.int + LG_NOM_ENTRIES_FIELD_NUMBER: builtins.int + @property + def storage_level(self) -> global___StorageLevel: ... + algorithm: builtins.str + lg_nom_entries: builtins.int + def __init__( + self, + *, + storage_level: global___StorageLevel | None = ..., + algorithm: builtins.str | None = ..., + lg_nom_entries: builtins.int | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_algorithm", + b"_algorithm", + "_lg_nom_entries", + b"_lg_nom_entries", + "_storage_level", + b"_storage_level", + "algorithm", + b"algorithm", + "lg_nom_entries", + b"lg_nom_entries", + "storage_level", + b"storage_level", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_algorithm", + b"_algorithm", + "_lg_nom_entries", + b"_lg_nom_entries", + "_storage_level", + b"_storage_level", + "algorithm", + b"algorithm", + "lg_nom_entries", + b"lg_nom_entries", + "storage_level", + b"storage_level", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_algorithm", b"_algorithm"] + ) -> typing_extensions.Literal["algorithm"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_lg_nom_entries", b"_lg_nom_entries"] + ) -> typing_extensions.Literal["lg_nom_entries"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_storage_level", b"_storage_level"] + ) -> typing_extensions.Literal["storage_level"] | None: ... + +global___TriangleCount = TriangleCount + +class Triplets(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + def __init__( + self, + ) -> None: ... + +global___Triplets = Triplets + +class MaximalIndependentSet(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + CHECKPOINT_INTERVAL_FIELD_NUMBER: builtins.int + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + USE_LOCAL_CHECKPOINTS_FIELD_NUMBER: builtins.int + SEED_FIELD_NUMBER: builtins.int + checkpoint_interval: builtins.int + @property + def storage_level(self) -> global___StorageLevel: ... + use_local_checkpoints: builtins.bool + seed: builtins.int + def __init__( + self, + *, + checkpoint_interval: builtins.int = ..., + storage_level: global___StorageLevel | None = ..., + use_local_checkpoints: builtins.bool = ..., + seed: builtins.int = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", b"_storage_level", "storage_level", b"storage_level" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", + b"_storage_level", + "checkpoint_interval", + b"checkpoint_interval", + "seed", + b"seed", + "storage_level", + b"storage_level", + "use_local_checkpoints", + b"use_local_checkpoints", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_storage_level", b"_storage_level"] + ) -> typing_extensions.Literal["storage_level"] | None: ... + +global___MaximalIndependentSet = MaximalIndependentSet + +class KCore(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + USE_LOCAL_CHECKPOINTS_FIELD_NUMBER: builtins.int + CHECKPOINT_INTERVAL_FIELD_NUMBER: builtins.int + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + use_local_checkpoints: builtins.bool + checkpoint_interval: builtins.int + @property + def storage_level(self) -> global___StorageLevel: ... + def __init__( + self, + *, + use_local_checkpoints: builtins.bool = ..., + checkpoint_interval: builtins.int = ..., + storage_level: global___StorageLevel | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", b"_storage_level", "storage_level", b"storage_level" + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_storage_level", + b"_storage_level", + "checkpoint_interval", + b"checkpoint_interval", + "storage_level", + b"storage_level", + "use_local_checkpoints", + b"use_local_checkpoints", + ], + ) -> None: ... + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_storage_level", b"_storage_level"] + ) -> typing_extensions.Literal["storage_level"] | None: ... + +global___KCore = KCore + +class AggregateNeighbors(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + STARTING_VERTICES_FIELD_NUMBER: builtins.int + MAX_HOPS_FIELD_NUMBER: builtins.int + ACCUMULATOR_NAMES_FIELD_NUMBER: builtins.int + ACCUMULATOR_INITS_FIELD_NUMBER: builtins.int + ACCUMULATOR_UPDATES_FIELD_NUMBER: builtins.int + STOPPING_CONDITION_FIELD_NUMBER: builtins.int + TARGET_CONDITION_FIELD_NUMBER: builtins.int + REQUIRED_VERTEX_ATTRIBUTES_FIELD_NUMBER: builtins.int + REQUIRED_EDGE_ATTRIBUTES_FIELD_NUMBER: builtins.int + EDGE_FILTER_FIELD_NUMBER: builtins.int + REMOVE_LOOPS_FIELD_NUMBER: builtins.int + CHECKPOINT_INTERVAL_FIELD_NUMBER: builtins.int + USE_LOCAL_CHECKPOINTS_FIELD_NUMBER: builtins.int + STORAGE_LEVEL_FIELD_NUMBER: builtins.int + @property + def starting_vertices(self) -> global___ColumnOrExpression: + """Starting vertices condition (Boolean column expression)""" + max_hops: builtins.int + """Maximum number of hops to explore""" + @property + def accumulator_names( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """Accumulator names""" + @property + def accumulator_inits( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___ColumnOrExpression + ]: + """Accumulator initial value expressions""" + @property + def accumulator_updates( + self, + ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[ + global___ColumnOrExpression + ]: + """Accumulator update expressions""" + @property + def stopping_condition(self) -> global___ColumnOrExpression: + """Optional stopping condition (Boolean column expression)""" + @property + def target_condition(self) -> global___ColumnOrExpression: + """Optional target condition (Boolean column expression)""" + @property + def required_vertex_attributes( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """Optional required vertex attributes to carry through traversal""" + @property + def required_edge_attributes( + self, + ) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[builtins.str]: + """Optional required edge attributes to carry through traversal""" + @property + def edge_filter(self) -> global___ColumnOrExpression: + """Optional edge filter condition (Boolean column expression)""" + remove_loops: builtins.bool + """Whether to remove self-loops""" + checkpoint_interval: builtins.int + """Checkpoint interval (0 means disabled)""" + use_local_checkpoints: builtins.bool + """Whether to use local checkpoints""" + @property + def storage_level(self) -> global___StorageLevel: + """Optional storage level for intermediate results""" + def __init__( + self, + *, + starting_vertices: global___ColumnOrExpression | None = ..., + max_hops: builtins.int = ..., + accumulator_names: collections.abc.Iterable[builtins.str] | None = ..., + accumulator_inits: collections.abc.Iterable[global___ColumnOrExpression] | None = ..., + accumulator_updates: collections.abc.Iterable[global___ColumnOrExpression] | None = ..., + stopping_condition: global___ColumnOrExpression | None = ..., + target_condition: global___ColumnOrExpression | None = ..., + required_vertex_attributes: collections.abc.Iterable[builtins.str] | None = ..., + required_edge_attributes: collections.abc.Iterable[builtins.str] | None = ..., + edge_filter: global___ColumnOrExpression | None = ..., + remove_loops: builtins.bool = ..., + checkpoint_interval: builtins.int = ..., + use_local_checkpoints: builtins.bool = ..., + storage_level: global___StorageLevel | None = ..., + ) -> None: ... + def HasField( + self, + field_name: typing_extensions.Literal[ + "_edge_filter", + b"_edge_filter", + "_stopping_condition", + b"_stopping_condition", + "_storage_level", + b"_storage_level", + "_target_condition", + b"_target_condition", + "edge_filter", + b"edge_filter", + "starting_vertices", + b"starting_vertices", + "stopping_condition", + b"stopping_condition", + "storage_level", + b"storage_level", + "target_condition", + b"target_condition", + ], + ) -> builtins.bool: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "_edge_filter", + b"_edge_filter", + "_stopping_condition", + b"_stopping_condition", + "_storage_level", + b"_storage_level", + "_target_condition", + b"_target_condition", + "accumulator_inits", + b"accumulator_inits", + "accumulator_names", + b"accumulator_names", + "accumulator_updates", + b"accumulator_updates", + "checkpoint_interval", + b"checkpoint_interval", + "edge_filter", + b"edge_filter", + "max_hops", + b"max_hops", + "remove_loops", + b"remove_loops", + "required_edge_attributes", + b"required_edge_attributes", + "required_vertex_attributes", + b"required_vertex_attributes", + "starting_vertices", + b"starting_vertices", + "stopping_condition", + b"stopping_condition", + "storage_level", + b"storage_level", + "target_condition", + b"target_condition", + "use_local_checkpoints", + b"use_local_checkpoints", + ], + ) -> None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_edge_filter", b"_edge_filter"] + ) -> typing_extensions.Literal["edge_filter"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_stopping_condition", b"_stopping_condition"] + ) -> typing_extensions.Literal["stopping_condition"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_storage_level", b"_storage_level"] + ) -> typing_extensions.Literal["storage_level"] | None: ... + @typing.overload + def WhichOneof( + self, oneof_group: typing_extensions.Literal["_target_condition", b"_target_condition"] + ) -> typing_extensions.Literal["target_condition"] | None: ... + +global___AggregateNeighbors = AggregateNeighbors + +class RandomWalkEmbeddings(google.protobuf.message.Message): + DESCRIPTOR: google.protobuf.descriptor.Descriptor + + USE_EDGE_DIRECTION_FIELD_NUMBER: builtins.int + RW_MODEL_FIELD_NUMBER: builtins.int + RW_MAX_NBRS_FIELD_NUMBER: builtins.int + RW_NUM_WALKS_PER_NODE_FIELD_NUMBER: builtins.int + RW_BATCH_SIZE_FIELD_NUMBER: builtins.int + RW_NUM_BATCHES_FIELD_NUMBER: builtins.int + RW_SEED_FIELD_NUMBER: builtins.int + RW_RESTART_PROBABILITY_FIELD_NUMBER: builtins.int + RW_TEMPORARY_PREFIX_FIELD_NUMBER: builtins.int + RW_CACHED_WALKS_FIELD_NUMBER: builtins.int + SEQUENCE_MODEL_FIELD_NUMBER: builtins.int + HASH2VEC_CONTEXT_SIZE_FIELD_NUMBER: builtins.int + HASH2VEC_NUM_PARTITIONS_FIELD_NUMBER: builtins.int + HASH2VEC_EMBEDDINGS_DIM_FIELD_NUMBER: builtins.int + HASH2VEC_DECAY_FUNCTION_FIELD_NUMBER: builtins.int + HASH2VEC_GAUSSIAN_SIGMA_FIELD_NUMBER: builtins.int + HASH2VEC_HASHING_SEED_FIELD_NUMBER: builtins.int + HASH2VEC_SIGN_SEED_FIELD_NUMBER: builtins.int + HASH2VEC_DO_L2_NORM_FIELD_NUMBER: builtins.int + HASH2VEC_SAFE_L2_FIELD_NUMBER: builtins.int + WORD2VEC_MAX_ITER_FIELD_NUMBER: builtins.int + WORD2VEC_EMBEDDINGS_DIM_FIELD_NUMBER: builtins.int + WORD2VEC_WINDOW_SIZE_FIELD_NUMBER: builtins.int + WORD2VEC_NUM_PARTITIONS_FIELD_NUMBER: builtins.int + WORD2VEC_MIN_COUNT_FIELD_NUMBER: builtins.int + WORD2VEC_MAX_SENTENCE_LENGTH_FIELD_NUMBER: builtins.int + WORD2VEC_SEED_FIELD_NUMBER: builtins.int + WORD2VEC_STEP_SIZE_FIELD_NUMBER: builtins.int + AGGREGATE_NEIGHBORS_FIELD_NUMBER: builtins.int + AGGREGATE_NEIGHBORS_MAX_NBRS_FIELD_NUMBER: builtins.int + AGGREGATE_NEIGHBORS_SEED_FIELD_NUMBER: builtins.int + CLEAN_UP_AFTER_RUN_FIELD_NUMBER: builtins.int + use_edge_direction: builtins.bool + rw_model: builtins.str + rw_max_nbrs: builtins.int + rw_num_walks_per_node: builtins.int + rw_batch_size: builtins.int + rw_num_batches: builtins.int + rw_seed: builtins.int + rw_restart_probability: builtins.float + rw_temporary_prefix: builtins.str + rw_cached_walks: builtins.str + sequence_model: builtins.str + hash2vec_context_size: builtins.int + hash2vec_num_partitions: builtins.int + hash2vec_embeddings_dim: builtins.int + hash2vec_decay_function: builtins.str + hash2vec_gaussian_sigma: builtins.float + hash2vec_hashing_seed: builtins.int + hash2vec_sign_seed: builtins.int + hash2vec_do_l2_norm: builtins.bool + hash2vec_safe_l2: builtins.bool + word2vec_max_iter: builtins.int + word2vec_embeddings_dim: builtins.int + word2vec_window_size: builtins.int + word2vec_num_partitions: builtins.int + word2vec_min_count: builtins.int + word2vec_max_sentence_length: builtins.int + word2vec_seed: builtins.int + word2vec_step_size: builtins.float + aggregate_neighbors: builtins.bool + aggregate_neighbors_max_nbrs: builtins.int + aggregate_neighbors_seed: builtins.int + clean_up_after_run: builtins.bool + def __init__( + self, + *, + use_edge_direction: builtins.bool = ..., + rw_model: builtins.str = ..., + rw_max_nbrs: builtins.int = ..., + rw_num_walks_per_node: builtins.int = ..., + rw_batch_size: builtins.int = ..., + rw_num_batches: builtins.int = ..., + rw_seed: builtins.int = ..., + rw_restart_probability: builtins.float = ..., + rw_temporary_prefix: builtins.str = ..., + rw_cached_walks: builtins.str = ..., + sequence_model: builtins.str = ..., + hash2vec_context_size: builtins.int = ..., + hash2vec_num_partitions: builtins.int = ..., + hash2vec_embeddings_dim: builtins.int = ..., + hash2vec_decay_function: builtins.str = ..., + hash2vec_gaussian_sigma: builtins.float = ..., + hash2vec_hashing_seed: builtins.int = ..., + hash2vec_sign_seed: builtins.int = ..., + hash2vec_do_l2_norm: builtins.bool = ..., + hash2vec_safe_l2: builtins.bool = ..., + word2vec_max_iter: builtins.int = ..., + word2vec_embeddings_dim: builtins.int = ..., + word2vec_window_size: builtins.int = ..., + word2vec_num_partitions: builtins.int = ..., + word2vec_min_count: builtins.int = ..., + word2vec_max_sentence_length: builtins.int = ..., + word2vec_seed: builtins.int = ..., + word2vec_step_size: builtins.float = ..., + aggregate_neighbors: builtins.bool = ..., + aggregate_neighbors_max_nbrs: builtins.int = ..., + aggregate_neighbors_seed: builtins.int = ..., + clean_up_after_run: builtins.bool = ..., + ) -> None: ... + def ClearField( + self, + field_name: typing_extensions.Literal[ + "aggregate_neighbors", + b"aggregate_neighbors", + "aggregate_neighbors_max_nbrs", + b"aggregate_neighbors_max_nbrs", + "aggregate_neighbors_seed", + b"aggregate_neighbors_seed", + "clean_up_after_run", + b"clean_up_after_run", + "hash2vec_context_size", + b"hash2vec_context_size", + "hash2vec_decay_function", + b"hash2vec_decay_function", + "hash2vec_do_l2_norm", + b"hash2vec_do_l2_norm", + "hash2vec_embeddings_dim", + b"hash2vec_embeddings_dim", + "hash2vec_gaussian_sigma", + b"hash2vec_gaussian_sigma", + "hash2vec_hashing_seed", + b"hash2vec_hashing_seed", + "hash2vec_num_partitions", + b"hash2vec_num_partitions", + "hash2vec_safe_l2", + b"hash2vec_safe_l2", + "hash2vec_sign_seed", + b"hash2vec_sign_seed", + "rw_batch_size", + b"rw_batch_size", + "rw_cached_walks", + b"rw_cached_walks", + "rw_max_nbrs", + b"rw_max_nbrs", + "rw_model", + b"rw_model", + "rw_num_batches", + b"rw_num_batches", + "rw_num_walks_per_node", + b"rw_num_walks_per_node", + "rw_restart_probability", + b"rw_restart_probability", + "rw_seed", + b"rw_seed", + "rw_temporary_prefix", + b"rw_temporary_prefix", + "sequence_model", + b"sequence_model", + "use_edge_direction", + b"use_edge_direction", + "word2vec_embeddings_dim", + b"word2vec_embeddings_dim", + "word2vec_max_iter", + b"word2vec_max_iter", + "word2vec_max_sentence_length", + b"word2vec_max_sentence_length", + "word2vec_min_count", + b"word2vec_min_count", + "word2vec_num_partitions", + b"word2vec_num_partitions", + "word2vec_seed", + b"word2vec_seed", + "word2vec_step_size", + b"word2vec_step_size", + "word2vec_window_size", + b"word2vec_window_size", + ], + ) -> None: ... + +global___RandomWalkEmbeddings = RandomWalkEmbeddings diff --git a/python/pyspark/sql/connect/proto/relations_pb2.py b/python/pyspark/sql/connect/proto/relations_pb2.py index 8b56455e69f2b..6f29be5c57e2b 100644 --- a/python/pyspark/sql/connect/proto/relations_pb2.py +++ b/python/pyspark/sql/connect/proto/relations_pb2.py @@ -40,11 +40,12 @@ from pyspark.sql.connect.proto import types_pb2 as spark_dot_connect_dot_types__pb2 from pyspark.sql.connect.proto import catalog_pb2 as spark_dot_connect_dot_catalog__pb2 from pyspark.sql.connect.proto import common_pb2 as spark_dot_connect_dot_common__pb2 +from pyspark.sql.connect.proto import graphframes_pb2 as spark_dot_connect_dot_graphframes__pb2 from pyspark.sql.connect.proto import ml_common_pb2 as spark_dot_connect_dot_ml__common__pb2 DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( - b'\n\x1dspark/connect/relations.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x1fspark/connect/expressions.proto\x1a\x19spark/connect/types.proto\x1a\x1bspark/connect/catalog.proto\x1a\x1aspark/connect/common.proto\x1a\x1dspark/connect/ml_common.proto"\xc9\x1f\n\x08Relation\x12\x35\n\x06\x63ommon\x18\x01 \x01(\x0b\x32\x1d.spark.connect.RelationCommonR\x06\x63ommon\x12)\n\x04read\x18\x02 \x01(\x0b\x32\x13.spark.connect.ReadH\x00R\x04read\x12\x32\n\x07project\x18\x03 \x01(\x0b\x32\x16.spark.connect.ProjectH\x00R\x07project\x12/\n\x06\x66ilter\x18\x04 \x01(\x0b\x32\x15.spark.connect.FilterH\x00R\x06\x66ilter\x12)\n\x04join\x18\x05 \x01(\x0b\x32\x13.spark.connect.JoinH\x00R\x04join\x12\x34\n\x06set_op\x18\x06 \x01(\x0b\x32\x1b.spark.connect.SetOperationH\x00R\x05setOp\x12)\n\x04sort\x18\x07 \x01(\x0b\x32\x13.spark.connect.SortH\x00R\x04sort\x12,\n\x05limit\x18\x08 \x01(\x0b\x32\x14.spark.connect.LimitH\x00R\x05limit\x12\x38\n\taggregate\x18\t \x01(\x0b\x32\x18.spark.connect.AggregateH\x00R\taggregate\x12&\n\x03sql\x18\n \x01(\x0b\x32\x12.spark.connect.SQLH\x00R\x03sql\x12\x45\n\x0elocal_relation\x18\x0b \x01(\x0b\x32\x1c.spark.connect.LocalRelationH\x00R\rlocalRelation\x12/\n\x06sample\x18\x0c \x01(\x0b\x32\x15.spark.connect.SampleH\x00R\x06sample\x12/\n\x06offset\x18\r \x01(\x0b\x32\x15.spark.connect.OffsetH\x00R\x06offset\x12>\n\x0b\x64\x65\x64uplicate\x18\x0e \x01(\x0b\x32\x1a.spark.connect.DeduplicateH\x00R\x0b\x64\x65\x64uplicate\x12,\n\x05range\x18\x0f \x01(\x0b\x32\x14.spark.connect.RangeH\x00R\x05range\x12\x45\n\x0esubquery_alias\x18\x10 \x01(\x0b\x32\x1c.spark.connect.SubqueryAliasH\x00R\rsubqueryAlias\x12>\n\x0brepartition\x18\x11 \x01(\x0b\x32\x1a.spark.connect.RepartitionH\x00R\x0brepartition\x12*\n\x05to_df\x18\x12 \x01(\x0b\x32\x13.spark.connect.ToDFH\x00R\x04toDf\x12U\n\x14with_columns_renamed\x18\x13 \x01(\x0b\x32!.spark.connect.WithColumnsRenamedH\x00R\x12withColumnsRenamed\x12<\n\x0bshow_string\x18\x14 \x01(\x0b\x32\x19.spark.connect.ShowStringH\x00R\nshowString\x12)\n\x04\x64rop\x18\x15 \x01(\x0b\x32\x13.spark.connect.DropH\x00R\x04\x64rop\x12)\n\x04tail\x18\x16 \x01(\x0b\x32\x13.spark.connect.TailH\x00R\x04tail\x12?\n\x0cwith_columns\x18\x17 \x01(\x0b\x32\x1a.spark.connect.WithColumnsH\x00R\x0bwithColumns\x12)\n\x04hint\x18\x18 \x01(\x0b\x32\x13.spark.connect.HintH\x00R\x04hint\x12\x32\n\x07unpivot\x18\x19 \x01(\x0b\x32\x16.spark.connect.UnpivotH\x00R\x07unpivot\x12\x36\n\tto_schema\x18\x1a \x01(\x0b\x32\x17.spark.connect.ToSchemaH\x00R\x08toSchema\x12\x64\n\x19repartition_by_expression\x18\x1b \x01(\x0b\x32&.spark.connect.RepartitionByExpressionH\x00R\x17repartitionByExpression\x12\x45\n\x0emap_partitions\x18\x1c \x01(\x0b\x32\x1c.spark.connect.MapPartitionsH\x00R\rmapPartitions\x12H\n\x0f\x63ollect_metrics\x18\x1d \x01(\x0b\x32\x1d.spark.connect.CollectMetricsH\x00R\x0e\x63ollectMetrics\x12,\n\x05parse\x18\x1e \x01(\x0b\x32\x14.spark.connect.ParseH\x00R\x05parse\x12\x36\n\tgroup_map\x18\x1f \x01(\x0b\x32\x17.spark.connect.GroupMapH\x00R\x08groupMap\x12=\n\x0c\x63o_group_map\x18 \x01(\x0b\x32\x19.spark.connect.CoGroupMapH\x00R\ncoGroupMap\x12\x45\n\x0ewith_watermark\x18! \x01(\x0b\x32\x1c.spark.connect.WithWatermarkH\x00R\rwithWatermark\x12\x63\n\x1a\x61pply_in_pandas_with_state\x18" \x01(\x0b\x32%.spark.connect.ApplyInPandasWithStateH\x00R\x16\x61pplyInPandasWithState\x12<\n\x0bhtml_string\x18# \x01(\x0b\x32\x19.spark.connect.HtmlStringH\x00R\nhtmlString\x12X\n\x15\x63\x61\x63hed_local_relation\x18$ \x01(\x0b\x32".spark.connect.CachedLocalRelationH\x00R\x13\x63\x61\x63hedLocalRelation\x12[\n\x16\x63\x61\x63hed_remote_relation\x18% \x01(\x0b\x32#.spark.connect.CachedRemoteRelationH\x00R\x14\x63\x61\x63hedRemoteRelation\x12\x8e\x01\n)common_inline_user_defined_table_function\x18& \x01(\x0b\x32\x33.spark.connect.CommonInlineUserDefinedTableFunctionH\x00R$commonInlineUserDefinedTableFunction\x12\x37\n\nas_of_join\x18\' \x01(\x0b\x32\x17.spark.connect.AsOfJoinH\x00R\x08\x61sOfJoin\x12\x85\x01\n&common_inline_user_defined_data_source\x18( \x01(\x0b\x32\x30.spark.connect.CommonInlineUserDefinedDataSourceH\x00R!commonInlineUserDefinedDataSource\x12\x45\n\x0ewith_relations\x18) \x01(\x0b\x32\x1c.spark.connect.WithRelationsH\x00R\rwithRelations\x12\x38\n\ttranspose\x18* \x01(\x0b\x32\x18.spark.connect.TransposeH\x00R\ttranspose\x12w\n unresolved_table_valued_function\x18+ \x01(\x0b\x32,.spark.connect.UnresolvedTableValuedFunctionH\x00R\x1dunresolvedTableValuedFunction\x12?\n\x0clateral_join\x18, \x01(\x0b\x32\x1a.spark.connect.LateralJoinH\x00R\x0blateralJoin\x12n\n\x1d\x63hunked_cached_local_relation\x18- \x01(\x0b\x32).spark.connect.ChunkedCachedLocalRelationH\x00R\x1a\x63hunkedCachedLocalRelation\x12K\n\x10relation_changes\x18. \x01(\x0b\x32\x1e.spark.connect.RelationChangesH\x00R\x0frelationChanges\x12\x46\n\x0fnearest_by_join\x18/ \x01(\x0b\x32\x1c.spark.connect.NearestByJoinH\x00R\rnearestByJoin\x12&\n\x03zip\x18\x30 \x01(\x0b\x32\x12.spark.connect.ZipH\x00R\x03zip\x12\x30\n\x07\x66ill_na\x18Z \x01(\x0b\x32\x15.spark.connect.NAFillH\x00R\x06\x66illNa\x12\x30\n\x07\x64rop_na\x18[ \x01(\x0b\x32\x15.spark.connect.NADropH\x00R\x06\x64ropNa\x12\x34\n\x07replace\x18\\ \x01(\x0b\x32\x18.spark.connect.NAReplaceH\x00R\x07replace\x12\x36\n\x07summary\x18\x64 \x01(\x0b\x32\x1a.spark.connect.StatSummaryH\x00R\x07summary\x12\x39\n\x08\x63rosstab\x18\x65 \x01(\x0b\x32\x1b.spark.connect.StatCrosstabH\x00R\x08\x63rosstab\x12\x39\n\x08\x64\x65scribe\x18\x66 \x01(\x0b\x32\x1b.spark.connect.StatDescribeH\x00R\x08\x64\x65scribe\x12*\n\x03\x63ov\x18g \x01(\x0b\x32\x16.spark.connect.StatCovH\x00R\x03\x63ov\x12-\n\x04\x63orr\x18h \x01(\x0b\x32\x17.spark.connect.StatCorrH\x00R\x04\x63orr\x12L\n\x0f\x61pprox_quantile\x18i \x01(\x0b\x32!.spark.connect.StatApproxQuantileH\x00R\x0e\x61pproxQuantile\x12=\n\nfreq_items\x18j \x01(\x0b\x32\x1c.spark.connect.StatFreqItemsH\x00R\tfreqItems\x12:\n\tsample_by\x18k \x01(\x0b\x32\x1b.spark.connect.StatSampleByH\x00R\x08sampleBy\x12\x33\n\x07\x63\x61talog\x18\xc8\x01 \x01(\x0b\x32\x16.spark.connect.CatalogH\x00R\x07\x63\x61talog\x12=\n\x0bml_relation\x18\xac\x02 \x01(\x0b\x32\x19.spark.connect.MlRelationH\x00R\nmlRelation\x12\x35\n\textension\x18\xe6\x07 \x01(\x0b\x32\x14.google.protobuf.AnyH\x00R\textension\x12\x33\n\x07unknown\x18\xe7\x07 \x01(\x0b\x32\x16.spark.connect.UnknownH\x00R\x07unknownB\n\n\x08rel_type"\xe4\x03\n\nMlRelation\x12\x43\n\ttransform\x18\x01 \x01(\x0b\x32#.spark.connect.MlRelation.TransformH\x00R\ttransform\x12,\n\x05\x66\x65tch\x18\x02 \x01(\x0b\x32\x14.spark.connect.FetchH\x00R\x05\x66\x65tch\x12P\n\x15model_summary_dataset\x18\x03 \x01(\x0b\x32\x17.spark.connect.RelationH\x01R\x13modelSummaryDataset\x88\x01\x01\x1a\xeb\x01\n\tTransform\x12\x33\n\x07obj_ref\x18\x01 \x01(\x0b\x32\x18.spark.connect.ObjectRefH\x00R\x06objRef\x12=\n\x0btransformer\x18\x02 \x01(\x0b\x32\x19.spark.connect.MlOperatorH\x00R\x0btransformer\x12-\n\x05input\x18\x03 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12/\n\x06params\x18\x04 \x01(\x0b\x32\x17.spark.connect.MlParamsR\x06paramsB\n\n\x08operatorB\t\n\x07ml_typeB\x18\n\x16_model_summary_dataset"\xcb\x02\n\x05\x46\x65tch\x12\x31\n\x07obj_ref\x18\x01 \x01(\x0b\x32\x18.spark.connect.ObjectRefR\x06objRef\x12\x35\n\x07methods\x18\x02 \x03(\x0b\x32\x1b.spark.connect.Fetch.MethodR\x07methods\x1a\xd7\x01\n\x06Method\x12\x16\n\x06method\x18\x01 \x01(\tR\x06method\x12\x34\n\x04\x61rgs\x18\x02 \x03(\x0b\x32 .spark.connect.Fetch.Method.ArgsR\x04\x61rgs\x1a\x7f\n\x04\x41rgs\x12\x39\n\x05param\x18\x01 \x01(\x0b\x32!.spark.connect.Expression.LiteralH\x00R\x05param\x12/\n\x05input\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationH\x00R\x05inputB\x0b\n\targs_type"\t\n\x07Unknown"\x8e\x01\n\x0eRelationCommon\x12#\n\x0bsource_info\x18\x01 \x01(\tB\x02\x18\x01R\nsourceInfo\x12\x1c\n\x07plan_id\x18\x02 \x01(\x03H\x00R\x06planId\x88\x01\x01\x12-\n\x06origin\x18\x03 \x01(\x0b\x32\x15.spark.connect.OriginR\x06originB\n\n\x08_plan_id"\xde\x03\n\x03SQL\x12\x14\n\x05query\x18\x01 \x01(\tR\x05query\x12\x34\n\x04\x61rgs\x18\x02 \x03(\x0b\x32\x1c.spark.connect.SQL.ArgsEntryB\x02\x18\x01R\x04\x61rgs\x12@\n\x08pos_args\x18\x03 \x03(\x0b\x32!.spark.connect.Expression.LiteralB\x02\x18\x01R\x07posArgs\x12O\n\x0fnamed_arguments\x18\x04 \x03(\x0b\x32&.spark.connect.SQL.NamedArgumentsEntryR\x0enamedArguments\x12>\n\rpos_arguments\x18\x05 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x0cposArguments\x1aZ\n\tArgsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x37\n\x05value\x18\x02 \x01(\x0b\x32!.spark.connect.Expression.LiteralR\x05value:\x02\x38\x01\x1a\\\n\x13NamedArgumentsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12/\n\x05value\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x05value:\x02\x38\x01"u\n\rWithRelations\x12+\n\x04root\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x04root\x12\x37\n\nreferences\x18\x02 \x03(\x0b\x32\x17.spark.connect.RelationR\nreferences"\xcd\x05\n\x04Read\x12\x41\n\x0bnamed_table\x18\x01 \x01(\x0b\x32\x1e.spark.connect.Read.NamedTableH\x00R\nnamedTable\x12\x41\n\x0b\x64\x61ta_source\x18\x02 \x01(\x0b\x32\x1e.spark.connect.Read.DataSourceH\x00R\ndataSource\x12!\n\x0cis_streaming\x18\x03 \x01(\x08R\x0bisStreaming\x1a\xc0\x01\n\nNamedTable\x12/\n\x13unparsed_identifier\x18\x01 \x01(\tR\x12unparsedIdentifier\x12\x45\n\x07options\x18\x02 \x03(\x0b\x32+.spark.connect.Read.NamedTable.OptionsEntryR\x07options\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x1a\xcb\x02\n\nDataSource\x12\x1b\n\x06\x66ormat\x18\x01 \x01(\tH\x00R\x06\x66ormat\x88\x01\x01\x12\x1b\n\x06schema\x18\x02 \x01(\tH\x01R\x06schema\x88\x01\x01\x12\x45\n\x07options\x18\x03 \x03(\x0b\x32+.spark.connect.Read.DataSource.OptionsEntryR\x07options\x12\x14\n\x05paths\x18\x04 \x03(\tR\x05paths\x12\x1e\n\npredicates\x18\x05 \x03(\tR\npredicates\x12$\n\x0bsource_name\x18\x06 \x01(\tH\x02R\nsourceName\x88\x01\x01\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x42\t\n\x07_formatB\t\n\x07_schemaB\x0e\n\x0c_source_nameB\x0b\n\tread_type"\xe8\x01\n\x0fRelationChanges\x12/\n\x13unparsed_identifier\x18\x01 \x01(\tR\x12unparsedIdentifier\x12\x45\n\x07options\x18\x02 \x03(\x0b\x32+.spark.connect.RelationChanges.OptionsEntryR\x07options\x12!\n\x0cis_streaming\x18\x03 \x01(\x08R\x0bisStreaming\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01"u\n\x07Project\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12;\n\x0b\x65xpressions\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x0b\x65xpressions"p\n\x06\x46ilter\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x37\n\tcondition\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionR\tcondition"\x95\x05\n\x04Join\x12+\n\x04left\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x04left\x12-\n\x05right\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationR\x05right\x12@\n\x0ejoin_condition\x18\x03 \x01(\x0b\x32\x19.spark.connect.ExpressionR\rjoinCondition\x12\x39\n\tjoin_type\x18\x04 \x01(\x0e\x32\x1c.spark.connect.Join.JoinTypeR\x08joinType\x12#\n\rusing_columns\x18\x05 \x03(\tR\x0cusingColumns\x12K\n\x0ejoin_data_type\x18\x06 \x01(\x0b\x32 .spark.connect.Join.JoinDataTypeH\x00R\x0cjoinDataType\x88\x01\x01\x1a\\\n\x0cJoinDataType\x12$\n\x0eis_left_struct\x18\x01 \x01(\x08R\x0cisLeftStruct\x12&\n\x0fis_right_struct\x18\x02 \x01(\x08R\risRightStruct"\xd0\x01\n\x08JoinType\x12\x19\n\x15JOIN_TYPE_UNSPECIFIED\x10\x00\x12\x13\n\x0fJOIN_TYPE_INNER\x10\x01\x12\x18\n\x14JOIN_TYPE_FULL_OUTER\x10\x02\x12\x18\n\x14JOIN_TYPE_LEFT_OUTER\x10\x03\x12\x19\n\x15JOIN_TYPE_RIGHT_OUTER\x10\x04\x12\x17\n\x13JOIN_TYPE_LEFT_ANTI\x10\x05\x12\x17\n\x13JOIN_TYPE_LEFT_SEMI\x10\x06\x12\x13\n\x0fJOIN_TYPE_CROSS\x10\x07\x42\x11\n\x0f_join_data_type"\xdf\x03\n\x0cSetOperation\x12\x36\n\nleft_input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\tleftInput\x12\x38\n\x0bright_input\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationR\nrightInput\x12\x45\n\x0bset_op_type\x18\x03 \x01(\x0e\x32%.spark.connect.SetOperation.SetOpTypeR\tsetOpType\x12\x1a\n\x06is_all\x18\x04 \x01(\x08H\x00R\x05isAll\x88\x01\x01\x12\x1c\n\x07\x62y_name\x18\x05 \x01(\x08H\x01R\x06\x62yName\x88\x01\x01\x12\x37\n\x15\x61llow_missing_columns\x18\x06 \x01(\x08H\x02R\x13\x61llowMissingColumns\x88\x01\x01"r\n\tSetOpType\x12\x1b\n\x17SET_OP_TYPE_UNSPECIFIED\x10\x00\x12\x19\n\x15SET_OP_TYPE_INTERSECT\x10\x01\x12\x15\n\x11SET_OP_TYPE_UNION\x10\x02\x12\x16\n\x12SET_OP_TYPE_EXCEPT\x10\x03\x42\t\n\x07_is_allB\n\n\x08_by_nameB\x18\n\x16_allow_missing_columns"L\n\x05Limit\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x14\n\x05limit\x18\x02 \x01(\x05R\x05limit"O\n\x06Offset\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x16\n\x06offset\x18\x02 \x01(\x05R\x06offset"K\n\x04Tail\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x14\n\x05limit\x18\x02 \x01(\x05R\x05limit"\xfe\x05\n\tAggregate\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x41\n\ngroup_type\x18\x02 \x01(\x0e\x32".spark.connect.Aggregate.GroupTypeR\tgroupType\x12L\n\x14grouping_expressions\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x13groupingExpressions\x12N\n\x15\x61ggregate_expressions\x18\x04 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x14\x61ggregateExpressions\x12\x34\n\x05pivot\x18\x05 \x01(\x0b\x32\x1e.spark.connect.Aggregate.PivotR\x05pivot\x12J\n\rgrouping_sets\x18\x06 \x03(\x0b\x32%.spark.connect.Aggregate.GroupingSetsR\x0cgroupingSets\x1ao\n\x05Pivot\x12+\n\x03\x63ol\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x03\x63ol\x12\x39\n\x06values\x18\x02 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x06values\x1aL\n\x0cGroupingSets\x12<\n\x0cgrouping_set\x18\x01 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x0bgroupingSet"\x9f\x01\n\tGroupType\x12\x1a\n\x16GROUP_TYPE_UNSPECIFIED\x10\x00\x12\x16\n\x12GROUP_TYPE_GROUPBY\x10\x01\x12\x15\n\x11GROUP_TYPE_ROLLUP\x10\x02\x12\x13\n\x0fGROUP_TYPE_CUBE\x10\x03\x12\x14\n\x10GROUP_TYPE_PIVOT\x10\x04\x12\x1c\n\x18GROUP_TYPE_GROUPING_SETS\x10\x05"\xa0\x01\n\x04Sort\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x39\n\x05order\x18\x02 \x03(\x0b\x32#.spark.connect.Expression.SortOrderR\x05order\x12 \n\tis_global\x18\x03 \x01(\x08H\x00R\x08isGlobal\x88\x01\x01\x42\x0c\n\n_is_global"\x8d\x01\n\x04\x44rop\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x33\n\x07\x63olumns\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x07\x63olumns\x12!\n\x0c\x63olumn_names\x18\x03 \x03(\tR\x0b\x63olumnNames"\xf0\x01\n\x0b\x44\x65\x64uplicate\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12!\n\x0c\x63olumn_names\x18\x02 \x03(\tR\x0b\x63olumnNames\x12\x32\n\x13\x61ll_columns_as_keys\x18\x03 \x01(\x08H\x00R\x10\x61llColumnsAsKeys\x88\x01\x01\x12.\n\x10within_watermark\x18\x04 \x01(\x08H\x01R\x0fwithinWatermark\x88\x01\x01\x42\x16\n\x14_all_columns_as_keysB\x13\n\x11_within_watermark"Y\n\rLocalRelation\x12\x17\n\x04\x64\x61ta\x18\x01 \x01(\x0cH\x00R\x04\x64\x61ta\x88\x01\x01\x12\x1b\n\x06schema\x18\x02 \x01(\tH\x01R\x06schema\x88\x01\x01\x42\x07\n\x05_dataB\t\n\x07_schema"H\n\x13\x43\x61\x63hedLocalRelation\x12\x12\n\x04hash\x18\x03 \x01(\tR\x04hashJ\x04\x08\x01\x10\x02J\x04\x08\x02\x10\x03R\x06userIdR\tsessionId"p\n\x1a\x43hunkedCachedLocalRelation\x12\x1e\n\ndataHashes\x18\x01 \x03(\tR\ndataHashes\x12#\n\nschemaHash\x18\x02 \x01(\tH\x00R\nschemaHash\x88\x01\x01\x42\r\n\x0b_schemaHash"7\n\x14\x43\x61\x63hedRemoteRelation\x12\x1f\n\x0brelation_id\x18\x01 \x01(\tR\nrelationId"\x91\x02\n\x06Sample\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x1f\n\x0blower_bound\x18\x02 \x01(\x01R\nlowerBound\x12\x1f\n\x0bupper_bound\x18\x03 \x01(\x01R\nupperBound\x12.\n\x10with_replacement\x18\x04 \x01(\x08H\x00R\x0fwithReplacement\x88\x01\x01\x12\x17\n\x04seed\x18\x05 \x01(\x03H\x01R\x04seed\x88\x01\x01\x12/\n\x13\x64\x65terministic_order\x18\x06 \x01(\x08R\x12\x64\x65terministicOrderB\x13\n\x11_with_replacementB\x07\n\x05_seed"\x91\x01\n\x05Range\x12\x19\n\x05start\x18\x01 \x01(\x03H\x00R\x05start\x88\x01\x01\x12\x10\n\x03\x65nd\x18\x02 \x01(\x03R\x03\x65nd\x12\x12\n\x04step\x18\x03 \x01(\x03R\x04step\x12*\n\x0enum_partitions\x18\x04 \x01(\x05H\x01R\rnumPartitions\x88\x01\x01\x42\x08\n\x06_startB\x11\n\x0f_num_partitions"r\n\rSubqueryAlias\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x14\n\x05\x61lias\x18\x02 \x01(\tR\x05\x61lias\x12\x1c\n\tqualifier\x18\x03 \x03(\tR\tqualifier"\x8e\x01\n\x0bRepartition\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12%\n\x0enum_partitions\x18\x02 \x01(\x05R\rnumPartitions\x12\x1d\n\x07shuffle\x18\x03 \x01(\x08H\x00R\x07shuffle\x88\x01\x01\x42\n\n\x08_shuffle"\x8e\x01\n\nShowString\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x19\n\x08num_rows\x18\x02 \x01(\x05R\x07numRows\x12\x1a\n\x08truncate\x18\x03 \x01(\x05R\x08truncate\x12\x1a\n\x08vertical\x18\x04 \x01(\x08R\x08vertical"r\n\nHtmlString\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x19\n\x08num_rows\x18\x02 \x01(\x05R\x07numRows\x12\x1a\n\x08truncate\x18\x03 \x01(\x05R\x08truncate"\\\n\x0bStatSummary\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x1e\n\nstatistics\x18\x02 \x03(\tR\nstatistics"Q\n\x0cStatDescribe\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols"e\n\x0cStatCrosstab\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ol1\x18\x02 \x01(\tR\x04\x63ol1\x12\x12\n\x04\x63ol2\x18\x03 \x01(\tR\x04\x63ol2"`\n\x07StatCov\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ol1\x18\x02 \x01(\tR\x04\x63ol1\x12\x12\n\x04\x63ol2\x18\x03 \x01(\tR\x04\x63ol2"\x89\x01\n\x08StatCorr\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ol1\x18\x02 \x01(\tR\x04\x63ol1\x12\x12\n\x04\x63ol2\x18\x03 \x01(\tR\x04\x63ol2\x12\x1b\n\x06method\x18\x04 \x01(\tH\x00R\x06method\x88\x01\x01\x42\t\n\x07_method"\xa4\x01\n\x12StatApproxQuantile\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols\x12$\n\rprobabilities\x18\x03 \x03(\x01R\rprobabilities\x12%\n\x0erelative_error\x18\x04 \x01(\x01R\rrelativeError"}\n\rStatFreqItems\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols\x12\x1d\n\x07support\x18\x03 \x01(\x01H\x00R\x07support\x88\x01\x01\x42\n\n\x08_support"\xb5\x02\n\x0cStatSampleBy\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12+\n\x03\x63ol\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x03\x63ol\x12\x42\n\tfractions\x18\x03 \x03(\x0b\x32$.spark.connect.StatSampleBy.FractionR\tfractions\x12\x17\n\x04seed\x18\x05 \x01(\x03H\x00R\x04seed\x88\x01\x01\x1a\x63\n\x08\x46raction\x12;\n\x07stratum\x18\x01 \x01(\x0b\x32!.spark.connect.Expression.LiteralR\x07stratum\x12\x1a\n\x08\x66raction\x18\x02 \x01(\x01R\x08\x66ractionB\x07\n\x05_seed"\x86\x01\n\x06NAFill\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols\x12\x39\n\x06values\x18\x03 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x06values"\x86\x01\n\x06NADrop\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols\x12\'\n\rmin_non_nulls\x18\x03 \x01(\x05H\x00R\x0bminNonNulls\x88\x01\x01\x42\x10\n\x0e_min_non_nulls"\xa8\x02\n\tNAReplace\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols\x12H\n\x0creplacements\x18\x03 \x03(\x0b\x32$.spark.connect.NAReplace.ReplacementR\x0creplacements\x1a\x8d\x01\n\x0bReplacement\x12>\n\told_value\x18\x01 \x01(\x0b\x32!.spark.connect.Expression.LiteralR\x08oldValue\x12>\n\tnew_value\x18\x02 \x01(\x0b\x32!.spark.connect.Expression.LiteralR\x08newValue"X\n\x04ToDF\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12!\n\x0c\x63olumn_names\x18\x02 \x03(\tR\x0b\x63olumnNames"\xfe\x02\n\x12WithColumnsRenamed\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12i\n\x12rename_columns_map\x18\x02 \x03(\x0b\x32\x37.spark.connect.WithColumnsRenamed.RenameColumnsMapEntryB\x02\x18\x01R\x10renameColumnsMap\x12\x42\n\x07renames\x18\x03 \x03(\x0b\x32(.spark.connect.WithColumnsRenamed.RenameR\x07renames\x1a\x43\n\x15RenameColumnsMapEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x1a\x45\n\x06Rename\x12\x19\n\x08\x63ol_name\x18\x01 \x01(\tR\x07\x63olName\x12 \n\x0cnew_col_name\x18\x02 \x01(\tR\nnewColName"w\n\x0bWithColumns\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x39\n\x07\x61liases\x18\x02 \x03(\x0b\x32\x1f.spark.connect.Expression.AliasR\x07\x61liases"\x86\x01\n\rWithWatermark\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x1d\n\nevent_time\x18\x02 \x01(\tR\teventTime\x12\'\n\x0f\x64\x65lay_threshold\x18\x03 \x01(\tR\x0e\x64\x65layThreshold"\x84\x01\n\x04Hint\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04name\x18\x02 \x01(\tR\x04name\x12\x39\n\nparameters\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\nparameters"\xc7\x02\n\x07Unpivot\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12+\n\x03ids\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x03ids\x12:\n\x06values\x18\x03 \x01(\x0b\x32\x1d.spark.connect.Unpivot.ValuesH\x00R\x06values\x88\x01\x01\x12\x30\n\x14variable_column_name\x18\x04 \x01(\tR\x12variableColumnName\x12*\n\x11value_column_name\x18\x05 \x01(\tR\x0fvalueColumnName\x1a;\n\x06Values\x12\x31\n\x06values\x18\x01 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x06valuesB\t\n\x07_values"z\n\tTranspose\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12>\n\rindex_columns\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x0cindexColumns"}\n\x1dUnresolvedTableValuedFunction\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12\x37\n\targuments\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\targuments"j\n\x08ToSchema\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12/\n\x06schema\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x06schema"\xcb\x01\n\x17RepartitionByExpression\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x42\n\x0fpartition_exprs\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x0epartitionExprs\x12*\n\x0enum_partitions\x18\x03 \x01(\x05H\x00R\rnumPartitions\x88\x01\x01\x42\x11\n\x0f_num_partitions"\xe8\x01\n\rMapPartitions\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x42\n\x04\x66unc\x18\x02 \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionR\x04\x66unc\x12"\n\nis_barrier\x18\x03 \x01(\x08H\x00R\tisBarrier\x88\x01\x01\x12"\n\nprofile_id\x18\x04 \x01(\x05H\x01R\tprofileId\x88\x01\x01\x42\r\n\x0b_is_barrierB\r\n\x0b_profile_id"\xd2\x06\n\x08GroupMap\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12L\n\x14grouping_expressions\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x13groupingExpressions\x12\x42\n\x04\x66unc\x18\x03 \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionR\x04\x66unc\x12J\n\x13sorting_expressions\x18\x04 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x12sortingExpressions\x12<\n\rinitial_input\x18\x05 \x01(\x0b\x32\x17.spark.connect.RelationR\x0cinitialInput\x12[\n\x1cinitial_grouping_expressions\x18\x06 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x1ainitialGroupingExpressions\x12;\n\x18is_map_groups_with_state\x18\x07 \x01(\x08H\x00R\x14isMapGroupsWithState\x88\x01\x01\x12$\n\x0boutput_mode\x18\x08 \x01(\tH\x01R\noutputMode\x88\x01\x01\x12&\n\x0ctimeout_conf\x18\t \x01(\tH\x02R\x0btimeoutConf\x88\x01\x01\x12?\n\x0cstate_schema\x18\n \x01(\x0b\x32\x17.spark.connect.DataTypeH\x03R\x0bstateSchema\x88\x01\x01\x12\x65\n\x19transform_with_state_info\x18\x0b \x01(\x0b\x32%.spark.connect.TransformWithStateInfoH\x04R\x16transformWithStateInfo\x88\x01\x01\x42\x1b\n\x19_is_map_groups_with_stateB\x0e\n\x0c_output_modeB\x0f\n\r_timeout_confB\x0f\n\r_state_schemaB\x1c\n\x1a_transform_with_state_info"\xdf\x01\n\x16TransformWithStateInfo\x12\x1b\n\ttime_mode\x18\x01 \x01(\tR\x08timeMode\x12\x38\n\x16\x65vent_time_column_name\x18\x02 \x01(\tH\x00R\x13\x65ventTimeColumnName\x88\x01\x01\x12\x41\n\routput_schema\x18\x03 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x01R\x0coutputSchema\x88\x01\x01\x42\x19\n\x17_event_time_column_nameB\x10\n\x0e_output_schema"\x8e\x04\n\nCoGroupMap\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12W\n\x1ainput_grouping_expressions\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x18inputGroupingExpressions\x12-\n\x05other\x18\x03 \x01(\x0b\x32\x17.spark.connect.RelationR\x05other\x12W\n\x1aother_grouping_expressions\x18\x04 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x18otherGroupingExpressions\x12\x42\n\x04\x66unc\x18\x05 \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionR\x04\x66unc\x12U\n\x19input_sorting_expressions\x18\x06 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x17inputSortingExpressions\x12U\n\x19other_sorting_expressions\x18\x07 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x17otherSortingExpressions"\xe5\x02\n\x16\x41pplyInPandasWithState\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12L\n\x14grouping_expressions\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x13groupingExpressions\x12\x42\n\x04\x66unc\x18\x03 \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionR\x04\x66unc\x12#\n\routput_schema\x18\x04 \x01(\tR\x0coutputSchema\x12!\n\x0cstate_schema\x18\x05 \x01(\tR\x0bstateSchema\x12\x1f\n\x0boutput_mode\x18\x06 \x01(\tR\noutputMode\x12!\n\x0ctimeout_conf\x18\x07 \x01(\tR\x0btimeoutConf"\xf4\x01\n$CommonInlineUserDefinedTableFunction\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12$\n\rdeterministic\x18\x02 \x01(\x08R\rdeterministic\x12\x37\n\targuments\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\targuments\x12<\n\x0bpython_udtf\x18\x04 \x01(\x0b\x32\x19.spark.connect.PythonUDTFH\x00R\npythonUdtfB\n\n\x08\x66unction"\xb1\x01\n\nPythonUDTF\x12=\n\x0breturn_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\nreturnType\x88\x01\x01\x12\x1b\n\teval_type\x18\x02 \x01(\x05R\x08\x65valType\x12\x18\n\x07\x63ommand\x18\x03 \x01(\x0cR\x07\x63ommand\x12\x1d\n\npython_ver\x18\x04 \x01(\tR\tpythonVerB\x0e\n\x0c_return_type"\x97\x01\n!CommonInlineUserDefinedDataSource\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12O\n\x12python_data_source\x18\x02 \x01(\x0b\x32\x1f.spark.connect.PythonDataSourceH\x00R\x10pythonDataSourceB\r\n\x0b\x64\x61ta_source"K\n\x10PythonDataSource\x12\x18\n\x07\x63ommand\x18\x01 \x01(\x0cR\x07\x63ommand\x12\x1d\n\npython_ver\x18\x02 \x01(\tR\tpythonVer"\x88\x01\n\x0e\x43ollectMetrics\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04name\x18\x02 \x01(\tR\x04name\x12\x33\n\x07metrics\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x07metrics"\x9a\x03\n\x05Parse\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x38\n\x06\x66ormat\x18\x02 \x01(\x0e\x32 .spark.connect.Parse.ParseFormatR\x06\x66ormat\x12\x34\n\x06schema\x18\x03 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\x06schema\x88\x01\x01\x12;\n\x07options\x18\x04 \x03(\x0b\x32!.spark.connect.Parse.OptionsEntryR\x07options\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01"n\n\x0bParseFormat\x12\x1c\n\x18PARSE_FORMAT_UNSPECIFIED\x10\x00\x12\x14\n\x10PARSE_FORMAT_CSV\x10\x01\x12\x15\n\x11PARSE_FORMAT_JSON\x10\x02\x12\x14\n\x10PARSE_FORMAT_XML\x10\x03\x42\t\n\x07_schema"\xdb\x03\n\x08\x41sOfJoin\x12+\n\x04left\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x04left\x12-\n\x05right\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationR\x05right\x12\x37\n\nleft_as_of\x18\x03 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x08leftAsOf\x12\x39\n\x0bright_as_of\x18\x04 \x01(\x0b\x32\x19.spark.connect.ExpressionR\trightAsOf\x12\x36\n\tjoin_expr\x18\x05 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x08joinExpr\x12#\n\rusing_columns\x18\x06 \x03(\tR\x0cusingColumns\x12\x1b\n\tjoin_type\x18\x07 \x01(\tR\x08joinType\x12\x37\n\ttolerance\x18\x08 \x01(\x0b\x32\x19.spark.connect.ExpressionR\ttolerance\x12.\n\x13\x61llow_exact_matches\x18\t \x01(\x08R\x11\x61llowExactMatches\x12\x1c\n\tdirection\x18\n \x01(\tR\tdirection"\xe6\x01\n\x0bLateralJoin\x12+\n\x04left\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x04left\x12-\n\x05right\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationR\x05right\x12@\n\x0ejoin_condition\x18\x03 \x01(\x0b\x32\x19.spark.connect.ExpressionR\rjoinCondition\x12\x39\n\tjoin_type\x18\x04 \x01(\x0e\x32\x1c.spark.connect.Join.JoinTypeR\x08joinType"\xa5\x02\n\rNearestByJoin\x12+\n\x04left\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x04left\x12-\n\x05right\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationR\x05right\x12H\n\x12ranking_expression\x18\x03 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x11rankingExpression\x12\x1f\n\x0bnum_results\x18\x04 \x01(\x05R\nnumResults\x12\x1b\n\tjoin_type\x18\x05 \x01(\tR\x08joinType\x12\x12\n\x04mode\x18\x06 \x01(\tR\x04mode\x12\x1c\n\tdirection\x18\x07 \x01(\tR\tdirection"a\n\x03Zip\x12+\n\x04left\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x04left\x12-\n\x05right\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationR\x05rightB6\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3' + b'\n\x1dspark/connect/relations.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x1fspark/connect/expressions.proto\x1a\x19spark/connect/types.proto\x1a\x1bspark/connect/catalog.proto\x1a\x1aspark/connect/common.proto\x1a\x1fspark/connect/graphframes.proto\x1a\x1dspark/connect/ml_common.proto"\x99 \n\x08Relation\x12\x35\n\x06\x63ommon\x18\x01 \x01(\x0b\x32\x1d.spark.connect.RelationCommonR\x06\x63ommon\x12)\n\x04read\x18\x02 \x01(\x0b\x32\x13.spark.connect.ReadH\x00R\x04read\x12\x32\n\x07project\x18\x03 \x01(\x0b\x32\x16.spark.connect.ProjectH\x00R\x07project\x12/\n\x06\x66ilter\x18\x04 \x01(\x0b\x32\x15.spark.connect.FilterH\x00R\x06\x66ilter\x12)\n\x04join\x18\x05 \x01(\x0b\x32\x13.spark.connect.JoinH\x00R\x04join\x12\x34\n\x06set_op\x18\x06 \x01(\x0b\x32\x1b.spark.connect.SetOperationH\x00R\x05setOp\x12)\n\x04sort\x18\x07 \x01(\x0b\x32\x13.spark.connect.SortH\x00R\x04sort\x12,\n\x05limit\x18\x08 \x01(\x0b\x32\x14.spark.connect.LimitH\x00R\x05limit\x12\x38\n\taggregate\x18\t \x01(\x0b\x32\x18.spark.connect.AggregateH\x00R\taggregate\x12&\n\x03sql\x18\n \x01(\x0b\x32\x12.spark.connect.SQLH\x00R\x03sql\x12\x45\n\x0elocal_relation\x18\x0b \x01(\x0b\x32\x1c.spark.connect.LocalRelationH\x00R\rlocalRelation\x12/\n\x06sample\x18\x0c \x01(\x0b\x32\x15.spark.connect.SampleH\x00R\x06sample\x12/\n\x06offset\x18\r \x01(\x0b\x32\x15.spark.connect.OffsetH\x00R\x06offset\x12>\n\x0b\x64\x65\x64uplicate\x18\x0e \x01(\x0b\x32\x1a.spark.connect.DeduplicateH\x00R\x0b\x64\x65\x64uplicate\x12,\n\x05range\x18\x0f \x01(\x0b\x32\x14.spark.connect.RangeH\x00R\x05range\x12\x45\n\x0esubquery_alias\x18\x10 \x01(\x0b\x32\x1c.spark.connect.SubqueryAliasH\x00R\rsubqueryAlias\x12>\n\x0brepartition\x18\x11 \x01(\x0b\x32\x1a.spark.connect.RepartitionH\x00R\x0brepartition\x12*\n\x05to_df\x18\x12 \x01(\x0b\x32\x13.spark.connect.ToDFH\x00R\x04toDf\x12U\n\x14with_columns_renamed\x18\x13 \x01(\x0b\x32!.spark.connect.WithColumnsRenamedH\x00R\x12withColumnsRenamed\x12<\n\x0bshow_string\x18\x14 \x01(\x0b\x32\x19.spark.connect.ShowStringH\x00R\nshowString\x12)\n\x04\x64rop\x18\x15 \x01(\x0b\x32\x13.spark.connect.DropH\x00R\x04\x64rop\x12)\n\x04tail\x18\x16 \x01(\x0b\x32\x13.spark.connect.TailH\x00R\x04tail\x12?\n\x0cwith_columns\x18\x17 \x01(\x0b\x32\x1a.spark.connect.WithColumnsH\x00R\x0bwithColumns\x12)\n\x04hint\x18\x18 \x01(\x0b\x32\x13.spark.connect.HintH\x00R\x04hint\x12\x32\n\x07unpivot\x18\x19 \x01(\x0b\x32\x16.spark.connect.UnpivotH\x00R\x07unpivot\x12\x36\n\tto_schema\x18\x1a \x01(\x0b\x32\x17.spark.connect.ToSchemaH\x00R\x08toSchema\x12\x64\n\x19repartition_by_expression\x18\x1b \x01(\x0b\x32&.spark.connect.RepartitionByExpressionH\x00R\x17repartitionByExpression\x12\x45\n\x0emap_partitions\x18\x1c \x01(\x0b\x32\x1c.spark.connect.MapPartitionsH\x00R\rmapPartitions\x12H\n\x0f\x63ollect_metrics\x18\x1d \x01(\x0b\x32\x1d.spark.connect.CollectMetricsH\x00R\x0e\x63ollectMetrics\x12,\n\x05parse\x18\x1e \x01(\x0b\x32\x14.spark.connect.ParseH\x00R\x05parse\x12\x36\n\tgroup_map\x18\x1f \x01(\x0b\x32\x17.spark.connect.GroupMapH\x00R\x08groupMap\x12=\n\x0c\x63o_group_map\x18 \x01(\x0b\x32\x19.spark.connect.CoGroupMapH\x00R\ncoGroupMap\x12\x45\n\x0ewith_watermark\x18! \x01(\x0b\x32\x1c.spark.connect.WithWatermarkH\x00R\rwithWatermark\x12\x63\n\x1a\x61pply_in_pandas_with_state\x18" \x01(\x0b\x32%.spark.connect.ApplyInPandasWithStateH\x00R\x16\x61pplyInPandasWithState\x12<\n\x0bhtml_string\x18# \x01(\x0b\x32\x19.spark.connect.HtmlStringH\x00R\nhtmlString\x12X\n\x15\x63\x61\x63hed_local_relation\x18$ \x01(\x0b\x32".spark.connect.CachedLocalRelationH\x00R\x13\x63\x61\x63hedLocalRelation\x12[\n\x16\x63\x61\x63hed_remote_relation\x18% \x01(\x0b\x32#.spark.connect.CachedRemoteRelationH\x00R\x14\x63\x61\x63hedRemoteRelation\x12\x8e\x01\n)common_inline_user_defined_table_function\x18& \x01(\x0b\x32\x33.spark.connect.CommonInlineUserDefinedTableFunctionH\x00R$commonInlineUserDefinedTableFunction\x12\x37\n\nas_of_join\x18\' \x01(\x0b\x32\x17.spark.connect.AsOfJoinH\x00R\x08\x61sOfJoin\x12\x85\x01\n&common_inline_user_defined_data_source\x18( \x01(\x0b\x32\x30.spark.connect.CommonInlineUserDefinedDataSourceH\x00R!commonInlineUserDefinedDataSource\x12\x45\n\x0ewith_relations\x18) \x01(\x0b\x32\x1c.spark.connect.WithRelationsH\x00R\rwithRelations\x12\x38\n\ttranspose\x18* \x01(\x0b\x32\x18.spark.connect.TransposeH\x00R\ttranspose\x12w\n unresolved_table_valued_function\x18+ \x01(\x0b\x32,.spark.connect.UnresolvedTableValuedFunctionH\x00R\x1dunresolvedTableValuedFunction\x12?\n\x0clateral_join\x18, \x01(\x0b\x32\x1a.spark.connect.LateralJoinH\x00R\x0blateralJoin\x12n\n\x1d\x63hunked_cached_local_relation\x18- \x01(\x0b\x32).spark.connect.ChunkedCachedLocalRelationH\x00R\x1a\x63hunkedCachedLocalRelation\x12K\n\x10relation_changes\x18. \x01(\x0b\x32\x1e.spark.connect.RelationChangesH\x00R\x0frelationChanges\x12\x46\n\x0fnearest_by_join\x18/ \x01(\x0b\x32\x1c.spark.connect.NearestByJoinH\x00R\rnearestByJoin\x12&\n\x03zip\x18\x30 \x01(\x0b\x32\x12.spark.connect.ZipH\x00R\x03zip\x12N\n\x0cgraph_frames\x18\x31 \x01(\x0b\x32).spark.connect.graphframes.GraphFramesAPIH\x00R\x0bgraphFrames\x12\x30\n\x07\x66ill_na\x18Z \x01(\x0b\x32\x15.spark.connect.NAFillH\x00R\x06\x66illNa\x12\x30\n\x07\x64rop_na\x18[ \x01(\x0b\x32\x15.spark.connect.NADropH\x00R\x06\x64ropNa\x12\x34\n\x07replace\x18\\ \x01(\x0b\x32\x18.spark.connect.NAReplaceH\x00R\x07replace\x12\x36\n\x07summary\x18\x64 \x01(\x0b\x32\x1a.spark.connect.StatSummaryH\x00R\x07summary\x12\x39\n\x08\x63rosstab\x18\x65 \x01(\x0b\x32\x1b.spark.connect.StatCrosstabH\x00R\x08\x63rosstab\x12\x39\n\x08\x64\x65scribe\x18\x66 \x01(\x0b\x32\x1b.spark.connect.StatDescribeH\x00R\x08\x64\x65scribe\x12*\n\x03\x63ov\x18g \x01(\x0b\x32\x16.spark.connect.StatCovH\x00R\x03\x63ov\x12-\n\x04\x63orr\x18h \x01(\x0b\x32\x17.spark.connect.StatCorrH\x00R\x04\x63orr\x12L\n\x0f\x61pprox_quantile\x18i \x01(\x0b\x32!.spark.connect.StatApproxQuantileH\x00R\x0e\x61pproxQuantile\x12=\n\nfreq_items\x18j \x01(\x0b\x32\x1c.spark.connect.StatFreqItemsH\x00R\tfreqItems\x12:\n\tsample_by\x18k \x01(\x0b\x32\x1b.spark.connect.StatSampleByH\x00R\x08sampleBy\x12\x33\n\x07\x63\x61talog\x18\xc8\x01 \x01(\x0b\x32\x16.spark.connect.CatalogH\x00R\x07\x63\x61talog\x12=\n\x0bml_relation\x18\xac\x02 \x01(\x0b\x32\x19.spark.connect.MlRelationH\x00R\nmlRelation\x12\x35\n\textension\x18\xe6\x07 \x01(\x0b\x32\x14.google.protobuf.AnyH\x00R\textension\x12\x33\n\x07unknown\x18\xe7\x07 \x01(\x0b\x32\x16.spark.connect.UnknownH\x00R\x07unknownB\n\n\x08rel_type"\xe4\x03\n\nMlRelation\x12\x43\n\ttransform\x18\x01 \x01(\x0b\x32#.spark.connect.MlRelation.TransformH\x00R\ttransform\x12,\n\x05\x66\x65tch\x18\x02 \x01(\x0b\x32\x14.spark.connect.FetchH\x00R\x05\x66\x65tch\x12P\n\x15model_summary_dataset\x18\x03 \x01(\x0b\x32\x17.spark.connect.RelationH\x01R\x13modelSummaryDataset\x88\x01\x01\x1a\xeb\x01\n\tTransform\x12\x33\n\x07obj_ref\x18\x01 \x01(\x0b\x32\x18.spark.connect.ObjectRefH\x00R\x06objRef\x12=\n\x0btransformer\x18\x02 \x01(\x0b\x32\x19.spark.connect.MlOperatorH\x00R\x0btransformer\x12-\n\x05input\x18\x03 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12/\n\x06params\x18\x04 \x01(\x0b\x32\x17.spark.connect.MlParamsR\x06paramsB\n\n\x08operatorB\t\n\x07ml_typeB\x18\n\x16_model_summary_dataset"\xcb\x02\n\x05\x46\x65tch\x12\x31\n\x07obj_ref\x18\x01 \x01(\x0b\x32\x18.spark.connect.ObjectRefR\x06objRef\x12\x35\n\x07methods\x18\x02 \x03(\x0b\x32\x1b.spark.connect.Fetch.MethodR\x07methods\x1a\xd7\x01\n\x06Method\x12\x16\n\x06method\x18\x01 \x01(\tR\x06method\x12\x34\n\x04\x61rgs\x18\x02 \x03(\x0b\x32 .spark.connect.Fetch.Method.ArgsR\x04\x61rgs\x1a\x7f\n\x04\x41rgs\x12\x39\n\x05param\x18\x01 \x01(\x0b\x32!.spark.connect.Expression.LiteralH\x00R\x05param\x12/\n\x05input\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationH\x00R\x05inputB\x0b\n\targs_type"\t\n\x07Unknown"\x8e\x01\n\x0eRelationCommon\x12#\n\x0bsource_info\x18\x01 \x01(\tB\x02\x18\x01R\nsourceInfo\x12\x1c\n\x07plan_id\x18\x02 \x01(\x03H\x00R\x06planId\x88\x01\x01\x12-\n\x06origin\x18\x03 \x01(\x0b\x32\x15.spark.connect.OriginR\x06originB\n\n\x08_plan_id"\xde\x03\n\x03SQL\x12\x14\n\x05query\x18\x01 \x01(\tR\x05query\x12\x34\n\x04\x61rgs\x18\x02 \x03(\x0b\x32\x1c.spark.connect.SQL.ArgsEntryB\x02\x18\x01R\x04\x61rgs\x12@\n\x08pos_args\x18\x03 \x03(\x0b\x32!.spark.connect.Expression.LiteralB\x02\x18\x01R\x07posArgs\x12O\n\x0fnamed_arguments\x18\x04 \x03(\x0b\x32&.spark.connect.SQL.NamedArgumentsEntryR\x0enamedArguments\x12>\n\rpos_arguments\x18\x05 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x0cposArguments\x1aZ\n\tArgsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x37\n\x05value\x18\x02 \x01(\x0b\x32!.spark.connect.Expression.LiteralR\x05value:\x02\x38\x01\x1a\\\n\x13NamedArgumentsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12/\n\x05value\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x05value:\x02\x38\x01"u\n\rWithRelations\x12+\n\x04root\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x04root\x12\x37\n\nreferences\x18\x02 \x03(\x0b\x32\x17.spark.connect.RelationR\nreferences"\xcd\x05\n\x04Read\x12\x41\n\x0bnamed_table\x18\x01 \x01(\x0b\x32\x1e.spark.connect.Read.NamedTableH\x00R\nnamedTable\x12\x41\n\x0b\x64\x61ta_source\x18\x02 \x01(\x0b\x32\x1e.spark.connect.Read.DataSourceH\x00R\ndataSource\x12!\n\x0cis_streaming\x18\x03 \x01(\x08R\x0bisStreaming\x1a\xc0\x01\n\nNamedTable\x12/\n\x13unparsed_identifier\x18\x01 \x01(\tR\x12unparsedIdentifier\x12\x45\n\x07options\x18\x02 \x03(\x0b\x32+.spark.connect.Read.NamedTable.OptionsEntryR\x07options\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x1a\xcb\x02\n\nDataSource\x12\x1b\n\x06\x66ormat\x18\x01 \x01(\tH\x00R\x06\x66ormat\x88\x01\x01\x12\x1b\n\x06schema\x18\x02 \x01(\tH\x01R\x06schema\x88\x01\x01\x12\x45\n\x07options\x18\x03 \x03(\x0b\x32+.spark.connect.Read.DataSource.OptionsEntryR\x07options\x12\x14\n\x05paths\x18\x04 \x03(\tR\x05paths\x12\x1e\n\npredicates\x18\x05 \x03(\tR\npredicates\x12$\n\x0bsource_name\x18\x06 \x01(\tH\x02R\nsourceName\x88\x01\x01\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x42\t\n\x07_formatB\t\n\x07_schemaB\x0e\n\x0c_source_nameB\x0b\n\tread_type"\xe8\x01\n\x0fRelationChanges\x12/\n\x13unparsed_identifier\x18\x01 \x01(\tR\x12unparsedIdentifier\x12\x45\n\x07options\x18\x02 \x03(\x0b\x32+.spark.connect.RelationChanges.OptionsEntryR\x07options\x12!\n\x0cis_streaming\x18\x03 \x01(\x08R\x0bisStreaming\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01"u\n\x07Project\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12;\n\x0b\x65xpressions\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x0b\x65xpressions"p\n\x06\x46ilter\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x37\n\tcondition\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionR\tcondition"\x95\x05\n\x04Join\x12+\n\x04left\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x04left\x12-\n\x05right\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationR\x05right\x12@\n\x0ejoin_condition\x18\x03 \x01(\x0b\x32\x19.spark.connect.ExpressionR\rjoinCondition\x12\x39\n\tjoin_type\x18\x04 \x01(\x0e\x32\x1c.spark.connect.Join.JoinTypeR\x08joinType\x12#\n\rusing_columns\x18\x05 \x03(\tR\x0cusingColumns\x12K\n\x0ejoin_data_type\x18\x06 \x01(\x0b\x32 .spark.connect.Join.JoinDataTypeH\x00R\x0cjoinDataType\x88\x01\x01\x1a\\\n\x0cJoinDataType\x12$\n\x0eis_left_struct\x18\x01 \x01(\x08R\x0cisLeftStruct\x12&\n\x0fis_right_struct\x18\x02 \x01(\x08R\risRightStruct"\xd0\x01\n\x08JoinType\x12\x19\n\x15JOIN_TYPE_UNSPECIFIED\x10\x00\x12\x13\n\x0fJOIN_TYPE_INNER\x10\x01\x12\x18\n\x14JOIN_TYPE_FULL_OUTER\x10\x02\x12\x18\n\x14JOIN_TYPE_LEFT_OUTER\x10\x03\x12\x19\n\x15JOIN_TYPE_RIGHT_OUTER\x10\x04\x12\x17\n\x13JOIN_TYPE_LEFT_ANTI\x10\x05\x12\x17\n\x13JOIN_TYPE_LEFT_SEMI\x10\x06\x12\x13\n\x0fJOIN_TYPE_CROSS\x10\x07\x42\x11\n\x0f_join_data_type"\xdf\x03\n\x0cSetOperation\x12\x36\n\nleft_input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\tleftInput\x12\x38\n\x0bright_input\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationR\nrightInput\x12\x45\n\x0bset_op_type\x18\x03 \x01(\x0e\x32%.spark.connect.SetOperation.SetOpTypeR\tsetOpType\x12\x1a\n\x06is_all\x18\x04 \x01(\x08H\x00R\x05isAll\x88\x01\x01\x12\x1c\n\x07\x62y_name\x18\x05 \x01(\x08H\x01R\x06\x62yName\x88\x01\x01\x12\x37\n\x15\x61llow_missing_columns\x18\x06 \x01(\x08H\x02R\x13\x61llowMissingColumns\x88\x01\x01"r\n\tSetOpType\x12\x1b\n\x17SET_OP_TYPE_UNSPECIFIED\x10\x00\x12\x19\n\x15SET_OP_TYPE_INTERSECT\x10\x01\x12\x15\n\x11SET_OP_TYPE_UNION\x10\x02\x12\x16\n\x12SET_OP_TYPE_EXCEPT\x10\x03\x42\t\n\x07_is_allB\n\n\x08_by_nameB\x18\n\x16_allow_missing_columns"L\n\x05Limit\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x14\n\x05limit\x18\x02 \x01(\x05R\x05limit"O\n\x06Offset\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x16\n\x06offset\x18\x02 \x01(\x05R\x06offset"K\n\x04Tail\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x14\n\x05limit\x18\x02 \x01(\x05R\x05limit"\xfe\x05\n\tAggregate\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x41\n\ngroup_type\x18\x02 \x01(\x0e\x32".spark.connect.Aggregate.GroupTypeR\tgroupType\x12L\n\x14grouping_expressions\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x13groupingExpressions\x12N\n\x15\x61ggregate_expressions\x18\x04 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x14\x61ggregateExpressions\x12\x34\n\x05pivot\x18\x05 \x01(\x0b\x32\x1e.spark.connect.Aggregate.PivotR\x05pivot\x12J\n\rgrouping_sets\x18\x06 \x03(\x0b\x32%.spark.connect.Aggregate.GroupingSetsR\x0cgroupingSets\x1ao\n\x05Pivot\x12+\n\x03\x63ol\x18\x01 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x03\x63ol\x12\x39\n\x06values\x18\x02 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x06values\x1aL\n\x0cGroupingSets\x12<\n\x0cgrouping_set\x18\x01 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x0bgroupingSet"\x9f\x01\n\tGroupType\x12\x1a\n\x16GROUP_TYPE_UNSPECIFIED\x10\x00\x12\x16\n\x12GROUP_TYPE_GROUPBY\x10\x01\x12\x15\n\x11GROUP_TYPE_ROLLUP\x10\x02\x12\x13\n\x0fGROUP_TYPE_CUBE\x10\x03\x12\x14\n\x10GROUP_TYPE_PIVOT\x10\x04\x12\x1c\n\x18GROUP_TYPE_GROUPING_SETS\x10\x05"\xa0\x01\n\x04Sort\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x39\n\x05order\x18\x02 \x03(\x0b\x32#.spark.connect.Expression.SortOrderR\x05order\x12 \n\tis_global\x18\x03 \x01(\x08H\x00R\x08isGlobal\x88\x01\x01\x42\x0c\n\n_is_global"\x8d\x01\n\x04\x44rop\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x33\n\x07\x63olumns\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x07\x63olumns\x12!\n\x0c\x63olumn_names\x18\x03 \x03(\tR\x0b\x63olumnNames"\xf0\x01\n\x0b\x44\x65\x64uplicate\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12!\n\x0c\x63olumn_names\x18\x02 \x03(\tR\x0b\x63olumnNames\x12\x32\n\x13\x61ll_columns_as_keys\x18\x03 \x01(\x08H\x00R\x10\x61llColumnsAsKeys\x88\x01\x01\x12.\n\x10within_watermark\x18\x04 \x01(\x08H\x01R\x0fwithinWatermark\x88\x01\x01\x42\x16\n\x14_all_columns_as_keysB\x13\n\x11_within_watermark"Y\n\rLocalRelation\x12\x17\n\x04\x64\x61ta\x18\x01 \x01(\x0cH\x00R\x04\x64\x61ta\x88\x01\x01\x12\x1b\n\x06schema\x18\x02 \x01(\tH\x01R\x06schema\x88\x01\x01\x42\x07\n\x05_dataB\t\n\x07_schema"H\n\x13\x43\x61\x63hedLocalRelation\x12\x12\n\x04hash\x18\x03 \x01(\tR\x04hashJ\x04\x08\x01\x10\x02J\x04\x08\x02\x10\x03R\x06userIdR\tsessionId"p\n\x1a\x43hunkedCachedLocalRelation\x12\x1e\n\ndataHashes\x18\x01 \x03(\tR\ndataHashes\x12#\n\nschemaHash\x18\x02 \x01(\tH\x00R\nschemaHash\x88\x01\x01\x42\r\n\x0b_schemaHash"7\n\x14\x43\x61\x63hedRemoteRelation\x12\x1f\n\x0brelation_id\x18\x01 \x01(\tR\nrelationId"\x91\x02\n\x06Sample\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x1f\n\x0blower_bound\x18\x02 \x01(\x01R\nlowerBound\x12\x1f\n\x0bupper_bound\x18\x03 \x01(\x01R\nupperBound\x12.\n\x10with_replacement\x18\x04 \x01(\x08H\x00R\x0fwithReplacement\x88\x01\x01\x12\x17\n\x04seed\x18\x05 \x01(\x03H\x01R\x04seed\x88\x01\x01\x12/\n\x13\x64\x65terministic_order\x18\x06 \x01(\x08R\x12\x64\x65terministicOrderB\x13\n\x11_with_replacementB\x07\n\x05_seed"\x91\x01\n\x05Range\x12\x19\n\x05start\x18\x01 \x01(\x03H\x00R\x05start\x88\x01\x01\x12\x10\n\x03\x65nd\x18\x02 \x01(\x03R\x03\x65nd\x12\x12\n\x04step\x18\x03 \x01(\x03R\x04step\x12*\n\x0enum_partitions\x18\x04 \x01(\x05H\x01R\rnumPartitions\x88\x01\x01\x42\x08\n\x06_startB\x11\n\x0f_num_partitions"r\n\rSubqueryAlias\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x14\n\x05\x61lias\x18\x02 \x01(\tR\x05\x61lias\x12\x1c\n\tqualifier\x18\x03 \x03(\tR\tqualifier"\x8e\x01\n\x0bRepartition\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12%\n\x0enum_partitions\x18\x02 \x01(\x05R\rnumPartitions\x12\x1d\n\x07shuffle\x18\x03 \x01(\x08H\x00R\x07shuffle\x88\x01\x01\x42\n\n\x08_shuffle"\x8e\x01\n\nShowString\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x19\n\x08num_rows\x18\x02 \x01(\x05R\x07numRows\x12\x1a\n\x08truncate\x18\x03 \x01(\x05R\x08truncate\x12\x1a\n\x08vertical\x18\x04 \x01(\x08R\x08vertical"r\n\nHtmlString\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x19\n\x08num_rows\x18\x02 \x01(\x05R\x07numRows\x12\x1a\n\x08truncate\x18\x03 \x01(\x05R\x08truncate"\\\n\x0bStatSummary\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x1e\n\nstatistics\x18\x02 \x03(\tR\nstatistics"Q\n\x0cStatDescribe\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols"e\n\x0cStatCrosstab\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ol1\x18\x02 \x01(\tR\x04\x63ol1\x12\x12\n\x04\x63ol2\x18\x03 \x01(\tR\x04\x63ol2"`\n\x07StatCov\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ol1\x18\x02 \x01(\tR\x04\x63ol1\x12\x12\n\x04\x63ol2\x18\x03 \x01(\tR\x04\x63ol2"\x89\x01\n\x08StatCorr\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ol1\x18\x02 \x01(\tR\x04\x63ol1\x12\x12\n\x04\x63ol2\x18\x03 \x01(\tR\x04\x63ol2\x12\x1b\n\x06method\x18\x04 \x01(\tH\x00R\x06method\x88\x01\x01\x42\t\n\x07_method"\xa4\x01\n\x12StatApproxQuantile\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols\x12$\n\rprobabilities\x18\x03 \x03(\x01R\rprobabilities\x12%\n\x0erelative_error\x18\x04 \x01(\x01R\rrelativeError"}\n\rStatFreqItems\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols\x12\x1d\n\x07support\x18\x03 \x01(\x01H\x00R\x07support\x88\x01\x01\x42\n\n\x08_support"\xb5\x02\n\x0cStatSampleBy\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12+\n\x03\x63ol\x18\x02 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x03\x63ol\x12\x42\n\tfractions\x18\x03 \x03(\x0b\x32$.spark.connect.StatSampleBy.FractionR\tfractions\x12\x17\n\x04seed\x18\x05 \x01(\x03H\x00R\x04seed\x88\x01\x01\x1a\x63\n\x08\x46raction\x12;\n\x07stratum\x18\x01 \x01(\x0b\x32!.spark.connect.Expression.LiteralR\x07stratum\x12\x1a\n\x08\x66raction\x18\x02 \x01(\x01R\x08\x66ractionB\x07\n\x05_seed"\x86\x01\n\x06NAFill\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols\x12\x39\n\x06values\x18\x03 \x03(\x0b\x32!.spark.connect.Expression.LiteralR\x06values"\x86\x01\n\x06NADrop\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols\x12\'\n\rmin_non_nulls\x18\x03 \x01(\x05H\x00R\x0bminNonNulls\x88\x01\x01\x42\x10\n\x0e_min_non_nulls"\xa8\x02\n\tNAReplace\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04\x63ols\x18\x02 \x03(\tR\x04\x63ols\x12H\n\x0creplacements\x18\x03 \x03(\x0b\x32$.spark.connect.NAReplace.ReplacementR\x0creplacements\x1a\x8d\x01\n\x0bReplacement\x12>\n\told_value\x18\x01 \x01(\x0b\x32!.spark.connect.Expression.LiteralR\x08oldValue\x12>\n\tnew_value\x18\x02 \x01(\x0b\x32!.spark.connect.Expression.LiteralR\x08newValue"X\n\x04ToDF\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12!\n\x0c\x63olumn_names\x18\x02 \x03(\tR\x0b\x63olumnNames"\xfe\x02\n\x12WithColumnsRenamed\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12i\n\x12rename_columns_map\x18\x02 \x03(\x0b\x32\x37.spark.connect.WithColumnsRenamed.RenameColumnsMapEntryB\x02\x18\x01R\x10renameColumnsMap\x12\x42\n\x07renames\x18\x03 \x03(\x0b\x32(.spark.connect.WithColumnsRenamed.RenameR\x07renames\x1a\x43\n\x15RenameColumnsMapEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01\x1a\x45\n\x06Rename\x12\x19\n\x08\x63ol_name\x18\x01 \x01(\tR\x07\x63olName\x12 \n\x0cnew_col_name\x18\x02 \x01(\tR\nnewColName"w\n\x0bWithColumns\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x39\n\x07\x61liases\x18\x02 \x03(\x0b\x32\x1f.spark.connect.Expression.AliasR\x07\x61liases"\x86\x01\n\rWithWatermark\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x1d\n\nevent_time\x18\x02 \x01(\tR\teventTime\x12\'\n\x0f\x64\x65lay_threshold\x18\x03 \x01(\tR\x0e\x64\x65layThreshold"\x84\x01\n\x04Hint\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04name\x18\x02 \x01(\tR\x04name\x12\x39\n\nparameters\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\nparameters"\xc7\x02\n\x07Unpivot\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12+\n\x03ids\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x03ids\x12:\n\x06values\x18\x03 \x01(\x0b\x32\x1d.spark.connect.Unpivot.ValuesH\x00R\x06values\x88\x01\x01\x12\x30\n\x14variable_column_name\x18\x04 \x01(\tR\x12variableColumnName\x12*\n\x11value_column_name\x18\x05 \x01(\tR\x0fvalueColumnName\x1a;\n\x06Values\x12\x31\n\x06values\x18\x01 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x06valuesB\t\n\x07_values"z\n\tTranspose\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12>\n\rindex_columns\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x0cindexColumns"}\n\x1dUnresolvedTableValuedFunction\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12\x37\n\targuments\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\targuments"j\n\x08ToSchema\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12/\n\x06schema\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataTypeR\x06schema"\xcb\x01\n\x17RepartitionByExpression\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x42\n\x0fpartition_exprs\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x0epartitionExprs\x12*\n\x0enum_partitions\x18\x03 \x01(\x05H\x00R\rnumPartitions\x88\x01\x01\x42\x11\n\x0f_num_partitions"\xe8\x01\n\rMapPartitions\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x42\n\x04\x66unc\x18\x02 \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionR\x04\x66unc\x12"\n\nis_barrier\x18\x03 \x01(\x08H\x00R\tisBarrier\x88\x01\x01\x12"\n\nprofile_id\x18\x04 \x01(\x05H\x01R\tprofileId\x88\x01\x01\x42\r\n\x0b_is_barrierB\r\n\x0b_profile_id"\xd2\x06\n\x08GroupMap\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12L\n\x14grouping_expressions\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x13groupingExpressions\x12\x42\n\x04\x66unc\x18\x03 \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionR\x04\x66unc\x12J\n\x13sorting_expressions\x18\x04 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x12sortingExpressions\x12<\n\rinitial_input\x18\x05 \x01(\x0b\x32\x17.spark.connect.RelationR\x0cinitialInput\x12[\n\x1cinitial_grouping_expressions\x18\x06 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x1ainitialGroupingExpressions\x12;\n\x18is_map_groups_with_state\x18\x07 \x01(\x08H\x00R\x14isMapGroupsWithState\x88\x01\x01\x12$\n\x0boutput_mode\x18\x08 \x01(\tH\x01R\noutputMode\x88\x01\x01\x12&\n\x0ctimeout_conf\x18\t \x01(\tH\x02R\x0btimeoutConf\x88\x01\x01\x12?\n\x0cstate_schema\x18\n \x01(\x0b\x32\x17.spark.connect.DataTypeH\x03R\x0bstateSchema\x88\x01\x01\x12\x65\n\x19transform_with_state_info\x18\x0b \x01(\x0b\x32%.spark.connect.TransformWithStateInfoH\x04R\x16transformWithStateInfo\x88\x01\x01\x42\x1b\n\x19_is_map_groups_with_stateB\x0e\n\x0c_output_modeB\x0f\n\r_timeout_confB\x0f\n\r_state_schemaB\x1c\n\x1a_transform_with_state_info"\xdf\x01\n\x16TransformWithStateInfo\x12\x1b\n\ttime_mode\x18\x01 \x01(\tR\x08timeMode\x12\x38\n\x16\x65vent_time_column_name\x18\x02 \x01(\tH\x00R\x13\x65ventTimeColumnName\x88\x01\x01\x12\x41\n\routput_schema\x18\x03 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x01R\x0coutputSchema\x88\x01\x01\x42\x19\n\x17_event_time_column_nameB\x10\n\x0e_output_schema"\x8e\x04\n\nCoGroupMap\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12W\n\x1ainput_grouping_expressions\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x18inputGroupingExpressions\x12-\n\x05other\x18\x03 \x01(\x0b\x32\x17.spark.connect.RelationR\x05other\x12W\n\x1aother_grouping_expressions\x18\x04 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x18otherGroupingExpressions\x12\x42\n\x04\x66unc\x18\x05 \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionR\x04\x66unc\x12U\n\x19input_sorting_expressions\x18\x06 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x17inputSortingExpressions\x12U\n\x19other_sorting_expressions\x18\x07 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x17otherSortingExpressions"\xe5\x02\n\x16\x41pplyInPandasWithState\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12L\n\x14grouping_expressions\x18\x02 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x13groupingExpressions\x12\x42\n\x04\x66unc\x18\x03 \x01(\x0b\x32..spark.connect.CommonInlineUserDefinedFunctionR\x04\x66unc\x12#\n\routput_schema\x18\x04 \x01(\tR\x0coutputSchema\x12!\n\x0cstate_schema\x18\x05 \x01(\tR\x0bstateSchema\x12\x1f\n\x0boutput_mode\x18\x06 \x01(\tR\noutputMode\x12!\n\x0ctimeout_conf\x18\x07 \x01(\tR\x0btimeoutConf"\xf4\x01\n$CommonInlineUserDefinedTableFunction\x12#\n\rfunction_name\x18\x01 \x01(\tR\x0c\x66unctionName\x12$\n\rdeterministic\x18\x02 \x01(\x08R\rdeterministic\x12\x37\n\targuments\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\targuments\x12<\n\x0bpython_udtf\x18\x04 \x01(\x0b\x32\x19.spark.connect.PythonUDTFH\x00R\npythonUdtfB\n\n\x08\x66unction"\xb1\x01\n\nPythonUDTF\x12=\n\x0breturn_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\nreturnType\x88\x01\x01\x12\x1b\n\teval_type\x18\x02 \x01(\x05R\x08\x65valType\x12\x18\n\x07\x63ommand\x18\x03 \x01(\x0cR\x07\x63ommand\x12\x1d\n\npython_ver\x18\x04 \x01(\tR\tpythonVerB\x0e\n\x0c_return_type"\x97\x01\n!CommonInlineUserDefinedDataSource\x12\x12\n\x04name\x18\x01 \x01(\tR\x04name\x12O\n\x12python_data_source\x18\x02 \x01(\x0b\x32\x1f.spark.connect.PythonDataSourceH\x00R\x10pythonDataSourceB\r\n\x0b\x64\x61ta_source"K\n\x10PythonDataSource\x12\x18\n\x07\x63ommand\x18\x01 \x01(\x0cR\x07\x63ommand\x12\x1d\n\npython_ver\x18\x02 \x01(\tR\tpythonVer"\x88\x01\n\x0e\x43ollectMetrics\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x12\n\x04name\x18\x02 \x01(\tR\x04name\x12\x33\n\x07metrics\x18\x03 \x03(\x0b\x32\x19.spark.connect.ExpressionR\x07metrics"\x9a\x03\n\x05Parse\x12-\n\x05input\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x05input\x12\x38\n\x06\x66ormat\x18\x02 \x01(\x0e\x32 .spark.connect.Parse.ParseFormatR\x06\x66ormat\x12\x34\n\x06schema\x18\x03 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x00R\x06schema\x88\x01\x01\x12;\n\x07options\x18\x04 \x03(\x0b\x32!.spark.connect.Parse.OptionsEntryR\x07options\x1a:\n\x0cOptionsEntry\x12\x10\n\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n\x05value\x18\x02 \x01(\tR\x05value:\x02\x38\x01"n\n\x0bParseFormat\x12\x1c\n\x18PARSE_FORMAT_UNSPECIFIED\x10\x00\x12\x14\n\x10PARSE_FORMAT_CSV\x10\x01\x12\x15\n\x11PARSE_FORMAT_JSON\x10\x02\x12\x14\n\x10PARSE_FORMAT_XML\x10\x03\x42\t\n\x07_schema"\xdb\x03\n\x08\x41sOfJoin\x12+\n\x04left\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x04left\x12-\n\x05right\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationR\x05right\x12\x37\n\nleft_as_of\x18\x03 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x08leftAsOf\x12\x39\n\x0bright_as_of\x18\x04 \x01(\x0b\x32\x19.spark.connect.ExpressionR\trightAsOf\x12\x36\n\tjoin_expr\x18\x05 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x08joinExpr\x12#\n\rusing_columns\x18\x06 \x03(\tR\x0cusingColumns\x12\x1b\n\tjoin_type\x18\x07 \x01(\tR\x08joinType\x12\x37\n\ttolerance\x18\x08 \x01(\x0b\x32\x19.spark.connect.ExpressionR\ttolerance\x12.\n\x13\x61llow_exact_matches\x18\t \x01(\x08R\x11\x61llowExactMatches\x12\x1c\n\tdirection\x18\n \x01(\tR\tdirection"\xe6\x01\n\x0bLateralJoin\x12+\n\x04left\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x04left\x12-\n\x05right\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationR\x05right\x12@\n\x0ejoin_condition\x18\x03 \x01(\x0b\x32\x19.spark.connect.ExpressionR\rjoinCondition\x12\x39\n\tjoin_type\x18\x04 \x01(\x0e\x32\x1c.spark.connect.Join.JoinTypeR\x08joinType"\xa5\x02\n\rNearestByJoin\x12+\n\x04left\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x04left\x12-\n\x05right\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationR\x05right\x12H\n\x12ranking_expression\x18\x03 \x01(\x0b\x32\x19.spark.connect.ExpressionR\x11rankingExpression\x12\x1f\n\x0bnum_results\x18\x04 \x01(\x05R\nnumResults\x12\x1b\n\tjoin_type\x18\x05 \x01(\tR\x08joinType\x12\x12\n\x04mode\x18\x06 \x01(\tR\x04mode\x12\x1c\n\tdirection\x18\x07 \x01(\tR\tdirection"a\n\x03Zip\x12+\n\x04left\x18\x01 \x01(\x0b\x32\x17.spark.connect.RelationR\x04left\x12-\n\x05right\x18\x02 \x01(\x0b\x32\x17.spark.connect.RelationR\x05rightB6\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3' ) _globals = globals() @@ -81,182 +82,182 @@ ]._serialized_options = b"\030\001" _globals["_PARSE_OPTIONSENTRY"]._loaded_options = None _globals["_PARSE_OPTIONSENTRY"]._serialized_options = b"8\001" - _globals["_RELATION"]._serialized_start = 224 - _globals["_RELATION"]._serialized_end = 4265 - _globals["_MLRELATION"]._serialized_start = 4268 - _globals["_MLRELATION"]._serialized_end = 4752 - _globals["_MLRELATION_TRANSFORM"]._serialized_start = 4480 - _globals["_MLRELATION_TRANSFORM"]._serialized_end = 4715 - _globals["_FETCH"]._serialized_start = 4755 - _globals["_FETCH"]._serialized_end = 5086 - _globals["_FETCH_METHOD"]._serialized_start = 4871 - _globals["_FETCH_METHOD"]._serialized_end = 5086 - _globals["_FETCH_METHOD_ARGS"]._serialized_start = 4959 - _globals["_FETCH_METHOD_ARGS"]._serialized_end = 5086 - _globals["_UNKNOWN"]._serialized_start = 5088 - _globals["_UNKNOWN"]._serialized_end = 5097 - _globals["_RELATIONCOMMON"]._serialized_start = 5100 - _globals["_RELATIONCOMMON"]._serialized_end = 5242 - _globals["_SQL"]._serialized_start = 5245 - _globals["_SQL"]._serialized_end = 5723 - _globals["_SQL_ARGSENTRY"]._serialized_start = 5539 - _globals["_SQL_ARGSENTRY"]._serialized_end = 5629 - _globals["_SQL_NAMEDARGUMENTSENTRY"]._serialized_start = 5631 - _globals["_SQL_NAMEDARGUMENTSENTRY"]._serialized_end = 5723 - _globals["_WITHRELATIONS"]._serialized_start = 5725 - _globals["_WITHRELATIONS"]._serialized_end = 5842 - _globals["_READ"]._serialized_start = 5845 - _globals["_READ"]._serialized_end = 6562 - _globals["_READ_NAMEDTABLE"]._serialized_start = 6023 - _globals["_READ_NAMEDTABLE"]._serialized_end = 6215 - _globals["_READ_NAMEDTABLE_OPTIONSENTRY"]._serialized_start = 6157 - _globals["_READ_NAMEDTABLE_OPTIONSENTRY"]._serialized_end = 6215 - _globals["_READ_DATASOURCE"]._serialized_start = 6218 - _globals["_READ_DATASOURCE"]._serialized_end = 6549 - _globals["_READ_DATASOURCE_OPTIONSENTRY"]._serialized_start = 6157 - _globals["_READ_DATASOURCE_OPTIONSENTRY"]._serialized_end = 6215 - _globals["_RELATIONCHANGES"]._serialized_start = 6565 - _globals["_RELATIONCHANGES"]._serialized_end = 6797 - _globals["_RELATIONCHANGES_OPTIONSENTRY"]._serialized_start = 6157 - _globals["_RELATIONCHANGES_OPTIONSENTRY"]._serialized_end = 6215 - _globals["_PROJECT"]._serialized_start = 6799 - _globals["_PROJECT"]._serialized_end = 6916 - _globals["_FILTER"]._serialized_start = 6918 - _globals["_FILTER"]._serialized_end = 7030 - _globals["_JOIN"]._serialized_start = 7033 - _globals["_JOIN"]._serialized_end = 7694 - _globals["_JOIN_JOINDATATYPE"]._serialized_start = 7372 - _globals["_JOIN_JOINDATATYPE"]._serialized_end = 7464 - _globals["_JOIN_JOINTYPE"]._serialized_start = 7467 - _globals["_JOIN_JOINTYPE"]._serialized_end = 7675 - _globals["_SETOPERATION"]._serialized_start = 7697 - _globals["_SETOPERATION"]._serialized_end = 8176 - _globals["_SETOPERATION_SETOPTYPE"]._serialized_start = 8013 - _globals["_SETOPERATION_SETOPTYPE"]._serialized_end = 8127 - _globals["_LIMIT"]._serialized_start = 8178 - _globals["_LIMIT"]._serialized_end = 8254 - _globals["_OFFSET"]._serialized_start = 8256 - _globals["_OFFSET"]._serialized_end = 8335 - _globals["_TAIL"]._serialized_start = 8337 - _globals["_TAIL"]._serialized_end = 8412 - _globals["_AGGREGATE"]._serialized_start = 8415 - _globals["_AGGREGATE"]._serialized_end = 9181 - _globals["_AGGREGATE_PIVOT"]._serialized_start = 8830 - _globals["_AGGREGATE_PIVOT"]._serialized_end = 8941 - _globals["_AGGREGATE_GROUPINGSETS"]._serialized_start = 8943 - _globals["_AGGREGATE_GROUPINGSETS"]._serialized_end = 9019 - _globals["_AGGREGATE_GROUPTYPE"]._serialized_start = 9022 - _globals["_AGGREGATE_GROUPTYPE"]._serialized_end = 9181 - _globals["_SORT"]._serialized_start = 9184 - _globals["_SORT"]._serialized_end = 9344 - _globals["_DROP"]._serialized_start = 9347 - _globals["_DROP"]._serialized_end = 9488 - _globals["_DEDUPLICATE"]._serialized_start = 9491 - _globals["_DEDUPLICATE"]._serialized_end = 9731 - _globals["_LOCALRELATION"]._serialized_start = 9733 - _globals["_LOCALRELATION"]._serialized_end = 9822 - _globals["_CACHEDLOCALRELATION"]._serialized_start = 9824 - _globals["_CACHEDLOCALRELATION"]._serialized_end = 9896 - _globals["_CHUNKEDCACHEDLOCALRELATION"]._serialized_start = 9898 - _globals["_CHUNKEDCACHEDLOCALRELATION"]._serialized_end = 10010 - _globals["_CACHEDREMOTERELATION"]._serialized_start = 10012 - _globals["_CACHEDREMOTERELATION"]._serialized_end = 10067 - _globals["_SAMPLE"]._serialized_start = 10070 - _globals["_SAMPLE"]._serialized_end = 10343 - _globals["_RANGE"]._serialized_start = 10346 - _globals["_RANGE"]._serialized_end = 10491 - _globals["_SUBQUERYALIAS"]._serialized_start = 10493 - _globals["_SUBQUERYALIAS"]._serialized_end = 10607 - _globals["_REPARTITION"]._serialized_start = 10610 - _globals["_REPARTITION"]._serialized_end = 10752 - _globals["_SHOWSTRING"]._serialized_start = 10755 - _globals["_SHOWSTRING"]._serialized_end = 10897 - _globals["_HTMLSTRING"]._serialized_start = 10899 - _globals["_HTMLSTRING"]._serialized_end = 11013 - _globals["_STATSUMMARY"]._serialized_start = 11015 - _globals["_STATSUMMARY"]._serialized_end = 11107 - _globals["_STATDESCRIBE"]._serialized_start = 11109 - _globals["_STATDESCRIBE"]._serialized_end = 11190 - _globals["_STATCROSSTAB"]._serialized_start = 11192 - _globals["_STATCROSSTAB"]._serialized_end = 11293 - _globals["_STATCOV"]._serialized_start = 11295 - _globals["_STATCOV"]._serialized_end = 11391 - _globals["_STATCORR"]._serialized_start = 11394 - _globals["_STATCORR"]._serialized_end = 11531 - _globals["_STATAPPROXQUANTILE"]._serialized_start = 11534 - _globals["_STATAPPROXQUANTILE"]._serialized_end = 11698 - _globals["_STATFREQITEMS"]._serialized_start = 11700 - _globals["_STATFREQITEMS"]._serialized_end = 11825 - _globals["_STATSAMPLEBY"]._serialized_start = 11828 - _globals["_STATSAMPLEBY"]._serialized_end = 12137 - _globals["_STATSAMPLEBY_FRACTION"]._serialized_start = 12029 - _globals["_STATSAMPLEBY_FRACTION"]._serialized_end = 12128 - _globals["_NAFILL"]._serialized_start = 12140 - _globals["_NAFILL"]._serialized_end = 12274 - _globals["_NADROP"]._serialized_start = 12277 - _globals["_NADROP"]._serialized_end = 12411 - _globals["_NAREPLACE"]._serialized_start = 12414 - _globals["_NAREPLACE"]._serialized_end = 12710 - _globals["_NAREPLACE_REPLACEMENT"]._serialized_start = 12569 - _globals["_NAREPLACE_REPLACEMENT"]._serialized_end = 12710 - _globals["_TODF"]._serialized_start = 12712 - _globals["_TODF"]._serialized_end = 12800 - _globals["_WITHCOLUMNSRENAMED"]._serialized_start = 12803 - _globals["_WITHCOLUMNSRENAMED"]._serialized_end = 13185 - _globals["_WITHCOLUMNSRENAMED_RENAMECOLUMNSMAPENTRY"]._serialized_start = 13047 - _globals["_WITHCOLUMNSRENAMED_RENAMECOLUMNSMAPENTRY"]._serialized_end = 13114 - _globals["_WITHCOLUMNSRENAMED_RENAME"]._serialized_start = 13116 - _globals["_WITHCOLUMNSRENAMED_RENAME"]._serialized_end = 13185 - _globals["_WITHCOLUMNS"]._serialized_start = 13187 - _globals["_WITHCOLUMNS"]._serialized_end = 13306 - _globals["_WITHWATERMARK"]._serialized_start = 13309 - _globals["_WITHWATERMARK"]._serialized_end = 13443 - _globals["_HINT"]._serialized_start = 13446 - _globals["_HINT"]._serialized_end = 13578 - _globals["_UNPIVOT"]._serialized_start = 13581 - _globals["_UNPIVOT"]._serialized_end = 13908 - _globals["_UNPIVOT_VALUES"]._serialized_start = 13838 - _globals["_UNPIVOT_VALUES"]._serialized_end = 13897 - _globals["_TRANSPOSE"]._serialized_start = 13910 - _globals["_TRANSPOSE"]._serialized_end = 14032 - _globals["_UNRESOLVEDTABLEVALUEDFUNCTION"]._serialized_start = 14034 - _globals["_UNRESOLVEDTABLEVALUEDFUNCTION"]._serialized_end = 14159 - _globals["_TOSCHEMA"]._serialized_start = 14161 - _globals["_TOSCHEMA"]._serialized_end = 14267 - _globals["_REPARTITIONBYEXPRESSION"]._serialized_start = 14270 - _globals["_REPARTITIONBYEXPRESSION"]._serialized_end = 14473 - _globals["_MAPPARTITIONS"]._serialized_start = 14476 - _globals["_MAPPARTITIONS"]._serialized_end = 14708 - _globals["_GROUPMAP"]._serialized_start = 14711 - _globals["_GROUPMAP"]._serialized_end = 15561 - _globals["_TRANSFORMWITHSTATEINFO"]._serialized_start = 15564 - _globals["_TRANSFORMWITHSTATEINFO"]._serialized_end = 15787 - _globals["_COGROUPMAP"]._serialized_start = 15790 - _globals["_COGROUPMAP"]._serialized_end = 16316 - _globals["_APPLYINPANDASWITHSTATE"]._serialized_start = 16319 - _globals["_APPLYINPANDASWITHSTATE"]._serialized_end = 16676 - _globals["_COMMONINLINEUSERDEFINEDTABLEFUNCTION"]._serialized_start = 16679 - _globals["_COMMONINLINEUSERDEFINEDTABLEFUNCTION"]._serialized_end = 16923 - _globals["_PYTHONUDTF"]._serialized_start = 16926 - _globals["_PYTHONUDTF"]._serialized_end = 17103 - _globals["_COMMONINLINEUSERDEFINEDDATASOURCE"]._serialized_start = 17106 - _globals["_COMMONINLINEUSERDEFINEDDATASOURCE"]._serialized_end = 17257 - _globals["_PYTHONDATASOURCE"]._serialized_start = 17259 - _globals["_PYTHONDATASOURCE"]._serialized_end = 17334 - _globals["_COLLECTMETRICS"]._serialized_start = 17337 - _globals["_COLLECTMETRICS"]._serialized_end = 17473 - _globals["_PARSE"]._serialized_start = 17476 - _globals["_PARSE"]._serialized_end = 17886 - _globals["_PARSE_OPTIONSENTRY"]._serialized_start = 6157 - _globals["_PARSE_OPTIONSENTRY"]._serialized_end = 6215 - _globals["_PARSE_PARSEFORMAT"]._serialized_start = 17765 - _globals["_PARSE_PARSEFORMAT"]._serialized_end = 17875 - _globals["_ASOFJOIN"]._serialized_start = 17889 - _globals["_ASOFJOIN"]._serialized_end = 18364 - _globals["_LATERALJOIN"]._serialized_start = 18367 - _globals["_LATERALJOIN"]._serialized_end = 18597 - _globals["_NEARESTBYJOIN"]._serialized_start = 18600 - _globals["_NEARESTBYJOIN"]._serialized_end = 18893 - _globals["_ZIP"]._serialized_start = 18895 - _globals["_ZIP"]._serialized_end = 18992 + _globals["_RELATION"]._serialized_start = 257 + _globals["_RELATION"]._serialized_end = 4378 + _globals["_MLRELATION"]._serialized_start = 4381 + _globals["_MLRELATION"]._serialized_end = 4865 + _globals["_MLRELATION_TRANSFORM"]._serialized_start = 4593 + _globals["_MLRELATION_TRANSFORM"]._serialized_end = 4828 + _globals["_FETCH"]._serialized_start = 4868 + _globals["_FETCH"]._serialized_end = 5199 + _globals["_FETCH_METHOD"]._serialized_start = 4984 + _globals["_FETCH_METHOD"]._serialized_end = 5199 + _globals["_FETCH_METHOD_ARGS"]._serialized_start = 5072 + _globals["_FETCH_METHOD_ARGS"]._serialized_end = 5199 + _globals["_UNKNOWN"]._serialized_start = 5201 + _globals["_UNKNOWN"]._serialized_end = 5210 + _globals["_RELATIONCOMMON"]._serialized_start = 5213 + _globals["_RELATIONCOMMON"]._serialized_end = 5355 + _globals["_SQL"]._serialized_start = 5358 + _globals["_SQL"]._serialized_end = 5836 + _globals["_SQL_ARGSENTRY"]._serialized_start = 5652 + _globals["_SQL_ARGSENTRY"]._serialized_end = 5742 + _globals["_SQL_NAMEDARGUMENTSENTRY"]._serialized_start = 5744 + _globals["_SQL_NAMEDARGUMENTSENTRY"]._serialized_end = 5836 + _globals["_WITHRELATIONS"]._serialized_start = 5838 + _globals["_WITHRELATIONS"]._serialized_end = 5955 + _globals["_READ"]._serialized_start = 5958 + _globals["_READ"]._serialized_end = 6675 + _globals["_READ_NAMEDTABLE"]._serialized_start = 6136 + _globals["_READ_NAMEDTABLE"]._serialized_end = 6328 + _globals["_READ_NAMEDTABLE_OPTIONSENTRY"]._serialized_start = 6270 + _globals["_READ_NAMEDTABLE_OPTIONSENTRY"]._serialized_end = 6328 + _globals["_READ_DATASOURCE"]._serialized_start = 6331 + _globals["_READ_DATASOURCE"]._serialized_end = 6662 + _globals["_READ_DATASOURCE_OPTIONSENTRY"]._serialized_start = 6270 + _globals["_READ_DATASOURCE_OPTIONSENTRY"]._serialized_end = 6328 + _globals["_RELATIONCHANGES"]._serialized_start = 6678 + _globals["_RELATIONCHANGES"]._serialized_end = 6910 + _globals["_RELATIONCHANGES_OPTIONSENTRY"]._serialized_start = 6270 + _globals["_RELATIONCHANGES_OPTIONSENTRY"]._serialized_end = 6328 + _globals["_PROJECT"]._serialized_start = 6912 + _globals["_PROJECT"]._serialized_end = 7029 + _globals["_FILTER"]._serialized_start = 7031 + _globals["_FILTER"]._serialized_end = 7143 + _globals["_JOIN"]._serialized_start = 7146 + _globals["_JOIN"]._serialized_end = 7807 + _globals["_JOIN_JOINDATATYPE"]._serialized_start = 7485 + _globals["_JOIN_JOINDATATYPE"]._serialized_end = 7577 + _globals["_JOIN_JOINTYPE"]._serialized_start = 7580 + _globals["_JOIN_JOINTYPE"]._serialized_end = 7788 + _globals["_SETOPERATION"]._serialized_start = 7810 + _globals["_SETOPERATION"]._serialized_end = 8289 + _globals["_SETOPERATION_SETOPTYPE"]._serialized_start = 8126 + _globals["_SETOPERATION_SETOPTYPE"]._serialized_end = 8240 + _globals["_LIMIT"]._serialized_start = 8291 + _globals["_LIMIT"]._serialized_end = 8367 + _globals["_OFFSET"]._serialized_start = 8369 + _globals["_OFFSET"]._serialized_end = 8448 + _globals["_TAIL"]._serialized_start = 8450 + _globals["_TAIL"]._serialized_end = 8525 + _globals["_AGGREGATE"]._serialized_start = 8528 + _globals["_AGGREGATE"]._serialized_end = 9294 + _globals["_AGGREGATE_PIVOT"]._serialized_start = 8943 + _globals["_AGGREGATE_PIVOT"]._serialized_end = 9054 + _globals["_AGGREGATE_GROUPINGSETS"]._serialized_start = 9056 + _globals["_AGGREGATE_GROUPINGSETS"]._serialized_end = 9132 + _globals["_AGGREGATE_GROUPTYPE"]._serialized_start = 9135 + _globals["_AGGREGATE_GROUPTYPE"]._serialized_end = 9294 + _globals["_SORT"]._serialized_start = 9297 + _globals["_SORT"]._serialized_end = 9457 + _globals["_DROP"]._serialized_start = 9460 + _globals["_DROP"]._serialized_end = 9601 + _globals["_DEDUPLICATE"]._serialized_start = 9604 + _globals["_DEDUPLICATE"]._serialized_end = 9844 + _globals["_LOCALRELATION"]._serialized_start = 9846 + _globals["_LOCALRELATION"]._serialized_end = 9935 + _globals["_CACHEDLOCALRELATION"]._serialized_start = 9937 + _globals["_CACHEDLOCALRELATION"]._serialized_end = 10009 + _globals["_CHUNKEDCACHEDLOCALRELATION"]._serialized_start = 10011 + _globals["_CHUNKEDCACHEDLOCALRELATION"]._serialized_end = 10123 + _globals["_CACHEDREMOTERELATION"]._serialized_start = 10125 + _globals["_CACHEDREMOTERELATION"]._serialized_end = 10180 + _globals["_SAMPLE"]._serialized_start = 10183 + _globals["_SAMPLE"]._serialized_end = 10456 + _globals["_RANGE"]._serialized_start = 10459 + _globals["_RANGE"]._serialized_end = 10604 + _globals["_SUBQUERYALIAS"]._serialized_start = 10606 + _globals["_SUBQUERYALIAS"]._serialized_end = 10720 + _globals["_REPARTITION"]._serialized_start = 10723 + _globals["_REPARTITION"]._serialized_end = 10865 + _globals["_SHOWSTRING"]._serialized_start = 10868 + _globals["_SHOWSTRING"]._serialized_end = 11010 + _globals["_HTMLSTRING"]._serialized_start = 11012 + _globals["_HTMLSTRING"]._serialized_end = 11126 + _globals["_STATSUMMARY"]._serialized_start = 11128 + _globals["_STATSUMMARY"]._serialized_end = 11220 + _globals["_STATDESCRIBE"]._serialized_start = 11222 + _globals["_STATDESCRIBE"]._serialized_end = 11303 + _globals["_STATCROSSTAB"]._serialized_start = 11305 + _globals["_STATCROSSTAB"]._serialized_end = 11406 + _globals["_STATCOV"]._serialized_start = 11408 + _globals["_STATCOV"]._serialized_end = 11504 + _globals["_STATCORR"]._serialized_start = 11507 + _globals["_STATCORR"]._serialized_end = 11644 + _globals["_STATAPPROXQUANTILE"]._serialized_start = 11647 + _globals["_STATAPPROXQUANTILE"]._serialized_end = 11811 + _globals["_STATFREQITEMS"]._serialized_start = 11813 + _globals["_STATFREQITEMS"]._serialized_end = 11938 + _globals["_STATSAMPLEBY"]._serialized_start = 11941 + _globals["_STATSAMPLEBY"]._serialized_end = 12250 + _globals["_STATSAMPLEBY_FRACTION"]._serialized_start = 12142 + _globals["_STATSAMPLEBY_FRACTION"]._serialized_end = 12241 + _globals["_NAFILL"]._serialized_start = 12253 + _globals["_NAFILL"]._serialized_end = 12387 + _globals["_NADROP"]._serialized_start = 12390 + _globals["_NADROP"]._serialized_end = 12524 + _globals["_NAREPLACE"]._serialized_start = 12527 + _globals["_NAREPLACE"]._serialized_end = 12823 + _globals["_NAREPLACE_REPLACEMENT"]._serialized_start = 12682 + _globals["_NAREPLACE_REPLACEMENT"]._serialized_end = 12823 + _globals["_TODF"]._serialized_start = 12825 + _globals["_TODF"]._serialized_end = 12913 + _globals["_WITHCOLUMNSRENAMED"]._serialized_start = 12916 + _globals["_WITHCOLUMNSRENAMED"]._serialized_end = 13298 + _globals["_WITHCOLUMNSRENAMED_RENAMECOLUMNSMAPENTRY"]._serialized_start = 13160 + _globals["_WITHCOLUMNSRENAMED_RENAMECOLUMNSMAPENTRY"]._serialized_end = 13227 + _globals["_WITHCOLUMNSRENAMED_RENAME"]._serialized_start = 13229 + _globals["_WITHCOLUMNSRENAMED_RENAME"]._serialized_end = 13298 + _globals["_WITHCOLUMNS"]._serialized_start = 13300 + _globals["_WITHCOLUMNS"]._serialized_end = 13419 + _globals["_WITHWATERMARK"]._serialized_start = 13422 + _globals["_WITHWATERMARK"]._serialized_end = 13556 + _globals["_HINT"]._serialized_start = 13559 + _globals["_HINT"]._serialized_end = 13691 + _globals["_UNPIVOT"]._serialized_start = 13694 + _globals["_UNPIVOT"]._serialized_end = 14021 + _globals["_UNPIVOT_VALUES"]._serialized_start = 13951 + _globals["_UNPIVOT_VALUES"]._serialized_end = 14010 + _globals["_TRANSPOSE"]._serialized_start = 14023 + _globals["_TRANSPOSE"]._serialized_end = 14145 + _globals["_UNRESOLVEDTABLEVALUEDFUNCTION"]._serialized_start = 14147 + _globals["_UNRESOLVEDTABLEVALUEDFUNCTION"]._serialized_end = 14272 + _globals["_TOSCHEMA"]._serialized_start = 14274 + _globals["_TOSCHEMA"]._serialized_end = 14380 + _globals["_REPARTITIONBYEXPRESSION"]._serialized_start = 14383 + _globals["_REPARTITIONBYEXPRESSION"]._serialized_end = 14586 + _globals["_MAPPARTITIONS"]._serialized_start = 14589 + _globals["_MAPPARTITIONS"]._serialized_end = 14821 + _globals["_GROUPMAP"]._serialized_start = 14824 + _globals["_GROUPMAP"]._serialized_end = 15674 + _globals["_TRANSFORMWITHSTATEINFO"]._serialized_start = 15677 + _globals["_TRANSFORMWITHSTATEINFO"]._serialized_end = 15900 + _globals["_COGROUPMAP"]._serialized_start = 15903 + _globals["_COGROUPMAP"]._serialized_end = 16429 + _globals["_APPLYINPANDASWITHSTATE"]._serialized_start = 16432 + _globals["_APPLYINPANDASWITHSTATE"]._serialized_end = 16789 + _globals["_COMMONINLINEUSERDEFINEDTABLEFUNCTION"]._serialized_start = 16792 + _globals["_COMMONINLINEUSERDEFINEDTABLEFUNCTION"]._serialized_end = 17036 + _globals["_PYTHONUDTF"]._serialized_start = 17039 + _globals["_PYTHONUDTF"]._serialized_end = 17216 + _globals["_COMMONINLINEUSERDEFINEDDATASOURCE"]._serialized_start = 17219 + _globals["_COMMONINLINEUSERDEFINEDDATASOURCE"]._serialized_end = 17370 + _globals["_PYTHONDATASOURCE"]._serialized_start = 17372 + _globals["_PYTHONDATASOURCE"]._serialized_end = 17447 + _globals["_COLLECTMETRICS"]._serialized_start = 17450 + _globals["_COLLECTMETRICS"]._serialized_end = 17586 + _globals["_PARSE"]._serialized_start = 17589 + _globals["_PARSE"]._serialized_end = 17999 + _globals["_PARSE_OPTIONSENTRY"]._serialized_start = 6270 + _globals["_PARSE_OPTIONSENTRY"]._serialized_end = 6328 + _globals["_PARSE_PARSEFORMAT"]._serialized_start = 17878 + _globals["_PARSE_PARSEFORMAT"]._serialized_end = 17988 + _globals["_ASOFJOIN"]._serialized_start = 18002 + _globals["_ASOFJOIN"]._serialized_end = 18477 + _globals["_LATERALJOIN"]._serialized_start = 18480 + _globals["_LATERALJOIN"]._serialized_end = 18710 + _globals["_NEARESTBYJOIN"]._serialized_start = 18713 + _globals["_NEARESTBYJOIN"]._serialized_end = 19006 + _globals["_ZIP"]._serialized_start = 19008 + _globals["_ZIP"]._serialized_end = 19105 # @@protoc_insertion_point(module_scope) diff --git a/python/pyspark/sql/connect/proto/relations_pb2.pyi b/python/pyspark/sql/connect/proto/relations_pb2.pyi index 2d17e88446d60..352d5bad68e53 100644 --- a/python/pyspark/sql/connect/proto/relations_pb2.pyi +++ b/python/pyspark/sql/connect/proto/relations_pb2.pyi @@ -44,6 +44,7 @@ import google.protobuf.message import pyspark.sql.connect.proto.catalog_pb2 import pyspark.sql.connect.proto.common_pb2 import pyspark.sql.connect.proto.expressions_pb2 +import pyspark.sql.connect.proto.graphframes_pb2 import pyspark.sql.connect.proto.ml_common_pb2 import pyspark.sql.connect.proto.types_pb2 import sys @@ -113,6 +114,7 @@ class Relation(google.protobuf.message.Message): RELATION_CHANGES_FIELD_NUMBER: builtins.int NEAREST_BY_JOIN_FIELD_NUMBER: builtins.int ZIP_FIELD_NUMBER: builtins.int + GRAPH_FRAMES_FIELD_NUMBER: builtins.int FILL_NA_FIELD_NUMBER: builtins.int DROP_NA_FIELD_NUMBER: builtins.int REPLACE_FIELD_NUMBER: builtins.int @@ -229,6 +231,8 @@ class Relation(google.protobuf.message.Message): @property def zip(self) -> global___Zip: ... @property + def graph_frames(self) -> pyspark.sql.connect.proto.graphframes_pb2.GraphFramesAPI: ... + @property def fill_na(self) -> global___NAFill: """NA functions""" @property @@ -318,6 +322,7 @@ class Relation(google.protobuf.message.Message): relation_changes: global___RelationChanges | None = ..., nearest_by_join: global___NearestByJoin | None = ..., zip: global___Zip | None = ..., + graph_frames: pyspark.sql.connect.proto.graphframes_pb2.GraphFramesAPI | None = ..., fill_na: global___NAFill | None = ..., drop_na: global___NADrop | None = ..., replace: global___NAReplace | None = ..., @@ -385,6 +390,8 @@ class Relation(google.protobuf.message.Message): b"filter", "freq_items", b"freq_items", + "graph_frames", + b"graph_frames", "group_map", b"group_map", "hint", @@ -518,6 +525,8 @@ class Relation(google.protobuf.message.Message): b"filter", "freq_items", b"freq_items", + "graph_frames", + b"graph_frames", "group_map", b"group_map", "hint", @@ -651,6 +660,7 @@ class Relation(google.protobuf.message.Message): "relation_changes", "nearest_by_join", "zip", + "graph_frames", "fill_na", "drop_na", "replace", diff --git a/sql/connect/common/src/main/protobuf/spark/connect/graphframes.proto b/sql/connect/common/src/main/protobuf/spark/connect/graphframes.proto new file mode 100644 index 0000000000000..1336a0be5a707 --- /dev/null +++ b/sql/connect/common/src/main/protobuf/spark/connect/graphframes.proto @@ -0,0 +1,331 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +syntax = 'proto3'; + +package spark.connect.graphframes; + +option java_multiple_files = true; +option java_package = "org.apache.spark.connect.proto.graphframes"; +option java_generate_equals_and_hash = true; +option optimize_for = SPEED; + + +// GraphFramesAPI represents the core message type for GraphFrames operations +// containing graph data and the specific graph algorithm to be executed +message GraphFramesAPI { + // Serialized vertex DataFrame containing node information + bytes vertices = 1; + // Serialized edge DataFrame containing relationship information + bytes edges = 2; + // Specifies which graph algorithm operation to perform + oneof method { + AggregateMessages aggregate_messages = 3; + BFS bfs = 4; + ConnectedComponents connected_components = 5; + DropIsolatedVertices drop_isolated_vertices = 6; + DetectingCycles detecting_cycles = 7; + FilterEdges filter_edges = 8; + FilterVertices filter_vertices = 9; + Find find = 10; + LabelPropagation label_propagation = 11; + PageRank page_rank = 12; + ParallelPersonalizedPageRank parallel_personalized_page_rank = 13; + PowerIterationClustering power_iteration_clustering = 14; + Pregel pregel = 15; + ShortestPaths shortest_paths = 16; + StronglyConnectedComponents strongly_connected_components = 17; + SVDPlusPlus svd_plus_plus = 18; + TriangleCount triangle_count = 19; + Triplets triplets = 20; + KCore kcore = 21; + MaximalIndependentSet mis = 22; + RandomWalkEmbeddings rw_embeddings = 23; + AggregateNeighbors aggregate_neighbors = 24; + NeighborhoodAwareCDLP neighborhood_aware_cdlp = 25; + AllPaths all_paths = 26; + HyperANF hyper_anf = 27; + } +} + +// Mapping follows PySpark Storage Levels! +// (not Scala-Spark Storage Levels) +message StorageLevel { + oneof storage_level { + bool disk_only = 1; + bool disk_only_2 = 2; + bool disk_only_3 = 3; + bool memory_and_disk = 4; + bool memory_and_disk_2 = 5; + bool memory_and_disk_deser = 6; + bool memory_only = 7; + bool memory_only_2 = 8; + } +} + +// String expression or serialized column +message ColumnOrExpression { + oneof col_or_expr { + bytes col = 1; + string expr = 2; + } +} + +// Connect supports only string or long-like IDs +message StringOrLongID { + oneof id { + int64 long_id = 1; + string string_id = 2; + } +} + +message AggregateMessages { + repeated ColumnOrExpression agg_col = 1; + repeated ColumnOrExpression send_to_src = 2; + repeated ColumnOrExpression send_to_dst = 3; + optional StorageLevel storage_level = 4; +} + +message BFS { + ColumnOrExpression from_expr = 1; + ColumnOrExpression to_expr = 2; + ColumnOrExpression edge_filter = 3; + int32 max_path_length = 4; +} + +message AllPaths { + ColumnOrExpression from_expr = 1; + ColumnOrExpression to_expr = 2; + ColumnOrExpression edge_filter = 3; + int32 max_path_length = 4; + bool is_directed = 5; + int32 checkpoint_interval = 6; + bool use_local_checkpoints = 7; + optional StorageLevel storage_level = 8; +} + +message HyperANF { + int32 n_hops = 1; + int32 lg_nom_entries = 2; + optional ColumnOrExpression edges_filter_expression = 3; + int32 checkpoint_interval = 4; + bool use_local_checkpoints = 5; + optional StorageLevel storage_level = 6; +} + +message ConnectedComponents { + string algorithm = 1; + int32 checkpoint_interval = 2; + int32 broadcast_threshold = 3; + bool use_labels_as_components = 4; + bool use_local_checkpoints = 5; + int32 max_iter = 6; + optional StorageLevel storage_level = 7; +} + +message DetectingCycles { + bool use_local_checkpoints = 1; + int32 checkpoint_interval = 2; + optional StorageLevel storage_level = 3; +} + +message DropIsolatedVertices {} + +message FilterEdges { + ColumnOrExpression condition = 1; +} + +message FilterVertices { + ColumnOrExpression condition = 2; +} + +message Find { + string pattern = 1; +} + +message LabelPropagation { + string algorithm = 1; + int32 max_iter = 2; + bool use_local_checkpoints = 3; + int32 checkpoint_interval = 4; + optional StorageLevel storage_level = 5; +} + +message NeighborhoodAwareCDLP { + int32 max_iter = 1; + bool ignore_direct_links = 2; + double structural_similarity_multiplier = 3; + bool use_local_checkpoints = 4; + int32 checkpoint_interval = 5; + optional StorageLevel storage_level = 6; + bool is_directed = 7; + int32 lg_nom_entries = 8; + optional string initial_label_col = 9; +} + +message PageRank { + double reset_probability = 1; + optional StringOrLongID source_id = 2; + optional int32 max_iter = 3; + optional double tol = 4; +} + +message ParallelPersonalizedPageRank { + double reset_probability = 1; + repeated StringOrLongID source_ids = 2; + int32 max_iter = 3; +} + +message PowerIterationClustering { + int32 k = 1; + int32 max_iter = 2; + optional string weight_col = 3; +} + +message Pregel { + ColumnOrExpression agg_msgs = 1; + repeated ColumnOrExpression send_msg_to_dst = 2; + repeated ColumnOrExpression send_msg_to_src = 3; + int32 checkpoint_interval = 4; + int32 max_iter = 5; + string additional_col_name = 6; + ColumnOrExpression additional_col_initial = 7; + ColumnOrExpression additional_col_upd = 8; + optional bool early_stopping = 9; + bool use_local_checkpoints = 10; + optional StorageLevel storage_level = 11; + optional bool stop_if_all_non_active = 12; + optional ColumnOrExpression initial_active_expr = 13; + optional ColumnOrExpression update_active_expr = 14; + optional bool skip_messages_from_non_active = 15; + // Required columns for triplet construction (memory optimization) + // Column names separated by comma + optional string required_src_columns = 16; + optional string required_dst_columns = 17; + optional string required_edge_columns = 18; +} + +message ShortestPaths { + repeated StringOrLongID landmarks = 1; + string algorithm = 2; + bool use_local_checkpoints = 3; + int32 checkpoint_interval = 4; + optional StorageLevel storage_level = 5; + optional bool is_directed = 6; +} + +message StronglyConnectedComponents { + int32 max_iter = 1; +} + +message SVDPlusPlus { + int32 rank = 1; + int32 max_iter = 2; + double min_value = 3; + double max_value = 4; + double gamma1 = 5; + double gamma2 = 6; + double gamma6 = 7; + double gamma7 = 8; +} + +message TriangleCount { + optional StorageLevel storage_level = 1; + optional string algorithm = 2; + optional int32 lg_nom_entries = 3; +} + +message Triplets {} + +message MaximalIndependentSet { + int32 checkpoint_interval = 1; + optional StorageLevel storage_level = 2; + bool use_local_checkpoints = 3; + int64 seed = 4; +} + +message KCore { + bool use_local_checkpoints = 1; + int32 checkpoint_interval = 2; + optional StorageLevel storage_level = 3; +} + +message AggregateNeighbors { + // Starting vertices condition (Boolean column expression) + ColumnOrExpression starting_vertices = 1; + // Maximum number of hops to explore + int32 max_hops = 2; + // Accumulator names + repeated string accumulator_names = 3; + // Accumulator initial value expressions + repeated ColumnOrExpression accumulator_inits = 4; + // Accumulator update expressions + repeated ColumnOrExpression accumulator_updates = 5; + // Optional stopping condition (Boolean column expression) + optional ColumnOrExpression stopping_condition = 6; + // Optional target condition (Boolean column expression) + optional ColumnOrExpression target_condition = 7; + // Optional required vertex attributes to carry through traversal + repeated string required_vertex_attributes = 8; + // Optional required edge attributes to carry through traversal + repeated string required_edge_attributes = 9; + // Optional edge filter condition (Boolean column expression) + optional ColumnOrExpression edge_filter = 10; + // Whether to remove self-loops + bool remove_loops = 11; + // Checkpoint interval (0 means disabled) + int32 checkpoint_interval = 12; + // Whether to use local checkpoints + bool use_local_checkpoints = 13; + // Optional storage level for intermediate results + optional StorageLevel storage_level = 14; +} + +message RandomWalkEmbeddings { + bool use_edge_direction = 1; + string rw_model = 2; + int32 rw_max_nbrs = 3; + int32 rw_num_walks_per_node = 4; + int32 rw_batch_size = 5; + int32 rw_num_batches = 6; + int64 rw_seed = 7; + double rw_restart_probability = 8; + string rw_temporary_prefix = 9; + string rw_cached_walks = 10; + string sequence_model = 11; + int32 hash2vec_context_size = 12; + int32 hash2vec_num_partitions = 13; + int32 hash2vec_embeddings_dim = 14; + string hash2vec_decay_function = 15; + double hash2vec_gaussian_sigma = 16; + int32 hash2vec_hashing_seed = 17; + int32 hash2vec_sign_seed = 18; + bool hash2vec_do_l2_norm = 19; + bool hash2vec_safe_l2 = 20; + int32 word2vec_max_iter = 21; + int32 word2vec_embeddings_dim = 22; + int32 word2vec_window_size = 23; + int32 word2vec_num_partitions = 24; + int32 word2vec_min_count = 25; + int32 word2vec_max_sentence_length = 26; + int64 word2vec_seed = 27; + double word2vec_step_size = 28; + bool aggregate_neighbors = 29; + int32 aggregate_neighbors_max_nbrs = 30; + int64 aggregate_neighbors_seed = 31; + bool clean_up_after_run = 32; +} diff --git a/sql/connect/common/src/main/protobuf/spark/connect/relations.proto b/sql/connect/common/src/main/protobuf/spark/connect/relations.proto index 1517197437032..90b0a7da5cd90 100644 --- a/sql/connect/common/src/main/protobuf/spark/connect/relations.proto +++ b/sql/connect/common/src/main/protobuf/spark/connect/relations.proto @@ -24,6 +24,7 @@ import "spark/connect/expressions.proto"; import "spark/connect/types.proto"; import "spark/connect/catalog.proto"; import "spark/connect/common.proto"; +import "spark/connect/graphframes.proto"; import "spark/connect/ml_common.proto"; option java_multiple_files = true; @@ -84,6 +85,7 @@ message Relation { RelationChanges relation_changes = 46; NearestByJoin nearest_by_join = 47; Zip zip = 48; + spark.connect.graphframes.GraphFramesAPI graph_frames = 49; // NA functions NAFill fill_na = 90; diff --git a/sql/connect/server/pom.xml b/sql/connect/server/pom.xml index e4b11686229cf..cb62c04c0b152 100644 --- a/sql/connect/server/pom.xml +++ b/sql/connect/server/pom.xml @@ -100,6 +100,12 @@ + + org.apache.spark + spark-graphframes_${scala.binary.version} + ${project.version} + provided + org.apache.spark spark-mllib_${scala.binary.version} diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala index 78562e8ebf92d..29dd92f7fa0bb 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala @@ -71,6 +71,7 @@ import org.apache.spark.sql.execution.QueryExecution import org.apache.spark.sql.execution.aggregate.{ScalaAggregator, TypedAggregateExpression} import org.apache.spark.sql.execution.arrow.ArrowConverters import org.apache.spark.sql.execution.command.{CreateViewCommand, ExternalCommandExecutor} +import org.apache.spark.sql.graphframes.{GraphFrameInternals, GraphFramesConnectUtils} import org.apache.spark.sql.execution.datasources.jdbc.JDBCOptions import org.apache.spark.sql.execution.datasources.v2.python.UserDefinedPythonDataSource import org.apache.spark.sql.execution.python.{UserDefinedPythonFunction, UserDefinedPythonTableFunction} @@ -234,6 +235,11 @@ class SparkConnectPlanner( case proto.Relation.RelTypeCase.ML_RELATION => MLHandler.transformMLRelation(rel.getMlRelation, sessionHolder).logicalPlan + // Built-in GraphFrames relation. + case proto.Relation.RelTypeCase.GRAPH_FRAMES => + GraphFrameInternals.planFromDataFrame( + GraphFramesConnectUtils.parseAPICall(rel.getGraphFrames, this)) + // Handle plugins for Spark Connect Relation types. case proto.Relation.RelTypeCase.EXTENSION => transformRelationPlugin(rel.getExtension) diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala new file mode 100644 index 0000000000000..c31369565aa45 --- /dev/null +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala @@ -0,0 +1,639 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Because Dataset.ofRows is private[sql], we are forced to use spark package; +// Same about a Column helper object. +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.connect.planner.SparkConnectPlanner +import org.apache.spark.sql.functions.expr +import org.apache.spark.sql.functions.lit +import org.apache.spark.storage.StorageLevel +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.connect.proto.{graphframes => proto} +import org.apache.spark.graphframes.embeddings.RandomWalkEmbeddings + +import scala.jdk.CollectionConverters._ + +/** + * Utility object providing helper methods for parsing and transforming data structures related to + * GraphFrames and enabling interaction between GraphFrames and Spark Connect APIs. + * + * The methods in this object are intended for internal use within the GraphFrames module + * (`private[graphframes]`) to support parsing, transformation, and execution of GraphFrame API + * calls based on serialized or protocol buffer inputs. + */ +object GraphFramesConnectUtils { + + /** + * Parses a protobuf StorageLevel object and converts it to a corresponding Spark StorageLevel. + * + * @param pbStorageLevel + * the protobuf StorageLevel object to be parsed + * @return + * the corresponding Spark StorageLevel + */ + private[graphframes] def parseStorageLevel(pbStorageLevel: proto.StorageLevel): StorageLevel = { + pbStorageLevel.getStorageLevelCase match { + case proto.StorageLevel.StorageLevelCase.DISK_ONLY => StorageLevel.DISK_ONLY + case proto.StorageLevel.StorageLevelCase.DISK_ONLY_2 => StorageLevel.DISK_ONLY_2 + case proto.StorageLevel.StorageLevelCase.DISK_ONLY_3 => StorageLevel.DISK_ONLY_3 + case proto.StorageLevel.StorageLevelCase.MEMORY_AND_DISK => StorageLevel.MEMORY_AND_DISK_SER + case proto.StorageLevel.StorageLevelCase.MEMORY_AND_DISK_2 => + StorageLevel.MEMORY_AND_DISK_SER_2 + case proto.StorageLevel.StorageLevelCase.MEMORY_AND_DISK_DESER => + StorageLevel.MEMORY_AND_DISK + case proto.StorageLevel.StorageLevelCase.MEMORY_ONLY => StorageLevel.MEMORY_ONLY_SER + case proto.StorageLevel.StorageLevelCase.MEMORY_ONLY_2 => StorageLevel.MEMORY_ONLY_SER_2 + case _ => throw new GraphFramesUnreachableException() + } + } + + /** + * Parses a proto.ColumnOrExpression object and converts it to a corresponding Spark Column. + * + * @param colOrExpr + * the proto.ColumnOrExpression object to be parsed + * @param planner + * the SparkConnectPlanner used for transforming expressions + * @return + * the resulting Spark Column + */ + private[graphframes] def parseColumnOrExpression( + colOrExpr: proto.ColumnOrExpression, + planner: SparkConnectPlanner): Column = { + colOrExpr.getColOrExprCase match { + case proto.ColumnOrExpression.ColOrExprCase.COL => + GraphFrameInternals.createColumn( + planner.transformExpression( + org.apache.spark.connect.proto.Expression.parseFrom(colOrExpr.getCol.toByteArray))) + case proto.ColumnOrExpression.ColOrExprCase.EXPR => expr(colOrExpr.getExpr) + case _ => + throw new GraphFramesUnreachableException() + } + } + + /** + * Converts a proto.StringOrLongID object to its corresponding Scala representation. + * + * @param id + * the proto.StringOrLongID object to be parsed + * @return + * the Scala representation of the ID (String or Long) + * @throws GraphFramesUnreachableException + * if the ID case is unrecognized + */ + private[graphframes] def parseLongOrStringID(id: proto.StringOrLongID): Any = { + id.getIdCase match { + case proto.StringOrLongID.IdCase.LONG_ID => id.getLongId + case proto.StringOrLongID.IdCase.STRING_ID => id.getStringId + case _ => + throw new GraphFramesUnreachableException() + } + } + + /** + * Parses the given serialized data to construct a Spark DataFrame. + * + * @param data + * the serialized representation of the DataFrame in ByteString format. Must not be empty. + * @param planner + * the SparkConnectPlanner instance used to transform the serialized plan into a Spark + * DataFrame. + * @return + * the resulting Spark DataFrame created from the provided data. + * @throws IllegalArgumentException + * if the given data is empty. + */ + private[graphframes] def parseDataFrame( + data: ByteString, + planner: SparkConnectPlanner): DataFrame = { + if (data.isEmpty) { + throw new IllegalArgumentException( + "Expected a serialized DataFrame but got an empty ByteString.") + } + GraphFrameInternals.createDataFrame( + planner.sessionHolder.session, + planner.transformRelation( + org.apache.spark.connect.proto.Plan.parseFrom(data.toByteArray).getRoot)) + } + + /** + * Extracts a GraphFrame from the provided GraphFramesAPI message using the specified planner. + * + * @param apiMessage + * the GraphFramesAPI protobuf message containing serialized vertices and edges + * @param planner + * the SparkConnectPlanner used for parsing and constructing DataFrames + * @return + * the constructed GraphFrame consisting of vertices and edges + */ + private[graphframes] def extractGraphFrame( + apiMessage: proto.GraphFramesAPI, + planner: SparkConnectPlanner): GraphFrame = { + val vertices = parseDataFrame(apiMessage.getVertices, planner) + val edges = parseDataFrame(apiMessage.getEdges, planner) + + GraphFrame(vertices, edges) + } + + /** + * Parses a GraphFrames API call from a protocol buffer message and executes the corresponding + * operation on the GraphFrame object obtained from the planner. + * + * @param apiMessage + * The protocol buffer message that defines the GraphFrames API operation and its parameters. + * @param planner + * A SparkConnectPlanner instance used to translate protocol buffer expressions into Spark SQL + * objects (e.g., DataFrame, Column). + * @return + * A DataFrame that represents the result of the executed GraphFrame operation. + */ + private[spark] def parseAPICall( + apiMessage: proto.GraphFramesAPI, + planner: SparkConnectPlanner): DataFrame = { + val graphFrame = extractGraphFrame(apiMessage, planner) + + apiMessage.getMethodCase match { + case proto.GraphFramesAPI.MethodCase.AGGREGATE_MESSAGES => { + val aggregateMessagesProto = apiMessage.getAggregateMessages + var aggregateMessages = graphFrame.aggregateMessages + if (aggregateMessagesProto.getSendToDstList.size() == 1) { + aggregateMessages = aggregateMessages.sendToDst( + parseColumnOrExpression(aggregateMessagesProto.getSendToDst(0), planner)) + } else if (aggregateMessagesProto.getSendToDstList.size() > 1) { + val sendToDst = aggregateMessagesProto.getSendToDstList.asScala.map( + parseColumnOrExpression(_, planner)) + aggregateMessages = + aggregateMessages.sendToDst(sendToDst.head, sendToDst.tail.toSeq: _*) + } + if (aggregateMessagesProto.getSendToSrcList.size() == 1) { + aggregateMessages = aggregateMessages.sendToSrc( + parseColumnOrExpression(aggregateMessagesProto.getSendToSrc(0), planner)) + } else if (aggregateMessagesProto.getSendToSrcList.size() > 1) { + val sendToSrc = aggregateMessagesProto.getSendToSrcList.asScala.map( + parseColumnOrExpression(_, planner)) + aggregateMessages = + aggregateMessages.sendToSrc(sendToSrc.head, sendToSrc.tail.toSeq: _*) + } + + if (aggregateMessagesProto.hasStorageLevel) { + aggregateMessages = aggregateMessages.setIntermediateStorageLevel( + parseStorageLevel(aggregateMessagesProto.getStorageLevel)) + } + + val aggCols = + aggregateMessagesProto.getAggColList.asScala.map(parseColumnOrExpression(_, planner)) + + // At least one agg col is required, and it is easier to check it on the client side + if (aggCols.size == 1) { + aggregateMessages.agg(aggCols.head) + } else { + aggregateMessages.agg(aggCols.head, aggCols.tail.toSeq: _*) + } + } + case proto.GraphFramesAPI.MethodCase.BFS => { + val bfsProto = apiMessage.getBfs + graphFrame.bfs + .toExpr(parseColumnOrExpression(bfsProto.getToExpr, planner)) + .fromExpr(parseColumnOrExpression(bfsProto.getFromExpr, planner)) + .edgeFilter(parseColumnOrExpression(bfsProto.getEdgeFilter, planner)) + .maxPathLength(bfsProto.getMaxPathLength) + .run() + } + case proto.GraphFramesAPI.MethodCase.ALL_PATHS => { + val allPathsProto = apiMessage.getAllPaths + var allPaths = graphFrame.allPaths + .toExpr(parseColumnOrExpression(allPathsProto.getToExpr, planner)) + .fromExpr(parseColumnOrExpression(allPathsProto.getFromExpr, planner)) + .edgeFilter(parseColumnOrExpression(allPathsProto.getEdgeFilter, planner)) + .maxPathLength(allPathsProto.getMaxPathLength) + .setIsDirected(allPathsProto.getIsDirected) + .setCheckpointInterval(allPathsProto.getCheckpointInterval) + .setUseLocalCheckpoints(allPathsProto.getUseLocalCheckpoints) + if (allPathsProto.hasStorageLevel) { + allPaths = + allPaths.setIntermediateStorageLevel(parseStorageLevel(allPathsProto.getStorageLevel)) + } + allPaths.run() + } + case proto.GraphFramesAPI.MethodCase.CONNECTED_COMPONENTS => { + val cc = apiMessage.getConnectedComponents + val ccBuilder = graphFrame.connectedComponents + .maxIter(cc.getMaxIter) + .setAlgorithm(cc.getAlgorithm) + .setCheckpointInterval(cc.getCheckpointInterval) + .setBroadcastThreshold(cc.getBroadcastThreshold) + .setUseLocalCheckpoints(cc.getUseLocalCheckpoints) + .setUseLabelsAsComponents(cc.getUseLabelsAsComponents) + + if (cc.hasStorageLevel) { + ccBuilder.setIntermediateStorageLevel(parseStorageLevel(cc.getStorageLevel)).run() + } else { + ccBuilder.run() + } + } + + case proto.GraphFramesAPI.MethodCase.DETECTING_CYCLES => { + val dc = apiMessage.getDetectingCycles + val dcBuilder = graphFrame.detectingCycles + .setCheckpointInterval(dc.getCheckpointInterval) + .setUseLocalCheckpoints(dc.getUseLocalCheckpoints) + if (dc.hasStorageLevel) { + dcBuilder.setIntermediateStorageLevel(parseStorageLevel(dc.getStorageLevel)).run() + } else { + dcBuilder.run() + } + } + + case proto.GraphFramesAPI.MethodCase.DROP_ISOLATED_VERTICES => { + graphFrame.dropIsolatedVertices().vertices + } + case proto.GraphFramesAPI.MethodCase.FILTER_EDGES => { + val condition = parseColumnOrExpression(apiMessage.getFilterEdges.getCondition, planner) + graphFrame.filterEdges(condition).edges + } + case proto.GraphFramesAPI.MethodCase.FILTER_VERTICES => { + val condition = + parseColumnOrExpression(apiMessage.getFilterVertices.getCondition, planner) + graphFrame.filterVertices(condition).vertices + } + case proto.GraphFramesAPI.MethodCase.FIND => { + graphFrame.find(apiMessage.getFind.getPattern) + } + case proto.GraphFramesAPI.MethodCase.LABEL_PROPAGATION => { + val lp = apiMessage.getLabelPropagation + val lpBuilder = graphFrame.labelPropagation + .maxIter(lp.getMaxIter) + .setAlgorithm(lp.getAlgorithm) + .setCheckpointInterval(lp.getCheckpointInterval) + .setUseLocalCheckpoints(lp.getUseLocalCheckpoints) + + if (lp.hasStorageLevel) { + lpBuilder.setIntermediateStorageLevel(parseStorageLevel(lp.getStorageLevel)).run() + } else { + lpBuilder.run() + } + } + case proto.GraphFramesAPI.MethodCase.NEIGHBORHOOD_AWARE_CDLP => { + val nc = apiMessage.getNeighborhoodAwareCdlp + val ncBuilder = graphFrame.structureAwareLabelPropagation + .maxIter(nc.getMaxIter) + .setIgnoreDirectLinks(nc.getIgnoreDirectLinks) + .setStructuralSimilarityMultiplier(nc.getStructuralSimilarityMultiplier) + .setUseLocalCheckpoints(nc.getUseLocalCheckpoints) + .setCheckpointInterval(nc.getCheckpointInterval) + .setIsDirected(nc.getIsDirected) + .setLgNomEntries(nc.getLgNomEntries) + + if (nc.hasInitialLabelCol) { + ncBuilder.setInitialLabelCol(nc.getInitialLabelCol) + } + + if (nc.hasStorageLevel) { + ncBuilder.setIntermediateStorageLevel(parseStorageLevel(nc.getStorageLevel)).run() + } else { + ncBuilder.run() + } + } + case proto.GraphFramesAPI.MethodCase.PAGE_RANK => { + val pageRankProto = apiMessage.getPageRank + val pageRank = graphFrame.pageRank.resetProbability(pageRankProto.getResetProbability) + + if (pageRankProto.hasMaxIter) { + pageRank.maxIter(pageRankProto.getMaxIter) + } else { + pageRank.tol(pageRankProto.getTol) + } + + if (pageRankProto.hasSourceId) { + pageRank.sourceId(parseLongOrStringID(pageRankProto.getSourceId)) + } + + // Edges should be updated on the client side + // TODO: do we really need an edge weights in that case? + // see comments in the Python API + pageRank.run().vertices + } + case proto.GraphFramesAPI.MethodCase.PARALLEL_PERSONALIZED_PAGE_RANK => { + val pPageRankProto = apiMessage.getParallelPersonalizedPageRank + val sourceIds = pPageRankProto.getSourceIdsList.asScala + .map(parseLongOrStringID) + .toArray + val pPageRank = graphFrame.parallelPersonalizedPageRank + pPageRank + .resetProbability(pPageRankProto.getResetProbability) + .maxIter(pPageRankProto.getMaxIter) + .sourceIds(sourceIds) + .run() + .vertices // See comment in the PageRank + } + case proto.GraphFramesAPI.MethodCase.POWER_ITERATION_CLUSTERING => { + val pic = apiMessage.getPowerIterationClustering + if (pic.hasWeightCol) { + graphFrame.powerIterationClustering(pic.getK, pic.getMaxIter, Some(pic.getWeightCol)) + } else { + graphFrame.powerIterationClustering(pic.getK, pic.getMaxIter, None) + } + } + case proto.GraphFramesAPI.MethodCase.PREGEL => { + val pregelProto = apiMessage.getPregel + var pregel = graphFrame.pregel + .aggMsgs(parseColumnOrExpression(pregelProto.getAggMsgs, planner)) + .setCheckpointInterval(pregelProto.getCheckpointInterval) + .withVertexColumn( + pregelProto.getAdditionalColName, + parseColumnOrExpression(pregelProto.getAdditionalColInitial, planner), + parseColumnOrExpression(pregelProto.getAdditionalColUpd, planner)) + .setMaxIter(pregelProto.getMaxIter) + .setUseLocalCheckpoints(pregelProto.getUseLocalCheckpoints) + + if (pregelProto.hasStorageLevel) { + pregel = + pregel.setIntermediateStorageLevel(parseStorageLevel(pregelProto.getStorageLevel)) + } + + if (pregelProto.hasInitialActiveExpr) { + // We are not checking here that all the attrs are present; + // Check should be done on the client side. + pregel = pregel + .setInitialActiveVertexExpression( + parseColumnOrExpression(pregelProto.getInitialActiveExpr, planner)) + .setUpdateActiveVertexExpression( + parseColumnOrExpression(pregelProto.getUpdateActiveExpr, planner)) + + if (pregelProto.hasSkipMessagesFromNonActive) { + pregel = pregel.setSkipMessagesFromNonActiveVertices( + pregelProto.getSkipMessagesFromNonActive) + } + + if (pregelProto.hasStopIfAllNonActive) { + pregel = pregel.setStopIfAllNonActiveVertices(pregelProto.getStopIfAllNonActive) + } + } + + pregel = pregelProto.getSendMsgToSrcList.asScala + .map(parseColumnOrExpression(_, planner)) + .foldLeft(pregel)((p, col) => p.sendMsgToSrc(col)) + pregel = pregelProto.getSendMsgToDstList.asScala + .map(parseColumnOrExpression(_, planner)) + .foldLeft(pregel)((p, col) => p.sendMsgToDst(col)) + + if (pregelProto.hasEarlyStopping) { + pregel = pregel.setEarlyStopping(pregelProto.getEarlyStopping) + } + + // Handle required columns for triplet optimization (comma-separated) + if (pregelProto.hasRequiredSrcColumns) { + val cols = + pregelProto.getRequiredSrcColumns.split(",").map(_.trim).filter(_.nonEmpty).toSeq + if (cols.nonEmpty) pregel = pregel.requiredSrcColumns(cols.head, cols.tail: _*) + } + + if (pregelProto.hasRequiredDstColumns) { + val cols = + pregelProto.getRequiredDstColumns.split(",").map(_.trim).filter(_.nonEmpty).toSeq + if (cols.nonEmpty) pregel = pregel.requiredDstColumns(cols.head, cols.tail: _*) + } + + if (pregelProto.hasRequiredEdgeColumns) { + val cols = + pregelProto.getRequiredEdgeColumns.split(",").map(_.trim).filter(_.nonEmpty).toSeq + if (cols.nonEmpty) pregel = pregel.requiredEdgeColumns(cols.head, cols.tail: _*) + } + + pregel.run() + } + case proto.GraphFramesAPI.MethodCase.SHORTEST_PATHS => { + val isDirected = if (apiMessage.getShortestPaths.hasIsDirected) { + apiMessage.getShortestPaths.getIsDirected + } else { + true + } + + val spBuilder = graphFrame.shortestPaths + .landmarks( + apiMessage.getShortestPaths.getLandmarksList.asScala.map(parseLongOrStringID).toSeq) + .setAlgorithm(apiMessage.getShortestPaths.getAlgorithm) + .setCheckpointInterval(apiMessage.getShortestPaths.getCheckpointInterval) + .setUseLocalCheckpoints(apiMessage.getShortestPaths.getUseLocalCheckpoints) + .setIsDirected(isDirected) + + if (apiMessage.getShortestPaths.hasStorageLevel) { + spBuilder + .setIntermediateStorageLevel( + parseStorageLevel(apiMessage.getShortestPaths.getStorageLevel)) + .run() + } else { + spBuilder.run() + } + } + case proto.GraphFramesAPI.MethodCase.STRONGLY_CONNECTED_COMPONENTS => { + graphFrame.stronglyConnectedComponents + .maxIter(apiMessage.getStronglyConnectedComponents.getMaxIter) + .run() + } + case proto.GraphFramesAPI.MethodCase.SVD_PLUS_PLUS => { + val svdPPProto = apiMessage.getSvdPlusPlus + val svd = graphFrame.svdPlusPlus + .maxIter(svdPPProto.getMaxIter) + .gamma1(svdPPProto.getGamma1) + .gamma2(svdPPProto.getGamma2) + .gamma6(svdPPProto.getGamma6) + .gamma7(svdPPProto.getGamma7) + .rank(svdPPProto.getRank) + .minValue(svdPPProto.getMinValue) + .maxValue(svdPPProto.getMaxValue) + val svdResult = svd.run() + svdResult.withColumn("loss", lit(svd.loss)) + } + case proto.GraphFramesAPI.MethodCase.TRIANGLE_COUNT => { + val message = apiMessage.getTriangleCount() + var trCounter = + graphFrame.triangleCount + + if (message.hasAlgorithm) { + trCounter = trCounter.setAlgorithm(message.getAlgorithm) + } + + if (message.hasLgNomEntries) { + trCounter = trCounter.setLgNomEntries(message.getLgNomEntries) + } + + if (message.hasStorageLevel) { + trCounter + .setIntermediateStorageLevel(parseStorageLevel(message.getStorageLevel)) + .run() + } else { + trCounter.run() + } + } + case proto.GraphFramesAPI.MethodCase.TRIPLETS => { + graphFrame.triplets + } + case proto.GraphFramesAPI.MethodCase.MIS => { + val mis = graphFrame.maximalIndependentSet + .setCheckpointInterval(apiMessage.getMis.getCheckpointInterval) + .setUseLocalCheckpoints(apiMessage.getMis.getUseLocalCheckpoints) + + if (apiMessage.getMis.hasStorageLevel) { + mis + .setIntermediateStorageLevel(parseStorageLevel(apiMessage.getMis.getStorageLevel)) + .run(apiMessage.getMis.getSeed) + } else { + mis.run(apiMessage.getMis.getSeed) + } + } + case proto.GraphFramesAPI.MethodCase.KCORE => { + var kCoreBuilder = + graphFrame.kCore + .setCheckpointInterval(apiMessage.getKcore.getCheckpointInterval) + .setUseLocalCheckpoints(apiMessage.getKcore.getUseLocalCheckpoints) + + if (apiMessage.getKcore.hasStorageLevel) { + kCoreBuilder = kCoreBuilder.setIntermediateStorageLevel( + parseStorageLevel(apiMessage.getKcore.getStorageLevel)) + } + + kCoreBuilder.run() + } + case proto.GraphFramesAPI.MethodCase.AGGREGATE_NEIGHBORS => { + val anProto = apiMessage.getAggregateNeighbors + var anBuilder = graphFrame.aggregateNeighbors + .setStartingVertices(parseColumnOrExpression(anProto.getStartingVertices, planner)) + .setMaxHops(anProto.getMaxHops) + + // Set accumulators + val accNames = anProto.getAccumulatorNamesList.asScala.toSeq + val accInits = anProto.getAccumulatorInitsList.asScala + .map(parseColumnOrExpression(_, planner)) + .toSeq + val accUpdates = anProto.getAccumulatorUpdatesList.asScala + .map(parseColumnOrExpression(_, planner)) + .toSeq + + if (accNames.nonEmpty) { + anBuilder = anBuilder.setAccumulators(accNames, accInits, accUpdates) + } + + // Optional parameters + if (anProto.hasStoppingCondition) { + anBuilder = anBuilder.setStoppingCondition( + parseColumnOrExpression(anProto.getStoppingCondition, planner)) + } + + if (anProto.hasTargetCondition) { + anBuilder = anBuilder.setTargetCondition( + parseColumnOrExpression(anProto.getTargetCondition, planner)) + } + + val reqVertexAttrs = anProto.getRequiredVertexAttributesList.asScala.toSeq + if (reqVertexAttrs.nonEmpty) { + anBuilder = anBuilder.setRequiredVertexAttributes(reqVertexAttrs) + } + + val reqEdgeAttrs = anProto.getRequiredEdgeAttributesList.asScala.toSeq + if (reqEdgeAttrs.nonEmpty) { + anBuilder = anBuilder.setRequiredEdgeAttributes(reqEdgeAttrs) + } + + if (anProto.hasEdgeFilter) { + anBuilder = + anBuilder.setEdgeFilter(parseColumnOrExpression(anProto.getEdgeFilter, planner)) + } + + anBuilder = anBuilder.setRemoveLoops(anProto.getRemoveLoops) + + if (anProto.getCheckpointInterval > 0) { + anBuilder = anBuilder.setCheckpointInterval(anProto.getCheckpointInterval) + } + + anBuilder = anBuilder.setUseLocalCheckpoints(anProto.getUseLocalCheckpoints) + + if (anProto.hasStorageLevel) { + anBuilder = + anBuilder.setIntermediateStorageLevel(parseStorageLevel(anProto.getStorageLevel)) + } + + anBuilder.run() + } + + case proto.GraphFramesAPI.MethodCase.RW_EMBEDDINGS => { + val message = apiMessage.getRwEmbeddings() + + RandomWalkEmbeddings.pythonAPI( + graph = graphFrame, + useEdgeDirection = message.getUseEdgeDirection(), + rwModel = message.getRwModel(), + rwMaxNbrs = message.getRwMaxNbrs(), + rwNumWalksPerNode = message.getRwNumWalksPerNode(), + rwBatchSize = message.getRwBatchSize(), + rwNumBatches = message.getRwNumBatches(), + rwSeed = message.getRwSeed(), + rwRestartProbability = message.getRwRestartProbability(), + rwTemporaryPrefix = message.getRwTemporaryPrefix(), + rwCachedWalks = message.getRwCachedWalks(), + sequenceModel = message.getSequenceModel(), + hash2vecContextSize = message.getHash2VecContextSize(), + hash2vecNumPartitions = message.getHash2VecNumPartitions(), + hash2vecEmbeddingsDim = message.getHash2VecEmbeddingsDim(), + hash2vecDecayFunction = message.getHash2VecDecayFunction(), + hash2vecGaussianSigma = message.getHash2VecGaussianSigma(), + hash2vecHashingSeed = message.getHash2VecHashingSeed(), + hash2vecSignSeed = message.getHash2VecSignSeed(), + hash2vecDoL2Norm = message.getHash2VecDoL2Norm(), + hash2vecSafeL2 = message.getHash2VecSafeL2(), + word2vecMaxIter = message.getWord2VecMaxIter(), + word2vecEmbeddingsDim = message.getWord2VecEmbeddingsDim(), + word2vecWindowSize = message.getWord2VecWindowSize(), + word2vecNumPartitions = message.getWord2VecNumPartitions(), + word2vecMinCount = message.getWord2VecMinCount(), + word2vecMaxSentenceLength = message.getWord2VecMaxSentenceLength(), + word2vecSeed = message.getWord2VecSeed(), + word2vecStepSize = message.getWord2VecStepSize(), + aggregateNeighbors = message.getAggregateNeighbors(), + aggregateNeighborsMaxNbrs = message.getAggregateNeighborsMaxNbrs(), + aggregateNeighborsSeed = message.getAggregateNeighborsSeed(), + cleanUpAfterRun = message.getCleanUpAfterRun()) + } + case proto.GraphFramesAPI.MethodCase.HYPER_ANF => { + val haProto = apiMessage.getHyperAnf + val haBuilder = graphFrame.hyperANF + .setNHops(haProto.getNHops) + .setLgNomEntries(haProto.getLgNomEntries) + .setCheckpointInterval(haProto.getCheckpointInterval) + .setUseLocalCheckpoints(haProto.getUseLocalCheckpoints) + + if (haProto.hasEdgesFilterExpression) { + haBuilder.setEdgesFilterExpression( + parseColumnOrExpression(haProto.getEdgesFilterExpression, planner)) + } + + if (haProto.hasStorageLevel) { + haBuilder.setIntermediateStorageLevel(parseStorageLevel(haProto.getStorageLevel)).run() + } else { + haBuilder.run() + } + } + case _ => throw new GraphFramesUnreachableException() // Unreachable + } + } +} From 0187810cecd7fce85322cdd7fce9769aeaba2baf Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Wed, 26 Aug 2026 05:46:12 +0000 Subject: [PATCH 3/8] [GRAPHFRAMES][PYTHON] Register GraphFrames tests --- dev/sparktestsupport/modules.py | 30 ++++++++++++++++++++++++++++-- dev/sparktestsupport/utils.py | 20 +++++++++++++------- 2 files changed, 41 insertions(+), 9 deletions(-) diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index 5d3efb657b9cd..2f966be70acfc 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -540,6 +540,17 @@ def __hash__(self): ], ) +graphframes = Module( + name="graphframes", + dependencies=[graphx, mllib], + source_file_regexes=[ + "graphframes/", + ], + sbt_test_goals=[ + "graphframes/test", + ], +) + pipelines = Module( name="pipelines", dependencies=[sql], @@ -551,7 +562,7 @@ def __hash__(self): connect = Module( name="connect", - dependencies=[hive, avro, protobuf, mllib], + dependencies=[hive, avro, protobuf, graphframes, mllib], source_file_regexes=[ "sql/connect", ], @@ -576,7 +587,9 @@ def __hash__(self): pyspark_core = Module( name="pyspark-core", dependencies=[core], - source_file_regexes=["python/(?!pyspark/(ml|mllib|sql|streaming|pandas|resource|testing))"], + source_file_regexes=[ + "python/(?!pyspark/(graphframes|ml|mllib|sql|streaming|pandas|resource|testing))" + ], python_test_goals=[ # doctests "pyspark.conf", @@ -1355,6 +1368,19 @@ def __hash__(self): ], ) +pyspark_graphframes = Module( + name="pyspark-graphframes", + dependencies=[pyspark_connect, graphframes], + source_file_regexes=[ + "python/pyspark/graphframes", + ], + python_test_goals=[ + "pyspark.graphframes.tests.test_graphframe", + "pyspark.graphframes.tests.connect.test_parity_graphframe", + "pyspark.graphframes.tests.connect.test_all_algorithms", + ], +) + pyspark_structured_streaming_connect = Module( name="pyspark-structured-streaming-connect", dependencies=[pyspark_connect, pyspark_structured_streaming], diff --git a/dev/sparktestsupport/utils.py b/dev/sparktestsupport/utils.py index b177952eb31dd..8d1218c06c511 100755 --- a/dev/sparktestsupport/utils.py +++ b/dev/sparktestsupport/utils.py @@ -105,11 +105,15 @@ def determine_modules_to_test(changed_modules, deduplicated=True): >>> [x.name for x in determine_modules_to_test([modules.launcher])] ['root'] >>> [x.name for x in determine_modules_to_test([modules.graphx])] - ['graphx', 'examples'] + ... # doctest: +NORMALIZE_WHITESPACE + ['graphx', 'examples', 'graphframes', 'connect', 'pyspark-connect', 'pyspark-graphframes', + 'pyspark-ml-connect', 'pyspark-pandas-connect', 'pyspark-pandas-slow-connect', + 'pyspark-pipelines', 'pyspark-structured-streaming-connect'] >>> sorted([x.name for x in determine_modules_to_test([modules.sql])]) ... # doctest: +NORMALIZE_WHITESPACE - ['avro', 'connect', 'docker-integration-tests', 'examples', 'hive', 'hive-thriftserver', - 'mllib', 'pipelines', 'protobuf', 'pyspark-connect', 'pyspark-ml', 'pyspark-ml-connect', + ['avro', 'connect', 'docker-integration-tests', 'examples', 'graphframes', 'hive', + 'hive-thriftserver', 'mllib', 'pipelines', 'protobuf', 'pyspark-connect', + 'pyspark-graphframes', 'pyspark-ml', 'pyspark-ml-connect', 'pyspark-mllib', 'pyspark-pandas', 'pyspark-pandas-connect', 'pyspark-pandas-slow', 'pyspark-pandas-slow-connect', 'pyspark-pipelines', 'pyspark-sql', 'pyspark-structured-streaming', 'pyspark-structured-streaming-connect', @@ -117,8 +121,9 @@ def determine_modules_to_test(changed_modules, deduplicated=True): >>> sorted([x.name for x in determine_modules_to_test( ... [modules.sparkr, modules.sql], deduplicated=False)]) ... # doctest: +NORMALIZE_WHITESPACE - ['avro', 'connect', 'docker-integration-tests', 'examples', 'hive', 'hive-thriftserver', - 'mllib', 'pipelines', 'protobuf', 'pyspark-connect', 'pyspark-ml', 'pyspark-ml-connect', + ['avro', 'connect', 'docker-integration-tests', 'examples', 'graphframes', 'hive', + 'hive-thriftserver', 'mllib', 'pipelines', 'protobuf', 'pyspark-connect', + 'pyspark-graphframes', 'pyspark-ml', 'pyspark-ml-connect', 'pyspark-mllib', 'pyspark-pandas', 'pyspark-pandas-connect', 'pyspark-pandas-slow', 'pyspark-pandas-slow-connect', 'pyspark-pipelines', 'pyspark-sql', 'pyspark-structured-streaming', 'pyspark-structured-streaming-connect', @@ -127,9 +132,10 @@ def determine_modules_to_test(changed_modules, deduplicated=True): ... [modules.sql, modules.core], deduplicated=False)]) ... # doctest: +NORMALIZE_WHITESPACE ['avro', 'catalyst', 'connect', 'core', 'credential-aws', 'docker-integration-tests', - 'examples', 'graphx', + 'examples', 'graphframes', 'graphx', 'hive', 'hive-thriftserver', 'mllib', 'mllib-local', 'pipelines', 'protobuf', - 'pyspark-connect', 'pyspark-core', 'pyspark-errors', 'pyspark-ml', 'pyspark-ml-connect', + 'pyspark-connect', 'pyspark-core', 'pyspark-errors', 'pyspark-graphframes', 'pyspark-ml', + 'pyspark-ml-connect', 'pyspark-mllib', 'pyspark-pandas', 'pyspark-pandas-connect', 'pyspark-pandas-slow', 'pyspark-pandas-slow-connect', 'pyspark-pipelines', 'pyspark-resource', 'pyspark-sql', 'pyspark-streaming', 'pyspark-structured-streaming', 'pyspark-structured-streaming-connect', From 52789020b7f8ebe3666cd92f8742a9fd22d6777a Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Wed, 26 Aug 2026 06:49:55 +0000 Subject: [PATCH 4/8] [GRAPHFRAMES][PYTHON] Run GraphFrames tests in build workflow --- .github/workflows/build_and_test.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index fd6aa5114add0..00d94293b20bb 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -625,6 +625,8 @@ jobs: pyspark-streaming, pyspark-structured-streaming, pyspark-structured-streaming-connect - >- pyspark-connect + - >- + pyspark-graphframes - >- pyspark-connect-old-client - >- @@ -642,6 +644,7 @@ jobs: - modules: ${{ fromJson(needs.precondition.outputs.required).pyspark != 'true' && 'pyspark-mllib, pyspark-ml, pyspark-ml-connect' }} - modules: ${{ fromJson(needs.precondition.outputs.required).pyspark != 'true' && 'pyspark-streaming, pyspark-structured-streaming, pyspark-structured-streaming-connect' }} - modules: ${{ fromJson(needs.precondition.outputs.required).pyspark != 'true' && 'pyspark-connect' }} + - modules: ${{ fromJson(needs.precondition.outputs.required).pyspark != 'true' && 'pyspark-graphframes' }} - modules: ${{ fromJson(needs.precondition.outputs.required).pyspark-connect-old-client != 'true' && 'pyspark-connect-old-client'}} # pyspark-install is very slow so we only run it when it's changed or explicity requested - modules: ${{ fromJson(needs.precondition.outputs.required).pyspark-install != 'true' && 'pyspark-install' }} From ed40a1117564590abf9aaa0a2b6cfd8403a92a56 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Wed, 26 Aug 2026 06:51:42 +0000 Subject: [PATCH 5/8] [GRAPHFRAMES] Run Scala tests in build workflow --- .github/workflows/build_and_test.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index 00d94293b20bb..6a18de6670032 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -164,7 +164,7 @@ jobs: docs=false java25=false fi - build=`./dev/is-changed.py -m "core,unsafe,kvstore,avro,utils,utils-java,network-common,network-shuffle,repl,launcher,examples,sketch,variant,api,catalyst,hive-thriftserver,mllib-local,mllib,graphx,streaming,sql-kafka-0-10,streaming-kafka-0-10,streaming-kinesis-asl,credential-aws,kubernetes,hadoop-cloud,spark-ganglia-lgpl,profiler,protobuf,yarn,connect,sql,hive,pipelines"` + build=`./dev/is-changed.py -m "core,unsafe,kvstore,avro,utils,utils-java,network-common,network-shuffle,repl,launcher,examples,sketch,variant,api,catalyst,hive-thriftserver,mllib-local,mllib,graphx,graphframes,streaming,sql-kafka-0-10,streaming-kafka-0-10,streaming-kinesis-asl,credential-aws,kubernetes,hadoop-cloud,spark-ganglia-lgpl,profiler,protobuf,yarn,connect,sql,hive,pipelines"` build_core_utils=`./dev/is-changed.py -m "core,unsafe,kvstore,utils,utils-java,network-common,network-shuffle,sketch,variant,launcher"` precondition=" { @@ -288,7 +288,7 @@ jobs: - >- api, catalyst, hive-thriftserver - >- - mllib-local, mllib, graphx, profiler, pipelines, repl, examples + mllib-local, mllib, graphx, graphframes, profiler, pipelines, repl, examples - >- streaming, sql-kafka-0-10, streaming-kafka-0-10, streaming-kinesis-asl, credential-aws, kubernetes, hadoop-cloud, spark-ganglia-lgpl, protobuf, connect, avro From 0cc556dfd3c87be281d3749df94d3c115a3a5966 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Fri, 28 Aug 2026 04:11:56 +0000 Subject: [PATCH 6/8] [GRAPHFRAMES][TESTS] Port upstream GraphFrames test coverage --- .github/workflows/build_and_test.yml | 4 +- .github/workflows/build_python_connect.yml | 2 +- NOTICE | 11 +- dev/sparktestsupport/modules.py | 18 +- dev/sparktestsupport/utils.py | 13 +- .../graphframes/GraphFramePythonAPI.scala | 4 +- .../examples/BeliefPropagation.scala | 274 ++++ .../spark/graphframes/examples/Graphs.scala | 8 +- .../propertygraph/PropertyGraphFrame.scala | 224 +++ .../property/EdgePropertyGroup.scala | 178 +++ .../property/PropertyGroup.scala | 47 + .../property/VertexPropertyGroup.scala | 127 ++ .../examples/BeliefPropagationSuite.scala | 68 + .../graphframes/examples/GraphsSuite.scala | 56 + .../graphframes/examples/LDBCUtils.scala | 129 ++ .../graphframes/ldbc/TestLDBCCases.scala | 258 ++++ .../PropertyGraphFrameSuite.scala | 299 ++++ .../expressions/KMinSamplingSuite.scala | 78 + python/packaging/classic/setup.py | 7 + python/packaging/client/setup.py | 9 + .../pyspark/graphframes/classic/graphframe.py | 2 +- python/pyspark/graphframes/classic/pregel.py | 342 +++++ .../pyspark/graphframes/examples/__init__.py | 21 + .../examples/belief_propagation.py | 177 +++ python/pyspark/graphframes/examples/graphs.py | 137 ++ python/pyspark/graphframes/lib/pregel.py | 307 +--- python/pyspark/graphframes/pg/__init__.py | 25 + .../graphframes/pg/property_graphframe.py | 381 +++++ .../pyspark/graphframes/pg/property_groups.py | 386 +++++ .../graphframes/tests/_upstream_test_utils.py | 109 ++ .../tests/connect/test_all_algorithms.py | 7 +- .../tests/connect/test_client_imports.py | 43 + .../tests/connect/test_property_graphframe.py | 119 ++ .../connect/test_upstream_graphframes.py | 176 +++ .../pyspark/graphframes/tests/pg/__init__.py | 16 + .../tests/pg/test_property_graphframe.py | 484 ++++++ .../graphframes/tests/test_graphframe.py | 6 + .../graphframes/tests/test_graphframes.py | 1311 +++++++++++++++++ 38 files changed, 5582 insertions(+), 281 deletions(-) create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/examples/BeliefPropagation.scala rename graphframes/src/{test => main}/scala/org/apache/spark/graphframes/examples/Graphs.scala (100%) create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/PropertyGraphFrame.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/property/EdgePropertyGroup.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/property/PropertyGroup.scala create mode 100644 graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/property/VertexPropertyGroup.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/examples/BeliefPropagationSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/examples/GraphsSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/examples/LDBCUtils.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/ldbc/TestLDBCCases.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/graphframes/propertygraph/PropertyGraphFrameSuite.scala create mode 100644 graphframes/src/test/scala/org/apache/spark/sql/graphframes/expressions/KMinSamplingSuite.scala create mode 100644 python/pyspark/graphframes/classic/pregel.py create mode 100644 python/pyspark/graphframes/examples/__init__.py create mode 100644 python/pyspark/graphframes/examples/belief_propagation.py create mode 100644 python/pyspark/graphframes/examples/graphs.py create mode 100644 python/pyspark/graphframes/pg/__init__.py create mode 100644 python/pyspark/graphframes/pg/property_graphframe.py create mode 100644 python/pyspark/graphframes/pg/property_groups.py create mode 100644 python/pyspark/graphframes/tests/_upstream_test_utils.py create mode 100644 python/pyspark/graphframes/tests/connect/test_client_imports.py create mode 100644 python/pyspark/graphframes/tests/connect/test_property_graphframe.py create mode 100644 python/pyspark/graphframes/tests/connect/test_upstream_graphframes.py create mode 100644 python/pyspark/graphframes/tests/pg/__init__.py create mode 100644 python/pyspark/graphframes/tests/pg/test_property_graphframe.py create mode 100644 python/pyspark/graphframes/tests/test_graphframes.py diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index 6a18de6670032..f044b42a14452 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -626,7 +626,7 @@ jobs: - >- pyspark-connect - >- - pyspark-graphframes + pyspark-graphframes,pyspark-graphframes-connect - >- pyspark-connect-old-client - >- @@ -644,7 +644,7 @@ jobs: - modules: ${{ fromJson(needs.precondition.outputs.required).pyspark != 'true' && 'pyspark-mllib, pyspark-ml, pyspark-ml-connect' }} - modules: ${{ fromJson(needs.precondition.outputs.required).pyspark != 'true' && 'pyspark-streaming, pyspark-structured-streaming, pyspark-structured-streaming-connect' }} - modules: ${{ fromJson(needs.precondition.outputs.required).pyspark != 'true' && 'pyspark-connect' }} - - modules: ${{ fromJson(needs.precondition.outputs.required).pyspark != 'true' && 'pyspark-graphframes' }} + - modules: ${{ fromJson(needs.precondition.outputs.required).pyspark != 'true' && 'pyspark-graphframes,pyspark-graphframes-connect' }} - modules: ${{ fromJson(needs.precondition.outputs.required).pyspark-connect-old-client != 'true' && 'pyspark-connect-old-client'}} # pyspark-install is very slow so we only run it when it's changed or explicity requested - modules: ${{ fromJson(needs.precondition.outputs.required).pyspark-install != 'true' && 'pyspark-install' }} diff --git a/.github/workflows/build_python_connect.yml b/.github/workflows/build_python_connect.yml index 5e1d269fc5f52..c66789e09e016 100644 --- a/.github/workflows/build_python_connect.yml +++ b/.github/workflows/build_python_connect.yml @@ -94,7 +94,7 @@ jobs: mv python/pyspark pyspark.back # Several tests related to catalog requires to run them sequencially, e.g., writing a table in a listener. - ./python/run-tests --parallelism=1 --python-executables=python3 --modules pyspark-connect,pyspark-ml-connect + ./python/run-tests --parallelism=1 --python-executables=python3 --modules pyspark-connect,pyspark-ml-connect,pyspark-graphframes-connect # None of tests are dependent on each other in Pandas API on Spark so run them in parallel ./python/run-tests --parallelism=1 --python-executables=python3 --modules pyspark-pandas-connect,pyspark-pandas-slow-connect diff --git a/NOTICE b/NOTICE index d5ea8dedb311b..6a8ad9bd74de2 100644 --- a/NOTICE +++ b/NOTICE @@ -38,4 +38,13 @@ LongAdder), which was released with the following comments: Written by Doug Lea with assistance from members of JCP JSR-166 Expert Group and released to the public domain, as explained at - http://creativecommons.org/publicdomain/zero/1.0/ \ No newline at end of file + http://creativecommons.org/publicdomain/zero/1.0/ + + +LDBC datasets +------------- + +The LDBC datasets used by GraphFrames integration tests are licensed under the Apache +Software License, Version 2.0. They are used for testing and evaluation purposes only. +Per the LDBC fair-use policy, results must not be described using the words "LDBC benchmark" +or any equivalent phrase. diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index 2f966be70acfc..9c12860ac2fd6 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -1370,14 +1370,30 @@ def __hash__(self): pyspark_graphframes = Module( name="pyspark-graphframes", - dependencies=[pyspark_connect, graphframes], + dependencies=[pyspark_sql, graphframes], source_file_regexes=[ "python/pyspark/graphframes", ], python_test_goals=[ "pyspark.graphframes.tests.test_graphframe", + "pyspark.graphframes.tests.test_graphframes", + "pyspark.graphframes.tests.pg.test_property_graphframe", + ], +) + +pyspark_graphframes_connect = Module( + name="pyspark-graphframes-connect", + dependencies=[pyspark_connect, pyspark_graphframes], + source_file_regexes=[ + "python/pyspark/graphframes/connect", + "python/pyspark/graphframes/tests/connect", + ], + python_test_goals=[ + "pyspark.graphframes.tests.connect.test_client_imports", "pyspark.graphframes.tests.connect.test_parity_graphframe", "pyspark.graphframes.tests.connect.test_all_algorithms", + "pyspark.graphframes.tests.connect.test_upstream_graphframes", + "pyspark.graphframes.tests.connect.test_property_graphframe", ], ) diff --git a/dev/sparktestsupport/utils.py b/dev/sparktestsupport/utils.py index 8d1218c06c511..f45fa4ecf30ea 100755 --- a/dev/sparktestsupport/utils.py +++ b/dev/sparktestsupport/utils.py @@ -106,14 +106,15 @@ def determine_modules_to_test(changed_modules, deduplicated=True): ['root'] >>> [x.name for x in determine_modules_to_test([modules.graphx])] ... # doctest: +NORMALIZE_WHITESPACE - ['graphx', 'examples', 'graphframes', 'connect', 'pyspark-connect', 'pyspark-graphframes', - 'pyspark-ml-connect', 'pyspark-pandas-connect', 'pyspark-pandas-slow-connect', + ['graphx', 'examples', 'graphframes', 'connect', 'pyspark-graphframes', 'pyspark-connect', + 'pyspark-graphframes-connect', 'pyspark-ml-connect', 'pyspark-pandas-connect', + 'pyspark-pandas-slow-connect', 'pyspark-pipelines', 'pyspark-structured-streaming-connect'] >>> sorted([x.name for x in determine_modules_to_test([modules.sql])]) ... # doctest: +NORMALIZE_WHITESPACE ['avro', 'connect', 'docker-integration-tests', 'examples', 'graphframes', 'hive', 'hive-thriftserver', 'mllib', 'pipelines', 'protobuf', 'pyspark-connect', - 'pyspark-graphframes', 'pyspark-ml', 'pyspark-ml-connect', + 'pyspark-graphframes', 'pyspark-graphframes-connect', 'pyspark-ml', 'pyspark-ml-connect', 'pyspark-mllib', 'pyspark-pandas', 'pyspark-pandas-connect', 'pyspark-pandas-slow', 'pyspark-pandas-slow-connect', 'pyspark-pipelines', 'pyspark-sql', 'pyspark-structured-streaming', 'pyspark-structured-streaming-connect', @@ -123,7 +124,7 @@ def determine_modules_to_test(changed_modules, deduplicated=True): ... # doctest: +NORMALIZE_WHITESPACE ['avro', 'connect', 'docker-integration-tests', 'examples', 'graphframes', 'hive', 'hive-thriftserver', 'mllib', 'pipelines', 'protobuf', 'pyspark-connect', - 'pyspark-graphframes', 'pyspark-ml', 'pyspark-ml-connect', + 'pyspark-graphframes', 'pyspark-graphframes-connect', 'pyspark-ml', 'pyspark-ml-connect', 'pyspark-mllib', 'pyspark-pandas', 'pyspark-pandas-connect', 'pyspark-pandas-slow', 'pyspark-pandas-slow-connect', 'pyspark-pipelines', 'pyspark-sql', 'pyspark-structured-streaming', 'pyspark-structured-streaming-connect', @@ -134,8 +135,8 @@ def determine_modules_to_test(changed_modules, deduplicated=True): ['avro', 'catalyst', 'connect', 'core', 'credential-aws', 'docker-integration-tests', 'examples', 'graphframes', 'graphx', 'hive', 'hive-thriftserver', 'mllib', 'mllib-local', 'pipelines', 'protobuf', - 'pyspark-connect', 'pyspark-core', 'pyspark-errors', 'pyspark-graphframes', 'pyspark-ml', - 'pyspark-ml-connect', + 'pyspark-connect', 'pyspark-core', 'pyspark-errors', 'pyspark-graphframes', + 'pyspark-graphframes-connect', 'pyspark-ml', 'pyspark-ml-connect', 'pyspark-mllib', 'pyspark-pandas', 'pyspark-pandas-connect', 'pyspark-pandas-slow', 'pyspark-pandas-slow-connect', 'pyspark-pipelines', 'pyspark-resource', 'pyspark-sql', 'pyspark-streaming', 'pyspark-structured-streaming', 'pyspark-structured-streaming-connect', diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFramePythonAPI.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFramePythonAPI.scala index f4b0676c01789..ede106de5d3ca 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFramePythonAPI.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFramePythonAPI.scala @@ -17,8 +17,9 @@ package org.apache.spark.graphframes -import org.apache.spark.sql.DataFrame +import org.apache.spark.graphframes.examples.Graphs import org.apache.spark.graphframes.lib.AggregateMessages +import org.apache.spark.sql.DataFrame private[graphframes] class GraphFramePythonAPI { @@ -31,4 +32,5 @@ private[graphframes] class GraphFramePythonAPI { val ATTR: String = GraphFrame.ATTR lazy val aggregateMessages: AggregateMessages.type = AggregateMessages + lazy val examples: Graphs = Graphs } diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/examples/BeliefPropagation.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/examples/BeliefPropagation.scala new file mode 100644 index 0000000000000..01b0e8154a1f2 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/examples/BeliefPropagation.scala @@ -0,0 +1,274 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.examples + +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.examples.Graphs.gridIsingModel +import org.apache.spark.graphframes.lib.AggregateMessages +import org.apache.spark.graphx +import org.apache.spark.sql.Column +import org.apache.spark.sql.Row +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.functions.sum +import org.apache.spark.sql.functions.udf +import org.apache.spark.sql.functions.when + +/** + * Example code for Belief Propagation (BP) + * + * This provides a template for building customized BP algorithms for different types of graphical + * models. + * + * This example: + * - Ising model on a grid + * - Parallel Belief Propagation using colored fields + * + * Ising models are probabilistic graphical models over binary variables x,,i,,. Each binary + * variable x,,i,, corresponds to one vertex, and it may take values -1 or +1. The probability + * distribution P(X) (over all x,,i,,) is parameterized by vertex factors a,,i,, and edge factors + * b,,ij,,: + * {{{ + * P(X) = (1/Z) * exp[ \sum_i a_i x_i + \sum_{ij} b_{ij} x_i x_j ] + * }}} + * where Z is the normalization constant (partition function). See + * [[https://en.wikipedia.org/wiki/Ising_model Wikipedia]] for more information on Ising models. + * + * Belief Propagation (BP) provides marginal probabilities of the values of the variables x,,i,,, + * i.e., P(x,,i,,) for each i. This allows a user to understand likely values of variables. See + * [[https://en.wikipedia.org/wiki/Belief_propagation Wikipedia]] for more information on BP. + * + * We use a batch synchronous BP algorithm, where batches of vertices are updated synchronously. + * We follow the mean field update algorithm in Slide 13 of the + * [[http://www.eecs.berkeley.edu/~wainwrig/Talks/A_GraphModel_Tutorial talk slides]] from: + * Wainwright. "Graphical models, message-passing algorithms, and convex optimization." + * + * The batches are chosen according to a coloring. For background on graph colorings for + * inference, see for example: Gonzalez et al. "Parallel Gibbs Sampling: From Colored Fields to + * Thin Junction Trees." AISTATS, 2011. + * + * The BP algorithm works by: + * - Coloring the graph by assigning a color to each vertex such that no neighboring vertices + * share the same color. + * - In each step of BP, update all vertices of a single color. Alternate colors. + */ +object BeliefPropagation { + + def main(args: Array[String]): Unit = { + val spark = SparkSession + .builder() + .appName("BeliefPropagation example") + .getOrCreate() + + // Create graphical model g of size 3 x 3. + val g = gridIsingModel(spark, 3) + + // scalastyle:off println + println("Original Ising model:") + g.vertices.show() + g.edges.show() + + // Run BP for 5 iterations. + val numIter = 5 + val results = runBPwithGraphX(g, numIter) + + // Display beliefs. + val beliefs = results.vertices.select("id", "belief") + println(s"Done with BP. Final beliefs after $numIter iterations:") + // scalastyle:on println + beliefs.show() + + spark.stop() + } + + /** + * Given a GraphFrame, choose colors for each vertex. No neighboring vertices will share the + * same color. The number of colors is minimized. + * + * This is written specifically for grid graphs. For non-grid graphs, it should be generalized, + * such as by using a greedy coloring scheme. + * + * @param g + * Grid graph generated by [[org.apache.spark.graphframes.examples.Graphs.gridIsingModel()]] + * @return + * Same graph, but with a new vertex column "color" of type Int (0 or 1) + */ + private def colorGraph(g: GraphFrame): GraphFrame = { + val colorUDF = udf { (i: Int, j: Int) => (i + j) % 2 } + val v = g.vertices.withColumn("color", colorUDF(col("i"), col("j"))) + GraphFrame(v, g.edges) + } + + /** + * Run Belief Propagation. + * + * This implementation of BP shows how to use GraphX's aggregateMessages method. It is simple to + * convert to and from GraphX format. This method does the following: + * - Color GraphFrame vertices for BP scheduling. + * - Convert GraphFrame to GraphX format. + * - Run BP using GraphX's aggregateMessages API. + * - Augment the original GraphFrame with the BP results (vertex beliefs). + * + * @param g + * Graphical model created by `org.apache.spark.graphframes.examples.Graphs.gridIsingModel()` + * @param numIter + * Number of iterations of BP to run. One iteration includes updating each vertex's belief + * once. + * @return + * Same graphical model, but with [[GraphFrame.vertices]] augmented with a new column "belief" + * containing P(x,,i,, = +1), the marginal probability of vertex i taking value +1 instead of + * -1. + */ + def runBPwithGraphX(g: GraphFrame, numIter: Int): GraphFrame = { + // Choose colors for vertices for BP scheduling. + val colorG = colorGraph(g) + val numColors: Int = colorG.vertices.select("color").distinct().count().toInt + + // Convert GraphFrame to GraphX, and initialize beliefs. + val gx0 = colorG.toGraphX + // Schema maps for extracting attributes + val vColsMap = colorG.vertexColumnMap + val eColsMap = colorG.edgeColumnMap + // Convert vertex attributes to nice case classes. + val gx1: graphx.Graph[VertexAttr, Row] = gx0.mapVertices { case (_, attr) => + // Initialize belief at 0.0 + VertexAttr(attr.getDouble(vColsMap("a")), 0.0, attr.getInt(vColsMap("color"))) + } + // Convert edge attributes to nice case classes. + val extractEdgeAttr: (graphx.Edge[Row] => EdgeAttr) = { e => + EdgeAttr(e.attr.getDouble(eColsMap("b"))) + } + var gx: graphx.Graph[VertexAttr, EdgeAttr] = gx1.mapEdges(extractEdgeAttr) + + // Run BP for numIter iterations. + for (_ <- Range(0, numIter)) { + // For each color, have that color receive messages from neighbors. + for (color <- Range(0, numColors)) { + // Send messages to vertices of the current color. + val msgs: graphx.VertexRDD[Double] = gx.aggregateMessages( + ctx => + // Can send to source or destination since edges are treated as undirected. + if (ctx.dstAttr.color == color) { + val msg = ctx.attr.b * ctx.srcAttr.belief + // Only send message if non-zero. + if (msg != 0) ctx.sendToDst(msg) + } else if (ctx.srcAttr.color == color) { + val msg = ctx.attr.b * ctx.dstAttr.belief + // Only send message if non-zero. + if (msg != 0) ctx.sendToSrc(msg) + }, + _ + _) + // Receive messages, and update beliefs for vertices of the current color. + gx = gx.outerJoinVertices(msgs) { case (_, vAttr, optMsg) => + if (vAttr.color == color) { + val x = vAttr.a + optMsg.getOrElse(0.0) + val newBelief = math.exp(-log1pExp(-x)) + VertexAttr(vAttr.a, newBelief, color) + } else { + vAttr + } + } + } + } + + // Convert back to GraphFrame with a new column "belief" for vertices DataFrame. + val gxFinal: graphx.Graph[Double, Unit] = + gx.mapVertices((_, attr) => attr.belief).mapEdges(_ => ()) + GraphFrame.fromGraphX(colorG, gxFinal, vertexNames = Seq("belief")) + } + + case class VertexAttr(a: Double, belief: Double, color: Int) + + case class EdgeAttr(b: Double) + + /** + * Run Belief Propagation. + * + * This implementation of BP shows how to use GraphFrame's aggregateMessages method. + * - Color GraphFrame vertices for BP scheduling. + * - Run BP using GraphFrame's aggregateMessages API. + * - Augment the original GraphFrame with the BP results (vertex beliefs). + * + * @param g + * Graphical model created by `org.apache.spark.graphframes.examples.Graphs.gridIsingModel()` + * @param numIter + * Number of iterations of BP to run. One iteration includes updating each vertex's belief + * once. + * @return + * Same graphical model, but with [[GraphFrame.vertices]] augmented with a new column "belief" + * containing P(x,,i,, = +1), the marginal probability of vertex i taking value +1 instead of + * -1. + */ + def runBPwithGraphFrames(g: GraphFrame, numIter: Int): GraphFrame = { + // Choose colors for vertices for BP scheduling. + val colorG = colorGraph(g) + val numColors: Int = colorG.vertices.select("color").distinct().count().toInt + + // TODO: Handle vertices without any edges. + + // Initialize vertex beliefs at 0.0. + var gx = GraphFrame(colorG.vertices.withColumn("belief", lit(0.0)), colorG.edges) + + // Run BP for numIter iterations. + for (_ <- Range(0, numIter)) { + // For each color, have that color receive messages from neighbors. + for (color <- Range(0, numColors)) { + // Define "AM" for shorthand for referring to the src, dst, edge, and msg fields. + // (See usage below.) + val AM = AggregateMessages + // Send messages to vertices of the current color. + // We may send to source or destination since edges are treated as undirected. + val msgForSrc: Column = when(AM.src("color") === color, AM.edge("b") * AM.dst("belief")) + val msgForDst: Column = when(AM.dst("color") === color, AM.edge("b") * AM.src("belief")) + val logistic = udf { (x: Double) => math.exp(-log1pExp(-x)) } + val aggregates = gx.aggregateMessages + .sendToSrc(msgForSrc) + .sendToDst(msgForDst) + .agg(sum(AM.msg).as("aggMess")) + val v = gx.vertices + // Receive messages, and update beliefs for vertices of the current color. + val newBeliefCol = when( + v("color") === color && aggregates("aggMess").isNotNull, + logistic(aggregates("aggMess") + v("a"))) + .otherwise(v("belief")) // keep old beliefs for other colors + val newVertices = v + .join(aggregates, v("id") === aggregates("id"), "left_outer") // join messages, vertices + .drop(aggregates("id")) // drop duplicate ID column (from outer join) + .withColumn("newBelief", newBeliefCol) // compute new beliefs + .drop("aggMess") // drop messages + .drop("belief") // drop old beliefs + .withColumnRenamed("newBelief", "belief") + val cachedNewVertices = newVertices.localCheckpoint() + gx = GraphFrame(cachedNewVertices, gx.edges) + } + } + + // Drop the "color" column from vertices + GraphFrame(gx.vertices.drop("color"), gx.edges) + } + + /** More numerically stable `log(1 + exp(x))` */ + private def log1pExp(x: Double): Double = { + if (x > 0) { + x + math.log1p(math.exp(-x)) + } else { + math.log1p(math.exp(x)) + } + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/examples/Graphs.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/examples/Graphs.scala similarity index 100% rename from graphframes/src/test/scala/org/apache/spark/graphframes/examples/Graphs.scala rename to graphframes/src/main/scala/org/apache/spark/graphframes/examples/Graphs.scala index a0fa223398b70..df4e504f9f7ce 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/examples/Graphs.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/examples/Graphs.scala @@ -17,15 +17,15 @@ package org.apache.spark.graphframes.examples +import scala.reflect.runtime.universe.TypeTag + +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrame._ import org.apache.spark.sql.SparkSession import org.apache.spark.sql.functions.col import org.apache.spark.sql.functions.lit import org.apache.spark.sql.functions.randn import org.apache.spark.sql.functions.udf -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFrame._ - -import scala.reflect.runtime.universe.TypeTag class Graphs private[graphframes] () { // Note: this cannot be values: we are creating and destroying spark contexts during the tests, diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/PropertyGraphFrame.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/PropertyGraphFrame.scala new file mode 100644 index 0000000000000..36f30d7b5cb4c --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/PropertyGraphFrame.scala @@ -0,0 +1,224 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.propertygraph + +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.propertygraph.property.EdgePropertyGroup +import org.apache.spark.graphframes.propertygraph.property.VertexPropertyGroup +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.lit + +/** + * A high-level abstraction for working with property graphs that simplifies interaction with the + * GraphFrames library. + * + * PropertyGraphFrame serves as a logical structure that manages collections of vertex and edge + * property groups, providing a user-friendly API for graph operations. It handles various + * internal complexities such as: + * - ID conversion and collision prevention + * - Management of directed/undirected graph representations + * - Handling of weighted/unweighted edges + * - Data consistency across different property groups + * + * The class maintains separate collections for vertex and edge properties, allowing for flexible + * graph construction while ensuring data integrity. Each property (vertex or edge) handles its + * data internally, while this class provides a simplified interface for working with the + * underlying GraphFrame structure. + * + * @param vertexPropertyGroups + * Sequence of vertex property groups that define the graph's vertices + * @param edgesPropertyGroups + * Sequence of edge property groups that define the graph's edges + */ +case class PropertyGraphFrame( + vertexPropertyGroups: Seq[VertexPropertyGroup], + edgesPropertyGroups: Seq[EdgePropertyGroup]) { + import PropertyGraphFrame._ + lazy private val vertexGroups: Map[String, VertexPropertyGroup] = + vertexPropertyGroups.map(pg => pg.name -> pg).toMap + lazy private val edgeGroups: Map[String, EdgePropertyGroup] = + edgesPropertyGroups.map(pg => pg.name -> pg).toMap + + /** + * Converts a heterogeneous property graph into a unified GraphFrame representation. + * + * This method transforms a property graph that may contain multiple vertex types and both + * directed and undirected edges into a single GraphFrame object where all vertices and edges + * share the same schema. The conversion process handles: + * + * - Internal ID generation and collision prevention by hashing vertex/edge IDs with their + * group names + * - Merging of different vertex types into a unified vertex DataFrame + * - Conversion of directed/undirected edge relationships into a consistent edge DataFrame + * - Filtering of vertices and edges based on provided predicates + * + * The method allows selecting a subset of property groups and applying filters to control which + * data is included in the final GraphFrame. + * + * @param vertexPropertyGroups + * Sequence of vertex property group names to include in the GraphFrame + * @param edgePropertyGroups + * Sequence of edge property group names to include in the GraphFrame + * @param edgeGroupFilters + * Map of edge property group names to filter predicates (Column expressions) + * @param vertexGroupFilters + * Map of vertex property group names to filter predicates (Column expressions) + * @return + * A GraphFrame containing the unified representation of the selected and filtered property + * groups + */ + def toGraphFrame( + vertexPropertyGroups: Seq[String], + edgePropertyGroups: Seq[String], + edgeGroupFilters: Map[String, Column], + vertexGroupFilters: Map[String, Column]): GraphFrame = { + vertexPropertyGroups.foreach(name => + require(vertexGroups.contains(name), s"Vertex property group $name does not exist")) + edgePropertyGroups.foreach(name => + require(edgeGroups.contains(name), s"Edge property group $name does not exist")) + + val vertices = vertexPropertyGroups + .map(name => vertexGroups(name).getData(vertexGroupFilters(name))) + .reduce(_ union _) + + val edges = edgePropertyGroups + .map(name => edgeGroups(name).getData(edgeGroupFilters(name))) + .reduce(_ union _) + + GraphFrame(vertices, edges) + } + + /** + * Projects a bipartite graph onto one of its parts, creating edges between vertices that share + * neighbors in the other part. Drops the property group used for projection through and returns + * a new property graph. + * + * @param leftBiGraphPart + * Name of the vertex property group to project onto + * @param rightBiGraphPart + * Name of the vertex property group to project through + * @param edgeGroup + * Name of the edge property group connecting the two parts + * @param newEdgeWeight + * Optional function that takes two weight columns (Column objects) of edges as input and + * returns a new weight column. If None, a default weight of 1.0 is used for all projected + * edges. + * @return + * A new PropertyGraphFrame containing the projected graph + */ + def projectionBy( + leftBiGraphPart: String, + rightBiGraphPart: String, + edgeGroup: String, + newEdgeWeight: Option[(Column, Column) => Column] = None): PropertyGraphFrame = { + require( + edgeGroups(edgeGroup).srcPropertyGroup.name == leftBiGraphPart, + s"Edge Property Group should have $leftBiGraphPart source group but has " + + edgeGroups(edgeGroup).srcPropertyGroup.name) + require( + edgeGroups(edgeGroup).dstPropertyGroup.name == rightBiGraphPart, + s"Edge Property Group should have $rightBiGraphPart destination group but has " + + edgeGroups(edgeGroup).dstPropertyGroup.name) + val keptVPropertyGroups = vertexPropertyGroups.filterNot(g => g.name == rightBiGraphPart) + val keptEPropertyGroups = edgesPropertyGroups.filterNot(g => g.name == edgeGroup) + val oldGroup = edgeGroups(edgeGroup) + val oldEdgesData = oldGroup.data + + // Create new edges by joining vertices through their common neighbors + val projectedEdges = oldEdgesData + .as("e1") + .join(oldEdgesData.as("e2"), col("e1.dst") === col("e2.dst")) + .where("e1.src < e2.src") + .select( + col("e1.src").alias(GraphFrame.SRC), + col("e2.src").alias(GraphFrame.DST), + newEdgeWeight match { + case Some(newEdgeFunc) => + newEdgeFunc( + col(s"e1.${oldGroup.weightColumnName}"), + col(s"e2.${oldGroup.weightColumnName}")).alias(GraphFrame.WEIGHT) + case None => lit(1.0).alias(GraphFrame.WEIGHT) + }) + + val newEdgeGroup = EdgePropertyGroup( + name = s"projected_$edgeGroup", + data = projectedEdges, + srcPropertyGroup = vertexGroups(leftBiGraphPart), + dstPropertyGroup = vertexGroups(leftBiGraphPart), + isDirected = false, + srcColumnName = GraphFrame.SRC, + dstColumnName = GraphFrame.DST, + weightColumnName = GraphFrame.WEIGHT) + + PropertyGraphFrame(keptVPropertyGroups, keptEPropertyGroups :+ newEdgeGroup) + } + + /** + * Joins the vertices data with the specified vertex property groups to produce a unified + * DataFrame. Each vertex property group defines how the data should be structured and filtered. + * + * @param verticesData + * The DataFrame containing the vertices data to join. It must include vertex properties and + * the group identifiers to filter and map. It is expected to be an output of calling graph + * algorithms on GraphFrame, made by the method toGraphFrame. + * @param vertexGroups + * A sequence of vertex group names that are to be joined. Each name must correspond to a + * valid vertex property group defined in the PropertyGraphFrame. + * @return + * A DataFrame representing the unified vertices data where each group has been appropriately + * filtered, joined, and processed based on its configuration. + */ + def joinVertices(verticesData: DataFrame, vertexGroups: Seq[String]): DataFrame = { + require(vertexGroups.forall(this.vertexGroups.contains)) + vertexGroups + .map { (vg: String) => + { + val associatedGroup = this.vertexGroups(vg) + val filteredForGroup = verticesData.filter(col(PROPERTY_GROUP_COL_NAME) === lit(vg)) + if (associatedGroup.applyMaskOnId) { + associatedGroup.internalIdMapping + .join(filteredForGroup, Seq(GraphFrame.ID), "left") + .drop(GraphFrame.ID) + } else { + associatedGroup + .getData() + .join(filteredForGroup, GraphFrame.ID, "left") + .withColumnRenamed(GraphFrame.ID, EXTERNAL_ID) + } + } + } + .reduce(_ union _) + } +} + +object PropertyGraphFrame { + + /** + * A constant representing the column name used for property grouping. It is used within the + * context of a property graph structure to manage or identify property group associations. + */ + val PROPERTY_GROUP_COL_NAME = "property_group" + + /** + * A constant representing the column name used for external identifiers. It serves as a key to + * associate external data or entities within the context of a property graph structure. + */ + val EXTERNAL_ID = "external_id" +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/property/EdgePropertyGroup.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/property/EdgePropertyGroup.scala new file mode 100644 index 0000000000000..e5907803a3fbf --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/property/EdgePropertyGroup.scala @@ -0,0 +1,178 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.propertygraph.property + +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.InvalidPropertyGroupException +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.concat +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.functions.sha2 +import org.apache.spark.sql.types._ + +/** + * Represents a logical group of edges in a property graph with associated metadata and data. + * + * EdgePropertyGroup encapsulates edge data stored in a DataFrame along with metadata describing + * how to interpret the data as graph edges. Each edge group has: + * + * - A unique name identifier + * - DataFrame containing the actual edge data + * - Source and destination vertex property groups + * - Direction flag indicating if edges are directed or undirected + * - Column names specifying source vertex, destination vertex and edge weight columns + * + * The class validates that required columns exist in the provided DataFrame on creation. Required + * columns are: + * - Source vertex column + * - Destination vertex column + * - Weight column + * + * @param name + * Unique identifier for this edge property group + * @param data + * DataFrame containing the edge data with required columns + * @param srcPropertyGroup + * Source vertex property group + * @param dstPropertyGroup + * Destination vertex property group + * @param isDirected + * Whether edges should be treated as directed (true) or undirected (false) + * @param srcColumnName + * Name of the source vertex column in the data + * @param dstColumnName + * Name of the destination vertex column in the data + * @param weightColumnName + * Name of the edge weight column in the data + * @note + * When edges from different groups are combined into a GraphFrame, their SRCs and DSTs are + * hashed with the group name to prevent collisions in the same way as ID of the corresponded + * vertex group is hashed. + */ +case class EdgePropertyGroup( + name: String, + data: DataFrame, + srcPropertyGroup: VertexPropertyGroup, + dstPropertyGroup: VertexPropertyGroup, + isDirected: Boolean, + srcColumnName: String, + dstColumnName: String, + weightColumnName: String) + extends PropertyGroup { + + override protected def validate(): this.type = { + if (!data.columns.contains(srcColumnName)) { + throw new InvalidPropertyGroupException( + s"source column $srcColumnName does not exist, existed columns " + + s"[${data.columns.mkString(", ")}]") + } + if (!data.columns.contains(dstColumnName)) { + throw new InvalidPropertyGroupException( + s"dest column $dstColumnName does not exist, existed columns " + + s"[${data.columns.mkString(", ")}]") + } + if (!data.columns.contains(weightColumnName)) { + throw new InvalidPropertyGroupException( + s"weight column $weightColumnName does not exist, existed columns " + + s"[${data.columns.mkString(", ")}]") + } + val weightColumnType = data.schema(weightColumnName).dataType + if (!weightColumnType.isInstanceOf[NumericType]) { + throw new InvalidPropertyGroupException( + s"weight column $weightColumnName must be numeric type, but was $weightColumnType") + } + + this + } + + private def hashSrcEdge: Column = if (srcPropertyGroup.applyMaskOnId) { + concat(lit(srcPropertyGroup.name), sha2(col(srcColumnName).cast(StringType), 256)) + } else { + col(srcColumnName).cast(StringType) + } + + private def hashDstEdge: Column = if (dstPropertyGroup.applyMaskOnId) { + concat(lit(dstPropertyGroup.name), sha2(col(dstColumnName).cast(StringType), 256)) + } else { + col(dstColumnName).cast(StringType) + } + + override protected[graphframes] def getData(filter: Column): DataFrame = { + val filteredData = data.filter(filter) + + val baseEdges = filteredData.select( + hashSrcEdge.alias(GraphFrame.SRC), + hashDstEdge.alias(GraphFrame.DST), + col(weightColumnName).alias(GraphFrame.WEIGHT)) + + if (isDirected) { + baseEdges + } else { + baseEdges.union( + baseEdges.select( + col(GraphFrame.DST).as(GraphFrame.SRC), + col(GraphFrame.SRC).as(GraphFrame.DST), + col(GraphFrame.WEIGHT).alias(GraphFrame.WEIGHT))) + } + } +} + +object EdgePropertyGroup { + def apply( + name: String, + data: DataFrame, + srcPropertyGroup: VertexPropertyGroup, + dstPropertyGroup: VertexPropertyGroup, + isDirected: Boolean, + srcColumnName: String, + dstColumnName: String, + weightColumnName: String): EdgePropertyGroup = { + new EdgePropertyGroup( + name, + data, + srcPropertyGroup, + dstPropertyGroup, + isDirected, + srcColumnName, + dstColumnName, + weightColumnName).validate() + } + + def apply( + name: String, + data: DataFrame, + srcPropertyGroup: VertexPropertyGroup, + dstPropertyGroup: VertexPropertyGroup, + isDirected: Boolean, + srcColumnName: String, + dstColumnName: String, + weightColumn: Column): EdgePropertyGroup = { + val dataWithWeight = data.withColumn(GraphFrame.WEIGHT, weightColumn) + apply( + name, + dataWithWeight, + srcPropertyGroup, + dstPropertyGroup, + isDirected, + srcColumnName, + dstColumnName, + GraphFrame.WEIGHT) + } +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/property/PropertyGroup.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/property/PropertyGroup.scala new file mode 100644 index 0000000000000..57ead8846cbc7 --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/property/PropertyGroup.scala @@ -0,0 +1,47 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.propertygraph.property + +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.lit + +trait PropertyGroup { + val name: String + val data: DataFrame + protected def validate(): this.type + + /** + * Returns a view of the data for the property group without applying any filter. + * + * @return + * A DataFrame containing the raw data. + */ + protected[graphframes] def getData(): DataFrame = getData(lit(true)) + + /** + * Returns a filtered view of the data for the property group, with an optional mask applied to + * IDs. + * + * @param filter + * A condition (Column) used to filter the data. + * @return + * A DataFrame containing the filtered and optionally transformed data. + */ + protected[graphframes] def getData(filter: Column): DataFrame +} diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/property/VertexPropertyGroup.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/property/VertexPropertyGroup.scala new file mode 100644 index 0000000000000..ab57c81fa312e --- /dev/null +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/propertygraph/property/VertexPropertyGroup.scala @@ -0,0 +1,127 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.propertygraph.property + +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.InvalidPropertyGroupException +import org.apache.spark.graphframes.propertygraph.PropertyGraphFrame +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.concat +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.functions.sha2 +import org.apache.spark.sql.types.StringType + +/** + * Represents a logical group of vertices in a property graph with associated data and + * identification. + * + * A VertexPropertyGroup is used to organize and manage vertices that share common characteristics + * or belong to the same logical group within a property graph. Each group maintains its own data + * in the form of a DataFrame and uses a primary key column for unique vertex identification. + * + * The class provides two ways to create a vertex property group: + * 1. With a specified primary key column: + * {{{ + * VertexPropertyGroup("users", userDataFrame, "userId") + * }}} + * 2. With the default primary key column ("id"): + * {{{ + * VertexPropertyGroup("users", userDataFrame) + * }}} + * + * @param name + * The unique identifier for this vertex property group + * @param data + * The DataFrame containing the vertex data + * @param primaryKeyColumn + * The column name used to uniquely identify vertices in this group + * @param applyMaskOnId + * A flag indicating whether to apply masking on vertex IDs. When false, uses raw IDs from + * primaryKeyColumn. When true, hashes IDs with group name. Defaults to true. + * @note + * When vertices from different groups are combined into a GraphFrame, their IDs are hashed with + * the group name to prevent collisions. + */ +case class VertexPropertyGroup( + name: String, + data: DataFrame, + primaryKeyColumn: String, + applyMaskOnId: Boolean = true) + extends PropertyGroup { + + override protected def validate(): this.type = { + if (!data.columns.contains(primaryKeyColumn)) { + throw new InvalidPropertyGroupException( + s"source column $primaryKeyColumn does not exist, existed columns " + + s"[${data.columns.mkString(", ")}]") + } + this + } + + private[graphframes] def internalIdMapping: DataFrame = data + .select(col(primaryKeyColumn).alias(PropertyGraphFrame.EXTERNAL_ID)) + .withColumn( + GraphFrame.ID, + concat(lit(name), sha2(col(PropertyGraphFrame.EXTERNAL_ID).cast(StringType), 256))) + + override protected[graphframes] def getData(filter: Column): DataFrame = { + val filteredData = data + .filter(filter) + val withId = if (applyMaskOnId) { + filteredData.select( + concat(lit(name), sha2(col(primaryKeyColumn).cast(StringType), 256)).alias(GraphFrame.ID)) + } else { + filteredData.select(col(primaryKeyColumn).cast(StringType).alias(GraphFrame.ID)) + } + + withId.select(col(GraphFrame.ID), lit(name).alias(PropertyGraphFrame.PROPERTY_GROUP_COL_NAME)) + } +} + +object VertexPropertyGroup { + + /** + * Creates a new VertexPropertyGroup with a specified primary key column. + * + * @param name + * Name of the vertex property group + * @param data + * DataFrame containing vertex data + * @param primaryKeyColumn + * Name of the column to be used as a primary key for vertex identification + * @return + * A validated VertexPropertyGroup instance + */ + def apply(name: String, data: DataFrame, primaryKeyColumn: String): VertexPropertyGroup = + new VertexPropertyGroup(name, data, primaryKeyColumn).validate() + + /** + * Creates a new VertexPropertyGroup using default a primary key column name. + * + * @param name + * Name of the vertex property group + * @param data + * DataFrame containing vertex data + * @return + * A validated VertexPropertyGroup instance + */ + def apply(name: String, data: DataFrame): VertexPropertyGroup = + new VertexPropertyGroup(name, data, GraphFrame.ID) +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/examples/BeliefPropagationSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/examples/BeliefPropagationSuite.scala new file mode 100644 index 0000000000000..61051ab681561 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/examples/BeliefPropagationSuite.scala @@ -0,0 +1,68 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.examples + +import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.examples.BeliefPropagation._ +import org.apache.spark.graphframes.examples.Graphs.gridIsingModel +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Row + +class BeliefPropagationSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + test("BP using GraphX and GraphFrame aggregateMessages") { + val n = 3 // graph is n x n + val numIter = 5 // iterations of BP + + // Create graphical model g. + val g = gridIsingModel(spark, n) + + // Run BP using GraphX + val gxResults = runBPwithGraphX(g, numIter) + // Run BP using GraphFrames + val gfResults = runBPwithGraphFrames(g, numIter) + + // Check beliefs. + def checkResults(v: DataFrame): Unit = { + v.select("belief").collect().foreach { + case Row(belief: Double) => + assert( + belief >= 0.0 && belief <= 1.0, + s"Expected belief to be probability in [0,1], but found $belief") + case _ => throw new GraphFramesUnreachableException() + } + } + checkResults(gxResults.vertices) + checkResults(gfResults.vertices) + + // Compare beliefs. + val gxBeliefs = gxResults.vertices.select("id", "belief") + val gfBeliefs = gfResults.vertices.select("id", "belief") + gxBeliefs + .join(gfBeliefs, "id") + .select(gxBeliefs("belief").as("gxBelief"), gfBeliefs("belief").as("gfBelief")) + .collect() + .foreach { + case Row(gxBelief: Double, gfBelief: Double) => + assert(math.abs(gxBelief - gfBelief) <= 1e-6) + case _ => throw new GraphFramesUnreachableException() + } + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/examples/GraphsSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/examples/GraphsSuite.scala new file mode 100644 index 0000000000000..979bd2f59472d --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/examples/GraphsSuite.scala @@ -0,0 +1,56 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.examples + +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite + +class GraphsSuite extends SparkFunSuite with GraphFrameTestSparkContext { + + test("empty graph") { + for (empty <- Seq(Graphs.empty[Int], Graphs.empty[Long], Graphs.empty[String])) { + assert(empty.vertices.count() === 0L) + assert(empty.edges.count() === 0L) + } + } + + test("chain graph") { + val spark = this.spark + import spark.implicits._ + + val chain0 = Graphs.chain(0L) + assert(chain0.vertices.count() === 0L) + assert(chain0.edges.count() === 0L) + + val chain1 = Graphs.chain(1L) + assert(chain1.vertices.as[Long].collect() === Array(0L)) + assert(chain1.edges.count() === 0L) + + val chain2 = Graphs.chain(2L) + assert(chain2.vertices.as[Long].collect().toSet === Set(0L, 1L)) + assert(chain2.edges.as[(Long, Long)].collect() === Array((0L, 1L))) + + val chain3 = Graphs.chain(3L) + assert(chain3.vertices.as[Long].collect().toSet === Set(0L, 1L, 2L)) + assert(chain3.edges.as[(Long, Long)].collect().toSet === Set((0L, 1L), (1L, 2L))) + + withClue("Constructing a large chain graph shouldn't OOM the driver.") { + Graphs.chain(1e10.toLong) + } + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/examples/LDBCUtils.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/examples/LDBCUtils.scala new file mode 100644 index 0000000000000..bbf188cfc5e04 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/examples/LDBCUtils.scala @@ -0,0 +1,129 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.examples + +import java.nio.file._ + +import scala.sys.process._ + +object LDBCUtils { + // scalastyle:off println + private val LDBC_URL_PREFIX = "https://datasets.ldbcouncil.org/graphalytics/" + + val TEST_BFS_DIRECTED = "test-bfs-directed" + val TEST_BFS_UNDIRECTED = "test-bfs-undirected" + val TEST_CDLP_DIRECTED = "test-cdlp-directed" + val TEST_CDLP_UNDIRECTED = "test-cdlp-undirected" + val TEST_PR_DIRECTED = "test-pr-directed" + val TEST_PR_UNDIRECTED = "test-pr-undirected" + val TEST_WCC_DIRECTED = "test-wcc-directed" + val TEST_WCC_UNDIRECTED = "test-wcc-undirected" + val KGS = "kgs" + val GRAPH500_22 = "graph500-22" + val GRAPH500_23 = "graph500-23" + val GRAPH500_24 = "graph500-24" + val GRAPH500_25 = "graph500-25" + val GRAPH500_26 = "graph500-26" + val GRAPH500_27 = "graph500-27" + val GRAPH500_28 = "graph500-28" + val GRAPH500_29 = "graph500-29" + val GRAPH500_30 = "graph500-30" + val CIT_PATENTS = "cit-Patents" + val WIKI_TALKS = "wiki-Talk" + + private val possibleCaseNames = Set( + TEST_BFS_DIRECTED, + TEST_BFS_UNDIRECTED, + TEST_CDLP_DIRECTED, + TEST_CDLP_UNDIRECTED, + TEST_PR_DIRECTED, + TEST_PR_UNDIRECTED, + TEST_WCC_DIRECTED, + TEST_WCC_UNDIRECTED, + KGS, + GRAPH500_22, + GRAPH500_23, + GRAPH500_24, + GRAPH500_25, + GRAPH500_26, + GRAPH500_27, + GRAPH500_28, + GRAPH500_29, + GRAPH500_30, + CIT_PATENTS, + WIKI_TALKS) + + private def ldbcURL(caseName: String): String = s"${LDBC_URL_PREFIX}${caseName}.tar.zst" + + private def checkZSTD(): Unit = { + try { + val version = "zstd --version".! + println(s"found zstd version: $version") + } catch { + case e: Exception => + throw new RuntimeException( + "zstd is not available or not found. Please install zstd and try again.", + e) + } + } + + private def checkName(name: String): Unit = { + require( + possibleCaseNames.contains(name), + s"Wrong ${name}, possible names: ${possibleCaseNames.mkString(", ")}") + } + + def downloadLDBCIfNotExists(path: Path, name: String): Unit = { + checkName(name) + val dir = path.resolve(name) + if (Files.notExists(dir) || (Files.list(dir).count() == 0L)) { + println(s"LDBC data for the case ${name} not found. Downloading...") + checkZSTD() + if (Files.notExists(dir)) { + Files.createDirectories(dir) + } + val archivePath = path.resolve(s"${name}.tar.zst") + // Use curl instead of Java's URLConnection because the LDBC CDN (Cloudflare) + // rejects Java 8's TLS fingerprint with HTTP 403. + // TODO: restore URLConnection after Spark 3.5.x EOL (~April 2026) when JDK 8 can be dropped: + // val connection = new java.net.URL(ldbcURL(name)).openConnection() + // val inputStream = connection.getInputStream + // val outputStream = Files.newOutputStream(archivePath) + // val buffer = new Array[Byte](8192) + // var bytesRead = 0 + // while ({ bytesRead = inputStream.read(buffer); bytesRead } != -1) { + // outputStream.write(buffer, 0, bytesRead) + // } + // inputStream.close() + // outputStream.close() + val curlExit = s"curl -fSL -o ${archivePath.toString} ${ldbcURL(name)}".! + if (curlExit != 0) { + throw new RuntimeException( + s"Failed to download ${ldbcURL(name)} (curl exit code: $curlExit)") + } + println(s"Uncompressing ${archivePath.toString} to ${dir.toString}...") + s"zstd -d ${archivePath.toString} -o ${archivePath.toString.replace(".zst", "")}".! + s"tar -xf ${archivePath.toString.replace(".zst", "")} -C ${dir.toString}".! + + // Clean up + Files.delete(archivePath) + Files.delete(Paths.get(archivePath.toString.replace(".zst", ""))) + } + } + // scalastyle:on println +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/ldbc/TestLDBCCases.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/ldbc/TestLDBCCases.scala new file mode 100644 index 0000000000000..ecd132455e0ab --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/ldbc/TestLDBCCases.scala @@ -0,0 +1,258 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.ldbc + +import java.io.File +import java.nio.file._ +import java.util.Properties + +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.examples.LDBCUtils +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.abs +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.functions.sum +import org.apache.spark.sql.types.DoubleType +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.types.LongType +import org.apache.spark.sql.types.StructField +import org.apache.spark.sql.types.StructType + +class TestLDBCCases extends SparkFunSuite with GraphFrameTestSparkContext { + private val resourcesPath = Paths.get(new File("target").toURI) + private val unreachableID = 9223372036854775807L + + // These upstream integration tests download Graphalytics fixtures and invoke curl, zstd, and tar. + private def ldbcTest(name: String)(body: => Any): Unit = { + if (sys.env.get("SPARK_RUN_LDBC_GRAPH_TESTS").contains("1")) { + test(name)(body) + } else { + ignore(name)(body) + } + } + + private def readUndirectedUnweighted(pathPrefix: String): GraphFrame = { + var edges = spark.read + .option("delimiter", " ") + .option("header", "false") + .schema(StructType(Seq(StructField("src", LongType), StructField("dst", LongType)))) + .csv(s"${pathPrefix}.e") + .toDF("src", "dst") + + // TODO: replace by symmetrize when graphframes/graphframes#548 is done! + edges = edges + .select("src", "dst") + .union(edges.select(col("dst").alias("src"), col("src").alias("dst"))) + + val nodes = spark.read + .text(s"${pathPrefix}.v") + .toDF("id") + .select(col("id").cast(LongType)) + + GraphFrame(nodes, edges) + } + + private def readDirectedUnweighted(pathPrefix: String): GraphFrame = { + val edges = spark.read + .option("delimiter", " ") + .option("header", "false") + .schema(StructType(Seq(StructField("src", LongType), StructField("dst", LongType)))) + .csv(s"${pathPrefix}.e") + .toDF("src", "dst") + + val nodes = spark.read + .text(s"${pathPrefix}.v") + .toDF("id") + .select(col("id").cast(LongType)) + + GraphFrame(nodes, edges) + } + + private def readProperties(path: Path): Properties = { + val props = new Properties() + val stream = Files.newInputStream(path) + props.load(stream) + stream.close() + props + } + + private lazy val ldbcTestBFSDirected: (GraphFrame, DataFrame, Long) = { + LDBCUtils.downloadLDBCIfNotExists(resourcesPath, LDBCUtils.TEST_BFS_UNDIRECTED) + val caseRoot = resourcesPath.resolve(LDBCUtils.TEST_BFS_UNDIRECTED) + + val expectedPath = caseRoot.resolve(s"${LDBCUtils.TEST_BFS_UNDIRECTED}-BFS") + + val expectedDistances = spark.read + .option("delimiter", " ") + .option("header", "false") + .schema(StructType(Seq(StructField("id", LongType), StructField("distance", IntegerType)))) + .csv(expectedPath.toString) + .toDF("id", "distance") + val props = readProperties(caseRoot.resolve(s"${LDBCUtils.TEST_BFS_UNDIRECTED}.properties")) + ( + readDirectedUnweighted(s"${caseRoot.toString}/${LDBCUtils.TEST_BFS_UNDIRECTED}"), + expectedDistances, + props.getProperty(s"graph.${LDBCUtils.TEST_BFS_UNDIRECTED}.bfs.source-vertex").toLong) + } + + Seq("graphframes", "graphx").foreach { algo => + ldbcTest(s"test undirected BFS with LDBC for impl ${algo}") { + val testCase = ldbcTestBFSDirected + val srcVertex = testCase._3 + + // this graph is undirected, but in GF direction exists + // only on the level of algorithms! + val spResult = testCase._1.shortestPaths + .landmarks(Seq(srcVertex)) + .setAlgorithm(algo) + .setIsDirected(false) + .run() + .select( + col(GraphFrame.ID), + col("distances").getItem(srcVertex).cast(LongType).alias("got_distance")) + .na + .fill(Map("got_distance" -> unreachableID)) + + assert(spResult.count() == testCase._1.vertices.count()) + assert( + spResult + .join(testCase._2, Seq("id"), "left") + .filter(col("got_distance") =!= col("distance")) + .collect() + .isEmpty) + + } + } + + private lazy val ldbcTestCDLPUndirected: (GraphFrame, DataFrame, Int) = { + LDBCUtils.downloadLDBCIfNotExists(resourcesPath, LDBCUtils.TEST_CDLP_UNDIRECTED) + val caseRoot = resourcesPath.resolve(LDBCUtils.TEST_CDLP_UNDIRECTED) + + val expectedPath = caseRoot.resolve(s"${LDBCUtils.TEST_CDLP_UNDIRECTED}-CDLP") + + val expectedCommunities = spark.read + .option("delimiter", " ") + .option("header", "false") + .schema(StructType(Seq(StructField("id", LongType), StructField("community", LongType)))) + .csv(expectedPath.toString) + .toDF("id", "community") + val props = readProperties(caseRoot.resolve(s"${LDBCUtils.TEST_CDLP_UNDIRECTED}.properties")) + ( + readUndirectedUnweighted(s"${caseRoot.toString}/${LDBCUtils.TEST_CDLP_UNDIRECTED}"), + expectedCommunities, + props.getProperty(s"graph.${LDBCUtils.TEST_CDLP_UNDIRECTED}.cdlp.max-iterations").toInt) + } + + Seq("graphx", "graphframes").foreach { algo => + ldbcTest(s"test undirected CDLP with LDBC for algo ${algo}") { + val testCase = ldbcTestCDLPUndirected + val cdlpResults = testCase._1.labelPropagation.setAlgorithm(algo).maxIter(testCase._3).run() + assert(cdlpResults.count() == testCase._1.vertices.count()) + assert( + cdlpResults + .join(testCase._2, Seq("id"), "left") + .filter(col("label") =!= col("community")) + .collect() + .isEmpty) + } + } + + private lazy val ldbcTestPageRankUndirected: (GraphFrame, DataFrame, Double, Int) = { + LDBCUtils.downloadLDBCIfNotExists(resourcesPath, LDBCUtils.TEST_PR_UNDIRECTED) + val caseRoot = resourcesPath.resolve(LDBCUtils.TEST_PR_UNDIRECTED) + + val expectedPath = caseRoot.resolve(s"${LDBCUtils.TEST_PR_UNDIRECTED}-PR") + + val expectedRanks = spark.read + .option("delimiter", " ") + .option("header", "false") + .schema(StructType(Seq(StructField("id", LongType), StructField("pr", DoubleType)))) + .csv(expectedPath.toString) + .toDF("id", "pr") + + val props = readProperties(caseRoot.resolve(s"${LDBCUtils.TEST_PR_UNDIRECTED}.properties")) + ( + readUndirectedUnweighted(s"${caseRoot.toString}/${LDBCUtils.TEST_PR_UNDIRECTED}"), + expectedRanks, + props.getProperty(s"graph.${LDBCUtils.TEST_PR_UNDIRECTED}.pr.damping-factor").toDouble, + props.getProperty(s"graph.${LDBCUtils.TEST_PR_UNDIRECTED}.pr.num-iterations").toInt) + } + + // TODO: add graphframes after finishing graphframes/graphframes#569 + Seq("graphx").foreach { algo => + ldbcTest(s"test undirected PR with LDBC for algo ${algo}") { + val testCase = ldbcTestPageRankUndirected + val prResults = testCase._1.pageRank + .resetProbability(1.0 - testCase._3) + .maxIter(testCase._4) + .run() + .vertices + + // Normalize?? + val sumPR = prResults.agg(sum(col("pagerank"))).collect().head.getAs[Double](0) + val prResultsNormalized = prResults.withColumn("pagerank", col("pagerank") / lit(sumPR)) + assert(prResults.count() == testCase._1.vertices.count()) + assert( + prResultsNormalized + .join(testCase._2, Seq("id"), "left") + .filter(abs(col("pagerank") - col("pr")) >= lit(1e-4)) + .collect() + .isEmpty) + } + } + + private lazy val ldbcTestWCCUndirected: (GraphFrame, DataFrame) = { + LDBCUtils.downloadLDBCIfNotExists(resourcesPath, LDBCUtils.TEST_WCC_UNDIRECTED) + val caseRoot = resourcesPath.resolve(LDBCUtils.TEST_WCC_UNDIRECTED) + + val expectedPath = caseRoot.resolve(s"${LDBCUtils.TEST_WCC_UNDIRECTED}-WCC") + + val expectedComponents = spark.read + .option("delimiter", " ") + .option("header", "false") + .schema(StructType(Seq(StructField("id", LongType), StructField("wcomp", LongType)))) + .csv(expectedPath.toString) + .toDF("id", "wcomp") + + ( + readUndirectedUnweighted(s"${caseRoot.toString}/${LDBCUtils.TEST_WCC_UNDIRECTED}"), + expectedComponents) + } + + Seq("two_phase", "graphx", "randomized_contraction").foreach { algo => + ldbcTest(s"test undirected WCC with LDBC for impl ${algo}") { + val testCase = ldbcTestWCCUndirected + var cc = testCase._1.connectedComponents.setAlgorithm(algo) + if (algo == "randomized_contraction") { + // RC is randomized by it's nature; + cc = cc.setUseLabelsAsComponents(true) + } + val ccResults = cc.run() + assert(ccResults.count() == testCase._1.vertices.count()) + assert( + ccResults + .join(testCase._2, Seq("id"), "left") + .filter(col("wcomp") =!= col("component")) + .collect() + .isEmpty) + } + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/propertygraph/PropertyGraphFrameSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/propertygraph/PropertyGraphFrameSuite.scala new file mode 100644 index 0000000000000..5f219ef8697a6 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/propertygraph/PropertyGraphFrameSuite.scala @@ -0,0 +1,299 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.graphframes.propertygraph + +import java.security.MessageDigest + +import org.scalatest.BeforeAndAfterAll + +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.propertygraph.property.EdgePropertyGroup +import org.apache.spark.graphframes.propertygraph.property.VertexPropertyGroup +import org.apache.spark.sql.Column +import org.apache.spark.sql.functions._ + +class PropertyGraphFrameSuite + extends SparkFunSuite + with GraphFrameTestSparkContext + with BeforeAndAfterAll { + var peopleMoviesGraph: PropertyGraphFrame = _ + + override def beforeAll(): Unit = { + super.beforeAll() + + // This graph represents a movie rating system with two types of vertices: 'people' + // (5 users: Alice, Bob, Charlie, David, Eve) and 'movies' (3 movies: Matrix, Inception, + // Interstellar). The graph has two types of edges: + // 1) 'likes' - undirected edges between people and movies with weight 1.0, representing + // movie preferences + // 2) 'messages' - directed edges between people with varying weights (0.3-0.9), representing + // communication patterns. The people-movie connections form a bipartite subgraph, while + // the messages form a directed cycle between users. + + val peopleData = spark + .createDataFrame( + Seq((1L, "Alice"), (2L, "Bob"), (3L, "Charlie"), (4L, "David"), (5L, "Eve"))) + .toDF("id", "name") + + val peopleGroup = VertexPropertyGroup("people", peopleData, "id") + + val moviesData = spark + .createDataFrame(Seq((1L, "Matrix"), (2L, "Inception"), (3L, "Interstellar"))) + .toDF("id", "title") + + val moviesGroup = VertexPropertyGroup("movies", moviesData, "id") + + val likesData = spark + .createDataFrame(Seq((1L, 1L), (1L, 2L), (2L, 1L), (3L, 2L), (4L, 3L), (5L, 2L))) + .toDF("src", "dst") + + val likesGroup = EdgePropertyGroup( + "likes", + likesData, + peopleGroup, + moviesGroup, + isDirected = false, + "src", + "dst", + lit(1.0)) + + val messagesData = spark + .createDataFrame( + Seq((1L, 2L, 5.0), (2L, 3L, 8.0), (3L, 4L, 3.0), (4L, 5L, 6.0), (5L, 1L, 9.0))) + .toDF("src", "dst", "weight") + + val messagesGroup = EdgePropertyGroup( + "messages", + messagesData, + peopleGroup, + peopleGroup, + isDirected = true, + "src", + "dst", + col("weight")) + + peopleMoviesGraph = + PropertyGraphFrame(Seq(peopleGroup, moviesGroup), Seq(likesGroup, messagesGroup)) + } + + test("projection by movies creates correct graph structure") { + val projectedGraph = peopleMoviesGraph.projectionBy("people", "movies", "likes") + + assert(projectedGraph.vertexPropertyGroups.length === 1) + assert(projectedGraph.vertexPropertyGroups.head.name === "people") + + assert(projectedGraph.edgesPropertyGroups.length === 2) + assert(projectedGraph.edgesPropertyGroups.exists(_.name === "messages")) + val projectedEdgesGroupOption = + projectedGraph.edgesPropertyGroups.find(_.name === "projected_likes") + + assert(projectedEdgesGroupOption.isDefined) + val projectedEdgesGroup = projectedEdgesGroupOption.get + + assert(projectedEdgesGroup.srcColumnName === GraphFrame.SRC) + assert(projectedEdgesGroup.dstColumnName === GraphFrame.DST) + assert(projectedEdgesGroup.weightColumnName === GraphFrame.WEIGHT) + assert(!projectedEdgesGroup.isDirected) + + val projectedEdges = projectedEdgesGroup.data + .collect() + .map(row => (row.getLong(0), row.getLong(1))) + .toSet + + // Expected edges between people who like the same movies + val expectedEdges = Set( + (1L, 2L), // Alice and Bob both like Matrix + (1L, 3L), // Alice and Charlie both like Inception + (1L, 5L), // Alice and Eve both like Inception + (3L, 5L) // Charlie and Eve both like Inception + ) + + assert(projectedEdges === expectedEdges) + } + + def sha256Hash(id: Long, groupName: String): String = { + val md = MessageDigest.getInstance("SHA-256") + val hash = md.digest(id.toString.getBytes("UTF-8")).map("%02x".format(_)).mkString + s"$groupName$hash" + } + + test("toGraphFrame with messages edges and people vertices only") { + val graph = peopleMoviesGraph.toGraphFrame( + Seq("people"), + Seq("messages"), + Map("messages" -> lit(true)), + Map("people" -> lit(true))) + + val vertices = graph.vertices.collect().map(row => row.getString(0)).toSet + val edges = graph.edges + .collect() + .map(row => (row.getString(0), row.getString(1), row.getDouble(2))) + .toSet + + // Verify vertices (all people) + val expectedVertices = Set(1L, 2L, 3L, 4L, 5L).map(sha256Hash(_, "people")) + assert(vertices === expectedVertices) + + // Verify directed message edges with weights + val expectedEdges = + Set((1L, 2L, 5.0), (2L, 3L, 8.0), (3L, 4L, 3.0), (4L, 5L, 6.0), (5L, 1L, 9.0)).map { + case (src, dst, weight) => (sha256Hash(src, "people"), sha256Hash(dst, "people"), weight) + } + assert(edges === expectedEdges) + } + + test("toGraphFrame with all groups and proper edge handling") { + val graph = peopleMoviesGraph.toGraphFrame( + Seq("people", "movies"), + Seq("messages", "likes"), + Map("messages" -> lit(true), "likes" -> lit(true)), + Map("people" -> lit(true), "movies" -> lit(true))) + + val vertices = graph.vertices.collect().toSet + val edges = graph.edges.collect().toSet + + // Verify all vertices are present + assert(vertices.size === 8) // 5 people + 3 movies + + // Verify vertex types are correctly preserved + assert(vertices.count(_.getString(0) == sha256Hash(1L, "movies")) === 1) + assert(vertices.count(_.getString(0) == sha256Hash(1L, "people")) === 1) + + // Verify edge counts and properties + val messageEdges = edges.filter(_.getDouble(2) != 1.0) + val likeEdges = edges.filter(_.getDouble(2) == 1.0) + + assert(messageEdges.size === 5) // Directed messages between people + assert(likeEdges.size === 12) // 6 original edges * 2 (undirected converted to directed) + + // Verify undirected edges were properly converted to directed pairs + val likesPairs = likeEdges.map(row => (row.getString(0), row.getString(1))).toSet + assert( + likesPairs.contains((sha256Hash(1, "people"), sha256Hash(1, "movies"))) && + likesPairs.contains((sha256Hash(1, "movies"), sha256Hash(1, "people")))) + assert( + likesPairs.contains((sha256Hash(1, "people"), sha256Hash(2, "movies"))) && + likesPairs.contains((sha256Hash(2, "movies"), sha256Hash(1, "people")))) + } + + test("toGraphFrame preserves original IDs when masking disabled for vertex group") { + // Create new movies group with masking disabled + val unmaskedMoviesGroup = VertexPropertyGroup( + "movies", + peopleMoviesGraph.vertexPropertyGroups.find(_.name == "movies").get.data, + "id", + applyMaskOnId = false) + + // Create new likes group with unmasked movies group + val oldLikesGroup = peopleMoviesGraph.edgesPropertyGroups.find(_.name == "likes").get + val newLikesGroup = EdgePropertyGroup( + "likes", + oldLikesGroup.data, + oldLikesGroup.srcPropertyGroup, + unmaskedMoviesGroup, + oldLikesGroup.isDirected, + oldLikesGroup.srcColumnName, + oldLikesGroup.dstColumnName, + oldLikesGroup.weightColumnName) + + // Create new graph with unmasked movies group and updated likes group + val modifiedGraph = PropertyGraphFrame( + peopleMoviesGraph.vertexPropertyGroups.filterNot(_.name == "movies") :+ unmaskedMoviesGroup, + peopleMoviesGraph.edgesPropertyGroups.filterNot(_.name == "likes") :+ newLikesGroup) + + val graph = modifiedGraph.toGraphFrame( + Seq("people", "movies"), + Seq("messages", "likes"), + Map("messages" -> lit(true), "likes" -> lit(true)), + Map("people" -> lit(true), "movies" -> lit(true))) + + val vertices = graph.vertices.collect().map(_.getString(0)).toSet + val edges = graph.edges.collect().toSet + + // Verify movies vertices have original IDs + assert(vertices.contains("1")) + assert(vertices.contains("2")) + assert(vertices.contains("3")) + + // Verify people vertices are masked + assert(vertices.contains(sha256Hash(1L, "people"))) + + // Verify edges have masked people IDs but original movie IDs + val likesEdges = edges.filter(_.getDouble(2) == 1.0) + assert( + likesEdges.exists(e => e.getString(0) == sha256Hash(1L, "people") && e.getString(1) == "1")) + assert( + likesEdges.exists(e => e.getString(0) == "1" && e.getString(1) == sha256Hash(1L, "people"))) + } + + test("projection with custom weight function") { + val projectedGraph = peopleMoviesGraph.projectionBy( + "people", + "movies", + "likes", + Some((leftWeight: Column, rightWeight: Column) => leftWeight + rightWeight)) + + val projectedEdgesGroupOption = + projectedGraph.edgesPropertyGroups.find(_.name === "projected_likes") + assert(projectedEdgesGroupOption.isDefined) + + val projectedEdges = projectedEdgesGroupOption.get.data + .collect() + .map(row => (row.getLong(0), row.getLong(1), row.getDouble(2))) + .toSet + + // Expected edges between people who like the same movies with sum of their weights + val expectedEdges = Set( + (1L, 2L, 2.0), // Alice and Bob both like Matrix (1.0 + 1.0) + (1L, 3L, 2.0), // Alice and Charlie both like Inception (1.0 + 1.0) + (1L, 5L, 2.0), // Alice and Eve both like Inception (1.0 + 1.0) + (3L, 5L, 2.0) // Charlie and Eve both like Inception (1.0 + 1.0) + ) + + assert(projectedEdges === expectedEdges) + } + + test("joinVertices withConnectedComponents") { + // Convert to GraphFrame with all vertices and edges + val graph = peopleMoviesGraph.toGraphFrame( + Seq("people", "movies"), + Seq("messages", "likes"), + Map("messages" -> lit(true), "likes" -> lit(true)), + Map("people" -> lit(true), "movies" -> lit(true))) + + // Compute connected components + val components = graph.connectedComponents.run() + + val joinedBack = peopleMoviesGraph + .joinVertices(components, Seq("people", "movies")) + .select( + PropertyGraphFrame.EXTERNAL_ID, + "component", + PropertyGraphFrame.PROPERTY_GROUP_COL_NAME) + .collect() + .map(r => Tuple3(r.getLong(0), r.getLong(1), r.getString(2))) + .groupBy(_._3) + + assert(joinedBack.contains("movies")) + assert(joinedBack.contains("people")) + assert(joinedBack("movies").length == 3) + assert(joinedBack("people").length == 5) + } +} diff --git a/graphframes/src/test/scala/org/apache/spark/sql/graphframes/expressions/KMinSamplingSuite.scala b/graphframes/src/test/scala/org/apache/spark/sql/graphframes/expressions/KMinSamplingSuite.scala new file mode 100644 index 0000000000000..c4156b2f794d6 --- /dev/null +++ b/graphframes/src/test/scala/org/apache/spark/sql/graphframes/expressions/KMinSamplingSuite.scala @@ -0,0 +1,78 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.graphframes.expressions + +import scala.util.Random + +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types.LongType + +class KMinSamplingSuite extends SparkFunSuite with GraphFrameTestSparkContext { + test("test kmin sampling") { + val data = Seq( + (1L, 2L, 1L), + (1L, 3L, 2L), + (1L, 4L, 3L), + (1L, 5L, 2L), + (2L, 1L, 1L), + (2L, 4L, 2L), + (3L, 1L, 1L), + (4L, 2L, 2L)) + + val toAgg = spark.createDataFrame(data).toDF("src", "dst", "weight") + val encoder = KMinSampling.getEncoder(spark, LongType, Seq("dst", "weight")) + val aggUDF = KMinSampling.fromSparkType(LongType, 3, encoder) + + val result = + toAgg.groupBy("src").agg(aggUDF(col("dst"), col("weight")).alias("r")) + + val collectedResult = result.collect() + val collectedMap = collectedResult.map(f => (f.getLong(0), f.getAs[Seq[Long]](1))).toMap + + assert(collectedMap.get(1L).get === Seq(2L, 3L, 5L)) + assert(collectedMap.get(2L).get === Seq(1L, 4L)) + assert(collectedMap.get(3L).get === Seq(1L)) + assert(collectedMap.get(4L).get === Seq(2L)) + } + + test("test kmin sampling many values") { + val random = new Random(42L) + val candidates = Array(1L, 2L, 3L, 4L, 5L, 6L) + val data = (1L to 10L) + .flatMap(id => (1 to 100).map(_ => (id, candidates(random.nextInt(5)), random.nextLong()))) + .toSeq + + val toAgg = spark.createDataFrame(data).toDF("src", "dst", "weight") + val encoder = KMinSampling.getEncoder(spark, LongType, Seq("dst", "weight")) + val aggUDF = KMinSampling.fromSparkType(LongType, 5, encoder) + + val result = + toAgg.groupBy("src").agg(aggUDF(col("dst"), col("weight")).alias("r")).collect() + + val collectedMap = result.map(f => (f.getLong(0), f.getAs[Seq[Long]](1))).toMap + // at least one should be full + assert(collectedMap.map(f => f._2.size).max == 5) + + // all should be within the limit + for (id <- (1L to 10L)) { + assert(collectedMap.get(id).get.size <= 5) + } + } +} diff --git a/python/packaging/classic/setup.py b/python/packaging/classic/setup.py index add11863f000a..9cd2aba29314b 100755 --- a/python/packaging/classic/setup.py +++ b/python/packaging/classic/setup.py @@ -285,6 +285,13 @@ def run(self): "pyspark.ml.param", "pyspark.ml.torch", "pyspark.ml.deepspeed", + "pyspark.graphframes", + "pyspark.graphframes.classic", + "pyspark.graphframes.connect", + "pyspark.graphframes.examples", + "pyspark.graphframes.internal", + "pyspark.graphframes.lib", + "pyspark.graphframes.pg", "pyspark.sql", "pyspark.sql.avro", "pyspark.sql.classic", diff --git a/python/packaging/client/setup.py b/python/packaging/client/setup.py index af9e9475018c8..9de1f26fea200 100755 --- a/python/packaging/client/setup.py +++ b/python/packaging/client/setup.py @@ -72,6 +72,9 @@ "pyspark.errors.tests.connect", "pyspark.tests", # for Memory profiler parity tests "pyspark.resource.tests", + "pyspark.graphframes.tests", + "pyspark.graphframes.tests.connect", + "pyspark.graphframes.tests.pg", "pyspark.sql.tests", "pyspark.sql.tests.arrow", "pyspark.sql.tests.connect", @@ -157,6 +160,12 @@ "pyspark.ml.param", "pyspark.ml.torch", "pyspark.ml.deepspeed", + "pyspark.graphframes", + "pyspark.graphframes.connect", + "pyspark.graphframes.examples", + "pyspark.graphframes.internal", + "pyspark.graphframes.lib", + "pyspark.graphframes.pg", "pyspark.sql", "pyspark.sql.avro", "pyspark.sql.connect", diff --git a/python/pyspark/graphframes/classic/graphframe.py b/python/pyspark/graphframes/classic/graphframe.py index edc5fcc737b5a..92ba9baeddf56 100644 --- a/python/pyspark/graphframes/classic/graphframe.py +++ b/python/pyspark/graphframes/classic/graphframe.py @@ -21,9 +21,9 @@ from py4j.java_gateway import JavaObject from pyspark import SparkContext +from pyspark.graphframes.classic.pregel import Pregel from pyspark.graphframes.classic.utils import storage_level_to_jvm from pyspark.graphframes.internal.utils import _RandomWalksEmbeddingsParameters -from pyspark.graphframes.lib import Pregel from pyspark.sql import SparkSession from pyspark.sql import functions as F from pyspark.sql.classic.column import Column, _to_seq diff --git a/python/pyspark/graphframes/classic/pregel.py b/python/pyspark/graphframes/classic/pregel.py new file mode 100644 index 0000000000000..1863bb9bb9849 --- /dev/null +++ b/python/pyspark/graphframes/classic/pregel.py @@ -0,0 +1,342 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from typing import TYPE_CHECKING, Any, final + +from typing_extensions import Self + +from pyspark.graphframes.classic.utils import storage_level_to_jvm +from pyspark.ml.wrapper import JavaWrapper +from pyspark.sql import Column, DataFrame, SparkSession +from pyspark.sql.classic.column import _to_seq +from pyspark.sql.functions import col +from pyspark.storagelevel import StorageLevel + +if TYPE_CHECKING: + from pyspark.graphframes.classic.graphframe import GraphFrame + + +@final +class Pregel(JavaWrapper): + """Implements a Pregel-like bulk-synchronous message-passing API based on DataFrame operations. + + See `Malewicz et al., Pregel: a system for large-scale graph processing `_ + for a detailed description of the Pregel algorithm. + + You can construct a Pregel instance using either this constructor or :attr:`graphframes.GraphFrame.pregel`, + then use builder pattern to describe the operations, and then call :func:`run` to start a run. + It returns a DataFrame of vertices from the last iteration. + + When a run starts, it expands the vertices DataFrame using column expressions defined by :func:`withVertexColumn`. + Those additional vertex properties can be changed during Pregel iterations. + In each Pregel iteration, there are three phases: + + - Given each edge triplet, generate messages and specify target vertices to send, described by :func:`sendMsgToDst` and :func:`sendMsgToSrc`. + - Aggregate messages by target vertex IDs, described by :func:`aggMsgs`. + - Update additional vertex properties based on aggregated messages and states from previous iteration, described by :func:`withVertexColumn`. + + Please find what columns you can reference at each phase in the method API docs. + + You can control the number of iterations by :func:`setMaxIter` and check API docs for advanced controls. + + :param graph: a :class:`graphframes.GraphFrame` object holding a graph with vertices and edges stored as DataFrames. + """ + + def __init__(self, graph: "GraphFrame") -> None: + super(Pregel, self).__init__() + + self.graph = graph + self._java_obj: Any = self._new_java_obj( + "org.apache.spark.graphframes.lib.Pregel", graph._jvm_graph + ) + + def setMaxIter(self, value: int) -> "Pregel": + """Sets the max number of iterations (default: 10). + + :param value: the number of Pregel iterations + """ + self._java_obj.setMaxIter(int(value)) + return self + + def setCheckpointInterval(self, value: int) -> "Pregel": + """Sets the number of iterations between two checkpoints (default: 2). + + This is an advanced control to balance query plan optimization and checkpoint data I/O cost. + In most cases, you should keep the default value. + + Checkpoint is disabled if this is set to 0. + """ + self._java_obj.setCheckpointInterval(int(value)) + return self + + def setEarlyStopping(self, value: bool) -> "Pregel": + """Set should Pregel stop earlier in case of no new messages to send or not. + + Early stopping allows to terminate Pregel before reaching maxIter by checking if there are any non-null messages. + While in some cases it may gain significant performance boost, in other cases it can lead to performance degradation, + because checking if the messages DataFrame is empty or not is an action and requires materialization of the Spark Plan + with some additional computations. + + In the case when the user can assume a good value of maxIter, it is recommended to leave this value to the default "false". + In the case when it is hard to estimate the number of iterations required for convergence, + it is recommended to set this value to "false" to avoid iterating over convergence until reaching maxIter. + When this value is "true", maxIter can be set to a bigger value without risks. + """ + self._java_obj.setEarlyStopping(bool(value)) + return self + + def withVertexColumn( + self, colName: str, initialExpr: Column, updateAfterAggMsgsExpr: Column + ) -> "Pregel": + """Defines an additional vertex column at the start of run and how to update it in each iteration. + + You can call it multiple times to add more than one additional vertex columns. + + :param colName: the name of the additional vertex column. + It cannot be an existing vertex column in the graph. + :param initialExpr: the expression to initialize the additional vertex column. + You can reference all original vertex columns in this expression. + :param updateAfterAggMsgsExpr: the expression to update the additional vertex column after messages aggregation. + You can reference all original vertex columns, additional vertex columns, and the + aggregated message column using :func:`msg`. + If the vertex received no messages, the message column would be null. + """ + self._java_obj.withVertexColumn(colName, initialExpr._jc, updateAfterAggMsgsExpr._jc) + return self + + def sendMsgToSrc(self, msgExpr: Column) -> "Pregel": + """Defines a message to send to the source vertex of each edge triplet. + + You can call it multiple times to send more than one messages. + + See method :func:`sendMsgToDst`. + + :param msgExpr: the expression of the message to send to the source vertex given a (src, edge, dst) triplet. + Source/destination vertex properties and edge properties are nested under columns `src`, `dst`, + and `edge`, respectively. + You can reference them using :func:`src`, :func:`dst`, and :func:`edge`. + Null messages are not included in message aggregation. + """ + self._java_obj.sendMsgToSrc(msgExpr._jc) + return self + + def sendMsgToDst(self, msgExpr: Column) -> "Pregel": + """Defines a message to send to the destination vertex of each edge triplet. + + You can call it multiple times to send more than one messages. + + See method :func:`sendMsgToSrc`. + + :param msgExpr: the message expression to send to the destination vertex given a (`src`, `edge`, `dst`) triplet. + Source/destination vertex properties and edge properties are nested under columns `src`, `dst`, + and `edge`, respectively. + You can reference them using :func:`src`, :func:`dst`, and :func:`edge`. + Null messages are not included in message aggregation. + """ + self._java_obj.sendMsgToDst(msgExpr._jc) + return self + + def aggMsgs(self, aggExpr: Column) -> "Pregel": + """Defines how messages are aggregated after grouped by target vertex IDs. + + :param aggExpr: the message aggregation expression, such as `sum(Pregel.msg())`. + You can reference the message column by :func:`msg` and the vertex ID by `col("id")`, + while the latter is usually not used. + """ + self._java_obj.aggMsgs(aggExpr._jc) + return self + + def setStopIfAllNonActiveVertices(self, value: bool) -> Self: + """Set should Pregel stop if all the vertices voted to halt. + + Activity (or vote) is determined based on the activity_col. + See methods :func:`setInitialActiveVertexExpression` and :func:`setUpdateActiveVertexExpression` for details + how to set and update activity_col. + + Be aware that checking of the vote is not free but a Spark Action. In case the + condition is not realistically reachable but set, it will just slow down the algorithm. + + :param value: the boolean value. + """ + self._java_obj.setStopIfAllNonActiveVertices(value) + return self + + def setInitialActiveVertexExpression(self, value: Column) -> Self: + """Sets the initial expression for the active vertex column. + + The active vertex column is used to determine if a vertices voting result on each iteration of Pregel. + This expression is evaluated on the initial vertices DataFrame to set the initial state of the activity column. + + :param value: expression to compute the initial active state of vertices. + You can reference all original vertex columns in this expression. + """ + self._java_obj.setInitialActiveVertexExpression(value._jc) + return self + + def setUpdateActiveVertexExpression(self, value: Column) -> Self: + """Sets the expression to update the active vertex column. + + The active vertex column is used to determine if a vertices voting result on each iteration of Pregel. + This expression is evaluated on the updated vertices DataFrame to set the new state of the activity column. + + :param value: expression to compute the new active state of vertices. + You can reference all original vertex columns and additional vertex columns in this expression. + """ + self._java_obj.setUpdateActiveVertexExpression(value._jc) + return self + + def setSkipMessagesFromNonActiveVertices(self, value: bool) -> Self: + """Set should Pregel skip sending messages from non-active vertices. + + When this option is enabled, messages will not be sent from vertices that are marked as inactive. + This can help optimize performance by avoiding unnecessary message propagation from inactive vertices. + + :param value: boolean value. + """ + self._java_obj.setSkipMessagesFromNonActiveVertices(value) + return self + + def setUseLocalCheckpoints(self, value: bool) -> Self: + """Set should Pregel use local checkpoints. + + Local checkpoints are faster and do not require configuring a persistent storage. + At the same time, local checkpoints are less reliable and may create a big load on local disks of executors. + + :param value: boolean value. + """ + self._java_obj.setUseLocalCheckpoints(value) + return self + + def setIntermediateStorageLevel(self, storage_level: StorageLevel) -> Self: + """Set the intermediate storage level. + On each iteration, Pregel cache results with a requested storage level. + + For very big graphs it is recommended to use DISK_ONLY. + + :param storage_level: storage level to use. + """ + self._java_obj.setIntermediateStorageLevel( + storage_level_to_jvm(storage_level, self.graph._spark) + ) + return self + + def required_src_columns(self, col_name: str, *col_names: str) -> Self: + """Specifies which source vertex columns are required when constructing triplets. + + By default, all source vertex columns are included in triplets, which can create large + intermediate datasets for algorithms with significant state (e.g., cycle detection, + random walks). Use this method to reduce memory usage by specifying only the columns + that are actually needed by the sendMsgToSrc and sendMsgToDst expressions. + + The ID column and the active flag column (if used) are always included automatically. + + :param col_name: the first required source vertex column name + :param col_names: additional required source vertex column names + + See also :func:`required_dst_columns` + """ + self._java_obj.requiredSrcColumns( + col_name, _to_seq(self.graph._spark.sparkContext, col_names) + ) + return self + + def required_dst_columns(self, col_name: str, *col_names: str) -> Self: + """Specifies which destination vertex columns are required when constructing triplets. + + By default, all destination vertex columns are included in triplets, which can create large + intermediate datasets for algorithms with significant state (e.g., cycle detection, + random walks). Use this method to reduce memory usage by specifying only the columns + that are actually needed by the sendMsgToSrc and sendMsgToDst expressions. + + The ID column and the active flag column (if used) are always included automatically. + + :param col_name: the first required destination vertex column name + :param col_names: additional required destination vertex column names + + See also :func:`required_src_columns` + """ + self._java_obj.requiredDstColumns( + col_name, _to_seq(self.graph._spark.sparkContext, col_names) + ) + return self + + def required_edge_columns(self, col_name: str, *col_names: str) -> Self: + """Specifies which edge columns are required when constructing triplets. + + By default, only src and dst columns are included from edges. Use this method to + specify additional edge columns that are needed by the sendMsgToSrc and sendMsgToDst + expressions. + + :param col_name: the first required edge column name + :param col_names: additional required edge column names + + See also :func:`required_src_columns` and :func:`required_dst_columns` + """ + self._java_obj.requiredEdgeColumns( + col_name, _to_seq(self.graph._spark.sparkContext, col_names) + ) + return self + + def run(self) -> DataFrame: + """Runs the defined Pregel algorithm. + + :return: the result vertex DataFrame from the final iteration including both original and additional columns. + """ + spark = SparkSession.getActiveSession() + if spark is None: + raise ValueError("SparkSession is dead or did not started.") + return DataFrame(self._java_obj.run(), spark) + + @staticmethod + def msg() -> Column: + """References the message column in aggregating messages and updating additional vertex columns. + + See :func:`aggMsgs` and :func:`withVertexColumn` + """ + return col("_pregel_msg_") + + @staticmethod + def src(colName: str) -> Column: + """References a source vertex column in generating messages to send. + + See :func:`sendMsgToSrc` and :func:`sendMsgToDst` + + :param colName: the vertex column name. + """ + return col("src." + colName) + + @staticmethod + def dst(colName: str) -> Column: + """ + References a destination vertex column in generating messages to send. + + See :func:`sendMsgToSrc` and :func:`sendMsgToDst` + + :param colName: the vertex column name. + """ + return col("dst." + colName) + + @staticmethod + def edge(colName: str) -> Column: + """ + References an edge column in generating messages to send. + + See :func:`sendMsgToSrc` and :func:`sendMsgToDst` + + :param colName: the edge column name. + """ + return col("edge." + colName) diff --git a/python/pyspark/graphframes/examples/__init__.py b/python/pyspark/graphframes/examples/__init__.py new file mode 100644 index 0000000000000..ed168e8d61f3c --- /dev/null +++ b/python/pyspark/graphframes/examples/__init__.py @@ -0,0 +1,21 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from .belief_propagation import BeliefPropagation +from .graphs import Graphs + +__all__ = ["BeliefPropagation", "Graphs"] diff --git a/python/pyspark/graphframes/examples/belief_propagation.py b/python/pyspark/graphframes/examples/belief_propagation.py new file mode 100644 index 0000000000000..b9aa54d48db49 --- /dev/null +++ b/python/pyspark/graphframes/examples/belief_propagation.py @@ -0,0 +1,177 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import math +from typing import Union + +# Import subpackage examples here explicitly so that +# this module can be run directly with spark-submit. +import pyspark.graphframes.examples +from pyspark.graphframes import GraphFrame +from pyspark.graphframes.lib import AggregateMessages as AM +from pyspark.sql import SparkSession, types +from pyspark.sql import functions as sqlfunctions + +__all__ = ["BeliefPropagation"] + + +class BeliefPropagation: + r"""Example code for Belief Propagation (BP) + + This provides a template for building customized BP algorithms for different types of graphical + models. + + This example: + + * Ising model on a grid + * Parallel Belief Propagation using colored fields + + Ising models are probabilistic graphical models over binary variables + (see :meth:`Graphs.gridIsingModel()`). + + Belief Propagation (BP) provides marginal probabilities of the values of the variables + x\ :sub:`i` i.e., P(x\ :sub:`i`) for each i. This allows a user to understand likely values of + variables. See `Wikipedia `__ for more + information on BP. + + We use a batch synchronous BP algorithm, where batches of vertices are updated synchronously. + We follow the mean field update algorithm in Slide 13 of the + `talk slides `__ from: + Wainwright. "Graphical models, message-passing algorithms, and convex optimization." + + The batches are chosen according to a coloring. For background on graph colorings for + inference, see for example: Gonzalez et al. "Parallel Gibbs Sampling: From Colored Fields to + Thin Junction Trees." AISTATS, 2011. + + The BP algorithm works by: + + * Coloring the graph by assigning a color to each vertex such that no neighboring vertices + share the same color. + * In each step of BP, update all vertices of a single color. Alternate colors. + """ + + @classmethod + def runBPwithGraphFrames(cls, g: GraphFrame, numIter: int) -> GraphFrame: + """Run Belief Propagation using GraphFrame. + + This implementation of BP shows how to use GraphFrame's aggregateMessages method. + """ + # choose colors for vertices for BP scheduling + colorG = cls._colorGraph(g) + numColors = colorG.vertices.select("color").distinct().count() + + # TODO: handle vertices without any edges + + # initialize vertex beliefs at 0.0 + gx = GraphFrame(colorG.vertices.withColumn("belief", sqlfunctions.lit(0.0)), colorG.edges) + + # run BP for numIter iterations + for iter_ in range(numIter): + # for each color, have that color receive messages from neighbors + for color in range(numColors): + # Send messages to vertices of the current color. + # We may send to source or destination since edges are treated as undirected. + msgForSrc = sqlfunctions.when( + AM.src["color"] == color, AM.edge["b"] * AM.dst["belief"] + ) + msgForDst = sqlfunctions.when( + AM.dst["color"] == color, AM.edge["b"] * AM.src["belief"] + ) + # numerically stable sigmoid + logistic = sqlfunctions.udf(cls._sigmoid, returnType=types.DoubleType()) + aggregates = gx.aggregateMessages( + sqlfunctions.sum(AM.msg).alias("aggMess"), + sendToSrc=msgForSrc, + sendToDst=msgForDst, + ) + v = gx.vertices + # receive messages and update beliefs for vertices of the current color + newBeliefCol = sqlfunctions.when( + (v["color"] == color) & (aggregates["aggMess"].isNotNull()), + logistic(aggregates["aggMess"] + v["a"]), + ).otherwise(v["belief"]) # keep old beliefs for other colors + newVertices = ( + v.join(aggregates, on=(v["id"] == aggregates["id"]), how="left_outer") + .drop(aggregates["id"]) # drop duplicate ID column (from outer join) + .withColumn("newBelief", newBeliefCol) # compute new beliefs + .drop("aggMess") # drop messages + .drop("belief") # drop old beliefs + .withColumnRenamed("newBelief", "belief") + ) + # cache new vertices using workaround for SPARK-1334 + cachedNewVertices = newVertices.localCheckpoint() + gx = GraphFrame(cachedNewVertices, gx.edges) + + # Drop the "color" column from vertices + return GraphFrame(gx.vertices.drop("color"), gx.edges) + + @staticmethod + def _colorGraph(g: GraphFrame) -> GraphFrame: + """Given a GraphFrame, choose colors for each vertex. + + No neighboring vertices will share the same color. The number of colors is minimized. + + This is written specifically for grid graphs. For non-grid graphs, it should be generalized, + such as by using a greedy coloring scheme. + + :param g: Grid graph generated by :meth:`Graphs.gridIsingModel()` + :return: Same graph, but with a new vertex column "color" of type Int (0 or 1) + + """ + + colorUDF = sqlfunctions.udf(lambda i, j: (i + j) % 2, returnType=types.IntegerType()) + v = g.vertices.withColumn("color", colorUDF(sqlfunctions.col("i"), sqlfunctions.col("j"))) + return GraphFrame(v, g.edges) + + @staticmethod + def _sigmoid(x: Union[int, float, None]) -> Union[float, None]: + """Numerically stable sigmoid function 1 / (1 + exp(-x))""" + if not x: + return None + if x >= 0: + z = math.exp(-x) + return 1 / (1 + z) + else: + z = math.exp(x) + return z / (1 + z) + + +def main() -> None: + """Run the belief propagation algorithm for an example problem.""" + # setup spark session + spark = SparkSession.builder.appName("BeliefPropagation example").getOrCreate() + + # create graphical model g of size 3 x 3 + g = pyspark.graphframes.examples.Graphs(spark).gridIsingModel(3) + print("Original Ising model:") + g.vertices.show() + g.edges.show() + + # run BP for 5 iterations + numIter = 5 + results = BeliefPropagation.runBPwithGraphFrames(g, numIter) + + # display beliefs + beliefs = results.vertices.select("id", "belief") + print("Done with BP. Final beliefs after {} iterations:".format(numIter)) + beliefs.show() + + spark.stop() + + +if __name__ == "__main__": + main() diff --git a/python/pyspark/graphframes/examples/graphs.py b/python/pyspark/graphframes/examples/graphs.py new file mode 100644 index 0000000000000..0372627a9df99 --- /dev/null +++ b/python/pyspark/graphframes/examples/graphs.py @@ -0,0 +1,137 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import itertools + +from pyspark.graphframes import GraphFrame +from pyspark.sql import SparkSession +from pyspark.sql import functions as sqlfunctions + +__all__ = ["Graphs"] + + +class Graphs: + """Example GraphFrames for testing the API + + :param spark: SparkSession + """ + + def __init__(self, spark: SparkSession) -> None: + self._spark = spark + + def friends(self) -> GraphFrame: + """A GraphFrame of friends in a (fake) social network.""" + # Vertex DataFrame + v = self._spark.createDataFrame( + [ + ("a", "Alice", 34), + ("b", "Bob", 36), + ("c", "Charlie", 30), + ("d", "David", 29), + ("e", "Esther", 32), + ("f", "Fanny", 36), + ], + ["id", "name", "age"], + ) + # Edge DataFrame + e = self._spark.createDataFrame( + [ + ("a", "b", "friend"), + ("b", "c", "follow"), + ("c", "b", "follow"), + ("f", "c", "follow"), + ("e", "f", "follow"), + ("e", "d", "friend"), + ("d", "a", "friend"), + ], + ["src", "dst", "relationship"], + ) + # Create a GraphFrame + return GraphFrame(v, e) + + def gridIsingModel(self, n: int, vStd: float = 1.0, eStd: float = 1.0) -> GraphFrame: + r"""Grid Ising model with random parameters. + + Ising models are probabilistic graphical models over binary variables x\ :sub:`i`. + Each binary variable x\ :sub:`i` corresponds to one vertex, and it may take values -1 or +1. + The probability distribution P(X) (over all x\ :sub:`i`) is parameterized by vertex factors + a\ :sub:`i` and edge factors b\ :sub:`ij`: + + P(X) = (1/Z) * exp[ \sum_i a_i x_i + \sum_{ij} b_{ij} x_i x_j ] + + where Z is the normalization constant (partition function). See `Wikipedia + `__ for more information on Ising models. + + Each vertex is parameterized by a single scalar a\ :sub:`i`. + Each edge is parameterized by a single scalar b\ :sub:`ij`. + + :param n: Length of one side of the grid. The grid will be of size n x n. + :param vStd: Standard deviation of normal distribution used to generate vertex factors "a". + Default of 1.0. + :param eStd: Standard deviation of normal distribution used to generate edge factors "b". + Default of 1.0. + :return: GraphFrame. Vertices have columns "id" and "a". Edges have columns "src", "dst", + and "b". Edges are directed, but they should be treated as undirected in any algorithms + run on this model. Vertex IDs are of the form "i,j". E.g., vertex "1,3" is in the + second row and fourth column of the grid. + """ + # check param n + if n < 1: + raise ValueError( + "Grid graph must have size >= 1, but was given invalid value n = {}".format(n) + ) + + # create coordinates grid + coordinates = self._spark.createDataFrame( + itertools.product(range(n), range(n)), schema=("i", "j") + ) + + # create SQL expression for converting coordinates (i,j) to a string ID "i,j" + # avoid Cartesian join due to SPARK-15425: use generator since n should be small + toIDudf = sqlfunctions.udf(lambda i, j: "{},{}".format(i, j)) + + # create the vertex DataFrame + # create SQL expression for converting coordinates (i,j) to a string ID "i,j" + vIDcol = toIDudf(sqlfunctions.col("i"), sqlfunctions.col("j")) + # add random parameters generated from a normal distribution + seed = 12345 + vertices = coordinates.withColumn("id", vIDcol).withColumn( + "a", sqlfunctions.randn(seed) * vStd + ) + + # create the edge DataFrame + # create SQL expression for converting coordinates (i,j+1) and (i+1,j) to string IDs + rightIDcol = toIDudf(sqlfunctions.col("i"), sqlfunctions.col("j") + 1) + downIDcol = toIDudf(sqlfunctions.col("i") + 1, sqlfunctions.col("j")) + horizontalEdges = coordinates.filter(sqlfunctions.col("j") != n - 1).select( + vIDcol.alias("src"), rightIDcol.alias("dst") + ) + verticalEdges = coordinates.filter(sqlfunctions.col("i") != n - 1).select( + vIDcol.alias("src"), downIDcol.alias("dst") + ) + allEdges = horizontalEdges.unionAll(verticalEdges) + # add random parameters from a normal distribution + edges = allEdges.withColumn("b", sqlfunctions.randn(seed + 1) * eStd) + + # create the GraphFrame + g = GraphFrame(vertices, edges) + + # materialize graph as workaround for SPARK-13333 + g.vertices.cache().count() + g.edges.cache().count() + + return g diff --git a/python/pyspark/graphframes/lib/pregel.py b/python/pyspark/graphframes/lib/pregel.py index 1863bb9bb9849..c88766a89e898 100644 --- a/python/pyspark/graphframes/lib/pregel.py +++ b/python/pyspark/graphframes/lib/pregel.py @@ -15,328 +15,109 @@ # limitations under the License. # -from typing import TYPE_CHECKING, Any, final +from __future__ import annotations + +from typing import Any from typing_extensions import Self -from pyspark.graphframes.classic.utils import storage_level_to_jvm -from pyspark.ml.wrapper import JavaWrapper -from pyspark.sql import Column, DataFrame, SparkSession -from pyspark.sql.classic.column import _to_seq -from pyspark.sql.functions import col +from pyspark.sql import Column, DataFrame +from pyspark.sql import functions as F from pyspark.storagelevel import StorageLevel -if TYPE_CHECKING: - from pyspark.graphframes.classic.graphframe import GraphFrame - - -@final -class Pregel(JavaWrapper): - """Implements a Pregel-like bulk-synchronous message-passing API based on DataFrame operations. - - See `Malewicz et al., Pregel: a system for large-scale graph processing `_ - for a detailed description of the Pregel algorithm. - - You can construct a Pregel instance using either this constructor or :attr:`graphframes.GraphFrame.pregel`, - then use builder pattern to describe the operations, and then call :func:`run` to start a run. - It returns a DataFrame of vertices from the last iteration. - - When a run starts, it expands the vertices DataFrame using column expressions defined by :func:`withVertexColumn`. - Those additional vertex properties can be changed during Pregel iterations. - In each Pregel iteration, there are three phases: - - - Given each edge triplet, generate messages and specify target vertices to send, described by :func:`sendMsgToDst` and :func:`sendMsgToSrc`. - - Aggregate messages by target vertex IDs, described by :func:`aggMsgs`. - - Update additional vertex properties based on aggregated messages and states from previous iteration, described by :func:`withVertexColumn`. - Please find what columns you can reference at each phase in the method API docs. - - You can control the number of iterations by :func:`setMaxIter` and check API docs for advanced controls. - - :param graph: a :class:`graphframes.GraphFrame` object holding a graph with vertices and edges stored as DataFrames. - """ - - def __init__(self, graph: "GraphFrame") -> None: - super(Pregel, self).__init__() +class Pregel: + """Mode-independent wrapper for the GraphFrames Pregel builder API.""" + def __init__(self, graph: Any) -> None: self.graph = graph - self._java_obj: Any = self._new_java_obj( - "org.apache.spark.graphframes.lib.Pregel", graph._jvm_graph - ) - - def setMaxIter(self, value: int) -> "Pregel": - """Sets the max number of iterations (default: 10). + graph_impl = getattr(graph, "_impl", graph) + self._impl = graph_impl.pregel - :param value: the number of Pregel iterations - """ - self._java_obj.setMaxIter(int(value)) + def setMaxIter(self, value: int) -> Self: + self._impl.setMaxIter(value) return self - def setCheckpointInterval(self, value: int) -> "Pregel": - """Sets the number of iterations between two checkpoints (default: 2). - - This is an advanced control to balance query plan optimization and checkpoint data I/O cost. - In most cases, you should keep the default value. - - Checkpoint is disabled if this is set to 0. - """ - self._java_obj.setCheckpointInterval(int(value)) + def setCheckpointInterval(self, value: int) -> Self: + self._impl.setCheckpointInterval(value) return self - def setEarlyStopping(self, value: bool) -> "Pregel": - """Set should Pregel stop earlier in case of no new messages to send or not. - - Early stopping allows to terminate Pregel before reaching maxIter by checking if there are any non-null messages. - While in some cases it may gain significant performance boost, in other cases it can lead to performance degradation, - because checking if the messages DataFrame is empty or not is an action and requires materialization of the Spark Plan - with some additional computations. - - In the case when the user can assume a good value of maxIter, it is recommended to leave this value to the default "false". - In the case when it is hard to estimate the number of iterations required for convergence, - it is recommended to set this value to "false" to avoid iterating over convergence until reaching maxIter. - When this value is "true", maxIter can be set to a bigger value without risks. - """ - self._java_obj.setEarlyStopping(bool(value)) + def setEarlyStopping(self, value: bool) -> Self: + self._impl.setEarlyStopping(value) return self def withVertexColumn( - self, colName: str, initialExpr: Column, updateAfterAggMsgsExpr: Column - ) -> "Pregel": - """Defines an additional vertex column at the start of run and how to update it in each iteration. - - You can call it multiple times to add more than one additional vertex columns. - - :param colName: the name of the additional vertex column. - It cannot be an existing vertex column in the graph. - :param initialExpr: the expression to initialize the additional vertex column. - You can reference all original vertex columns in this expression. - :param updateAfterAggMsgsExpr: the expression to update the additional vertex column after messages aggregation. - You can reference all original vertex columns, additional vertex columns, and the - aggregated message column using :func:`msg`. - If the vertex received no messages, the message column would be null. - """ - self._java_obj.withVertexColumn(colName, initialExpr._jc, updateAfterAggMsgsExpr._jc) + self, + colName: str, + initialExpr: Column | str, + updateAfterAggMsgsExpr: Column | str, + ) -> Self: + self._impl.withVertexColumn(colName, initialExpr, updateAfterAggMsgsExpr) return self - def sendMsgToSrc(self, msgExpr: Column) -> "Pregel": - """Defines a message to send to the source vertex of each edge triplet. - - You can call it multiple times to send more than one messages. - - See method :func:`sendMsgToDst`. - - :param msgExpr: the expression of the message to send to the source vertex given a (src, edge, dst) triplet. - Source/destination vertex properties and edge properties are nested under columns `src`, `dst`, - and `edge`, respectively. - You can reference them using :func:`src`, :func:`dst`, and :func:`edge`. - Null messages are not included in message aggregation. - """ - self._java_obj.sendMsgToSrc(msgExpr._jc) + def sendMsgToSrc(self, msgExpr: Column | str) -> Self: + self._impl.sendMsgToSrc(msgExpr) return self - def sendMsgToDst(self, msgExpr: Column) -> "Pregel": - """Defines a message to send to the destination vertex of each edge triplet. - - You can call it multiple times to send more than one messages. - - See method :func:`sendMsgToSrc`. - - :param msgExpr: the message expression to send to the destination vertex given a (`src`, `edge`, `dst`) triplet. - Source/destination vertex properties and edge properties are nested under columns `src`, `dst`, - and `edge`, respectively. - You can reference them using :func:`src`, :func:`dst`, and :func:`edge`. - Null messages are not included in message aggregation. - """ - self._java_obj.sendMsgToDst(msgExpr._jc) + def sendMsgToDst(self, msgExpr: Column | str) -> Self: + self._impl.sendMsgToDst(msgExpr) return self - def aggMsgs(self, aggExpr: Column) -> "Pregel": - """Defines how messages are aggregated after grouped by target vertex IDs. - - :param aggExpr: the message aggregation expression, such as `sum(Pregel.msg())`. - You can reference the message column by :func:`msg` and the vertex ID by `col("id")`, - while the latter is usually not used. - """ - self._java_obj.aggMsgs(aggExpr._jc) + def aggMsgs(self, aggExpr: Column) -> Self: + self._impl.aggMsgs(aggExpr) return self def setStopIfAllNonActiveVertices(self, value: bool) -> Self: - """Set should Pregel stop if all the vertices voted to halt. - - Activity (or vote) is determined based on the activity_col. - See methods :func:`setInitialActiveVertexExpression` and :func:`setUpdateActiveVertexExpression` for details - how to set and update activity_col. - - Be aware that checking of the vote is not free but a Spark Action. In case the - condition is not realistically reachable but set, it will just slow down the algorithm. - - :param value: the boolean value. - """ - self._java_obj.setStopIfAllNonActiveVertices(value) + self._impl.setStopIfAllNonActiveVertices(value) return self - def setInitialActiveVertexExpression(self, value: Column) -> Self: - """Sets the initial expression for the active vertex column. - - The active vertex column is used to determine if a vertices voting result on each iteration of Pregel. - This expression is evaluated on the initial vertices DataFrame to set the initial state of the activity column. - - :param value: expression to compute the initial active state of vertices. - You can reference all original vertex columns in this expression. - """ - self._java_obj.setInitialActiveVertexExpression(value._jc) + def setInitialActiveVertexExpression(self, value: Column | str) -> Self: + self._impl.setInitialActiveVertexExpression(value) return self - def setUpdateActiveVertexExpression(self, value: Column) -> Self: - """Sets the expression to update the active vertex column. - - The active vertex column is used to determine if a vertices voting result on each iteration of Pregel. - This expression is evaluated on the updated vertices DataFrame to set the new state of the activity column. - - :param value: expression to compute the new active state of vertices. - You can reference all original vertex columns and additional vertex columns in this expression. - """ - self._java_obj.setUpdateActiveVertexExpression(value._jc) + def setUpdateActiveVertexExpression(self, value: Column | str) -> Self: + self._impl.setUpdateActiveVertexExpression(value) return self def setSkipMessagesFromNonActiveVertices(self, value: bool) -> Self: - """Set should Pregel skip sending messages from non-active vertices. - - When this option is enabled, messages will not be sent from vertices that are marked as inactive. - This can help optimize performance by avoiding unnecessary message propagation from inactive vertices. - - :param value: boolean value. - """ - self._java_obj.setSkipMessagesFromNonActiveVertices(value) + self._impl.setSkipMessagesFromNonActiveVertices(value) return self def setUseLocalCheckpoints(self, value: bool) -> Self: - """Set should Pregel use local checkpoints. - - Local checkpoints are faster and do not require configuring a persistent storage. - At the same time, local checkpoints are less reliable and may create a big load on local disks of executors. - - :param value: boolean value. - """ - self._java_obj.setUseLocalCheckpoints(value) + self._impl.setUseLocalCheckpoints(value) return self def setIntermediateStorageLevel(self, storage_level: StorageLevel) -> Self: - """Set the intermediate storage level. - On each iteration, Pregel cache results with a requested storage level. - - For very big graphs it is recommended to use DISK_ONLY. - - :param storage_level: storage level to use. - """ - self._java_obj.setIntermediateStorageLevel( - storage_level_to_jvm(storage_level, self.graph._spark) - ) + self._impl.setIntermediateStorageLevel(storage_level) return self def required_src_columns(self, col_name: str, *col_names: str) -> Self: - """Specifies which source vertex columns are required when constructing triplets. - - By default, all source vertex columns are included in triplets, which can create large - intermediate datasets for algorithms with significant state (e.g., cycle detection, - random walks). Use this method to reduce memory usage by specifying only the columns - that are actually needed by the sendMsgToSrc and sendMsgToDst expressions. - - The ID column and the active flag column (if used) are always included automatically. - - :param col_name: the first required source vertex column name - :param col_names: additional required source vertex column names - - See also :func:`required_dst_columns` - """ - self._java_obj.requiredSrcColumns( - col_name, _to_seq(self.graph._spark.sparkContext, col_names) - ) + self._impl.required_src_columns(col_name, *col_names) return self def required_dst_columns(self, col_name: str, *col_names: str) -> Self: - """Specifies which destination vertex columns are required when constructing triplets. - - By default, all destination vertex columns are included in triplets, which can create large - intermediate datasets for algorithms with significant state (e.g., cycle detection, - random walks). Use this method to reduce memory usage by specifying only the columns - that are actually needed by the sendMsgToSrc and sendMsgToDst expressions. - - The ID column and the active flag column (if used) are always included automatically. - - :param col_name: the first required destination vertex column name - :param col_names: additional required destination vertex column names - - See also :func:`required_src_columns` - """ - self._java_obj.requiredDstColumns( - col_name, _to_seq(self.graph._spark.sparkContext, col_names) - ) + self._impl.required_dst_columns(col_name, *col_names) return self def required_edge_columns(self, col_name: str, *col_names: str) -> Self: - """Specifies which edge columns are required when constructing triplets. - - By default, only src and dst columns are included from edges. Use this method to - specify additional edge columns that are needed by the sendMsgToSrc and sendMsgToDst - expressions. - - :param col_name: the first required edge column name - :param col_names: additional required edge column names - - See also :func:`required_src_columns` and :func:`required_dst_columns` - """ - self._java_obj.requiredEdgeColumns( - col_name, _to_seq(self.graph._spark.sparkContext, col_names) - ) + self._impl.required_edge_columns(col_name, *col_names) return self def run(self) -> DataFrame: - """Runs the defined Pregel algorithm. - - :return: the result vertex DataFrame from the final iteration including both original and additional columns. - """ - spark = SparkSession.getActiveSession() - if spark is None: - raise ValueError("SparkSession is dead or did not started.") - return DataFrame(self._java_obj.run(), spark) + return self._impl.run() @staticmethod def msg() -> Column: - """References the message column in aggregating messages and updating additional vertex columns. - - See :func:`aggMsgs` and :func:`withVertexColumn` - """ - return col("_pregel_msg_") + return F.col("_pregel_msg_") @staticmethod def src(colName: str) -> Column: - """References a source vertex column in generating messages to send. - - See :func:`sendMsgToSrc` and :func:`sendMsgToDst` - - :param colName: the vertex column name. - """ - return col("src." + colName) + return F.col("src." + colName) @staticmethod def dst(colName: str) -> Column: - """ - References a destination vertex column in generating messages to send. - - See :func:`sendMsgToSrc` and :func:`sendMsgToDst` - - :param colName: the vertex column name. - """ - return col("dst." + colName) + return F.col("dst." + colName) @staticmethod def edge(colName: str) -> Column: - """ - References an edge column in generating messages to send. - - See :func:`sendMsgToSrc` and :func:`sendMsgToDst` - - :param colName: the edge column name. - """ - return col("edge." + colName) + return F.col("edge." + colName) diff --git a/python/pyspark/graphframes/pg/__init__.py b/python/pyspark/graphframes/pg/__init__.py new file mode 100644 index 0000000000000..f1badcee91d04 --- /dev/null +++ b/python/pyspark/graphframes/pg/__init__.py @@ -0,0 +1,25 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from pyspark.graphframes.pg.property_graphframe import PropertyGraphFrame +from pyspark.graphframes.pg.property_groups import EdgePropertyGroup, VertexPropertyGroup + +__all__ = [ + "VertexPropertyGroup", + "EdgePropertyGroup", + "PropertyGraphFrame", +] diff --git a/python/pyspark/graphframes/pg/property_graphframe.py b/python/pyspark/graphframes/pg/property_graphframe.py new file mode 100644 index 0000000000000..899c324ee32d2 --- /dev/null +++ b/python/pyspark/graphframes/pg/property_graphframe.py @@ -0,0 +1,381 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""PropertyGraphFrame implementation for PySpark.""" + +from __future__ import annotations + +from collections.abc import Callable, Sequence +from typing import TYPE_CHECKING + +from pyspark.graphframes.pg.property_groups import EdgePropertyGroup, VertexPropertyGroup +from pyspark.sql.functions import col, lit + +if TYPE_CHECKING: + from pyspark.sql import Column, DataFrame + +from pyspark.graphframes import GraphFrame + + +class PropertyGraphFrame: + """ + A high-level abstraction for working with property graphs in PySpark. + + PropertyGraphFrame serves as a logical structure that manages collections of vertex and edge + property groups, providing a user-friendly API for graph operations. It handles various + internal complexities such as: + + - ID conversion and collision prevention + - Management of directed/undirected graph representations + - Handling of weighted/unweighted edges + - Data consistency across different property groups + + The class maintains separate collections for vertex and edge properties, allowing for flexible + graph construction while ensuring data integrity. + + Example: + >>> from pyspark.graphframes.pg import VertexPropertyGroup, EdgePropertyGroup, PropertyGraphFrame + >>> from pyspark.graphframes import GraphFrame + >>> + >>> # Create vertex groups + >>> people_data = spark.createDataFrame([(1, "Alice"), (2, "Bob")], ["id", "name"]) + >>> people_group = VertexPropertyGroup("people", people_data, "id") + >>> + >>> movies_data = spark.createDataFrame([(1, "Matrix"), (2, "Inception")], ["id", "title"]) + >>> movies_group = VertexPropertyGroup("movies", movies_data, "id") + >>> + >>> # Create edge group + >>> likes_data = spark.createDataFrame([(1, 1, 1.0)], ["src", "dst", "weight"]) + >>> likes_group = EdgePropertyGroup( + ... "likes", likes_data, people_group, movies_group, + ... is_directed=False, src_column_name="src", dst_column_name="dst", + ... weight_column_name="weight" + ... ) + >>> + >>> # Create property graph + >>> pg = PropertyGraphFrame([people_group, movies_group], [likes_group]) + + :param vertex_property_groups: Sequence of vertex property groups + :param edges_property_groups: Sequence of edge property groups + """ + + PROPERTY_GROUP_COL_NAME = "property_group" + EXTERNAL_ID = "external_id" + + def __init__( + self, + vertex_property_groups: Sequence, + edges_property_groups: Sequence, + ) -> None: + """ + Initialize a PropertyGraphFrame. + + :param vertex_property_groups: Sequence of vertex property groups + :param edges_property_groups: Sequence of edge property groups + """ + + # Validate input types + for group in vertex_property_groups: + if not isinstance(group, VertexPropertyGroup): + raise TypeError( + f"All vertex_property_groups must be VertexPropertyGroup instances, " + f"got {type(group)}" + ) + + for group in edges_property_groups: + if not isinstance(group, EdgePropertyGroup): + raise TypeError( + f"All edges_property_groups must be EdgePropertyGroup instances, " + f"got {type(group)}" + ) + + self._vertex_property_groups = list(vertex_property_groups) + self._edges_property_groups = list(edges_property_groups) + + # Create lookup maps + self._vertex_groups: dict[str, VertexPropertyGroup] = { + group.name: group for group in self._vertex_property_groups + } + self._edge_groups: dict[str, EdgePropertyGroup] = { + group.name: group for group in self._edges_property_groups + } + + @property + def vertex_property_groups(self) -> list[VertexPropertyGroup]: + """Return the list of vertex property groups.""" + + return self._vertex_property_groups + + @property + def edges_property_groups(self) -> list[EdgePropertyGroup]: + """Return the list of edge property groups.""" + + return self._edges_property_groups + + def to_graphframe( + self, + vertex_property_groups: Sequence[str], + edge_property_groups: Sequence[str], + edge_group_filters: dict[str, Column] | None = None, + vertex_group_filters: dict[str, Column] | None = None, + ) -> GraphFrame: + """ + Convert the property graph to a unified GraphFrame representation. + + This method transforms a property graph that may contain multiple vertex types and both + directed and undirected edges into a single GraphFrame object where all vertices and edges + share the same schema. The conversion process handles: + + - Internal ID generation and collision prevention by hashing vertex/edge IDs with their + group names + - Merging of different vertex types into a unified vertex DataFrame + - Conversion of directed/undirected edge relationships into a consistent edge DataFrame + - Filtering of vertices and edges based on provided predicates + + :param vertex_property_groups: Sequence of vertex property group names to include + :param edge_property_groups: Sequence of edge property group names to include + :param edge_group_filters: Optional dict mapping edge group names to filter predicates + :param vertex_group_filters: Optional dict mapping vertex group names to filter predicates + :return: A GraphFrame containing the unified representation + :raises ValueError: If a specified group name does not exist + + Example: + >>> from pyspark.sql.functions import lit + >>> graph = pg.to_graph_frame( + ... vertex_property_groups=["people", "movies"], + ... edge_property_groups=["likes", "messages"], + ... edge_group_filters={"likes": lit(True), "messages": lit(True)}, + ... vertex_group_filters={"people": lit(True), "movies": lit(True)} + ... ) + """ + # Set default filters if not provided + if edge_group_filters is None: + edge_group_filters = {} + if vertex_group_filters is None: + vertex_group_filters = {} + + # Validate group names + for name in vertex_property_groups: + if name not in self._vertex_groups: + raise ValueError(f"Vertex property group '{name}' does not exist") + + for name in edge_property_groups: + if name not in self._edge_groups: + raise ValueError(f"Edge property group '{name}' does not exist") + + # Combine vertices from all specified groups + if not vertex_property_groups: + raise ValueError("At least one vertex property group must be specified") + + vertices_list = [] + for name in vertex_property_groups: + filter_col = vertex_group_filters.get(name, lit(True)) + group_data = self._vertex_groups[name].get_data(filter_col) + vertices_list.append(group_data) + + vertices = vertices_list[0] + for v in vertices_list[1:]: + vertices = vertices.union(v) + + # Combine edges from all specified groups + if not edge_property_groups: + raise ValueError("At least one edge property group must be specified") + + edges_list = [] + for name in edge_property_groups: + filter_col = edge_group_filters.get(name, lit(True)) + group_data = self._edge_groups[name].get_data(filter_col) + edges_list.append(group_data) + + edges = edges_list[0] + for e in edges_list[1:]: + edges = edges.union(e) + + return GraphFrame(vertices, edges) + + def projection_by( + self, + left_bi_graph_part: str, + right_bi_graph_part: str, + edge_group: str, + new_edge_weight: Callable[[Column, Column], Column] | None = None, + ) -> "PropertyGraphFrame": + """ + Project a bipartite graph onto one of its parts. + + Creates edges between vertices that share neighbors in the other part. Drops the property + group used for projection and returns a new property graph. + + :param left_bi_graph_part: Name of the vertex property group to project onto + :param right_bi_graph_part: Name of the vertex property group to project through + :param edge_group: Name of the edge property group connecting the two parts + :param new_edge_weight: Optional function that takes two weight columns and returns + a new weight column. If None, uses weight 1.0 for all edges. + :return: A new PropertyGraphFrame containing the projected graph + :raises ValueError: If group names are invalid or edge group doesn't connect the parts + + Example: + >>> # Project people through movies they both like + >>> projected = pg.projection_by("people", "movies", "likes") + >>> # Custom weight function + >>> from pyspark.sql.functions import col + >>> projected = pg.projection_by( + ... "people", "movies", "likes", + ... new_edge_weight=lambda w1, w2: w1 + w2 + ... ) + """ + # Validate inputs + if edge_group not in self._edge_groups: + raise ValueError(f"Edge property group '{edge_group}' does not exist") + + if left_bi_graph_part not in self._vertex_groups: + raise ValueError(f"Vertex property group '{left_bi_graph_part}' does not exist") + + if right_bi_graph_part not in self._vertex_groups: + raise ValueError(f"Vertex property group '{right_bi_graph_part}' does not exist") + + old_group = self._edge_groups[edge_group] + + # Validate edge group connects the specified parts + if old_group.src_property_group.name != left_bi_graph_part: + raise ValueError( + f"Edge property group should have '{left_bi_graph_part}' as source " + f"but has '{old_group.src_property_group.name}'" + ) + + if old_group.dst_property_group.name != right_bi_graph_part: + raise ValueError( + f"Edge property group should have '{right_bi_graph_part}' as destination " + f"but has '{old_group.dst_property_group.name}'" + ) + + # Get vertex groups to keep + kept_v_property_groups = [ + g for g in self._vertex_property_groups if g.name != right_bi_graph_part + ] + + # Get edge groups to keep (excluding the one being projected) + kept_e_property_groups = [g for g in self._edges_property_groups if g.name != edge_group] + + # Create projected edges by joining edges through common neighbors + old_edges_data = old_group.data + + e1 = old_edges_data.alias("e1") + e2 = old_edges_data.alias("e2") + + # Join edges on common destination (the right part) + joined = e1.join( + e2, col("e1." + old_group.dst_column_name) == col("e2." + old_group.dst_column_name) + ) + + # Filter to avoid duplicates (e1.src < e2.src) + joined = joined.filter( + col("e1." + old_group.src_column_name) < col("e2." + old_group.src_column_name) + ) + + # Add weight column + if new_edge_weight is not None: + w1 = col(f"e1.{old_group.weight_column_name}") + w2 = col(f"e2.{old_group.weight_column_name}") + weight_col = new_edge_weight(w1, w2) + else: + weight_col = lit(1.0) + + # Select source and destination for new edges + projected_edges = joined.select( + col("e1." + old_group.src_column_name).alias(GraphFrame.SRC), + col("e2." + old_group.src_column_name).alias(GraphFrame.DST), + weight_col.alias(GraphFrame.WEIGHT), + ) + + # Create new edge property group + left_group = self._vertex_groups[left_bi_graph_part] + + new_edge_group = EdgePropertyGroup( + name=f"projected_{edge_group}", + data=projected_edges, + src_property_group=left_group, + dst_property_group=left_group, + is_directed=False, + src_column_name=GraphFrame.SRC, + dst_column_name=GraphFrame.DST, + weight_column_name=GraphFrame.WEIGHT, + ) + + return PropertyGraphFrame(kept_v_property_groups, kept_e_property_groups + [new_edge_group]) + + def join_vertices( + self, + vertices_data: DataFrame, + vertex_groups: Sequence[str], + ) -> DataFrame: + """ + Join algorithm results back to the original vertex data. + + Joins the vertices data (typically output from graph algorithms) with the specified + vertex property groups to produce a unified DataFrame with original vertex attributes. + + :param vertices_data: DataFrame containing vertex algorithm results (from to_graph_frame) + :param vertex_groups: Sequence of vertex group names to join + :return: A DataFrame with joined vertex data + :raises ValueError: If a specified group name does not exist + + Example: + >>> # Run connected components and join results back + >>> graph = pg.to_graph_frame(["people"], ["messages"], {}, {}) + >>> components = graph.connectedComponents() + >>> joined = pg.join_vertices(components, ["people"]) + """ + # Validate group names + for name in vertex_groups: + if name not in self._vertex_groups: + raise ValueError(f"Vertex property group '{name}' does not exist") + + if not vertex_groups: + raise ValueError("At least one vertex group must be specified") + + # Join each group separately + result_dfs = [] + + for vg_name in vertex_groups: + group = self._vertex_groups[vg_name] + + # Filter vertices data for this group + filtered = vertices_data.filter( + col(PropertyGraphFrame.PROPERTY_GROUP_COL_NAME) == lit(vg_name) + ) + + if group.apply_mask_on_id: + # Use internal ID mapping to join back to original data + id_mapping = group._get_internal_id_mapping() + joined = id_mapping.join(filtered, [GraphFrame.ID], "left").drop(GraphFrame.ID) + else: + # Direct join on ID + joined = ( + group.get_data() + .join(filtered, GraphFrame.ID, "left") + .withColumnRenamed(GraphFrame.ID, PropertyGraphFrame.EXTERNAL_ID) + ) + + result_dfs.append(joined) + + # Union all results + result = result_dfs[0] + for df in result_dfs[1:]: + result = result.union(df) + + return result diff --git a/python/pyspark/graphframes/pg/property_groups.py b/python/pyspark/graphframes/pg/property_groups.py new file mode 100644 index 0000000000000..45931474c624e --- /dev/null +++ b/python/pyspark/graphframes/pg/property_groups.py @@ -0,0 +1,386 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Property group classes for property graphs.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +from pyspark.graphframes import GraphFrame +from pyspark.sql.functions import col, concat, lit, sha2 +from pyspark.sql.types import ( + ByteType, + DecimalType, + DoubleType, + FloatType, + IntegerType, + LongType, + ShortType, + StringType, +) + +if TYPE_CHECKING: + from pyspark.sql import Column, DataFrame + + +class InvalidPropertyGroupException(Exception): + """Exception raised when a property group is invalid.""" + + pass + + +class PropertyGroup(ABC): + """Abstract base class for property groups.""" + + def __init__(self, name: str, data: DataFrame) -> None: + """ + Initialize a property group. + + :param name: The unique identifier for this property group + :param data: The DataFrame containing the property data + """ + self._name = name + self._data = data + self._validate() + + @property + def name(self) -> str: + """Return the name of the property group.""" + return self._name + + @property + def data(self) -> DataFrame: + """Return the DataFrame containing the property data.""" + return self._data + + @abstractmethod + def _validate(self) -> None: + """Validate the property group. Must be implemented by subclasses.""" + pass + + def get_data(self, filter_col: Column | None = None) -> DataFrame: + """ + Return a view of the data for the property group. + + :param filter_col: An optional filter condition (Column) to apply to the data + :return: A DataFrame containing the filtered and optionally transformed data + """ + + if filter_col is None: + filter_col = lit(True) + return self._get_data(filter_col) + + @abstractmethod + def _get_data(self, filter_col: Column) -> DataFrame: + """Internal method to get filtered data. Must be implemented by subclasses.""" + pass + + +class VertexPropertyGroup(PropertyGroup): + """ + Represents a logical group of vertices in a property graph. + + A VertexPropertyGroup organizes and manages vertices that share common characteristics + or belong to the same logical group within a property graph. Each group maintains its + own data in the form of a DataFrame and uses a primary key column for unique vertex + identification. + + When vertices from different groups are combined into a GraphFrame, their IDs are + hashed with the group name to prevent collisions. + + Example: + >>> people_data = spark.createDataFrame([(1, "Alice"), (2, "Bob")], ["id", "name"]) + >>> people_group = VertexPropertyGroup("people", people_data, "id") + + :param name: The unique identifier for this vertex property group + :param data: The DataFrame containing the vertex data + :param primary_key_column: The column name used to uniquely identify vertices + :param apply_mask_on_id: Whether to hash IDs with group name (default: True) + """ + + def __init__( + self, + name: str, + data: DataFrame, + primary_key_column: str = "id", + apply_mask_on_id: bool = True, + ) -> None: + """ + Initialize a VertexPropertyGroup. + + :param name: Name of the vertex property group + :param data: DataFrame containing vertex data + :param primary_key_column: Name of the column to use as primary key (default: "id") + :param apply_mask_on_id: Whether to apply masking on vertex IDs (default: True) + """ + self._primary_key_column = primary_key_column + self._apply_mask_on_id = apply_mask_on_id + super().__init__(name, data) + + @property + def primary_key_column(self) -> str: + """Return the primary key column name.""" + return self._primary_key_column + + @property + def apply_mask_on_id(self) -> bool: + """Return whether ID masking is applied.""" + return self._apply_mask_on_id + + def _validate(self) -> None: + """Validate that the primary key column exists in the data.""" + if self._primary_key_column not in self._data.columns: + raise InvalidPropertyGroupException( + f"source column {self._primary_key_column} does not exist, " + f"existed columns [{', '.join(self._data.columns)}]" + ) + + def _get_internal_id_mapping(self) -> DataFrame: + """ + Create a mapping from external IDs to internal hashed IDs. + + :return: DataFrame with columns 'external_id' and 'id' + """ + + EXTERNAL_ID = "external_id" + + return self._data.select(col(self._primary_key_column).alias(EXTERNAL_ID)).withColumn( + GraphFrame.ID, + concat( + lit(self._name), + sha2(col(EXTERNAL_ID).cast(StringType()), 256), + ), + ) + + def _get_data(self, filter_col: Column) -> DataFrame: + """ + Return filtered vertex data with internal IDs and property group column. + + :param filter_col: Filter condition to apply + :return: DataFrame with columns 'id' and 'property_group' + """ + PROPERTY_GROUP_COL_NAME = "property_group" + + filtered_data = self._data.filter(filter_col) + + if self._apply_mask_on_id: + result = filtered_data.select( + concat( + lit(self._name), + sha2(col(self._primary_key_column).cast(StringType()), 256), + ).alias(GraphFrame.ID) + ) + else: + result = filtered_data.select( + col(self._primary_key_column).cast(StringType()).alias(GraphFrame.ID) + ) + + return result.select( + col(GraphFrame.ID), + lit(self._name).alias(PROPERTY_GROUP_COL_NAME), + ) + + +class EdgePropertyGroup(PropertyGroup): + """ + Represents a logical group of edges in a property graph. + + EdgePropertyGroup encapsulates edge data stored in a DataFrame along with metadata + describing how to interpret the data as graph edges. Each edge group has: + + - A unique name identifier + - DataFrame containing the actual edge data + - Source and destination vertex property groups + - Direction flag indicating if edges are directed or undirected + - Column names specifying source vertex, destination vertex, and edge weight + + When edges from different groups are combined into a GraphFrame, their src and dst + are hashed with the group name to prevent ID collisions. + + Example: + >>> edges_data = spark.createDataFrame([(1, 2, 1.0)], ["src", "dst", "weight"]) + >>> edges_group = EdgePropertyGroup( + ... "likes", edges_data, people_group, movies_group, + ... is_directed=False, src_column="src", dst_column="dst", weight_column="weight" + ... ) + + :param name: Unique identifier for this edge property group + :param data: DataFrame containing the edge data + :param src_property_group: Source vertex property group + :param dst_property_group: Destination vertex property group + :param is_directed: Whether edges should be treated as directed + :param src_column_name: Name of the source vertex column in the data + :param dst_column_name: Name of the destination vertex column in the data + :param weight_column_name: Name of the edge weight column in the data + """ + + def __init__( + self, + name: str, + data: DataFrame, + src_property_group: VertexPropertyGroup, + dst_property_group: VertexPropertyGroup, + is_directed: bool, + src_column_name: str, + dst_column_name: str, + weight_column_name: str | None = None, + ) -> None: + """ + Initialize an EdgePropertyGroup. + + :param name: Unique identifier for this edge property group + :param data: DataFrame containing the edge data with required columns + :param src_property_group: Source vertex property group + :param dst_property_group: Destination vertex property group + :param is_directed: Whether edges are directed (True) or undirected (False) + :param src_column_name: Name of the source vertex column + :param dst_column_name: Name of the destination vertex column + :param weight_column_name: Name of the edge weight column + (None means the lit(1).alias("weight") will be used) + """ + if weight_column_name is None: + data = data.withColumn("weight", lit(1.0)) + weight_column_name = "weight" + + self._src_property_group = src_property_group + self._dst_property_group = dst_property_group + self._is_directed = is_directed + self._src_column_name = src_column_name + self._dst_column_name = dst_column_name + self._weight_column_name = weight_column_name + super().__init__(name, data) + + @property + def src_property_group(self) -> VertexPropertyGroup: + """Return the source vertex property group.""" + return self._src_property_group + + @property + def dst_property_group(self) -> VertexPropertyGroup: + """Return the destination vertex property group.""" + return self._dst_property_group + + @property + def is_directed(self) -> bool: + """Return whether edges are directed.""" + return self._is_directed + + @property + def src_column_name(self) -> str: + """Return the source column name.""" + return self._src_column_name + + @property + def dst_column_name(self) -> str: + """Return the destination column name.""" + return self._dst_column_name + + @property + def weight_column_name(self) -> str: + """Return the weight column name.""" + return self._weight_column_name + + def _validate(self) -> None: + """Validate that required columns exist and weight column is numeric.""" + if self._src_column_name not in self._data.columns: + raise InvalidPropertyGroupException( + f"source column {self._src_column_name} does not exist, " + f"existed columns [{', '.join(self._data.columns)}]" + ) + if self._dst_column_name not in self._data.columns: + raise InvalidPropertyGroupException( + f"dest column {self._dst_column_name} does not exist, " + f"existed columns [{', '.join(self._data.columns)}]" + ) + if self._weight_column_name not in self._data.columns: + raise InvalidPropertyGroupException( + f"weight column {self._weight_column_name} does not exist, " + f"existed columns [{', '.join(self._data.columns)}]" + ) + + # Check weight column type + weight_column_type = self._data.schema[self._weight_column_name].dataType + if not self._is_numeric_type(weight_column_type): + _msg = "weight column {} must be numeric type, but was {}" + raise InvalidPropertyGroupException( + _msg.format(self._weight_column_name, weight_column_type) + ) + + def _is_numeric_type(self, data_type) -> bool: + """Check if a Spark data type is numeric.""" + + numeric_types = ( + ByteType, + ShortType, + IntegerType, + LongType, + FloatType, + DoubleType, + DecimalType, + ) + return isinstance(data_type, numeric_types) + + def _hash_src_edge(self) -> Column: + """Hash the source edge ID based on the source property group settings.""" + + if self._src_property_group.apply_mask_on_id: + return concat( + lit(self._src_property_group.name), + sha2(col(self._src_column_name).cast(StringType()), 256), + ) + else: + return col(self._src_column_name).cast(StringType()) + + def _hash_dst_edge(self) -> Column: + """Hash the destination edge ID based on the destination property group settings.""" + if self._dst_property_group.apply_mask_on_id: + return concat( + lit(self._dst_property_group.name), + sha2(col(self._dst_column_name).cast(StringType()), 256), + ) + else: + return col(self._dst_column_name).cast(StringType()) + + def _get_data(self, filter_col: Column) -> DataFrame: + """ + Return filtered edge data with hashed IDs and weights. + + For undirected edges, creates bidirectional edges. + + :param filter_col: Filter condition to apply + :return: DataFrame with columns 'src', 'dst', and 'weight' + """ + filtered_data = self._data.filter(filter_col) + + base_edges = filtered_data.select( + self._hash_src_edge().alias(GraphFrame.SRC), + self._hash_dst_edge().alias(GraphFrame.DST), + col(self._weight_column_name).alias(GraphFrame.WEIGHT), + ) + + if self._is_directed: + return base_edges + else: + # For undirected edges, create bidirectional edges + reverse_edges = base_edges.select( + col(GraphFrame.DST).alias(GraphFrame.SRC), + col(GraphFrame.SRC).alias(GraphFrame.DST), + col(GraphFrame.WEIGHT).alias(GraphFrame.WEIGHT), + ) + return base_edges.union(reverse_edges) diff --git a/python/pyspark/graphframes/tests/_upstream_test_utils.py b/python/pyspark/graphframes/tests/_upstream_test_utils.py new file mode 100644 index 0000000000000..cc489ac282a35 --- /dev/null +++ b/python/pyspark/graphframes/tests/_upstream_test_utils.py @@ -0,0 +1,109 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from __future__ import annotations + +import math +import re +import unittest +from collections.abc import Callable +from typing import Any + + +class _Approx: + def __init__(self, expected: float, abs: float = 1e-12, rel: float = 1e-6) -> None: + self.expected = expected + self.abs = abs + self.rel = rel + + def __eq__(self, actual: object) -> bool: + return isinstance(actual, (int, float)) and math.isclose( + actual, self.expected, abs_tol=self.abs, rel_tol=self.rel + ) + + +class _Raises: + def __init__(self, exception: type[BaseException], match: str | None = None) -> None: + self.exception = exception + self.match = match + + def __enter__(self) -> "_Raises": + return self + + def __exit__( + self, + exception_type: type[BaseException] | None, + exception: BaseException | None, + traceback: object, + ) -> bool: + if exception_type is None: + raise AssertionError(f"{self.exception.__name__} was not raised") + if not issubclass(exception_type, self.exception): + return False + if self.match is not None and ( + exception is None or re.search(self.match, str(exception)) is None + ): + raise AssertionError( + f"{self.match!r} does not match exception message {str(exception)!r}" + ) + return True + + +class _Mark: + @staticmethod + def parametrize( + *args: Any, **kwargs: Any + ) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + return lambda function: function + + @staticmethod + def skipif( + condition: bool, *, reason: str + ) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + def decorate(function: Callable[..., Any]) -> Callable[..., Any]: + if not condition: + return function + + def skipped(*args: Any, **kwargs: Any) -> None: + raise unittest.SkipTest(reason) + + return skipped + + return decorate + + +class _PytestCompatibility: + mark = _Mark() + + @staticmethod + def approx(expected: float, *, abs: float = 1e-12, rel: float = 1e-6) -> _Approx: + return _Approx(expected, abs=abs, rel=rel) + + @staticmethod + def fixture(*args: Any, **kwargs: Any) -> Callable[[Callable[..., Any]], Callable[..., Any]]: + return lambda function: function + + @staticmethod + def raises(exception: type[BaseException], *, match: str | None = None) -> _Raises: + return _Raises(exception, match) + + @staticmethod + def skip(reason: str) -> None: + raise unittest.SkipTest(reason) + + +pytest = _PytestCompatibility() diff --git a/python/pyspark/graphframes/tests/connect/test_all_algorithms.py b/python/pyspark/graphframes/tests/connect/test_all_algorithms.py index dae483040f891..afa8f0f9463cf 100644 --- a/python/pyspark/graphframes/tests/connect/test_all_algorithms.py +++ b/python/pyspark/graphframes/tests/connect/test_all_algorithms.py @@ -17,7 +17,7 @@ from pyspark.graphframes import GraphFrame from pyspark.graphframes.graphframe import AggregateNeighbors, RandomWalkEmbeddings -from pyspark.graphframes.lib import AggregateMessages +from pyspark.graphframes.lib import AggregateMessages, Pregel from pyspark.sql import functions as F from pyspark.storagelevel import StorageLevel from pyspark.testing.connectutils import ReusedConnectTestCase @@ -82,6 +82,11 @@ def test_message_aggregation_and_pregel(self) -> None: self.assertIn("value", result.columns) result.unpersist() + def test_direct_pregel_construction(self) -> None: + pregel = Pregel(self.graph) + self.assertIs(pregel.setMaxIter(1), pregel) + self.assertIsInstance(pregel, Pregel) + def test_component_and_community_algorithms(self) -> None: components = self.graph.connectedComponents( algorithm="two_phase", diff --git a/python/pyspark/graphframes/tests/connect/test_client_imports.py b/python/pyspark/graphframes/tests/connect/test_client_imports.py new file mode 100644 index 0000000000000..421549b6d952b --- /dev/null +++ b/python/pyspark/graphframes/tests/connect/test_client_imports.py @@ -0,0 +1,43 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import sys +import unittest + +from pyspark.graphframes import GraphFrame +from pyspark.graphframes.examples import BeliefPropagation, Graphs +from pyspark.graphframes.lib import AggregateMessages, Pregel +from pyspark.graphframes.lib.pregel import Pregel as PregelFromSubmodule +from pyspark.util import is_remote_only + + +class GraphFramesClientImportTests(unittest.TestCase): + def test_public_imports(self) -> None: + self.assertTrue(GraphFrame) + self.assertTrue(AggregateMessages) + self.assertTrue(BeliefPropagation) + self.assertTrue(Graphs) + self.assertTrue(Pregel) + self.assertIs(Pregel, PregelFromSubmodule) + if is_remote_only(): + self.assertNotIn("pyspark.graphframes.classic.pregel", sys.modules) + + +if __name__ == "__main__": + from pyspark.testing.unittestutils import main + + main() diff --git a/python/pyspark/graphframes/tests/connect/test_property_graphframe.py b/python/pyspark/graphframes/tests/connect/test_property_graphframe.py new file mode 100644 index 0000000000000..d8f6cc67f168f --- /dev/null +++ b/python/pyspark/graphframes/tests/connect/test_property_graphframe.py @@ -0,0 +1,119 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os +import shutil +import tempfile + +from pyspark import SparkConf +from pyspark.graphframes.tests.pg import test_property_graphframe as upstream +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class PropertyGraphFrameConnectTests(ReusedConnectTestCase): + _checkpoint_dir = tempfile.mkdtemp(prefix="spark-property-graphframe-connect-") + + @classmethod + def conf(cls) -> SparkConf: + return ( + super() + .conf() + .set("spark.sql.shuffle.partitions", "4") + .set("spark.checkpoint.dir", cls._checkpoint_dir) + ) + + @classmethod + def setUpClass(cls) -> None: + super().setUpClass() + if cls._legacy_sc is not None: + cls._legacy_sc.setCheckpointDir(cls._checkpoint_dir) + + @classmethod + def tearDownClass(cls) -> None: + super().tearDownClass() + if os.path.exists(cls._checkpoint_dir): + shutil.rmtree(cls._checkpoint_dir) + + def groups(self): + people = upstream.people_group(self.spark) + movies = upstream.movies_group(self.spark) + likes = upstream.likes_group(self.spark, people, movies) + messages = upstream.messages_group(self.spark, people) + graph = upstream.people_movies_graph(people, movies, likes, messages) + return people, movies, likes, messages, graph + + def test_property_graph_frame_constructor(self) -> None: + *_, graph = self.groups() + upstream.test_property_graph_frame_constructor(graph) + + def test_vertex_property_group_creation(self) -> None: + people, *_ = self.groups() + upstream.test_vertex_property_group_creation(people) + + def test_edge_property_group_creation(self) -> None: + _, _, likes, _, _ = self.groups() + upstream.test_edge_property_group_creation(likes) + + def test_projection_by_movies(self) -> None: + *_, graph = self.groups() + upstream.test_projection_by_movies(graph) + + def test_projection_with_custom_weight(self) -> None: + *_, graph = self.groups() + upstream.test_projection_with_custom_weight(graph) + + def test_to_graph_frame_messages_only(self) -> None: + *_, graph = self.groups() + upstream.test_to_graph_frame_messages_only(graph) + + def test_to_graph_frame_all_groups(self) -> None: + *_, graph = self.groups() + upstream.test_to_graph_frame_all_groups(graph) + + def test_to_graph_frame_unmasked_ids(self) -> None: + people, _, likes, messages, _ = self.groups() + upstream.test_to_graph_frame_unmasked_ids(self.spark, people, likes, messages) + + def test_join_vertices_with_connected_components(self) -> None: + *_, graph = self.groups() + upstream.test_join_vertices_with_connected_components(graph) + + def test_vertex_property_group_validation(self) -> None: + people, *_ = self.groups() + upstream.test_vertex_property_group_validation(people) + + def test_edge_property_group_validation(self) -> None: + people, movies, likes, _, _ = self.groups() + upstream.test_edge_property_group_validation(people, movies, likes) + + def test_to_graph_frame_invalid_group(self) -> None: + *_, graph = self.groups() + upstream.test_to_graph_frame_invalid_group(graph) + + def test_projection_by_invalid_group(self) -> None: + *_, graph = self.groups() + upstream.test_projection_by_invalid_group(graph) + + def test_property_graph_frame_to_graph_frame_conversion(self) -> None: + *_, graph = self.groups() + upstream.test_property_graph_frame_to_graph_frame_conversion(graph) + + +if __name__ == "__main__": + from pyspark.testing.unittestutils import main + + main() diff --git a/python/pyspark/graphframes/tests/connect/test_upstream_graphframes.py b/python/pyspark/graphframes/tests/connect/test_upstream_graphframes.py new file mode 100644 index 0000000000000..f1f89938ef607 --- /dev/null +++ b/python/pyspark/graphframes/tests/connect/test_upstream_graphframes.py @@ -0,0 +1,176 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os +import shutil +import tempfile + +from pyspark import SparkConf +from pyspark.graphframes import GraphFrame +from pyspark.graphframes.tests import test_graphframes as upstream +from pyspark.testing.connectutils import ReusedConnectTestCase + + +class GraphFramesUpstreamConnectTests(ReusedConnectTestCase): + _checkpoint_dir = tempfile.mkdtemp(prefix="spark-graphframes-connect-") + + @classmethod + def conf(cls) -> SparkConf: + return ( + super() + .conf() + .set("spark.driver.memory", "4g") + .set("spark.sql.shuffle.partitions", "4") + .set("spark.checkpoint.dir", cls._checkpoint_dir) + ) + + @classmethod + def setUpClass(cls) -> None: + super().setUpClass() + if cls._legacy_sc is not None: + cls._legacy_sc.setCheckpointDir(cls._checkpoint_dir) + + @classmethod + def tearDownClass(cls) -> None: + super().tearDownClass() + if os.path.exists(cls._checkpoint_dir): + shutil.rmtree(cls._checkpoint_dir) + + def local_graph(self) -> GraphFrame: + vertices = self.spark.createDataFrame([(1, "A"), (2, "B"), (3, "C")], ["id", "name"]) + edges = self.spark.createDataFrame( + [(1, 2, "love"), (2, 1, "hate"), (2, 3, "follow")], + ["src", "dst", "action"], + ) + return GraphFrame(vertices, edges) + + def test_construction(self) -> None: + upstream.test_construction(self.spark, self.local_graph()) + + def test_page_rank(self) -> None: + for args in upstream.PREGEL_ARGUMENTS: + with self.subTest(args=args): + upstream.test_page_rank(self.spark, args) + + def test_pregel_early_stopping(self) -> None: + for args in upstream.PREGEL_ARGUMENTS: + with self.subTest(args=args): + upstream.test_pregel_early_stopping(self.spark, args) + + def test_connected_components(self) -> None: + for args in upstream.PREGEL_ARGUMENTS: + for cc_args in [(-1, True), (10000, True), (-1, False), (10000, False)]: + with self.subTest(args=args, cc_args=cc_args): + upstream.test_connected_components(self.spark, args, cc_args) + + def test_connected_components2(self) -> None: + for args in upstream.PREGEL_ARGUMENTS: + for cc_args in [(-1, True), (10000, True), (-1, False), (10000, False)]: + with self.subTest(args=args, cc_args=cc_args): + upstream.test_connected_components2(self.spark, args, cc_args) + + def test_shortest_paths(self) -> None: + for args in upstream.PREGEL_ARGUMENTS: + with self.subTest(args=args): + upstream.test_shortest_paths(self.spark, args) + + def test_triangle_counts(self) -> None: + for storage_level in upstream.STORAGE_LEVELS: + with self.subTest(storage_level=storage_level): + upstream.test_triangle_counts(self.spark, storage_level) + + def test_cycles_finding(self) -> None: + for args in upstream.PREGEL_ARGUMENTS: + with self.subTest(args=args): + upstream.test_cycles_finding(self.spark, args) + + def test_mis(self) -> None: + for storage_level in upstream.STORAGE_LEVELS: + with self.subTest(storage_level=storage_level): + upstream.test_mis(self.spark, storage_level) + + def test_kcore(self) -> None: + for args in upstream.PREGEL_ARGUMENTS: + with self.subTest(args=args): + upstream.test_kcore(self.spark, args) + + +def _spark_test(function): + def run(self) -> None: + function(self.spark) + + run.__name__ = function.__name__ + return run + + +def _local_graph_test(function): + def run(self) -> None: + function(self.local_graph()) + + run.__name__ = function.__name__ + return run + + +for _test_function in [ + upstream.test_validate, + upstream.test_as_undirected, + upstream.test_as_reversed, + upstream.test_power_iteration_clustering, + upstream.test_graphframes_pagerank, + upstream.test_pregel_required_edge_columns, + upstream.test_connected_components_example, + upstream.test_shortest_paths2, + upstream.test_neighborhood_aware_cdlp_api_defaults, + upstream.test_neighborhood_aware_cdlp_api_with_all_args, + upstream.test_neighborhood_aware_cdlp_api_rejects_invalid_multiplier_combination, + upstream.test_strongly_connected_components, + upstream.test_approx_triangle_counts, + upstream.test_aggregate_neighbors_basic, + upstream.test_aggregate_neighbors_with_edge_filter, + upstream.test_aggregate_neighbors_multiple_accumulators, + upstream.test_hyper_anf_basic, + upstream.test_hyper_anf_args_passed, + upstream.test_hyper_anf_invalid_args, +]: + setattr(GraphFramesUpstreamConnectTests, _test_function.__name__, _spark_test(_test_function)) + + +for _test_function in [ + upstream.test_cache, + upstream.test_degrees, + upstream.test_type_degrees, + upstream.test_type_degrees_with_explicit_types, + upstream.test_motif_finding, + upstream.test_filterVertices, + upstream.test_filterEdges, + upstream.test_dropIsolatedVertices, + upstream.test_bfs, + upstream.test_all_paths, + upstream.test_random_walk_embeddings_api, + upstream.test_random_walk_embeddings_invalid_args, +]: + setattr( + GraphFramesUpstreamConnectTests, + _test_function.__name__, + _local_graph_test(_test_function), + ) + + +if __name__ == "__main__": + from pyspark.testing.unittestutils import main + + main() diff --git a/python/pyspark/graphframes/tests/pg/__init__.py b/python/pyspark/graphframes/tests/pg/__init__.py new file mode 100644 index 0000000000000..cce3acad34a49 --- /dev/null +++ b/python/pyspark/graphframes/tests/pg/__init__.py @@ -0,0 +1,16 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# diff --git a/python/pyspark/graphframes/tests/pg/test_property_graphframe.py b/python/pyspark/graphframes/tests/pg/test_property_graphframe.py new file mode 100644 index 0000000000000..1da5ed93eac22 --- /dev/null +++ b/python/pyspark/graphframes/tests/pg/test_property_graphframe.py @@ -0,0 +1,484 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import hashlib +import tempfile + +from pyspark.graphframes import GraphFrame +from pyspark.graphframes.pg import EdgePropertyGroup, PropertyGraphFrame, VertexPropertyGroup +from pyspark.graphframes.tests._upstream_test_utils import pytest +from pyspark.sql import SparkSession +from pyspark.sql.functions import lit +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +def sha256_hash(id_val, group_name): + """Helper to compute SHA256 hash like Scala does.""" + hash_val = hashlib.sha256(str(id_val).encode("utf-8")).hexdigest() + return f"{group_name}{hash_val}" + + +@pytest.fixture(scope="module") +def people_group(spark: SparkSession): + people_data = spark.createDataFrame( + [(1, "Alice"), (2, "Bob"), (3, "Charlie"), (4, "David"), (5, "Eve")], + ["id", "name"], + ) + return VertexPropertyGroup("people", people_data, "id") + + +@pytest.fixture(scope="module") +def movies_group(spark: SparkSession): + movies_data = spark.createDataFrame( + [(1, "Matrix"), (2, "Inception"), (3, "Interstellar")], + ["id", "title"], + ) + return VertexPropertyGroup("movies", movies_data, "id") + + +@pytest.fixture(scope="module") +def likes_group( + spark: SparkSession, people_group: VertexPropertyGroup, movies_group: VertexPropertyGroup +): + likes_data = spark.createDataFrame( + [(1, 1), (1, 2), (2, 1), (3, 2), (4, 3), (5, 2)], + ["src", "dst"], + ) + likes_data_with_weight = likes_data.withColumn("weight", lit(1.0)) + return EdgePropertyGroup( + "likes", + likes_data_with_weight, + people_group, + movies_group, + is_directed=False, + src_column_name="src", + dst_column_name="dst", + weight_column_name="weight", + ) + + +@pytest.fixture(scope="module") +def messages_group(spark: SparkSession, people_group: VertexPropertyGroup): + messages_data = spark.createDataFrame( + [(1, 2, 5.0), (2, 3, 8.0), (3, 4, 3.0), (4, 5, 6.0), (5, 1, 9.0)], + ["src", "dst", "weight"], + ) + return EdgePropertyGroup( + "messages", + messages_data, + people_group, + people_group, + is_directed=True, + src_column_name="src", + dst_column_name="dst", + weight_column_name="weight", + ) + + +@pytest.fixture(scope="module") +def people_movies_graph( + people_group: VertexPropertyGroup, + movies_group: VertexPropertyGroup, + likes_group: EdgePropertyGroup, + messages_group: EdgePropertyGroup, +): + return PropertyGraphFrame( + [people_group, movies_group], + [likes_group, messages_group], + ) + + +def test_property_graph_frame_constructor(people_movies_graph: PropertyGraphFrame) -> None: + assert len(people_movies_graph.vertex_property_groups) == 2 + assert len(people_movies_graph.edges_property_groups) == 2 + + +def test_vertex_property_group_creation(people_group: VertexPropertyGroup) -> None: + assert people_group.name == "people" + assert people_group.primary_key_column == "id" + assert people_group.apply_mask_on_id + + +def test_edge_property_group_creation( + likes_group: EdgePropertyGroup, +) -> None: + assert likes_group.name == "likes" + assert likes_group.src_property_group.name == "people" + assert likes_group.dst_property_group.name == "movies" + assert not likes_group.is_directed + + +def test_projection_by_movies(people_movies_graph: PropertyGraphFrame) -> None: + projected_graph = people_movies_graph.projection_by("people", "movies", "likes") + + assert len(projected_graph.vertex_property_groups) == 1 + assert projected_graph.vertex_property_groups[0].name == "people" + + assert len(projected_graph.edges_property_groups) == 2 + assert any(group.name == "messages" for group in projected_graph.edges_property_groups) + + projected_edges_group = next( + ( + group + for group in projected_graph.edges_property_groups + if group.name == "projected_likes" + ), + None, + ) + assert projected_edges_group is not None + assert projected_edges_group.src_column_name == GraphFrame.SRC + assert projected_edges_group.dst_column_name == GraphFrame.DST + assert projected_edges_group.weight_column_name == GraphFrame.WEIGHT + assert not projected_edges_group.is_directed + + projected_edges = projected_edges_group.data.collect() + edge_pairs = {(row.src, row.dst) for row in projected_edges} + + expected_edges = { + (1, 2), # Alice and Bob both like Matrix + (1, 3), # Alice and Charlie both like Inception + (1, 5), # Alice and Eve both like Inception + (3, 5), # Charlie and Eve both like Inception + } + assert edge_pairs == expected_edges + + +def test_projection_with_custom_weight(people_movies_graph: PropertyGraphFrame) -> None: + projected_graph = people_movies_graph.projection_by( + "people", "movies", "likes", new_edge_weight=lambda w1, w2: w1 + w2 + ) + + projected_edges_group = next( + ( + group + for group in projected_graph.edges_property_groups + if group.name == "projected_likes" + ), + None, + ) + assert projected_edges_group is not None + + projected_edges = projected_edges_group.data.collect() + edge_triples = {(row.src, row.dst, row.weight) for row in projected_edges} + + expected_edges = { + (1, 2, 2.0), + (1, 3, 2.0), + (1, 5, 2.0), + (3, 5, 2.0), + } + assert edge_triples == expected_edges + + +def test_to_graph_frame_messages_only(people_movies_graph: PropertyGraphFrame) -> None: + graph = people_movies_graph.to_graphframe( + vertex_property_groups=["people"], + edge_property_groups=["messages"], + edge_group_filters={"messages": lit(True)}, + vertex_group_filters={"people": lit(True)}, + ) + + vertices = {row.id for row in graph.vertices.collect()} + edges = {(row.src, row.dst, row.weight) for row in graph.edges.collect()} + + expected_vertices = {sha256_hash(i, "people") for i in range(1, 6)} + assert vertices == expected_vertices + + expected_edges = { + (sha256_hash(1, "people"), sha256_hash(2, "people"), 5.0), + (sha256_hash(2, "people"), sha256_hash(3, "people"), 8.0), + (sha256_hash(3, "people"), sha256_hash(4, "people"), 3.0), + (sha256_hash(4, "people"), sha256_hash(5, "people"), 6.0), + (sha256_hash(5, "people"), sha256_hash(1, "people"), 9.0), + } + assert edges == expected_edges + + +def test_to_graph_frame_all_groups(people_movies_graph: PropertyGraphFrame) -> None: + graph = people_movies_graph.to_graphframe( + vertex_property_groups=["people", "movies"], + edge_property_groups=["messages", "likes"], + edge_group_filters={"messages": lit(True), "likes": lit(True)}, + vertex_group_filters={"people": lit(True), "movies": lit(True)}, + ) + + vertices = graph.vertices.collect() + edges = graph.edges.collect() + + assert len(vertices) == 8 # 5 people + 3 movies + + vertex_ids = {row.id for row in vertices} + assert sha256_hash(1, "movies") in vertex_ids + assert sha256_hash(1, "people") in vertex_ids + + message_edges = [e for e in edges if e.weight != 1.0] + like_edges = [e for e in edges if e.weight == 1.0] + + assert len(message_edges) == 5 # Directed messages + assert len(like_edges) == 12 # 6 undirected edges * 2 + + +def test_to_graph_frame_unmasked_ids( + spark: SparkSession, + people_group: VertexPropertyGroup, + likes_group: EdgePropertyGroup, + messages_group: EdgePropertyGroup, +) -> None: + movies_data = spark.createDataFrame( + [(1, "Matrix"), (2, "Inception"), (3, "Interstellar")], + ["id", "title"], + ) + unmasked_movies_group = VertexPropertyGroup("movies", movies_data, "id", apply_mask_on_id=False) + + new_likes_group = EdgePropertyGroup( + "likes", + likes_group.data, + likes_group.src_property_group, + unmasked_movies_group, + likes_group.is_directed, + likes_group.src_column_name, + likes_group.dst_column_name, + likes_group.weight_column_name, + ) + + modified_graph = PropertyGraphFrame( + [people_group, unmasked_movies_group], + [new_likes_group, messages_group], + ) + + graph = modified_graph.to_graphframe( + vertex_property_groups=["people", "movies"], + edge_property_groups=["messages", "likes"], + edge_group_filters={"messages": lit(True), "likes": lit(True)}, + vertex_group_filters={"people": lit(True), "movies": lit(True)}, + ) + + vertices = {row.id for row in graph.vertices.collect()} + edges = graph.edges.collect() + + assert "1" in vertices + assert "2" in vertices + assert "3" in vertices + assert sha256_hash(1, "people") in vertices + + likes_edges = [e for e in edges if e.weight == 1.0] + assert any(e.src == sha256_hash(1, "people") and e.dst == "1" for e in likes_edges) + assert any(e.src == "1" and e.dst == sha256_hash(1, "people") for e in likes_edges) + + +def test_join_vertices_with_connected_components( + people_movies_graph: PropertyGraphFrame, +) -> None: + graph = people_movies_graph.to_graphframe( + vertex_property_groups=["people", "movies"], + edge_property_groups=["messages", "likes"], + edge_group_filters={"messages": lit(True), "likes": lit(True)}, + vertex_group_filters={"people": lit(True), "movies": lit(True)}, + ) + + components = graph.connectedComponents() + + joined_back = people_movies_graph.join_vertices(components, vertex_groups=["people", "movies"]) + + joined_data = joined_back.collect() + + by_group = {} + for row in joined_data: + group = row.property_group + if group not in by_group: + by_group[group] = [] + by_group[group].append(row) + + assert "movies" in by_group + assert "people" in by_group + assert len(by_group["movies"]) == 3 + assert len(by_group["people"]) == 5 + + +def test_vertex_property_group_validation(people_group: VertexPropertyGroup) -> None: + from pyspark.graphframes.pg.property_groups import InvalidPropertyGroupException + + with pytest.raises(InvalidPropertyGroupException): + VertexPropertyGroup("test", people_group.data, "nonexistent_column") + + +def test_edge_property_group_validation( + people_group: VertexPropertyGroup, + movies_group: VertexPropertyGroup, + likes_group: EdgePropertyGroup, +) -> None: + from pyspark.graphframes.pg.property_groups import InvalidPropertyGroupException + + with pytest.raises(InvalidPropertyGroupException): + EdgePropertyGroup( + "test", + likes_group.data, + people_group, + movies_group, + is_directed=True, + src_column_name="nonexistent", + dst_column_name="dst", + weight_column_name="weight", + ) + + with pytest.raises(InvalidPropertyGroupException): + EdgePropertyGroup( + "test", + likes_group.data, + people_group, + movies_group, + is_directed=True, + src_column_name="src", + dst_column_name="nonexistent", + weight_column_name="weight", + ) + + with pytest.raises(InvalidPropertyGroupException): + EdgePropertyGroup( + "test", + likes_group.data, + people_group, + movies_group, + is_directed=True, + src_column_name="src", + dst_column_name="dst", + weight_column_name="nonexistent", + ) + + +def test_to_graph_frame_invalid_group(people_movies_graph: PropertyGraphFrame) -> None: + with pytest.raises(ValueError): + people_movies_graph.to_graphframe( + vertex_property_groups=["nonexistent"], + edge_property_groups=["likes"], + ) + + with pytest.raises(ValueError): + people_movies_graph.to_graphframe( + vertex_property_groups=["people"], + edge_property_groups=["nonexistent"], + ) + + +def test_projection_by_invalid_group(people_movies_graph: PropertyGraphFrame) -> None: + with pytest.raises(ValueError): + people_movies_graph.projection_by("nonexistent", "movies", "likes") + + with pytest.raises(ValueError): + people_movies_graph.projection_by("people", "nonexistent", "likes") + + with pytest.raises(ValueError): + people_movies_graph.projection_by("people", "movies", "nonexistent") + + +def test_property_graph_frame_to_graph_frame_conversion( + people_movies_graph: PropertyGraphFrame, +) -> None: + graph = people_movies_graph.to_graphframe( + vertex_property_groups=["people"], + edge_property_groups=["messages"], + ) + + assert isinstance(graph, GraphFrame) + assert GraphFrame.ID in graph.vertices.columns + assert GraphFrame.SRC in graph.edges.columns + assert GraphFrame.DST in graph.edges.columns + assert GraphFrame.WEIGHT in graph.edges.columns + + +class PropertyGraphFrameTests(ReusedSQLTestCase): + @classmethod + def setUpClass(cls) -> None: + super().setUpClass() + cls._checkpoint_dir = tempfile.TemporaryDirectory() + cls.spark.sparkContext.setCheckpointDir(cls._checkpoint_dir.name) + cls.spark.conf.set("spark.sql.shuffle.partitions", "4") + + @classmethod + def tearDownClass(cls) -> None: + super().tearDownClass() + cls._checkpoint_dir.cleanup() + + def groups(self): + people = people_group(self.spark) + movies = movies_group(self.spark) + likes = likes_group(self.spark, people, movies) + messages = messages_group(self.spark, people) + graph = people_movies_graph(people, movies, likes, messages) + return people, movies, likes, messages, graph + + def test_property_graph_frame_constructor(self) -> None: + *_, graph = self.groups() + test_property_graph_frame_constructor(graph) + + def test_vertex_property_group_creation(self) -> None: + people, *_ = self.groups() + test_vertex_property_group_creation(people) + + def test_edge_property_group_creation(self) -> None: + _, _, likes, _, _ = self.groups() + test_edge_property_group_creation(likes) + + def test_projection_by_movies(self) -> None: + *_, graph = self.groups() + test_projection_by_movies(graph) + + def test_projection_with_custom_weight(self) -> None: + *_, graph = self.groups() + test_projection_with_custom_weight(graph) + + def test_to_graph_frame_messages_only(self) -> None: + *_, graph = self.groups() + test_to_graph_frame_messages_only(graph) + + def test_to_graph_frame_all_groups(self) -> None: + *_, graph = self.groups() + test_to_graph_frame_all_groups(graph) + + def test_to_graph_frame_unmasked_ids(self) -> None: + people, _, likes, messages, _ = self.groups() + test_to_graph_frame_unmasked_ids(self.spark, people, likes, messages) + + def test_join_vertices_with_connected_components(self) -> None: + *_, graph = self.groups() + test_join_vertices_with_connected_components(graph) + + def test_vertex_property_group_validation(self) -> None: + people, *_ = self.groups() + test_vertex_property_group_validation(people) + + def test_edge_property_group_validation(self) -> None: + people, movies, likes, _, _ = self.groups() + test_edge_property_group_validation(people, movies, likes) + + def test_to_graph_frame_invalid_group(self) -> None: + *_, graph = self.groups() + test_to_graph_frame_invalid_group(graph) + + def test_projection_by_invalid_group(self) -> None: + *_, graph = self.groups() + test_projection_by_invalid_group(graph) + + def test_property_graph_frame_to_graph_frame_conversion(self) -> None: + *_, graph = self.groups() + test_property_graph_frame_to_graph_frame_conversion(graph) + + +if __name__ == "__main__": + from pyspark.testing.unittestutils import main + + main() diff --git a/python/pyspark/graphframes/tests/test_graphframe.py b/python/pyspark/graphframes/tests/test_graphframe.py index 339f109d70e76..3d005aba3adfa 100644 --- a/python/pyspark/graphframes/tests/test_graphframe.py +++ b/python/pyspark/graphframes/tests/test_graphframe.py @@ -16,6 +16,7 @@ # from pyspark.graphframes import GraphFrame +from pyspark.graphframes.lib import Pregel from pyspark.sql import Row from pyspark.sql import functions as F from pyspark.testing.sqlutils import ReusedSQLTestCase @@ -109,6 +110,11 @@ def test_validate(self) -> None: with self.assertRaisesRegex(ValueError, "duplicate vertices"): GraphFrame(vertices, edges).validate() + def test_direct_pregel_construction(self) -> None: + pregel = Pregel(self.graph) + self.assertIs(pregel.setMaxIter(1), pregel) + self.assertIsInstance(pregel, Pregel) + if __name__ == "__main__": from pyspark.graphframes.tests.test_graphframe import * # noqa: F403 diff --git a/python/pyspark/graphframes/tests/test_graphframes.py b/python/pyspark/graphframes/tests/test_graphframes.py new file mode 100644 index 0000000000000..4a96a8a7f9465 --- /dev/null +++ b/python/pyspark/graphframes/tests/test_graphframes.py @@ -0,0 +1,1311 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + + +import tempfile +from dataclasses import dataclass + +from pyspark import SparkConf +from pyspark.graphframes.graphframe import AggregateNeighbors, GraphFrame, RandomWalkEmbeddings +from pyspark.graphframes.tests._upstream_test_utils import pytest +from pyspark.sql import DataFrame, SparkSession +from pyspark.sql import functions as sqlfunctions +from pyspark.sql.utils import is_remote +from pyspark.storagelevel import StorageLevel +from pyspark.testing.sqlutils import ReusedSQLTestCase + + +@dataclass +class PregelArguments: + algorithm: str + use_local_checkpoints: bool + checkpoint_interval: int + storage_level: StorageLevel + + +PREGEL_ARGUMENTS = [ + PregelArguments("graphframes", True, 5, StorageLevel.MEMORY_AND_DISK), + PregelArguments("graphx", False, 3, StorageLevel.DISK_ONLY), + PregelArguments("graphframes", False, 7, StorageLevel.MEMORY_ONLY), + PregelArguments("graphframes", True, 1, StorageLevel.DISK_ONLY_3), +] +PREGEL_IDS: list[str] = [ + "graphframes,local,5,MEMORY_AND_DISK", + "graphx,global,3,DISK_ONLY", + "graphframes,global,7,MEMORY_ONLY", + "graphframes,local,1,DISK_ONLY_3", +] +STORAGE_LEVELS = [ + StorageLevel.MEMORY_AND_DISK_2, + StorageLevel.DISK_ONLY, + StorageLevel.MEMORY_ONLY, +] +STORAGE_LEVELS_IDS = [ + "MEMORY_AND_DISK_2", + "DISK_ONLY", + "MEMORY_ONLY", +] + + +def test_construction(spark: SparkSession, local_g: GraphFrame) -> None: + vertexIDs = [row[0] for row in local_g.vertices.select("id").collect()] + assert sorted(vertexIDs) == [1, 2, 3] + + edgeActions = [row[0] for row in local_g.edges.select("action").collect()] + assert sorted(edgeActions) == ["follow", "hate", "love"] + tripletsFirst = list( + map( + lambda x: (x[0][1], x[1][1], x[2][2]), + local_g.triplets.sort("src.id").select("src", "dst", "edge").take(1), + ) + ) + assert tripletsFirst == [("A", "B", "love")], tripletsFirst + + # Try with invalid vertices and edges DataFrames + v_invalid = spark.createDataFrame( + [(1, "A"), (2, "B"), (3, "C")], ["invalid_colname_1", "invalid_colname_2"] + ) + e_invalid = spark.createDataFrame( + [(1, 2), (2, 3), (3, 1)], ["invalid_colname_3", "invalid_colname_4"] + ) + with pytest.raises(ValueError): + _ = GraphFrame(v_invalid, e_invalid) + + +def test_validate(spark: SparkSession) -> None: + good_g = GraphFrame( + spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")]).toDF("id", "attr"), + spark.createDataFrame([(1, 2), (2, 1), (2, 3)]).toDF("src", "dst"), + ) + good_g.validate() # no exception should be thrown + + not_distinct_vertices = GraphFrame( + spark.createDataFrame([(1, "a"), (2, "b"), (3, "c"), (1, "d")]).toDF("id", "attr"), + spark.createDataFrame([(1, 2), (2, 1), (2, 3)]).toDF("src", "dst"), + ) + with pytest.raises(ValueError): + not_distinct_vertices.validate() + + missing_vertices = GraphFrame( + spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")]).toDF("id", "attr"), + spark.createDataFrame([(1, 2), (2, 1), (2, 3), (1, 4)]).toDF("src", "dst"), + ) + with pytest.raises(ValueError): + missing_vertices.validate() + + +def test_as_undirected(spark: SparkSession) -> None: + # Test without edge attributes + v = spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")]).toDF("id", "name") + e = spark.createDataFrame([(1, 2), (2, 3)]).toDF("src", "dst") + g = GraphFrame(v, e) + undirected = g.as_undirected() + + # Check edge count doubled + assert undirected.edges.count() == 2 * g.edges.count() + + # Verify reverse edges exist + edges = undirected.edges.sort("src", "dst").collect() + assert len(edges) == 4 + assert edges[0][0] == 1 + assert edges[0][1] == 2 + assert edges[1][0] == 2 + assert edges[1][1] == 1 + assert edges[2][0] == 2 + assert edges[2][1] == 3 + assert edges[3][0] == 3 + assert edges[3][1] == 2 + + # Test with edge attributes + v2 = spark.createDataFrame([(1, "a"), (2, "b")]).toDF("id", "name") + e2 = spark.createDataFrame([(1, 2, "edge1")]).toDF("src", "dst", "attr") + g2 = GraphFrame(v2, e2) + undirected2 = g2.as_undirected() + + edges2 = undirected2.edges.collect() + assert len(edges2) == 2 + assert any(row[0] == 1 and row[1] == 2 and row[2] == "edge1" for row in edges2) + assert any(row[0] == 2 and row[1] == 1 and row[2] == "edge1" for row in edges2) + + +def test_as_reversed(spark: SparkSession) -> None: + # Test without edge attributes + v = spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")]).toDF("id", "name") + e = spark.createDataFrame([(1, 2), (2, 3)]).toDF("src", "dst") + g = GraphFrame(v, e) + reversed_g = g.as_reversed() + + # Check edge count is the same + assert reversed_g.edges.count() == g.edges.count() + + # Verify edges are reversed + edges = reversed_g.edges.sort("src", "dst").collect() + assert len(edges) == 2 + assert edges[0][0] == 2 + assert edges[0][1] == 1 + assert edges[1][0] == 3 + assert edges[1][1] == 2 + + # Test with edge attributes + v2 = spark.createDataFrame([(1, "a"), (2, "b")]).toDF("id", "name") + e2 = spark.createDataFrame([(1, 2, "edge1")]).toDF("src", "dst", "attr") + g2 = GraphFrame(v2, e2) + reversed2 = g2.as_reversed() + + edges2 = reversed2.edges.collect() + assert len(edges2) == 1 + assert edges2[0][0] == 2 + assert edges2[0][1] == 1 + assert edges2[0][2] == "edge1" + + +def test_cache(local_g: GraphFrame) -> None: + _ = local_g.cache() + _ = local_g.unpersist() + + +def test_degrees(local_g: GraphFrame) -> None: + outDeg = local_g.outDegrees + assert set(outDeg.columns) == {"id", "outDegree"} + inDeg = local_g.inDegrees + assert set(inDeg.columns) == {"id", "inDegree"} + deg = local_g.degrees + assert set(deg.columns) == {"id", "degree"} + + +def test_type_degrees(local_g: GraphFrame) -> None: + type_out_degree = local_g.type_out_degree("action") + assert set(type_out_degree.columns) == {"id", "outDegrees"} + + schema = type_out_degree.schema["outDegrees"].dataType + field_names = {field.name for field in schema.fields} + assert field_names == {"love", "hate", "follow"} + + results = {row.id: row.outDegrees for row in type_out_degree.collect()} + assert results[1].love == 1 + assert results[1].hate == 0 + assert results[1].follow == 0 + assert results[2].love == 0 + assert results[2].hate == 1 + assert results[2].follow == 1 + + type_in_degree = local_g.type_in_degree("action") + assert set(type_in_degree.columns) == {"id", "inDegrees"} + + schema = type_in_degree.schema["inDegrees"].dataType + field_names = {field.name for field in schema.fields} + assert field_names == {"love", "hate", "follow"} + + results = {row.id: row.inDegrees for row in type_in_degree.collect()} + assert results[1].love == 0 + assert results[1].hate == 1 + assert results[1].follow == 0 + assert results[2].love == 1 + assert results[2].hate == 0 + assert results[2].follow == 0 + assert results[3].love == 0 + assert results[3].hate == 0 + assert results[3].follow == 1 + + type_degree = local_g.type_degree("action") + assert set(type_degree.columns) == {"id", "degrees"} + + schema = type_degree.schema["degrees"].dataType + field_names = {field.name for field in schema.fields} + assert field_names == {"love", "hate", "follow"} + + results = {row.id: row.degrees for row in type_degree.collect()} + assert results[1].love == 1 + assert results[1].hate == 1 + assert results[1].follow == 0 + assert results[2].love == 1 + assert results[2].hate == 1 + assert results[2].follow == 1 + assert results[3].love == 0 + assert results[3].hate == 0 + assert results[3].follow == 1 + + +def test_type_degrees_with_explicit_types(local_g: GraphFrame) -> None: + edge_types = ["love", "hate", "follow"] + type_out_degree = local_g.type_out_degree("action", edge_types) + assert set(type_out_degree.columns) == {"id", "outDegrees"} + + schema = type_out_degree.schema["outDegrees"].dataType + field_names = {field.name for field in schema.fields} + assert field_names == {"love", "hate", "follow"} + + results = {row.id: row.outDegrees for row in type_out_degree.collect()} + assert results[1].love == 1 + assert results[1].hate == 0 + assert results[1].follow == 0 + assert results[2].love == 0 + assert results[2].hate == 1 + assert results[2].follow == 1 + + type_in_degree = local_g.type_in_degree("action", edge_types) + assert set(type_in_degree.columns) == {"id", "inDegrees"} + + results = {row.id: row.inDegrees for row in type_in_degree.collect()} + assert results[1].love == 0 + assert results[1].hate == 1 + assert results[1].follow == 0 + assert results[2].love == 1 + assert results[2].hate == 0 + assert results[2].follow == 0 + assert results[3].love == 0 + assert results[3].hate == 0 + assert results[3].follow == 1 + + type_degree = local_g.type_degree("action", edge_types) + assert set(type_degree.columns) == {"id", "degrees"} + + results = {row.id: row.degrees for row in type_degree.collect()} + assert results[1].love == 1 + assert results[1].hate == 1 + assert results[1].follow == 0 + assert results[2].love == 1 + assert results[2].hate == 1 + assert results[2].follow == 1 + assert results[3].love == 0 + assert results[3].hate == 0 + assert results[3].follow == 1 + + +def test_motif_finding(local_g: GraphFrame) -> None: + motifs = local_g.find("(a)-[e]->(b)") + assert motifs.count() == 3 + assert set(motifs.columns) == {"a", "e", "b"} + + +def test_filterVertices(local_g: GraphFrame) -> None: + conditions = ["id < 3", local_g.vertices.id < 3] + expected_v = [(1, "A"), (2, "B")] + expected_e = [(1, 2, "love"), (2, 1, "hate")] + for cond in conditions: + g2 = local_g.filterVertices(cond) + v2 = g2.vertices.select("id", "name").collect() + e2 = g2.edges.select("src", "dst", "action").collect() + assert len(v2) == len(expected_v) + assert len(e2) == len(expected_e) + assert set(v2) == set(expected_v) + assert set(e2) == set(expected_e) + + +def test_filterEdges(local_g: GraphFrame) -> None: + conditions = ["dst > 2", local_g.edges.dst > 2] + expected_v = [(1, "A"), (2, "B"), (3, "C")] + expected_e = [(2, 3, "follow")] + for cond in conditions: + g2 = local_g.filterEdges(cond) + v2 = g2.vertices.select("id", "name").collect() + e2 = g2.edges.select("src", "dst", "action").collect() + assert len(v2) == len(expected_v) + assert len(e2) == len(expected_e) + assert set(v2) == set(expected_v) + assert set(e2) == set(expected_e) + + +def test_dropIsolatedVertices(local_g: GraphFrame) -> None: + g2 = local_g.filterEdges("dst > 2").dropIsolatedVertices() + v2 = g2.vertices.select("id", "name").collect() + e2 = g2.edges.select("src", "dst", "action").collect() + expected_v = [(2, "B"), (3, "C")] + expected_e = [(2, 3, "follow")] + assert len(v2) == len(expected_v) + assert len(e2) == len(expected_e) + assert set(v2) == set(expected_v) + assert set(e2) == set(expected_e) + + +def test_bfs(local_g: GraphFrame) -> None: + paths = local_g.bfs("name='A'", "name='C'") + assert paths is not None + assert paths.count() == 1 + # Expecting that the first intermediary vertex in the BFS is "B" + head = paths.select("v1.name").head() + assert head is not None + assert head[0] == "B" + + paths2 = local_g.bfs("name='A'", "name='C'", edgeFilter="action!='follow'") + assert paths2.count() == 0 + + paths3 = local_g.bfs("name='A'", "name='C'", maxPathLength=1) + assert paths3.count() == 0 + + +def test_all_paths(local_g: GraphFrame) -> None: + # local_g: A->B (love), B->A (hate), B->C (follow) + # Directed: A can reach C via A->B->C (1 path) + paths = local_g.all_paths("name='A'", "name='C'", use_local_checkpoints=True) + assert paths is not None + assert {"path", "len"}.issubset(set(paths.columns)) + paths_list = paths.collect() + assert len(paths_list) == 1 + assert paths_list[0]["len"] == 2 + + # With edge filter that removes the 'follow' edge: no path A->C + paths_filtered = local_g.all_paths( + "name='A'", "name='C'", edge_filter="action!='follow'", use_local_checkpoints=True + ) + assert paths_filtered.count() == 0 + + # Undirected: A->B->C and A->B->A (no simple path to C via A again), + # but also C->B->A is now possible, still only A->B->C from A to C + paths_undirected = local_g.all_paths( + "name='A'", "name='C'", is_directed=False, use_local_checkpoints=True + ) + assert paths_undirected.count() >= 1 + + # max_path_length too short: no paths + paths_short = local_g.all_paths( + "name='A'", "name='C'", max_path_length=1, use_local_checkpoints=True + ) + assert paths_short.count() == 0 + + +def test_power_iteration_clustering(spark: SparkSession) -> None: + vertices = [ + (1, 0, 0.5), + (2, 0, 0.5), + (2, 1, 0.7), + (3, 0, 0.5), + (3, 1, 0.7), + (3, 2, 0.9), + (4, 0, 0.5), + (4, 1, 0.7), + (4, 2, 0.9), + (4, 3, 1.1), + (5, 0, 0.5), + (5, 1, 0.7), + (5, 2, 0.9), + (5, 3, 1.1), + (5, 4, 1.3), + ] + edges = [(0,), (1,), (2,), (3,), (4,), (5,)] + g = GraphFrame( + v=spark.createDataFrame(edges).toDF("id"), + e=spark.createDataFrame(vertices).toDF("src", "dst", "weight"), + ) + clusters_df = g.powerIterationClustering(k=2, maxIter=40, weightCol="weight") + + clusters = [r["cluster"] for r in clusters_df.sort("id").collect()] + + if is_remote(): + # It returns different results on Connect/Classic; + # For connect mode it works like a smoke-test + assert len(clusters) == 6 + else: + assert clusters == [0, 0, 0, 0, 1, 0] + _ = clusters_df.unpersist() + + +@pytest.mark.parametrize("args", PREGEL_ARGUMENTS, ids=PREGEL_IDS) +def test_page_rank(spark: SparkSession, args: PregelArguments) -> None: + edges = spark.createDataFrame( + [ + [0, 1], + [1, 2], + [2, 4], + [2, 0], + [3, 4], # 3 has no in-links + [4, 0], + [4, 2], + ], + ["src", "dst"], + ) + _ = edges.cache() + vertices = spark.createDataFrame([[0], [1], [2], [3], [4]], ["id"]) + numVertices = vertices.count() + + vertices = GraphFrame(vertices, edges).outDegrees + _ = vertices.toPandas().head() + _ = vertices.cache() + + # Construct a new GraphFrame with the updated vertices DataFrame. + graph = GraphFrame(vertices, edges) + alpha = 0.15 + pregel = graph.pregel + ranks = ( + graph.pregel.setMaxIter(5) + .withVertexColumn( + "rank", + sqlfunctions.lit(1.0 / numVertices), + sqlfunctions.coalesce(pregel.msg(), sqlfunctions.lit(0.0)) + * sqlfunctions.lit(1.0 - alpha) + + sqlfunctions.lit(alpha / numVertices), + ) + .sendMsgToDst(pregel.src("rank") / pregel.src("outDegree")) + .aggMsgs(sqlfunctions.sum(pregel.msg())) + .run() + ) + resultRows = ranks.sort("id").collect() + result = map(lambda x: x.rank, resultRows) + expected = [0.245, 0.224, 0.303, 0.03, 0.197] + + # Compare each result with its expected value using a tolerance of 1e-3. + for a, b in zip(result, expected): + assert a == pytest.approx(b, abs=1e-3) + _ = ranks.unpersist() + + +def test_graphframes_pagerank(spark: SparkSession) -> None: + """Regression test for graphframes/graphframes#889: pageRank fails on Spark Connect due to + AttributeError accessing self.edges and self.outDegrees on GraphFrameConnect.""" + edges = spark.createDataFrame( + [ + [0, 1], + [1, 2], + [2, 4], + [2, 0], + [3, 4], + [4, 0], + [4, 2], + ], + ["src", "dst"], + ) + vertices = spark.createDataFrame([[0], [1], [2], [3], [4]], ["id"]) + g = GraphFrame(vertices, edges) + + result = g.pageRank(resetProbability=0.15, maxIter=3) + + assert "pagerank" in result.vertices.columns + assert "weight" in result.edges.columns + assert result.vertices.count() == 5 + assert result.edges.count() == edges.count() + + +@pytest.mark.parametrize("args", PREGEL_ARGUMENTS, ids=PREGEL_IDS) +def test_pregel_early_stopping(spark: SparkSession, args: PregelArguments) -> None: + edges = spark.createDataFrame( + [ + [0, 1], + [1, 2], + [2, 4], + [2, 0], + [3, 4], # 3 has no in-links + [4, 0], + [4, 2], + ], + ["src", "dst"], + ) + _ = edges.cache() + vertices = spark.createDataFrame([[0], [1], [2], [3], [4]], ["id"]) + numVertices = vertices.count() + + vertices = GraphFrame(vertices, edges).outDegrees + _ = vertices.toPandas().head() + _ = vertices.cache() + + # Construct a new GraphFrame with the updated vertices DataFrame. + graph = GraphFrame(vertices, edges) + alpha = 0.15 + pregel = graph.pregel + ranks = ( + graph.pregel.setMaxIter(5) + .setUseLocalCheckpoints(args.use_local_checkpoints) + .setIntermediateStorageLevel(args.storage_level) + .setCheckpointInterval(args.checkpoint_interval) + .setEarlyStopping(True) + .setUseLocalCheckpoints(args.use_local_checkpoints) + .setIntermediateStorageLevel(args.storage_level) + .setCheckpointInterval(args.checkpoint_interval) + .withVertexColumn( + "rank", + sqlfunctions.lit(1.0 / numVertices), + sqlfunctions.coalesce(pregel.msg(), sqlfunctions.lit(0.0)) + * sqlfunctions.lit(1.0 - alpha) + + sqlfunctions.lit(alpha / numVertices), + ) + .sendMsgToDst(pregel.src("rank") / pregel.src("outDegree")) + .aggMsgs(sqlfunctions.sum(pregel.msg())) + .run() + ) + resultRows = ranks.sort("id").collect() + result = map(lambda x: x.rank, resultRows) + expected = [0.245, 0.224, 0.303, 0.03, 0.197] + + # Compare each result with its expected value using a tolerance of 1e-3. + for a, b in zip(result, expected): + assert a == pytest.approx(b, abs=1e-3) + _ = ranks.unpersist() + + +def test_pregel_required_edge_columns(spark: SparkSession) -> None: + edges = spark.createDataFrame( + [(0, 1, 0.5), (1, 2, 1.0), (2, 0, 0.3)], + ["src", "dst", "weight"], + ) + vertices = spark.createDataFrame([(0,), (1,), (2,)], ["id"]) + graph = GraphFrame(vertices, edges) + pregel = graph.pregel + + result = ( + graph.pregel.setMaxIter(2) + .withVertexColumn( + "value", + sqlfunctions.lit(0.0), + sqlfunctions.coalesce(pregel.msg(), sqlfunctions.lit(0.0)), + ) + .sendMsgToDst(pregel.src("value") + pregel.edge("weight")) + .aggMsgs(sqlfunctions.sum(pregel.msg())) + .required_edge_columns("weight") + .run() + ) + assert "value" in result.columns + assert result.count() == 3 + _ = result.unpersist() + + +def _df_hasCols(df: DataFrame, vcols: list[str] = []) -> None: + for c in vcols: + assert c in df.columns, f"DataFrame missing column: {c}" + + +@pytest.mark.parametrize("args", PREGEL_ARGUMENTS, ids=PREGEL_IDS) +@pytest.mark.parametrize( + "cc_args", + [(-1, True), (10000, True), (-1, False), (10000, False)], + ids=["aqe,local", "skewed,local", "aqe,checkpoints", "skewed,checkpoints"], +) +def test_connected_components( + spark: SparkSession, args: PregelArguments, cc_args: tuple[int, bool] +) -> None: + v = spark.createDataFrame([(0, "a", "b")], ["id", "vattr", "gender"]) + e = spark.createDataFrame([(0, 0, 1)], ["src", "dst", "test"]) + g = GraphFrame(v, e) + comps = g.connectedComponents( + algorithm=args.algorithm, + checkpointInterval=args.checkpoint_interval, + use_local_checkpoints=args.use_local_checkpoints, + storage_level=args.storage_level, + broadcastThreshold=cc_args[0], + useLabelsAsComponents=cc_args[1], + ) + _df_hasCols(comps, vcols=["id", "component", "vattr", "gender"]) + assert comps.count() == 1 + _ = comps.unpersist() + + +@pytest.mark.parametrize("args", PREGEL_ARGUMENTS, ids=PREGEL_IDS) +@pytest.mark.parametrize( + "cc_args", + [(-1, True), (10000, True), (-1, False), (10000, False)], + ids=["aqe,local", "skewed,local", "aqe,checkpoints", "skewed,checkpoints"], +) +def test_connected_components2( + spark: SparkSession, args: PregelArguments, cc_args: tuple[int, bool] +) -> None: + v = spark.createDataFrame([(0, "a0", "b0"), (1, "a1", "b1")], ["id", "A", "B"]) + e = spark.createDataFrame([(0, 1, "a01", "b01")], ["src", "dst", "A", "B"]) + g = GraphFrame(v, e) + comps = g.connectedComponents( + algorithm=args.algorithm, + checkpointInterval=args.checkpoint_interval, + use_local_checkpoints=args.use_local_checkpoints, + storage_level=args.storage_level, + broadcastThreshold=cc_args[0], + useLabelsAsComponents=cc_args[1], + ) + _df_hasCols(comps, vcols=["id", "component", "A", "B"]) + assert comps.count() == 2 + _ = comps.unpersist() + + +def test_connected_components_example(spark: SparkSession) -> None: + nodes = [(1, "Alice", 30), (2, "Bob", 25), (3, "Charlie", 35)] + nodes_df = spark.createDataFrame(nodes, ["id", "name", "age"]) + + edges = [ + (1, 2, "friend"), + (2, 1, "friend"), + (2, 3, "friend"), + (3, 2, "enemy"), # eek! + ] + edges_df = spark.createDataFrame(edges, ["src", "dst", "relationship"]) + + g = GraphFrame(nodes_df, edges_df) + cc = g.connectedComponents() + cc.write.mode("overwrite").format("noop").save() + res = cc.collect() + assert len(res) == 3 + _ = cc.unpersist() + + +@pytest.mark.parametrize("args", PREGEL_ARGUMENTS, ids=PREGEL_IDS) +def test_shortest_paths(spark: SparkSession, args: PregelArguments) -> None: + edges = [(1, 2), (1, 5), (2, 3), (2, 5), (3, 4), (4, 5), (4, 6)] + # Create bidirectional edges. + all_edges = [z for (a, b) in edges for z in [(a, b), (b, a)]] + edges = spark.createDataFrame(all_edges, ["src", "dst"]) + edges = spark.createDataFrame(all_edges, ["src", "dst"]) + edgesDF = spark.createDataFrame(all_edges, ["src", "dst"]) + vertices = spark.createDataFrame([(i,) for i in range(1, 7)], ["id"]) + g = GraphFrame(vertices, edgesDF) + landmarks: list[str | int] = [1, 4] + v2 = g.shortestPaths( + landmarks=landmarks, + algorithm=args.algorithm, + use_local_checkpoints=args.use_local_checkpoints, + checkpoint_interval=args.checkpoint_interval, + storage_level=args.storage_level, + ) + _df_hasCols(v2, vcols=["id", "distances"]) + _ = v2.unpersist() + + +def test_shortest_paths2(spark: SparkSession) -> None: + # Create an undirected graph + vertices = spark.createDataFrame([(i,) for i in range(1, 6)], ["id"]) + edges = spark.createDataFrame([(1, 2), (2, 3), (3, 4), (4, 5)], ["src", "dst"]) + g = GraphFrame(vertices, edges) + landmarks = [1] + result = g.shortestPaths(landmarks=landmarks, is_directed=False) + + # Check that distances are correct + distances = result.sort("id").select("id", "distances").collect() + + assert distances[0]["distances"] == {1: 0} + assert distances[1]["distances"] == {1: 1} + assert distances[2]["distances"] == {1: 2} + assert distances[3]["distances"] == {1: 3} + assert distances[4]["distances"] == {1: 4} + + _ = result.unpersist() + + +def test_neighborhood_aware_cdlp_api_defaults(spark: SparkSession) -> None: + if spark.version[:3] < "4.1": + pytest.skip("NeighborhoodAwareCDLP requires Spark >= 4.1") + + vertices = spark.createDataFrame([(1,), (2,), (3,)], ["id"]) + edges = spark.createDataFrame([(1, 2), (2, 3), (3, 1)], ["src", "dst"]) + g = GraphFrame(vertices, edges) + + result = g.neighborhood_aware_cdlp(max_iter=1) + _df_hasCols(result, vcols=["id", "label"]) + _ = result.unpersist() + + +def test_neighborhood_aware_cdlp_api_with_all_args(spark: SparkSession) -> None: + if spark.version[:3] < "4.1": + pytest.skip("NeighborhoodAwareCDLP requires Spark >= 4.1") + + vertices = spark.createDataFrame( + [(1, "A"), (2, "B"), (3, "C"), (4, "D")], + ["id", "seed_label"], + ) + edges = spark.createDataFrame([(1, 2), (2, 3), (3, 4), (4, 1)], ["src", "dst"]) + g = GraphFrame(vertices, edges) + + result = g.neighborhood_aware_cdlp( + max_iter=2, + structural_similarity_multiplier=0.25, + ignore_direct_links=False, + initial_label_col="seed_label", + is_directed=False, + lg_nom_entries=12, + use_local_checkpoints=False, + checkpoint_interval=2, + storage_level=StorageLevel.MEMORY_AND_DISK_DESER, + ) + _df_hasCols(result, vcols=["id", "label"]) + _ = result.unpersist() + + +def test_neighborhood_aware_cdlp_api_rejects_invalid_multiplier_combination( + spark: SparkSession, +) -> None: + vertices = spark.createDataFrame([(1,), (2,)], ["id"]) + edges = spark.createDataFrame([(1, 2)], ["src", "dst"]) + g = GraphFrame(vertices, edges) + + with pytest.raises(ValueError, match="must be > 0 when ignore_direct_links is True"): + _ = g.neighborhood_aware_cdlp( + max_iter=1, + structural_similarity_multiplier=0.0, + ignore_direct_links=True, + ) + + +def test_random_walk_embeddings_api(local_g: GraphFrame) -> None: + rwe = RandomWalkEmbeddings(local_g) + rwe.set_rw_model("/tmp/") + rwe.set_hash2vec() + + result = rwe.run() + result.write.mode("overwrite").format("noop").save() + + +def test_random_walk_embeddings_invalid_args(local_g: GraphFrame) -> None: + rwe = RandomWalkEmbeddings(local_g) + + with pytest.raises(ValueError, match="supported decay functions are"): + rwe.set_hash2vec(decay_function="invalid_function") + + with pytest.raises(ValueError, match="TMP path or cached walks path should be provided!"): + rwe.run() + + +def test_strongly_connected_components(spark: SparkSession) -> None: + # Simple island test + vertices = spark.createDataFrame([(i,) for i in range(1, 6)], ["id"]) + edges = spark.createDataFrame([(7, 8)], ["src", "dst"]) + g = GraphFrame(vertices, edges) + c = g.stronglyConnectedComponents(5) + for row in c.collect(): + assert row.id == row.component, ( + f"Vertex {row.id} not equal to its component {row.component}" + ) + _ = c.unpersist() + + +@pytest.mark.parametrize("storage_level", STORAGE_LEVELS, ids=STORAGE_LEVELS_IDS) +def test_triangle_counts(spark: SparkSession, storage_level: StorageLevel) -> None: + edges = spark.createDataFrame([(0, 1), (1, 2), (2, 0)], ["src", "dst"]) + vertices = spark.createDataFrame([(0,), (1,), (2,)], ["id"]) + g = GraphFrame(vertices, edges) + c = g.triangleCount(storage_level=storage_level) + for row in c.select("id", "count").collect(): + assert row.asDict()["count"] == 1, f"Triangle count for vertex {row.id} is not 1" + _ = c.unpersist() + + +def test_approx_triangle_counts(spark: SparkSession) -> None: + edges = spark.createDataFrame([(0, 1), (1, 2), (2, 0)], ["src", "dst"]) + vertices = spark.createDataFrame([(0,), (1,), (2,)], ["id"]) + g = GraphFrame(vertices, edges) + + if spark.version[:3] >= "4.1": + c = g.triangleCount(storage_level=StorageLevel.MEMORY_AND_DISK, algorithm="approx") + for row in c.select("id", "count").collect(): + assert row.asDict()["count"] == 1, f"Triangle count for vertex {row.id} is not 1" + _ = c.unpersist() + else: + with pytest.raises(ValueError, match=".*requires Spark.*"): + c = g.triangleCount(storage_level=StorageLevel.MEMORY_AND_DISK, algorithm="approx") + + +@pytest.mark.parametrize("args", PREGEL_ARGUMENTS, ids=PREGEL_IDS) +def test_cycles_finding(spark: SparkSession, args: PregelArguments) -> None: + vertices = spark.createDataFrame( + [(1, "a"), (2, "b"), (3, "c"), (4, "d"), (5, "e")], ["id", "attr"] + ) + edges = spark.createDataFrame([(1, 2), (2, 3), (3, 1), (1, 4), (2, 5)], ["src", "dst"]) + graph = GraphFrame(vertices, edges) + res = graph.detectingCycles( + checkpoint_interval=args.checkpoint_interval, + use_local_checkpoints=args.use_local_checkpoints, + storage_level=args.storage_level, + ) + assert res.count() == 1 + collected = res.sort("id").select("found_cycles").collect() + assert collected[0][0] == [1, 2, 3, 1] + _ = res.unpersist() + + +@pytest.mark.parametrize("storage_level", STORAGE_LEVELS, ids=STORAGE_LEVELS_IDS) +def test_mis(spark: SparkSession, storage_level: StorageLevel) -> None: + # Create a graph with isolated vertices + vertices = spark.createDataFrame([(0, "a"), (1, "b"), (2, "c"), (3, "d")], ["id", "name"]) + + # Only connect vertices 0 and 1 + edges = spark.createDataFrame([(0, 1, "edge1")], ["src", "dst", "name"]) + + graph = GraphFrame(vertices, edges) + mis = graph.maximal_independent_set(storage_level=storage_level, seed=12345) + + # Check that all vertices are in the MIS (since 2 and 3 are isolated) + mis_ids = set(row[0] for row in mis.select("id").collect()) + assert len(mis_ids) == 3, "MIS should contain 2 isolated vertices and one of linked" + assert 2 in mis_ids, "Isolated vertex 2 should be in MIS" + assert 3 in mis_ids, "Isolated vertex 3 should be in MIS" + + _ = mis.unpersist() + + +@pytest.mark.skipif(is_remote(), reason="DISABLE FOR CONNECT") +def test_svd_plus_plus(examples, spark: SparkSession): + from pyspark.graphframes.classic.graphframe import _from_java_gf + + g = _from_java_gf(getattr(examples, "ALSSyntheticData")(), spark) + (v2, cost) = g.svdPlusPlus() + _df_hasCols(v2, vcols=["id", "column1", "column2", "column3", "column4"]) + + +@pytest.mark.skipif(is_remote(), reason="DISABLE FOR CONNECT") +def test_mutithreaded_sparksession_usage(spark: SparkSession): + # Test that the GraphFrame API works correctly from multiple threads. + 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"]) + + exc = None + + def run_graphframe() -> None: + nonlocal exc + try: + GraphFrame(v, e) + except Exception as _e: + exc = _e + + import threading + + thread = threading.Thread(target=run_graphframe) + thread.start() + thread.join() + assert exc is None, f"Exception was raised in thread: {exc}" + + +@pytest.mark.skipif(is_remote(), reason="DISABLE FOR CONNECT") +def test_belief_propagation(spark: SparkSession): + from pyspark.graphframes.examples import BeliefPropagation, Graphs + + # Create a graphical model g of size 3x3. + g = Graphs(spark).gridIsingModel(3) + # Run Belief Propagation (BP) for 5 iterations. + numIter = 5 + results = BeliefPropagation.runBPwithGraphFrames(g, numIter) + # Check that each belief is a valid probability in [0, 1]. + for row in results.vertices.select("belief").collect(): + belief = row["belief"] + assert 0 <= belief <= 1, f"Expected belief to be probability in [0,1], but found {belief}" + + +@pytest.mark.skipif(is_remote(), reason="DISABLE FOR CONNECT") +def test_graph_friends(spark: SparkSession): + from pyspark.graphframes.examples import Graphs + + # Construct the graph. + g = Graphs(spark).friends() + # Check that the result is an instance of GraphFrame. + assert isinstance(g, GraphFrame) + + +@pytest.mark.skipif(is_remote(), reason="DISABLE FOR CONNECT") +def test_graph_grid_ising_model(spark: SparkSession): + from pyspark.graphframes.examples import Graphs + + # Construct a grid Ising model graph. + n = 3 + g = Graphs(spark).gridIsingModel(n) + # Collect the vertex ids + ids = [v["id"] for v in g.vertices.collect()] + # Verify that every expected vertex id appears. + for i in range(n): + for j in range(n): + assert f"{i},{j}" in ids + + +@pytest.mark.parametrize("args", PREGEL_ARGUMENTS, ids=PREGEL_IDS) +def test_kcore(spark: SparkSession, args: PregelArguments) -> None: + # Create a graph designed to have clear k-core layers + v = spark.createDataFrame([(i, f"v{i}") for i in range(30)], ["id", "name"]) + + # Build edges to create a hierarchical structure: + # Core (k=5): vertices 0-4 - fully connected + core_edges = [(i, j) for i in range(5) for j in range(i + 1, 5)] + + # Next layer (k=3): vertices 5-14 - each connects to multiple core vertices + mid_layer_edges = [ + (5, 0), + (5, 1), + (5, 2), # Connect to core + (6, 0), + (6, 1), + (6, 3), + (7, 1), + (7, 2), + (7, 4), + (8, 0), + (8, 3), + (8, 4), + (9, 1), + (9, 2), + (9, 3), + (10, 0), + (10, 4), + (11, 2), + (11, 3), + (12, 1), + (12, 4), + (13, 0), + (13, 2), + (14, 3), + (14, 4), + ] + + # Outer layer (k=1): vertices 15-29 - sparse connections + outer_edges = [ + (15, 5), + (16, 6), + (17, 7), + (18, 8), + (19, 9), + (20, 10), + (21, 11), + (22, 12), + (23, 13), + (24, 14), + (25, 15), + (26, 16), + (27, 17), + (28, 18), + (29, 19), + ] + + all_edges = core_edges + mid_layer_edges + outer_edges + e = spark.createDataFrame(all_edges, ["src", "dst"]) + g = GraphFrame(v, e) + result = g.k_core( + checkpoint_interval=args.checkpoint_interval, + use_local_checkpoints=args.use_local_checkpoints, + storage_level=args.storage_level, + ) + + assert result.count() == 30 + + rows = result.collect() + kcore_map = {row["id"]: row["kcore"] for row in rows} + + # Validate hierarchical structure + # Core vertices (0-4) should have highest k-core + for i in range(5): + assert kcore_map[i] >= 4, f"Core vertex {i} should have high k-core, got {kcore_map[i]}" + + # Mid-layer vertices (5-14) should have medium k-core + for i in range(5, 15): + assert 2 <= kcore_map[i] <= 4, ( + f"Mid-layer vertex {i} should have medium k-core, got {kcore_map[i]}" + ) + + # Outer vertices (15-29) should have low k-core + for i in range(15, 30): + assert kcore_map[i] <= 2, f"Outer vertex {i} should have low k-core, got {kcore_map[i]}" + + _ = result.unpersist() + + +def test_aggregate_neighbors_basic(spark: SparkSession) -> None: + """Test basic AggregateNeighbors functionality with argument verification.""" + # Create a simple graph: 1 -> 2 -> 3 + v = spark.createDataFrame([(1, "A"), (2, "B"), (3, "C")], ["id", "name"]) + e = spark.createDataFrame([(1, 2), (2, 3)], ["src", "dst"]) + g = GraphFrame(v, e) + + # Test basic path finding from vertex 1 to vertex 3 + result = g.aggregate_neighbors( + starting_vertices=sqlfunctions.col("id") == 1, + max_hops=3, + accumulator_names=["path_length"], + accumulator_inits=[sqlfunctions.lit(0)], + accumulator_updates=[sqlfunctions.col("path_length") + 1], + target_condition=AggregateNeighbors.dst_attr("id") == 3, + required_vertex_attributes=["id"], + ) + + # Verify result structure + assert "id" in result.columns + assert "hop" in result.columns + assert "path_length" in result.columns + + # Should find one path: 1 -> 2 -> 3 + rows = result.collect() + assert len(rows) == 1 + assert rows[0]["id"] == 3 + assert rows[0]["hop"] == 2 + assert rows[0]["path_length"] == 2 + + _ = result.unpersist() + + +def test_aggregate_neighbors_with_edge_filter(spark: SparkSession) -> None: + """Test AggregateNeighbors with edge filtering.""" + # Create a graph with different edge types + v = spark.createDataFrame([(1, "A"), (2, "B"), (3, "C"), (4, "D")], ["id", "name"]) + e = spark.createDataFrame( + [(1, 2, "allowed"), (2, 3, "allowed"), (1, 3, "blocked")], + ["src", "dst", "edge_type"], + ) + g = GraphFrame(v, e) + + result = g.aggregate_neighbors( + starting_vertices=sqlfunctions.col("id") == 1, + max_hops=3, + accumulator_names=["count"], + accumulator_inits=[sqlfunctions.lit(0)], + accumulator_updates=[sqlfunctions.col("count") + 1], + target_condition=AggregateNeighbors.dst_attr("id") == 3, + edge_filter=AggregateNeighbors.edge_attr("edge_type") == "allowed", + required_edge_attributes=["edge_type"], + ) + + rows = result.collect() + # Should only find path 1 -> 2 -> 3 (not 1 -> 3 directly due to filter) + assert len(rows) == 1 + assert rows[0]["count"] == 2 # Two hops + + _ = result.unpersist() + + +def test_aggregate_neighbors_multiple_accumulators(spark: SparkSession) -> None: + """Test AggregateNeighbors with multiple accumulators.""" + v = spark.createDataFrame([(1, 10), (2, 20), (3, 30)], ["id", "value"]) + e = spark.createDataFrame([(1, 2, 5.0), (2, 3, 6.0)], ["src", "dst", "weight"]) + g = GraphFrame(v, e) + + result = g.aggregate_neighbors( + starting_vertices=sqlfunctions.col("id") == 1, + max_hops=3, + accumulator_names=["sum_values", "sum_weights"], + accumulator_inits=[sqlfunctions.lit(0), sqlfunctions.lit(0.0)], + accumulator_updates=[ + sqlfunctions.col("sum_values") + AggregateNeighbors.dst_attr("value"), + sqlfunctions.col("sum_weights") + AggregateNeighbors.edge_attr("weight"), + ], + target_condition=AggregateNeighbors.dst_attr("id") == 3, + required_vertex_attributes=["id", "value"], + required_edge_attributes=["weight"], + ) + + rows = result.collect() + assert len(rows) == 1 + # sum_values: 20 + 30 = 50 + assert rows[0]["sum_values"] == 50 + # sum_weights: 5.0 + 6.0 = 11.0 + assert abs(rows[0]["sum_weights"] - 11.0) < 0.001 + + _ = result.unpersist() + + +def test_hyper_anf_basic(spark: SparkSession) -> None: + """Smoke test: verify hyper_anf returns correct columns and basic functionality.""" + # Directed graph: 1 -> 2, 2 -> 3, 3 -> 1 (cycle) + v = spark.createDataFrame([(1,), (2,), (3,)], ["id"]) + e = spark.createDataFrame([(1, 2), (2, 3), (3, 1)], ["src", "dst"]) + g = GraphFrame(v, e) + + result = g.hyper_anf(n_hops=2, use_local_checkpoints=True) + + # Verify columns: id, hop_0, hop_1, hop_2 + assert "id" in result.columns + assert "hop_0" in result.columns + assert "hop_1" in result.columns + assert "hop_2" in result.columns + + # Every vertex with an outgoing edge should be present (all 3) + assert result.count() == 3 + + _ = result.unpersist() + + +def test_hyper_anf_args_passed(spark: SparkSession) -> None: + """Verify that non-default args (lg_nom_entries, edge_filter) are passed correctly.""" + v = spark.createDataFrame([(1,), (2,), (3,), (4,)], ["id"]) + e = spark.createDataFrame([(1, 2), (2, 3), (3, 4), (1, 3)], ["src", "dst"]) + g = GraphFrame(v, e) + + # Use non-default lg_nom_entries to verify the arg is forwarded + result_default = g.hyper_anf(n_hops=1, lg_nom_entries=10, use_local_checkpoints=True) + assert "hop_0" in result_default.columns + assert "hop_1" in result_default.columns + assert result_default.count() > 0 + _ = result_default.unpersist() + + # Use edge_filter to restrict computation to edges where src == 1 + result_filtered = g.hyper_anf( + n_hops=1, + edge_filter=sqlfunctions.col("src") == 1, + use_local_checkpoints=True, + ) + rows = result_filtered.collect() + # Only vertex 1 has outgoing edges with src == 1 + ids = {row["id"] for row in rows} + assert ids == {1} + + _ = result_filtered.unpersist() + + +def test_hyper_anf_invalid_args(spark: SparkSession) -> None: + """Verify that invalid arguments raise ValueError on the Python side.""" + v = spark.createDataFrame([(1,), (2,)], ["id"]) + e = spark.createDataFrame([(1, 2)], ["src", "dst"]) + g = GraphFrame(v, e) + + with pytest.raises(ValueError, match="n_hops must be a positive integer"): + g.hyper_anf(n_hops=0) + + with pytest.raises(ValueError, match="n_hops must be a positive integer"): + g.hyper_anf(n_hops=-1) + + with pytest.raises(ValueError, match="lg_nom_entries must be between 4 and 21"): + g.hyper_anf(lg_nom_entries=3) + + with pytest.raises(ValueError, match="lg_nom_entries must be between 4 and 21"): + g.hyper_anf(lg_nom_entries=22) + + +class GraphFramesUpstreamTests(ReusedSQLTestCase): + @classmethod + def conf(cls) -> SparkConf: + return SparkConf().set("spark.driver.memory", "4g") + + @classmethod + def setUpClass(cls) -> None: + super().setUpClass() + cls._checkpoint_dir = tempfile.TemporaryDirectory() + cls.spark.sparkContext.setCheckpointDir(cls._checkpoint_dir.name) + cls.spark.conf.set("spark.sql.shuffle.partitions", "4") + + @classmethod + def tearDownClass(cls) -> None: + super().tearDownClass() + cls._checkpoint_dir.cleanup() + + def local_graph(self) -> GraphFrame: + vertices = self.spark.createDataFrame([(1, "A"), (2, "B"), (3, "C")], ["id", "name"]) + edges = self.spark.createDataFrame( + [(1, 2, "love"), (2, 1, "hate"), (2, 3, "follow")], + ["src", "dst", "action"], + ) + return GraphFrame(vertices, edges) + + def examples(self): + from pyspark.graphframes.classic.graphframe import _java_api + + return _java_api(self.spark.sparkContext).examples() + + def test_construction(self) -> None: + test_construction(self.spark, self.local_graph()) + + def test_svd_plus_plus(self) -> None: + test_svd_plus_plus(self.examples(), self.spark) + + def test_page_rank(self) -> None: + for args in PREGEL_ARGUMENTS: + with self.subTest(args=args): + test_page_rank(self.spark, args) + + def test_pregel_early_stopping(self) -> None: + for args in PREGEL_ARGUMENTS: + with self.subTest(args=args): + test_pregel_early_stopping(self.spark, args) + + def test_connected_components(self) -> None: + for args in PREGEL_ARGUMENTS: + for cc_args in [(-1, True), (10000, True), (-1, False), (10000, False)]: + with self.subTest(args=args, cc_args=cc_args): + test_connected_components(self.spark, args, cc_args) + + def test_connected_components2(self) -> None: + for args in PREGEL_ARGUMENTS: + for cc_args in [(-1, True), (10000, True), (-1, False), (10000, False)]: + with self.subTest(args=args, cc_args=cc_args): + test_connected_components2(self.spark, args, cc_args) + + def test_shortest_paths(self) -> None: + for args in PREGEL_ARGUMENTS: + with self.subTest(args=args): + test_shortest_paths(self.spark, args) + + def test_triangle_counts(self) -> None: + for storage_level in STORAGE_LEVELS: + with self.subTest(storage_level=storage_level): + test_triangle_counts(self.spark, storage_level) + + def test_cycles_finding(self) -> None: + for args in PREGEL_ARGUMENTS: + with self.subTest(args=args): + test_cycles_finding(self.spark, args) + + def test_mis(self) -> None: + for storage_level in STORAGE_LEVELS: + with self.subTest(storage_level=storage_level): + test_mis(self.spark, storage_level) + + def test_kcore(self) -> None: + for args in PREGEL_ARGUMENTS: + with self.subTest(args=args): + test_kcore(self.spark, args) + + +def _spark_test(function): + def run(self) -> None: + function(self.spark) + + run.__name__ = function.__name__ + return run + + +def _local_graph_test(function): + def run(self) -> None: + function(self.local_graph()) + + run.__name__ = function.__name__ + return run + + +for _test_function in [ + test_validate, + test_as_undirected, + test_as_reversed, + test_power_iteration_clustering, + test_graphframes_pagerank, + test_pregel_required_edge_columns, + test_connected_components_example, + test_shortest_paths2, + test_neighborhood_aware_cdlp_api_defaults, + test_neighborhood_aware_cdlp_api_with_all_args, + test_neighborhood_aware_cdlp_api_rejects_invalid_multiplier_combination, + test_strongly_connected_components, + test_approx_triangle_counts, + test_mutithreaded_sparksession_usage, + test_belief_propagation, + test_graph_friends, + test_graph_grid_ising_model, + test_aggregate_neighbors_basic, + test_aggregate_neighbors_with_edge_filter, + test_aggregate_neighbors_multiple_accumulators, + test_hyper_anf_basic, + test_hyper_anf_args_passed, + test_hyper_anf_invalid_args, +]: + setattr(GraphFramesUpstreamTests, _test_function.__name__, _spark_test(_test_function)) + + +for _test_function in [ + test_cache, + test_degrees, + test_type_degrees, + test_type_degrees_with_explicit_types, + test_motif_finding, + test_filterVertices, + test_filterEdges, + test_dropIsolatedVertices, + test_bfs, + test_all_paths, + test_random_walk_embeddings_api, + test_random_walk_embeddings_invalid_args, +]: + setattr(GraphFramesUpstreamTests, _test_function.__name__, _local_graph_test(_test_function)) + + +if __name__ == "__main__": + from pyspark.testing.unittestutils import main + + main() From 30d4fc5a4324d3fb43508e3a56eaecbef4c878ef Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Fri, 28 Aug 2026 08:16:31 +0000 Subject: [PATCH 7/8] [GRAPHFRAMES][FOLLOWUP] Fix lint and API documentation --- .../apache/spark/graphframes/GraphFrame.scala | 110 +++++++++--------- .../convolutions/SamplingConvolution.scala | 3 +- .../graphframes/embeddings/Hash2Vec.scala | 46 ++++---- .../embeddings/RandomWalkEmbeddings.scala | 22 ++-- .../examples/BeliefPropagation.scala | 13 ++- .../spark/graphframes/examples/Graphs.scala | 3 +- .../graphframes/lib/AggregateMessages.scala | 22 ++-- .../graphframes/lib/AggregateNeighbors.scala | 15 +-- .../spark/graphframes/lib/AllPaths.scala | 10 +- .../apache/spark/graphframes/lib/BFS.scala | 10 +- .../graphframes/lib/ConnectedComponents.scala | 23 ++-- .../graphframes/lib/DetectingCycles.scala | 10 +- .../graphframes/lib/GraphXConversions.scala | 10 +- .../spark/graphframes/lib/HyperANF.scala | 17 +-- .../apache/spark/graphframes/lib/KCore.scala | 15 +-- .../graphframes/lib/LabelPropagation.scala | 19 +-- .../lib/MaximalIndependentSet.scala | 12 +- .../spark/graphframes/lib/PageRank.scala | 2 +- .../lib/ParallelPersonalizedPageRank.scala | 4 +- .../apache/spark/graphframes/lib/Pregel.scala | 76 +++++------- .../lib/RandomizedContraction.scala | 28 +++-- .../spark/graphframes/lib/SVDPlusPlus.scala | 16 +-- .../spark/graphframes/lib/ShortestPaths.scala | 25 ++-- .../lib/StronglyConnectedComponents.scala | 4 +- .../lib/StructureAwareLabelPropagation.scala | 11 +- .../spark/graphframes/lib/TriangleCount.scala | 6 +- .../spark/graphframes/lib/TwoPhase.scala | 25 ++-- .../spark/graphframes/pattern/patterns.scala | 10 +- .../spark/graphframes/rw/RandomWalkBase.scala | 14 +-- .../rw/RandomWalkWithRestart.scala | 2 +- .../sql/graphframes/GraphFrameInternals.scala | 9 +- .../sql/graphframes/GraphFramesConf.scala | 65 ++++++----- .../graphframes/expressions/KCoreMerge.scala | 4 +- .../expressions/KMinSampling.scala | 11 +- .../GraphFrameInternalsSuite.scala | 2 +- .../spark/graphframes/GraphFrameSuite.scala | 13 ++- .../GraphFrameTestSparkContext.scala | 17 +-- .../spark/graphframes/PatternMatchSuite.scala | 2 + .../spark/graphframes/SparkFunSuite.scala | 24 +--- .../apache/spark/graphframes/TestUtils.scala | 2 +- .../SamplingConvolutionSuite.scala | 9 +- .../embeddings/Hash2VecSuite.scala | 73 ++++++------ .../lib/AggregateMessagesSuite.scala | 12 +- .../lib/AggregateNeighborsSuite.scala | 2 +- .../spark/graphframes/lib/AllPathsSuite.scala | 6 +- .../spark/graphframes/lib/BFSSuite.scala | 8 +- .../lib/ConnectedComponentsSuite.scala | 12 +- .../lib/DetectingCyclesSuite.scala | 6 +- .../spark/graphframes/lib/HyperANFSuite.scala | 12 +- .../spark/graphframes/lib/KCoreSuite.scala | 6 +- .../lib/LabelPropagationSuite.scala | 2 +- .../lib/MaximalIndependentSetSuite.scala | 4 +- .../spark/graphframes/lib/PageRankSuite.scala | 4 +- .../ParallelPersonalizedPageRankSuite.scala | 17 +-- .../spark/graphframes/lib/PregelSuite.scala | 5 +- .../lib/RandomizedContractionSuite.scala | 12 +- .../graphframes/lib/SVDPlusPlusSuite.scala | 8 +- .../graphframes/lib/ShortestPathsSuite.scala | 4 +- .../StronglyConnectedComponentsSuite.scala | 4 +- .../lib/StructureAwareLabelPropagation.scala | 10 +- .../graphframes/lib/TriangleCountSuite.scala | 6 +- .../graphframes/pattern/PatternSuite.scala | 3 +- .../rw/RandomWalkWithRestartSuite.scala | 22 ++-- .../connect/planner/SparkConnectPlanner.scala | 2 +- .../graphframes/GraphFramesConnectUtils.scala | 109 +++++++++-------- 65 files changed, 537 insertions(+), 553 deletions(-) diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFrame.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFrame.scala index cc7f6ac32865c..4f739e12756a0 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFrame.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/GraphFrame.scala @@ -17,6 +17,13 @@ package org.apache.spark.graphframes +import java.util.Random + +import scala.reflect.runtime.universe.TypeTag + +import org.apache.spark.graphframes.embeddings.RandomWalkEmbeddings +import org.apache.spark.graphframes.lib._ +import org.apache.spark.graphframes.pattern._ import org.apache.spark.graphx.Edge import org.apache.spark.graphx.Graph import org.apache.spark.ml.clustering.PowerIterationClustering @@ -33,12 +40,6 @@ import org.apache.spark.sql.functions.monotonically_increasing_id import org.apache.spark.sql.functions.struct import org.apache.spark.sql.types._ import org.apache.spark.storage.StorageLevel -import org.apache.spark.graphframes.embeddings.RandomWalkEmbeddings -import org.apache.spark.graphframes.lib._ -import org.apache.spark.graphframes.pattern._ - -import java.util.Random -import scala.reflect.runtime.universe.TypeTag /** * A representation of a graph using `DataFrame`s. @@ -148,12 +149,9 @@ class GraphFrame private ( * Validates the consistency and integrity of a graph by performing checks on the vertices and * edges. * - * @return - * Unit, as the method, performs validation checks and throws an exception if validation - * fails. - * @throws InvalidGraphException - * if there are any inconsistencies in the graph, such as duplicate vertices, mismatched - * vertices between edges and vertex DataFrames or missing connections. + * This method returns `Unit`. It throws `InvalidGraphException` if there are any inconsistencies + * in the graph, such as duplicate vertices, mismatched vertices between edges and vertex + * DataFrames, or missing connections. * * @group utils */ @@ -171,12 +169,9 @@ class GraphFrame private ( * @param intermediateStorageLevel * the storage level to be used when persisting intermediate DataFrame computations during the * validation process. - * @return - * Unit, as the method, performs validation checks and throws an exception if validation - * fails. - * @throws InvalidGraphException - * if there are any inconsistencies in the graph, such as duplicate vertices, mismatched - * vertices between edges and vertex DataFrames or missing connections. + * This method returns `Unit`. It throws `InvalidGraphException` if there are any inconsistencies + * in the graph, such as duplicate vertices, mismatched vertices between edges and vertex + * DataFrames, or missing connections. * * @group utils */ @@ -261,7 +256,7 @@ class GraphFrame private ( /** * The dataframe representation of the vertices of the graph. * - * It contains a column called [[GraphFrame.ID]] with the id of the vertex, and various other + * It contains a column called `id` with the id of the vertex, and various other * user-defined attributes with other attributes. * * The order of the columns is available in [[vertexColumns]]. @@ -278,7 +273,7 @@ class GraphFrame private ( /** * The dataframe representation of the edges of the graph. * - * It contains two columns called [[GraphFrame.SRC]] and [[GraphFrame.DST]] that contain the ids + * It contains two columns called `src` and `dst` that contain the ids * of the source vertex and the destination vertex of each edge, respectively. It may also * contain various other columns with user-defined attributes for each edge. * @@ -299,9 +294,9 @@ class GraphFrame private ( /** * Returns triplets: (source vertex)-[edge]->(destination vertex) for all edges in the graph. - * The DataFrame returned has 3 columns, with names: [[GraphFrame.SRC]], [[GraphFrame.EDGE]], - * and [[GraphFrame.DST]]. Each column is a struct. The 2 vertex columns have schema matching - * [[GraphFrame.vertices]], and the edge column has a schema matching [[GraphFrame.edges]]. For + * The DataFrame returned has 3 columns, with names: `src`, `edge`, and `dst`. Each column is a + * struct. The 2 vertex columns have schema matching the vertices DataFrame, and the edge column + * has a schema matching the edges DataFrame. For * example, `triplets.select(col(SRC)(ID))` selects ID of the source column. * * @group structure @@ -409,7 +404,7 @@ class GraphFrame private ( /** * The out-degree of each vertex in the graph, returned as a DataFrame with two columns: - * - [[GraphFrame.ID]] the ID of the vertex + * - `id` the ID of the vertex * - "outDegree" (integer) storing the out-degree of the vertex Note that vertices with 0 * out-edges are not returned in the result. * @@ -421,7 +416,7 @@ class GraphFrame private ( /** * The in-degree of each vertex in the graph, returned as a DataFame with two columns: - * - [[GraphFrame.ID]] the ID of the vertex "- "inDegree" (int) storing the in-degree of the + * - `id` the ID of the vertex "- "inDegree" (int) storing the in-degree of the * vertex Note that vertices with 0 in-edges are not returned in the result. * * @group degree @@ -432,7 +427,7 @@ class GraphFrame private ( /** * The degree of each vertex in the graph, returned as a DataFrame with two columns: - * - [[GraphFrame.ID]] the ID of the vertex + * - `id` the ID of the vertex * - 'degree' (integer) the degree of the vertex Note that vertices with 0 edges are not * returned in the result. * @@ -447,7 +442,7 @@ class GraphFrame private ( /** * The out-degree of each vertex per edge type, returned as a DataFrame with two columns: - * - [[GraphFrame.ID]] the ID of the vertex + * - `id` the ID of the vertex * - "outDegrees" a struct with a field for each edge type, storing the out-degree count * * @param edgeTypeCol @@ -476,7 +471,7 @@ class GraphFrame private ( /** * The in-degree of each vertex per edge type, returned as a DataFrame with two columns: - * - [[GraphFrame.ID]] the ID of the vertex + * - `id` the ID of the vertex * - "inDegrees" a struct with a field for each edge type, storing the in-degree count * * @param edgeTypeCol @@ -506,7 +501,7 @@ class GraphFrame private ( /** * The total degree of each vertex per edge type (both in and out), returned as a DataFrame with * two columns: - * - [[GraphFrame.ID]] the ID of the vertex + * - `id` the ID of the vertex * - "degrees" a struct with a field for each edge type, storing the total degree count * * @param edgeTypeCol @@ -562,9 +557,9 @@ class GraphFrame private ( * - The names are used as column names in the result `DataFrame`. If a motif contains named * vertex `a`, then the result `DataFrame` will contain a column "a" which is a * `StructType` with sub-fields equivalent to the schema (columns) of - * [[GraphFrame.vertices]]. Similarly, an edge `e` in a motif will produce a column "e" in + * the vertices DataFrame. Similarly, an edge `e` in a motif will produce a column "e" in * the result `DataFrame` with sub-fields equivalent to the schema (columns) of - * [[GraphFrame.edges]]. + * the edges DataFrame. * - Be aware that names do *not* identify *distinct* elements: two elements with different * names may refer to the same graph element. For example, in the motif `"(a)-[e]->(b); * (b)-[e2]->(c)"`, the names `a` and `c` could refer to the same vertex. To restrict @@ -606,7 +601,7 @@ class GraphFrame private ( * {{{ * spark.conf.set("spark.sql.cbo.joinReorder.dp.threshold", "20") * }}} - * CBO relies on table statistics, so run `ANALYZE TABLE COMPUTE STATISTICS` on the + * CBO relies on table statistics, so run `ANALYZE TABLE table_name COMPUTE STATISTICS` on the * vertices and edges tables to ensure accurate statistics are available. * * @param pattern @@ -754,8 +749,7 @@ class GraphFrame private ( * graphs or large maxHops values. Consider using appropriate storage levels and checkpoint * intervals for stability. * - * @see - * [[org.apache.spark.graphframes.lib.AggregateNeighbors]] for implementation details + * See `AggregateNeighbors` for implementation details. * @return * an [[org.apache.spark.graphframes.lib.AggregateNeighbors]] instance for configuration * @group stdlib @@ -871,11 +865,10 @@ class GraphFrame private ( /** * Pregel algorithm. * - * @see - * [[org.apache.spark.graphframes.lib.Pregel]] + * See `Pregel` for more details. * @group stdlib */ - def pregel = new Pregel(this) + def pregel: Pregel = new Pregel(this) /** * Shortest paths algorithm. @@ -977,12 +970,12 @@ class GraphFrame private ( def kCore: KCore = new KCore(this) /** - * Find all cycles in the graph. An implementation of the Rocha–Thatte cycle detection + * Find all cycles in the graph. An implementation of the Rocha-Thatte cycle detection * algorithm. * * Rocha, Rodrigo Caetano, and Bhalchandra D. Thatte. "Distributed cycle detection in - * large-scale sparse graphs." Proceedings of Simpósio Brasileiro de Pesquisa Operacional - * (SBPO’15) (2015): 1-11. + * large-scale sparse graphs." Proceedings of Simposio Brasileiro de Pesquisa Operacional + * (SBPO '15) (2015): 1-11. * * Returns a DataFrame with unique cycles. * @@ -1173,7 +1166,7 @@ object GraphFrame extends Serializable with Logging { } /** - * Column name for vertex IDs in [[GraphFrame.vertices]] Note that GraphFrame assigns a unique + * Column name for vertex IDs in the vertices DataFrame. Note that GraphFrame assigns a unique * long ID to each vertex, If the vertex ID type is one of byte / int / long / short type, * GraphFrame casts the original IDs to long as the unique long ID, otherwise GraphFrame * generates the unique long ID by Spark function ``monotonically_increasing_id`` which is less @@ -1183,23 +1176,23 @@ object GraphFrame extends Serializable with Logging { /** * Column name for source vertices of edges. - * - In [[GraphFrame.edges]], this is a column of vertex IDs. - * - In [[GraphFrame.triplets]], this is a column of vertices with schema matching - * [[GraphFrame.vertices]]. + * - In the edges DataFrame, this is a column of vertex IDs. + * - In the triplets DataFrame, this is a column of vertices with schema matching the vertices + * DataFrame. */ val SRC: String = "src" /** * Column name for destination vertices of edges. - * - In [[GraphFrame.edges]], this is a column of vertex IDs. - * - In [[GraphFrame.triplets]], this is a column of vertices with schema matching - * [[GraphFrame.vertices]]. + * - In the edges DataFrame, this is a column of vertex IDs. + * - In the triplets DataFrame, this is a column of vertices with schema matching the vertices + * DataFrame. */ val DST: String = "dst" /** - * Column name for edge in [[GraphFrame.triplets]]. In [[GraphFrame.triplets]], this is a column - * of edges with schema matching [[GraphFrame.edges]]. + * Column name for edge in the triplets DataFrame. In the triplets DataFrame, this is a column of + * edges with schema matching the edges DataFrame. */ val EDGE: String = "edge" @@ -1244,10 +1237,10 @@ object GraphFrame extends Serializable with Logging { } /** - * Create a new [[GraphFrame]] from an edge `DataFrame`. The resulting [[GraphFrame]] will have - * [[GraphFrame.vertices]] with a single "id" column. + * Create a new [[GraphFrame]] from an edge `DataFrame`. The resulting [[GraphFrame]] will have a + * vertices DataFrame with a single "id" column. * - * Note: The [[GraphFrame.vertices]] DataFrame will be persisted at level + * Note: The vertices DataFrame will be persisted at level * `StorageLevel.MEMORY_AND_DISK`. * @param e * Edge DataFrame. This must include columns "src" and "dst" containing source and destination @@ -1262,10 +1255,10 @@ object GraphFrame extends Serializable with Logging { } /** - * Create a new [[GraphFrame]] from an edge `DataFrame`. The resulting [[GraphFrame]] will have - * [[GraphFrame.vertices]] with a single "id" column. + * Create a new [[GraphFrame]] from an edge `DataFrame`. The resulting [[GraphFrame]] will have a + * vertices DataFrame with a single "id" column. * - * Note: The [[GraphFrame.vertices]] DataFrame will be persisted at level + * Note: The vertices DataFrame will be persisted at level * `StorageLevel.MEMORY_AND_DISK`. * @param e * Edge DataFrame. This must include columns "src" and "dst" containing source and destination @@ -1279,7 +1272,8 @@ object GraphFrame extends Serializable with Logging { */ def fromEdges(e: DataFrame, storageLevel: StorageLevel): GraphFrame = { logWarn( - s"this method persists graph vertices with storage level ${storageLevel.toString()}, users should manually unpersist it when the graph is not needed!") + s"this method persists graph vertices with storage level ${storageLevel.toString()}, " + + "users should manually unpersist it when the graph is not needed!") val srcs = e.select(e("src").as("id")) val dsts = e.select(e("dst").as("id")) val v = srcs.unionAll(dsts).distinct().persist(storageLevel) @@ -1307,7 +1301,7 @@ object GraphFrame extends Serializable with Logging { /** * Given: * - a GraphFrame `originalGraph` - * - a GraphX graph derived from the GraphFrame using [[GraphFrame.toGraphX]] this method + * - a GraphX graph derived from the GraphFrame using `toGraphX`; this method * merges attributes from the GraphX graph into the original GraphFrame. * * This method is useful for doing computations using the GraphX API and then merging the @@ -1317,7 +1311,7 @@ object GraphFrame extends Serializable with Logging { * "category" and an Int edge attribute we want to call "count" We can call * `fromGraphX(originalGraph, graph, Seq("category"), Seq("count"))` to produce a new * GraphFrame. The new GraphFrame will be an augmented version of `originalGraph`, with new - * [[GraphFrame.vertices]] column "category" and new [[GraphFrame.edges]] column "count" + * vertices column "category" and new edges column "count" * added. * * See [[org.apache.spark.graphframes.examples.BeliefPropagation]] for example usage. diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/convolutions/SamplingConvolution.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/convolutions/SamplingConvolution.scala index a4e4b1a90b79a..86902b3623ef8 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/convolutions/SamplingConvolution.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/convolutions/SamplingConvolution.scala @@ -17,13 +17,12 @@ package org.apache.spark.graphframes.convolutions +import org.apache.spark.graphframes.{GraphFrame, Logging} import org.apache.spark.ml.functions._ import org.apache.spark.ml.stat.Summarizer import org.apache.spark.sql.DataFrame import org.apache.spark.sql.functions._ import org.apache.spark.sql.graphframes.expressions.KMinSampling -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.Logging /** * A convolution operation on graph data that aggregates features from sampled neighbors using diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/Hash2Vec.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/Hash2Vec.scala index e443bec75fd0d..24b6a6545219e 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/Hash2Vec.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/Hash2Vec.scala @@ -17,7 +17,13 @@ package org.apache.spark.graphframes.embeddings +import scala.reflect.ClassTag +import scala.util.hashing.MurmurHash3 + import dev.ludovic.netlib.blas.BLAS + +import org.apache.spark.graphframes.GraphFramesUnsupportedVertexTypeException +import org.apache.spark.graphframes.rw.RandomWalkBase import org.apache.spark.ml.linalg import org.apache.spark.ml.linalg.SQLDataTypes.VectorType import org.apache.spark.ml.linalg.Vectors @@ -41,15 +47,10 @@ import org.apache.spark.sql.types.StringType import org.apache.spark.sql.types.StructField import org.apache.spark.sql.types.StructType import org.apache.spark.unsafe.hash.Murmur3_x86_32._ -import org.apache.spark.graphframes.GraphFramesUnsupportedVertexTypeException -import org.apache.spark.graphframes.rw.RandomWalkBase - -import scala.reflect.ClassTag -import scala.util.hashing.MurmurHash3 /** * Implementation of Hash2Vec, an efficient word embedding technique using feature hashing. Based - * on: Argerich, Luis, Joaquín Torré Zaffaroni, and Matías J. Cano. "Hash2vec, feature hashing for + * on: Argerich, Luis, Joaquin Torre Zaffaroni, and Matias J. Cano. "Hash2vec, feature hashing for * word embeddings." arXiv preprint arXiv:1608.08940 (2016). * * Produces embeddings for elements in sequences using a hash-based approach to avoid storing a @@ -81,11 +82,11 @@ class Hash2Vec extends Serializable { private var maxVectorsPerPartition: Int = 100000 /** - * Sets whether final vectors are L2‑normalized after aggregation across partitions. When + * Sets whether final vectors are L2-normalized after aggregation across partitions. When * normalization is enabled, each vector is scaled to unit length (L2 norm = 1). * * When `safeNorm` is true (default), the method adds an extra channel to the vector equal to - * `log(L2‑norm + 1) / sqrt(dim)`. This preserves some information about the original magnitude + * `log(L2-norm + 1) / sqrt(dim)`. This preserves some information about the original magnitude * while still making vectors comparable via cosine similarity. * * When `safeNorm` is false, normalizes without adding an extra channel, discarding magnitude @@ -97,7 +98,7 @@ class Hash2Vec extends Serializable { * If true, output vectors are normalized. * @param safeNorm * If true (and doNorm is true), retains magnitude information in an extra dimension. If - * false, performs standard L2‑normalization. + * false, performs standard L2-normalization. * @return * This Hash2Vec instance for method chaining. */ @@ -131,11 +132,11 @@ class Hash2Vec extends Serializable { } /** - * Convenience overload for `setDoNormalization(doNorm, safeNorm)` that uses safe‑mode (extra + * Convenience overload for `setDoNormalization(doNorm, safeNorm)` that uses safe-mode (extra * channel) by default. Equivalent to `setDoNormalization(value, true)`. * * @param value - * If true, output vectors are L2‑normalized with safe (extra‑channel) semantics. + * If true, output vectors are L2-normalized with safe (extra-channel) semantics. * @return * This Hash2Vec instance for method chaining. */ @@ -295,7 +296,8 @@ class Hash2Vec extends Serializable { StructType(Seq(StructField("id", LongType), StructField("vector", VectorType)))) case _ => throw new GraphFramesUnsupportedVertexTypeException( - s"Hash2vec supports only string or numeric types of elements but got ${elDataType.toString()}") + "Hash2vec supports only string or numeric types of elements but got " + + elDataType.toString) } val embeddings = spark @@ -510,22 +512,22 @@ object Hash2Vec { /** * A paged matrix of double-precision vectors that stores vectors contiguously in large - * fixed‑sized pages, each holding PAGE_SIZE (4096) vectors of dimension `dim`. + * fixed-sized pages, each holding PAGE_SIZE (4096) vectors of dimension `dim`. * * This layout replaces a HashMap[T, Array[Double]] with two separate structures: * 1. A mapping from element identifier (T) to a vector ID (Int), maintained by the caller. * 2. The actual vector data stored in a few large arrays (pages) instead of many small - * per‑element arrays. + * per-element arrays. * * Advantages over a HashMap-of-arrays: - * 1. Eliminates per‑vector Array object overhead (object + * 1. Eliminates per-vector Array object overhead (object * header, reference, GC metadata). - * 2. Reduces GC pressure because the backing store is a small number of large long‑lived - * arrays, not many short‑lived small arrays that become garbage as the map is updated. + * 2. Reduces GC pressure because the backing store is a small number of large long-lived + * arrays, not many short-lived small arrays that become garbage as the map is updated. * 3. Better memory locality: vectors of the same dimension are stored consecutively, * improving cache line utilisation during sequential access (e.g., inside a page). * 4. Predictable memory growth: pages are allocated only when the current page is full, - * avoiding repeated resizing of a hash‑map and associated re‑hashing / copying. + * avoiding repeated resizing of a hash-map and associated re-hashing / copying. * * The cost is an extra indirection to compute the page index and offset, which is cheap (bit * shifts and masks) compared to the GC and memory overhead it saves. @@ -534,8 +536,8 @@ object Hash2Vec { * 1. PAGE_BITS = 12, PAGE_SIZE = 4096 (2^12). This keeps pageIdx = * vectorId >>> PAGE_BITS and localRow = vectorId & PAGE_MASK cheap, while limiting page memory * to PAGE_SIZE * dim doubles. - * 2. The first page is pre‑allocated in the constructor; subsequent - * pages are added on‑demand when allocateVector() crosses a page boundary. + * 2. The first page is pre-allocated in the constructor; subsequent + * pages are added on-demand when allocateVector() crosses a page boundary. * 3. allocateVector() * returns a monotonically increasing integer ID, which is the index of the vector across all * pages. The caller stores this ID in a HashMap[T, Int] instead of storing the whole array. @@ -561,7 +563,7 @@ object Hash2Vec { pages += new Array[Double](size.toInt) } - /** Allocate a new zero‑initialized vector and return its unique integer ID. */ + /** Allocate a new zero-initialized vector and return its unique integer ID. */ def allocateVector(): Int = { val id = vectorCount val localIdx = id & PAGE_MASK // ~id % 4096 @@ -574,7 +576,7 @@ object Hash2Vec { id } - /** Accumulate `value` into the component `offset` (0‑based) of vector `vectorId`. */ + /** Accumulate `value` into the component `offset` (0-based) of vector `vectorId`. */ @inline def add(vectorId: Int, offset: Int, value: Double): Unit = { // vectorId / PAGE_SIZE using unsigned shift (page index) diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/RandomWalkEmbeddings.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/RandomWalkEmbeddings.scala index 6772b8f7cccd8..a32fb312dca0e 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/RandomWalkEmbeddings.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/embeddings/RandomWalkEmbeddings.scala @@ -17,20 +17,17 @@ package org.apache.spark.graphframes.embeddings +import org.apache.spark.graphframes.{GraphFrame, GraphFramesW2VException, Logging} +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.convolutions.SamplingConvolution +import org.apache.spark.graphframes.embeddings.RandomWalkEmbeddings.rwModels +import org.apache.spark.graphframes.rw.{RandomWalkBase, RandomWalkWithRestart} import org.apache.spark.ml.feature.Word2Vec import org.apache.spark.sql.Column import org.apache.spark.sql.DataFrame import org.apache.spark.sql.functions.col import org.apache.spark.sql.functions.transform import org.apache.spark.sql.types.StringType -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFramesW2VException -import org.apache.spark.graphframes.Logging -import org.apache.spark.graphframes.WithIntermediateStorageLevel -import org.apache.spark.graphframes.convolutions.SamplingConvolution -import org.apache.spark.graphframes.embeddings.RandomWalkEmbeddings.rwModels -import org.apache.spark.graphframes.rw.RandomWalkBase -import org.apache.spark.graphframes.rw.RandomWalkWithRestart /** * RandomWalkEmbeddings is a class for generating node embeddings in a graph using random walks @@ -206,7 +203,7 @@ class RandomWalkEmbeddings private[graphframes] (private val graph: GraphFrame) } val embeddings = sequenceModel match { - case Left(w2v) => { + case Left(w2v) => val model = w2v.setInputCol(RandomWalkBase.rwColName) val preProcessedSequences = if (graph.vertices.schema(GraphFrame.ID).dataType != StringType) { @@ -220,10 +217,8 @@ class RandomWalkEmbeddings private[graphframes] (private val graph: GraphFrame) val fittedW2V = model.fit(preProcessedSequences) fittedW2V.getVectors.withColumnsRenamed( Map("word" -> GraphFrame.ID, "vector" -> RandomWalkEmbeddings.embeddingColName)) - } - case Right(h2v) => { + case Right(h2v) => h2v.run(walks).withColumnRenamed("vector", RandomWalkEmbeddings.embeddingColName) - } } val persistedEmbeddings = if (aggregateNeighbors) { @@ -288,6 +283,8 @@ object RandomWalkEmbeddings extends Serializable { * * Instead of this API it is recommended to use new + setters of the class! */ + // This compatibility entry point mirrors the flat Python API. + // scalastyle:off argcount def pythonAPI( graph: GraphFrame, useEdgeDirection: Boolean, @@ -381,4 +378,5 @@ object RandomWalkEmbeddings extends Serializable { embeddingsGenerator.useCachedRandomWalks(rwCachedWalks).run() } } + // scalastyle:on argcount } diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/examples/BeliefPropagation.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/examples/BeliefPropagation.scala index 01b0e8154a1f2..31c08051145ea 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/examples/BeliefPropagation.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/examples/BeliefPropagation.scala @@ -48,15 +48,18 @@ import org.apache.spark.sql.functions.when * P(X) = (1/Z) * exp[ \sum_i a_i x_i + \sum_{ij} b_{ij} x_i x_j ] * }}} * where Z is the normalization constant (partition function). See - * [[https://en.wikipedia.org/wiki/Ising_model Wikipedia]] for more information on Ising models. + * See Wikipedia for more information on + * Ising models. * * Belief Propagation (BP) provides marginal probabilities of the values of the variables x,,i,,, * i.e., P(x,,i,,) for each i. This allows a user to understand likely values of variables. See - * [[https://en.wikipedia.org/wiki/Belief_propagation Wikipedia]] for more information on BP. + * See Wikipedia for more + * information on BP. * * We use a batch synchronous BP algorithm, where batches of vertices are updated synchronously. * We follow the mean field update algorithm in Slide 13 of the - * [[http://www.eecs.berkeley.edu/~wainwrig/Talks/A_GraphModel_Tutorial talk slides]] from: + * See the graphical + * model tutorial slides from: * Wainwright. "Graphical models, message-passing algorithms, and convex optimization." * * The batches are chosen according to a coloring. For background on graph colorings for @@ -131,7 +134,7 @@ object BeliefPropagation { * Number of iterations of BP to run. One iteration includes updating each vertex's belief * once. * @return - * Same graphical model, but with [[GraphFrame.vertices]] augmented with a new column "belief" + * Same graphical model, but with its vertices DataFrame augmented with a new column "belief" * containing P(x,,i,, = +1), the marginal probability of vertex i taking value +1 instead of * -1. */ @@ -211,7 +214,7 @@ object BeliefPropagation { * Number of iterations of BP to run. One iteration includes updating each vertex's belief * once. * @return - * Same graphical model, but with [[GraphFrame.vertices]] augmented with a new column "belief" + * Same graphical model, but with its vertices DataFrame augmented with a new column "belief" * containing P(x,,i,, = +1), the marginal probability of vertex i taking value +1 instead of * -1. */ diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/examples/Graphs.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/examples/Graphs.scala index df4e504f9f7ce..bf07102346f65 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/examples/Graphs.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/examples/Graphs.scala @@ -174,7 +174,8 @@ class Graphs private[graphframes] () { * P(X) = (1/Z) * exp[ \sum_i a_i x_i + \sum_{ij} b_{ij} x_i x_j ] * }}} * where Z is the normalization constant (partition function). See - * [[https://en.wikipedia.org/wiki/Ising_model Wikipedia]] for more information on Ising models. + * See Wikipedia for more information on + * Ising models. * * Each vertex is parameterized by a single scalar a,,i,,. Each edge is parameterized by a * single scalar b,,ij,,. diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateMessages.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateMessages.scala index fb2be56278fa9..2c09f2f569e83 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateMessages.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateMessages.scala @@ -17,35 +17,35 @@ package org.apache.spark.graphframes.lib +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithIntermediateStorageLevel import org.apache.spark.sql.Column import org.apache.spark.sql.DataFrame import org.apache.spark.sql.functions.col import org.apache.spark.sql.functions.expr import org.apache.spark.sql.functions.struct -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.Logging -import org.apache.spark.graphframes.WithIntermediateStorageLevel /** * This is a primitive for implementing graph algorithms. This method aggregates messages from the * neighboring edges and vertices of each vertex. * - * For each triplet (source vertex, edge, destination vertex) in [[GraphFrame.triplets]], this can + * For each triplet (source vertex, edge, destination vertex) in the triplets DataFrame, this can * send a message to the source and/or destination vertices. * - `AggregateMessages.sendToSrc()` sends a message to the source vertex of each triplet * - `AggregateMessages.sendToDst()` sends a message to the destination vertex of each triplet * - `AggregateMessages.agg` specifies an aggregation function for aggregating the messages sent * to each vertex. It also runs the aggregation, computing a DataFrame with one row for each * vertex which receives > 0 messages. The DataFrame has 2 columns: - * - vertex column ID (named [[GraphFrame.ID]]) + * - vertex column ID (named `id`) * - aggregate from messages sent to vertex (with the name given to the `Column` specified in * `AggregateMessages.agg()`) * * When specifying the messages and aggregation function, the user may reference columns using: - * - [[AggregateMessages.src]]: column for source vertex of edge - * - [[AggregateMessages.edge]]: column for edge - * - [[AggregateMessages.dst]]: column for destination vertex of edge - * - [[AggregateMessages.msg]]: message sent to vertex (for aggregation function) + * - `AggregateMessages.src`: column for source vertex of edge + * - `AggregateMessages.edge`: column for edge + * - `AggregateMessages.dst`: column for destination vertex of edge + * - `AggregateMessages.msg`: message sent to vertex (for aggregation function) * * Note: If you use this operation to write an iterative algorithm, you may want to use * `checkpoint()` (`localCheckpoint()`) as a workaround for caching issues. @@ -120,7 +120,7 @@ class AggregateMessages private[graphframes] (private val g: GraphFrame) * - column "id": vertex ID * - aggCol: aggregate result * - aggCols: one column with the result of each additional defined aggregation - * If you need to join this with the original [[GraphFrame.vertices]], you can run an inner join + * If you need to join this with the original vertices DataFrame, you can run an inner join * of the form: * {{{ * val g: GraphFrame = ... @@ -190,7 +190,7 @@ class AggregateMessages private[graphframes] (private val g: GraphFrame) object AggregateMessages extends Logging with Serializable { - /** Column name for aggregated messages, used in [[AggregateMessages.msg]] */ + /** Column name for aggregated messages, used in `AggregateMessages.msg` */ val MSG_COL_NAME: String = "MSG" /** Reference for source column, used for specifying messages */ diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateNeighbors.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateNeighbors.scala index ddf21ceb541f0..e9c3f664c019c 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateNeighbors.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AggregateNeighbors.scala @@ -17,14 +17,14 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.Column -import org.apache.spark.sql.DataFrame -import org.apache.spark.sql.functions._ import org.apache.spark.graphframes.GraphFrame import org.apache.spark.graphframes.Logging import org.apache.spark.graphframes.WithCheckpointInterval import org.apache.spark.graphframes.WithIntermediateStorageLevel import org.apache.spark.graphframes.WithLocalCheckpoints +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ /** * A class for performing multi-hop neighbor aggregation on a graph. @@ -33,7 +33,7 @@ import org.apache.spark.graphframes.WithLocalCheckpoints * accumulating values along paths using customizable accumulator expressions. It supports both * stopping conditions (when to stop exploring) and target conditions (when to collect a result). * If no target condition is provided, collect all traversals that reached stopping condition. The - * algorithm processes the graph in a breadth‑first manner. + * algorithm processes the graph in a breadth-first manner. * * Use the builder pattern to configure parameters, then call `run()` to execute. * @@ -238,10 +238,10 @@ class AggregateNeighbors private[graphframes] (graph: GraphFrame) } /** - * Controls whether self‑loops (edges where src == dst) are excluded. + * Controls whether self-loops (edges where src == dst) are excluded. * * @param value - * if true, self‑loop edges are filtered out + * if true, self-loop edges are filtered out * @return * this AggregateNeighbors instance for method chaining */ @@ -264,8 +264,9 @@ class AggregateNeighbors private[graphframes] (graph: GraphFrame) */ def run(): DataFrame = { require(maxHops > 0, "maxHops must be greater than 0") - if (maxHops > 10) + if (maxHops > 10) { logWarn(s"maxHops is very large ($maxHops). This might be performance-intensive.") + } require(accumulatorsNames.nonEmpty, "At least one accumulator must be added") require( stoppingCondition.orElse(targetCondition).isDefined, diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AllPaths.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AllPaths.scala index 784438440201d..eade74ac0196d 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AllPaths.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/AllPaths.scala @@ -17,6 +17,11 @@ package org.apache.spark.graphframes.lib +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.WithCheckpointInterval +import org.apache.spark.graphframes.WithDirection +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLocalCheckpoints import org.apache.spark.sql.Column import org.apache.spark.sql.DataFrame import org.apache.spark.sql.functions.array @@ -25,11 +30,6 @@ import org.apache.spark.sql.functions.col import org.apache.spark.sql.functions.concat import org.apache.spark.sql.functions.expr import org.apache.spark.sql.graphframes.GraphFrameInternals -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.WithCheckpointInterval -import org.apache.spark.graphframes.WithDirection -import org.apache.spark.graphframes.WithIntermediateStorageLevel -import org.apache.spark.graphframes.WithLocalCheckpoints /** * Computes all simple paths between source and destination vertices. diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/BFS.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/BFS.scala index 1eefebe749bd3..057119ece1d3e 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/BFS.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/BFS.scala @@ -17,15 +17,15 @@ package org.apache.spark.graphframes.lib +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrame.nestAsCol +import org.apache.spark.graphframes.Logging import org.apache.spark.sql.Column import org.apache.spark.sql.DataFrame import org.apache.spark.sql.Row import org.apache.spark.sql.functions.col import org.apache.spark.sql.functions.expr import org.apache.spark.sql.graphframes.GraphFrameInternals -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFrame.nestAsCol -import org.apache.spark.graphframes.Logging /** * Breadth-first search (BFS) @@ -43,7 +43,7 @@ import org.apache.spark.graphframes.Logging * - `v[i]` intermediate vertex i in the path, indexed from 1 * - `to` end vertex of path * Each of these columns is a StructType whose fields are the same as the columns of - * [[GraphFrame.vertices]] or [[GraphFrame.edges]]. + * the vertices or edges DataFrame. * * For example, suppose we have a graph g. Say the vertices DataFrame of g has columns "id" and * "job", and the edges DataFrame of g has columns "src", "dst", and "relation". @@ -61,7 +61,7 @@ import org.apache.spark.graphframes.Logging * If one or more vertices match both the from and to conditions, then there is a 0-hop path. The * returned DataFrame will have the "from" and "to" columns (as above); however, the "from" and * "to" columns will be exactly the same. There will be one row for each vertex in - * [[GraphFrame.vertices]] matching both `fromExpr` and `toExpr`. + * the vertices DataFrame matching both `fromExpr` and `toExpr`. * * Parameters: * diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ConnectedComponents.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ConnectedComponents.scala index 360b65776f6ed..6e2dbad698d0d 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ConnectedComponents.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ConnectedComponents.scala @@ -17,10 +17,8 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.graphx -import org.apache.spark.sql.DataFrame -import org.apache.spark.sql.graphframes.GraphFramesConf -import org.apache.spark.storage.StorageLevel +import java.util.Locale + import org.apache.spark.graphframes.GraphFrame import org.apache.spark.graphframes.GraphFramesUnreachableException import org.apache.spark.graphframes.Logging @@ -30,6 +28,10 @@ import org.apache.spark.graphframes.WithIntermediateStorageLevel import org.apache.spark.graphframes.WithLocalCheckpoints import org.apache.spark.graphframes.WithMaxIter import org.apache.spark.graphframes.WithUseLabelsAsComponents +import org.apache.spark.graphx +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.graphframes.GraphFramesConf +import org.apache.spark.storage.StorageLevel /** * Connected Components algorithm. @@ -69,16 +71,15 @@ class ConnectedComponents private[graphframes] (private val graph: GraphFrame) /** * Sets the algorithm to use for computing connected components. Supported values: - * - [[ConnectedComponents.ALGO_GRAPHX]]: use the GraphX implementation - * - [[ConnectedComponents.ALGO_GRAPHFRAMES]]: deprecated alias for - * [[ConnectedComponents.ALGO_TWO_PHASE]] - * - [[ConnectedComponents.ALGO_TWO_PHASE]]: use the two-phase label propagation + * - `graphx`: use the GraphX implementation + * - `graphframes`: deprecated alias for `two_phase` + * - `two_phase`: use the two-phase label propagation * implementation - * - [[ConnectedComponents.ALGO_RANDOMIZED_CONTRACTION]]: use the randomized contraction + * - `randomized_contraction`: use the randomized contraction * implementation */ def setAlgorithm(value: String): this.type = { - val normalized = value.toLowerCase + val normalized = value.toLowerCase(Locale.ROOT) normalized match { case ALGO_GRAPHX | ALGO_TWO_PHASE | ALGO_RANDOMIZED_CONTRACTION => algorithm = normalized @@ -102,7 +103,7 @@ class ConnectedComponents private[graphframes] (private val graph: GraphFrame) def getAlgorithm: String = algorithm /** - * !! WARNING: INTERNAL API — FOR VERY EXPERIENCED USERS ONLY !! + * !! WARNING: INTERNAL API - FOR VERY EXPERIENCED USERS ONLY !! * * Sets whether the graph has already been prepared before being passed to the algorithm, * skipping the internal graph preparation step. The default is `false`, meaning the algorithm diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/DetectingCycles.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/DetectingCycles.scala index 0204abcf90fa2..15987ae15432d 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/DetectingCycles.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/DetectingCycles.scala @@ -17,16 +17,16 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.Column -import org.apache.spark.sql.DataFrame -import org.apache.spark.sql.functions._ -import org.apache.spark.sql.types.ArrayType -import org.apache.spark.storage.StorageLevel import org.apache.spark.graphframes.GraphFrame import org.apache.spark.graphframes.Logging import org.apache.spark.graphframes.WithCheckpointInterval import org.apache.spark.graphframes.WithIntermediateStorageLevel import org.apache.spark.graphframes.WithLocalCheckpoints +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types.ArrayType +import org.apache.spark.storage.StorageLevel class DetectingCycles private[graphframes] (private val graph: GraphFrame) extends Arguments diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/GraphXConversions.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/GraphXConversions.scala index a4442a307f192..9223a50b01361 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/GraphXConversions.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/GraphXConversions.scala @@ -17,16 +17,15 @@ package org.apache.spark.graphframes.lib +import scala.reflect.runtime.universe._ + +import org.apache.spark.graphframes.{GraphFrame, NoSuchVertexException} import org.apache.spark.graphx.Graph import org.apache.spark.sql.DataFrame import org.apache.spark.sql.Row import org.apache.spark.sql.functions._ import org.apache.spark.sql.types.StructField import org.apache.spark.sql.types.StructType -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.NoSuchVertexException - -import scala.reflect.runtime.universe._ /** * Convenience functions to map GraphX graphs to GraphFrames, checking for the types expected by @@ -45,7 +44,8 @@ private[graphframes] object GraphXConversions { /** Indicates if T is a Product type */ private def isProductType[T: TypeTag]: Boolean = { val t = typeOf[T] - // See http://stackoverflow.com/questions/21209006/how-to-check-if-reflected-type-represents-a-tuple + // See https://stackoverflow.com/questions/21209006/ + // how-to-check-if-reflected-type-represents-a-tuple t.typeSymbol.fullName.startsWith("scala.Tuple") } diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/HyperANF.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/HyperANF.scala index fa4d8629c54fb..30d7081021361 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/HyperANF.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/HyperANF.scala @@ -17,6 +17,12 @@ package org.apache.spark.graphframes.lib +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFramesUnsupportedVertexTypeException +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithCheckpointInterval +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLocalCheckpoints import org.apache.spark.sql.Column import org.apache.spark.sql.DataFrame import org.apache.spark.sql.functions.col @@ -29,18 +35,13 @@ import org.apache.spark.sql.types.IntegerType import org.apache.spark.sql.types.LongType import org.apache.spark.sql.types.ShortType import org.apache.spark.sql.types.StringType -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFramesUnsupportedVertexTypeException -import org.apache.spark.graphframes.Logging -import org.apache.spark.graphframes.WithCheckpointInterval -import org.apache.spark.graphframes.WithIntermediateStorageLevel -import org.apache.spark.graphframes.WithLocalCheckpoints /** * HyperANF-style approximation of the neighbourhood function on top of GraphFrames. * - * This implementation is inspired by - * [[https://arxiv.org/pdf/1011.5599 Vigna, Paolo; Boldi, Marco; Rosa, Sebastiano. "HyperANF: Approximating the Neighbourhood Function of Very Large Graphs on a Budget." arXiv preprint arXiv:1011.5599 (2010)]]. + * This implementation is inspired by Vigna, Boldi, and Rosa, + * "HyperANF: Approximating the Neighbourhood Function of + * Very Large Graphs on a Budget" (2010). * * The input graph is treated as directed: for each vertex, reachability is computed by following * outgoing edges from `src` to `dst`. diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/KCore.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/KCore.scala index 9227601377fbe..57dd95c72cd1b 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/KCore.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/KCore.scala @@ -17,6 +17,11 @@ package org.apache.spark.graphframes.lib +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithCheckpointInterval +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLocalCheckpoints import org.apache.spark.sql.DataFrame import org.apache.spark.sql.catalyst.FunctionIdentifier import org.apache.spark.sql.catalyst.expressions.Expression @@ -28,11 +33,6 @@ import org.apache.spark.sql.functions.when import org.apache.spark.sql.graphframes.expressions.KCoreMerge import org.apache.spark.sql.types.IntegerType import org.apache.spark.storage.StorageLevel -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.Logging -import org.apache.spark.graphframes.WithCheckpointInterval -import org.apache.spark.graphframes.WithIntermediateStorageLevel -import org.apache.spark.graphframes.WithLocalCheckpoints /** * K-Core decomposition algorithm implementation for GraphFrames. @@ -48,7 +48,7 @@ import org.apache.spark.graphframes.WithLocalCheckpoints * * '''Edge representation''': K-core decomposition is defined for undirected graphs. Since * GraphFrames represents edges as directed, each undirected edge `{u, v}` should be supplied as a - * single directed edge in either direction — the algorithm symmetrizes internally. Supplying both + * single directed edge in either direction; the algorithm symmetrizes internally. Supplying both * `(u, v)` and `(v, u)` will double-count the edge and produce incorrect results. */ class KCore private[graphframes] (private val graph: GraphFrame) @@ -119,7 +119,8 @@ object KCore extends Serializable with Logging { new FunctionIdentifier("_kcoreMerge", Some("builtin"), Some("system"))) if (!dereg) { logWarn( - "graphframes faced an internal error and was not able to de-register function _kcoreMerge; Spark' functionRegistry is in a bad state") + "GraphFrames faced an internal error and could not de-register function _kcoreMerge; " + + "Spark's function registry is in a bad state") } } } diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/LabelPropagation.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/LabelPropagation.scala index becc77e6ff78b..520d953eef644 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/LabelPropagation.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/LabelPropagation.scala @@ -17,13 +17,6 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.graphx -import org.apache.spark.sql.Column -import org.apache.spark.sql.DataFrame -import org.apache.spark.sql.functions._ -import org.apache.spark.sql.types.IntegerType -import org.apache.spark.sql.types.MapType -import org.apache.spark.storage.StorageLevel import org.apache.spark.graphframes.GraphFrame import org.apache.spark.graphframes.Logging import org.apache.spark.graphframes.WithAlgorithmChoice @@ -31,6 +24,13 @@ import org.apache.spark.graphframes.WithCheckpointInterval import org.apache.spark.graphframes.WithIntermediateStorageLevel import org.apache.spark.graphframes.WithLocalCheckpoints import org.apache.spark.graphframes.WithMaxIter +import org.apache.spark.graphx +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.types.MapType +import org.apache.spark.storage.StorageLevel /** * Run static Label Propagation for detecting communities in networks. @@ -84,8 +84,9 @@ private object LabelPropagation { private def keyWithMaxValue(column: Column): Column = { // Get the key with the highest value, using the key to break a tie. To do this, simply get - // map entries, swap the value and key columns to create the natural ordering, multiply key by -1 and then - // take the key from the max entry (multiply it by -1 again to get the original key). + // map entries, swap the value and key columns to create the natural ordering, multiply key by + // -1 and then take the key from the max entry (multiply it by -1 again to get the original + // key). array_max( transform( map_entries(column), diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/MaximalIndependentSet.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/MaximalIndependentSet.scala index 86785afcff2a0..fef0d3062eccf 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/MaximalIndependentSet.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/MaximalIndependentSet.scala @@ -17,17 +17,17 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.DataFrame -import org.apache.spark.sql.functions._ -import org.apache.spark.sql.types.DoubleType -import org.apache.spark.storage.StorageLevel +import java.io.IOException + import org.apache.spark.graphframes.GraphFrame import org.apache.spark.graphframes.Logging import org.apache.spark.graphframes.WithCheckpointInterval import org.apache.spark.graphframes.WithIntermediateStorageLevel import org.apache.spark.graphframes.WithLocalCheckpoints - -import java.io.IOException +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types.DoubleType +import org.apache.spark.storage.StorageLevel /** * This class implements a distributed algorithm for finding a Maximal Independent Set (MIS) in a diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/PageRank.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/PageRank.scala index faf92ed4cfcc9..66bc2959aadce 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/PageRank.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/PageRank.scala @@ -17,9 +17,9 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.graphx.{lib => graphxlib} import org.apache.spark.graphframes.GraphFrame import org.apache.spark.graphframes.Logging +import org.apache.spark.graphx.{lib => graphxlib} /** * PageRank algorithm implementation. There are two implementations of PageRank. diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRank.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRank.scala index 6e183eae9ef24..3da7236ec2c9b 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRank.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRank.scala @@ -17,15 +17,15 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.graphx import org.apache.spark.graphframes.GraphFrame import org.apache.spark.graphframes.Logging import org.apache.spark.graphframes.WithMaxIter +import org.apache.spark.graphx /** * Parallel Personalized PageRank algorithm implementation. * - * This implementation uses the standalone [[GraphFrame]] interface and runs personalized PageRank + * This implementation uses the standalone `GraphFrame` interface and runs personalized PageRank * in parallel for a fixed number of iterations. This can be run by setting `maxIter`. The source * vertex Ids are set in `sourceIds`. A simple local implementation of this algorithm is as * follows. diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/Pregel.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/Pregel.scala index 3d4d492bea7c8..4dd3c87a5315c 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/Pregel.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/Pregel.scala @@ -17,6 +17,16 @@ package org.apache.spark.graphframes.lib +import java.io.IOException + +import scala.util.control.Breaks.break +import scala.util.control.Breaks.breakable + +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrame._ +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLocalCheckpoints import org.apache.spark.sql.Column import org.apache.spark.sql.DataFrame import org.apache.spark.sql.functions.array @@ -25,15 +35,6 @@ import org.apache.spark.sql.functions.explode import org.apache.spark.sql.functions.lit import org.apache.spark.sql.functions.struct import org.apache.spark.sql.graphframes.GraphFrameInternals -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFrame._ -import org.apache.spark.graphframes.Logging -import org.apache.spark.graphframes.WithIntermediateStorageLevel -import org.apache.spark.graphframes.WithLocalCheckpoints - -import java.io.IOException -import scala.util.control.Breaks.break -import scala.util.control.Breaks.breakable /** * Implements a Pregel-like bulk-synchronous message-passing API based on DataFrame operations. @@ -42,9 +43,8 @@ import scala.util.control.Breaks.breakable * large-scale graph processing for a detailed description of the Pregel algorithm. * * You can construct a Pregel instance using either this constructor or - * [[org.apache.spark.graphframes.GraphFrame#pregel]], then use builder pattern to describe the - * operations, and then call [[run]] to start a run. It returns a DataFrame of vertices from the - * last iteration. + * `GraphFrame.pregel`, then use builder pattern to describe the operations, and then call [[run]] + * to start a run. It returns a DataFrame of vertices from the last iteration. * * When a run starts, it expands the vertices DataFrame using column expressions defined by * [[withVertexColumn]]. Those additional vertex properties can be changed during Pregel @@ -83,13 +83,12 @@ import scala.util.control.Breaks.breakable * `StructType`). That behavior is considered as bug and starting from 0.12 edge columns are not * kept by default. * + * See `GraphFrame.pregel` and + * Malewicz et al., Pregel: a system for + * large-scale graph processing. + * * @param graph * The graph that Pregel will run on. - * @see - * [[org.apache.spark.graphframes.GraphFrame#pregel]] - * @see - * Malewicz et al., Pregel: a system for - * large-scale graph processing. */ class Pregel(val graph: GraphFrame) extends Logging @@ -250,7 +249,7 @@ class Pregel(val graph: GraphFrame) * @param updateAfterAggMsgsExpr * the expression to update the additional vertex column after messages aggregation. You can * reference all original vertex columns, additional vertex columns, and the aggregated - * message column using [[Pregel$#msg]]. If the vertex received no messages, the message + * message column using `Pregel.msg`. If the vertex received no messages, the message * column would be null. */ def withVertexColumn( @@ -279,11 +278,9 @@ class Pregel(val graph: GraphFrame) * @param msgExpr * the expression of the message to send to the source vertex given a (src, edge, dst) * triplet. Source/destination vertex properties and edge properties are nested under columns - * `src`, `dst`, and `edge`, respectively. You can reference them using [[Pregel$#src]], - * [[Pregel$#dst]], and [[Pregel$#edge]]. Null messages are not included in message + * `src`, `dst`, and `edge`, respectively. You can reference them using `Pregel.src`, + * `Pregel.dst`, and `Pregel.edge`. Null messages are not included in message * aggregation. - * @see - * [[sendMsgToDst]] */ def sendMsgToSrc(msgExpr: Column): this.type = { sendMsgs += Tuple2(Pregel.src(ID), msgExpr) @@ -298,11 +295,9 @@ class Pregel(val graph: GraphFrame) * @param msgExpr * the message expression to send to the destination vertex given a (`src`, `edge`, `dst`) * triplet. Source/destination vertex properties and edge properties are nested under columns - * `src`, `dst`, and `edge`, respectively. You can reference them using [[Pregel$#src]], - * [[Pregel$#dst]], and [[Pregel$#edge]]. Null messages are not included in message + * `src`, `dst`, and `edge`, respectively. You can reference them using `Pregel.src`, + * `Pregel.dst`, and `Pregel.edge`. Null messages are not included in message * aggregation. - * @see - * [[sendMsgToSrc]] */ def sendMsgToDst(msgExpr: Column): this.type = { sendMsgs += Tuple2(Pregel.dst(ID), msgExpr) @@ -323,8 +318,6 @@ class Pregel(val graph: GraphFrame) * the first required source vertex column name * @param colNames * additional required source vertex column names - * @see - * [[requiredDstColumns]] */ def requiredSrcColumns(colName: String, colNames: String*): this.type = { requiredSrcColumnsList.clear() @@ -347,8 +340,6 @@ class Pregel(val graph: GraphFrame) * the first required destination vertex column name * @param colNames * additional required destination vertex column names - * @see - * [[requiredSrcColumns]] */ def requiredDstColumns(colName: String, colNames: String*): this.type = { requiredDstColumnsList.clear() @@ -368,8 +359,6 @@ class Pregel(val graph: GraphFrame) * the first required edge column name * @param colNames * additional required edge column names - * @see - * [[requiredSrcColumns]] and [[requiredDstColumns]] */ def requiredEdgeColumns(colName: String, colNames: String*): this.type = { requiredEdgeColumnsList.clear() @@ -383,8 +372,8 @@ class Pregel(val graph: GraphFrame) * * @param aggExpr * the message aggregation expression, such as `sum(Pregel.msg)`. You can reference the - * message column by [[Pregel$#msg]] and the vertex ID by [[GraphFrame$#ID]], while the latter - * is usually not used. + * message column by `Pregel.msg` and the vertex ID by `GraphFrame.ID`, while the latter is + * usually not used. */ def aggMsgs(aggExpr: Column): this.type = { aggMsgsCol = aggExpr @@ -469,7 +458,10 @@ class Pregel(val graph: GraphFrame) val shouldCheckpoint = checkpointInterval > 0 - if (shouldCheckpoint && graph.spark.sparkContext.getCheckpointDir.isEmpty && !useLocalCheckpoints) { + if ( + shouldCheckpoint && + graph.spark.sparkContext.getCheckpointDir.isEmpty && + !useLocalCheckpoints) { // Spark Connect workaround graph.spark.conf.getOption("spark.checkpoint.dir") match { case Some(d) => graph.spark.sparkContext.setCheckpointDir(d) @@ -497,9 +489,11 @@ class Pregel(val graph: GraphFrame) // Prune non-active vertices early if skipMessagesFromNonActiveVertices // is enabled and we don't need the dst state. val srcVertices = - if (!needsDstState && skipMessagesFromNonActiveVertices) + if (!needsDstState && skipMessagesFromNonActiveVertices) { currentVertices.filter(col(Pregel.ACTIVE_FLAG_COL)) - else currentVertices + } else { + currentVertices + } // Build triplets: start with src vertex state joined with edges val srcWithEdges = srcVertices @@ -621,8 +615,6 @@ object Pregel extends Serializable { /** * References the message column in aggregating messages and updating additional vertex columns. * - * @see - * [[Pregel.aggMsgs]] and [[Pregel.withVertexColumn]] */ val msg: Column = col(MSG_COL_NAME) @@ -631,8 +623,6 @@ object Pregel extends Serializable { * * @param colName * the vertex column name. - * @see - * [[Pregel.sendMsgToSrc]] and [[Pregel.sendMsgToDst]] */ def src(colName: String): Column = col(GraphFrame.SRC + "." + colName) @@ -641,8 +631,6 @@ object Pregel extends Serializable { * * @param colName * the vertex column name. - * @see - * [[Pregel.sendMsgToSrc]] and [[Pregel.sendMsgToDst]] */ def dst(colName: String): Column = col(GraphFrame.DST + "." + colName) @@ -651,8 +639,6 @@ object Pregel extends Serializable { * * @param colName * the edge column name. - * @see - * [[Pregel.sendMsgToSrc]] and [[Pregel.sendMsgToDst]] */ def edge(colName: String): Column = col(GraphFrame.EDGE + "." + colName) } diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/RandomizedContraction.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/RandomizedContraction.scala index e42c1cba259fe..8a85b4f3f2fdb 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/RandomizedContraction.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/RandomizedContraction.scala @@ -17,7 +17,17 @@ package org.apache.spark.graphframes.lib +import java.io.IOException +import java.util.UUID + +import scala.collection.mutable.Stack +import scala.util.Random + import org.apache.hadoop.fs.Path + +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrame.{DST, ID, LONG_DST, LONG_ID, LONG_SRC, SRC} +import org.apache.spark.graphframes.Logging import org.apache.spark.sql.Column import org.apache.spark.sql.DataFrame import org.apache.spark.sql.catalyst.FunctionIdentifier @@ -25,23 +35,10 @@ import org.apache.spark.sql.catalyst.expressions.Expression import org.apache.spark.sql.functions._ import org.apache.spark.sql.graphframes.expressions.FiniteAXPlusB import org.apache.spark.storage.StorageLevel -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFrame.DST -import org.apache.spark.graphframes.GraphFrame.ID -import org.apache.spark.graphframes.GraphFrame.LONG_DST -import org.apache.spark.graphframes.GraphFrame.LONG_ID -import org.apache.spark.graphframes.GraphFrame.LONG_SRC -import org.apache.spark.graphframes.GraphFrame.SRC -import org.apache.spark.graphframes.Logging - -import java.io.IOException -import java.util.UUID -import scala.collection.mutable.Stack -import scala.util.Random /** * Implementation of parallel connected components algorithm using randomized contraction, based - * on Bögeholz, Harald, Michael Brand, and Radu-Alexandru Todor. "In-database connected component + * on Boegeholz, Harald, Michael Brand, and Radu-Alexandru Todor. "In-database connected component * analysis." 2020 IEEE 36th International Conference on Data Engineering (ICDE). IEEE, 2020. * * The algorithm contracts the graph iteratively using random linear functions, until no edges @@ -284,7 +281,8 @@ private[graphframes] object RandomizedContraction extends Logging with Serializa new FunctionIdentifier("_axpb", Some("builtin"), Some("system"))) if (!dereg) { logWarn( - "graphframes faced an internal error and was not able to de-register function _axpb; Spark' functionRegistry is in a bad state") + "GraphFrames faced an internal error and could not de-register function _axpb; " + + "Spark's function registry is in a bad state") } } } diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/SVDPlusPlus.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/SVDPlusPlus.scala index 5b61a7393c54d..2d18814b51ceb 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/SVDPlusPlus.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/SVDPlusPlus.scala @@ -17,14 +17,14 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.graphx.Edge -import org.apache.spark.graphx.{lib => graphxlib} -import org.apache.spark.sql.DataFrame -import org.apache.spark.sql.functions.col import org.apache.spark.graphframes.GraphFrame import org.apache.spark.graphframes.InvalidGraphException import org.apache.spark.graphframes.Logging import org.apache.spark.graphframes.WithMaxIter +import org.apache.spark.graphx.{lib => graphxlib} +import org.apache.spark.graphx.Edge +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.col /** * Arguments for SVD++ algorithm. @@ -32,9 +32,9 @@ import org.apache.spark.graphframes.WithMaxIter * This class implements the SVD++ algorithm for Collaborative Filtering, primarily used for * Recommender Systems (Link Prediction). * - * Based on the paper "Factorization Meets the Neighborhood: a Multifaceted Collaborative - * Filtering Model" by Yehuda Koren (2008), available at - * [[https://dl.acm.org/citation.cfm?id=1401944]]. + * Based on Yehuda Koren's paper + * "Factorization Meets the Neighborhood: a + * Multifaceted Collaborative Filtering Model" (2008). * * ==Problem Definition== * The algorithm predicts unknown ratings in a user-item system. It accounts for: @@ -44,7 +44,7 @@ import org.apache.spark.graphframes.WithMaxIter * * The prediction rule for a rating `r_ui` (user `u`, item `i`) is: * {{{ - * r_ui = µ + b_u + b_i + q_i^T * (p_u + |N(u)|^-0.5 * sum(y_j for j in N(u))) + * r_ui = mu + b_u + b_i + q_i^T * (p_u + |N(u)|^-0.5 * sum(y_j for j in N(u))) * }}} * Where `N(u)` is the set of items user `u` has interacted with (implicit feedback). * diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ShortestPaths.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ShortestPaths.scala index 40acfc814721b..f30d575f299d3 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ShortestPaths.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/ShortestPaths.scala @@ -17,6 +17,19 @@ package org.apache.spark.graphframes.lib +import java.util + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrame.quote +import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.graphframes.Logging +import org.apache.spark.graphframes.WithAlgorithmChoice +import org.apache.spark.graphframes.WithCheckpointInterval +import org.apache.spark.graphframes.WithDirection +import org.apache.spark.graphframes.WithIntermediateStorageLevel +import org.apache.spark.graphframes.WithLocalCheckpoints import org.apache.spark.graphx import org.apache.spark.sql.Column import org.apache.spark.sql.DataFrame @@ -33,18 +46,6 @@ import org.apache.spark.sql.functions.when import org.apache.spark.sql.types.IntegerType import org.apache.spark.sql.types.MapType import org.apache.spark.storage.StorageLevel -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFrame.quote -import org.apache.spark.graphframes.GraphFramesUnreachableException -import org.apache.spark.graphframes.Logging -import org.apache.spark.graphframes.WithAlgorithmChoice -import org.apache.spark.graphframes.WithCheckpointInterval -import org.apache.spark.graphframes.WithDirection -import org.apache.spark.graphframes.WithIntermediateStorageLevel -import org.apache.spark.graphframes.WithLocalCheckpoints - -import java.util -import scala.jdk.CollectionConverters._ /** * Computes shortest paths from every vertex to the given set of landmark vertices. Note that this diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponents.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponents.scala index aa11d4ec3c666..42618a09246f0 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponents.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponents.scala @@ -17,12 +17,10 @@ package org.apache.spark.graphframes.lib +import org.apache.spark.graphframes.{GraphFrame, Logging, WithMaxIter} import org.apache.spark.graphx.{lib => graphxlib} import org.apache.spark.sql.DataFrame import org.apache.spark.storage.StorageLevel -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.Logging -import org.apache.spark.graphframes.WithMaxIter /** * Compute the strongly connected component (SCC) of each vertex and return a DataFrame with each diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala index 724b98191ec7d..e0be7affda7d7 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala @@ -17,10 +17,6 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.Column -import org.apache.spark.sql.DataFrame -import org.apache.spark.sql.functions._ -import org.apache.spark.sql.types._ import org.apache.spark.graphframes.GraphFrame import org.apache.spark.graphframes.GraphFrame._ import org.apache.spark.graphframes.GraphFramesSparkVersionException @@ -31,6 +27,10 @@ import org.apache.spark.graphframes.WithIntermediateStorageLevel import org.apache.spark.graphframes.WithLgNomEntries import org.apache.spark.graphframes.WithLocalCheckpoints import org.apache.spark.graphframes.WithMaxIter +import org.apache.spark.sql.Column +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types._ /** * Neighborhood-aware community detection via weighted label propagation. @@ -132,7 +132,8 @@ class StructureAwareLabelPropagation private[graphframes] (private val graph: Gr def setInitialLabelCol(col: String): this.type = { require( graph.vertices.columns.contains(col), - s"Initial label column '$col' does not exist in vertex columns: ${graph.vertices.columns.mkString(", ")}") + s"Initial label column '$col' does not exist in vertex columns: " + + graph.vertices.columns.mkString(", ")) initialLabelCol = Some(col) this } diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/TriangleCount.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/TriangleCount.scala index 069c994c4dc83..b739ee113f65e 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/TriangleCount.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/TriangleCount.scala @@ -17,14 +17,14 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.DataFrame -import org.apache.spark.sql.functions._ -import org.apache.spark.storage.StorageLevel import org.apache.spark.graphframes.GraphFrame import org.apache.spark.graphframes.GraphFramesSparkVersionException import org.apache.spark.graphframes.Logging import org.apache.spark.graphframes.WithIntermediateStorageLevel import org.apache.spark.graphframes.WithLgNomEntries +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions._ +import org.apache.spark.storage.StorageLevel /** * Triangle count implementation. diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/TwoPhase.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/TwoPhase.scala index 4cf3e70fb11eb..91c60235a72a7 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/lib/TwoPhase.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/lib/TwoPhase.scala @@ -17,25 +17,20 @@ package org.apache.spark.graphframes.lib +import java.io.IOException +import java.math.BigDecimal +import java.util.UUID + import org.apache.hadoop.fs.Path + +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrame.{ATTR, DST, ID, LONG_DST, LONG_ID, LONG_SRC, SRC} +import org.apache.spark.graphframes.Logging import org.apache.spark.sql.Column import org.apache.spark.sql.DataFrame import org.apache.spark.sql.functions._ import org.apache.spark.sql.types.DecimalType import org.apache.spark.storage.StorageLevel -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFrame.ATTR -import org.apache.spark.graphframes.GraphFrame.DST -import org.apache.spark.graphframes.GraphFrame.ID -import org.apache.spark.graphframes.GraphFrame.LONG_DST -import org.apache.spark.graphframes.GraphFrame.LONG_ID -import org.apache.spark.graphframes.GraphFrame.LONG_SRC -import org.apache.spark.graphframes.GraphFrame.SRC -import org.apache.spark.graphframes.Logging - -import java.io.IOException -import java.math.BigDecimal -import java.util.UUID /** * Two-phase label propagation implementation of connected components. @@ -279,8 +274,8 @@ private[graphframes] object TwoPhase extends Logging { case Some(d) => new Path(d, s"$CHECKPOINT_NAME_PREFIX-$runId").toString case None => throw new IOException( - "Checkpoint directory is not set. Please set it first using sc.setCheckpointDir()" + - "or by specifying the conf 'spark.checkpoint.dir'.") + "Checkpoint directory is not set. Please set it first using " + + "sc.setCheckpointDir() or by specifying the conf 'spark.checkpoint.dir'.") } } logInfo(s"$logPrefix Using $dir for checkpointing with interval $checkpointInterval.") diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/pattern/patterns.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/pattern/patterns.scala index ee95049ec64ee..c15169112fe1e 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/pattern/patterns.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/pattern/patterns.scala @@ -17,12 +17,12 @@ package org.apache.spark.graphframes.pattern -import org.apache.spark.graphframes.GraphFramesUnreachableException -import org.apache.spark.graphframes.InvalidParseException - import scala.collection.mutable import scala.util.parsing.combinator._ +import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.graphframes.InvalidParseException + /** * Parser for graph patterns for motif finding. Copied from GraphFrames with minor modification. */ @@ -195,8 +195,8 @@ private[graphframes] object Pattern { case AnonymousEdge(AnonymousVertex, AnonymousVertex) => throw new InvalidParseException( "Motif finding does not support completely " + - "anonymous negated edges !()-[]-(). Users can check for the existence of edges in the " + - "graph using the edges DataFrame.") + "anonymous negated edges !()-[]-(). Users can check for the existence of " + + "edges in the graph using the edges DataFrame.") case _ => addEdge(e) } case e @ AnonymousEdge(_, _) => diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkBase.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkBase.scala index 7aafb15e0add1..d6790f70ea86e 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkBase.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkBase.scala @@ -17,6 +17,11 @@ package org.apache.spark.graphframes.rw +import scala.util.Random + +import org.apache.hadoop.fs.Path + +import org.apache.spark.graphframes.{GraphFrame, Logging, WithIntermediateStorageLevel} import org.apache.spark.sql.Column import org.apache.spark.sql.DataFrame import org.apache.spark.sql.functions.array @@ -32,11 +37,6 @@ import org.apache.spark.sql.functions.xxhash64 import org.apache.spark.sql.graphframes.expressions.KMinSampling import org.apache.spark.sql.types.ArrayType import org.apache.spark.sql.types.DataType -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.Logging -import org.apache.spark.graphframes.WithIntermediateStorageLevel - -import scala.util.Random /** * Base trait for implementing random walk algorithms on graph data. Provides common functionality @@ -416,16 +416,16 @@ object RandomWalkBase extends Serializable { spark: org.apache.spark.sql.SparkSession): Unit = { val sc = spark.sparkContext val hadoopConf = sc.hadoopConfiguration - val fs = org.apache.hadoop.fs.FileSystem.get(hadoopConf) val basePath = temporaryPrefix val runPath = if (basePath.endsWith("/")) { s"${basePath}${runID}_batch_" } else { s"${basePath}/${runID}_batch_" } + val fs = new Path(runPath).getFileSystem(hadoopConf) // Delete all batch directories (1 to numBatches) for (i <- 1 to numBatches) { - val path = new org.apache.hadoop.fs.Path(s"${runPath}${i}") + val path = new Path(s"${runPath}${i}") if (fs.exists(path)) { fs.delete(path, true) // recursive delete } diff --git a/graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestart.scala b/graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestart.scala index 06302a15a8455..b96e461a64045 100644 --- a/graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestart.scala +++ b/graphframes/src/main/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestart.scala @@ -17,10 +17,10 @@ package org.apache.spark.graphframes.rw +import org.apache.spark.graphframes.GraphFrame import org.apache.spark.sql.DataFrame import org.apache.spark.sql.functions._ import org.apache.spark.sql.types.ArrayType -import org.apache.spark.graphframes.GraphFrame /** * An implementation of random walk with restart. At each step of the walk, there is a probability diff --git a/graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFrameInternals.scala b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFrameInternals.scala index 308ade66076e8..14f0e406b3f34 100644 --- a/graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFrameInternals.scala +++ b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFrameInternals.scala @@ -17,6 +17,8 @@ package org.apache.spark.sql.graphframes +import scala.collection.mutable + import org.apache.spark.sql.Column import org.apache.spark.sql.DataFrame import org.apache.spark.sql.SparkSession @@ -27,13 +29,10 @@ import org.apache.spark.sql.catalyst.expressions.Expression import org.apache.spark.sql.catalyst.expressions.GetStructField import org.apache.spark.sql.catalyst.expressions.Literal import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan -import org.apache.spark.sql.classic.ClassicConversions._ import org.apache.spark.sql.classic.{DataFrame => ClassicDataFrame} -import org.apache.spark.sql.classic.Dataset -import org.apache.spark.sql.classic.ExpressionUtils +import org.apache.spark.sql.classic.{Dataset, ExpressionUtils} import org.apache.spark.sql.classic.{SparkSession => ClassicSparkSession} - -import scala.collection.mutable +import org.apache.spark.sql.classic.ClassicConversions._ object GraphFrameInternals { diff --git a/graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala index 38e3e19ee8e14..14cd6bface0e7 100644 --- a/graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala +++ b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala @@ -17,6 +17,8 @@ package org.apache.spark.sql.graphframes +import java.util.Locale + import org.apache.spark.internal.config.ConfigEntry import org.apache.spark.sql.SparkSession import org.apache.spark.sql.internal.SQLConf @@ -26,8 +28,9 @@ object GraphFramesConf { private val USE_LOCAL_CHECKPOINTS = SQLConf .buildConf("spark.graphframes.useLocalCheckpoints") - .doc(""" Tells the connected components algorithm to use local checkpoints (default: "false"). - | If set to "true", iterative algorithm will use the checkpointing mechanism to the persistent storage. + .doc(""" Tells the connected components algorithm to use local checkpoints + | (default: "false"). If set to "true", iterative algorithms will use the + | checkpointing mechanism to persistent storage. | Local checkpoints are faster but can make the whole job less prone to errors. | @note This option may become default "true" in the future. |""".stripMargin) @@ -38,8 +41,9 @@ object GraphFramesConf { private val USE_LABELS_AS_COMPONENTS = SQLConf .buildConf("spark.graphframes.useLabelsAsComponents") - .doc(""" Tells the connected components algorithm to use (default: "true") labels as components in the output - | DataFrame. If set to "false", randomly generated labels with the data type LONG will returned. + .doc(""" Tells the connected components algorithm to use labels as components in the + | output DataFrame (default: "true"). If set to "false", randomly generated labels + | with the data type LONG will be returned. |""".stripMargin) .version("0.9.0") .booleanConf @@ -48,15 +52,19 @@ object GraphFramesConf { private val CONNECTED_COMPONENTS_ALGORITHM = SQLConf .buildConf("spark.graphframes.connectedComponents.algorithm") - .doc(""" Sets the connected components algorithm to use (default: "graphframes"). Supported algorithms + .doc(""" Sets the connected components algorithm to use (default: "graphframes"). + | Supported algorithms: | - "two_phase": Uses alternating large star and small star iterations proposed in - | [[http://dx.doi.org/10.1145/2670979.2670997 Connected Components in MapReduce and Beyond]] + | "Connected Components in MapReduce and Beyond" + | (http://dx.doi.org/10.1145/2670979.2670997). | - "randomized_contraction": Uses randomized algorithm proposed in - | [[https://arxiv.org/pdf/1802.09478 In-database connected component analysis]] + | "In-database connected component analysis" + | (https://arxiv.org/pdf/1802.09478). | - "graphframes": Deprecated alias for "two_phase" - | - "graphx": Converts the graph to a GraphX graph and then uses the connected components - | implementation in GraphX. - | @see org.apache.spark.graphframes.lib.ConnectedComponents.supportedAlgorithms""".stripMargin) + | - "graphx": Converts the graph to a GraphX graph and then uses the connected + | components implementation in GraphX. + | @see org.apache.spark.graphframes.lib.ConnectedComponents.supportedAlgorithms + |""".stripMargin) .version("0.9.0") .stringConf .createOptional @@ -64,11 +72,13 @@ object GraphFramesConf { private val CONNECTED_COMPONENTS_BROADCAST_THRESHOLD = SQLConf .buildConf("spark.graphframes.connectedComponents.broadcastthreshold") - .doc(""" Sets broadcast threshold in propagating component assignments (default: 1000000). If a node - | degree is greater than this threshold at some iteration, its component assignment will be - | collected and then broadcasted back to propagate the assignment to its neighbors. Otherwise, - | the assignment propagation is done by a normal Spark join. This parameter is only used when - | the algorithm is set to "graphframes".""".stripMargin) + .doc(""" Sets broadcast threshold in propagating component assignments + | (default: 1000000). If a node degree is greater than this threshold at some + | iteration, its component assignment will be collected and then broadcast back to + | propagate the assignment to its neighbors. Otherwise, the assignment propagation + | is done by a normal Spark join. This parameter is only used when the algorithm is + | set to "graphframes". + |""".stripMargin) .version("0.9.0") .intConf .createOptional @@ -79,14 +89,15 @@ object GraphFramesConf { .doc(""" Sets checkpoint interval in terms of number of iterations (default: 2). Checkpointing | regularly helps recover from failures, clean shuffle files, shorten the lineage of the | computation graph, and reduce the complexity of plan optimization. As of Spark 2.0, the - | complexity of plan optimization would grow exponentially without checkpointing. Hence, - | disabling or setting longer-than-default checkpoint intervals are not recommended. Checkpoint - | data is saved under `org.apache.spark.SparkContext.getCheckpointDir` with prefix - | "connected-components". If the checkpoint directory is not set, this throws a - | `java.io.IOException`. Set a nonpositive value to disable checkpointing. This parameter is - | only used when the algorithm is set to "graphframes". Its default value might change in the - | future. - | @see `org.apache.spark.SparkContext.setCheckpointDir` in Spark API doc""".stripMargin) + | complexity of plan optimization would grow exponentially without checkpointing. + | Hence, disabling or setting longer-than-default checkpoint intervals are not + | recommended. Checkpoint data is saved under + | `org.apache.spark.SparkContext.getCheckpointDir` with prefix "connected-components". + | If the checkpoint directory is not set, this throws a `java.io.IOException`. Set a + | nonpositive value to disable checkpointing. This parameter is only used when the + | algorithm is set to "graphframes". Its default value might change in the future. + | @see `org.apache.spark.SparkContext.setCheckpointDir` in Spark API doc + |""".stripMargin) .version("0.9.0") .intConf .createOptional @@ -94,7 +105,9 @@ object GraphFramesConf { private val CONNECTED_COMPONENTS_INTERMEDIATE_STORAGE_LEVEL = SQLConf .buildConf("spark.graphframes.connectedComponents.intermediatestoragelevel") - .doc("Sets storage level for intermediate datasets that require multiple passes (default: ``MEMORY_AND_DISK``).") + .doc( + "Sets storage level for intermediate datasets that require multiple passes " + + "(default: ``MEMORY_AND_DISK``).") .version("0.9.0") .stringConf .createOptional @@ -109,7 +122,7 @@ object GraphFramesConf { def getConnectedComponentsAlgorithm: Option[String] = { get(CONNECTED_COMPONENTS_ALGORITHM) match { - case Some(threshold) => Some(threshold.toLowerCase) + case Some(threshold) => Some(threshold.toLowerCase(Locale.ROOT)) case _ => None } } @@ -130,7 +143,7 @@ object GraphFramesConf { def getConnectedComponentsStorageLevel: Option[StorageLevel] = { get(CONNECTED_COMPONENTS_INTERMEDIATE_STORAGE_LEVEL) match { - case Some(level) => Some(StorageLevel.fromString(level.toUpperCase)) + case Some(level) => Some(StorageLevel.fromString(level.toUpperCase(Locale.ROOT))) case _ => None } } diff --git a/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KCoreMerge.scala b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KCoreMerge.scala index 32d637c516d59..0189fb560b793 100644 --- a/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KCoreMerge.scala +++ b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KCoreMerge.scala @@ -36,7 +36,7 @@ import org.apache.spark.sql.types.IntegerType * @param right * core of the vertex */ -case class KCoreMerge(left: Expression, right: Expression) +private[spark] case class KCoreMerge(left: Expression, right: Expression) extends BinaryExpression with CodegenFallback { override protected def withNewChildrenInternal( @@ -47,7 +47,7 @@ case class KCoreMerge(left: Expression, right: Expression) /** * Each node initializes its core value with the degree of itself. Each node (say u) then sends - * messages to its neighbors v ∈ N (u) with the current estimate of its (u’s) core value. For an + * messages to its neighbors v in N(u) with the current estimate of its (u's) core value. For an * undirected graph with m edges, there can be at most a total of 2m messages that have been * sent during a message passing session. Upon receiving all the messages from its neighbors, * the vertex u computes the largest value l such that the number of neighbors of u whose diff --git a/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KMinSampling.scala b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KMinSampling.scala index 0401ac8ff779e..b317fa22d70e9 100644 --- a/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KMinSampling.scala +++ b/graphframes/src/main/scala/org/apache/spark/sql/graphframes/expressions/KMinSampling.scala @@ -17,6 +17,11 @@ package org.apache.spark.sql.graphframes.expressions +import scala.annotation.nowarn +import scala.reflect.ClassTag +import scala.reflect.runtime.universe.TypeTag + +import org.apache.spark.graphframes.GraphFramesUnsupportedVertexTypeException import org.apache.spark.sql.Encoder import org.apache.spark.sql.Encoders import org.apache.spark.sql.Row @@ -26,12 +31,6 @@ import org.apache.spark.sql.expressions.Aggregator import org.apache.spark.sql.expressions.UserDefinedFunction import org.apache.spark.sql.functions.udaf import org.apache.spark.sql.types._ -import org.apache.spark.sql.types.DataType -import org.apache.spark.graphframes.GraphFramesUnsupportedVertexTypeException - -import scala.annotation.nowarn -import scala.reflect.ClassTag -import scala.reflect.runtime.universe.TypeTag case class KMinAccum[T](values: Array[T], weights: Array[Long], var cnt: Int) extends Serializable diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameInternalsSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameInternalsSuite.scala index 430cfe27588a9..8b44d836770da 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameInternalsSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameInternalsSuite.scala @@ -17,9 +17,9 @@ package org.apache.spark.graphframes +import org.apache.spark.graphframes.lib.Pregel import org.apache.spark.sql.functions._ import org.apache.spark.sql.graphframes.GraphFrameInternals -import org.apache.spark.graphframes.lib.Pregel /** * Unit tests for GraphFrameInternals.extractColumnReferences. diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameSuite.scala index 159bbf88c6304..3cdf9e7c5fc65 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameSuite.scala @@ -17,8 +17,12 @@ package org.apache.spark.graphframes -import org.apache.commons.io.FileUtils +import java.io.File +import java.nio.file.Files + import org.apache.hadoop.fs.Path + +import org.apache.spark.graphframes.examples.Graphs import org.apache.spark.graphx.Edge import org.apache.spark.graphx.Graph import org.apache.spark.rdd.RDD @@ -31,10 +35,7 @@ import org.apache.spark.sql.types.StringType import org.apache.spark.sql.types.StructField import org.apache.spark.sql.types.StructType import org.apache.spark.storage.StorageLevel -import org.apache.spark.graphframes.examples.Graphs - -import java.io.File -import java.nio.file.Files +import org.apache.spark.util.SparkFileUtils class GraphFrameSuite extends SparkFunSuite with GraphFrameTestSparkContext { @@ -59,7 +60,7 @@ class GraphFrameSuite extends SparkFunSuite with GraphFrameTestSparkContext { } override def afterAll(): Unit = { - FileUtils.deleteQuietly(tempDir) + SparkFileUtils.deleteQuietly(tempDir) super.afterAll() } diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameTestSparkContext.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameTestSparkContext.scala index 06dddd1387707..5a1a213ab0540 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameTestSparkContext.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/GraphFrameTestSparkContext.scala @@ -17,16 +17,17 @@ package org.apache.spark.graphframes -import org.apache.commons.io.FileUtils -import org.apache.spark.SparkContext -import org.apache.spark.sql.SQLContext -import org.apache.spark.sql.SQLImplicits -import org.apache.spark.sql.SparkSession +import java.io.File +import java.nio.file.Files + import org.scalatest.BeforeAndAfterAll import org.scalatest.Suite -import java.io.File -import java.nio.file.Files +import org.apache.spark.SparkContext +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.SQLContext +import org.apache.spark.sql.SQLImplicits +import org.apache.spark.util.SparkFileUtils trait GraphFrameTestSparkContext extends BeforeAndAfterAll { self: Suite => @transient var spark: SparkSession = _ @@ -79,7 +80,7 @@ trait GraphFrameTestSparkContext extends BeforeAndAfterAll { self: Suite => sc = null checkpointDir.foreach { dir => - FileUtils.deleteQuietly(new File(dir)) + SparkFileUtils.deleteQuietly(new File(dir)) } super.afterAll() } diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/PatternMatchSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/PatternMatchSuite.scala index 47d2700c3db0b..9fb6321aac204 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/PatternMatchSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/PatternMatchSuite.scala @@ -71,12 +71,14 @@ class PatternMatchSuite extends SparkFunSuite with GraphFrameTestSparkContext { private def compareResultToExpected[A](result: Set[A], expected: Set[A]): Unit = { if (result !== expected) { + // scalastyle:off throwerror throw new AssertionError( "result !== expected.\n" + s"Result contained additional values: ${result.diff(expected)}\n" + s"Expected contained additional values: ${expected.diff(result)}\n" + s"Result: $result\n" + s"Expected: $expected") + // scalastyle:on throwerror } } diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/SparkFunSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/SparkFunSuite.scala index b0ed7c370a48a..6055a550a6f85 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/SparkFunSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/SparkFunSuite.scala @@ -17,29 +17,7 @@ package org.apache.spark.graphframes -import org.scalatest.Outcome -import org.scalatest.funsuite.AnyFunSuite - /** * Base abstract class for all unit tests in Spark for handling common functionality. */ -abstract class SparkFunSuite extends AnyFunSuite with Logging { - - /** - * Log the suite name and the test name before and after each test. - * - * Subclasses should never override this method. If they wish to run custom code before and - * after each test, they should mix in the {{org.scalatest.BeforeAndAfter}} trait instead. - */ - final protected override def withFixture(test: NoArgTest): Outcome = { - val testName = test.text - val suiteName = this.getClass.getName - val shortSuiteName = suiteName.replaceAll("org.apache.spark", "o.a.s") - try { - logInfo(s"\n\n===== TEST OUTPUT FOR $shortSuiteName: '$testName' =====\n") - test() - } finally { - logInfo(s"\n\n===== FINISHED $shortSuiteName: '$testName' =====\n") - } - } -} +abstract class SparkFunSuite extends org.apache.spark.SparkFunSuite diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/TestUtils.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/TestUtils.scala index 1302c424561c1..e7b034a8abc75 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/TestUtils.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/TestUtils.scala @@ -17,10 +17,10 @@ package org.apache.spark.graphframes +import org.apache.spark.graphframes.GraphFrame._ import org.apache.spark.sql.DataFrame import org.apache.spark.sql.types.DataType import org.apache.spark.sql.types.StructType -import org.apache.spark.graphframes.GraphFrame._ object TestUtils { diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/convolutions/SamplingConvolutionSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/convolutions/SamplingConvolutionSuite.scala index 4ddb1dd395957..88c2b0ee1bf45 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/convolutions/SamplingConvolutionSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/convolutions/SamplingConvolutionSuite.scala @@ -17,15 +17,16 @@ package org.apache.spark.graphframes.convolutions +import org.scalatest.BeforeAndAfterAll + +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite import org.apache.spark.ml.linalg.DenseVector import org.apache.spark.ml.linalg.Vector import org.apache.spark.ml.linalg.Vectors import org.apache.spark.sql.DataFrame import org.apache.spark.sql.functions._ -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFrameTestSparkContext -import org.apache.spark.graphframes.SparkFunSuite -import org.scalatest.BeforeAndAfterAll class SamplingConvolutionSuite extends SparkFunSuite diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/embeddings/Hash2VecSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/embeddings/Hash2VecSuite.scala index daaaca149f866..70108b56e34e4 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/embeddings/Hash2VecSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/embeddings/Hash2VecSuite.scala @@ -17,17 +17,18 @@ package org.apache.spark.graphframes.embeddings +import scala.util.Random + +import org.scalatest.BeforeAndAfterAll + +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite import org.apache.spark.ml.linalg.DenseVector import org.apache.spark.ml.linalg.SQLDataTypes.VectorType import org.apache.spark.sql.DataFrame import org.apache.spark.sql.functions.col import org.apache.spark.sql.types.LongType import org.apache.spark.sql.types.StringType -import org.apache.spark.graphframes.GraphFrameTestSparkContext -import org.apache.spark.graphframes.SparkFunSuite -import org.scalatest.BeforeAndAfterAll - -import scala.util.Random class Hash2VecSuite extends SparkFunSuite with GraphFrameTestSparkContext with BeforeAndAfterAll { private var longSequences: DataFrame = _ @@ -232,8 +233,8 @@ class Hash2VecSuite extends SparkFunSuite with GraphFrameTestSparkContext with B } } - test("Hash2Vec - cosine distances reflect co‑occurrence patterns") { - // Create a tiny dataset where some words co‑occur often, others rarely. + test("Hash2Vec - cosine distances reflect co-occurrence patterns") { + // Create a tiny dataset where some words co-occur often, others rarely. // We'll use string sequences for simplicity. val sequences = Seq( Seq("apple", "banana", "apple", "cherry", "banana"), @@ -284,46 +285,46 @@ class Hash2VecSuite extends SparkFunSuite with GraphFrameTestSparkContext with B dot / (math.sqrt(norm1) * math.sqrt(norm2)) } - // apple and banana co‑occur very frequently → high similarity + // apple and banana co-occur very frequently -> high similarity val appleBananaSim = cosineSimilarity(embMap("apple"), embMap("banana")) - // cherry and date also co‑occur frequently (in the fourth sequence) + // cherry and date also co-occur frequently (in the fourth sequence) val cherryDateSim = cosineSimilarity(embMap("cherry"), embMap("date")) - // apple and fig almost never appear together → low similarity + // apple and fig almost never appear together -> low similarity val appleFigSim = cosineSimilarity(embMap("apple"), embMap("fig")) // banana and fig also rarely together val bananaFigSim = cosineSimilarity(embMap("banana"), embMap("fig")) - // elderberry and fig co‑occur (sixth sequence) + // elderberry and fig co-occur (sixth sequence) val elderberryFigSim = cosineSimilarity(embMap("elderberry"), embMap("fig")) - // Assert ordering of similarities matches expected co‑occurrence patterns - // apple‑banana should be among the highest similarities - assert(appleBananaSim > 0.3, s"apple‑banana similarity $appleBananaSim should be > 0.3") - // apple‑fig should be low (close to zero or negative) + // Assert ordering of similarities matches expected co-occurrence patterns + // apple-banana should be among the highest similarities + assert(appleBananaSim > 0.3, s"apple-banana similarity $appleBananaSim should be > 0.3") + // apple-fig should be low (close to zero or negative) assert( appleFigSim < appleBananaSim, - s"apple‑fig ($appleFigSim) should be < apple‑banana ($appleBananaSim)") + s"apple-fig ($appleFigSim) should be < apple-banana ($appleBananaSim)") assert( bananaFigSim < appleBananaSim, - s"banana‑fig ($bananaFigSim) should be < apple‑banana ($appleBananaSim)") - // cherry‑date similarity should be relatively high (they co‑occur exclusively) - assert(cherryDateSim > 0.2, s"cherry‑date similarity $cherryDateSim should be > 0.2") - // elderberry‑fig should be higher than apple‑fig (because they co‑occur) + s"banana-fig ($bananaFigSim) should be < apple-banana ($appleBananaSim)") + // cherry-date similarity should be relatively high (they co-occur exclusively) + assert(cherryDateSim > 0.2, s"cherry-date similarity $cherryDateSim should be > 0.2") + // elderberry-fig should be higher than apple-fig (because they co-occur) assert( elderberryFigSim > appleFigSim, - s"elderberry‑fig ($elderberryFigSim) should be > apple‑fig ($appleFigSim)") + s"elderberry-fig ($elderberryFigSim) should be > apple-fig ($appleFigSim)") - // Self‑similarity should be 1.0 (or close after normalization) + // Self-similarity should be 1.0 (or close after normalization) val appleSelf = cosineSimilarity(embMap("apple"), embMap("apple")) - assert(math.abs(appleSelf - 1.0) < 1e-6, s"self‑similarity should be ~1.0, got $appleSelf") + assert(math.abs(appleSelf - 1.0) < 1e-6, s"self-similarity should be ~1.0, got $appleSelf") } - test("Hash2Vec - long‑typed co‑occurrence") { + test("Hash2Vec - long-typed co-occurrence") { // Use numeric ids to test long sequences. val sequences = Seq( - Seq(1L, 2L, 1L, 3L, 2L), // 1‑2 frequent, 3 appears with 2 + Seq(1L, 2L, 1L, 3L, 2L), // 1-2 frequent, 3 appears with 2 Seq(1L, 2L, 3L, 2L), Seq(1L, 2L, 1L, 2L, 2L), - Seq(3L, 4L, 3L, 4L), // 3‑4 frequent pair + Seq(3L, 4L, 3L, 4L), // 3-4 frequent pair Seq(4L, 5L, 4L), Seq(5L, 6L, 5L), Seq(6L, 6L, 6L)) @@ -370,19 +371,19 @@ class Hash2VecSuite extends SparkFunSuite with GraphFrameTestSparkContext with B val sim16 = cosineSimilarity(embMap(1L), embMap(6L)) val sim56 = cosineSimilarity(embMap(5L), embMap(6L)) - // 1‑2 co‑occur very often - assert(sim12 > 0.3, s"1‑2 similarity $sim12 should be > 0.3") - // 1‑3 appear together less often than 1‑2 - assert(sim13 < sim12, s"1‑3 ($sim13) should be < 1‑2 ($sim12)") - // 3‑4 are exclusive pair - assert(sim34 > 0.2, s"3‑4 similarity $sim34 should be > 0.2") + // 1-2 co-occur very often + assert(sim12 > 0.3, s"1-2 similarity $sim12 should be > 0.3") + // 1-3 appear together less often than 1-2 + assert(sim13 < sim12, s"1-3 ($sim13) should be < 1-2 ($sim12)") + // 3-4 are exclusive pair + assert(sim34 > 0.2, s"3-4 similarity $sim34 should be > 0.2") // 1 and 6 almost never together - assert(sim16 < 0.1, s"1‑6 similarity $sim16 should be near zero") - // 5‑6 co‑occur in a sequence - assert(sim56 > sim16, s"5‑6 ($sim56) should be > 1‑6 ($sim16)") + assert(sim16 < 0.1, s"1-6 similarity $sim16 should be near zero") + // 5-6 co-occur in a sequence + assert(sim56 > sim16, s"5-6 ($sim56) should be > 1-6 ($sim16)") // Self similarity val self = cosineSimilarity(embMap(1L), embMap(1L)) - assert(math.abs(self - 1.0) < 1e-6, s"self‑similarity should be ~1.0, got $self") + assert(math.abs(self - 1.0) < 1e-6, s"self-similarity should be ~1.0, got $self") } } diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateMessagesSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateMessagesSuite.scala index 5ee33cd5530a9..91e44d045d870 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateMessagesSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateMessagesSuite.scala @@ -17,17 +17,17 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.Row -import org.apache.spark.sql.functions._ -import org.apache.spark.sql.types._ +import scala.collection.mutable + import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFrameTestSparkContext import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.graphframes.GraphFrameTestSparkContext import org.apache.spark.graphframes.SparkFunSuite import org.apache.spark.graphframes.TestUtils import org.apache.spark.graphframes.examples.Graphs - -import scala.collection.mutable +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types._ class AggregateMessagesSuite extends SparkFunSuite with GraphFrameTestSparkContext { diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateNeighborsSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateNeighborsSuite.scala index ccaa45d96ef0b..d832b09b28382 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateNeighborsSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AggregateNeighborsSuite.scala @@ -17,10 +17,10 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.functions._ import org.apache.spark.graphframes.GraphFrame import org.apache.spark.graphframes.GraphFrameTestSparkContext import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.sql.functions._ class AggregateNeighborsSuite extends SparkFunSuite with GraphFrameTestSparkContext { diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AllPathsSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AllPathsSuite.scala index 63abcbfdad7cd..1347eaa686e5e 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AllPathsSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/AllPathsSuite.scala @@ -17,13 +17,13 @@ package org.apache.spark.graphframes.lib +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite import org.apache.spark.sql.functions.col import org.apache.spark.sql.types.ArrayType import org.apache.spark.sql.types.IntegerType import org.apache.spark.sql.types.LongType -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFrameTestSparkContext -import org.apache.spark.graphframes.SparkFunSuite class AllPathsSuite extends SparkFunSuite with GraphFrameTestSparkContext { diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/BFSSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/BFSSuite.scala index c67204fffe70f..31ae3719d3258 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/BFSSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/BFSSuite.scala @@ -17,13 +17,13 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.DataFrame -import org.apache.spark.sql.Row -import org.apache.spark.sql.functions.col import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFrameTestSparkContext import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.graphframes.GraphFrameTestSparkContext import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions.col class BFSSuite extends SparkFunSuite with GraphFrameTestSparkContext { diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ConnectedComponentsSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ConnectedComponentsSuite.scala index 78f9c1602b8ea..39cd8ccba1563 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ConnectedComponentsSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ConnectedComponentsSuite.scala @@ -17,6 +17,12 @@ package org.apache.spark.graphframes.lib +import scala.reflect.ClassTag +import scala.reflect.runtime.universe.TypeTag + +import org.apache.spark.graphframes._ +import org.apache.spark.graphframes.GraphFrame._ +import org.apache.spark.graphframes.examples.Graphs import org.apache.spark.sql.DataFrame import org.apache.spark.sql.Row import org.apache.spark.sql.functions.col @@ -24,12 +30,6 @@ import org.apache.spark.sql.functions.lit import org.apache.spark.sql.types.DataTypes import org.apache.spark.sql.types.LongType import org.apache.spark.storage.StorageLevel -import org.apache.spark.graphframes._ -import org.apache.spark.graphframes.GraphFrame._ -import org.apache.spark.graphframes.examples.Graphs - -import scala.reflect.ClassTag -import scala.reflect.runtime.universe.TypeTag class ConnectedComponentsSuite extends SparkFunSuite with GraphFrameTestSparkContext { diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/DetectingCyclesSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/DetectingCyclesSuite.scala index 1a72c9d3223e0..cbcc1be0fbd4d 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/DetectingCyclesSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/DetectingCyclesSuite.scala @@ -17,13 +17,13 @@ package org.apache.spark.graphframes.lib +import scala.annotation.nowarn +import scala.collection.mutable + import org.apache.spark.graphframes.GraphFrame import org.apache.spark.graphframes.GraphFrameTestSparkContext import org.apache.spark.graphframes.SparkFunSuite -import scala.annotation.nowarn -import scala.collection.mutable - class DetectingCyclesSuite extends SparkFunSuite with GraphFrameTestSparkContext { test("test detecting cycles") { val graph = GraphFrame( diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/HyperANFSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/HyperANFSuite.scala index 6d6e5f4070580..182de1546aee7 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/HyperANFSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/HyperANFSuite.scala @@ -17,15 +17,15 @@ package org.apache.spark.graphframes.lib +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.TestUtils import org.apache.spark.sql.DataFrame import org.apache.spark.sql.functions.col import org.apache.spark.sql.functions.expr import org.apache.spark.sql.functions.hll_sketch_estimate import org.apache.spark.sql.types.DataTypes -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFrameTestSparkContext -import org.apache.spark.graphframes.SparkFunSuite -import org.apache.spark.graphframes.TestUtils class HyperANFSuite extends SparkFunSuite with GraphFrameTestSparkContext { @@ -99,8 +99,8 @@ class HyperANFSuite extends SparkFunSuite with GraphFrameTestSparkContext { result.unpersist() } - test( - "HyperANF starting vertices expression limits output to selected vertices with outgoing edges") { + test("HyperANF starting vertices expression limits output to selected vertices " + + "with outgoing edges") { val graph = diamondCycleGraph val result = new HyperANF(graph) .setEdgesFilterExpression(expr("src IN (1, 3, 42)")) diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/KCoreSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/KCoreSuite.scala index 71b4a7b631b44..4ce11a7cfbf60 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/KCoreSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/KCoreSuite.scala @@ -17,9 +17,9 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.types.DataTypes import org.apache.spark.graphframes._ import org.apache.spark.graphframes.examples.Graphs +import org.apache.spark.sql.types.DataTypes class KCoreSuite extends SparkFunSuite with GraphFrameTestSparkContext { test("empty graph") { @@ -310,8 +310,8 @@ class KCoreSuite extends SparkFunSuite with GraphFrameTestSparkContext { test("triangle with tail - exact kcore values") { // This graph has vertices where degree != kcore, which is important to test correctness: - // it would catch a buggy implementation that converges too early (e.g. after one superstep), which - // would return kcore = degree for all vertices and pass simpler tests. + // it would catch a buggy implementation that converges too early (e.g. after one superstep), + // which would return kcore = degree for all vertices and pass simpler tests. // // Undirected graph: // diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/LabelPropagationSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/LabelPropagationSuite.scala index c2036ee7b64dc..c4883c9c611d2 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/LabelPropagationSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/LabelPropagationSuite.scala @@ -17,11 +17,11 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.types.DataTypes import org.apache.spark.graphframes.GraphFrameTestSparkContext import org.apache.spark.graphframes.SparkFunSuite import org.apache.spark.graphframes.TestUtils import org.apache.spark.graphframes.examples.Graphs +import org.apache.spark.sql.types.DataTypes class LabelPropagationSuite extends SparkFunSuite with GraphFrameTestSparkContext { diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/MaximalIndependentSetSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/MaximalIndependentSetSuite.scala index 335d114562b53..da48c0cab3b1f 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/MaximalIndependentSetSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/MaximalIndependentSetSuite.scala @@ -17,10 +17,10 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.DataFrame -import org.apache.spark.sql.functions.col import org.apache.spark.graphframes._ import org.apache.spark.graphframes.examples.Graphs +import org.apache.spark.sql.DataFrame +import org.apache.spark.sql.functions.col class MaximalIndependentSetSuite extends SparkFunSuite with GraphFrameTestSparkContext { test("isolated vertices should be included in MIS") { diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/PageRankSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/PageRankSuite.scala index ba70519d1e172..224b55def32bd 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/PageRankSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/PageRankSuite.scala @@ -17,12 +17,12 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.functions.col -import org.apache.spark.sql.types.DataTypes import org.apache.spark.graphframes.GraphFrameTestSparkContext import org.apache.spark.graphframes.SparkFunSuite import org.apache.spark.graphframes.TestUtils import org.apache.spark.graphframes.examples.Graphs +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.types.DataTypes class PageRankSuite extends SparkFunSuite with GraphFrameTestSparkContext { diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRankSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRankSuite.scala index 0976ad97e9836..9d954926a9cdd 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRankSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ParallelPersonalizedPageRankSuite.scala @@ -16,15 +16,16 @@ */ package org.apache.spark.graphframes.lib -import org.apache.spark.ml.linalg.SQLDataTypes -import org.apache.spark.ml.linalg.SparseVector -import org.apache.spark.sql.Row -import org.apache.spark.sql.functions.col -import org.apache.spark.sql.types.DataTypes + import org.apache.spark.graphframes.GraphFrameTestSparkContext import org.apache.spark.graphframes.SparkFunSuite import org.apache.spark.graphframes.TestUtils import org.apache.spark.graphframes.examples.Graphs +import org.apache.spark.ml.linalg.SparseVector +import org.apache.spark.ml.linalg.SQLDataTypes +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.types.DataTypes class ParallelPersonalizedPageRankSuite extends SparkFunSuite with GraphFrameTestSparkContext { @@ -87,7 +88,8 @@ class ParallelPersonalizedPageRankSuite extends SparkFunSuite with GraphFrameTes } assert( prInvalid.size === 0, - s"found ${prInvalid.size} entries with invalid number of returned personalized pagerank vector") + s"found ${prInvalid.size} entries with invalid number of returned personalized " + + "pagerank vector") val gRank = pr.vertices .filter(col("id") === "g") @@ -96,7 +98,8 @@ class ParallelPersonalizedPageRankSuite extends SparkFunSuite with GraphFrameTes .getAs[SparseVector](0) assert( gRank.numNonzeros === 0, - s"User g (Gabby) doesn't connect with a. So its pagerank should be 0 but we got ${gRank.numNonzeros}.") + "User g (Gabby) doesn't connect with a. So its pagerank should be 0 but we got " + + s"${gRank.numNonzeros}.") pr.unpersist() } } diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/PregelSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/PregelSuite.scala index e9b80471ea593..7ed6aac7555bc 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/PregelSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/PregelSuite.scala @@ -17,10 +17,11 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.functions._ -import org.apache.spark.graphframes._ import org.scalactic.Tolerance._ +import org.apache.spark.graphframes._ +import org.apache.spark.sql.functions._ + class PregelSuite extends SparkFunSuite with GraphFrameTestSparkContext { import sqlImplicits._ diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/RandomizedContractionSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/RandomizedContractionSuite.scala index 3a51cc815da50..b2f7bb89d9eb8 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/RandomizedContractionSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/RandomizedContractionSuite.scala @@ -17,15 +17,15 @@ package org.apache.spark.graphframes.lib +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.examples.Graphs import org.apache.spark.sql.Row import org.apache.spark.sql.catalyst.FunctionIdentifier import org.apache.spark.sql.functions.col import org.apache.spark.sql.functions.lit import org.apache.spark.storage.StorageLevel -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFrameTestSparkContext -import org.apache.spark.graphframes.SparkFunSuite -import org.apache.spark.graphframes.examples.Graphs class RandomizedContractionSuite extends SparkFunSuite with GraphFrameTestSparkContext { @@ -273,9 +273,9 @@ class RandomizedContractionSuite extends SparkFunSuite with GraphFrameTestSparkC } private def listParquetFiles(): Set[String] = { - val hadoopConf = spark.sparkContext.hadoopConfiguration - val fs = org.apache.hadoop.fs.FileSystem.get(hadoopConf) + val hadoopConf = spark.sessionState.newHadoopConf() val rootPath = new org.apache.hadoop.fs.Path(spark.conf.get("spark.sql.warehouse.dir")) + val fs = rootPath.getFileSystem(hadoopConf) def listFiles(path: org.apache.hadoop.fs.Path): Set[String] = { if (fs.exists(path)) { diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/SVDPlusPlusSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/SVDPlusPlusSuite.scala index c4b4957548ffb..a8a4034af783d 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/SVDPlusPlusSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/SVDPlusPlusSuite.scala @@ -17,15 +17,15 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.Row -import org.apache.spark.sql.functions.col -import org.apache.spark.sql.types.DataTypes import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFrameTestSparkContext import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.graphframes.GraphFrameTestSparkContext import org.apache.spark.graphframes.SparkFunSuite import org.apache.spark.graphframes.TestUtils import org.apache.spark.graphframes.examples.Graphs +import org.apache.spark.sql.Row +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.types.DataTypes class SVDPlusPlusSuite extends SparkFunSuite with GraphFrameTestSparkContext { diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ShortestPathsSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ShortestPathsSuite.scala index 7c5a708a6ed7d..9acd7c9976b2f 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ShortestPathsSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/ShortestPathsSuite.scala @@ -17,12 +17,12 @@ package org.apache.spark.graphframes.lib +import org.apache.spark.graphframes._ +import org.apache.spark.graphframes.GraphFrame.quote import org.apache.spark.sql.DataFrame import org.apache.spark.sql.Row import org.apache.spark.sql.functions.col import org.apache.spark.sql.types.DataTypes -import org.apache.spark.graphframes._ -import org.apache.spark.graphframes.GraphFrame.quote class ShortestPathsSuite extends SparkFunSuite with GraphFrameTestSparkContext { diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponentsSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponentsSuite.scala index b7c9811754c09..e4f19b4301319 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponentsSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/StronglyConnectedComponentsSuite.scala @@ -17,12 +17,12 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.Row -import org.apache.spark.sql.types.DataTypes import org.apache.spark.graphframes.GraphFrame import org.apache.spark.graphframes.GraphFrameTestSparkContext import org.apache.spark.graphframes.SparkFunSuite import org.apache.spark.graphframes.TestUtils +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.DataTypes class StronglyConnectedComponentsSuite extends SparkFunSuite with GraphFrameTestSparkContext { test("Island Strongly Connected Components") { diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala index 6c12ff4be9e3b..8c63f3ff4dacf 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/StructureAwareLabelPropagation.scala @@ -17,13 +17,13 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.functions.col -import org.apache.spark.sql.functions.lit -import org.apache.spark.sql.types.DataTypes import org.apache.spark.graphframes.GraphFrame import org.apache.spark.graphframes.GraphFrameTestSparkContext import org.apache.spark.graphframes.SparkFunSuite import org.apache.spark.graphframes.TestUtils +import org.apache.spark.sql.functions.col +import org.apache.spark.sql.functions.lit +import org.apache.spark.sql.types.DataTypes class StructureAwareLabelPropagationSuite extends SparkFunSuite with GraphFrameTestSparkContext { @@ -56,8 +56,8 @@ class StructureAwareLabelPropagationSuite extends SparkFunSuite with GraphFrameT result.unpersist() } - test( - "different structuralSimilarityMultiplier values can change winner between direct-link mass and common-neighbor overlap") { + test("different structuralSimilarityMultiplier values can change winner between direct-link " + + "mass and common-neighbor overlap") { assume(TestUtils.requireSparkVersionGE(4, 1, spark.version)) val vertices = spark diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/TriangleCountSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/TriangleCountSuite.scala index 5147e9844d8d2..c117bb7d8a868 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/lib/TriangleCountSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/lib/TriangleCountSuite.scala @@ -17,14 +17,14 @@ package org.apache.spark.graphframes.lib -import org.apache.spark.sql.Row -import org.apache.spark.sql.types.DataTypes import org.apache.spark.graphframes.GraphFrame import org.apache.spark.graphframes.GraphFrame.quote -import org.apache.spark.graphframes.GraphFrameTestSparkContext import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.graphframes.GraphFrameTestSparkContext import org.apache.spark.graphframes.SparkFunSuite import org.apache.spark.graphframes.TestUtils +import org.apache.spark.sql.Row +import org.apache.spark.sql.types.DataTypes class TriangleCountSuite extends SparkFunSuite with GraphFrameTestSparkContext { diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/pattern/PatternSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/pattern/PatternSuite.scala index 7c1c9b200c57e..d23ccbad6728c 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/pattern/PatternSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/pattern/PatternSuite.scala @@ -199,7 +199,8 @@ class PatternSuite extends SparkFunSuite { } } withClue( - "Failed to catch parse error with completely anonymous negated and undirected edge !()-[]-()") { + "Failed to catch parse error with completely anonymous negated and undirected " + + "edge !()-[]-()") { intercept[InvalidParseException] { Pattern.parse("!()-[]-()") } diff --git a/graphframes/src/test/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestartSuite.scala b/graphframes/src/test/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestartSuite.scala index f0bbe33860ecf..4aa1916e09cb3 100644 --- a/graphframes/src/test/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestartSuite.scala +++ b/graphframes/src/test/scala/org/apache/spark/graphframes/rw/RandomWalkWithRestartSuite.scala @@ -17,14 +17,16 @@ package org.apache.spark.graphframes.rw +import org.apache.hadoop.fs.Path + +import org.apache.spark.graphframes.GraphFrameTestSparkContext +import org.apache.spark.graphframes.SparkFunSuite +import org.apache.spark.graphframes.examples.Graphs import org.apache.spark.sql.functions.array_size import org.apache.spark.sql.functions.col import org.apache.spark.sql.functions.lit import org.apache.spark.sql.types.ArrayType import org.apache.spark.sql.types.StringType -import org.apache.spark.graphframes.GraphFrameTestSparkContext -import org.apache.spark.graphframes.SparkFunSuite -import org.apache.spark.graphframes.examples.Graphs class RandomWalkWithRestartSuite extends SparkFunSuite with GraphFrameTestSparkContext { test("test RW base") { @@ -65,12 +67,14 @@ class RandomWalkWithRestartSuite extends SparkFunSuite with GraphFrameTestSparkC rwRunner.cleanUp() // Verify that all temporary files have been deleted + // scalastyle:off hadoopconfiguration val hadoopConf = spark.sparkContext.hadoopConfiguration - val fs = org.apache.hadoop.fs.FileSystem.get(hadoopConf) + // scalastyle:on hadoopconfiguration val basePath = "/tmp" + val fs = new Path(basePath).getFileSystem(hadoopConf) val runPath = s"$basePath/${runId}_batch_" (1 to numBatches).foreach { i => - val path = new org.apache.hadoop.fs.Path(s"${runPath}${i}") + val path = new Path(s"${runPath}${i}") assert(!fs.exists(path), s"Temporary file not deleted: $path") } } @@ -94,7 +98,7 @@ class RandomWalkWithRestartSuite extends SparkFunSuite with GraphFrameTestSparkC .setTemporaryPrefix("/tmp") val runId = rwRunner1.getRunId() - println(s"Using runId: $runId") + logInfo(s"Using runId: $runId") // Run and persist the result val walks1 = rwRunner1.run() @@ -162,12 +166,14 @@ class RandomWalkWithRestartSuite extends SparkFunSuite with GraphFrameTestSparkC rwRunner2.cleanUp() // Verify cleanup + // scalastyle:off hadoopconfiguration val hadoopConf = spark.sparkContext.hadoopConfiguration - val fs = org.apache.hadoop.fs.FileSystem.get(hadoopConf) + // scalastyle:on hadoopconfiguration val basePath = "/tmp" + val fs = new Path(basePath).getFileSystem(hadoopConf) val runPath = s"$basePath/${runId}_batch_" (1 to numBatches).foreach { i => - val path = new org.apache.hadoop.fs.Path(s"${runPath}${i}") + val path = new Path(s"${runPath}${i}") assert(!fs.exists(path), s"Temporary file not deleted: $path") } diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala index 29dd92f7fa0bb..a18166e5848d8 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala @@ -71,7 +71,6 @@ import org.apache.spark.sql.execution.QueryExecution import org.apache.spark.sql.execution.aggregate.{ScalaAggregator, TypedAggregateExpression} import org.apache.spark.sql.execution.arrow.ArrowConverters import org.apache.spark.sql.execution.command.{CreateViewCommand, ExternalCommandExecutor} -import org.apache.spark.sql.graphframes.{GraphFrameInternals, GraphFramesConnectUtils} import org.apache.spark.sql.execution.datasources.jdbc.JDBCOptions import org.apache.spark.sql.execution.datasources.v2.python.UserDefinedPythonDataSource import org.apache.spark.sql.execution.python.{UserDefinedPythonFunction, UserDefinedPythonTableFunction} @@ -80,6 +79,7 @@ import org.apache.spark.sql.execution.stat.StatFunctions import org.apache.spark.sql.execution.streaming.operators.stateful.flatmapgroupswithstate.GroupStateImpl.groupStateTimeoutFromString import org.apache.spark.sql.execution.streaming.runtime.StreamingQueryWrapper import org.apache.spark.sql.expressions.{Aggregator, ReduceAggregator, SparkUserDefinedFunction, UserDefinedAggregator, UserDefinedFunction} +import org.apache.spark.sql.graphframes.{GraphFrameInternals, GraphFramesConnectUtils} import org.apache.spark.sql.streaming.{GroupStateTimeout, OutputMode, StatefulProcessor, StatefulProcessorWithInitialState, StreamingQuery, StreamingQueryListener, StreamingQueryProgress, Trigger} import org.apache.spark.sql.types._ import org.apache.spark.sql.util.{ArrowUtils, CaseInsensitiveStringMap} diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala index c31369565aa45..a164b2a72ba77 100644 --- a/sql/connect/server/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala +++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala @@ -19,19 +19,20 @@ // Same about a Column helper object. package org.apache.spark.sql.graphframes +import scala.jdk.CollectionConverters._ + import com.google.protobuf.ByteString + +import org.apache.spark.connect.proto.{graphframes => proto} +import org.apache.spark.graphframes.GraphFrame +import org.apache.spark.graphframes.GraphFramesUnreachableException +import org.apache.spark.graphframes.embeddings.RandomWalkEmbeddings import org.apache.spark.sql.Column import org.apache.spark.sql.DataFrame import org.apache.spark.sql.connect.planner.SparkConnectPlanner import org.apache.spark.sql.functions.expr import org.apache.spark.sql.functions.lit import org.apache.spark.storage.StorageLevel -import org.apache.spark.graphframes.GraphFrame -import org.apache.spark.graphframes.GraphFramesUnreachableException -import org.apache.spark.connect.proto.{graphframes => proto} -import org.apache.spark.graphframes.embeddings.RandomWalkEmbeddings - -import scala.jdk.CollectionConverters._ /** * Utility object providing helper methods for parsing and transforming data structures related to @@ -173,7 +174,7 @@ object GraphFramesConnectUtils { val graphFrame = extractGraphFrame(apiMessage, planner) apiMessage.getMethodCase match { - case proto.GraphFramesAPI.MethodCase.AGGREGATE_MESSAGES => { + case proto.GraphFramesAPI.MethodCase.AGGREGATE_MESSAGES => val aggregateMessagesProto = apiMessage.getAggregateMessages var aggregateMessages = graphFrame.aggregateMessages if (aggregateMessagesProto.getSendToDstList.size() == 1) { @@ -209,8 +210,8 @@ object GraphFramesConnectUtils { } else { aggregateMessages.agg(aggCols.head, aggCols.tail.toSeq: _*) } - } - case proto.GraphFramesAPI.MethodCase.BFS => { + + case proto.GraphFramesAPI.MethodCase.BFS => val bfsProto = apiMessage.getBfs graphFrame.bfs .toExpr(parseColumnOrExpression(bfsProto.getToExpr, planner)) @@ -218,8 +219,8 @@ object GraphFramesConnectUtils { .edgeFilter(parseColumnOrExpression(bfsProto.getEdgeFilter, planner)) .maxPathLength(bfsProto.getMaxPathLength) .run() - } - case proto.GraphFramesAPI.MethodCase.ALL_PATHS => { + + case proto.GraphFramesAPI.MethodCase.ALL_PATHS => val allPathsProto = apiMessage.getAllPaths var allPaths = graphFrame.allPaths .toExpr(parseColumnOrExpression(allPathsProto.getToExpr, planner)) @@ -234,8 +235,8 @@ object GraphFramesConnectUtils { allPaths.setIntermediateStorageLevel(parseStorageLevel(allPathsProto.getStorageLevel)) } allPaths.run() - } - case proto.GraphFramesAPI.MethodCase.CONNECTED_COMPONENTS => { + + case proto.GraphFramesAPI.MethodCase.CONNECTED_COMPONENTS => val cc = apiMessage.getConnectedComponents val ccBuilder = graphFrame.connectedComponents .maxIter(cc.getMaxIter) @@ -250,9 +251,8 @@ object GraphFramesConnectUtils { } else { ccBuilder.run() } - } - case proto.GraphFramesAPI.MethodCase.DETECTING_CYCLES => { + case proto.GraphFramesAPI.MethodCase.DETECTING_CYCLES => val dc = apiMessage.getDetectingCycles val dcBuilder = graphFrame.detectingCycles .setCheckpointInterval(dc.getCheckpointInterval) @@ -262,24 +262,23 @@ object GraphFramesConnectUtils { } else { dcBuilder.run() } - } - case proto.GraphFramesAPI.MethodCase.DROP_ISOLATED_VERTICES => { + case proto.GraphFramesAPI.MethodCase.DROP_ISOLATED_VERTICES => graphFrame.dropIsolatedVertices().vertices - } - case proto.GraphFramesAPI.MethodCase.FILTER_EDGES => { + + case proto.GraphFramesAPI.MethodCase.FILTER_EDGES => val condition = parseColumnOrExpression(apiMessage.getFilterEdges.getCondition, planner) graphFrame.filterEdges(condition).edges - } - case proto.GraphFramesAPI.MethodCase.FILTER_VERTICES => { + + case proto.GraphFramesAPI.MethodCase.FILTER_VERTICES => val condition = parseColumnOrExpression(apiMessage.getFilterVertices.getCondition, planner) graphFrame.filterVertices(condition).vertices - } - case proto.GraphFramesAPI.MethodCase.FIND => { + + case proto.GraphFramesAPI.MethodCase.FIND => graphFrame.find(apiMessage.getFind.getPattern) - } - case proto.GraphFramesAPI.MethodCase.LABEL_PROPAGATION => { + + case proto.GraphFramesAPI.MethodCase.LABEL_PROPAGATION => val lp = apiMessage.getLabelPropagation val lpBuilder = graphFrame.labelPropagation .maxIter(lp.getMaxIter) @@ -292,8 +291,8 @@ object GraphFramesConnectUtils { } else { lpBuilder.run() } - } - case proto.GraphFramesAPI.MethodCase.NEIGHBORHOOD_AWARE_CDLP => { + + case proto.GraphFramesAPI.MethodCase.NEIGHBORHOOD_AWARE_CDLP => val nc = apiMessage.getNeighborhoodAwareCdlp val ncBuilder = graphFrame.structureAwareLabelPropagation .maxIter(nc.getMaxIter) @@ -313,8 +312,8 @@ object GraphFramesConnectUtils { } else { ncBuilder.run() } - } - case proto.GraphFramesAPI.MethodCase.PAGE_RANK => { + + case proto.GraphFramesAPI.MethodCase.PAGE_RANK => val pageRankProto = apiMessage.getPageRank val pageRank = graphFrame.pageRank.resetProbability(pageRankProto.getResetProbability) @@ -332,8 +331,8 @@ object GraphFramesConnectUtils { // TODO: do we really need an edge weights in that case? // see comments in the Python API pageRank.run().vertices - } - case proto.GraphFramesAPI.MethodCase.PARALLEL_PERSONALIZED_PAGE_RANK => { + + case proto.GraphFramesAPI.MethodCase.PARALLEL_PERSONALIZED_PAGE_RANK => val pPageRankProto = apiMessage.getParallelPersonalizedPageRank val sourceIds = pPageRankProto.getSourceIdsList.asScala .map(parseLongOrStringID) @@ -345,16 +344,16 @@ object GraphFramesConnectUtils { .sourceIds(sourceIds) .run() .vertices // See comment in the PageRank - } - case proto.GraphFramesAPI.MethodCase.POWER_ITERATION_CLUSTERING => { + + case proto.GraphFramesAPI.MethodCase.POWER_ITERATION_CLUSTERING => val pic = apiMessage.getPowerIterationClustering if (pic.hasWeightCol) { graphFrame.powerIterationClustering(pic.getK, pic.getMaxIter, Some(pic.getWeightCol)) } else { graphFrame.powerIterationClustering(pic.getK, pic.getMaxIter, None) } - } - case proto.GraphFramesAPI.MethodCase.PREGEL => { + + case proto.GraphFramesAPI.MethodCase.PREGEL => val pregelProto = apiMessage.getPregel var pregel = graphFrame.pregel .aggMsgs(parseColumnOrExpression(pregelProto.getAggMsgs, planner)) @@ -421,8 +420,8 @@ object GraphFramesConnectUtils { } pregel.run() - } - case proto.GraphFramesAPI.MethodCase.SHORTEST_PATHS => { + + case proto.GraphFramesAPI.MethodCase.SHORTEST_PATHS => val isDirected = if (apiMessage.getShortestPaths.hasIsDirected) { apiMessage.getShortestPaths.getIsDirected } else { @@ -445,13 +444,13 @@ object GraphFramesConnectUtils { } else { spBuilder.run() } - } - case proto.GraphFramesAPI.MethodCase.STRONGLY_CONNECTED_COMPONENTS => { + + case proto.GraphFramesAPI.MethodCase.STRONGLY_CONNECTED_COMPONENTS => graphFrame.stronglyConnectedComponents .maxIter(apiMessage.getStronglyConnectedComponents.getMaxIter) .run() - } - case proto.GraphFramesAPI.MethodCase.SVD_PLUS_PLUS => { + + case proto.GraphFramesAPI.MethodCase.SVD_PLUS_PLUS => val svdPPProto = apiMessage.getSvdPlusPlus val svd = graphFrame.svdPlusPlus .maxIter(svdPPProto.getMaxIter) @@ -464,8 +463,8 @@ object GraphFramesConnectUtils { .maxValue(svdPPProto.getMaxValue) val svdResult = svd.run() svdResult.withColumn("loss", lit(svd.loss)) - } - case proto.GraphFramesAPI.MethodCase.TRIANGLE_COUNT => { + + case proto.GraphFramesAPI.MethodCase.TRIANGLE_COUNT => val message = apiMessage.getTriangleCount() var trCounter = graphFrame.triangleCount @@ -485,11 +484,11 @@ object GraphFramesConnectUtils { } else { trCounter.run() } - } - case proto.GraphFramesAPI.MethodCase.TRIPLETS => { + + case proto.GraphFramesAPI.MethodCase.TRIPLETS => graphFrame.triplets - } - case proto.GraphFramesAPI.MethodCase.MIS => { + + case proto.GraphFramesAPI.MethodCase.MIS => val mis = graphFrame.maximalIndependentSet .setCheckpointInterval(apiMessage.getMis.getCheckpointInterval) .setUseLocalCheckpoints(apiMessage.getMis.getUseLocalCheckpoints) @@ -501,8 +500,8 @@ object GraphFramesConnectUtils { } else { mis.run(apiMessage.getMis.getSeed) } - } - case proto.GraphFramesAPI.MethodCase.KCORE => { + + case proto.GraphFramesAPI.MethodCase.KCORE => var kCoreBuilder = graphFrame.kCore .setCheckpointInterval(apiMessage.getKcore.getCheckpointInterval) @@ -514,8 +513,8 @@ object GraphFramesConnectUtils { } kCoreBuilder.run() - } - case proto.GraphFramesAPI.MethodCase.AGGREGATE_NEIGHBORS => { + + case proto.GraphFramesAPI.MethodCase.AGGREGATE_NEIGHBORS => val anProto = apiMessage.getAggregateNeighbors var anBuilder = graphFrame.aggregateNeighbors .setStartingVertices(parseColumnOrExpression(anProto.getStartingVertices, planner)) @@ -574,9 +573,8 @@ object GraphFramesConnectUtils { } anBuilder.run() - } - case proto.GraphFramesAPI.MethodCase.RW_EMBEDDINGS => { + case proto.GraphFramesAPI.MethodCase.RW_EMBEDDINGS => val message = apiMessage.getRwEmbeddings() RandomWalkEmbeddings.pythonAPI( @@ -613,8 +611,8 @@ object GraphFramesConnectUtils { aggregateNeighborsMaxNbrs = message.getAggregateNeighborsMaxNbrs(), aggregateNeighborsSeed = message.getAggregateNeighborsSeed(), cleanUpAfterRun = message.getCleanUpAfterRun()) - } - case proto.GraphFramesAPI.MethodCase.HYPER_ANF => { + + case proto.GraphFramesAPI.MethodCase.HYPER_ANF => val haProto = apiMessage.getHyperAnf val haBuilder = graphFrame.hyperANF .setNHops(haProto.getNHops) @@ -632,7 +630,6 @@ object GraphFramesConnectUtils { } else { haBuilder.run() } - } case _ => throw new GraphFramesUnreachableException() // Unreachable } } From 93ef5a21636a573cd764999933440f7e81768237 Mon Sep 17 00:00:00 2001 From: Ruifeng Zheng Date: Fri, 28 Aug 2026 12:25:08 +0000 Subject: [PATCH 8/8] [GRAPHFRAMES][PYTHON] Add missing property type annotation --- python/pyspark/graphframes/pg/property_groups.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/pyspark/graphframes/pg/property_groups.py b/python/pyspark/graphframes/pg/property_groups.py index 45931474c624e..555e89b6be548 100644 --- a/python/pyspark/graphframes/pg/property_groups.py +++ b/python/pyspark/graphframes/pg/property_groups.py @@ -26,6 +26,7 @@ from pyspark.sql.functions import col, concat, lit, sha2 from pyspark.sql.types import ( ByteType, + DataType, DecimalType, DoubleType, FloatType, @@ -322,7 +323,7 @@ def _validate(self) -> None: _msg.format(self._weight_column_name, weight_column_type) ) - def _is_numeric_type(self, data_type) -> bool: + def _is_numeric_type(self, data_type: DataType) -> bool: """Check if a Spark data type is numeric.""" numeric_types = (