diff --git a/README.md b/README.md index 7ed472e..ac05cad 100644 --- a/README.md +++ b/README.md @@ -64,10 +64,14 @@ Input format is parquet by default; `--format csv` and `--format json` are suppo The work directory holds two kinds of data: the algorithms' parquet checkpoints (written and re-read between iterations) and DataFusion's spill files. Both are high-throughput and latency-sensitive, and on large graphs the algorithms can stream large amounts of data through them. On my experiments during processing billion-scale graphs with limited memory (~5-6GB limit) the process read from disk more than 200GBs. Point this at the fastest local storage available, preferably a dedicated NVMe or SSD. Network filesystems and spinning disks will make spill and checkpoint latency dominate runtime. The directory is created if missing; relative paths resolve against the current directory. `--max-temp-file` (env `GRAPHFRAMES_MAX_TEMP_FILE`, default `200G`) bounds the total size of the spill subdirectory. For graphs of the size of few billions edges and vertices the default may be not enough. -### Libraray +### Library TBD +## Known Limitations + +- Some algorithms (K-Core, Label Propagation, etc.) are implemented with an assumption that the single node degree in the graph is always less than `i32::MAX`. In other words it is assummed that each node has less than ~2 billions neighbors / adjacent edges. While it is possible to support graph with hubs with more than 2B neighbors it will introduce a performance degradation for anything less (99.999% of all the graphs in the world). If you are working with a graph where a hub has degree gretater than 2B I would recommend to take a look on edges prunning technique if you want to use this library. + ## References ### Core diff --git a/src/algorithm/centrality/k_core.rs b/src/algorithm/centrality/k_core.rs index 0444692..02bd826 100644 --- a/src/algorithm/centrality/k_core.rs +++ b/src/algorithm/centrality/k_core.rs @@ -1,11 +1,11 @@ use crate::algorithm::pregel::{MessageDirection, pregel_default_msg, pregel_src}; -use crate::expressions::kcore_merge_expr; +use crate::expressions::kcore_reduce; use crate::memory::CheckpointConfig; use crate::utils::symmetrize; use crate::{EDGE_DST, EDGE_SRC, GraphFrame, VERTEX_ID}; +use datafusion::arrow::datatypes::DataType; use datafusion::error::Result; use datafusion::execution::object_store::ObjectStoreUrl; -use datafusion::functions_aggregate::array_agg::array_agg; use datafusion::functions_aggregate::count::count; use datafusion::object_store::path::Path; use datafusion::prelude::*; @@ -108,14 +108,17 @@ impl<'a> KCoreBuilder<'a> { .select(vec![col(VERTEX_ID)])? .join(degrees, JoinType::Left, &[VERTEX_ID], &[EDGE_SRC], None)? .with_column(DEGREE, coalesce(vec![col(DEGREE), lit(0i64)]))? - .select(vec![col(VERTEX_ID), col(DEGREE)])?; + .select(vec![ + col(VERTEX_ID), + cast(col(DEGREE), DataType::Int32).alias(DEGREE), + ])?; let prepared_graph = GraphFrame { vertices: prepared_vertices, edges: prepared_edges, }; - let new_core = kcore_merge_expr(pregel_default_msg(), col(KCORE)); + let new_core = coalesce(vec![pregel_default_msg(), lit(0)]); let mut pregel_builder = prepared_graph .pregel() @@ -127,7 +130,7 @@ impl<'a> KCoreBuilder<'a> { // corrupt their estimates. Early stopping therefore relies on // the voting column alone, never on participation pruning. .add_message(pregel_src(KCORE), MessageDirection::SrcToDst) - .add_aggregate_expr(array_agg(pregel_default_msg())) + .add_aggregate_expr(kcore_reduce(pregel_default_msg())) // A vertex votes "still active" exactly while its core number // changed this iteration; the run stops once nobody changed. .with_vertex_voting("active", col(KCORE).not_eq(new_core.clone())) @@ -156,7 +159,7 @@ impl GraphFrame { mod tests { use super::*; use crate::utils::collect_to_i64; - use datafusion::arrow::array::Int64Array; + use datafusion::arrow::array::{Int32Array, Int64Array}; use std::collections::HashMap; use std::fs; use std::path::PathBuf; @@ -224,13 +227,14 @@ mod tests { .as_any() .downcast_ref::() .unwrap(); + // KCORE is Int32 (see `kcore_reduce`); vertex id stays Int64. let cores = batch .column(1) .as_any() - .downcast_ref::() + .downcast_ref::() .unwrap(); for i in 0..ids.len() { - map.insert(ids.value(i), cores.value(i)); + map.insert(ids.value(i), cores.value(i) as i64); } } Ok(map) @@ -512,9 +516,10 @@ mod tests { .run(&ctx, &output_uri, false) .await?; + // KCORE is Int32; cast to Int64 so `collect_to_i64` can read it. let result = read_result(&ctx, &output_uri) .await? - .select(vec![col(VERTEX_ID), col(KCORE)])?; + .select(vec![col(VERTEX_ID), cast(col(KCORE), DataType::Int64)])?; let values = collect_to_i64(&result.sort_by(vec![col(VERTEX_ID)])?, 1).await?; assert_eq!(values, &[2, 2, 2]); Ok(()) diff --git a/src/algorithm/community/classical_lp.rs b/src/algorithm/community/classical_lp.rs index cfbf88a..034624d 100644 --- a/src/algorithm/community/classical_lp.rs +++ b/src/algorithm/community/classical_lp.rs @@ -93,7 +93,7 @@ impl<'a> ClassicalLPBuilder<'a> { // Aggregate the neighbour labels carried by the *messages* // (`__pregel_msg_msg`), not the `community` vertex column, which // does not exist in the aggregated-messages frame. - .add_aggregate_expr(most_common_by(pregel_default_msg(), lit(1.0))) + .add_aggregate_expr(most_common_by(pregel_default_msg(), lit(1.0f32))) .max_iterations(self.max_iter) .skip_dest_state() .with_checkpoint_store(self.checkpoint_config.store_url.clone()) diff --git a/src/expressions.rs b/src/expressions.rs index bd6b36b..3c7a4c4 100644 --- a/src/expressions.rs +++ b/src/expressions.rs @@ -1,10 +1,10 @@ mod common; mod finite_axpb; mod hll; -mod kcore_merge; +mod kcore_reduce; mod most_common_by; pub(crate) use finite_axpb::{axpb, finite_axpb}; pub(crate) use hll::{hll_long, hll_long_aggregate, hll_long_estimate, hll_long_union}; -pub(crate) use kcore_merge::kcore_merge_expr; +pub(crate) use kcore_reduce::kcore_reduce; pub(crate) use most_common_by::most_common_by; diff --git a/src/expressions/common.rs b/src/expressions/common.rs index 291791a..fead671 100644 --- a/src/expressions/common.rs +++ b/src/expressions/common.rs @@ -1,4 +1,6 @@ -use datafusion::arrow::array::{Array, ArrayRef, BinaryArray, BinaryViewArray, Int64Array}; +use datafusion::arrow::array::{ + Array, ArrayRef, BinaryArray, BinaryViewArray, Int32Array, Int64Array, +}; use datafusion::error::{DataFusionError, Result}; /// Helper for other UDFs: vertexId and edgeSrc / edgeDst are all i64 @@ -15,6 +17,21 @@ pub(crate) fn downcast_int64<'a>( }) } +/// Helper for other UDFs: with assumption degree < i32::MAX +/// this one is a common thing. +pub(crate) fn downcast_int32<'a>( + array: &'a ArrayRef, + fname: &str, + label: &str, +) -> Result<&'a Int32Array> { + array.as_any().downcast_ref::().ok_or_else(|| { + DataFusionError::Plan(format!( + "{fname} {label} argument must be Int32, got: {:?}", + array.data_type() + )) + }) +} + /// Read-only accessor over either a [`BinaryArray`] or a [`BinaryViewArray`]. /// /// Parquet spills round-trip `Binary` sketch columns as `BinaryView` (DataFusion diff --git a/src/expressions/kcore_merge.rs b/src/expressions/kcore_merge.rs deleted file mode 100644 index 9f94aab..0000000 --- a/src/expressions/kcore_merge.rs +++ /dev/null @@ -1,327 +0,0 @@ -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; - -/// Per-vertex update rule of the distributed k-core decomposition. -/// -/// Given a vertex's neighbours' current core estimates and the vertex's own -/// current core, returns the largest value `l` such that at least `l` -/// neighbours have a core estimate `>= l`. The result is bounded above by -/// `current_core`, so core values are monotonically non-increasing. An empty -/// neighbour list yields `0`. -/// -/// This is the hot-path implementation used by the UDF below: the histogram -/// is written into a caller-provided buffer (reused across rows) and the -/// neighbour values are pulled from a lazy iterator, so evaluating a whole -/// batch performs no per-row heap allocation. `num_neighbors` is only an -/// upper bound on the real neighbour count (e.g. the list slot count, nulls -/// included); a looser bound merely widens the top of the histogram with zero -/// buckets and cannot change the result. -/// -/// 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 kcore_merge_into( - counts: &mut Vec, - num_neighbors: usize, - current_core: i64, - neighbors: impl Iterator, -) -> i64 { - // The new core can never exceed the number of neighbours (nor - // current_core), so the counts table only needs to be that large. This - // bounds the allocation independently of `current_core` (seeded from the - // degree), defending against a pathological value such as `i64::MAX`. - let cap = (current_core.max(0) as usize).min(num_neighbors); - counts.clear(); - counts.resize(cap + 1, 0); - for el in neighbors { - // Anything above `cap` lands in the top bucket; negatives clamp to 0. - let bucket = (el.max(0) as usize).min(cap); - counts[bucket] += 1; - } - let mut current_weight = 0i64; - for i in (1..=cap).rev() { - current_weight += counts[i]; - if (i as i64) <= current_weight { - return i as i64; - } - } - 0 -} - -fn list_int64_type() -> DataType { - DataType::List(Arc::new(Field::new("item", DataType::Int64, true))) -} - -/// DataFusion scalar UDF implementing the k-core update rule -/// ([`kcore_merge_into`]) over `(List, Int64) -> Int64`. -#[derive(Debug, PartialEq, Eq, Hash)] -pub(crate) struct KCoreMerge { - signature: Signature, -} - -impl KCoreMerge { - pub(crate) fn new() -> Self { - Self { - signature: Signature::exact( - vec![list_int64_type(), DataType::Int64], - 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), DataType::Int64] if f.data_type() == &DataType::Int64 => { - Ok(DataType::Int64) - } - _ => Err(DataFusionError::Plan(format!( - "kcore_merge expects (List, Int64), got: {arg_types:?}" - ))), - } - } - - fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let arrays = ColumnarValue::values_to_arrays(&args.args)?; - if arrays.len() != 2 { - return Err(DataFusionError::Plan(format!( - "kcore_merge expects exactly two arguments, got: {}", - arrays.len() - ))); - } - let list = downcast_list(&arrays[0], "first")?; - let cores = downcast_int64(&arrays[1], "second")?; - - let values = list.values(); - let values = values - .as_any() - .downcast_ref::() - .ok_or_else(|| { - DataFusionError::Plan(format!( - "kcore_merge list elements must be Int64, got: {:?}", - list.value_type() - )) - })?; - - let offsets = list.value_offsets(); - let len = arrays - .iter() - .map(|arr| arr.len()) - .max() - .unwrap_or(0) - .max(args.number_rows); - - // Reuse a single histogram buffer across rows so the hot path performs - // no per-row heap allocation; the neighbour values are read lazily - // straight off the Arrow array (skipping nulls) instead of being - // collected into a fresh `Vec` per row. - let mut counts: Vec = Vec::new(); - let result: Int64Array = (0..len) - .map(|i| { - let core_idx = i % cores.len(); - if cores.is_null(core_idx) { - return Some(0i64); - } - let current_core = cores.value(core_idx); - let row = i % list.len(); - if list.is_null(row) { - // No message this iteration: keep the current core. - return Some(current_core); - } - 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, - current_core, - neighbors, - )) - }) - .collect(); - - Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef)) - } -} - -fn downcast_list<'a>(array: &'a ArrayRef, label: &str) -> Result<&'a ListArray> { - array.as_any().downcast_ref::().ok_or_else(|| { - DataFusionError::Plan(format!( - "kcore_merge {label} argument must be a list of Int64, got: {:?}", - array.data_type() - )) - }) -} - -fn downcast_int64<'a>(array: &'a ArrayRef, label: &str) -> Result<&'a Int64Array> { - array.as_any().downcast_ref::().ok_or_else(|| { - DataFusionError::Plan(format!( - "kcore_merge {label} argument must be Int64, got: {:?}", - array.data_type() - )) - }) -} - -/// Builds an [`Expr`] that applies `kcore_merge(neighbors, current_core)`. -pub(crate) fn kcore_merge_expr(neighbors: Expr, current_core: Expr) -> Expr { - ScalarUDF::from(KCoreMerge::new()).call(vec![neighbors, current_core]) -} - -#[cfg(test)] -mod tests { - use super::*; - use datafusion::prelude::*; - - /// Thin slice-based wrapper around [`kcore_merge_into`] for the pure-function - /// tests below: hides the buffer plumbing behind the natural `(&[i64], i64)` - /// signature. - fn kcore_merge(neighbor_cores: &[i64], current_core: i64) -> i64 { - let mut counts = Vec::new(); - kcore_merge_into( - &mut counts, - neighbor_cores.len(), - current_core, - neighbor_cores.iter().copied(), - ) - } - - #[test] - fn test_merge_empty_returns_zero() { - assert_eq!(kcore_merge(&[], 5), 0); - assert_eq!(kcore_merge(&[], 0), 0); - } - - #[test] - fn test_merge_two_connected() { - // Single neighbour with core 1, current core 1 -> 1. - assert_eq!(kcore_merge(&[1], 1), 1); - } - - #[test] - fn test_merge_triangle() { - // Two neighbours both with core 2, current core 2 -> 2. - assert_eq!(kcore_merge(&[2, 2], 2), 2); - } - - #[test] - fn test_merge_star_center() { - // Center with degree 3: three leaves with core 1. - // largest l with >= l neighbours >= l: l=1 (3 neighbours >= 1). - assert_eq!(kcore_merge(&[1, 1, 1], 3), 1); - } - - #[test] - fn test_merge_capped_values() { - // Neighbour estimates above the current core are capped into the - // top bucket: [5, 5, 5], current 3 -> 3 neighbours >= 3 -> l=3. - assert_eq!(kcore_merge(&[5, 5, 5], 3), 3); - } - - #[test] - fn test_merge_decreases_core() { - // current core 4, neighbours [2, 2, 1]: - // i=4: weight 0; i=3: weight 0; i=2: weight 2 -> 2<=2 -> 2. - assert_eq!(kcore_merge(&[2, 2, 1], 4), 2); - } - - #[test] - fn test_merge_ignores_negatives() { - // Negative estimates are clamped to bucket 0 and cannot satisfy the - // condition for i >= 1. - assert_eq!(kcore_merge(&[-1, -1, 3], 3), 1); - } - - #[test] - fn test_merge_pathological_current_core_does_not_oom() { - // A corrupted / huge `current_core` must not allocate a buffer of that - // size. The result is bounded by the neighbour count, so a tiny - // allocation suffices and the answer is still correct. - assert_eq!(kcore_merge(&[3, 3, 3], i64::MAX), 3); - assert_eq!(kcore_merge(&[i64::MAX, i64::MAX], i64::MAX), 2); - assert_eq!(kcore_merge(&[], i64::MAX), 0); - } - - #[tokio::test] - async fn test_merge_udf_dataframe() -> Result<()> { - use datafusion::arrow::array::{Int64Builder, ListBuilder, RecordBatch}; - use datafusion::arrow::datatypes::{DataType, Field, Schema}; - use std::sync::Arc; - - // Build a List column: row 0 = [2,2], row 1 = [1,1,1]. - let mut list_builder = ListBuilder::new(Int64Builder::new()); - list_builder.values().append_value(2); - list_builder.values().append_value(2); - list_builder.append(true); - list_builder.values().append_value(1); - list_builder.values().append_value(1); - list_builder.values().append_value(1); - list_builder.append(true); - let list_arr = list_builder.finish(); - - let core_arr = Int64Array::from(vec![2, 3]); - - let schema = Arc::new(Schema::new(vec![ - Field::new("neighbors", list_int64_type(), false), - Field::new("core", DataType::Int64, false), - ])); - let batch = RecordBatch::try_new( - schema, - vec![ - Arc::new(list_arr) as ArrayRef, - Arc::new(core_arr) as ArrayRef, - ], - )?; - - let ctx = SessionContext::new(); - let df = ctx.read_batch(batch)?; - let result = df - .select(vec![ - kcore_merge_expr(col("neighbors"), col("core")).alias("k"), - ])? - .collect() - .await?; - - let arr = result[0] - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - - assert_eq!(arr.value(0), 2); - assert_eq!(arr.value(1), 1); - Ok(()) - } - - #[tokio::test] - async fn test_merge_udf_return_type_validation() { - let udf = KCoreMerge::new(); - assert_eq!( - udf.return_type(&[list_int64_type(), DataType::Int64]) - .unwrap(), - DataType::Int64 - ); - // Wrong arity / wrong type must error. - assert!( - udf.return_type(&[DataType::Int64, DataType::Int64]) - .is_err() - ); - assert!(udf.return_type(&[list_int64_type()]).is_err()); - } -} diff --git a/src/expressions/kcore_reduce.rs b/src/expressions/kcore_reduce.rs new file mode 100644 index 0000000..80e8ff6 --- /dev/null +++ b/src/expressions/kcore_reduce.rs @@ -0,0 +1,410 @@ +use std::sync::Arc; + +use crate::expressions::common::{as_binary_like, downcast_int32}; +use datafusion::arrow::array::ArrayRef; +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 KCoreReduceAccumulator { + // histogram of neighbors cores + counts: HashMap, +} + +impl KCoreReduceAccumulator { + pub(crate) fn new() -> Self { + Self { + counts: 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 + 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 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(()) + } + + /// 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; + } + } + + 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 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))); + } + + /// 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_by.rs b/src/expressions/most_common_by.rs index 12d781e..112ec30 100644 --- a/src/expressions/most_common_by.rs +++ b/src/expressions/most_common_by.rs @@ -1,7 +1,13 @@ +//! 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, Float64Array}; +use datafusion::arrow::array::{ArrayRef, Float32Array}; use datafusion::arrow::datatypes::{DataType, Field}; use datafusion::common::HashMap; use datafusion::error::{DataFusionError, Result}; @@ -14,7 +20,7 @@ use datafusion::scalar::ScalarValue; #[derive(Debug)] pub(crate) struct MostCommonByAccumulator { - sums: HashMap, + sums: HashMap, } impl MostCommonByAccumulator { @@ -25,10 +31,10 @@ impl MostCommonByAccumulator { } } -fn se_map(m: &HashMap) -> Vec { +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 + 16usize * (n as usize)); + 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() { @@ -39,23 +45,23 @@ fn se_map(m: &HashMap) -> Vec { buf } -fn de_map_and_insert(buf: &[u8], map: &mut HashMap) -> Result<()> { +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 + 16 * n { + 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 + 16 * i; + let o = 4 + 12 * i; let k = i64::from_le_bytes(buf[o..o + 8].try_into().unwrap()); - let v = f64::from_le_bytes(buf[o + 8..o + 16].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; } @@ -65,9 +71,9 @@ fn de_map_and_insert(buf: &[u8], map: &mut HashMap) -> Result<()> { 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( + let weights = &values[1].as_any().downcast_ref::().ok_or( DataFusionError::Execution( - "expected most_common_by secobd argument be f64 array".to_string(), + "expected most_common_by second argument be f32 array".to_string(), ), )?; @@ -92,7 +98,7 @@ impl Accumulator for MostCommonByAccumulator { return Ok(ScalarValue::Int64(None)); } let mut max = -1; - let mut max_value = f64::MIN; + let mut max_value = f32::MIN; for (k, v) in self.sums.iter() { if v > &max_value { @@ -112,7 +118,7 @@ impl Accumulator for MostCommonByAccumulator { } fn size(&self) -> usize { - size_of::() + self.sums.capacity() * (size_of::() + size_of::()) + size_of::() + self.sums.capacity() * (size_of::() + size_of::()) } fn state(&mut self) -> Result> { @@ -141,7 +147,7 @@ impl MostCommonBy { pub(crate) fn new() -> Self { Self { signature: Signature::exact( - vec![DataType::Int64, DataType::Float64], + vec![DataType::Int64, DataType::Float32], Volatility::Immutable, ), } @@ -161,11 +167,11 @@ impl AggregateUDFImpl for MostCommonBy { if (arg_types.len() != 2) || !matches!( (&arg_types[0], &arg_types[1]), - (DataType::Int64, DataType::Float64) + (DataType::Int64, DataType::Float32) ) { return Err(DataFusionError::Plan(format!( - "most_common_by expets exactly two arguments of types i64 and f64 but got {arg_types:?}" + "most_common_by expets exactly two arguments of types i64 and f32 but got {arg_types:?}" ))); } @@ -199,7 +205,7 @@ pub(crate) fn most_common_by(a: Expr, b: Expr) -> Expr { #[cfg(test)] mod tests { use super::*; - use datafusion::arrow::array::{BinaryArray, Float64Array, Int64Array, RecordBatch}; + use datafusion::arrow::array::{BinaryArray, Float32Array, Int64Array, RecordBatch}; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::datasource::MemTable; use datafusion::prelude::SessionConfig; @@ -212,7 +218,7 @@ mod tests { // 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(Float64Array::from(vec![2.0f64, 10.0, 3.0])) 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))); @@ -231,7 +237,7 @@ mod tests { let mut acc = MostCommonByAccumulator::new(); acc.update_batch(&[ Arc::new(Int64Array::from(vec![5i64, 2])) as ArrayRef, - Arc::new(Float64Array::from(vec![5.0f64, 5.0])) as ArrayRef, + Arc::new(Float32Array::from(vec![5.0f32, 5.0])) as ArrayRef, ]) .unwrap(); assert_eq!(acc.evaluate().unwrap(), ScalarValue::Int64(Some(2))); @@ -242,12 +248,12 @@ mod tests { /// bogus `n` and the length check rejects the blob). #[test] fn test_se_de_map_roundtrip_preserves_sums() { - let mut original: HashMap = HashMap::new(); + 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(); + 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 { @@ -263,9 +269,9 @@ mod tests { /// (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(); + let mut m: HashMap = HashMap::new(); m.insert(1, 3.0); - let mut partial: HashMap = HashMap::new(); + 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); @@ -283,7 +289,7 @@ mod tests { /// 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(); + 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. @@ -300,13 +306,13 @@ mod tests { // a: {1->3, 2->2} local winner 1 a.update_batch(&[ Arc::new(Int64Array::from(vec![1i64, 2])) as ArrayRef, - Arc::new(Float64Array::from(vec![3.0f64, 2.0])) 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(Float64Array::from(vec![4.0f64])) as ArrayRef, + Arc::new(Float32Array::from(vec![4.0f32])) as ArrayRef, ]) .unwrap(); @@ -330,7 +336,7 @@ mod tests { let df = dataframe!( "g" => vec![0i64, 0, 0, 0], "a" => vec![1i64, 2, 1, 2], - "b" => vec![5.0f64, 1.0, 1.0, 2.0], + "b" => vec![5.0f32, 1.0, 1.0, 2.0], )?; let out = df .aggregate( @@ -354,7 +360,7 @@ mod tests { async fn test_most_common_by_tie_picks_min_label() -> Result<()> { let df = dataframe!( "a" => vec![5i64, 2], - "b" => vec![5.0f64, 5.0], + "b" => vec![5.0f32, 5.0], )?; let out = df .aggregate(vec![], vec![most_common_by(col("a"), col("b")).alias("m")])? @@ -375,7 +381,7 @@ mod tests { let df = dataframe!( "g" => vec![0i64, 0, 0, 1, 1, 1], "a" => vec![1i64, 2, 1, 10, 20, 20], - "b" => vec![1.0f64, 3.0, 1.0, 5.0, 2.0, 2.0], + "b" => vec![1.0f32, 3.0, 1.0, 5.0, 2.0, 2.0], )?; let out = df .aggregate( @@ -415,9 +421,9 @@ mod tests { let schema = Arc::new(Schema::new(vec![ Field::new("g", DataType::Int64, false), Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Float64, false), + Field::new("b", DataType::Float32, false), ])); - let mk = |labels: Vec, weights: Vec| { + let mk = |labels: Vec, weights: Vec| { let n = labels.len(); RecordBatch::try_new( schema.clone(), @@ -426,14 +432,14 @@ mod tests { std::iter::repeat(0i64).take(n), )) as ArrayRef, Arc::new(Int64Array::from(labels)) as ArrayRef, - Arc::new(Float64Array::from(weights)) 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.0f64; 100]); - let p2 = mk(vec![2i64; 150], vec![1.0f64; 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)); @@ -460,12 +466,12 @@ mod tests { assert_eq!(m.value(0), 2); } - /// `return_type` must accept exactly (Int64, Float64) and reject the rest. + /// `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::Float64]) + udf.return_type(&[DataType::Int64, DataType::Float32]) .unwrap(), DataType::Int64 ); @@ -475,11 +481,11 @@ mod tests { .is_err() ); assert!( - udf.return_type(&[DataType::Float64, DataType::Float64]) + udf.return_type(&[DataType::Float64, DataType::Float32]) .is_err() ); assert!( - udf.return_type(&[DataType::Int64, DataType::Float64, DataType::Int64]) + udf.return_type(&[DataType::Int64, DataType::Float32, DataType::Int64]) .is_err() ); }