allow useLabelsAsComponents to be set locally on the ConnectedComponents instance (#632)

james-willis · web-flow · commit f3b3b208d2ca · 2025-07-15T11:07:57.000+02:00
* allow useLabelsAsComponents to be set locally on the ConnectedComponents instance

* allow useLabelsAsComponents to be set locally on the ConnectedComponents instance

* typo
diff --git a/graphframes-connect/src/main/protobuf/graphframes.proto b/graphframes-connect/src/main/protobuf/graphframes.proto
@@ -63,6 +63,7 @@ message ConnectedComponents {
   string algorithm = 1;
   int32 checkpoint_interval = 2;
   int32 broadcast_threshold = 3;
+  bool use_labels_as_components = 4;
 }
 
 message DropIsolatedVertices {}
diff --git a/graphframes-connect/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala b/graphframes-connect/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConnectUtils.scala
@@ -102,6 +102,7 @@ object GraphFramesConnectUtils {
           .setAlgorithm(cc.getAlgorithm)
           .setCheckpointInterval(cc.getCheckpointInterval)
           .setBroadcastThreshold(cc.getBroadcastThreshold)
+          .setUseLabelsAsComponents(cc.getUseLabelsAsComponents)
           .run()
       }
       case MethodCase.DROP_ISOLATED_VERTICES => {
diff --git a/python/graphframes/classic/graphframe.py b/python/graphframes/classic/graphframe.py
@@ -206,12 +206,14 @@ def connectedComponents(
         algorithm: str = "graphframes",
         checkpointInterval: int = 2,
         broadcastThreshold: int = 1000000,
+        useLabelsAsComponents: bool = False,
     ) -> DataFrame:
         jdf = (
             self._jvm_graph.connectedComponents()
             .setAlgorithm(algorithm)
             .setCheckpointInterval(checkpointInterval)
             .setBroadcastThreshold(broadcastThreshold)
+            .setUseLabelsAsComponents(useLabelsAsComponents)
             .run()
         )
         return DataFrame(jdf, self._spark)
diff --git a/python/graphframes/connect/graphframe_client.py b/python/graphframes/connect/graphframe_client.py
@@ -504,6 +504,7 @@ def connectedComponents(
         algorithm: str = "graphframes",
         checkpointInterval: int = 2,
         broadcastThreshold: int = 1000000,
+        useLabelsAsComponents: bool = False,
     ) -> DataFrame:
         class ConnectedComponents(LogicalPlan):
             def __init__(
@@ -513,13 +514,15 @@ def __init__(
                 algorithm: str,
                 checkpoint_interval: int,
                 broadcast_threshold: int,
+                use_labels_as_components: bool,
             ) -> 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
 
             def plan(self, session: SparkConnectClient) -> proto.Relation:
                 graphframes_api_call = GraphFrameConnect._get_pb_api_message(
@@ -530,6 +533,7 @@ def plan(self, session: SparkConnectClient) -> proto.Relation:
                         algorithm=self.algorithm,
                         checkpoint_interval=self.checkpoint_interval,
                         broadcast_threshold=self.broadcast_threshold,
+                        use_labels_as_components=self.use_labels_as_components,
                     )
                 )
                 plan = self._create_proto_relation()
@@ -543,6 +547,7 @@ def plan(self, session: SparkConnectClient) -> proto.Relation:
                 algorithm,
                 checkpointInterval,
                 broadcastThreshold,
+                useLabelsAsComponents,
             ),
             self._spark,
         )
diff --git a/python/graphframes/connect/proto/graphframes_pb2.py b/python/graphframes/connect/proto/graphframes_pb2.py
diff --git a/python/graphframes/connect/proto/graphframes_pb2.pyi b/python/graphframes/connect/proto/graphframes_pb2.pyi
@@ -147,18 +147,26 @@ class BFS(_message.Message):
     ) -> None: ...
 
 class ConnectedComponents(_message.Message):
-    __slots__ = ("algorithm", "checkpoint_interval", "broadcast_threshold")
+    __slots__ = (
+        "algorithm",
+        "checkpoint_interval",
+        "broadcast_threshold",
+        "use_labels_as_components",
+    )
     ALGORITHM_FIELD_NUMBER: _ClassVar[int]
     CHECKPOINT_INTERVAL_FIELD_NUMBER: _ClassVar[int]
     BROADCAST_THRESHOLD_FIELD_NUMBER: _ClassVar[int]
+    USE_LABELS_AS_COMPONENTS_FIELD_NUMBER: _ClassVar[int]
     algorithm: str
     checkpoint_interval: int
     broadcast_threshold: int
+    use_labels_as_components: bool
     def __init__(
         self,
         algorithm: _Optional[str] = ...,
         checkpoint_interval: _Optional[int] = ...,
         broadcast_threshold: _Optional[int] = ...,
+        use_labels_as_components: bool = ...,
     ) -> None: ...
 
 class DropIsolatedVertices(_message.Message):
diff --git a/python/graphframes/graphframe.py b/python/graphframes/graphframe.py
@@ -267,6 +267,7 @@ def connectedComponents(
         algorithm: str = "graphframes",
         checkpointInterval: int = 2,
         broadcastThreshold: int = 1000000,
+        useLabelsAsComponents: bool = False,
     ) -> DataFrame:
         """
         Computes the connected components of the graph.
@@ -278,13 +279,16 @@ def connectedComponents(
         :param checkpointInterval: checkpoint interval in terms of number of iterations (default: 2)
         :param broadcastThreshold: broadcast threshold in propagating component assignments
           (default: 1000000)
+        :param useLabelsAsComponents: if True, uses the vertex labels as components, otherwise will
+          use longs
 
         :return: DataFrame with new vertices column "component"
         """
         return self._impl.connectedComponents(
             algorithm=algorithm,
             checkpointInterval=checkpointInterval,
             broadcastThreshold=broadcastThreshold,
+            useLabelsAsComponents=useLabelsAsComponents,
         )
 
     def labelPropagation(self, maxIter: int) -> DataFrame:
diff --git a/python/tests/conftest.py b/python/tests/conftest.py
@@ -17,7 +17,7 @@
 
     def is_remote() -> bool:
         return False
-    
+
 spark_major_version = __version__[:1]
 scala_version = os.environ.get("SCALA_VERSION", "2.12" if __version__ < "4" else "2.13")
 
diff --git a/python/tests/test_graphframes.py b/python/tests/test_graphframes.py
@@ -333,6 +333,7 @@ def test_connected_components_friends(examples, spark):
         g.connectedComponents(checkpointInterval=0),
         g.connectedComponents(checkpointInterval=10),
         g.connectedComponents(algorithm="graphx"),
+        g.connectedComponents(useLabelsAsComponents=True),
     ]
     for c in comps_tests:
         assert c.groupBy("component").count().count() == 2
diff --git a/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala b/src/main/scala/org/apache/spark/sql/graphframes/GraphFramesConf.scala
@@ -14,7 +14,7 @@ object GraphFramesConf {
           |""".stripMargin)
       .version("0.9.0")
       .booleanConf
-      .createWithDefault(true)
+      .createOptional
 
   private val CONNECTED_COMPONENTS_ALGORITHM =
     SQLConf
@@ -104,5 +104,8 @@ object GraphFramesConf {
     }
   }
 
-  def getUseLabelsAsComponents: Boolean = get(USE_LABELS_AS_COMPONENTS).get.toBoolean
+  def getUseLabelsAsComponents: Option[Boolean] = get(USE_LABELS_AS_COMPONENTS) match {
+    case Some(use) => Some(use.toBoolean)
+    case _ => None
+  }
 }
diff --git a/src/main/scala/org/graphframes/lib/ConnectedComponents.scala b/src/main/scala/org/graphframes/lib/ConnectedComponents.scala
@@ -31,6 +31,7 @@ import org.graphframes.WithBroadcastThreshold
 import org.graphframes.WithCheckpointInterval
 import org.graphframes.WithIntermediateStorageLevel
 import org.graphframes.WithMaxIter
+import org.graphframes.WithUseLabelsAsComponents
 
 import java.io.IOException
 import java.math.BigDecimal
@@ -52,6 +53,7 @@ class ConnectedComponents private[graphframes] (private val graph: GraphFrame)
     with WithCheckpointInterval
     with WithBroadcastThreshold
     with WithIntermediateStorageLevel
+    with WithUseLabelsAsComponents
     with WithMaxIter {
 
   setAlgorithm(GraphFramesConf.getConnectedComponentsAlgorithm.getOrElse(ALGO_GRAPHFRAMES))
@@ -61,6 +63,8 @@ class ConnectedComponents private[graphframes] (private val graph: GraphFrame)
     GraphFramesConf.getConnectedComponentsBroadcastThreshold.getOrElse(broadcastThreshold))
   setIntermediateStorageLevel(
     GraphFramesConf.getConnectedComponentsStorageLevel.getOrElse(intermediateStorageLevel))
+  setUseLabelsAsComponents(
+    GraphFramesConf.getUseLabelsAsComponents.getOrElse(useLabelsAsComponents))
 
   /**
    * Runs the algorithm.
@@ -72,6 +76,7 @@ class ConnectedComponents private[graphframes] (private val graph: GraphFrame)
       broadcastThreshold = broadcastThreshold,
       checkpointInterval = checkpointInterval,
       intermediateStorageLevel = intermediateStorageLevel,
+      useLabelsAsComponents = useLabelsAsComponents,
       maxIter = maxIter)
   }
 }
@@ -184,6 +189,7 @@ object ConnectedComponents extends Logging {
       broadcastThreshold: Int,
       checkpointInterval: Int,
       intermediateStorageLevel: StorageLevel,
+      useLabelsAsComponents: Boolean,
       maxIter: Option[Int]): DataFrame = {
     if (runInGraphX) {
       return runGraphX(graph, maxIter.getOrElse(Int.MaxValue))
@@ -340,7 +346,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 || !GraphFramesConf.getUseLabelsAsComponents) {
+      val output = if (graph.hasIntegralIdType || !useLabelsAsComponents) {
         indexedLabel
           .select(col(s"$ATTR.*"), col(COMPONENT))
           .persist(intermediateStorageLevel)
diff --git a/src/main/scala/org/graphframes/mixins.scala b/src/main/scala/org/graphframes/mixins.scala
@@ -122,3 +122,22 @@ private[graphframes] trait WithMaxIter {
     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
+}
diff --git a/src/test/scala/org/graphframes/lib/ConnectedComponentsSuite.scala b/src/test/scala/org/graphframes/lib/ConnectedComponentsSuite.scala
@@ -40,6 +40,7 @@ class ConnectedComponentsSuite extends SparkFunSuite with GraphFrameTestSparkCon
     assert(cc.getAlgorithm === "graphframes")
     assert(cc.getBroadcastThreshold === 1000000)
     assert(cc.getCheckpointInterval === 2)
+    assert(!cc.getUseLabelsAsComponents)
   }
 
   test("empty graph") {

Original file line number	Diff line number	Diff line change
`@@ -63,6 +63,7 @@ message ConnectedComponents {`
`63`	`63`	`string algorithm = 1;`
`64`	`64`	`int32 checkpoint_interval = 2;`
`65`	`65`	`int32 broadcast_threshold = 3;`
	`66`	`+ bool use_labels_as_components = 4;`
`66`	`67`	`}`
`67`	`68`
`68`	`69`	`message DropIsolatedVertices {}`
Original file line number	Diff line number	Diff line change
`@@ -102,6 +102,7 @@ object GraphFramesConnectUtils {`
`102`	`102`	`.setAlgorithm(cc.getAlgorithm)`
`103`	`103`	`.setCheckpointInterval(cc.getCheckpointInterval)`
`104`	`104`	`.setBroadcastThreshold(cc.getBroadcastThreshold)`
	`105`	`+ .setUseLabelsAsComponents(cc.getUseLabelsAsComponents)`
`105`	`106`	`.run()`
`106`	`107`	`}`
`107`	`108`	`case MethodCase.DROP_ISOLATED_VERTICES => {`
Original file line number	Diff line number	Diff line change
`@@ -333,6 +333,7 @@ def test_connected_components_friends(examples, spark):`
`333`	`333`	`g.connectedComponents(checkpointInterval=0),`
`334`	`334`	`g.connectedComponents(checkpointInterval=10),`
`335`	`335`	`g.connectedComponents(algorithm="graphx"),`
	`336`	`+ g.connectedComponents(useLabelsAsComponents=True),`
`336`	`337`	`]`
`337`	`338`	`for c in comps_tests:`
`338`	`339`	`assert c.groupBy("component").count().count() == 2`
Original file line number	Diff line number	Diff line change
`@@ -14,7 +14,7 @@ object GraphFramesConf {`
`14`	`14`	`\|""".stripMargin)`
`15`	`15`	`.version("0.9.0")`
`16`	`16`	`.booleanConf`
`17`		`- .createWithDefault(true)`
	`17`	`+ .createOptional`
`18`	`18`
`19`	`19`	`private val CONNECTED_COMPONENTS_ALGORITHM =`
`20`	`20`	`SQLConf`
`@@ -104,5 +104,8 @@ object GraphFramesConf {`
`104`	`104`	`}`
`105`	`105`	`}`
`106`	`106`
`107`		`- def getUseLabelsAsComponents: Boolean = get(USE_LABELS_AS_COMPONENTS).get.toBoolean`
	`107`	`+ def getUseLabelsAsComponents: Option[Boolean] = get(USE_LABELS_AS_COMPONENTS) match {`
	`108`	`+ case Some(use) => Some(use.toBoolean)`
	`109`	`+ case _ => None`
	`110`	`+ }`
`108`	`111`	`}`