diff --git a/src/algorithm/centrality/pagerank.rs b/src/algorithm/centrality/pagerank.rs index 9d33416..b55eebc 100644 --- a/src/algorithm/centrality/pagerank.rs +++ b/src/algorithm/centrality/pagerank.rs @@ -86,7 +86,10 @@ impl<'a> PageRankBuilder<'a> { let reset_prob_per_vertices = self.reset_prob; // PageRank needs the out-degree of each vertex to distribute its rank. - let vertices_with_degrees = self.graph.out_degrees().await?; + // `include_zero_deg = true` keeps zero-out-degree vertices (sinks) in the + // vertex set with `out_degree = 0`: dropping them would silently remove + // sinks from the result and lose the rank mass they should accumulate. + let vertices_with_degrees = self.graph.out_degrees(true).await?; // Create a temp graph that has vertices with degrees that will be used in Pregel Execution let graph_with_degrees = GraphFrame { @@ -215,6 +218,7 @@ impl GraphFrame { mod tests { use super::*; use crate::utils::create_ldbc_test_graph; + use crate::{EDGE_DST, EDGE_SRC}; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use std::fs; use std::path::PathBuf; @@ -304,6 +308,15 @@ mod tests { .await?; let ldbc_page_rank = get_ldbc_pr_results(test_dataset).await?; + // The result must contain exactly the same number of rows as the LDBC + // ground truth: a vertex silently dropped along the way (e.g. a + // zero-out-degree sink) would otherwise never be compared below. + assert_eq!( + calculated_page_rank.clone().count().await?, + ldbc_page_rank.clone().count().await?, + "result row count must match the LDBC ground truth row count" + ); + let comparison_df = calculated_page_rank .join( ldbc_page_rank, @@ -348,6 +361,15 @@ mod tests { .await?; let ldbc_page_rank = get_ldbc_pr_results(test_dataset).await?; + // The result must contain exactly the same number of rows as the LDBC + // ground truth: a vertex silently dropped along the way (e.g. a + // zero-out-degree sink) would otherwise never be compared below. + assert_eq!( + calculated_page_rank.clone().count().await?, + ldbc_page_rank.clone().count().await?, + "result row count must match the LDBC ground truth row count" + ); + let comparison_df = calculated_page_rank .join( ldbc_page_rank, @@ -364,4 +386,101 @@ mod tests { Ok(()) } + + /// Builds a tiny in-memory graph from vertex ids and (src, dst) edge pairs. + fn create_graph(vertices: Vec, edges: Vec<(i64, i64)>) -> Result { + let vertices_df = dataframe!(VERTEX_ID => Vec::::from(vertices))?; + let edges_df = dataframe!( + EDGE_SRC => Vec::::from(edges.iter().map(|(s, _)| *s).collect::>()), + EDGE_DST => Vec::::from(edges.iter().map(|(_, d)| *d).collect::>()), + )?; + Ok(GraphFrame { + vertices: vertices_df, + edges: edges_df, + }) + } + + /// PageRank must handle "sinks" — vertices with incoming edges but zero outgoing + /// edges. Such vertices receive rank from their in-neighbors every iteration but + /// never emit messages, so they are genuine members of the stationary distribution + /// and must appear in the output. + /// + /// Regression test for the bug where the intermediate graph used by the Pregel + /// engine was built from `GraphFrame::out_degrees(false)`, which aggregates *edges* + /// grouped by source. Vertices with out-degree 0 never appear in that aggregation, + /// so they were silently dropped from the vertex set, from every Pregel iteration, + /// and from the final output. The masses they should have accumulated were also + /// lost, skewing the ranks of the remaining vertices. + /// + /// Graph under test (vertex 4 is the sink, with two incoming edges): + /// 1 -> 2 -> 4 + /// 1 -> 3 -> 4 + /// + /// Expected values are the exact stationary distribution of PageRank with + /// alpha = 0.85 (reset_prob = 0.15), normalized to sum to 1 — the same convention + /// as the LDBC reference used by `test_pagerank_run` (verified analytically: + /// sinks either absorbing or teleporting mass uniformly yield the same result + /// after normalization). Values computed by solving the linear system + /// PR = (1 - alpha)/N + alpha * P^T PR. + #[tokio::test] + async fn test_pagerank_with_sink_vertices() -> Result<()> { + // Vertex 4 is a sink: two incoming edges, zero outgoing edges. + let graph = create_graph(vec![1, 2, 3, 4], vec![(1, 2), (1, 3), (2, 4), (3, 4)])?; + + let (ctx, checkpoint_dir, output_uri, _guard) = setup("pagerank_sink")?; + graph + .pagerank() + .max_iter(14) + .reset_prob(0.15) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri, false) + .await?; + + let calculated_page_rank = ctx + .read_parquet(&output_uri, ParquetReadOptions::default()) + .await?; + + // Every vertex of the graph must survive to the output — in particular the + // sink (vertex 4), which the bug drops entirely. + assert_eq!( + calculated_page_rank.clone().count().await?, + 4, + "output must contain all 4 vertices, including the sink (vertex 4)" + ); + + let expected = dataframe!( + "vertex_id" => vec![1i64, 2i64, 3i64, 4i64], + "expected_pr" => vec![ + 0.1375042970f64, // vertex 1 + 0.1959436232, // vertex 2 + 0.1959436232, // vertex 3 + 0.4706084565, // vertex 4 (sink) + ], + )?; + + // LEFT join from the *expected* side: a vertex missing from the calculated + // output yields a NULL pagerank, which the filter below catches explicitly. + let comparison_df = expected + .join( + calculated_page_rank, + JoinType::Left, + &["vertex_id"], + &[VERTEX_ID], + None, + )? + .with_column("difference", abs(col(PAGERANK) - col("expected_pr")))? + .filter( + col("difference") + .gt(lit(0.0015)) + .or(col(PAGERANK).is_null()), + )?; + + assert_eq!( + comparison_df.count().await?, + 0, + "every vertex (incl. the sink) must match the reference PageRank" + ); + + Ok(()) + } } diff --git a/src/algorithm/community/classical_lp.rs b/src/algorithm/community/classical_lp.rs index 7712419..ad63cff 100644 --- a/src/algorithm/community/classical_lp.rs +++ b/src/algorithm/community/classical_lp.rs @@ -215,6 +215,15 @@ mod tests { "expected all 8 vertices" ); + // The result must contain exactly the same number of rows as the LDBC + // ground truth: a vertex silently dropped along the way would otherwise + // never be compared below. + assert_eq!( + calculated.clone().count().await?, + expected.clone().count().await?, + "result row count must match the LDBC ground truth row count" + ); + // Exact match: no vertex may disagree with the reference community. let mismatches = calculated .clone() diff --git a/src/algorithm/connectivity/connected_components.rs b/src/algorithm/connectivity/connected_components.rs index 884aaae..de245c3 100644 --- a/src/algorithm/connectivity/connected_components.rs +++ b/src/algorithm/connectivity/connected_components.rs @@ -682,6 +682,15 @@ mod tests { .await?; let results = read_result(&ctx, &output_uri).await?; + // The result must contain exactly the same number of rows as the LDBC + // ground truth: a vertex silently dropped along the way would otherwise + // never be compared below. + assert_eq!( + results.clone().count().await?, + expected_components.clone().count().await?, + "result row count must match the LDBC ground truth row count" + ); + let diff = results .clone() .join( @@ -725,6 +734,15 @@ mod tests { .await?; let results = read_result(&ctx, &output_uri).await?; + // The result must contain exactly the same number of rows as the LDBC + // ground truth: a vertex silently dropped along the way would otherwise + // never be compared below. + assert_eq!( + results.clone().count().await?, + expected_components.clone().count().await?, + "result row count must match the LDBC ground truth row count" + ); + let diff = results .clone() .join( diff --git a/src/algorithm/connectivity/shortest_paths.rs b/src/algorithm/connectivity/shortest_paths.rs index e9db83d..d007fd3 100644 --- a/src/algorithm/connectivity/shortest_paths.rs +++ b/src/algorithm/connectivity/shortest_paths.rs @@ -414,6 +414,15 @@ mod tests { .read_parquet(&output_uri, ParquetReadOptions::default()) .await?; + // The result must contain exactly the same number of rows as the LDBC + // ground truth: a vertex silently dropped along the way would otherwise + // never be compared below. + assert_eq!( + results.clone().count().await?, + expected_distances.clone().count().await?, + "result row count must match the LDBC ground truth row count" + ); + let diff = results .join( expected_distances, diff --git a/src/lib.rs b/src/lib.rs index 244d85d..8ed0bbd 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -168,6 +168,14 @@ impl GraphFrame { /// - `VERTEX_ID`: The unique identifier of the vertex (derived from the destination of the edges). /// - `in_degree`: The count of incoming edges (in-degrees) for each vertex. /// + /// # Arguments + /// + /// * `include_zero_deg` - When `true`, the computed degrees are left-joined + /// back onto the graph's full vertex set and missing counts are coalesced + /// to `0`, so vertices with no incoming edges are still present with + /// `in_degree = 0`. When `false`, only vertices with at least one incoming + /// edge appear in the result. + /// /// # Returns /// An asynchronous function that returns: /// - `Ok(DataFrame)` containing the vertex IDs and their corresponding in-degrees. @@ -188,26 +196,35 @@ impl GraphFrame { /// ).unwrap(); /// /// let graph = GraphFrame::try_new(vertices, edges).unwrap(); - /// let edge_count = graph.in_degrees(); + /// let edge_count = graph.in_degrees(false); /// ``` - pub async fn in_degrees(&self) -> Result { + pub async fn in_degrees(&self, include_zero_deg: bool) -> Result { let df = self.edges.clone().aggregate( vec![col(EDGE_DST)], vec![count(col(EDGE_SRC)).alias("in_degree")], )?; - Ok(df.select(vec![col(EDGE_DST).alias(VERTEX_ID), col("in_degree")])?) + let df = df.select(vec![col(EDGE_DST).alias(VERTEX_ID), col("in_degree")])?; + self.with_zero_degree_vertices(df, "in_degree", include_zero_deg) } /// Computes the out-degrees for each vertex in the graph. /// /// This function calculates the out-degree of each vertex by counting the number of - /// outcoming edges. It returns a `DataFrame` + /// outgoing edges. It returns a `DataFrame` /// containing two columns: - /// - `VERTEX_ID`: The unique identifier of the vertex (derived from the destination of the edges). - /// - `in_degree`: The count of incoming edges (in-degrees) for each vertex. + /// - `VERTEX_ID`: The unique identifier of the vertex (derived from the source of the edges). + /// - `out_degree`: The count of outgoing edges (out-degrees) for each vertex. + /// + /// # Arguments + /// + /// * `include_zero_deg` - When `true`, the computed degrees are left-joined + /// back onto the graph's full vertex set and missing counts are coalesced + /// to `0`, so vertices with no outgoing edges (sinks) are still present + /// with `out_degree = 0`. When `false`, only vertices with at least one + /// outgoing edge appear in the result. /// /// # Returns /// An asynchronous function that returns: - /// - `Ok(DataFrame)` containing the vertex IDs and their corresponding in-degrees. + /// - `Ok(DataFrame)` containing the vertex IDs and their corresponding out-degrees. /// - `Err` if the aggregation or selection operation fails. /// /// # Example @@ -225,16 +242,46 @@ impl GraphFrame { /// ).unwrap(); /// /// let graph = GraphFrame::try_new(vertices, edges).unwrap(); - /// let edge_count = graph.in_degrees(); + /// let edge_count = graph.out_degrees(false); /// ``` - pub async fn out_degrees(&self) -> Result { + pub async fn out_degrees(&self, include_zero_deg: bool) -> Result { let df = self.edges.clone().aggregate( vec![col(EDGE_SRC)], vec![count(col(EDGE_DST)).alias("out_degree")], )?; - Ok(df.select(vec![col(EDGE_SRC).alias(VERTEX_ID), col("out_degree")])?) + let df = df.select(vec![col(EDGE_SRC).alias(VERTEX_ID), col("out_degree")])?; + self.with_zero_degree_vertices(df, "out_degree", include_zero_deg) } + /// When `include_zero_deg` is `true`, left-joins the computed degree counts + /// back onto the full vertex set and coalesces missing counts to `0`, so + /// vertices without a single incident edge in the counted direction still + /// appear in the result with degree `0` (e.g. sinks for out-degrees). + /// Otherwise the counts are returned as-is. + fn with_zero_degree_vertices( + &self, + degrees: DataFrame, + degree_col: &str, + include_zero_deg: bool, + ) -> Result { + if !include_zero_deg { + return Ok(degrees); + } + let degrees = degrees.with_column_renamed(VERTEX_ID, "__degree_vid")?; + let all_vertices = self.vertices.clone().select(vec![col(VERTEX_ID)])?; + Ok(all_vertices + .join( + degrees, + JoinType::Left, + &[VERTEX_ID], + &["__degree_vid"], + None, + )? + .select(vec![ + col(VERTEX_ID), + coalesce(vec![col(degree_col), lit(0)]).alias(degree_col), + ])?) + } /// Creates a symmetric graph by duplicating all edges in the reverse direction. /// For each edge (a,b) in the graph, adds the edge (b,a) if it doesn't exist. /// Any additional edge attributes are preserved in the reversed edges. @@ -417,7 +464,7 @@ mod tests { #[tokio::test] async fn test_in_degrees() -> Result<()> { let graph = create_test_graph()?; - let df = graph.in_degrees().await?; + let df = graph.in_degrees(false).await?; let batches = df.collect().await?; let mut degree_map = HashMap::new(); @@ -455,7 +502,7 @@ mod tests { #[tokio::test] async fn test_out_degrees() -> Result<()> { let graph = create_test_graph()?; - let df = graph.out_degrees().await?; + let df = graph.out_degrees(false).await?; let batches = df.collect().await?; let mut degree_map = HashMap::new(); @@ -490,6 +537,53 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_degrees_include_zero() -> Result<()> { + // 1 is a pure source (no in-edges), 4 is a pure sink (no out-edges). + let vertices = dataframe!(VERTEX_ID => vec![1i64, 2i64, 3i64, 4i64])?; + let edges = dataframe!( + EDGE_SRC => vec![1i64, 1i64, 2i64, 3i64], + EDGE_DST => vec![2i64, 3i64, 4i64, 4i64], + )?; + let graph = GraphFrame { vertices, edges }; + + async fn degree_map(df: DataFrame) -> Result> { + let mut map = HashMap::new(); + for batch in df.collect().await? { + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let degrees = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..ids.len() { + map.insert(ids.value(i), degrees.value(i)); + } + } + Ok(map) + } + + // Without the flag, only vertices with at least one incident edge appear. + let in_deg = degree_map(graph.in_degrees(false).await?).await?; + assert_eq!(in_deg, HashMap::from([(2, 1), (3, 1), (4, 2)])); + + let out_deg = degree_map(graph.out_degrees(false).await?).await?; + assert_eq!(out_deg, HashMap::from([(1, 2), (2, 1), (3, 1)])); + + // With the flag, every vertex is present and missing degrees are zero. + let in_deg = degree_map(graph.in_degrees(true).await?).await?; + assert_eq!(in_deg, HashMap::from([(1, 0), (2, 1), (3, 1), (4, 2)])); + + let out_deg = degree_map(graph.out_degrees(true).await?).await?; + assert_eq!(out_deg, HashMap::from([(1, 2), (2, 1), (3, 1), (4, 0)])); + + Ok(()) + } + #[tokio::test] async fn test_triplets() -> Result<()> { let vertices = @@ -551,8 +645,8 @@ mod tests { assert_eq!(sym_edges, orig_edges * 2); // In and out degrees should be equal for all vertices in symmetric graph - let in_degrees = sym_graph.in_degrees().await?.collect().await?; - let out_degrees = sym_graph.out_degrees().await?.collect().await?; + let in_degrees = sym_graph.in_degrees(false).await?.collect().await?; + let out_degrees = sym_graph.out_degrees(false).await?.collect().await?; let mut in_degree_map = HashMap::new(); let mut out_degree_map = HashMap::new();