From d4df714aa2886e2af89c017bd123b2e7699dc861 Mon Sep 17 00:00:00 2001 From: semyonsinchenko Date: Tue, 11 Aug 2026 12:43:09 +0200 Subject: [PATCH 1/2] 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<()> { From c122df07f49ed7c675992da3c6f751649775ce58 Mon Sep 17 00:00:00 2001 From: semyonsinchenko Date: Tue, 11 Aug 2026 17:43:22 +0200 Subject: [PATCH 2/2] perf: rollback the old style After a lot of tests I found that the best way is still to use collect list; While it may be counterintuitive, it is 3x times faster than custom maps based aggregations. --- .gitignore | 1 + src/algorithm/centrality/k_core.rs | 16 +- src/algorithm/community/classical_lp.rs | 11 +- src/expressions.rs | 8 +- src/expressions/kcore_merge.rs | 197 +++++++ src/expressions/kcore_reduce.rs | 682 ------------------------ src/expressions/most_common.rs | 192 +++++++ src/expressions/most_common_by.rs | 492 ----------------- 8 files changed, 413 insertions(+), 1186 deletions(-) create mode 100644 src/expressions/kcore_merge.rs delete mode 100644 src/expressions/kcore_reduce.rs create mode 100644 src/expressions/most_common.rs delete mode 100644 src/expressions/most_common_by.rs diff --git a/.gitignore b/.gitignore index 45323e9..f3886d6 100644 --- a/.gitignore +++ b/.gitignore @@ -17,6 +17,7 @@ gf_workdir* test-test* .#* graph500* +wiki-Talk* twitter_mpi* diff --git a/src/algorithm/centrality/k_core.rs b/src/algorithm/centrality/k_core.rs index 02bd826..6af56d0 100644 --- a/src/algorithm/centrality/k_core.rs +++ b/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::() .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() diff --git a/src/algorithm/community/classical_lp.rs b/src/algorithm/community/classical_lp.rs index 034624d..7712419 100644 --- a/src/algorithm/community/classical_lp.rs +++ b/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()) diff --git a/src/expressions.rs b/src/expressions.rs index 3c7a4c4..78a35c2 100644 --- a/src/expressions.rs +++ b/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; diff --git a/src/expressions/kcore_merge.rs b/src/expressions/kcore_merge.rs new file mode 100644 index 0000000..b617f2b --- /dev/null +++ b/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, + num_neighbors: usize, + neighbors: impl Iterator, +) -> 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`. +#[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 { + match arg_types { + [DataType::List(f)] if f.data_type() == &DataType::Int32 => Ok(DataType::Int32), + _ => Err(DataFusionError::Plan(format!( + "kcore_merge expects (List), got: {arg_types:?}" + ))), + } + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + 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::() + .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 = 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::().unwrap(); + let k = b.column(1).as_any().downcast_ref::().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(()) + } +} diff --git a/src/expressions/kcore_reduce.rs b/src/expressions/kcore_reduce.rs deleted file mode 100644 index 2298aea..0000000 --- a/src/expressions/kcore_reduce.rs +++ /dev/null @@ -1,682 +0,0 @@ -use std::sync::Arc; - -use crate::expressions::common::{as_binary_like, downcast_int32}; -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, EmitTo, Expr, GroupsAccumulator, Signature, - Volatility, -}; -use datafusion::scalar::ScalarValue; - -#[derive(Debug)] -pub(crate) struct KCoreReduceAccumulator { - // histogram of neighbors cores - counts: HashMap, -} - -impl KCoreReduceAccumulator { - pub(crate) fn new() -> Self { - Self { - counts: HashMap::new(), - } - } -} - -#[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; - let mut buf = Vec::with_capacity(4 + 8usize * (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( - "k_core_reduce: 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 + 8 * n { - return Err(DataFusionError::Execution( - "k_core_reduce: corrupt state (length mismatch)".to_string(), - )); - } - - // Each entry is 8 bytes: i32 key | u32 value (see `se_map`). - for i in 0..n { - let o = 4 + 8 * i; - let k = i32::from_le_bytes(buf[o..o + 4].try_into().unwrap()); - let v = u32::from_le_bytes(buf[o + 4..o + 8].try_into().unwrap()); - *map.entry(k).or_insert(0u32) += v; - } - - 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")?; - - // No nulls expected: this is an internal aggregator over neighbour core - // estimates, which are non-null by the Pregel message contract. - for i in 0..labels.len() { - let l = labels.value(i); - - let cur = self.counts.entry(l).or_insert(0u32); - *cur += 1u32; - } - - Ok(()) - } - - fn evaluate(&mut self) -> Result { - let best = uncapped_core(&self.counts); - Ok(ScalarValue::Int32(Some(best))) - } - - fn size(&self) -> usize { - size_of::() + self.counts.capacity() * (size_of::() + size_of::()) - } - - fn state(&mut self) -> Result> { - Ok(vec![ScalarValue::Binary(Some(se_map(&self.counts)))]) - } - - fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - let v = as_binary_like(&states[0], "k_core_reduce", "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.counts)?; - } - - Ok(()) - } -} - -#[derive(Debug, PartialEq, Eq, Hash)] -pub(crate) struct KCoreReduce { - signature: Signature, -} - -impl KCoreReduce { - pub(crate) fn new() -> Self { - Self { - signature: Signature::exact(vec![DataType::Int32], Volatility::Immutable), - } - } -} - -impl AggregateUDFImpl for KCoreReduce { - fn name(&self) -> &str { - "k_core_reduce" - } - - fn signature(&self) -> &Signature { - &self.signature - } - - fn return_type(&self, arg_types: &[DataType]) -> Result { - if (arg_types.len() != 1) || (arg_types[0] != DataType::Int32) { - return Err(DataFusionError::Plan(format!( - "k_core_reduce expects exactly one argument of type i32 but got {arg_types:?}" - ))); - } - - Ok(DataType::Int32) - } - - fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result> { - 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"), - DataType::Binary, - true, - ))]) - } -} - -/// Builds an [`Expr`] that applies `k_core_reduce(a)`. -pub(crate) fn kcore_reduce(a: Expr) -> Expr { - AggregateUDF::from(KCoreReduce::new()).call(vec![a]) -} - -#[cfg(test)] -mod tests { - use super::*; - use datafusion::arrow::array::{ArrayRef, BinaryArray, Int32Array, Int64Array, RecordBatch}; - use datafusion::arrow::datatypes::{DataType, Field, Schema}; - use datafusion::datasource::MemTable; - use datafusion::prelude::SessionConfig; - use datafusion::prelude::*; - - /// Feed a batch of i32 neighbour cores into a fresh accumulator. - fn reduce(cores: &[i32]) -> KCoreReduceAccumulator { - let mut acc = KCoreReduceAccumulator::new(); - acc.update_batch(&[Arc::new(Int32Array::from(cores.to_vec())) as ArrayRef]) - .unwrap(); - acc - } - - /// Direct accumulator: uncapped_A over a neighbour multiset. - #[test] - fn test_accumulator_picks_uncapped_core() { - // {3:3}: ge(3)=3 -> 3 - assert_eq!( - reduce(&[3, 3, 3]).evaluate().unwrap(), - ScalarValue::Int32(Some(3)) - ); - // {2:2, 1:1}: ge(2)=2 -> 2 - assert_eq!( - reduce(&[2, 2, 1]).evaluate().unwrap(), - ScalarValue::Int32(Some(2)) - ); - // {5:3}: only 3 neighbours, so uncapped capped at 3 - assert_eq!( - reduce(&[5, 5, 5]).evaluate().unwrap(), - ScalarValue::Int32(Some(3)) - ); - // {1:1}: single neighbour -> 1 - assert_eq!( - reduce(&[1]).evaluate().unwrap(), - ScalarValue::Int32(Some(1)) - ); - // {10:1, 1:1}: ge(1)=2, ge(2)=1 -> 1 - assert_eq!( - reduce(&[10, 1]).evaluate().unwrap(), - ScalarValue::Int32(Some(1)) - ); - } - - /// No neighbours -> uncapped_A = 0 (the empty multiset). - #[test] - fn test_accumulator_empty_returns_zero() { - let mut acc = KCoreReduceAccumulator::new(); - assert_eq!(acc.evaluate().unwrap(), ScalarValue::Int32(Some(0))); - } - - /// All neighbours core 0 -> no support -> uncapped 0. - #[test] - fn test_accumulator_all_zero_returns_zero() { - assert_eq!( - reduce(&[0, 0, 0, 0, 0]).evaluate().unwrap(), - ScalarValue::Int32(Some(0)) - ); - } - - /// se_map/de_map round-trip preserves the histogram exactly. - #[test] - fn test_se_de_map_roundtrip_preserves_counts() { - let mut original: HashMap = HashMap::new(); - original.insert(1, 3); - original.insert(2, 5); - original.insert(7, 1); - 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_eq!(restored.get(k).copied().unwrap(), *v, "key {k}"); - } - } - - /// Merge semantics: `de_map_and_insert` must ADD counts for shared keys, - /// not overwrite. This is the property that makes partial->final merging - /// correct across partitions. - #[test] - fn test_de_map_and_insert_accumulates_not_overwrites() { - let mut m: HashMap = HashMap::new(); - m.insert(3, 2); - let mut partial: HashMap = HashMap::new(); - partial.insert(3, 1); // shared key - partial.insert(5, 4); // new key - let bytes = se_map(&partial); - de_map_and_insert(&bytes, &mut m).unwrap(); - assert_eq!(m.get(&3).copied().unwrap(), 3, "shared key must sum"); - assert_eq!(m.get(&5).copied().unwrap(), 4, "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. - /// a={3:2} (uncapped 2), b={3:1} (uncapped 1); merged {3:3} -> uncapped 3. - /// A no-op `merge_batch` yields 2; keeping only the local winner yields 1. - #[test] - fn test_accumulator_merge_unions_partial_states() { - let mut a = reduce(&[3, 3]); - let mut b = reduce(&[3]); - - 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(); - - 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<()> { - // g=0: cores [1,2,1] -> {1:2,2:1} -> uncapped 1 - // g=1: cores [10,20,20] -> {10:1,20:2} -> uncapped 3 - let df = dataframe!( - "g" => vec![0i64, 0, 0, 1, 1, 1], - "a" => vec![1i32, 2, 1, 10, 20, 20], - )?; - let out = df - .aggregate(vec![col("g")], vec![kcore_reduce(col("a")).alias("k")])? - .sort(vec![col("g").sort(true, true)])? - .collect() - .await?; - let g = out[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - let k = out[0] - .column(1) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(g.len(), 2); - assert_eq!(g.value(0), 0); - assert_eq!(k.value(0), 1); - assert_eq!(g.value(1), 1); - assert_eq!(k.value(1), 3); - Ok(()) - } - - /// Multi-partition: force the partial->final path via two `MemTable` - /// partitions and `target_partitions(2)`. Vertex 0 receives core 3 twice in - /// P1 and once in P2; only a correct `state()` + `merge_batch` (summing the - /// shared key 3 -> 3) yields uncapped 3. This is the test that catches a - /// broken `de_map` stride/offset or a wrong `state_fields`. - #[tokio::test] - async fn test_kcore_reduce_multi_partition_merge() { - let schema = Arc::new(Schema::new(vec![ - Field::new("g", DataType::Int64, false), - Field::new("a", DataType::Int32, false), - ])); - let mk = |cores: Vec| { - let n = cores.len(); - RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(Int64Array::from_iter_values( - std::iter::repeat(0i64).take(n), - )) as ArrayRef, - Arc::new(Int32Array::from(cores)) as ArrayRef, - ], - ) - .unwrap() - }; - // P1: {3:2} -> uncapped 2 alone; P2: {3:1} -> uncapped 1 alone. - let p1 = mk(vec![3, 3]); - let p2 = mk(vec![3]); - - 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![kcore_reduce(col("a")).alias("k")]) - .unwrap() - .collect() - .await - .unwrap(); - let k = out[0] - .column(1) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(k.len(), 1); - // merged {3:3} -> uncapped 3 - assert_eq!(k.value(0), 3); - } - - /// `return_type` must accept exactly (Int32) and reject the rest. - #[test] - fn test_kcore_reduce_return_type_validation() { - let udf = KCoreReduce::new(); - assert_eq!( - udf.return_type(&[DataType::Int32]).unwrap(), - DataType::Int32 - ); - assert!(udf.return_type(&[DataType::Int64]).is_err()); - assert!(udf.return_type(&[]).is_err()); - assert!( - udf.return_type(&[DataType::Int32, DataType::Int32]) - .is_err() - ); - } -} diff --git a/src/expressions/most_common.rs b/src/expressions/most_common.rs new file mode 100644 index 0000000..ae691d0 --- /dev/null +++ b/src/expressions/most_common.rs @@ -0,0 +1,192 @@ +//! Mode-with-minimal-label scalar UDF for classical label propagation. +//! +//! LDBC CDLP semantics: the label with the largest total weight among a +//! vertex's neighbours wins; ties break toward the smallest label. With the +//! unit weights used by classical_lp this is exactly the mode. + +use crate::expressions::common::downcast_int64; +use datafusion::arrow::array::{Array, ArrayRef, Int64Array, 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; + +/// Most frequent label, ties broken by the smallest label. Returns `None` for +/// an empty multiset; +fn mode_min_label(scratch: &mut Vec) -> Option { + if scratch.is_empty() { + return None; + } + scratch.sort_unstable(); + let mut best: Option = None; + let mut best_count: i64 = 0; + let mut i = 0; + while i < scratch.len() { + let label = scratch[i]; + let mut count: i64 = 0; + while i < scratch.len() && scratch[i] == label { + count += 1; + i += 1; + } + if count > best_count || (count == best_count && best.map_or(true, |b| label < b)) { + best = Some(label); + best_count = count; + } + } + best +} + +fn list_int64_type() -> DataType { + DataType::List(Arc::new(Field::new("item", DataType::Int64, true))) +} + +/// Scalar UDF `most_common(List) -> Int64` (nullable). +#[derive(Debug, PartialEq, Eq, Hash)] +pub(crate) struct MostCommon { + signature: Signature, +} + +impl MostCommon { + pub(crate) fn new() -> Self { + Self { + signature: Signature::exact(vec![list_int64_type()], Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for MostCommon { + fn name(&self) -> &str { + "most_common" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + match arg_types { + [DataType::List(f)] if f.data_type() == &DataType::Int64 => Ok(DataType::Int64), + _ => Err(DataFusionError::Plan(format!( + "most_common expects (List), got: {arg_types:?}" + ))), + } + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + if arrays.len() != 1 { + return Err(DataFusionError::Plan(format!( + "most_common expects exactly one argument, got: {}", + arrays.len() + ))); + } + let list = arrays[0] + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Plan(format!( + "most_common argument must be List, got: {:?}", + arrays[0].data_type() + )) + })?; + let values = downcast_int64(list.values(), "most_common", "list elements")?; + let offsets = list.offsets(); + let len = args.number_rows.max(list.len()); + + // Reuse the sort scratch across rows: no per-row allocation. + // One scratch buffer reused across all rows: no per-row allocation. + let mut scratch: Vec = Vec::new(); + let result: Int64Array = (0..len) + .map(|i| { + let row = i % list.len(); + if list.is_null(row) { + return None; + } + let start = offsets[row] as usize; + let end = offsets[row + 1] as usize; + scratch.clear(); + for j in start..end { + if !values.is_null(j) { + scratch.push(values.value(j)); + } + } + mode_min_label(&mut scratch) + }) + .collect(); + + Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef)) + } +} + +/// Builds an [`Expr`] that applies `most_common(labels)`. +pub(crate) fn most_common_expr(labels: Expr) -> Expr { + ScalarUDF::from(MostCommon::new()).call(vec![labels]) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::prelude::*; + + fn mode(labels: &[i64]) -> Option { + let mut scratch = labels.to_vec(); + mode_min_label(&mut scratch) + } + + /// LDBC semantics: most frequent label wins. + #[test] + fn test_mode_picks_most_frequent() { + assert_eq!(mode(&[5, 5, 5, 3, 3]), Some(5)); + assert_eq!(mode(&[7]), Some(7)); + assert_eq!(mode(&[1, 2, 3, 4, 4, 4, 2, 2]), Some(2)); + } + + /// Ties break toward the smallest label. + #[test] + fn test_mode_tie_picks_min_label() { + assert_eq!(mode(&[5, 5, 2, 2]), Some(2)); + assert_eq!(mode(&[10, 10, 1, 1]), Some(1)); + assert_eq!(mode(&[3, 3, 3, 7, 7, 7]), Some(3)); + } + + /// Empty multiset -> None (SQL NULL), matching the old accumulator. + #[test] + fn test_mode_empty_returns_none() { + assert_eq!(mode(&[]), None); + } + + /// Nested scalar-over-aggregate wiring (the `classical_lp.rs` shape): + /// most_common(array_agg(label)) per group. + #[tokio::test] + async fn test_nested_most_common_aggregate() -> Result<()> { + use datafusion::arrow::array::Int64Array as I64A; + use datafusion::functions_aggregate::array_agg::array_agg; + let df = dataframe!( + "g" => vec![0i64, 0, 0, 1, 1, 1], + "a" => vec![5i64, 5, 3, 10, 20, 20], + )?; + + let out = df + .aggregate( + vec![col("g")], + vec![most_common_expr(array_agg(col("a"))).alias("c")], + )? + .collect() + .await?; + let mut pairs: Vec<(i64, i64)> = Vec::new(); + for b in &out { + let g = b.column(0).as_any().downcast_ref::().unwrap(); + let c = b.column(1).as_any().downcast_ref::().unwrap(); + for r in 0..g.len() { + pairs.push((g.value(r), c.value(r))); + } + } + pairs.sort_unstable(); + // g=0: {5:2, 3:1} -> 5; g=1: {10:1, 20:2} -> 20 + assert_eq!(pairs, vec![(0, 5), (1, 20)]); + Ok(()) + } +} diff --git a/src/expressions/most_common_by.rs b/src/expressions/most_common_by.rs deleted file mode 100644 index 112ec30..0000000 --- a/src/expressions/most_common_by.rs +++ /dev/null @@ -1,492 +0,0 @@ -//! Most Commont By aggregation. -//! -//! Limited by design. -//! For the given (a: i64, b: f64) returns a value of a -//! for which the sum(b) is maximal. Tie-breaking is based -//! on value of a itself (minimal, LDBC Label Propagation semantics) -use std::sync::Arc; - -use crate::expressions::common::{as_binary_like, downcast_int64}; -use datafusion::arrow::array::{ArrayRef, Float32Array}; -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 + 12usize * (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 + 12 * n { - return Err(DataFusionError::Execution( - "most_common_by: corrupt state (length mismatch)".to_string(), - )); - } - - for i in 0..n { - let o = 4 + 12 * i; - let k = i64::from_le_bytes(buf[o..o + 8].try_into().unwrap()); - let v = f32::from_le_bytes(buf[o + 8..o + 12].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 second argument be f32 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 = f32::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::Float32], - 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::Float32) - ) - { - return Err(DataFusionError::Plan(format!( - "most_common_by expets exactly two arguments of types i64 and f32 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, Float32Array, 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(Float32Array::from(vec![2.0f32, 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(Float32Array::from(vec![5.0f32, 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(Float32Array::from(vec![3.0f32, 2.0])) as ArrayRef, - ]) - .unwrap(); - // b: {2->4} local winner 2 - b.update_batch(&[ - Arc::new(Int64Array::from(vec![2i64])) as ArrayRef, - Arc::new(Float32Array::from(vec![4.0f32])) 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.0f32, 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.0f32, 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.0f32, 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::Float32, 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(Float32Array::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.0f32; 100]); - let p2 = mk(vec![2i64; 150], vec![1.0f32; 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, Float32) 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::Float32]) - .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::Float32]) - .is_err() - ); - assert!( - udf.return_type(&[DataType::Int64, DataType::Float32, DataType::Int64]) - .is_err() - ); - } -}