feat: perf

README.md

@@ -64,10 +64,14 @@ Input format is parquet by default; `--format csv` and `--format json` are suppo
 
 The work directory holds two kinds of data: the algorithms' parquet checkpoints (written and re-read between iterations) and DataFusion's spill files. Both are high-throughput and latency-sensitive, and on large graphs the algorithms can stream large amounts of data through them. On my experiments during processing billion-scale graphs with limited memory (~5-6GB limit) the process read from disk more than 200GBs. Point this at the fastest local storage available, preferably a dedicated NVMe or SSD. Network filesystems and spinning disks will make spill and checkpoint latency dominate runtime. The directory is created if missing; relative paths resolve against the current directory. `--max-temp-file` (env `GRAPHFRAMES_MAX_TEMP_FILE`, default `200G`) bounds the total size of the spill subdirectory. For graphs of the size of few billions edges and vertices the default may be not enough.
 
-### Libraray
+### Library
 
 TBD
 
+## Known Limitations
+
+- Some algorithms (K-Core, Label Propagation, etc.) are implemented with an assumption that the single node degree in the graph is always less than `i32::MAX`. In other words it is assummed that each node has less than ~2 billions neighbors / adjacent edges. While it is possible to support graph with hubs with more than 2B neighbors it will introduce a performance degradation for anything less (99.999% of all the graphs in the world). If you are working with a graph where a hub has degree gretater than 2B I would recommend to take a look on edges prunning technique if you want to use this library.
+
 ## References
 
 ### Core

src/algorithm/centrality/k_core.rs

@@ -1,11 +1,11 @@
 use crate::algorithm::pregel::{MessageDirection, pregel_default_msg, pregel_src};
-use crate::expressions::kcore_merge_expr;
+use crate::expressions::kcore_reduce;
 use crate::memory::CheckpointConfig;
 use crate::utils::symmetrize;
 use crate::{EDGE_DST, EDGE_SRC, GraphFrame, VERTEX_ID};
+use datafusion::arrow::datatypes::DataType;
 use datafusion::error::Result;
 use datafusion::execution::object_store::ObjectStoreUrl;
-use datafusion::functions_aggregate::array_agg::array_agg;
 use datafusion::functions_aggregate::count::count;
 use datafusion::object_store::path::Path;
 use datafusion::prelude::*;
@@ -108,14 +108,17 @@ impl<'a> KCoreBuilder<'a> {
             .select(vec![col(VERTEX_ID)])?
             .join(degrees, JoinType::Left, &[VERTEX_ID], &[EDGE_SRC], None)?
             .with_column(DEGREE, coalesce(vec![col(DEGREE), lit(0i64)]))?
-            .select(vec![col(VERTEX_ID), col(DEGREE)])?;
+            .select(vec![
+                col(VERTEX_ID),
+                cast(col(DEGREE), DataType::Int32).alias(DEGREE),
+            ])?;
 
         let prepared_graph = GraphFrame {
             vertices: prepared_vertices,
             edges: prepared_edges,
         };
 
-        let new_core = kcore_merge_expr(pregel_default_msg(), col(KCORE));
+        let new_core = coalesce(vec![pregel_default_msg(), lit(0)]);
 
         let mut pregel_builder = prepared_graph
             .pregel()
@@ -127,7 +130,7 @@ impl<'a> KCoreBuilder<'a> {
             // corrupt their estimates. Early stopping therefore relies on
             // the voting column alone, never on participation pruning.
             .add_message(pregel_src(KCORE), MessageDirection::SrcToDst)
-            .add_aggregate_expr(array_agg(pregel_default_msg()))
+            .add_aggregate_expr(kcore_reduce(pregel_default_msg()))
             // A vertex votes "still active" exactly while its core number
             // changed this iteration; the run stops once nobody changed.
             .with_vertex_voting("active", col(KCORE).not_eq(new_core.clone()))
@@ -156,7 +159,7 @@ impl GraphFrame {
 mod tests {
     use super::*;
     use crate::utils::collect_to_i64;
-    use datafusion::arrow::array::Int64Array;
+    use datafusion::arrow::array::{Int32Array, Int64Array};
     use std::collections::HashMap;
     use std::fs;
     use std::path::PathBuf;
@@ -224,13 +227,14 @@ mod tests {
                 .as_any()
                 .downcast_ref::<Int64Array>()
                 .unwrap();
+            // KCORE is Int32 (see `kcore_reduce`); vertex id stays Int64.
             let cores = batch
                 .column(1)
                 .as_any()
-                .downcast_ref::<Int64Array>()
+                .downcast_ref::<Int32Array>()
                 .unwrap();
             for i in 0..ids.len() {
-                map.insert(ids.value(i), cores.value(i));
+                map.insert(ids.value(i), cores.value(i) as i64);
             }
         }
         Ok(map)
@@ -512,9 +516,10 @@ mod tests {
             .run(&ctx, &output_uri, false)
             .await?;
 
+        // KCORE is Int32; cast to Int64 so `collect_to_i64` can read it.
         let result = read_result(&ctx, &output_uri)
             .await?
-            .select(vec![col(VERTEX_ID), col(KCORE)])?;
+            .select(vec![col(VERTEX_ID), cast(col(KCORE), DataType::Int64)])?;
         let values = collect_to_i64(&result.sort_by(vec![col(VERTEX_ID)])?, 1).await?;
         assert_eq!(values, &[2, 2, 2]);
         Ok(())

src/algorithm/community/classical_lp.rs

@@ -93,7 +93,7 @@ impl<'a> ClassicalLPBuilder<'a> {
             // 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)))
+            .add_aggregate_expr(most_common_by(pregel_default_msg(), lit(1.0f32)))
             .max_iterations(self.max_iter)
             .skip_dest_state()
             .with_checkpoint_store(self.checkpoint_config.store_url.clone())

src/expressions.rs

@@ -1,10 +1,10 @@
 mod common;
 mod finite_axpb;
 mod hll;
-mod kcore_merge;
+mod kcore_reduce;
 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 kcore_reduce::kcore_reduce;
 pub(crate) use most_common_by::most_common_by;

src/expressions/common.rs

@@ -1,4 +1,6 @@
-use datafusion::arrow::array::{Array, ArrayRef, BinaryArray, BinaryViewArray, Int64Array};
+use datafusion::arrow::array::{
+    Array, ArrayRef, BinaryArray, BinaryViewArray, Int32Array, Int64Array,
+};
 use datafusion::error::{DataFusionError, Result};
 
 /// Helper for other UDFs: vertexId and edgeSrc / edgeDst are all i64
@@ -15,6 +17,21 @@ pub(crate) fn downcast_int64<'a>(
     })
 }
 
+/// Helper for other UDFs: with assumption degree < i32::MAX
+/// this one is a common thing.
+pub(crate) fn downcast_int32<'a>(
+    array: &'a ArrayRef,
+    fname: &str,
+    label: &str,
+) -> Result<&'a Int32Array> {
+    array.as_any().downcast_ref::<Int32Array>().ok_or_else(|| {
+        DataFusionError::Plan(format!(
+            "{fname} {label} argument must be Int32, got: {:?}",
+            array.data_type()
+        ))
+    })
+}
+
 /// Read-only accessor over either a [`BinaryArray`] or a [`BinaryViewArray`].
 ///
 /// Parquet spills round-trip `Binary` sketch columns as `BinaryView` (DataFusion