From 10715e28d9f7c450e74881bcd4acce8dc99a250f Mon Sep 17 00:00:00 2001 From: semyonsinchenko Date: Tue, 18 Aug 2026 17:03:46 +0200 Subject: [PATCH] feat: x2 better perf for WCC --- .../connectivity/connected_components.rs | 57 ++++++++++--------- 1 file changed, 29 insertions(+), 28 deletions(-) diff --git a/src/algorithm/connectivity/connected_components.rs b/src/algorithm/connectivity/connected_components.rs index de245c3..56a36cb 100644 --- a/src/algorithm/connectivity/connected_components.rs +++ b/src/algorithm/connectivity/connected_components.rs @@ -14,7 +14,7 @@ use crate::expressions::{axpb, finite_axpb}; use crate::memory::{CheckpointConfig, ParquetCheckpointer}; -use crate::utils::{GraphFramesConfig, scoped_ctx, symmetrize}; +use crate::utils::{GraphFramesConfig, scoped_ctx}; use crate::{EDGE_DST, EDGE_SRC, GraphFrame, VERTEX_ID}; use datafusion::dataframe::DataFrameWriteOptions; use datafusion::error::Result; @@ -44,18 +44,25 @@ const CC_NEW_COMPONENT: &str = "__cc_new_component"; /// `rep(v) = least(axpb(rA, v, rB), min over neighbours of axpb(rA, u, rB))`. /// /// Returns a 2-column frame `[CC_V, CC_REP]`. +/// +/// Perf: fusing edges symmetrization with representative compute here; +/// Union is faster than unnest (~20%) fn compute_cc_reps(edges: &DataFrame, r_a: i64, r_b: i64) -> Result { - edges - .clone() - .aggregate( - vec![col(EDGE_SRC)], - vec![min(finite_axpb(lit(r_a), col(EDGE_DST), lit(r_b))).alias(CC_NBR_REP)], - )? - .with_column(CC_SELF_REP, finite_axpb(lit(r_a), col(EDGE_SRC), lit(r_b)))? + let fwd = edges.clone().select(vec![ + col(EDGE_SRC).alias(CC_V), + finite_axpb(lit(r_a), col(EDGE_DST), lit(r_b)).alias(CC_NBR_REP), + ])?; + let rev = edges.clone().select(vec![ + col(EDGE_DST).alias(CC_V), + finite_axpb(lit(r_a), col(EDGE_SRC), lit(r_b)).alias(CC_NBR_REP), + ])?; + fwd.union(rev)? + .aggregate(vec![col(CC_V)], vec![min(col(CC_NBR_REP)).alias(CC_REP)])? + .with_column(CC_SELF_REP, finite_axpb(lit(r_a), col(CC_V), lit(r_b)))? .select(vec![ - col(EDGE_SRC).alias(CC_V), - when(col(CC_SELF_REP).lt(col(CC_NBR_REP)), col(CC_SELF_REP)) - .otherwise(col(CC_NBR_REP))? + col(CC_V), + when(col(CC_SELF_REP).lt(col(CC_REP)), col(CC_SELF_REP)) + .otherwise(col(CC_REP))? .alias(CC_REP), ]) } @@ -205,23 +212,10 @@ impl<'a> ConnectedComponentsBuilder<'a> { .vertices .clone() .select_columns(&vec![VERTEX_ID])?; - let original_edges = self - .graph_frame - .edges - .clone() - .select_columns(&vec![EDGE_SRC, EDGE_DST])?; // ---- prepare: drop self-loops, symmetrize, deduplicate ---- let run_id = Uuid::new_v4().to_string(); log::info!("start WCC with run-id {run_id}"); - // it may look conter-intuitive - // but IRL distinct is more expensive here. - // - // With distinct: full-scan over huge edges, slightly less on the first iter; - // Without distinct: fast flow here, slightly slower on the first iter. - // - // TODO: benchmark it and see which one is better IRL, not in theory - let prepared_edges = symmetrize(&original_edges, false)?; let ckpt_base = self.checkpoint_config.dir.clone().join(run_id.clone()); let store_url = self.checkpoint_config.store_url.clone(); @@ -238,10 +232,17 @@ impl<'a> ConnectedComponentsBuilder<'a> { let mut back_ckptr = ParquetCheckpointer::new(store_url.clone(), ckpt_base.clone().join("backward")); - // Offload the initial edges; from here on `edges` is disk-backed. - let mut edges = edges_ckptr.push(ctx, "initial", prepared_edges).await?; + // Create a mutable link to the edges + let mut edges = self + .graph_frame + .edges + .clone() + .select_columns(&vec![EDGE_SRC, EDGE_DST])?; + + // It should be cheap: just scan parquet metadata; + // CSV/JSON are suffering but they are not a target case. let mut graph_size = edges.clone().count().await?; - log::info!("after preparation graph has {graph_size} edges"); + log::info!("starting iterations over the graph with {graph_size} edges"); // ---- forward pass: iterative randomized contraction ---- let mut rng = StdRng::seed_from_u64(self.random_seed); @@ -281,7 +282,7 @@ impl<'a> ConnectedComponentsBuilder<'a> { // first (while the old parquet still exists) avoids a NotFound. graph_size = edges.clone().count().await?; // Keep only the latest edges checkpoint on disk. - edges_ckptr.evict(ctx, 1).await?; + edges_ckptr.evict_all_but_latest_n(ctx, 1).await?; log::info!("cc forward iteration {iteration}, edges remaining: {graph_size}"); }