graphframes · SemyonSinchenko · Aug 27, 2026 · Aug 24, 2026 · Aug 24, 2026
diff --git a/core/src/main/scala/org/apache/spark/sql/graphframes/expressions/FiniteAXPlusB.scala b/core/src/main/scala/org/apache/spark/sql/graphframes/expressions/FiniteAXPlusB.scala
@@ -4,14 +4,12 @@ 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 {
+    extends TernaryExpression {
   override def dataType: DataType = LongType
 
   override protected def withNewChildrenInternal(
@@ -32,32 +30,40 @@ case class FiniteAXPlusB(first: Expression, second: Expression, third: Expressio
     val x = ctx.freshName("x")
     val b = ctx.freshName("b")
     val r = ctx.freshName("r")
+    val irrpoly = ctx.freshName("irrpoly")
+    val nullFlag = ctx.freshName("nullFlag")
 
     val aGenCode = first.genCode(ctx)
     val xGenCode = second.genCode(ctx)
     val bGenCode = third.genCode(ctx)
 
     ev.copy(code = code"""
-      ${aGenCode.code}
       ${xGenCode.code}
+      ${aGenCode.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;
+      long $irrpoly = 0x1bL;
+      boolean $nullFlag = false;
+
+      if (!${xGenCode.isNull}) {
+        long $x = ${xGenCode.value};
+        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;
+          }
         }
+      } else {
+        $nullFlag = true;
       }
-      boolean ${ev.isNull} = false;
+      boolean ${ev.isNull} = $nullFlag;
       long ${ev.value} = $r ^ $b;
     """)
   }

diff --git a/core/src/main/scala/org/graphframes/lib/RandomizedContraction.scala b/core/src/main/scala/org/graphframes/lib/RandomizedContraction.scala
@@ -29,6 +29,10 @@ import scala.util.Random
  *
  * The algorithm contracts the graph iteratively using random linear functions, until no edges
  * remain, then reconstructs the component identifiers.
+ *
+ * Improvements compared to the original paper:
+ *   - fusing symmetrization to the representative choosing
+ *   - remove the first heavy distinct op
  */
 private[graphframes] object RandomizedContraction extends Logging with Serializable {
   private val CHECKPOINT_NAME_PREFIX = "randomized-contraction"
@@ -39,11 +43,8 @@ private[graphframes] object RandomizedContraction extends Logging with Serializa
 
     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)
 
+    GraphFrame(vertices, edges)
   }
 
   def run(
@@ -126,10 +127,23 @@ private[graphframes] object RandomizedContraction extends Logging with Serializa
         stackA.push(rA)
         stackB.push(rB)
 
+        /// trick:
+        /// instead of symmetrization (heavy)
+        /// we sent bi-directional messages
+        ///
+        /// CSE will resolve this into 2 calls of axpb and one least call
+        /// per both directions;
+        ///
+        /// instead of scanning double-edges and 3 calls per both directions.
         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"))
+          .withColumn("msg", least(axpb(rA, col(SRC), rB), axpb(rA, col(DST), rB)))
+          .select(
+            explode(
+              array(
+                struct(col(SRC).alias("v"), col("msg")),
+                struct(col(DST).alias("v"), col("msg")))).alias("arr"))
+          .groupBy(col("arr.v").alias("v"))
+          .agg(min(col("arr.msg")).alias("rep"))
 
         // "free" checkpointing
         ccRepresentatives.write.parquet(tableName(iter))

diff --git a/dev/run_doc_benchmarks.py b/dev/run_doc_benchmarks.py
@@ -42,7 +42,7 @@ def main() -> None:
         (
             results_dir / "connected-components.json",
             "benchmarks/Jmh/run -rf json -p graphName=wiki-Talk -p useLocalCheckpoints=true "
-            "-p algorithm=graphframes,graphx -p broadcastThreshold=-1 "
+            "-p algorithm=randomized_contraction,graphx "
             "org.graphframes.benchmarks.ConnectedComponentsBenchmark",
         ),
         (

diff --git a/docs/src/04-user-guide/05-traversals.md b/docs/src/04-user-guide/05-traversals.md
@@ -230,6 +230,10 @@ _With GraphFrames 0.3.0 and later releases, the default Connected Components alg
 
 _Be aware, that returned `DataFrame` is persistent and should be unpersisted manually after processing to avoid memory leaks!_
 
+**NOTE**
+
+_From all the provided algorithms the "randomized_contraction" is the best one by all the characteristics. It is not the default only because of back-compatibility concerns._
+
 ---
 
 ### Python API
@@ -286,7 +290,7 @@ A DataFrame-native implementation based on randomized graph contraction, describ
 
 > 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.
 
-This algorithm iteratively contracts the graph using random linear functions until no edges remain, then reconstructs component identifiers in a reverse pass. It has similar convergence characteristics to `two_phase` (AQE mode) and performs comparably on benchmarks — slightly worse than `two_phase` with AQE, but significantly better than `two_phase` with manual skewed joins.
+This algorithm iteratively contracts the graph using random linear functions until no edges remain, then reconstructs component identifiers in a reverse pass. It has similar convergence characteristics to `two_phase` (AQE mode) and performs better on benchmarks than `two_phase` with AQE. This algorithm requires around 2x less amount of memory to perform compared to "two_phase"
 
 Unlike `two_phase`, `randomized_contraction` **always** produces random `Long` component IDs regardless of the input vertex ID type, unless `use_labels_as_components=True` is set.