From ce03bf3bfd7772958215b6e5ff8ab592ff17c632 Mon Sep 17 00:00:00 2001 From: semyonsinchenko Date: Thu, 3 Jul 2025 17:28:16 +0200 Subject: [PATCH 1/2] Add `useLabelsAsComponents` configuration for connected components Introduced a new configuration option `spark.graphframes.useLabelsAsComponents` to allow connected components to use labels as components. Updated the logic in `ConnectedComponents` and added relevant test cases for both enabled and disabled configurations. --- .../sql/graphframes/GraphFramesConf.scala | 12 +++++++++++ .../graphframes/lib/ConnectedComponents.scala | 2 +- .../lib/ConnectedComponentsSuite.scala | 20 +++++++++++++++++++ 3 files changed, 33 insertions(+), 1 deletion(-) diff --git a/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala b/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala index 2de8d89c0..78203a9a5 100644 --- a/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala +++ b/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala @@ -6,6 +6,16 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.spark.storage.StorageLevel 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 will returned. + |""".stripMargin) + .version("0.9.0") + .booleanConf + .createWithDefault(true) + private val CONNECTED_COMPONENTS_ALGORITHM = SQLConf .buildConf("spark.graphframes.connectedComponents.algorithm") @@ -93,4 +103,6 @@ object GraphFramesConf { case _ => None } } + + def getUseLabelsAsComponents: Boolean = get(USE_LABELS_AS_COMPONENTS).get.toBoolean } diff --git a/src/main/scala/org/graphframes/lib/ConnectedComponents.scala b/src/main/scala/org/graphframes/lib/ConnectedComponents.scala index 5439c6214..a3dd60009 100644 --- a/src/main/scala/org/graphframes/lib/ConnectedComponents.scala +++ b/src/main/scala/org/graphframes/lib/ConnectedComponents.scala @@ -340,7 +340,7 @@ object ConnectedComponents extends Logging { vv(ATTR), when(ee(SRC).isNull, vv(ID)).otherwise(ee(SRC)).as(COMPONENT), col(ATTR + "." + ID).as(ID)) - val output = if (graph.hasIntegralIdType) { + val output = if (graph.hasIntegralIdType || !GraphFramesConf.getUseLabelsAsComponents) { indexedLabel .select(col(s"$ATTR.*"), col(COMPONENT)) .persist(intermediateStorageLevel) diff --git a/src/test/scala/org/graphframes/lib/ConnectedComponentsSuite.scala b/src/test/scala/org/graphframes/lib/ConnectedComponentsSuite.scala index 3811a1252..39f4d4cd6 100644 --- a/src/test/scala/org/graphframes/lib/ConnectedComponentsSuite.scala +++ b/src/test/scala/org/graphframes/lib/ConnectedComponentsSuite.scala @@ -22,6 +22,7 @@ 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.graphframes.GraphFrame._ import org.graphframes._ @@ -75,6 +76,25 @@ class ConnectedComponentsSuite extends SparkFunSuite with GraphFrameTestSparkCon assertComponents(components, expected) } + test("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() + val expected = Seq("a", "b", "c", "d", "e").map(Set(_)).toSet + assertComponents(components, expected) + } + + test("don't using labels as components") { + spark.conf.set("spark.graphframes.useLabelsAsComponents", "false") + 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) + spark.conf.set("spark.graphframes.useLabelsAsComponents", "true") + } + test("two connected vertices") { 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") From a521af3cc6b365dc94fd01cb6786a9c3dc86d299 Mon Sep 17 00:00:00 2001 From: semyonsinchenko Date: Thu, 3 Jul 2025 18:00:19 +0200 Subject: [PATCH 2/2] updates from comments --- .../org/apache/spark/sql/graphframes/GraphFramesConf.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala b/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala index 78203a9a5..3dace99e1 100644 --- a/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala +++ b/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala @@ -10,7 +10,7 @@ object GraphFramesConf { 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 will returned. + | DataFrame. If set to "false", randomly generated labels with the data type LONG will returned. |""".stripMargin) .version("0.9.0") .booleanConf