graphframes · SemyonSinchenko · Jun 11, 2026 · Jun 9, 2026 · Jun 10, 2026 · Jun 10, 2026
diff --git a/connect/src/main/protobuf/graphframes.proto b/connect/src/main/protobuf/graphframes.proto
@@ -189,6 +189,7 @@ message Pregel {
   // Column names separated by comma
   optional string required_src_columns = 16;
   optional string required_dst_columns = 17;
+  optional string required_edge_columns = 18;
 }
 
 message ShortestPaths {

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
@@ -392,6 +392,12 @@ object GraphFramesConnectUtils {
           if (cols.nonEmpty) pregel = pregel.requiredDstColumns(cols.head, cols.tail: _*)
         }
 
+        if (pregelProto.hasRequiredEdgeColumns) {
+          val cols =
+            pregelProto.getRequiredEdgeColumns.split(",").map(_.trim).filter(_.nonEmpty).toSeq
+          if (cols.nonEmpty) pregel = pregel.requiredEdgeColumns(cols.head, cols.tail: _*)
+        }
+
         pregel.run()
       }
       case proto.GraphFramesAPI.MethodCase.SHORTEST_PATHS => {

diff --git a/core/src/main/scala/org/graphframes/lib/Pregel.scala b/core/src/main/scala/org/graphframes/lib/Pregel.scala
@@ -75,6 +75,13 @@ import scala.util.control.Breaks.breakable
  *     .run()
  * }}}
  *
+ * Migration note: pre 0.12 users that used edge columns in Pregel expressions should explicitly
+ * specify these columns using [[org.graphframes.lib.Pregel#requiredDstColumns]]. In 0.11 and
+ * earlier there was an unspecified bug that leads all the edge columns are always kept and
+ * persisted that created a bug memory pressure (2 columns in O(|E|) rows in the form of
+ * `StructType`). That behavior is considered as bug and starting from 0.12 edge columns are not
+ * kept by default.
+ *
  * @param graph
  *   The graph that Pregel will run on.
  * @see
@@ -106,6 +113,10 @@ class Pregel(val graph: GraphFrame)
   private val requiredSrcColumnsList = collection.mutable.ListBuffer.empty[String]
   private val requiredDstColumnsList = collection.mutable.ListBuffer.empty[String]
 
+  // Required columns for edges
+  // When empty, only src and dst are selected
+  private val requiredEdgeColumnsList = collection.mutable.ListBuffer.empty[String]
+
   /** Sets the max number of iterations (default: 10). */
   def setMaxIter(value: Int): this.type = {
     maxIter = value
@@ -345,6 +356,27 @@ class Pregel(val graph: GraphFrame)
     this
   }
 
+  /**
+   * Specifies which edge columns are required when constructing triplets.
+   *
+   * By default, only the source and destination ID columns from edges are included in triplets.
+   * Use this method to include additional edge properties that are needed by the sendMsgToSrc and
+   * sendMsgToDst expressions.
+   *
+   * @param colName
+   *   the first required edge column name
+   * @param colNames
+   *   additional required edge column names
+   * @see
+   *   [[requiredSrcColumns]] and [[requiredDstColumns]]
+   */
+  def requiredEdgeColumns(colName: String, colNames: String*): this.type = {
+    requiredEdgeColumnsList.clear()
+    requiredEdgeColumnsList += colName
+    requiredEdgeColumnsList ++= colNames
+    this
+  }
+
   /**
    * Defines how messages are aggregated after grouped by target vertex IDs.
    *
@@ -419,10 +451,18 @@ class Pregel(val graph: GraphFrame)
         "Optimization: skipping second join (dst state not required by message expressions)")
     }
 
-    val edges = graph.edges
-      .select(col(SRC).alias("edge_src"), col(DST).alias("edge_dst"), struct(col("*")).as(EDGE))
-      .repartition(col("edge_src"))
-      .persist(intermediateStorageLevel)
+    val edges = (if (requiredEdgeColumnsList.isEmpty) {
+                   graph.edges
+                     .select(col(SRC).alias("edge_src"), col(DST).alias("edge_dst"))
+                 } else {
+                   graph.edges
+                     .select(
+                       col(SRC).alias("edge_src"),
+                       col(DST).alias("edge_dst"),
+                       struct(
+                         requiredEdgeColumnsList.head,
+                         requiredEdgeColumnsList.tail.toSeq: _*).as(EDGE))
+                 }).repartition(col("edge_src")).persist(intermediateStorageLevel)
 
     var iteration = 1
 

diff --git a/core/src/main/scala/org/graphframes/lib/StructureAwareLabelPropagation.scala b/core/src/main/scala/org/graphframes/lib/StructureAwareLabelPropagation.scala
@@ -185,6 +185,7 @@ class StructureAwareLabelPropagation private[graphframes] (private val graph: Gr
       .setCheckpointInterval(checkpointInterval)
       .setUseLocalCheckpoints(useLocalCheckpoints)
       .setIntermediateStorageLevel(intermediateStorageLevel)
+      .requiredEdgeColumns(EDGE_WEIGHT_COL)
       .sendMsgToDst(struct(Pregel.src(LABEL_COL), Pregel.edge(EDGE_WEIGHT_COL)))
       .aggMsgs(aggregateMessages(Pregel.msg, vertices.schema(INITIAL_LABEL_COL).dataType))
       .withVertexColumn(

diff --git a/core/src/test/scala/org/graphframes/lib/PregelSuite.scala b/core/src/test/scala/org/graphframes/lib/PregelSuite.scala
@@ -400,6 +400,7 @@ class PregelSuite extends SparkFunSuite with GraphFrameTestSparkContext {
 
     // Only uses Pregel.edge("weight") - dst join should be skipped
     val resultDF = graph.pregel
+      .requiredEdgeColumns("weight")
       .setMaxIter(1)
       .withVertexColumn("received", lit(0L), coalesce(Pregel.msg, col("received")))
       .sendMsgToSrc(Pregel.edge("weight"))
@@ -425,6 +426,7 @@ class PregelSuite extends SparkFunSuite with GraphFrameTestSparkContext {
 
     // Only uses Pregel.edge("weight") - dst join should be skipped
     val result = graph.pregel
+      .requiredEdgeColumns("weight")
       .setMaxIter(1) // Single iteration to simplify testing
       .withVertexColumn("total", lit(0.0), coalesce(Pregel.msg, col("total")))
       .sendMsgToDst(Pregel.edge("weight"))
@@ -523,6 +525,7 @@ class PregelSuite extends SparkFunSuite with GraphFrameTestSparkContext {
     val graph = GraphFrame(vertices, edges)
 
     val result = graph.pregel
+      .requiredEdgeColumns("weights")
       .setMaxIter(1)
       .withVertexColumn("received", lit(0L), coalesce(Pregel.msg, col("received")))
       // Use dst.key to look up value in edge.weights map
@@ -545,6 +548,7 @@ class PregelSuite extends SparkFunSuite with GraphFrameTestSparkContext {
     val graph = GraphFrame(vertices, edges)
 
     val result = graph.pregel
+      .requiredEdgeColumns("values")
       .setMaxIter(1)
       .withVertexColumn("received", lit(0L), coalesce(Pregel.msg, col("received")))
       // Use dst.idx to index into edge.values array (element_at is 1-based)

diff --git a/docs/src/04-user-guide/10-pregel.md b/docs/src/04-user-guide/10-pregel.md
@@ -6,7 +6,7 @@ Pregel API is one of the core backbones of GraphFrames. It is based on the imple
 
 **NOTE**
 
-*Be aware, that returned `DataFrame` is persistent and should be unpersisted manually after processing to avoid memory leaks!*
+_Be aware, that returned `DataFrame` is persistent and should be unpersisted manually after processing to avoid memory leaks!_
 
 ---
 
@@ -52,14 +52,16 @@ Under the hood, the passed name of the column will be resolved to get the corres
 
 By default, all vertex columns are included when constructing triplets. For algorithms with large per-vertex state (e.g., cycle detection storing sequences, random walks), this can create huge intermediate datasets in memory.
 
-To reduce memory usage, you can specify only the columns that are actually needed using `requiredSrcColumns` and `requiredDstColumns`:
+To reduce memory usage, you can specify only the columns that are actually needed using `requiredSrcColumns`, `requiredDstColumns`, and `requiredEdgeColumns`:
 
 ```scala
 graph.pregel
   .withVertexColumn("distances", ...)
   .sendMsgToDst(Pregel.src("distances"))  // Only needs "distances" from source
+  .sendMsgToDst(Pregel.edge("weight"))    // Needs "weight" from edge
   .requiredSrcColumns("distances")         // Only include "distances" in src struct
   .requiredDstColumns("distances")         // Only include "distances" in dst struct
+  .requiredEdgeColumns("weight")           // Only include "weight" in edge struct
   .aggMsgs(...)
   .run()
 ```
@@ -70,13 +72,23 @@ In Python:
 graph.pregel \
   .withVertexColumn("distances", ...) \
   .sendMsgToDst(Pregel.src("distances")) \
+  .sendMsgToDst(Pregel.edge("weight")) \
   .required_src_columns("distances") \
   .required_dst_columns("distances") \
+  .required_edge_columns("weight") \
   .aggMsgs(...) \
   .run()
 ```
 
-The `id` column and the active flag column (if used) are always included automatically, so you don't need to specify them.
+The `id` column (and `src`/`dst` for edges) and the active flag column (if used) are always included automatically, so you don't need to specify them.
+
+---
+
+**NOTE**
+
+_Before 0.12.0, all edge columns were always included in triplets by default. This was caused by an unspecified bug that kept and persisted all edge columns as a `StructType`, creating significant memory pressure. Starting from 0.12.0, only the source and destination ID columns from edges are included by default. If your message expressions reference edge properties (e.g., `Pregel.edge("weight")`), you must explicitly specify them using `requiredEdgeColumns`._
+
+---
 
 ### Sending Messages
 
@@ -103,5 +115,5 @@ graph.pregel.aggMsgs(sum(Pregel.msg))
 GraphFrames Pregel API provides the following termination conditions:
 
 - **By a number of iterations.** Users can specify the maximum number of iterations with `setMaxIter(value: Int)`.
-- **In case of no new messages are sent.** User can say GF to terminate the computations if all the messages sent on the iteration are empty (`null`). To do this, user should specify `setEarlyStopping(value: Boolean)`. **Be careful, because the checking of nullity is a not free operation, but Apache Spark action!** So, for example, if messages cannot be empty, this condition should be set to `false`. For example, in algorithms like `ShortestPaths`, this condition should be set to `true`, but for algorithms like `PageRank`, this condition should be set to `false`  because the messages cannot be empty.
+- **In case of no new messages are sent.** User can say GF to terminate the computations if all the messages sent on the iteration are empty (`null`). To do this, user should specify `setEarlyStopping(value: Boolean)`. **Be careful, because the checking of nullity is a not free operation, but Apache Spark action!** So, for example, if messages cannot be empty, this condition should be set to `false`. For example, in algorithms like `ShortestPaths`, this condition should be set to `true`, but for algorithms like `PageRank`, this condition should be set to `false` because the messages cannot be empty.
 - **By vertex voting.** Users can specify the participation condition per vertex with `setInitialActiveVertexExpression(expression: Column` and `setUpdateActiveVertexExpression(expression: Column)`. In the case if `stopIfAllNonActiveVertices(value: Boolean)` is set to `true`, the computation will stop if all the vertices are inactive. This is useful for algorithms like `LabelPropagation`, when messages are always not `null`, but if no vertex changed a label on the last iteration, the computation should stop. **Be careful, because the checking of vertex status is a not free operation, but Apache Spark action!**
diff --git a/python/graphframes/connect/graphframes_client.py b/python/graphframes/connect/graphframes_client.py
@@ -60,6 +60,7 @@ def __init__(self, graph: "GraphFrameConnect") -> None:
         self._skip_messages_from_non_active = False
         self._required_src_columns: list[str] = []
         self._required_dst_columns: list[str] = []
+        self._required_edge_columns: list[str] = []
 
     def setMaxIter(self, value: int) -> Self:
         self._max_iter = value
@@ -146,6 +147,19 @@ def required_dst_columns(self, col_name: str, *col_names: str) -> Self:
         self._required_dst_columns = [col_name] + list(col_names)
         return self
 
+    def required_edge_columns(self, col_name: str, *col_names: str) -> Self:
+        """Specifies which edge columns are required when constructing triplets.
+
+        By default, only src and dst columns are included. Use this method to specify
+        additional edge columns that are needed by the sendMsgToSrc and sendMsgToDst
+        expressions.
+
+        :param col_name: the first required edge column name
+        :param col_names: additional required edge column names
+        """
+        self._required_edge_columns = [col_name] + list(col_names)
+        return self
+
     def run(self) -> DataFrame:
         @final
         class Pregel(LogicalPlan):
@@ -168,6 +182,7 @@ def __init__(
                 skip_message_from_non_active: bool,
                 required_src_columns: list[str],
                 required_dst_columns: list[str],
+                required_edge_columns: list[str],
                 vertices: DataFrame,
                 edges: DataFrame,
             ) -> None:
@@ -189,6 +204,7 @@ def __init__(
                 self.skip_message_from_non_active = skip_message_from_non_active
                 self.required_src_columns = required_src_columns
                 self.required_dst_columns = required_dst_columns
+                self.required_edge_columns = required_edge_columns
                 self.vertices = vertices
                 self.edges = edges
 
@@ -224,6 +240,9 @@ def plan(self, session: SparkConnectClient) -> proto.Relation:
                     required_dst_columns=",".join(self.required_dst_columns)
                     if self.required_dst_columns
                     else None,
+                    required_edge_columns=",".join(self.required_edge_columns)
+                    if self.required_edge_columns
+                    else None,
                 )
                 pb_message = pb.GraphFramesAPI(
                     vertices=dataframe_to_proto(self.vertices, session),
@@ -262,6 +281,7 @@ def plan(self, session: SparkConnectClient) -> proto.Relation:
                 skip_message_from_non_active=self._skip_messages_from_non_active,
                 required_src_columns=self._required_src_columns,
                 required_dst_columns=self._required_dst_columns,
+                required_edge_columns=self._required_edge_columns,
                 storage_level=self._storage_level,
                 vertices=self.graph._vertices,
                 edges=self.graph._edges,

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
@@ -468,6 +468,7 @@ class Pregel(_message.Message):
         "skip_messages_from_non_active",
         "required_src_columns",
         "required_dst_columns",
+        "required_edge_columns",
     )
     AGG_MSGS_FIELD_NUMBER: _ClassVar[int]
     SEND_MSG_TO_DST_FIELD_NUMBER: _ClassVar[int]
@@ -486,6 +487,7 @@ class Pregel(_message.Message):
     SKIP_MESSAGES_FROM_NON_ACTIVE_FIELD_NUMBER: _ClassVar[int]
     REQUIRED_SRC_COLUMNS_FIELD_NUMBER: _ClassVar[int]
     REQUIRED_DST_COLUMNS_FIELD_NUMBER: _ClassVar[int]
+    REQUIRED_EDGE_COLUMNS_FIELD_NUMBER: _ClassVar[int]
     agg_msgs: ColumnOrExpression
     send_msg_to_dst: _containers.RepeatedCompositeFieldContainer[ColumnOrExpression]
     send_msg_to_src: _containers.RepeatedCompositeFieldContainer[ColumnOrExpression]
@@ -503,6 +505,7 @@ class Pregel(_message.Message):
     skip_messages_from_non_active: bool
     required_src_columns: str
     required_dst_columns: str
+    required_edge_columns: str
     def __init__(
         self,
         agg_msgs: _Optional[_Union[ColumnOrExpression, _Mapping]] = ...,
@@ -522,6 +525,7 @@ class Pregel(_message.Message):
         skip_messages_from_non_active: _Optional[bool] = ...,
         required_src_columns: _Optional[str] = ...,
         required_dst_columns: _Optional[str] = ...,
+        required_edge_columns: _Optional[str] = ...,
     ) -> None: ...
 
 class ShortestPaths(_message.Message):

diff --git a/python/graphframes/lib/pregel.py b/python/graphframes/lib/pregel.py
@@ -276,6 +276,23 @@ def required_dst_columns(self, col_name: str, *col_names: str) -> Self:
         )
         return self
 
+    def required_edge_columns(self, col_name: str, *col_names: str) -> Self:
+        """Specifies which edge columns are required when constructing triplets.
+
+        By default, only src and dst columns are included from edges. Use this method to
+        specify additional edge columns that are needed by the sendMsgToSrc and sendMsgToDst
+        expressions.
+
+        :param col_name: the first required edge column name
+        :param col_names: additional required edge column names
+
+        See also :func:`required_src_columns` and :func:`required_dst_columns`
+        """
+        self._java_obj.requiredEdgeColumns(
+            col_name, _to_seq(self.graph._spark.sparkContext, col_names)
+        )
+        return self
+
     def run(self) -> DataFrame:
         """Runs the defined Pregel algorithm.
 

diff --git a/python/tests/test_graphframes.py b/python/tests/test_graphframes.py
@@ -511,6 +511,32 @@ def test_pregel_early_stopping(spark: SparkSession, args: PregelArguments) -> No
     _ = ranks.unpersist()
 
 
+def test_pregel_required_edge_columns(spark: SparkSession) -> None:
+    edges = spark.createDataFrame(
+        [(0, 1, 0.5), (1, 2, 1.0), (2, 0, 0.3)],
+        ["src", "dst", "weight"],
+    )
+    vertices = spark.createDataFrame([(0,), (1,), (2,)], ["id"])
+    graph = GraphFrame(vertices, edges)
+    pregel = graph.pregel
+
+    result = (
+        graph.pregel.setMaxIter(2)
+        .withVertexColumn(
+            "value",
+            sqlfunctions.lit(0.0),
+            sqlfunctions.coalesce(pregel.msg(), sqlfunctions.lit(0.0)),
+        )
+        .sendMsgToDst(pregel.src("value") + pregel.edge("weight"))
+        .aggMsgs(sqlfunctions.sum(pregel.msg()))
+        .required_edge_columns("weight")
+        .run()
+    )
+    assert "value" in result.columns
+    assert result.count() == 3
+    _ = result.unpersist()
+
+
 def _df_hasCols(df: DataFrame, vcols: list[str] = []) -> None:
     for c in vcols:
         assert c in df.columns, f"DataFrame missing column: {c}"