diff --git a/src/algorithm/centrality/pagerank.rs b/src/algorithm/centrality/pagerank.rs index b55eebc..0cef7e9 100644 --- a/src/algorithm/centrality/pagerank.rs +++ b/src/algorithm/centrality/pagerank.rs @@ -26,6 +26,10 @@ pub struct PageRankBuilder<'a> { reset_prob: f64, tol: f64, + /// Optional source vertex ids for personalized (GraphX-style) PageRank. + /// `None` (or an empty vec) selects the uniform seeding. + sources: Option>, + /// Storage options checkpoint_config: CheckpointConfig, } @@ -37,6 +41,7 @@ impl<'a> PageRankBuilder<'a> { max_iter: 0, reset_prob: 0.15, tol: 0.01, + sources: None, checkpoint_config: CheckpointConfig::default_local_fs(), } } @@ -56,6 +61,22 @@ impl<'a> PageRankBuilder<'a> { self } + /// Set the source vertices for personalized (GraphX-style) PageRank. + /// + /// When set and non-empty, only these vertices are seeded: their initial + /// rank and delta are `1.0`, every other vertex starts at `0.0`, and only + /// the sources are active in the first iteration (see the seeding logic in + /// [`PageRankBuilder::run`]). When unset (or empty), the uniform seeding + /// with `reset_prob` is used. + /// + /// Only `i64` ids are accepted: vertex ids in a `GraphFrame` must be + /// `Int64` (enforced by `GraphFrame::try_new`), so the type is validated + /// at compile time. + pub fn sources(mut self, sources: Vec) -> Self { + self.sources = Some(sources); + self + } + /// Set the object store URL pub fn with_checkpoint_store(mut self, store_url: ObjectStoreUrl) -> Self { self.checkpoint_config.store_url = store_url; @@ -123,18 +144,35 @@ impl<'a> PageRankBuilder<'a> { // stationary distribution, which is wrong whenever the graph has sinks. let new_delta = lit(alpha) * coalesce(vec![pregel_default_msg(), lit(0.0)]); - let pregel_builder = graph_with_degrees - .pregel() - .add_vertex_column( - PAGERANK, + // Seeding. The uniform case seeds both the rank and the delta with the reset + // probability, which reproduces the bootstrap of the GraphX dynamic PageRank + // (its initial message is resetProb / damping). The personalized case seeds + // the rank and the delta of the source vertices with 1.0, like the GraphX + // personalized vertex program; only the source vertices are active initially. + let (init_rank, init_delta, init_active) = match self.sources.as_deref() { + Some(sources) if !sources.is_empty() => { + let is_source = in_list( + col(VERTEX_ID), + sources.iter().map(|&id| lit(id)).collect::>(), + false, + ); + ( + when(is_source.clone(), lit(1.0)).otherwise(lit(0.0))?, + when(is_source.clone(), lit(1.0)).otherwise(lit(0.0))?, + is_source, + ) + } + _ => ( lit(reset_prob_per_vertices), - col(PAGERANK) + new_delta.clone(), - ) - .add_vertex_column( - PAGERANK_DELTA, lit(reset_prob_per_vertices), - new_delta.clone(), - ) + lit(true), + ), + }; + + let pregel_builder = graph_with_degrees + .pregel() + .add_vertex_column(PAGERANK, init_rank, col(PAGERANK) + new_delta.clone()) + .add_vertex_column(PAGERANK_DELTA, init_delta, new_delta.clone()) .add_vertex_column("out_degree", col("out_degree"), col("out_degree")) .add_message( pregel_src(PAGERANK_DELTA) / pregel_src("out_degree"), @@ -146,7 +184,7 @@ impl<'a> PageRankBuilder<'a> { // message generation every iteration, while voting only decides when to stop. .with_participation_column( "participates", - lit(true), + init_active, new_delta.clone().gt(lit(self.tol)), ) .skip_dest_state() @@ -483,4 +521,114 @@ mod tests { Ok(()) } + + /// Joins the calculated ranks with expected `(vertex_id, pagerank)` pairs + /// (LEFT join from the expected side, so a vertex missing from the output + /// yields a NULL rank) and asserts every rank is within `tol`. + async fn assert_ranks_close( + calculated: DataFrame, + expected: Vec<(i64, f64)>, + tol: f64, + ) -> Result<()> { + let expected_df = dataframe!( + "vertex_id" => expected.iter().map(|(id, _)| *id).collect::>(), + "expected_pr" => expected.iter().map(|(_, rank)| *rank).collect::>(), + )?; + let mismatches = expected_df + .join( + calculated, + JoinType::Left, + &["vertex_id"], + &[VERTEX_ID], + None, + )? + .with_column("difference", abs(col("expected_pr") - col(PAGERANK)))? + .filter(col("difference").gt(lit(tol)).or(col(PAGERANK).is_null()))?; + assert_eq!( + mismatches.count().await?, + 0, + "every vertex rank must match the expected value within {tol}" + ); + Ok(()) + } + + /// Personalized PageRank with GraphX-style seeding on a graph whose every + /// iteration can be verified with a calculator. + /// + /// Graph (a 2-cycle fed by vertex 1): `1 -> 2 -> 3 -> 2` + /// + /// With `sources = [1]`, `reset_prob = 0.15` (damping alpha = 0.85) and + /// `max_iter = 4`, only vertex 1 is seeded and active initially: + /// init: PR = (1, 0, 0), delta = (1, 0, 0) + /// it 1: v2 += 0.85 * 1 -> PR = (1, 0.85, 0) + /// it 2: v3 += 0.85 * 0.85 -> PR = (1, 0.85, 0.7225) + /// it 3: v2 += 0.85 * 0.7225 -> PR = (1, 1.464125, 0.7225) + /// it 4: v3 += 0.85 * 0.614125 -> PR = (1, 1.464125, 1.24450625) + /// The raw sum is 3.70863125, so the normalized result is + /// v1 = 0.269641259157, v2 = 0.394788508564, v3 = 0.335570232279. + #[tokio::test] + async fn test_pagerank_personalized_single_source() -> Result<()> { + let graph = create_graph(vec![1, 2, 3], vec![(1, 2), (2, 3), (3, 2)])?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("pagerank_personalized_single")?; + graph + .pagerank() + .max_iter(4) + .reset_prob(0.15) + .sources(vec![1]) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri, false) + .await?; + + let calculated_page_rank = ctx + .read_parquet(&output_uri, ParquetReadOptions::default()) + .await?; + assert_ranks_close( + calculated_page_rank, + vec![ + (1, 0.269641259157), + (2, 0.394788508564), + (3, 0.335570232279), + ], + 1e-9, + ) + .await?; + Ok(()) + } + + /// Personalized PageRank with multiple sources on the same 2-cycle graph. + /// + /// `sources = [2, 3]` seeds vertices 2 and 3 with rank/delta 1.0, while + /// vertex 1 starts at 0.0, never receives a message (no incoming edges) + /// and must end at exactly 0 after normalization. Vertices 2 and 3 stay + /// symmetric — each feeds the other with its full delta every iteration: + /// init: PR = (0, 1, 1), delta = (0, 1, 1) + /// it 1: PR = (0, 1.85, 1.85) + /// it 2: PR = (0, 2.5725, 2.5725) + /// it 3: PR = (0, 3.186625, 3.186625) + /// it 4: PR = (0, 3.70863125, 3.70863125) + /// so both normalize to exactly 0.5. + #[tokio::test] + async fn test_pagerank_personalized_multi_source() -> Result<()> { + let graph = create_graph(vec![1, 2, 3], vec![(1, 2), (2, 3), (3, 2)])?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("pagerank_personalized_multi")?; + graph + .pagerank() + .max_iter(4) + .reset_prob(0.15) + .sources(vec![2, 3]) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri, false) + .await?; + + let calculated_page_rank = ctx + .read_parquet(&output_uri, ParquetReadOptions::default()) + .await?; + assert_ranks_close( + calculated_page_rank, + vec![(1, 0.0), (2, 0.5), (3, 0.5)], + 1e-9, + ) + .await?; + Ok(()) + } } diff --git a/src/main.rs b/src/main.rs index 0743f43..5fa6bda 100644 --- a/src/main.rs +++ b/src/main.rs @@ -170,6 +170,14 @@ enum Command { #[arg(long, default_value_t = 0)] max_iter: usize, + + /// Personalized PageRank (GraphX-style): comma-separated source vertex + /// ids seeded with rank/delta 1.0, e.g. --personalized 1,4. Other + /// vertices start from 0.0 and only the sources are active initially. + /// Requires Int64 id columns: string ids are remapped to synthetic ids, + /// so personalized source ids cannot be resolved. + #[arg(long, value_delimiter = ',')] + personalized: Option>, }, /// Weakly connected components via randomized contraction. @@ -553,6 +561,36 @@ async fn read_data(ctx: &SessionContext, path: &str, format: Format) -> Result Result<()> { + let ctx = SessionContext::new(); + let vertices = read_data(&ctx, &common.vertices, common.format.clone()).await?; + let vid_type = vertices + .schema() + .field_with_unqualified_name(&common.id_col_name)? + .data_type() + .clone(); + if matches!( + vid_type, + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) { + return Err(DataFusionError::NotImplemented( + "personalized PageRank cannot be combined with string vertex ids: string ids \ + are remapped to synthetic Int64 ids and no reverse mapping is performed, so \ + the requested --personalized source ids would not refer to any vertex. \ + Provide Int64 id columns or run without --personalized." + .to_string(), + )); + } + Ok(()) +} + /// Renames the listed columns to their canonical names, keeping every other /// column (and its order) unchanged. fn rename_to_canonical(df: DataFrame, pairs: &[(&str, &str)]) -> Result { @@ -705,13 +743,23 @@ async fn main() -> Result<()> { tol, reset_prob, max_iter, + personalized, } => { + // Fast-fail before any setup work: string ids are remapped to + // synthetic Int64 ids, so personalized source ids cannot be resolved. + if personalized.is_some() { + ensure_personalized_ids_supported(&common).await?; + } let (ctx, g, ckpt) = setup(&common).await?; - let _ = g + let mut builder = g .pagerank() .reset_prob(reset_prob) .tol(tol) - .max_iter(max_iter) + .max_iter(max_iter); + if let Some(sources) = personalized { + builder = builder.sources(sources); + } + let _ = builder .set_checkpoint_dir(ckpt) .run(&ctx, &common.output, false) .await?;