perf: rollback collect list

.gitignore

@@ -17,6 +17,7 @@ gf_workdir*
 test-test*
 .#*
 graph500*
+wiki-Talk*
 twitter_mpi*
 
 

src/algorithm/centrality/k_core.rs

@@ -1,11 +1,12 @@
 use crate::algorithm::pregel::{MessageDirection, pregel_default_msg, pregel_src};
-use crate::expressions::kcore_reduce;
+use crate::expressions::kcore_merge_expr;
 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::*;
@@ -118,6 +119,10 @@ impl<'a> KCoreBuilder<'a> {
             edges: prepared_edges,
         };
 
+        // The aggregate spills only the reduced per-vertex core (O(|V|)), not
+        // the raw O(|E|) list: DataFusion splits the nested aggregate
+        // expression below, so the kcore_merge projection runs before the
+        // checkpoint write.
         let new_core = coalesce(vec![pregel_default_msg(), lit(0)]);
 
         let mut pregel_builder = prepared_graph
@@ -130,7 +135,12 @@ 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(kcore_reduce(pregel_default_msg()))
+            // Nested expression: DataFusion splits `kcore_merge(array_agg(...))`
+            // into the array_agg aggregate plus a kcore_merge projection on top,
+            // so the checkpoint spill writes only the reduced O(|V|) result.
+            // `kcore_merge` is uncapped by design (see kcore_merge.rs): the cap
+            // never binds in a degree-seeded run.
+            .add_aggregate_expr(kcore_merge_expr(array_agg(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()))
@@ -227,7 +237,7 @@ mod tests {
                 .as_any()
                 .downcast_ref::<Int64Array>()
                 .unwrap();
-            // KCORE is Int32 (see `kcore_reduce`); vertex id stays Int64.
+            // KCORE is Int32 (see `kcore_merge`); vertex id stays Int64.
             let cores = batch
                 .column(1)
                 .as_any()

src/algorithm/community/classical_lp.rs

@@ -5,10 +5,11 @@ use datafusion::{
 use crate::{
     EDGE_DST, EDGE_SRC, GraphFrame, VERTEX_ID,
     algorithm::pregel::{MessageDirection, pregel_default_msg, pregel_src},
-    expressions::most_common_by,
+    expressions::most_common_expr,
     memory::CheckpointConfig,
     utils::symmetrize,
 };
+use datafusion::functions_aggregate::array_agg::array_agg;
 
 pub const COMMUNITY: &str = "community";
 
@@ -90,10 +91,10 @@ impl<'a> ClassicalLPBuilder<'a> {
                 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.0f32)))
+            // Nested expression: DataFusion splits `most_common(array_agg(...))`
+            // into the array_agg aggregate plus a most_common projection on top,
+            // so the checkpoint spill writes only the reduced O(|V|) result.
+            .add_aggregate_expr(most_common_expr(array_agg(pregel_default_msg())))
             .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_reduce;
-mod most_common_by;
+mod kcore_merge;
+mod most_common;
 
 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_reduce::kcore_reduce;
-pub(crate) use most_common_by::most_common_by;
+pub(crate) use kcore_merge::kcore_merge_expr;
+pub(crate) use most_common::most_common_expr;

src/expressions/kcore_merge.rs

@@ -0,0 +1,197 @@
+//! Per-vertex k-core update as a scalar UDF over the collected neighbour-core
+//! list (`array_agg` result).
+
+use crate::expressions::common::downcast_int32;
+use datafusion::arrow::array::{Array, ArrayRef, Int32Array, ListArray};
+use datafusion::arrow::datatypes::{DataType, Field};
+use datafusion::common::DataFusionError;
+use datafusion::error::Result;
+use datafusion::logical_expr::{
+    ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility,
+};
+use std::sync::Arc;
+
+/// Counts are `u32`:
+/// each bucket holds at most `num_neighbors` entries, bounded by the "degree <
+/// i32::MAX" assumption shared with the rest of the codebase.
+fn kcore_merge_into(
+    counts: &mut Vec<u32>,
+    num_neighbors: usize,
+    neighbors: impl Iterator<Item = i32>,
+) -> i32 {
+    let cap = num_neighbors;
+    counts.clear();
+    counts.resize(cap + 1, 0);
+    for el in neighbors {
+        let bucket = (el.max(0) as usize).min(cap);
+        counts[bucket] += 1;
+    }
+    let mut current_weight = 0u32;
+    for i in (1..=cap).rev() {
+        current_weight += counts[i];
+        if (i as u32) <= current_weight {
+            return i as i32;
+        }
+    }
+    0
+}
+
+fn list_int32_type() -> DataType {
+    DataType::List(Arc::new(Field::new("item", DataType::Int32, true)))
+}
+
+/// Scalar UDF `kcore_merge(List<Int32>) -> Int32`.
+#[derive(Debug, PartialEq, Eq, Hash)]
+pub(crate) struct KCoreMerge {
+    signature: Signature,
+}
+
+impl KCoreMerge {
+    pub(crate) fn new() -> Self {
+        Self {
+            signature: Signature::exact(vec![list_int32_type()], Volatility::Immutable),
+        }
+    }
+}
+
+impl ScalarUDFImpl for KCoreMerge {
+    fn name(&self) -> &str {
+        "kcore_merge"
+    }
+
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+
+    fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+        match arg_types {
+            [DataType::List(f)] if f.data_type() == &DataType::Int32 => Ok(DataType::Int32),
+            _ => Err(DataFusionError::Plan(format!(
+                "kcore_merge expects (List<Int32>), got: {arg_types:?}"
+            ))),
+        }
+    }
+
+    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        let arrays = ColumnarValue::values_to_arrays(&args.args)?;
+        if arrays.len() != 1 {
+            return Err(DataFusionError::Plan(format!(
+                "kcore_merge expects exactly one argument, got: {}",
+                arrays.len()
+            )));
+        }
+        let list = arrays[0]
+            .as_any()
+            .downcast_ref::<ListArray>()
+            .ok_or_else(|| {
+                DataFusionError::Plan(format!(
+                    "kcore_merge argument must be List, got: {:?}",
+                    arrays[0].data_type()
+                ))
+            })?;
+        let values = downcast_int32(list.values(), "kcore_merge", "list elements")?;
+        let offsets = list.offsets();
+        let len = args.number_rows.max(list.len());
+
+        let mut counts: Vec<u32> = Vec::new();
+        let result: Int32Array = (0..len)
+            .map(|i| {
+                let row = i % list.len();
+                if list.is_null(row) {
+                    // No neighbours: nothing can support l >= 1.
+                    return Some(0i32);
+                }
+                let start = offsets[row] as usize;
+                let end = offsets[row + 1] as usize;
+                let neighbors = (start..end)
+                    .filter(|&j| !values.is_null(j))
+                    .map(|j| values.value(j));
+                Some(kcore_merge_into(&mut counts, end - start, neighbors))
+            })
+            .collect();
+
+        Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef))
+    }
+}
+
+/// Builds an [`Expr`] that applies `kcore_merge(neighbors)`.
+pub(crate) fn kcore_merge_expr(neighbors: Expr) -> Expr {
+    ScalarUDF::from(KCoreMerge::new()).call(vec![neighbors])
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use datafusion::prelude::*;
+
+    /// Thin slice wrapper over the counting-array reducer.
+    fn reduce(neighbors: &[i32]) -> i32 {
+        let mut counts = Vec::new();
+        kcore_merge_into(&mut counts, neighbors.len(), neighbors.iter().copied())
+    }
+
+    /// Uncapped_A over a neighbour multiset (mirrors the old accumulator tests).
+    #[test]
+    fn test_kcore_merge_picks_uncapped_core() {
+        // {3:3}: ge(3)=3 -> 3
+        assert_eq!(reduce(&[3, 3, 3]), 3);
+        // {2:2, 1:1}: ge(2)=2 -> 2
+        assert_eq!(reduce(&[2, 2, 1]), 2);
+        // {5:3}: only 3 neighbours, so uncapped capped at 3
+        assert_eq!(reduce(&[5, 5, 5]), 3);
+        // {1:1}: single neighbour -> 1
+        assert_eq!(reduce(&[1]), 1);
+        // {10:1, 1:1}: ge(1)=2, ge(2)=1 -> 1
+        assert_eq!(reduce(&[10, 1]), 1);
+        // neighbours' cores above the degree clamp to the top bucket
+        assert_eq!(reduce(&[1000, 1000, 1000, 1000]), 4);
+    }
+
+    /// No neighbours -> uncapped_A = 0; all-zero neighbours -> 0.
+    #[test]
+    fn test_kcore_merge_empty_and_zero() {
+        assert_eq!(reduce(&[]), 0);
+        assert_eq!(reduce(&[0, 0, 0, 0, 0]), 0);
+    }
+
+    /// Negative neighbour cores clamp to bucket 0 and cannot support l >= 1.
+    #[test]
+    fn test_kcore_merge_negative_clamps() {
+        assert_eq!(reduce(&[-5, -5]), 0);
+        assert_eq!(reduce(&[-5, 3]), 1);
+    }
+
+    /// DataFusion accepts a scalar-over-aggregate expression directly in
+    /// `DataFrame::aggregate` and splits it into the array_agg aggregate plus
+    /// a kcore_merge projection on top — so k-core spills only the reduced
+    /// O(|V|) result. Regression test for the `k_core.rs` wiring.
+    #[tokio::test]
+    async fn test_nested_scalar_over_aggregate() -> Result<()> {
+        use datafusion::arrow::array::{Int32Array as I32A, Int64Array as I64A};
+        use datafusion::functions_aggregate::array_agg::array_agg;
+        let df = dataframe!(
+            "g" => vec![0i64, 0, 1, 1, 1],
+            "a" => vec![3i32, 3, 10, 20, 20],
+        )?;
+
+        let out = df
+            .aggregate(
+                vec![col("g")],
+                vec![kcore_merge_expr(array_agg(col("a"))).alias("k")],
+            )?
+            .collect()
+            .await?;
+        let mut pairs: Vec<(i64, i32)> = Vec::new();
+        for b in &out {
+            let g = b.column(0).as_any().downcast_ref::<I64A>().unwrap();
+            let k = b.column(1).as_any().downcast_ref::<I32A>().unwrap();
+            for r in 0..g.len() {
+                pairs.push((g.value(r), k.value(r)));
+            }
+        }
+        pairs.sort_unstable();
+        // g=0: {3:2} -> uncapped 2; g=1: {10:1,20:2} -> uncapped 3
+        assert_eq!(pairs, vec![(0, 2), (1, 3)]);
+        Ok(())
+    }
+}