From 46cb47457424ec74465668a0ff2f0dfe84af0b86 Mon Sep 17 00:00:00 2001 From: semyonsinchenko Date: Tue, 28 Jul 2026 17:03:10 +0200 Subject: [PATCH] feat: classical label propagaion --- .github/workflows/ci.yml | 1 + .gitignore | 7 + README.md | 32 +- src/algorithm.rs | 1 + src/algorithm/community.rs | 1 + src/algorithm/community/classical_lp.rs | 404 +++++++++++++++ src/expressions.rs | 2 + src/expressions/common.rs | 60 ++- src/expressions/hll.rs | 56 +- src/expressions/most_common_by.rs | 486 ++++++++++++++++++ src/main.rs | 32 ++ .../test-cdlp-directed-CDLP.csv | 8 + .../test-cdlp-directed.e.csv | 18 + .../test-cdlp-directed.properties | 17 + .../test-cdlp-directed.v.csv | 8 + 15 files changed, 1071 insertions(+), 62 deletions(-) create mode 100644 src/algorithm/community/classical_lp.rs create mode 100644 src/expressions/most_common_by.rs create mode 100644 testing/data/ldbc/test-cdlp-directed/test-cdlp-directed-CDLP.csv create mode 100644 testing/data/ldbc/test-cdlp-directed/test-cdlp-directed.e.csv create mode 100644 testing/data/ldbc/test-cdlp-directed/test-cdlp-directed.properties create mode 100644 testing/data/ldbc/test-cdlp-directed/test-cdlp-directed.v.csv diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 63914aa..51e3f67 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -57,6 +57,7 @@ jobs: $BIN kcore $COMMON --output file:///tmp/ci_kcore/ $BIN hyperanf $COMMON --n-hops 3 --output file:///tmp/ci_hyperanf/ $BIN shortest-path $COMMON --landmarks 1 --output file:///tmp/ci_sp/ + $BIN classical-lp $COMMON --output file:///tmp/clp/ # Each command must have produced parquet output. for d in pr wcc mis kcore hyperanf sp; do diff --git a/.gitignore b/.gitignore index d47656c..45323e9 100644 --- a/.gitignore +++ b/.gitignore @@ -13,8 +13,15 @@ release/ *# gf_df_tmp* gf_checkpoints* +gf_workdir* test-test* .#* graph500* twitter_mpi* + +# Some things I'm using in development +*.gnuplot +*.txt +watch_size.sh +*.png \ No newline at end of file diff --git a/README.md b/README.md index 0a5f209..7ed472e 100644 --- a/README.md +++ b/README.md @@ -36,6 +36,7 @@ Available algorithms: - `kcore`: K-Core decomposition; - `hyperanf`: Approximate Neighbor Function; - `shortest-path`: Multi-Source Shortest Path; +- `classical-lp`: Classical (Raghavan) Label Propagaion Each has its own arguments — run `graphframes --help` for the full list. For what each algorithm computes, see [References](#references). @@ -69,10 +70,27 @@ TBD ## References -1. **Pregel**: _Malewicz, Grzegorz, et al. "Pregel: a system for large-scale graph processing." Proceedings of the 2010 ACM SIGMOD International Conference on Management of data. 2010._ - - **PageRank**: _Zadeh, R., et al. "Cme 323: Distributed algorithms and optimization, spring 2015." University Lecture (2015)._ - - **Multi Source Shortest Paths**: _Malewicz, Grzegorz, et al. "Pregel: a system for large-scale graph processing." Proceedings of the 2010 ACM SIGMOD International Conference on Management of data. 2010._ - - **K-Core**: _Mandal, Aritra, and Mohammad Al Hasan. "A distributed k-core decomposition algorithm on spark." 2017 IEEE International Conference on Big Data (Big Data). IEEE, 2017._ -2. **Weakly Connected Components**: _Bögeholz, Harald, Michael Brand, and Radu-Alexandru Todor. "In-database connected component analysis." 2020 IEEE 36th International Conference on Data Engineering (ICDE). IEEE, 2020._ -3. **Maximal Independent Set**: _Ghaffari, Mohsen. "An improved distributed algorithm for maximal independent set." Proceedings of the twenty-seventh annual ACM-SIAM symposium on Discrete algorithms. Society for Industrial and Applied Mathematics, 2016._ -4. **Approximate Neighbor Function**: _Boldi, Paolo, Marco Rosa, and Sebastiano Vigna. "HyperANF: Approximating the neighbourhood function of very large graphs on a budget." Proceedings of the 20th international conference on World Wide Web. 2011._ +### Core + +- **Pregel**: + - _Malewicz, Grzegorz, et al. "Pregel: a system for large-scale graph processing." Proceedings of the 2010 ACM SIGMOD International Conference on Management of data. 2010._ + - _Gonzalez, Joseph E., et al. "{GraphX}: Graph processing in a distributed dataflow framework." 11th USENIX symposium on operating systems design and implementation (OSDI 14). 2014._ + +### Centrality Metrics + +- **PageRank**: _Zadeh, R., et al. "Cme 323: Distributed algorithms and optimization, spring 2015." University Lecture (2015)._ +- **K-Core**: _Mandal, Aritra, and Mohammad Al Hasan. "A distributed k-core decomposition algorithm on spark." 2017 IEEE International Conference on Big Data (Big Data). IEEE, 2017._ +- **Approximate Neighbor Function**: _Boldi, Paolo, Marco Rosa, and Sebastiano Vigna. "HyperANF: Approximating the neighbourhood function of very large graphs on a budget." Proceedings of the 20th international conference on World Wide Web. 2011._ + +### Connectivity + +- **Multi Source Shortest Paths**: _Malewicz, Grzegorz, et al. "Pregel: a system for large-scale graph processing." Proceedings of the 2010 ACM SIGMOD International Conference on Management of data. 2010._ +- **Weakly Connected Components**: _Bögeholz, Harald, Michael Brand, and Radu-Alexandru Todor. "In-database connected component analysis." 2020 IEEE 36th International Conference on Data Engineering (ICDE). IEEE, 2020._ + +### Community Detection + +- **Classical Label Propagation**: _Raghavan, Usha Nandini, Réka Albert, and Soundar Kumara. "Near linear time algorithm to detect community structures in large-scale networks." Physical Review E—Statistical, Nonlinear, and Soft Matter Physics 76.3 (2007): 036106._ + +### Subgraphs + +- **Maximal Independent Set**: _Ghaffari, Mohsen. "An improved distributed algorithm for maximal independent set." Proceedings of the twenty-seventh annual ACM-SIAM symposium on Discrete algorithms. Society for Industrial and Applied Mathematics, 2016._ diff --git a/src/algorithm.rs b/src/algorithm.rs index b7a9f9f..95fe876 100644 --- a/src/algorithm.rs +++ b/src/algorithm.rs @@ -1,4 +1,5 @@ mod centrality; +mod community; mod connectivity; mod pregel; mod subgraph; diff --git a/src/algorithm/community.rs b/src/algorithm/community.rs index e69de29..9dda726 100644 --- a/src/algorithm/community.rs +++ b/src/algorithm/community.rs @@ -0,0 +1 @@ +mod classical_lp; diff --git a/src/algorithm/community/classical_lp.rs b/src/algorithm/community/classical_lp.rs new file mode 100644 index 0000000..cfbf88a --- /dev/null +++ b/src/algorithm/community/classical_lp.rs @@ -0,0 +1,404 @@ +use datafusion::{ + error::Result, execution::object_store::ObjectStoreUrl, object_store::path::Path, prelude::*, +}; + +use crate::{ + EDGE_DST, EDGE_SRC, GraphFrame, VERTEX_ID, + algorithm::pregel::{MessageDirection, pregel_default_msg, pregel_src}, + expressions::most_common_by, + memory::CheckpointConfig, + utils::symmetrize, +}; + +pub const COMMUNITY: &str = "community"; + +pub struct ClassicalLPBuilder<'a> { + graph: &'a GraphFrame, + directed: bool, + max_iter: usize, + checkpoint_config: CheckpointConfig, +} + +impl<'a> ClassicalLPBuilder<'a> { + pub fn new(graph: &'a GraphFrame) -> Self { + ClassicalLPBuilder { + graph: graph, + directed: true, // LDBC canonical implementation + max_iter: 10, // LDBC default + checkpoint_config: CheckpointConfig::default_local_fs(), + } + } + + pub fn max_iter(mut self, iter: usize) -> Self { + self.max_iter = iter; + self + } + + /// Whether to treat the graph as directed. + /// By default this library follows the LDBC semantic + /// and treat all edges as bidirectional for directed graph. + /// + /// For undirected graph symmetrization is skipped. + pub fn directed(mut self, directed: bool) -> Self { + self.directed = directed; + self + } + + pub fn with_checkpoint_store(mut self, store_url: ObjectStoreUrl) -> Self { + self.checkpoint_config.store_url = store_url; + self + } + + pub fn set_checkpoint_dir(mut self, dir: Path) -> Self { + self.checkpoint_config.dir = dir; + self + } + + pub async fn run( + self, + ctx: &SessionContext, + output: &str, + _include_debug_columns: bool, + ) -> Result { + let edges = if self.directed { + self.graph + .edges + .clone() + .select_columns(&[EDGE_SRC, EDGE_DST])? + } else { + symmetrize( + &self + .graph + .edges + .clone() + .select_columns(&[EDGE_SRC, EDGE_DST])?, + false, + )? + }; + + let vertices = self.graph.vertices.clone().select_columns(&[VERTEX_ID])?; + let g = GraphFrame { vertices, edges }; + + let pregel = g + .pregel() + // Label starts as the vertex's own id. Each iteration a vertex + // adopts the most common label among its neighbours; a vertex that + // received no message (e.g. isolated) keeps its current label. + .add_vertex_column( + COMMUNITY, + col(VERTEX_ID), + coalesce(vec![pregel_default_msg(), col(COMMUNITY)]), + ) + .add_message(pregel_src(COMMUNITY), MessageDirection::SrcToDst) + // Aggregate the neighbour labels carried by the *messages* + // (`__pregel_msg_msg`), not the `community` vertex column, which + // does not exist in the aggregated-messages frame. + .add_aggregate_expr(most_common_by(pregel_default_msg(), lit(1.0))) + .max_iterations(self.max_iter) + .skip_dest_state() + .with_checkpoint_store(self.checkpoint_config.store_url.clone()) + .set_checkpoint_dir(self.checkpoint_config.dir.clone()); + + let num_iterations = pregel.run(ctx, output, false).await?; + Ok(num_iterations) + } +} + +impl GraphFrame { + /// Constructs a [`ClassicalLPBuilder`] running classical (Raghavan-style) + /// Label Propagation. + pub fn classical_lp(&self) -> ClassicalLPBuilder<'_> { + ClassicalLPBuilder::new(self) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::utils::create_ldbc_test_graph; + use datafusion::arrow::array::Int64Array; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use std::fs; + use std::path::PathBuf; + use std::process::id; + use std::sync::atomic::{AtomicU64, Ordering}; + use url::Url; + + static COUNTER: AtomicU64 = AtomicU64::new(0); + + fn unique_temp_dir(label: &str) -> PathBuf { + let n = COUNTER.fetch_add(1, Ordering::SeqCst); + let dir = std::env::temp_dir().join(format!("graphframes_cdlp_test_{}_{n}_{label}", id())); + fs::create_dir_all(&dir).expect("failed to create unique temp dir"); + dir + } + + struct TempGuard(PathBuf); + impl Drop for TempGuard { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } + } + + fn setup(label: &str) -> Result<(SessionContext, Path, String, TempGuard)> { + let parent = unique_temp_dir(label); + let checkpoint_root = parent.join("checkpoints"); + let output_root = parent.join("output"); + fs::create_dir_all(&checkpoint_root).expect("failed to create checkpoint dir"); + fs::create_dir_all(&output_root).expect("failed to create output dir"); + let checkpoint_dir = Path::from_filesystem_path(&checkpoint_root) + .expect("checkpoint dir must be convertible to object_store path"); + let output_uri = Url::from_directory_path(&output_root) + .expect("output dir must be convertible to file:// URL") + .to_string(); + let ctx = SessionContext::new(); + Ok((ctx, checkpoint_dir, output_uri, TempGuard(parent))) + } + + /// Reads the LDBC reference CDLP communities (`-CDLP.csv`, matching + /// the `-WCC.csv` / `-BFS.csv` / `-PR.csv` convention): `vertex_id + /// expected_community`, space-delimited. + async fn get_ldbc_cdlp_results(dataset: &str) -> Result { + let ctx = SessionContext::new(); + let manifest_dir = env!("CARGO_MANIFEST_DIR"); + let schema = Schema::new(vec![ + Field::new("vertex_id", DataType::Int64, false), + Field::new("expected_community", DataType::Int64, false), + ]); + let path = format!( + "{}/testing/data/ldbc/{}/{}-CDLP.csv", + manifest_dir, dataset, dataset + ); + Ok(ctx + .read_csv( + &path, + CsvReadOptions::new() + .delimiter(b' ') + .has_header(false) + .schema(&schema), + ) + .await?) + } + + /// Correctness against the LDBC `test-cdlp-directed` reference. CDLP runs + /// on the undirected graph (`.directed(false)`), and crucially the + /// symmetrization does NOT deduplicate: a mutual edge `u <-> v` contributes + /// twice to each endpoint's label tally, which is what matches the LDBC + /// reference (see `symmetrize(.., false)`). Labels are exact integers, so + /// the check is an exact per-vertex match (no tolerance, unlike PageRank). + #[tokio::test] + async fn test_classical_lp_ldbc_directed() -> Result<()> { + let test_dataset = "test-cdlp-directed"; + let graph = create_ldbc_test_graph(test_dataset, false, false).await?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("classical_lp_ldbc")?; + graph + .classical_lp() + .directed(false) + .max_iter(5) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri, false) + .await?; + + let calculated = ctx + .read_parquet(&output_uri, ParquetReadOptions::default()) + .await? + .sort(vec![col(VERTEX_ID).sort(true, true)])? + .cache() + .await?; + let expected = get_ldbc_cdlp_results(test_dataset).await?; + + // Every vertex must be present. + assert_eq!( + calculated.clone().count().await?, + 8, + "expected all 8 vertices" + ); + + // Exact match: no vertex may disagree with the reference community. + let mismatches = calculated + .clone() + .join( + expected, + JoinType::Inner, + &[VERTEX_ID], + &["vertex_id"], + None, + )? + .filter(col(COMMUNITY).not_eq(col("expected_community")))? + .count() + .await?; + assert_eq!( + mismatches, 0, + "{mismatches} vertices have a wrong community" + ); + Ok(()) + } + + /// Builds a small `GraphFrame` from vertex ids and `(src, dst)` edges. + fn create_graph(vertices: Vec, edges: Vec<(i64, i64)>) -> Result { + let vertices_df = dataframe!(VERTEX_ID => vertices)?; + let (srcs, dsts): (Vec, Vec) = edges.into_iter().unzip(); + let edges_df = dataframe!(EDGE_SRC => srcs, EDGE_DST => dsts)?; + Ok(GraphFrame { + vertices: vertices_df, + edges: edges_df, + }) + } + + /// Smoke test: two disconnected triangles. Synchronous CDLP with + /// smallest-label tie-break converges each component to its minimum id, so + /// triangle {1,2,3} -> 1 and triangle {4,5,6} -> 4. Exercises the engine + /// end-to-end independently of the LDBC fixture. + #[tokio::test] + async fn test_classical_lp_two_triangles() -> Result<()> { + let graph = create_graph( + vec![1, 2, 3, 4, 5, 6], + vec![ + (1, 2), + (2, 1), + (2, 3), + (3, 2), + (1, 3), + (3, 1), + (4, 5), + (5, 4), + (5, 6), + (6, 5), + (4, 6), + (6, 4), + ], + )?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("two_triangles")?; + graph + .classical_lp() + .directed(false) + .max_iter(3) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri, false) + .await?; + + let out = ctx + .read_parquet(&output_uri, ParquetReadOptions::default()) + .await? + .select(vec![col(VERTEX_ID), col(COMMUNITY)])? + .sort(vec![col(VERTEX_ID).sort(true, true)])? + .collect() + .await?; + let ids = out[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let comms = out[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let want = [1i64, 1, 1, 4, 4, 4]; + for i in 0..want.len() { + assert_eq!(comms.value(i), want[i], "vertex {}", ids.value(i)); + } + Ok(()) + } + + /// No edges at all: every vertex is isolated and must keep its own id as its + /// community, regardless of iteration count (the `coalesce` fallback in the + /// update rule is what guarantees this). + #[tokio::test] + async fn test_classical_lp_isolated_vertices() -> Result<()> { + let graph = create_graph(vec![1, 2, 3, 4], vec![])?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("isolated")?; + graph + .classical_lp() + .directed(false) + .max_iter(5) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri, false) + .await?; + + let out = ctx + .read_parquet(&output_uri, ParquetReadOptions::default()) + .await? + .select(vec![col(VERTEX_ID), col(COMMUNITY)])? + .sort(vec![col(VERTEX_ID).sort(true, true)])? + .collect() + .await?; + let ids = out[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let comms = out[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..ids.len() { + assert_eq!( + comms.value(i), + ids.value(i), + "isolated vertex {} must keep its own id as community", + ids.value(i) + ); + } + Ok(()) + } + + /// Runs the chain `1-2-3-4-5-6` (undirected) for `max_iter` iterations and + /// returns the per-vertex community in id order. + async fn run_chain(max_iter: usize) -> Result> { + let graph = create_graph( + vec![1, 2, 3, 4, 5, 6], + vec![(1, 2), (2, 3), (3, 4), (4, 5), (5, 6)], + )?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup(&format!("chain_{max_iter}"))?; + graph + .classical_lp() + .directed(false) + .max_iter(max_iter) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri, false) + .await?; + let out = ctx + .read_parquet(&output_uri, ParquetReadOptions::default()) + .await? + .select(vec![col(VERTEX_ID), col(COMMUNITY)])? + .sort(vec![col(VERTEX_ID).sort(true, true)])? + .collect() + .await?; + let comms = out[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + Ok((0..comms.len()).map(|i| comms.value(i)).collect()) + } + + /// On a path the minimum label (1) propagates one hop per iteration. After + /// `max_iter = k`, vertex `k+1` carries label 1 while vertex `k+2` (when it + /// exists) does not yet — i.e. the label-1 frontier sits at depth `k`. + /// Increasing `max_iter` pushes that frontier one vertex further each time. + #[tokio::test] + async fn test_classical_lp_chain_propagation_depth() -> Result<()> { + for max_iter in 1..=4 { + let comms = run_chain(max_iter).await?; + // vertex (max_iter + 1) is at index max_iter and must be labelled 1 + assert_eq!( + comms[max_iter], + 1, + "max_iter={max_iter}: vertex {} should carry label 1, got {comms:?}", + max_iter + 1 + ); + // vertex (max_iter + 2), if present, must NOT yet carry label 1 + if max_iter + 1 < comms.len() { + assert_ne!( + comms[max_iter + 1], + 1, + "max_iter={max_iter}: vertex {} should not yet carry label 1, got {comms:?}", + max_iter + 2 + ); + } + } + Ok(()) + } +} diff --git a/src/expressions.rs b/src/expressions.rs index 519e9b7..bd6b36b 100644 --- a/src/expressions.rs +++ b/src/expressions.rs @@ -2,7 +2,9 @@ mod common; mod finite_axpb; mod hll; mod kcore_merge; +mod most_common_by; pub(crate) use finite_axpb::{axpb, finite_axpb}; pub(crate) use hll::{hll_long, hll_long_aggregate, hll_long_estimate, hll_long_union}; pub(crate) use kcore_merge::kcore_merge_expr; +pub(crate) use most_common_by::most_common_by; diff --git a/src/expressions/common.rs b/src/expressions/common.rs index b043f82..291791a 100644 --- a/src/expressions/common.rs +++ b/src/expressions/common.rs @@ -1,4 +1,4 @@ -use datafusion::arrow::array::{ArrayRef, Int64Array}; +use datafusion::arrow::array::{Array, ArrayRef, BinaryArray, BinaryViewArray, Int64Array}; use datafusion::error::{DataFusionError, Result}; /// Helper for other UDFs: vertexId and edgeSrc / edgeDst are all i64 @@ -14,3 +14,61 @@ pub(crate) fn downcast_int64<'a>( )) }) } + +/// Read-only accessor over either a [`BinaryArray`] or a [`BinaryViewArray`]. +/// +/// Parquet spills round-trip `Binary` sketch columns as `BinaryView` (DataFusion +/// reads parquet `Binary` as `BinaryView`), so every sketch-consuming UDF must +/// accept both representations. +pub(crate) enum BinaryLike<'a> { + Fixed(&'a BinaryArray), + View(&'a BinaryViewArray), +} + +impl<'a> BinaryLike<'a> { + pub(crate) fn len(&self) -> usize { + match self { + BinaryLike::Fixed(a) => a.len(), + BinaryLike::View(a) => a.len(), + } + } + + pub(crate) fn null_count(&self) -> usize { + match self { + BinaryLike::Fixed(a) => a.null_count(), + BinaryLike::View(a) => a.null_count(), + } + } + + pub(crate) fn is_null(&self, i: usize) -> bool { + match self { + BinaryLike::Fixed(a) => a.is_null(i), + BinaryLike::View(a) => a.is_null(i), + } + } + + pub(crate) fn value(&self, i: usize) -> &[u8] { + match self { + BinaryLike::Fixed(a) => a.value(i), + BinaryLike::View(a) => a.value(i), + } + } +} + +/// Helper for binary vectors +pub(crate) fn as_binary_like<'a>( + array: &'a ArrayRef, + fname: &str, + label: &str, +) -> Result> { + if let Some(a) = array.as_any().downcast_ref::() { + return Ok(BinaryLike::Fixed(a)); + } + if let Some(a) = array.as_any().downcast_ref::() { + return Ok(BinaryLike::View(a)); + } + Err(DataFusionError::Plan(format!( + "{fname} {label} argument must be Binary or BinaryView, got: {:?}", + array.data_type() + ))) +} diff --git a/src/expressions/hll.rs b/src/expressions/hll.rs index 90e3deb..7218bee 100644 --- a/src/expressions/hll.rs +++ b/src/expressions/hll.rs @@ -27,7 +27,7 @@ use std::sync::Arc; use datafusion::{ arrow::{ - array::{Array, ArrayRef, BinaryArray, BinaryViewArray, Float64Array}, + array::{Array, ArrayRef, BinaryArray, Float64Array}, datatypes::DataType, }, error::{DataFusionError, Result}, @@ -39,64 +39,12 @@ use datafusion::{ }; use datasketches::hll::{HllSketch, HllType}; +use crate::expressions::common::as_binary_like; use crate::expressions::common::downcast_int64; const DEFAULT_HLL_TYPE: HllType = HllType::Hll8; const DEFAULT_LG_K: u8 = 12; -/// Read-only accessor over either a [`BinaryArray`] or a [`BinaryViewArray`]. -/// -/// Parquet spills round-trip `Binary` sketch columns as `BinaryView` (DataFusion -/// reads parquet `Binary` as `BinaryView`), so every sketch-consuming UDF must -/// accept both representations. -enum BinaryLike<'a> { - Fixed(&'a BinaryArray), - View(&'a BinaryViewArray), -} - -impl<'a> BinaryLike<'a> { - fn len(&self) -> usize { - match self { - BinaryLike::Fixed(a) => a.len(), - BinaryLike::View(a) => a.len(), - } - } - - fn null_count(&self) -> usize { - match self { - BinaryLike::Fixed(a) => a.null_count(), - BinaryLike::View(a) => a.null_count(), - } - } - - fn is_null(&self, i: usize) -> bool { - match self { - BinaryLike::Fixed(a) => a.is_null(i), - BinaryLike::View(a) => a.is_null(i), - } - } - - fn value(&self, i: usize) -> &[u8] { - match self { - BinaryLike::Fixed(a) => a.value(i), - BinaryLike::View(a) => a.value(i), - } - } -} - -fn as_binary_like<'a>(array: &'a ArrayRef, fname: &str, label: &str) -> Result> { - if let Some(a) = array.as_any().downcast_ref::() { - return Ok(BinaryLike::Fixed(a)); - } - if let Some(a) = array.as_any().downcast_ref::() { - return Ok(BinaryLike::View(a)); - } - Err(DataFusionError::Plan(format!( - "{fname} {label} argument must be Binary or BinaryView, got: {:?}", - array.data_type() - ))) -} - #[derive(Debug, PartialEq, Eq, Hash)] pub(crate) struct HllLong { signature: Signature, diff --git a/src/expressions/most_common_by.rs b/src/expressions/most_common_by.rs new file mode 100644 index 0000000..12d781e --- /dev/null +++ b/src/expressions/most_common_by.rs @@ -0,0 +1,486 @@ +use std::sync::Arc; + +use crate::expressions::common::{as_binary_like, downcast_int64}; +use datafusion::arrow::array::{ArrayRef, Float64Array}; +use datafusion::arrow::datatypes::{DataType, Field}; +use datafusion::common::HashMap; +use datafusion::error::{DataFusionError, Result}; +use datafusion::logical_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion::logical_expr::utils::format_state_name; +use datafusion::logical_expr::{ + Accumulator, AggregateUDF, AggregateUDFImpl, Expr, Signature, Volatility, +}; +use datafusion::scalar::ScalarValue; + +#[derive(Debug)] +pub(crate) struct MostCommonByAccumulator { + sums: HashMap, +} + +impl MostCommonByAccumulator { + pub(crate) fn new() -> Self { + Self { + sums: HashMap::new(), + } + } +} + +fn se_map(m: &HashMap) -> Vec { + // We are assumming that a single node degree < i32::MAX + let n = m.len() as u32; + let mut buf = Vec::with_capacity(4 + 16usize * (n as usize)); + buf.extend_from_slice(&n.to_le_bytes()); + + for (&k, &v) in m.iter() { + buf.extend_from_slice(&k.to_le_bytes()); + buf.extend_from_slice(&v.to_le_bytes()); + } + + buf +} + +fn de_map_and_insert(buf: &[u8], map: &mut HashMap) -> Result<()> { + if buf.len() < 4 { + return Err(DataFusionError::Execution( + "most_common_by: corrupt state (length less 4)".to_string(), + )); + } + let n = u32::from_le_bytes(buf[0..4].try_into().unwrap()) as usize; + if buf.len() != 4 + 16 * n { + return Err(DataFusionError::Execution( + "most_common_by: corrupt state (length mismatch)".to_string(), + )); + } + + for i in 0..n { + let o = 4 + 16 * i; + let k = i64::from_le_bytes(buf[o..o + 8].try_into().unwrap()); + let v = f64::from_le_bytes(buf[o + 8..o + 16].try_into().unwrap()); + *map.entry(k).or_insert(0.0) += v; + } + + Ok(()) +} + +impl Accumulator for MostCommonByAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + let labels = downcast_int64(&values[0], "most_common_by", "first")?; + let weights = &values[1].as_any().downcast_ref::().ok_or( + DataFusionError::Execution( + "expected most_common_by secobd argument be f64 array".to_string(), + ), + )?; + + // we are assumming that there won't be nulls; + // it is internal function that assumet to aggregate neighbors: + // 1) nbr ID is not null by the DataFusion contract + // 2) weights are not null by the contract of this aggregator + for i in 0..labels.len() { + let l = labels.value(i); + let v = weights.value(i); + + let cur = self.sums.entry(l).or_insert(0.0); + *cur += v; + } + + Ok(()) + } + + fn evaluate(&mut self) -> Result { + // That can be a possible case. No neighbors we should return null. + if self.sums.is_empty() { + return Ok(ScalarValue::Int64(None)); + } + let mut max = -1; + let mut max_value = f64::MIN; + + for (k, v) in self.sums.iter() { + if v > &max_value { + max_value = *v; + max = *k; + } else if v == &max_value { + // tie-breaking: minimal key value: + // LDBC's Label Propagation semantics + if k < &max { + max_value = *v; + max = *k; + } + } + } + + Ok(ScalarValue::Int64(Some(max))) + } + + fn size(&self) -> usize { + size_of::() + self.sums.capacity() * (size_of::() + size_of::()) + } + + fn state(&mut self) -> Result> { + Ok(vec![ScalarValue::Binary(Some(se_map(&self.sums)))]) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + let v = as_binary_like(&states[0], "most_common_by", "argument")?; + + // we are assumming no null-maps here; + // null can be only an output of evaluate, not state + for i in 0..v.len() { + de_map_and_insert(v.value(i), &mut self.sums)?; + } + + Ok(()) + } +} + +#[derive(Debug, PartialEq, Eq, Hash)] +pub(crate) struct MostCommonBy { + signature: Signature, +} + +impl MostCommonBy { + pub(crate) fn new() -> Self { + Self { + signature: Signature::exact( + vec![DataType::Int64, DataType::Float64], + Volatility::Immutable, + ), + } + } +} + +impl AggregateUDFImpl for MostCommonBy { + fn name(&self) -> &str { + "most_common_by" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + if (arg_types.len() != 2) + || !matches!( + (&arg_types[0], &arg_types[1]), + (DataType::Int64, DataType::Float64) + ) + { + return Err(DataFusionError::Plan(format!( + "most_common_by expets exactly two arguments of types i64 and f64 but got {arg_types:?}" + ))); + } + + Ok(DataType::Int64) + } + + fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result> { + Ok(Box::new(MostCommonByAccumulator::new())) + } + + /// The intermediate state is a serialized `HashMap` carried as an + /// opaque `Binary` blob — a *different* type from the `Int64` final return + /// value. The default `state_fields` mirrors `return_type`, so it would + /// declare an `Int64` state column and the partial->final merge would never + /// receive the serialized maps, silently breaking multi-partition + /// aggregation. Override to declare the true `Binary` intermediate state. + fn state_fields(&self, args: StateFieldsArgs) -> Result>> { + Ok(vec![Arc::new(Field::new( + format_state_name(args.name, "value"), + DataType::Binary, + true, + ))]) + } +} + +/// Builds an [`Expr`] that applies `most_common_by(a, b)`. +pub(crate) fn most_common_by(a: Expr, b: Expr) -> Expr { + AggregateUDF::from(MostCommonBy::new()).call(vec![a, b]) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::{BinaryArray, Float64Array, Int64Array, RecordBatch}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::MemTable; + use datafusion::prelude::SessionConfig; + use datafusion::prelude::*; + + /// Direct accumulator: weights summed per label, argmax picked. + #[test] + fn test_accumulator_update_picks_argmax() { + let mut acc = MostCommonByAccumulator::new(); + // label 1 -> 5.0 (2+3), label 2 -> 10.0 -> winner 2 + acc.update_batch(&[ + Arc::new(Int64Array::from(vec![1i64, 2, 1])) as ArrayRef, + Arc::new(Float64Array::from(vec![2.0f64, 10.0, 3.0])) as ArrayRef, + ]) + .unwrap(); + assert_eq!(acc.evaluate().unwrap(), ScalarValue::Int64(Some(2))); + } + + /// No inputs -> null result (vertex received no messages this iteration). + #[test] + fn test_accumulator_empty_returns_null() { + let mut acc = MostCommonByAccumulator::new(); + assert_eq!(acc.evaluate().unwrap(), ScalarValue::Int64(None)); + } + + /// Tie on summed weight -> smallest label wins (LDBC LP semantics). + #[test] + fn test_accumulator_tie_picks_min_label() { + let mut acc = MostCommonByAccumulator::new(); + acc.update_batch(&[ + Arc::new(Int64Array::from(vec![5i64, 2])) as ArrayRef, + Arc::new(Float64Array::from(vec![5.0f64, 5.0])) as ArrayRef, + ]) + .unwrap(); + assert_eq!(acc.evaluate().unwrap(), ScalarValue::Int64(Some(2))); + } + + /// se_map/de_map round-trip preserves the map exactly. This is the test + /// that catches the LE/BE count mismatch (a swapped endianness yields a + /// bogus `n` and the length check rejects the blob). + #[test] + fn test_se_de_map_roundtrip_preserves_sums() { + let mut original: HashMap = HashMap::new(); + original.insert(1, 3.0); + original.insert(2, 5.5); + original.insert(-7, 0.25); + let bytes = se_map(&original); + let mut restored: HashMap = HashMap::new(); + de_map_and_insert(&bytes, &mut restored).unwrap(); + assert_eq!(restored.len(), original.len()); + for (k, v) in &original { + assert!( + (restored.get(k).copied().unwrap() - v).abs() < 1e-12, + "key {k}" + ); + } + } + + /// Merge semantics: `de_map_and_insert` must ADD values for shared keys, not + /// overwrite. This is the property that makes partial->final merging correct + /// (keeping only the local winner would not be associative across partitions). + #[test] + fn test_de_map_and_insert_accumulates_not_overwrites() { + let mut m: HashMap = HashMap::new(); + m.insert(1, 3.0); + let mut partial: HashMap = HashMap::new(); + partial.insert(1, 2.0); // shared key + partial.insert(5, 1.0); // new key + let bytes = se_map(&partial); + de_map_and_insert(&bytes, &mut m).unwrap(); + assert!( + (m.get(&1).copied().unwrap() - 5.0).abs() < 1e-12, + "shared key must sum" + ); + assert!( + (m.get(&5).copied().unwrap() - 1.0).abs() < 1e-12, + "new key must appear" + ); + } + + /// A corrupt/truncated state blob must surface as a query error, not a panic. + #[test] + fn test_de_map_and_insert_malformed_errors() { + let mut m: HashMap = HashMap::new(); + // Header claims 1 entry but there is no payload -> length mismatch. + assert!(de_map_and_insert(&1u32.to_le_bytes(), &mut m).is_err()); + // Empty buffer -> length < 4. + assert!(de_map_and_insert(&[], &mut m).is_err()); + } + + /// Direct partial->final merge: serialize b's state and fold it into a. The + /// correct result is the argmax over the UNION; a no-op `merge_batch` yields + /// null, and keeping only local winners yields a different key. + #[test] + fn test_accumulator_merge_unions_partial_states() { + let mut a = MostCommonByAccumulator::new(); + let mut b = MostCommonByAccumulator::new(); + // a: {1->3, 2->2} local winner 1 + a.update_batch(&[ + Arc::new(Int64Array::from(vec![1i64, 2])) as ArrayRef, + Arc::new(Float64Array::from(vec![3.0f64, 2.0])) as ArrayRef, + ]) + .unwrap(); + // b: {2->4} local winner 2 + b.update_batch(&[ + Arc::new(Int64Array::from(vec![2i64])) as ArrayRef, + Arc::new(Float64Array::from(vec![4.0f64])) as ArrayRef, + ]) + .unwrap(); + + // Serialize b's state and feed it back through merge_batch (final path). + let b_state = b.state().unwrap(); + let bytes = match &b_state[0] { + ScalarValue::Binary(Some(x)) => x.clone(), + other => panic!("expected Binary state, got {other:?}"), + }; + a.merge_batch(&[Arc::new(BinaryArray::from(vec![Some(bytes.as_slice())])) as ArrayRef]) + .unwrap(); + + // merged: {1->3, 2->6} -> winner 2 + assert_eq!(a.evaluate().unwrap(), ScalarValue::Int64(Some(2))); + } + + /// Single-group SQL aggregate with a clear winner. + #[tokio::test] + async fn test_most_common_by_clear_winner() -> Result<()> { + // sums: 1 -> 6.0, 2 -> 3.0 -> winner 1 + let df = dataframe!( + "g" => vec![0i64, 0, 0, 0], + "a" => vec![1i64, 2, 1, 2], + "b" => vec![5.0f64, 1.0, 1.0, 2.0], + )?; + let out = df + .aggregate( + vec![col("g")], + vec![most_common_by(col("a"), col("b")).alias("m")], + )? + .collect() + .await?; + let m = out[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(m.len(), 1); + assert_eq!(m.value(0), 1); + Ok(()) + } + + /// Tie at the SQL level -> minimal label. + #[tokio::test] + async fn test_most_common_by_tie_picks_min_label() -> Result<()> { + let df = dataframe!( + "a" => vec![5i64, 2], + "b" => vec![5.0f64, 5.0], + )?; + let out = df + .aggregate(vec![], vec![most_common_by(col("a"), col("b")).alias("m")])? + .collect() + .await?; + let m = out[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(m.value(0), 2); + Ok(()) + } + + /// GROUP BY: each group resolves to its own argmax independently. + #[tokio::test] + async fn test_most_common_by_grouped() -> Result<()> { + let df = dataframe!( + "g" => vec![0i64, 0, 0, 1, 1, 1], + "a" => vec![1i64, 2, 1, 10, 20, 20], + "b" => vec![1.0f64, 3.0, 1.0, 5.0, 2.0, 2.0], + )?; + let out = df + .aggregate( + vec![col("g")], + vec![most_common_by(col("a"), col("b")).alias("m")], + )? + .sort(vec![col("g").sort(true, true)])? + .collect() + .await?; + let g = out[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let m = out[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(g.len(), 2); + assert_eq!(g.value(0), 0); + assert_eq!(g.value(1), 1); + // g=0: 1->2.0, 2->3.0 -> 2 + assert_eq!(m.value(0), 2); + // g=1: 10->5.0, 20->4.0 -> 10 + assert_eq!(m.value(1), 10); + Ok(()) + } + + /// Multi-partition: force the partial->final path via two `MemTable` + /// partitions and `target_partitions(2)`. Each partition favors a different + /// label; only a correct merge (state() + merge_batch) yields the global + /// argmax. A missing/wrong `state_fields` errors the query; a no-op + /// `merge_batch` yields null; keeping only one partition yields key 1 not 2. + #[tokio::test] + async fn test_most_common_by_multi_partition_merge() { + let schema = Arc::new(Schema::new(vec![ + Field::new("g", DataType::Int64, false), + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Float64, false), + ])); + let mk = |labels: Vec, weights: Vec| { + let n = labels.len(); + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from_iter_values( + std::iter::repeat(0i64).take(n), + )) as ArrayRef, + Arc::new(Int64Array::from(labels)) as ArrayRef, + Arc::new(Float64Array::from(weights)) as ArrayRef, + ], + ) + .unwrap() + }; + // P1: 100 edges label 1 -> {1: 100}; P2: 150 edges label 2 -> {2: 150}. + let p1 = mk(vec![1i64; 100], vec![1.0f64; 100]); + let p2 = mk(vec![2i64; 150], vec![1.0f64; 150]); + + let table = MemTable::try_new(schema, vec![vec![p1], vec![p2]]).unwrap(); + let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(2)); + ctx.register_table("t", Arc::new(table)).unwrap(); + let out = ctx + .table("t") + .await + .unwrap() + .aggregate( + vec![col("g")], + vec![most_common_by(col("a"), col("b")).alias("m")], + ) + .unwrap() + .collect() + .await + .unwrap(); + let m = out[0] + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(m.len(), 1); + // merged {1:100, 2:150} -> 2 + assert_eq!(m.value(0), 2); + } + + /// `return_type` must accept exactly (Int64, Float64) and reject the rest. + #[test] + fn test_most_common_by_return_type_validation() { + let udf = MostCommonBy::new(); + assert_eq!( + udf.return_type(&[DataType::Int64, DataType::Float64]) + .unwrap(), + DataType::Int64 + ); + assert!(udf.return_type(&[DataType::Int64]).is_err()); + assert!( + udf.return_type(&[DataType::Int64, DataType::Int64]) + .is_err() + ); + assert!( + udf.return_type(&[DataType::Float64, DataType::Float64]) + .is_err() + ); + assert!( + udf.return_type(&[DataType::Int64, DataType::Float64, DataType::Int64]) + .is_err() + ); + } +} diff --git a/src/main.rs b/src/main.rs index 8286c1f..ed1cd7a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -170,6 +170,24 @@ enum Command { #[arg(long)] to_landmarks: bool, }, + + /// Classical Label Propagation (CDLP). + ClassicalLp { + #[command(flatten)] + common: CommonArgs, + + /// Maximum iterations (LDBC default). + #[arg(long, default_value_t = 10)] + max_iter: usize, + + /// Treat the graph as undirected. + /// CDLP is defined on the undirected graph: + /// the default (`false`) treat graph as directed + /// symmetrizes the edge set (LDBC semantics); + /// pass `true` to skip symmetrization. + #[arg(long, default_value_t = false)] + undirected: bool, + }, } #[derive(Parser, Debug)] @@ -397,6 +415,20 @@ async fn main() -> Result<()> { .run(&ctx, &common.output, false) .await?; } + Command::ClassicalLp { + common, + max_iter, + undirected, + } => { + let (ctx, g, ckpt) = setup(&common).await?; + let _ = g + .classical_lp() + .directed(undirected) + .max_iter(max_iter) + .set_checkpoint_dir(ckpt) + .run(&ctx, &common.output, false) + .await?; + } } Ok(()) diff --git a/testing/data/ldbc/test-cdlp-directed/test-cdlp-directed-CDLP.csv b/testing/data/ldbc/test-cdlp-directed/test-cdlp-directed-CDLP.csv new file mode 100644 index 0000000..f30545d --- /dev/null +++ b/testing/data/ldbc/test-cdlp-directed/test-cdlp-directed-CDLP.csv @@ -0,0 +1,8 @@ +1 1 +2 1 +3 1 +4 5 +5 4 +6 4 +7 4 +8 4 diff --git a/testing/data/ldbc/test-cdlp-directed/test-cdlp-directed.e.csv b/testing/data/ldbc/test-cdlp-directed/test-cdlp-directed.e.csv new file mode 100644 index 0000000..347bcc8 --- /dev/null +++ b/testing/data/ldbc/test-cdlp-directed/test-cdlp-directed.e.csv @@ -0,0 +1,18 @@ +1 2 +1 3 +1 7 +2 1 +2 3 +3 1 +3 2 +4 5 +4 6 +5 4 +5 6 +5 7 +6 5 +6 7 +7 5 +7 6 +7 8 +8 6 diff --git a/testing/data/ldbc/test-cdlp-directed/test-cdlp-directed.properties b/testing/data/ldbc/test-cdlp-directed/test-cdlp-directed.properties new file mode 100644 index 0000000..34b09c6 --- /dev/null +++ b/testing/data/ldbc/test-cdlp-directed/test-cdlp-directed.properties @@ -0,0 +1,17 @@ +# Filenames of graph on local filesystem +graph.test-cdlp-directed.vertex-file = test-cdlp-directed.v +graph.test-cdlp-directed.edge-file = test-cdlp-directed.e + +# Graph metadata for reporting purposes +graph.test-cdlp-directed.meta.vertices = 8 +graph.test-cdlp-directed.meta.edges = 18 + +# Properties describing the graph format +graph.test-cdlp-directed.directed = true + +# List of supported algorithms on the graph +graph.test-cdlp-directed.algorithms = cdlp + + +# Parameters for CDLP +graph.test-cdlp-directed.cdlp.max-iterations = 5 diff --git a/testing/data/ldbc/test-cdlp-directed/test-cdlp-directed.v.csv b/testing/data/ldbc/test-cdlp-directed/test-cdlp-directed.v.csv new file mode 100644 index 0000000..535d2b0 --- /dev/null +++ b/testing/data/ldbc/test-cdlp-directed/test-cdlp-directed.v.csv @@ -0,0 +1,8 @@ +1 +2 +3 +4 +5 +6 +7 +8