Skip to content

feat: x2 better perf for WCC #56

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
SemyonSinchenko merged 1 commit into from
Aug 18, 2026
Merged

feat: x2 better perf for WCC #56

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 29 additions & 28 deletions src/algorithm/connectivity/connected_components.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<DataFrame> {
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),
])
}
Expand Down Expand Up @@ -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();
Expand All @@ -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);
Expand Down Expand Up @@ -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}");
}

Expand Down