perf: switch k-core to groups accumulator

src/expressions/kcore_reduce.rs

@@ -1,14 +1,17 @@
 use std::sync::Arc;
 
 use crate::expressions::common::{as_binary_like, downcast_int32};
-use datafusion::arrow::array::ArrayRef;
+use datafusion::arrow::array::{
+    Array, ArrayRef, BinaryArray, BinaryBuilder, BooleanArray, Int32Array,
+};
 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,
+    Accumulator, AggregateUDF, AggregateUDFImpl, EmitTo, Expr, GroupsAccumulator, Signature,
+    Volatility,
 };
 use datafusion::scalar::ScalarValue;
 
@@ -26,6 +29,54 @@ impl KCoreReduceAccumulator {
     }
 }
 
+#[derive(Debug)]
+pub(crate) struct KCoreReduceGroupsAccumulator {
+    // accumulator for all groups at once;
+    // fast path.
+    counts: Vec<HashMap<i32, u32>>,
+}
+
+impl KCoreReduceGroupsAccumulator {
+    pub(crate) fn new() -> Self {
+        Self { counts: Vec::new() }
+    }
+}
+
+/// The main logic of choosing a new (uncapped) core.
+///
+/// Returns `uncapped_A = max{ l : #{neighbours with core >= l} >= l }`.
+/// The `min(., current_core)` cap is applied later, in the Pregel vertex
+/// update expression, because `current_core` is not visible to the
+/// aggregate (see `skip_dest_state`).
+///
+/// Based on: 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.
+///
+/// Shared between Accumulator and GroupsAccumulator
+fn uncapped_core(m: &HashMap<i32, u32>) -> i32 {
+    if m.is_empty() {
+        return 0i32;
+    }
+
+    let mut entries: Vec<(i32, u32)> = m.iter().map(|(&k, &v)| (k, v)).collect();
+    // Descending by value so the running sum is the "# neighbours >= value".
+    entries.sort_unstable_by(|a, b| b.0.cmp(&a.0));
+
+    let mut cum: u32 = 0; // running #{neighbours with value >= current}; bounded by degree < i32::MAX
+    let mut best: i32 = 0; // uncapped_A >= 0 always (l=0 always satisfies ge(0)=total>=0)
+    for (value, count) in entries {
+        cum += count;
+        // candidate = the largest l <= value with ge(l) >= l on this step
+        let candidate = value.min(cum as i32);
+        if candidate > best {
+            best = candidate;
+        }
+    }
+
+    best
+}
+
 fn se_map(m: &HashMap<i32, u32>) -> Vec<u8> {
     // We are assumming that a single node degree < i32::MAX
     let n = m.len() as u32;
@@ -64,6 +115,111 @@ fn de_map_and_insert(buf: &[u8], map: &mut HashMap<i32, u32>) -> Result<()> {
     Ok(())
 }
 
+impl GroupsAccumulator for KCoreReduceGroupsAccumulator {
+    fn update_batch(
+        &mut self,
+        values: &[ArrayRef],
+        group_indices: &[usize],
+        opt_filter: Option<&BooleanArray>,
+        total_num_groups: usize,
+    ) -> Result<()> {
+        let v = downcast_int32(&values[0], "k_core_reduce", "first")?;
+        self.counts.resize(total_num_groups, HashMap::new());
+
+        // Nulls are not expected by the Pregel message contract, but skip them
+        // (and filtered-out rows) defensively: counting a null slot's raw bits
+        // would silently corrupt the histogram.
+        for i in 0..v.len() {
+            if v.is_null(i) || opt_filter.is_some_and(|f| !f.value(i)) {
+                continue;
+            }
+            let l = v.value(i);
+            let cur = self.counts[group_indices[i]].entry(l).or_insert(0u32);
+            *cur += 1u32;
+        }
+
+        Ok(())
+    }
+
+    fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
+        let maps = emit_to.take_needed(&mut self.counts);
+        let result: Int32Array = (0..maps.len()).map(|i| uncapped_core(&maps[i])).collect();
+
+        Ok(Arc::new(result) as ArrayRef)
+    }
+
+    fn size(&self) -> usize {
+        // size of self
+        // + size of hashmap struct * num groups
+        // + sum of sizes of each map (i32 + u32) * capacity of each
+        let mut r = size_of::<Self>();
+        r += self.counts.capacity() * size_of::<HashMap<i32, u32>>();
+        for i in 0..self.counts.len() {
+            r += self.counts[i].capacity() * (size_of::<i32>() + size_of::<u32>());
+        }
+
+        r
+    }
+
+    fn state(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+        let maps = emit_to.take_needed(&mut self.counts);
+        let result = BinaryArray::from_iter_values((0..maps.len()).map(|i| se_map(&maps[i])));
+
+        Ok(vec![Arc::new(result) as ArrayRef])
+    }
+
+    fn merge_batch(
+        &mut self,
+        values: &[ArrayRef],
+        group_indices: &[usize],
+        opt_filter: Option<&BooleanArray>,
+        total_num_groups: usize,
+    ) -> Result<()> {
+        let v = as_binary_like(&values[0], "k_core_reduce", "argument")?;
+
+        self.counts.resize(total_num_groups, HashMap::new());
+
+        // Null state rows appear when the skip-aggregation path filtered rows
+        // out (`convert_to_state` emits nulls for them); they must be ignored,
+        // otherwise the empty blob would be rejected as corrupt state.
+        for i in 0..v.len() {
+            if v.is_null(i) || opt_filter.is_some_and(|f| !f.value(i)) {
+                continue;
+            }
+            de_map_and_insert(v.value(i), &mut self.counts[group_indices[i]])?;
+        }
+
+        Ok(())
+    }
+
+    fn convert_to_state(
+        &self,
+        values: &[ArrayRef],
+        opt_filter: Option<&BooleanArray>,
+    ) -> Result<Vec<ArrayRef>> {
+        let v = downcast_int32(&values[0], "k_core_reduce", "first")?;
+
+        // Each input row becomes its own single-entry histogram state.
+        // Filtered-out (and null, not expected by contract) rows become null
+        // states so the Final phase ignores them in `merge_batch`.
+        let mut builder = BinaryBuilder::new();
+        for i in 0..v.len() {
+            if v.is_null(i) || opt_filter.is_some_and(|f| !f.value(i)) {
+                builder.append_null();
+            } else {
+                let m = HashMap::from([(v.value(i), 1u32)]);
+                builder.append_value(se_map(&m));
+            }
+        }
+
+        Ok(vec![Arc::new(builder.finish()) as ArrayRef])
+    }
+
+    fn supports_convert_to_state(&self) -> bool {
+        true
+    }
+}
+
 impl Accumulator for KCoreReduceAccumulator {
     fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
         let labels = downcast_int32(&values[0], "k_core_reduce", "first")?;
@@ -80,38 +236,8 @@ impl Accumulator for KCoreReduceAccumulator {
         Ok(())
     }
 
-    /// The main logic of choosing a new (uncapped) core.
-    ///
-    /// Returns `uncapped_A = max{ l : #{neighbours with core >= l} >= l }`.
-    /// The `min(., current_core)` cap is applied later, in the Pregel vertex
-    /// update expression, because `current_core` is not visible to the
-    /// aggregate (see `skip_dest_state`).
-    ///
-    /// Based on: 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.
     fn evaluate(&mut self) -> Result<ScalarValue> {
-        // Defensive: a group only exists if it received >= 1 row. The empty
-        // multiset has uncapped_A = 0 (no neighbour can support any l >= 1).
-        if self.counts.is_empty() {
-            return Ok(ScalarValue::Int32(Some(0i32)));
-        }
-
-        let mut entries: Vec<(i32, u32)> = self.counts.iter().map(|(&k, &v)| (k, v)).collect();
-        // Descending by value so the running sum is the "# neighbours >= value".
-        entries.sort_unstable_by(|a, b| b.0.cmp(&a.0));
-
-        let mut cum: u32 = 0; // running #{neighbours with value >= current}; bounded by degree < i32::MAX
-        let mut best: i32 = 0; // uncapped_A >= 0 always (l=0 always satisfies ge(0)=total>=0)
-        for (value, count) in entries {
-            cum += count;
-            // candidate = the largest l <= value with ge(l) >= l on this step
-            let candidate = value.min(cum as i32);
-            if candidate > best {
-                best = candidate;
-            }
-        }
-
+        let best = uncapped_core(&self.counts);
         Ok(ScalarValue::Int32(Some(best)))
     }
 
@@ -172,6 +298,17 @@ impl AggregateUDFImpl for KCoreReduce {
         Ok(Box::new(KCoreReduceAccumulator::new()))
     }
 
+    fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool {
+        true
+    }
+
+    fn create_groups_accumulator(
+        &self,
+        _args: AccumulatorArgs,
+    ) -> Result<Box<dyn GroupsAccumulator>> {
+        Ok(Box::new(KCoreReduceGroupsAccumulator::new()))
+    }
+
     fn state_fields(&self, args: StateFieldsArgs) -> Result<Vec<Arc<Field>>> {
         Ok(vec![Arc::new(Field::new(
             format_state_name(args.name, "value"),
@@ -310,6 +447,141 @@ mod tests {
         assert_eq!(a.evaluate().unwrap(), ScalarValue::Int32(Some(3)));
     }
 
+    /// GroupsAccumulator: multiple groups updated in a single batch evaluate
+    /// to their own uncapped_A, in group-index order, and `EmitTo::All`
+    /// resets the internal state.
+    #[test]
+    fn test_groups_accumulator_update_evaluate_multi_group() {
+        let mut acc = KCoreReduceGroupsAccumulator::new();
+        // g0: {3:3} -> 3; g1: {2:2, 1:1} -> 2; g2: {5:3} -> 3
+        acc.update_batch(
+            &[Arc::new(Int32Array::from(vec![3i32, 3, 3, 2, 2, 1, 5, 5, 5])) as ArrayRef],
+            &[0, 0, 0, 1, 1, 1, 2, 2, 2],
+            None,
+            3,
+        )
+        .unwrap();
+
+        let out = acc.evaluate(EmitTo::All).unwrap();
+        let out = out.as_any().downcast_ref::<Int32Array>().unwrap();
+        assert_eq!(out.len(), 3);
+        assert_eq!(out.value(0), 3);
+        assert_eq!(out.value(1), 2);
+        assert_eq!(out.value(2), 3);
+
+        // EmitTo::All released the state; a fresh batch starts from scratch.
+        assert!(acc.counts.is_empty());
+    }
+
+    /// GroupsAccumulator: `state()`/`merge_batch()` round-trip must SUM the
+    /// partial histograms for shared keys — the associativity property that
+    /// makes two-phase (Partial -> Final) aggregation correct. This is the
+    /// groups-level analogue of `test_accumulator_merge_unions_partial_states`.
+    #[test]
+    fn test_groups_accumulator_state_merge_roundtrip() {
+        let mut a = KCoreReduceGroupsAccumulator::new();
+        a.update_batch(
+            &[Arc::new(Int32Array::from(vec![3i32, 3])) as ArrayRef],
+            &[0, 0],
+            None,
+            1,
+        )
+        .unwrap();
+
+        let mut b = KCoreReduceGroupsAccumulator::new();
+        b.update_batch(
+            &[Arc::new(Int32Array::from(vec![3i32, 5])) as ArrayRef],
+            &[0, 0],
+            None,
+            1,
+        )
+        .unwrap();
+
+        let states = b.state(EmitTo::All).unwrap();
+        assert_eq!(states.len(), 1);
+        // a={3:2}, b={3:1, 5:1}; merged {3:3, 5:1} -> ge(3)=3 -> uncapped 3
+        a.merge_batch(&states, &[0], None, 1).unwrap();
+
+        let out = a.evaluate(EmitTo::All).unwrap();
+        let out = out.as_any().downcast_ref::<Int32Array>().unwrap();
+        assert_eq!(out.len(), 1);
+        assert_eq!(out.value(0), 3);
+    }
+
+    /// GroupsAccumulator: `EmitTo::First(n)` emits the first n groups and
+    /// shifts the remaining groups' indices down by n.
+    #[test]
+    fn test_groups_accumulator_emit_to_first() {
+        let mut acc = KCoreReduceGroupsAccumulator::new();
+        acc.update_batch(
+            &[Arc::new(Int32Array::from(vec![1i32, 10, 20, 20])) as ArrayRef],
+            &[0, 1, 1, 1],
+            None,
+            2,
+        )
+        .unwrap();
+
+        let out = acc.evaluate(EmitTo::First(1)).unwrap();
+        let out = out.as_any().downcast_ref::<Int32Array>().unwrap();
+        // g0: {1:1} -> 1
+        assert_eq!(out.len(), 1);
+        assert_eq!(out.value(0), 1);
+
+        // g1 (now index 0): {10:1, 20:2} -> 3
+        let out = acc.evaluate(EmitTo::All).unwrap();
+        let out = out.as_any().downcast_ref::<Int32Array>().unwrap();
+        assert_eq!(out.len(), 1);
+        assert_eq!(out.value(0), 3);
+    }
+
+    /// GroupsAccumulator: `convert_to_state` turns each input row into its own
+    /// single-entry histogram state, and merging those states reproduces the
+    /// histogram of the original input. Filtered-out rows must become null
+    /// states and `merge_batch` must skip them.
+    #[test]
+    fn test_groups_accumulator_convert_to_state() {
+        let acc = KCoreReduceGroupsAccumulator::new();
+
+        let states = acc
+            .convert_to_state(
+                &[Arc::new(Int32Array::from(vec![3i32, 3, 5])) as ArrayRef],
+                None,
+            )
+            .unwrap();
+        assert_eq!(states.len(), 1);
+        let bin = states[0].as_any().downcast_ref::<BinaryArray>().unwrap();
+        assert_eq!(bin.len(), 3);
+        assert_eq!(bin.null_count(), 0);
+
+        let mut merged = KCoreReduceGroupsAccumulator::new();
+        merged.merge_batch(&states, &[0, 0, 0], None, 1).unwrap();
+        let out = merged.evaluate(EmitTo::All).unwrap();
+        let out = out.as_any().downcast_ref::<Int32Array>().unwrap();
+        // {3:2, 5:1} -> ge(3)=3 -> uncapped 3
+        assert_eq!(out.len(), 1);
+        assert_eq!(out.value(0), 3);
+
+        // Filtered-out row (index 1) must become a null state and be skipped.
+        let filter = Arc::new(BooleanArray::from(vec![true, false, true]));
+        let states = acc
+            .convert_to_state(
+                &[Arc::new(Int32Array::from(vec![3i32, 3, 5])) as ArrayRef],
+                Some(filter.as_ref()),
+            )
+            .unwrap();
+        let bin = states[0].as_any().downcast_ref::<BinaryArray>().unwrap();
+        assert_eq!(bin.null_count(), 1);
+        assert!(bin.is_null(1));
+
+        let mut merged = KCoreReduceGroupsAccumulator::new();
+        merged.merge_batch(&states, &[0, 0, 0], None, 1).unwrap();
+        let out = merged.evaluate(EmitTo::All).unwrap();
+        let out = out.as_any().downcast_ref::<Int32Array>().unwrap();
+        // only rows 0 and 2 counted: {3:1, 5:1} -> ge(2)=2 -> uncapped 2
+        assert_eq!(out.len(), 1);
+        assert_eq!(out.value(0), 2);
+    }
+
     /// GROUP BY: each group resolves to its own uncapped_A independently.
     #[tokio::test]
     async fn test_kcore_reduce_grouped() -> Result<()> {