From d4df714aa2886e2af89c017bd123b2e7699dc861 Mon Sep 17 00:00:00 2001 From: semyonsinchenko Date: Tue, 11 Aug 2026 12:43:09 +0200 Subject: [PATCH] perf: switch k-core to groups accumulator --- src/expressions/kcore_reduce.rs | 338 ++++++++++++++++++++++++++++---- 1 file changed, 305 insertions(+), 33 deletions(-) diff --git a/src/expressions/kcore_reduce.rs b/src/expressions/kcore_reduce.rs index 80e8ff6..2298aea 100644 --- a/src/expressions/kcore_reduce.rs +++ b/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>, +} + +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 { + 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) -> Vec { // 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) -> 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 { + 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::(); + r += self.counts.capacity() * size_of::>(); + for i in 0..self.counts.len() { + r += self.counts[i].capacity() * (size_of::() + size_of::()); + } + + r + } + + fn state(&mut self, emit_to: EmitTo) -> Result> { + 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> { + 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 { - // 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> { + Ok(Box::new(KCoreReduceGroupsAccumulator::new())) + } + fn state_fields(&self, args: StateFieldsArgs) -> Result>> { 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::().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::().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::().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::().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::().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::().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::().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::().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<()> {