graphframes
diff --git a/‎python/graphframes/classic/graphframe.py‎
Lines changed: 2 additions & 0 deletions b/‎python/graphframes/classic/graphframe.py‎
Lines changed: 2 additions & 0 deletions
diff --git a/‎python/graphframes/connect/graphframe_client.py‎
Lines changed: 16 additions & 13 deletions b/‎python/graphframes/connect/graphframe_client.py‎
Lines changed: 16 additions & 13 deletions
@@ -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)
 
@@ -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,28 +514,29 @@ def __init__(
                 algorithm: str,
                 checkpoint_interval: int,
                 broadcast_threshold: int,
+                useLabelsAsComponents: 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.useLabelsAsComponents = (useLabelsAsComponents,)
 
-            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,
-                    )
+        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.useLabelsAsComponents,
                 )
-                plan = self._create_proto_relation()
-                plan.extension.Pack(graphframes_api_call)
-                return plan
+            )
+            plan = self._create_proto_relation()
+            plan.extension.Pack(graphframes_api_call)
+            return plan
 
         return _dataframe_from_plan(
             ConnectedComponents(
@@ -543,6 +545,7 @@ def plan(self, session: SparkConnectClient) -> proto.Relation:
                 algorithm,
                 checkpointInterval,
                 broadcastThreshold,
+                useLabelsAsComponents,
             ),
             self._spark,
         )