From fd25a08463b655a554211f97010552dd17b83df6 Mon Sep 17 00:00:00 2001 From: semyonsinchenko Date: Sat, 18 Oct 2025 08:45:12 +0200 Subject: [PATCH] K-Core --- connect/src/main/protobuf/graphframes.proto | 7 + .../graphframes/GraphFramesConnectUtils.scala | 13 + .../graphframes/expressions/KCoreMerge.scala | 101 ++++++ .../scala/org/graphframes/GraphFrame.scala | 9 + .../scala/org/graphframes/lib/KCore.scala | 103 ++++++ .../org/graphframes/lib/KCoreSuite.scala | 297 ++++++++++++++++++ docs/src/04-user-guide/03-centralities.md | 161 +++++++++- python/graphframes/classic/graphframe.py | 15 + .../graphframes/connect/graphframes_client.py | 50 +++ .../connect/proto/graphframes_pb2.py | 90 +++--- .../connect/proto/graphframes_pb2.pyi | 133 +++----- python/graphframes/graphframe.py | 24 ++ python/tests/test_graphframes.py | 93 ++++++ 13 files changed, 956 insertions(+), 140 deletions(-) create mode 100644 core/src/main/scala/org/apache/spark/sql/graphframes/expressions/KCoreMerge.scala create mode 100644 core/src/main/scala/org/graphframes/lib/KCore.scala create mode 100644 core/src/test/scala/org/graphframes/lib/KCoreSuite.scala diff --git a/connect/src/main/protobuf/graphframes.proto b/connect/src/main/protobuf/graphframes.proto index c22223941..6a92ec2aa 100644 --- a/connect/src/main/protobuf/graphframes.proto +++ b/connect/src/main/protobuf/graphframes.proto @@ -35,6 +35,7 @@ message GraphFramesAPI { SVDPlusPlus svd_plus_plus = 18; TriangleCount triangle_count = 19; Triplets triplets = 20; + KCore kcore = 21; } } @@ -186,3 +187,9 @@ message TriangleCount { } message Triplets {} + +message KCore { + bool use_local_checkpoints = 1; + int32 checkpoint_interval = 2; + optional StorageLevel storage_level = 3; +} diff --git a/connect/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala b/connect/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala index ea63ae281..f9018959b 100644 --- a/connect/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala +++ b/connect/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala @@ -399,6 +399,19 @@ object GraphFramesConnectUtils { case proto.GraphFramesAPI.MethodCase.TRIPLETS => { graphFrame.triplets } + 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 _ => throw new GraphFramesUnreachableException() // Unreachable } } diff --git a/core/src/main/scala/org/apache/spark/sql/graphframes/expressions/KCoreMerge.scala b/core/src/main/scala/org/apache/spark/sql/graphframes/expressions/KCoreMerge.scala new file mode 100644 index 000000000..82765c070 --- /dev/null +++ b/core/src/main/scala/org/apache/spark/sql/graphframes/expressions/KCoreMerge.scala @@ -0,0 +1,101 @@ +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/core/src/main/scala/org/graphframes/GraphFrame.scala b/core/src/main/scala/org/graphframes/GraphFrame.scala index af2a1ae1e..ccb2f8aea 100644 --- a/core/src/main/scala/org/graphframes/GraphFrame.scala +++ b/core/src/main/scala/org/graphframes/GraphFrame.scala @@ -616,6 +616,15 @@ class GraphFrame private ( } } + /** + * K-Core decomposition. + * + * See [[org.graphframes.lib.KCore]] for more details. + * + * @group stdlib + */ + def kCore: KCore = new KCore(this) + /** * Validates the consistency and integrity of a graph by performing checks on the vertices and * edges. diff --git a/core/src/main/scala/org/graphframes/lib/KCore.scala b/core/src/main/scala/org/graphframes/lib/KCore.scala new file mode 100644 index 000000000..3db02ab28 --- /dev/null +++ b/core/src/main/scala/org/graphframes/lib/KCore.scala @@ -0,0 +1,103 @@ +package org.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.graphframes.GraphFrame +import org.graphframes.Logging +import org.graphframes.WithCheckpointInterval +import org.graphframes.WithIntermediateStorageLevel +import org.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. + */ +class KCore private[graphframes] (private val graph: GraphFrame) + extends Serializable + with WithIntermediateStorageLevel + with WithCheckpointInterval + with WithLocalCheckpoints { + import org.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( + FunctionIdentifier("_kcoreMerge"), + (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.src(kCoreColumnName)) + .sendMsgToDst(Pregel.dst(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(FunctionIdentifier("_kcoreMerge")) + 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/core/src/test/scala/org/graphframes/lib/KCoreSuite.scala b/core/src/test/scala/org/graphframes/lib/KCoreSuite.scala new file mode 100644 index 000000000..1a3e27376 --- /dev/null +++ b/core/src/test/scala/org/graphframes/lib/KCoreSuite.scala @@ -0,0 +1,297 @@ +package org.graphframes.lib + +import org.apache.spark.sql.types.DataTypes +import org.graphframes.* +import org.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() + // Center vertex should have k-core value of 3, leaf vertices should have k-core value of 1 + rows.foreach { row => + val id = row.getAs[Long]("id") + val kcore = row.getAs[Int]("kcore") + if (id == 0L) { + assert(kcore === 3) + } else { + assert(kcore === 1) + } + } + result.unpersist() + } + + test("chain 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 + // because graph is cosidered as undirected + rows.foreach { row => + assert(row.getAs[Int]("kcore") === 2) + } + 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 > 3, "Should have more than 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) >= 4, "Central vertex should have high k-core") + assert(kcoreMap(1L) >= 3, "Well-connected vertex should have medium-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() + } +} diff --git a/docs/src/04-user-guide/03-centralities.md b/docs/src/04-user-guide/03-centralities.md index 1d7882b41..e3e66f81a 100644 --- a/docs/src/04-user-guide/03-centralities.md +++ b/docs/src/04-user-guide/03-centralities.md @@ -117,4 +117,163 @@ GraphFrames also supports parallel personalized PageRank that allows users to co For the API details refer to: * Scala API: @:scaladoc(org.graphframes.lib.ParallelPersonalizedPageRank) -* Python API: @:pydoc(graphframes.GraphFrame.parallelPersonalizedPageRank) \ No newline at end of file +* Python API: @:pydoc(graphframes.GraphFrame.parallelPersonalizedPageRank) + +## K-Core + +K-Core decomposition is a method used to identify the most tightly connected subgraphs within a network. A k-core is a maximal subgraph where every vertex has at least degree k. This metric helps in understanding the inner structure of networks by filtering out less connected nodes, revealing cores of highly interconnected entities. K-Core centrality can be applied in various domains such as social network analysis to find influential users, in biology to detect stable protein complexes, or in infrastructure networks to assess robustness and vulnerability. + +The provided implementation of K-Core decomposition in GraphFrames is based on the research described in the paper available at [IEEE Xplore](https://ieeexplore.ieee.org/abstract/document/8258018). Using think-like-a-vertex paradigm, the proposed method utilizes a message passing paradigm for solving k-core decomposition, thus reducing the I/O cost substantially. + +For more information, see: + +> A. Farajollahi, S. G. Khaki, and L. Wang, "Efficient distributed k-core decomposition for large-scale graphs," *2017 IEEE International Conference on Big Data (Big Data)*, Boston, MA, USA, 2017, pp. 1430-1435. + +### Arguments + +- `checkpoint_interval` + +For `graphframes` only. To avoid exponential growing of the Spark' Logical Plan, DataFrame lineage and query optimization time, it is required to do checkpointing periodically. While checkpoint itself is not free, it is still recommended to set this value to something less than `5`. + +- `use_local_checkpoints` + +For `graphframes` only. By default, GraphFrames uses persistent checkpoints. They are realiable and reduce the errors rate. The downside of the persistent checkpoints is that they are requiride to set up a `checkpointDir` in persistent storage like `S3` or `HDFS`. By providing `use_local_checkpoints=True`, user can say GraphFrames to use local disks of Spark' executurs for checkpointing. Local checkpoints are faster, but they are less reliable: if the executur lost, for example, is taking by the higher priority job, checkpoints will be lost and the whole job fails. + +- `storage_level` + +The level of storage for intermediate results and the output `DataFrame` with components. By default it is memory and disk deserialized as a good balance between performance and reliability. For very big graphs and out-of-core scenarious, using `DISK_ONLY` may be faster. + +### Python API + +```python +import org.graphframes.GraphFrame + +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, +) +``` + +### Scala API + +```scala +import org.graphframes.GraphFrame + +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() +``` diff --git a/python/graphframes/classic/graphframe.py b/python/graphframes/classic/graphframe.py index ab9618906..9623c1c69 100644 --- a/python/graphframes/classic/graphframe.py +++ b/python/graphframes/classic/graphframe.py @@ -341,3 +341,18 @@ def powerIterationClustering( weightCol = self._spark._jvm.scala.Option.empty() jdf = self._jvm_graph.powerIterationClustering(k, maxIter, weightCol) return DataFrame(jdf, self._spark) + + def k_core( + self, + checkpoint_interval: int, + use_local_checkpoints: bool, + storage_level: StorageLevel, + ) -> DataFrame: + jdf = ( + self._jvm_graph.kCore() + .setUseLocalCheckpoints(use_local_checkpoints) + .setCheckpointInterval(checkpoint_interval) + .setIntermediateStorageLevel(storage_level_to_jvm(storage_level, self._spark)) + .run() + ) + return DataFrame(jdf, self._spark) diff --git a/python/graphframes/connect/graphframes_client.py b/python/graphframes/connect/graphframes_client.py index 2e64d7a41..ce9315774 100644 --- a/python/graphframes/connect/graphframes_client.py +++ b/python/graphframes/connect/graphframes_client.py @@ -1064,3 +1064,53 @@ def plan(self, session: SparkConnectClient) -> proto.Relation: return _dataframe_from_plan( TriangleCount(self._vertices, self._edges, storage_level), 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.extension.Pack(graphframes_api_call) + return plan + + return _dataframe_from_plan( + KCore( + self._vertices, + self._edges, + checkpoint_interval, + use_local_checkpoints, + storage_level, + ), + self._spark, + ) diff --git a/python/graphframes/connect/proto/graphframes_pb2.py b/python/graphframes/connect/proto/graphframes_pb2.py index a34f42a42..30f006066 100644 --- a/python/graphframes/connect/proto/graphframes_pb2.py +++ b/python/graphframes/connect/proto/graphframes_pb2.py @@ -19,7 +19,7 @@ DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( - b'\n\x11graphframes.proto\x12\x1dorg.graphframes.connect.proto"\xb3\r\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\x61\n\x12\x61ggregate_messages\x18\x03 \x01(\x0b\x32\x30.org.graphframes.connect.proto.AggregateMessagesH\x00R\x11\x61ggregateMessages\x12\x36\n\x03\x62\x66s\x18\x04 \x01(\x0b\x32".org.graphframes.connect.proto.BFSH\x00R\x03\x62\x66s\x12g\n\x14\x63onnected_components\x18\x05 \x01(\x0b\x32\x32.org.graphframes.connect.proto.ConnectedComponentsH\x00R\x13\x63onnectedComponents\x12k\n\x16\x64rop_isolated_vertices\x18\x06 \x01(\x0b\x32\x33.org.graphframes.connect.proto.DropIsolatedVerticesH\x00R\x14\x64ropIsolatedVertices\x12[\n\x10\x64\x65tecting_cycles\x18\x07 \x01(\x0b\x32..org.graphframes.connect.proto.DetectingCyclesH\x00R\x0f\x64\x65tectingCycles\x12O\n\x0c\x66ilter_edges\x18\x08 \x01(\x0b\x32*.org.graphframes.connect.proto.FilterEdgesH\x00R\x0b\x66ilterEdges\x12X\n\x0f\x66ilter_vertices\x18\t \x01(\x0b\x32-.org.graphframes.connect.proto.FilterVerticesH\x00R\x0e\x66ilterVertices\x12\x39\n\x04\x66ind\x18\n \x01(\x0b\x32#.org.graphframes.connect.proto.FindH\x00R\x04\x66ind\x12^\n\x11label_propagation\x18\x0b \x01(\x0b\x32/.org.graphframes.connect.proto.LabelPropagationH\x00R\x10labelPropagation\x12\x46\n\tpage_rank\x18\x0c \x01(\x0b\x32\'.org.graphframes.connect.proto.PageRankH\x00R\x08pageRank\x12\x84\x01\n\x1fparallel_personalized_page_rank\x18\r \x01(\x0b\x32;.org.graphframes.connect.proto.ParallelPersonalizedPageRankH\x00R\x1cparallelPersonalizedPageRank\x12w\n\x1apower_iteration_clustering\x18\x0e \x01(\x0b\x32\x37.org.graphframes.connect.proto.PowerIterationClusteringH\x00R\x18powerIterationClustering\x12?\n\x06pregel\x18\x0f \x01(\x0b\x32%.org.graphframes.connect.proto.PregelH\x00R\x06pregel\x12U\n\x0eshortest_paths\x18\x10 \x01(\x0b\x32,.org.graphframes.connect.proto.ShortestPathsH\x00R\rshortestPaths\x12\x80\x01\n\x1dstrongly_connected_components\x18\x11 \x01(\x0b\x32:.org.graphframes.connect.proto.StronglyConnectedComponentsH\x00R\x1bstronglyConnectedComponents\x12P\n\rsvd_plus_plus\x18\x12 \x01(\x0b\x32*.org.graphframes.connect.proto.SVDPlusPlusH\x00R\x0bsvdPlusPlus\x12U\n\x0etriangle_count\x18\x13 \x01(\x0b\x32,.org.graphframes.connect.proto.TriangleCountH\x00R\rtriangleCount\x12\x45\n\x08triplets\x18\x14 \x01(\x0b\x32\'.org.graphframes.connect.proto.TripletsH\x00R\x08tripletsB\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"\xee\x02\n\x11\x41ggregateMessages\x12J\n\x07\x61gg_col\x18\x01 \x03(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\x06\x61ggCol\x12Q\n\x0bsend_to_src\x18\x02 \x03(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\tsendToSrc\x12Q\n\x0bsend_to_dst\x18\x03 \x03(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\tsendToDst\x12U\n\rstorage_level\x18\x04 \x01(\x0b\x32+.org.graphframes.connect.proto.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\x9d\x02\n\x03\x42\x46S\x12N\n\tfrom_expr\x18\x01 \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\x08\x66romExpr\x12J\n\x07to_expr\x18\x02 \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\x06toExpr\x12R\n\x0b\x65\x64ge_filter\x18\x03 \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\nedgeFilter\x12&\n\x0fmax_path_length\x18\x04 \x01(\x05R\rmaxPathLength"\x86\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\x12U\n\rstorage_level\x18\x07 \x01(\x0b\x32+.org.graphframes.connect.proto.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\xdf\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\x12U\n\rstorage_level\x18\x03 \x01(\x0b\x32+.org.graphframes.connect.proto.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\x16\n\x14\x44ropIsolatedVertices"^\n\x0b\x46ilterEdges\x12O\n\tcondition\x18\x01 \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\tcondition"a\n\x0e\x46ilterVertices\x12O\n\tcondition\x18\x02 \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\tcondition" \n\x04\x46ind\x12\x18\n\x07pattern\x18\x01 \x01(\tR\x07pattern"\x99\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\x12U\n\rstorage_level\x18\x05 \x01(\x0b\x32+.org.graphframes.connect.proto.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\xe2\x01\n\x08PageRank\x12+\n\x11reset_probability\x18\x01 \x01(\x01R\x10resetProbability\x12O\n\tsource_id\x18\x02 \x01(\x0b\x32-.org.graphframes.connect.proto.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"\xb4\x01\n\x1cParallelPersonalizedPageRank\x12+\n\x11reset_probability\x18\x01 \x01(\x01R\x10resetProbability\x12L\n\nsource_ids\x18\x02 \x03(\x0b\x32-.org.graphframes.connect.proto.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"\xe6\t\n\x06Pregel\x12L\n\x08\x61gg_msgs\x18\x01 \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\x07\x61ggMsgs\x12X\n\x0fsend_msg_to_dst\x18\x02 \x03(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\x0csendMsgToDst\x12X\n\x0fsend_msg_to_src\x18\x03 \x03(\x0b\x32\x31.org.graphframes.connect.proto.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\x12g\n\x16\x61\x64\x64itional_col_initial\x18\x07 \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\x14\x61\x64\x64itionalColInitial\x12_\n\x12\x61\x64\x64itional_col_upd\x18\x08 \x01(\x0b\x32\x31.org.graphframes.connect.proto.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\x12U\n\rstorage_level\x18\x0b \x01(\x0b\x32+.org.graphframes.connect.proto.StorageLevelH\x01R\x0cstorageLevel\x88\x01\x01\x12\x37\n\x16stop_if_all_non_active\x18\x0c \x01(\x08H\x02R\x12stopIfAllNonActive\x88\x01\x01\x12\x66\n\x13initial_active_expr\x18\r \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionH\x03R\x11initialActiveExpr\x88\x01\x01\x12\x64\n\x12update_active_expr\x18\x0e \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionH\x04R\x10updateActiveExpr\x88\x01\x01\x12\x45\n\x1dskip_messages_from_non_active\x18\x0f \x01(\x08H\x05R\x19skipMessagesFromNonActive\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_active"\xc8\x02\n\rShortestPaths\x12K\n\tlandmarks\x18\x01 \x03(\x0b\x32-.org.graphframes.connect.proto.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\x12U\n\rstorage_level\x18\x05 \x01(\x0b\x32+.org.graphframes.connect.proto.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"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"x\n\rTriangleCount\x12U\n\rstorage_level\x18\x01 \x01(\x0b\x32+.org.graphframes.connect.proto.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\n\n\x08TripletsB\xd2\x01\n!com.org.graphframes.connect.protoB\x10GraphframesProtoH\x01P\x01\xa0\x01\x01\xa2\x02\x04OGCP\xaa\x02\x1dOrg.Graphframes.Connect.Proto\xca\x02\x1dOrg\\Graphframes\\Connect\\Proto\xe2\x02)Org\\Graphframes\\Connect\\Proto\\GPBMetadata\xea\x02 Org::Graphframes::Connect::Protob\x06proto3' + b'\n\x11graphframes.proto\x12\x1dorg.graphframes.connect.proto"\xf1\r\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\x61\n\x12\x61ggregate_messages\x18\x03 \x01(\x0b\x32\x30.org.graphframes.connect.proto.AggregateMessagesH\x00R\x11\x61ggregateMessages\x12\x36\n\x03\x62\x66s\x18\x04 \x01(\x0b\x32".org.graphframes.connect.proto.BFSH\x00R\x03\x62\x66s\x12g\n\x14\x63onnected_components\x18\x05 \x01(\x0b\x32\x32.org.graphframes.connect.proto.ConnectedComponentsH\x00R\x13\x63onnectedComponents\x12k\n\x16\x64rop_isolated_vertices\x18\x06 \x01(\x0b\x32\x33.org.graphframes.connect.proto.DropIsolatedVerticesH\x00R\x14\x64ropIsolatedVertices\x12[\n\x10\x64\x65tecting_cycles\x18\x07 \x01(\x0b\x32..org.graphframes.connect.proto.DetectingCyclesH\x00R\x0f\x64\x65tectingCycles\x12O\n\x0c\x66ilter_edges\x18\x08 \x01(\x0b\x32*.org.graphframes.connect.proto.FilterEdgesH\x00R\x0b\x66ilterEdges\x12X\n\x0f\x66ilter_vertices\x18\t \x01(\x0b\x32-.org.graphframes.connect.proto.FilterVerticesH\x00R\x0e\x66ilterVertices\x12\x39\n\x04\x66ind\x18\n \x01(\x0b\x32#.org.graphframes.connect.proto.FindH\x00R\x04\x66ind\x12^\n\x11label_propagation\x18\x0b \x01(\x0b\x32/.org.graphframes.connect.proto.LabelPropagationH\x00R\x10labelPropagation\x12\x46\n\tpage_rank\x18\x0c \x01(\x0b\x32\'.org.graphframes.connect.proto.PageRankH\x00R\x08pageRank\x12\x84\x01\n\x1fparallel_personalized_page_rank\x18\r \x01(\x0b\x32;.org.graphframes.connect.proto.ParallelPersonalizedPageRankH\x00R\x1cparallelPersonalizedPageRank\x12w\n\x1apower_iteration_clustering\x18\x0e \x01(\x0b\x32\x37.org.graphframes.connect.proto.PowerIterationClusteringH\x00R\x18powerIterationClustering\x12?\n\x06pregel\x18\x0f \x01(\x0b\x32%.org.graphframes.connect.proto.PregelH\x00R\x06pregel\x12U\n\x0eshortest_paths\x18\x10 \x01(\x0b\x32,.org.graphframes.connect.proto.ShortestPathsH\x00R\rshortestPaths\x12\x80\x01\n\x1dstrongly_connected_components\x18\x11 \x01(\x0b\x32:.org.graphframes.connect.proto.StronglyConnectedComponentsH\x00R\x1bstronglyConnectedComponents\x12P\n\rsvd_plus_plus\x18\x12 \x01(\x0b\x32*.org.graphframes.connect.proto.SVDPlusPlusH\x00R\x0bsvdPlusPlus\x12U\n\x0etriangle_count\x18\x13 \x01(\x0b\x32,.org.graphframes.connect.proto.TriangleCountH\x00R\rtriangleCount\x12\x45\n\x08triplets\x18\x14 \x01(\x0b\x32\'.org.graphframes.connect.proto.TripletsH\x00R\x08triplets\x12<\n\x05kcore\x18\x15 \x01(\x0b\x32$.org.graphframes.connect.proto.KCoreH\x00R\x05kcoreB\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"\xee\x02\n\x11\x41ggregateMessages\x12J\n\x07\x61gg_col\x18\x01 \x03(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\x06\x61ggCol\x12Q\n\x0bsend_to_src\x18\x02 \x03(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\tsendToSrc\x12Q\n\x0bsend_to_dst\x18\x03 \x03(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\tsendToDst\x12U\n\rstorage_level\x18\x04 \x01(\x0b\x32+.org.graphframes.connect.proto.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\x9d\x02\n\x03\x42\x46S\x12N\n\tfrom_expr\x18\x01 \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\x08\x66romExpr\x12J\n\x07to_expr\x18\x02 \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\x06toExpr\x12R\n\x0b\x65\x64ge_filter\x18\x03 \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\nedgeFilter\x12&\n\x0fmax_path_length\x18\x04 \x01(\x05R\rmaxPathLength"\x86\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\x12U\n\rstorage_level\x18\x07 \x01(\x0b\x32+.org.graphframes.connect.proto.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\xdf\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\x12U\n\rstorage_level\x18\x03 \x01(\x0b\x32+.org.graphframes.connect.proto.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\x16\n\x14\x44ropIsolatedVertices"^\n\x0b\x46ilterEdges\x12O\n\tcondition\x18\x01 \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\tcondition"a\n\x0e\x46ilterVertices\x12O\n\tcondition\x18\x02 \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\tcondition" \n\x04\x46ind\x12\x18\n\x07pattern\x18\x01 \x01(\tR\x07pattern"\x99\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\x12U\n\rstorage_level\x18\x05 \x01(\x0b\x32+.org.graphframes.connect.proto.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\xe2\x01\n\x08PageRank\x12+\n\x11reset_probability\x18\x01 \x01(\x01R\x10resetProbability\x12O\n\tsource_id\x18\x02 \x01(\x0b\x32-.org.graphframes.connect.proto.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"\xb4\x01\n\x1cParallelPersonalizedPageRank\x12+\n\x11reset_probability\x18\x01 \x01(\x01R\x10resetProbability\x12L\n\nsource_ids\x18\x02 \x03(\x0b\x32-.org.graphframes.connect.proto.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"\xe6\t\n\x06Pregel\x12L\n\x08\x61gg_msgs\x18\x01 \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\x07\x61ggMsgs\x12X\n\x0fsend_msg_to_dst\x18\x02 \x03(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\x0csendMsgToDst\x12X\n\x0fsend_msg_to_src\x18\x03 \x03(\x0b\x32\x31.org.graphframes.connect.proto.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\x12g\n\x16\x61\x64\x64itional_col_initial\x18\x07 \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionR\x14\x61\x64\x64itionalColInitial\x12_\n\x12\x61\x64\x64itional_col_upd\x18\x08 \x01(\x0b\x32\x31.org.graphframes.connect.proto.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\x12U\n\rstorage_level\x18\x0b \x01(\x0b\x32+.org.graphframes.connect.proto.StorageLevelH\x01R\x0cstorageLevel\x88\x01\x01\x12\x37\n\x16stop_if_all_non_active\x18\x0c \x01(\x08H\x02R\x12stopIfAllNonActive\x88\x01\x01\x12\x66\n\x13initial_active_expr\x18\r \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionH\x03R\x11initialActiveExpr\x88\x01\x01\x12\x64\n\x12update_active_expr\x18\x0e \x01(\x0b\x32\x31.org.graphframes.connect.proto.ColumnOrExpressionH\x04R\x10updateActiveExpr\x88\x01\x01\x12\x45\n\x1dskip_messages_from_non_active\x18\x0f \x01(\x08H\x05R\x19skipMessagesFromNonActive\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_active"\xc8\x02\n\rShortestPaths\x12K\n\tlandmarks\x18\x01 \x03(\x0b\x32-.org.graphframes.connect.proto.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\x12U\n\rstorage_level\x18\x05 \x01(\x0b\x32+.org.graphframes.connect.proto.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"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"x\n\rTriangleCount\x12U\n\rstorage_level\x18\x01 \x01(\x0b\x32+.org.graphframes.connect.proto.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_level"\n\n\x08Triplets"\xd5\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\x12U\n\rstorage_level\x18\x03 \x01(\x0b\x32+.org.graphframes.connect.proto.StorageLevelH\x00R\x0cstorageLevel\x88\x01\x01\x42\x10\n\x0e_storage_levelB\xd2\x01\n!com.org.graphframes.connect.protoB\x10GraphframesProtoH\x01P\x01\xa0\x01\x01\xa2\x02\x04OGCP\xaa\x02\x1dOrg.Graphframes.Connect.Proto\xca\x02\x1dOrg\\Graphframes\\Connect\\Proto\xe2\x02)Org\\Graphframes\\Connect\\Proto\\GPBMetadata\xea\x02 Org::Graphframes::Connect::Protob\x06proto3' ) _globals = globals() @@ -31,47 +31,49 @@ "DESCRIPTOR" ]._serialized_options = b"\n!com.org.graphframes.connect.protoB\020GraphframesProtoH\001P\001\240\001\001\242\002\004OGCP\252\002\035Org.Graphframes.Connect.Proto\312\002\035Org\\Graphframes\\Connect\\Proto\342\002)Org\\Graphframes\\Connect\\Proto\\GPBMetadata\352\002 Org::Graphframes::Connect::Proto" _globals["_GRAPHFRAMESAPI"]._serialized_start = 53 - _globals["_GRAPHFRAMESAPI"]._serialized_end = 1768 - _globals["_STORAGELEVEL"]._serialized_start = 1771 - _globals["_STORAGELEVEL"]._serialized_end = 2114 - _globals["_COLUMNOREXPRESSION"]._serialized_start = 2116 - _globals["_COLUMNOREXPRESSION"]._serialized_end = 2193 - _globals["_STRINGORLONGID"]._serialized_start = 2195 - _globals["_STRINGORLONGID"]._serialized_end = 2275 - _globals["_AGGREGATEMESSAGES"]._serialized_start = 2278 - _globals["_AGGREGATEMESSAGES"]._serialized_end = 2644 - _globals["_BFS"]._serialized_start = 2647 - _globals["_BFS"]._serialized_end = 2932 - _globals["_CONNECTEDCOMPONENTS"]._serialized_start = 2935 - _globals["_CONNECTEDCOMPONENTS"]._serialized_end = 3325 - _globals["_DETECTINGCYCLES"]._serialized_start = 3328 - _globals["_DETECTINGCYCLES"]._serialized_end = 3551 - _globals["_DROPISOLATEDVERTICES"]._serialized_start = 3553 - _globals["_DROPISOLATEDVERTICES"]._serialized_end = 3575 - _globals["_FILTEREDGES"]._serialized_start = 3577 - _globals["_FILTEREDGES"]._serialized_end = 3671 - _globals["_FILTERVERTICES"]._serialized_start = 3673 - _globals["_FILTERVERTICES"]._serialized_end = 3770 - _globals["_FIND"]._serialized_start = 3772 - _globals["_FIND"]._serialized_end = 3804 - _globals["_LABELPROPAGATION"]._serialized_start = 3807 - _globals["_LABELPROPAGATION"]._serialized_end = 4088 - _globals["_PAGERANK"]._serialized_start = 4091 - _globals["_PAGERANK"]._serialized_end = 4317 - _globals["_PARALLELPERSONALIZEDPAGERANK"]._serialized_start = 4320 - _globals["_PARALLELPERSONALIZEDPAGERANK"]._serialized_end = 4500 - _globals["_POWERITERATIONCLUSTERING"]._serialized_start = 4502 - _globals["_POWERITERATIONCLUSTERING"]._serialized_end = 4620 - _globals["_PREGEL"]._serialized_start = 4623 - _globals["_PREGEL"]._serialized_end = 5877 - _globals["_SHORTESTPATHS"]._serialized_start = 5880 - _globals["_SHORTESTPATHS"]._serialized_end = 6208 - _globals["_STRONGLYCONNECTEDCOMPONENTS"]._serialized_start = 6210 - _globals["_STRONGLYCONNECTEDCOMPONENTS"]._serialized_end = 6266 - _globals["_SVDPLUSPLUS"]._serialized_start = 6269 - _globals["_SVDPLUSPLUS"]._serialized_end = 6483 - _globals["_TRIANGLECOUNT"]._serialized_start = 6485 - _globals["_TRIANGLECOUNT"]._serialized_end = 6605 - _globals["_TRIPLETS"]._serialized_start = 6607 - _globals["_TRIPLETS"]._serialized_end = 6617 + _globals["_GRAPHFRAMESAPI"]._serialized_end = 1830 + _globals["_STORAGELEVEL"]._serialized_start = 1833 + _globals["_STORAGELEVEL"]._serialized_end = 2176 + _globals["_COLUMNOREXPRESSION"]._serialized_start = 2178 + _globals["_COLUMNOREXPRESSION"]._serialized_end = 2255 + _globals["_STRINGORLONGID"]._serialized_start = 2257 + _globals["_STRINGORLONGID"]._serialized_end = 2337 + _globals["_AGGREGATEMESSAGES"]._serialized_start = 2340 + _globals["_AGGREGATEMESSAGES"]._serialized_end = 2706 + _globals["_BFS"]._serialized_start = 2709 + _globals["_BFS"]._serialized_end = 2994 + _globals["_CONNECTEDCOMPONENTS"]._serialized_start = 2997 + _globals["_CONNECTEDCOMPONENTS"]._serialized_end = 3387 + _globals["_DETECTINGCYCLES"]._serialized_start = 3390 + _globals["_DETECTINGCYCLES"]._serialized_end = 3613 + _globals["_DROPISOLATEDVERTICES"]._serialized_start = 3615 + _globals["_DROPISOLATEDVERTICES"]._serialized_end = 3637 + _globals["_FILTEREDGES"]._serialized_start = 3639 + _globals["_FILTEREDGES"]._serialized_end = 3733 + _globals["_FILTERVERTICES"]._serialized_start = 3735 + _globals["_FILTERVERTICES"]._serialized_end = 3832 + _globals["_FIND"]._serialized_start = 3834 + _globals["_FIND"]._serialized_end = 3866 + _globals["_LABELPROPAGATION"]._serialized_start = 3869 + _globals["_LABELPROPAGATION"]._serialized_end = 4150 + _globals["_PAGERANK"]._serialized_start = 4153 + _globals["_PAGERANK"]._serialized_end = 4379 + _globals["_PARALLELPERSONALIZEDPAGERANK"]._serialized_start = 4382 + _globals["_PARALLELPERSONALIZEDPAGERANK"]._serialized_end = 4562 + _globals["_POWERITERATIONCLUSTERING"]._serialized_start = 4564 + _globals["_POWERITERATIONCLUSTERING"]._serialized_end = 4682 + _globals["_PREGEL"]._serialized_start = 4685 + _globals["_PREGEL"]._serialized_end = 5939 + _globals["_SHORTESTPATHS"]._serialized_start = 5942 + _globals["_SHORTESTPATHS"]._serialized_end = 6270 + _globals["_STRONGLYCONNECTEDCOMPONENTS"]._serialized_start = 6272 + _globals["_STRONGLYCONNECTEDCOMPONENTS"]._serialized_end = 6328 + _globals["_SVDPLUSPLUS"]._serialized_start = 6331 + _globals["_SVDPLUSPLUS"]._serialized_end = 6545 + _globals["_TRIANGLECOUNT"]._serialized_start = 6547 + _globals["_TRIANGLECOUNT"]._serialized_end = 6667 + _globals["_TRIPLETS"]._serialized_start = 6669 + _globals["_TRIPLETS"]._serialized_end = 6679 + _globals["_KCORE"]._serialized_start = 6682 + _globals["_KCORE"]._serialized_end = 6895 # @@protoc_insertion_point(module_scope) diff --git a/python/graphframes/connect/proto/graphframes_pb2.pyi b/python/graphframes/connect/proto/graphframes_pb2.pyi index ffe59932d..ce14ed8ca 100644 --- a/python/graphframes/connect/proto/graphframes_pb2.pyi +++ b/python/graphframes/connect/proto/graphframes_pb2.pyi @@ -11,28 +11,7 @@ from google.protobuf.internal import containers as _containers DESCRIPTOR: _descriptor.FileDescriptor class GraphFramesAPI(_message.Message): - __slots__ = ( - "vertices", - "edges", - "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", - ) + __slots__ = () VERTICES_FIELD_NUMBER: _ClassVar[int] EDGES_FIELD_NUMBER: _ClassVar[int] AGGREGATE_MESSAGES_FIELD_NUMBER: _ClassVar[int] @@ -53,6 +32,7 @@ class GraphFramesAPI(_message.Message): SVD_PLUS_PLUS_FIELD_NUMBER: _ClassVar[int] TRIANGLE_COUNT_FIELD_NUMBER: _ClassVar[int] TRIPLETS_FIELD_NUMBER: _ClassVar[int] + KCORE_FIELD_NUMBER: _ClassVar[int] vertices: bytes edges: bytes aggregate_messages: AggregateMessages @@ -73,6 +53,7 @@ class GraphFramesAPI(_message.Message): svd_plus_plus: SVDPlusPlus triangle_count: TriangleCount triplets: Triplets + kcore: KCore def __init__( self, vertices: _Optional[bytes] = ..., @@ -99,19 +80,11 @@ class GraphFramesAPI(_message.Message): svd_plus_plus: _Optional[_Union[SVDPlusPlus, _Mapping]] = ..., triangle_count: _Optional[_Union[TriangleCount, _Mapping]] = ..., triplets: _Optional[_Union[Triplets, _Mapping]] = ..., + kcore: _Optional[_Union[KCore, _Mapping]] = ..., ) -> None: ... class StorageLevel(_message.Message): - __slots__ = ( - "disk_only", - "disk_only_2", - "disk_only_3", - "memory_and_disk", - "memory_and_disk_2", - "memory_and_disk_deser", - "memory_only", - "memory_only_2", - ) + __slots__ = () DISK_ONLY_FIELD_NUMBER: _ClassVar[int] DISK_ONLY_2_FIELD_NUMBER: _ClassVar[int] DISK_ONLY_3_FIELD_NUMBER: _ClassVar[int] @@ -141,7 +114,7 @@ class StorageLevel(_message.Message): ) -> None: ... class ColumnOrExpression(_message.Message): - __slots__ = ("col", "expr") + __slots__ = () COL_FIELD_NUMBER: _ClassVar[int] EXPR_FIELD_NUMBER: _ClassVar[int] col: bytes @@ -149,7 +122,7 @@ class ColumnOrExpression(_message.Message): def __init__(self, col: _Optional[bytes] = ..., expr: _Optional[str] = ...) -> None: ... class StringOrLongID(_message.Message): - __slots__ = ("long_id", "string_id") + __slots__ = () LONG_ID_FIELD_NUMBER: _ClassVar[int] STRING_ID_FIELD_NUMBER: _ClassVar[int] long_id: int @@ -157,7 +130,7 @@ class StringOrLongID(_message.Message): def __init__(self, long_id: _Optional[int] = ..., string_id: _Optional[str] = ...) -> None: ... class AggregateMessages(_message.Message): - __slots__ = ("agg_col", "send_to_src", "send_to_dst", "storage_level") + __slots__ = () AGG_COL_FIELD_NUMBER: _ClassVar[int] SEND_TO_SRC_FIELD_NUMBER: _ClassVar[int] SEND_TO_DST_FIELD_NUMBER: _ClassVar[int] @@ -175,7 +148,7 @@ class AggregateMessages(_message.Message): ) -> None: ... class BFS(_message.Message): - __slots__ = ("from_expr", "to_expr", "edge_filter", "max_path_length") + __slots__ = () FROM_EXPR_FIELD_NUMBER: _ClassVar[int] TO_EXPR_FIELD_NUMBER: _ClassVar[int] EDGE_FILTER_FIELD_NUMBER: _ClassVar[int] @@ -193,15 +166,7 @@ class BFS(_message.Message): ) -> None: ... class ConnectedComponents(_message.Message): - __slots__ = ( - "algorithm", - "checkpoint_interval", - "broadcast_threshold", - "use_labels_as_components", - "use_local_checkpoints", - "max_iter", - "storage_level", - ) + __slots__ = () ALGORITHM_FIELD_NUMBER: _ClassVar[int] CHECKPOINT_INTERVAL_FIELD_NUMBER: _ClassVar[int] BROADCAST_THRESHOLD_FIELD_NUMBER: _ClassVar[int] @@ -228,7 +193,7 @@ class ConnectedComponents(_message.Message): ) -> None: ... class DetectingCycles(_message.Message): - __slots__ = ("use_local_checkpoints", "checkpoint_interval", "storage_level") + __slots__ = () USE_LOCAL_CHECKPOINTS_FIELD_NUMBER: _ClassVar[int] CHECKPOINT_INTERVAL_FIELD_NUMBER: _ClassVar[int] STORAGE_LEVEL_FIELD_NUMBER: _ClassVar[int] @@ -247,7 +212,7 @@ class DropIsolatedVertices(_message.Message): def __init__(self) -> None: ... class FilterEdges(_message.Message): - __slots__ = ("condition",) + __slots__ = () CONDITION_FIELD_NUMBER: _ClassVar[int] condition: ColumnOrExpression def __init__( @@ -255,7 +220,7 @@ class FilterEdges(_message.Message): ) -> None: ... class FilterVertices(_message.Message): - __slots__ = ("condition",) + __slots__ = () CONDITION_FIELD_NUMBER: _ClassVar[int] condition: ColumnOrExpression def __init__( @@ -263,19 +228,13 @@ class FilterVertices(_message.Message): ) -> None: ... class Find(_message.Message): - __slots__ = ("pattern",) + __slots__ = () PATTERN_FIELD_NUMBER: _ClassVar[int] pattern: str def __init__(self, pattern: _Optional[str] = ...) -> None: ... class LabelPropagation(_message.Message): - __slots__ = ( - "algorithm", - "max_iter", - "use_local_checkpoints", - "checkpoint_interval", - "storage_level", - ) + __slots__ = () ALGORITHM_FIELD_NUMBER: _ClassVar[int] MAX_ITER_FIELD_NUMBER: _ClassVar[int] USE_LOCAL_CHECKPOINTS_FIELD_NUMBER: _ClassVar[int] @@ -296,7 +255,7 @@ class LabelPropagation(_message.Message): ) -> None: ... class PageRank(_message.Message): - __slots__ = ("reset_probability", "source_id", "max_iter", "tol") + __slots__ = () RESET_PROBABILITY_FIELD_NUMBER: _ClassVar[int] SOURCE_ID_FIELD_NUMBER: _ClassVar[int] MAX_ITER_FIELD_NUMBER: _ClassVar[int] @@ -314,7 +273,7 @@ class PageRank(_message.Message): ) -> None: ... class ParallelPersonalizedPageRank(_message.Message): - __slots__ = ("reset_probability", "source_ids", "max_iter") + __slots__ = () RESET_PROBABILITY_FIELD_NUMBER: _ClassVar[int] SOURCE_IDS_FIELD_NUMBER: _ClassVar[int] MAX_ITER_FIELD_NUMBER: _ClassVar[int] @@ -329,7 +288,7 @@ class ParallelPersonalizedPageRank(_message.Message): ) -> None: ... class PowerIterationClustering(_message.Message): - __slots__ = ("k", "max_iter", "weight_col") + __slots__ = () K_FIELD_NUMBER: _ClassVar[int] MAX_ITER_FIELD_NUMBER: _ClassVar[int] WEIGHT_COL_FIELD_NUMBER: _ClassVar[int] @@ -344,23 +303,7 @@ class PowerIterationClustering(_message.Message): ) -> None: ... class Pregel(_message.Message): - __slots__ = ( - "agg_msgs", - "send_msg_to_dst", - "send_msg_to_src", - "checkpoint_interval", - "max_iter", - "additional_col_name", - "additional_col_initial", - "additional_col_upd", - "early_stopping", - "use_local_checkpoints", - "storage_level", - "stop_if_all_non_active", - "initial_active_expr", - "update_active_expr", - "skip_messages_from_non_active", - ) + __slots__ = () AGG_MSGS_FIELD_NUMBER: _ClassVar[int] SEND_MSG_TO_DST_FIELD_NUMBER: _ClassVar[int] SEND_MSG_TO_SRC_FIELD_NUMBER: _ClassVar[int] @@ -411,13 +354,7 @@ class Pregel(_message.Message): ) -> None: ... class ShortestPaths(_message.Message): - __slots__ = ( - "landmarks", - "algorithm", - "use_local_checkpoints", - "checkpoint_interval", - "storage_level", - ) + __slots__ = () LANDMARKS_FIELD_NUMBER: _ClassVar[int] ALGORITHM_FIELD_NUMBER: _ClassVar[int] USE_LOCAL_CHECKPOINTS_FIELD_NUMBER: _ClassVar[int] @@ -438,22 +375,13 @@ class ShortestPaths(_message.Message): ) -> None: ... class StronglyConnectedComponents(_message.Message): - __slots__ = ("max_iter",) + __slots__ = () MAX_ITER_FIELD_NUMBER: _ClassVar[int] max_iter: int def __init__(self, max_iter: _Optional[int] = ...) -> None: ... class SVDPlusPlus(_message.Message): - __slots__ = ( - "rank", - "max_iter", - "min_value", - "max_value", - "gamma1", - "gamma2", - "gamma6", - "gamma7", - ) + __slots__ = () RANK_FIELD_NUMBER: _ClassVar[int] MAX_ITER_FIELD_NUMBER: _ClassVar[int] MIN_VALUE_FIELD_NUMBER: _ClassVar[int] @@ -483,7 +411,7 @@ class SVDPlusPlus(_message.Message): ) -> None: ... class TriangleCount(_message.Message): - __slots__ = ("storage_level",) + __slots__ = () STORAGE_LEVEL_FIELD_NUMBER: _ClassVar[int] storage_level: StorageLevel def __init__(self, storage_level: _Optional[_Union[StorageLevel, _Mapping]] = ...) -> None: ... @@ -491,3 +419,18 @@ class TriangleCount(_message.Message): class Triplets(_message.Message): __slots__ = () def __init__(self) -> None: ... + +class KCore(_message.Message): + __slots__ = () + USE_LOCAL_CHECKPOINTS_FIELD_NUMBER: _ClassVar[int] + CHECKPOINT_INTERVAL_FIELD_NUMBER: _ClassVar[int] + STORAGE_LEVEL_FIELD_NUMBER: _ClassVar[int] + use_local_checkpoints: bool + checkpoint_interval: int + storage_level: StorageLevel + def __init__( + self, + use_local_checkpoints: _Optional[bool] = ..., + checkpoint_interval: _Optional[int] = ..., + storage_level: _Optional[_Union[StorageLevel, _Mapping]] = ..., + ) -> None: ... diff --git a/python/graphframes/graphframe.py b/python/graphframes/graphframe.py index d2cae7b9e..f9466d38a 100644 --- a/python/graphframes/graphframe.py +++ b/python/graphframes/graphframe.py @@ -456,6 +456,30 @@ def connectedComponents( storage_level=storage_level, ) + 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") + """ # noqa: E501 + return self._impl.k_core(checkpoint_interval, use_local_checkpoints, storage_level) + def labelPropagation( self, maxIter: int, diff --git a/python/tests/test_graphframes.py b/python/tests/test_graphframes.py index f5aa0eace..35af4da0b 100644 --- a/python/tests/test_graphframes.py +++ b/python/tests/test_graphframes.py @@ -553,3 +553,96 @@ def test_graph_grid_ising_model(spark: SparkSession): 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()