diff --git a/.gitignore b/.gitignore index f973868..76dc5c7 100644 --- a/.gitignore +++ b/.gitignore @@ -15,6 +15,7 @@ gf_df_tmp* gf_checkpoints* gf_workdir* test-test* +email-* .#* /graph500* /wiki-Talk* diff --git a/Cargo.lock b/Cargo.lock index 46940b1..a4bcbd8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -513,6 +513,12 @@ version = "3.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "46c5e41b57b8bba42a04676d81cb89e9ee8e859a1a66f80a5a72e1cb76b34d43" +[[package]] +name = "bytemuck" +version = "1.25.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" + [[package]] name = "byteorder" version = "1.5.0" @@ -1777,6 +1783,7 @@ dependencies = [ "tokio", "url", "uuid", + "wide", ] [[package]] @@ -2707,6 +2714,15 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f" +[[package]] +name = "safe_arch" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96b02de82ddbe1b636e6170c21be622223aea188ef2e139be0a5b219ec215323" +dependencies = [ + "bytemuck", +] + [[package]] name = "same-file" version = "1.0.6" @@ -3338,6 +3354,16 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "wide" +version = "0.7.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ce5da8ecb62bcd8ec8b7ea19f69a51275e91299be594ea5cc6ef7819e16cd03" +dependencies = [ + "bytemuck", + "safe_arch", +] + [[package]] name = "winapi-util" version = "0.1.9" diff --git a/Cargo.toml b/Cargo.toml index 0bfea02..477f4d2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,6 +19,7 @@ futures = "0.3" uuid = {version = "1", features = ["v4"] } log = "0.4" rand = { version = "0.9", features = ["std_rng"] } +wide = "0.7" datasketches = { version = "0.3.0", default-features = false, features = ["hll"] } # CLI dependencies diff --git a/README.md b/README.md index c4b098e..e2a8943 100644 --- a/README.md +++ b/README.md @@ -45,7 +45,9 @@ Available algorithms: - `kcore`: K-Core decomposition; - `hyperanf`: Approximate Neighbor Function; - `shortest-path`: Multi-Source Shortest Path; -- `classical-lp`: Classical (Raghavan) Label Propagaion +- `classical-lp`: Classical (Raghavan) Label Propagaion; +- `pic`: Power Iteration Clustering; +- `mllib kmeans`: raw K-Means over a `List` feature column of the vertices file (no edges read). Each has its own arguments — run `graphframes --help` for the full list. For what each algorithm computes, see [References](#references). @@ -103,6 +105,7 @@ TBD ### Community Detection - **Classical Label Propagation**: _Raghavan, Usha Nandini, Réka Albert, and Soundar Kumara. "Near linear time algorithm to detect community structures in large-scale networks." Physical Review E—Statistical, Nonlinear, and Soft Matter Physics 76.3 (2007): 036106._ +- **Power Iteration Clustering**: _Lin, Frank, and William W. Cohen. "Power iteration clustering." (2010)._ ### Subgraphs diff --git a/benches/python/main.py b/benches/python/main.py index 0a2db61..159a1db 100644 --- a/benches/python/main.py +++ b/benches/python/main.py @@ -57,6 +57,9 @@ "sp": {"cli": "shortest-path", "args": [], "undirected": True}, "cdlp": {"cli": "classical-lp", "args": ["--max-iter", "10"], "undirected": False}, "mis": {"cli": "mis", "args": [], "undirected": False}, + # PIC symmetrizes internally (the affinity matrix is symmetric by + # definition), so it needs no --symmetrize input handling. + "pic": {"cli": "pic", "args": ["--k", "2", "--max-iter", "20"], "undirected": False}, } diff --git a/src/algorithm.rs b/src/algorithm.rs index 95fe876..20d3e35 100644 --- a/src/algorithm.rs +++ b/src/algorithm.rs @@ -1,5 +1,5 @@ mod centrality; -mod community; +pub(crate) mod community; mod connectivity; mod pregel; mod subgraph; diff --git a/src/algorithm/centrality/hyperanf.rs b/src/algorithm/centrality/hyperanf.rs index eeac01e..9aaab34 100644 --- a/src/algorithm/centrality/hyperanf.rs +++ b/src/algorithm/centrality/hyperanf.rs @@ -95,6 +95,7 @@ impl<'a> HyperANFBuilder<'a> { .clone() .select_columns(&[EDGE_SRC, EDGE_DST])?, true, + None, )? }; diff --git a/src/algorithm/centrality/k_core.rs b/src/algorithm/centrality/k_core.rs index 6af56d0..77c0c25 100644 --- a/src/algorithm/centrality/k_core.rs +++ b/src/algorithm/centrality/k_core.rs @@ -87,6 +87,7 @@ impl<'a> KCoreBuilder<'a> { .clone() .select_columns(&[EDGE_SRC, EDGE_DST])?, true, + None, )?; // The undirected degree is the out-degree of the symmetrized graph. diff --git a/src/algorithm/community.rs b/src/algorithm/community.rs index 9dda726..9dda77f 100644 --- a/src/algorithm/community.rs +++ b/src/algorithm/community.rs @@ -1 +1,2 @@ mod classical_lp; +pub(crate) mod power_iteration_clustering; diff --git a/src/algorithm/community/classical_lp.rs b/src/algorithm/community/classical_lp.rs index ad63cff..a5580d9 100644 --- a/src/algorithm/community/classical_lp.rs +++ b/src/algorithm/community/classical_lp.rs @@ -74,6 +74,7 @@ impl<'a> ClassicalLPBuilder<'a> { .clone() .select_columns(&[EDGE_SRC, EDGE_DST])?, false, + None, )? }; diff --git a/src/algorithm/community/power_iteration_clustering.rs b/src/algorithm/community/power_iteration_clustering.rs new file mode 100644 index 0000000..6ac1bbc --- /dev/null +++ b/src/algorithm/community/power_iteration_clustering.rs @@ -0,0 +1,1202 @@ +//! Power Iteration Clustering (PIC). +//! +//! Truncated power iteration on the row-normalized affinity matrix: +//! starting from a non-constant `v_0`, the update `v_{t+1} = (D^-1 A) v_t`. +//! +//! References: +//! * Frank Lin and William W. Cohen, *Power Iteration Clustering*, ICML 2010 +//! (). +//! * Apache Spark MLlib +//! [`PowerIterationClustering`](https://github.com/apache/spark/blob/master/mllib/src/main/scala/org/apache/spark/mllib/clustering/PowerIterationClustering.scala) +//! — a direct inspiration for this implementation: the affinity contract +//! (non-negative similarities, self-loops dropped, symmetrized internally), +//! the convergence criterion, and the `degree` init vector the paper recommends. +//! +//! Edge weights: an optional `f64` weight column, else unit weights; an +//! optional PPMI (positive pointwise mutual information) transform replaces +//! the weights with `max(0, ln(w_ij·W / (d_i·d_j)))` computed from the raw +//! symmetrized statistics. Non-positive weights are dropped before use. +//! +//! Vertex coverage contract: the output contains exactly the vertices with +//! at least one incident edge of positive weight *after* symmetrization and +//! the selected weight transform. + +use std::sync::Arc; + +use datafusion::arrow::array::Float64Array; +use datafusion::arrow::datatypes::{DataType, Field}; +use datafusion::dataframe::DataFrameWriteOptions; +use datafusion::error::{DataFusionError, Result}; +use datafusion::execution::object_store::ObjectStoreUrl; +use datafusion::functions_aggregate::average::avg; +use datafusion::functions_aggregate::min_max::min as agg_min; +use datafusion::functions_aggregate::sum::sum; +use datafusion::object_store::path::Path; +use datafusion::prelude::*; +use rand::rngs::StdRng; +use rand::{Rng, RngCore, SeedableRng}; +use uuid::Uuid; + +use crate::expressions::{finite_axpb, kmeans_assign_expr}; +use crate::memory::ParquetCheckpointer; +use crate::ml::{DistanceMetric, KMeansBuilder}; +use crate::utils::scoped_ctx; +use crate::{ + EDGE_DST, EDGE_SRC, GraphFrame, memory::CheckpointConfig, ml::KMeansResult, utils::symmetrize, +}; +use crate::{EDGE_WEIGHT, GraphFramesConfig, VERTEX_ID}; + +async fn unsafe_first_f64(df: &DataFrame, column: usize) -> Result { + // do not use this outside PIC + let r = df + .clone() + .collect() + .await? + .first() + .unwrap() + .column(column) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + + Ok(r) +} + +/// This function is unsafe to use outside of PIC; +/// It mutates the checkpointer using an implicit contract with a single caller. +async fn ppmi( + edges: &DataFrame, + weight_col: &str, + checkpointer: &mut ParquetCheckpointer, + ctx: &SessionContext, +) -> Result { + // validation of the input is not the responsibility of this function; + let cached_edges = checkpointer.push(ctx, "__raw_edges", edges.clone()).await?; + let d_raw = cached_edges.clone().aggregate( + vec![col(EDGE_SRC).alias(VERTEX_ID)], + vec![sum(col(EDGE_WEIGHT)).alias("d_raw")], + )?; + let w = unsafe_first_f64( + &d_raw.clone().aggregate(vec![], vec![sum(col("d_raw"))])?, + 0, + ) + .await?; + + let ee = cached_edges + .join( + d_raw.clone().select(vec![ + col(VERTEX_ID).alias("__src_id"), + col("d_raw").alias("src_d_raw"), + ])?, + JoinType::Left, + &vec![EDGE_SRC], + &vec!["__src_id"], + None, + )? + .join( + d_raw.select(vec![ + col(VERTEX_ID).alias("__dst_id"), + col("d_raw").alias("dst_d_raw"), + ])?, + JoinType::Left, + &vec![EDGE_DST], + &vec!["__dst_id"], + None, + )? + .select(vec![ + col(EDGE_SRC), + col(EDGE_DST), + // ppmi_ij = max(0, ln(w_ij·W/(d_raw_i·d_raw_j))) + greatest(vec![ + lit(0f64), + ln(col(weight_col).mul(lit(w).div(col("src_d_raw").mul(col("dst_d_raw"))))), + ]) + .alias(EDGE_WEIGHT), + ])? + .filter(col(EDGE_WEIGHT).gt(lit(0.0f64)))?; + + let r = checkpointer + .push_pre_sorted(ctx, "edges", ee, EDGE_SRC) + .await?; + checkpointer.evict_all_but_latest_n(ctx, 1).await?; + + Ok(r) +} + +#[derive(Debug, Copy, Eq, PartialEq, Hash, Clone)] +pub enum InitStrategy { + Random, + DegreeBased, +} + +#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)] +pub enum WeightsStrategy { + None, + PPMI, +} + +/// What KMeans consumes and what the `embedding` output column holds. +/// +/// The paper (Lin & Cohen, ICML 2010) clusters on the *final* +/// iterate only: "k-means to cluster points on vt". +/// +/// `FullTrajectory` is the extended mode: the iterate history `[v1..vm]` as +/// one vector (embedding width = executed iterations). +#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)] +pub enum EmbeddingMode { + /// Paper mode (default): the final iterate `v_m` — a 1-wide embedding. + LastIterate, + /// Extended mode: the full iterate history `[v1..vm]` + FullTrajectory, +} + +#[derive(Debug)] +pub struct PICBuilder { + graph: GraphFrame, + max_iterations: usize, + /// Convergence threshold on |delta_t − delta_{t−1}| for the (already + /// mass-normalized) relative delta — see the module docs. + tol: f64, + init_strategy: InitStrategy, + weight_col: Option, + weights_strategy: WeightsStrategy, + embedding_mode: EmbeddingMode, + k: Vec, + checkpoint_config: CheckpointConfig, + seed: u64, +} + +impl PICBuilder { + pub fn new(graph: GraphFrame) -> Self { + Self { + graph: graph, + tol: 1e-5, + // Spark MLlib defaults to 100 for its dense affinity setting; the + // paper's datasets average 13. On large sparse graphs λ2 ~= 1 and + // the acceleration criterion cannot fire before the cluster + // signal decays into the uniform mode. + max_iterations: 20usize, + init_strategy: InitStrategy::DegreeBased, + weight_col: None, // unweighted graph + weights_strategy: WeightsStrategy::None, + embedding_mode: EmbeddingMode::LastIterate, // the paper's mode + k: vec![2], // mirroring SparkML' default + checkpoint_config: CheckpointConfig::default_local_fs(), + seed: 42u64, + } + } + + /// Set maximal amount of iterations. + /// PIC should converge fast by itself, + /// so this is more than a safeguard than a control parameter. + pub fn set_max_iterations(mut self, v: usize) -> Self { + self.max_iterations = v; + self + } + + /// Convergence threshold for the acceleration criterion. The delta is + /// mass-normalized (relative), so it does not depend on the graph size. + /// + /// Note: you probably want to increase the default value, not decrease it. + pub fn set_tol(mut self, v: f64) -> Self { + self.tol = v; + self + } + + /// Set a weights column name + pub fn set_edge_weight_col(mut self, c: &str) -> Self { + self.weight_col = Some(c.to_string()); + self + } + + /// Set a strategy of weights transformation. + /// Be aware that while PPMI may be better, it is more expensive. + pub fn set_weights_strategy(mut self, s: WeightsStrategy) -> Self { + self.weights_strategy = s; + self + } + + /// Set init strategy (random or degree based) + pub fn set_init_strategy(mut self, v: InitStrategy) -> Self { + self.init_strategy = v; + self + } + + /// Select what the `embedding` column holds and KMeans clusters on. + /// Default: last iteration like in the paper; + pub fn set_embedding_mode(mut self, v: EmbeddingMode) -> Self { + self.embedding_mode = v; + self + } + + /// Set multiple k: to check different number of clusters in parallel + pub fn set_multiple_k(mut self, kk: Vec) -> Self { + self.k = kk; + self + } + + /// Set a single k: number of expected clusters + pub fn set_k(mut self, k: usize) -> Self { + self.k = vec![k]; + self + } + + /// Set the object store URL + pub fn with_checkpoint_store(mut self, store_url: ObjectStoreUrl) -> Self { + self.checkpoint_config.store_url = store_url; + self + } + + /// Set the checkpoint directory + pub fn set_checkpoint_dir(mut self, dir: Path) -> Self { + self.checkpoint_config.dir = dir; + self + } + + /// Set random seed + pub fn set_seet(mut self, seed: u64) -> Self { + self.seed = seed; + self + } + + pub async fn run(self, ctx: &SessionContext, output: &str) -> Result { + if self.max_iterations < 2 { + // that is a very specific usecase; + return Err(DataFusionError::Plan( + "max_iteartiaons should be greater than 2".to_string(), + )); + } + let gf_config = ctx + .state() + .config() + .options() + .extensions + .get::() + .cloned() + .unwrap_or_default(); + + let ctx = &scoped_ctx(ctx, gf_config.prefer_smj); + self.checkpoint_config.validate_output(output)?; + + let run_id = Uuid::new_v4().to_string(); + log::info!("start PIC with ID {run_id}"); + + // original vertices should come from disk + // so the op is on original, not filtered; + // this may differ up to the normalization term + // on graphs with dangling nodes; + let n = self.graph.vertices.clone().count().await?; + let tol = if 1e-5f64 / (n as f64) > 1e-8 { + 1e-5f64 / (n as f64) + } else { + 1e-8f64 + }; + + let mut edges_checkpointer = ParquetCheckpointer::new( + self.checkpoint_config.store_url.clone(), + self.checkpoint_config + .dir + .clone() + .join(run_id.clone()) + .join("edges_state"), + ); + + let mut state_checkpointer = ParquetCheckpointer::new( + self.checkpoint_config.store_url.clone(), + self.checkpoint_config + .dir + .clone() + .join(run_id.clone()) + .join("vertex_state"), + ); + + // we will keep links there; + let mut states = Vec::::new(); + + // weight columns should: + // a) exists + // b) be an f32 data type + let w_col = match self.weight_col { + Some(c) => { + if !self + .graph + .edges + .schema() + .has_column_with_unqualified_name(&c) + { + return Err(DataFusionError::Plan(format!( + "column {} does not exist in edges", + c + ))); + } + + let resolved = self.graph.edges.schema().field_with_unqualified_name(&c)?; + + if resolved.data_type() != &DataType::Float64 { + return Err(DataFusionError::Plan(format!( + "weight column {} has data type {} file expected float64", + c, + resolved.data_type() + ))); + } + + col(c).alias(EDGE_WEIGHT) + } + None => lit(1.0f64).alias(EDGE_WEIGHT), + }; + + let symmetrized_edges = symmetrize( + &self + .graph + .edges + .clone() + .select(vec![col(EDGE_SRC), col(EDGE_DST), w_col])?, + false, + Some(vec![EDGE_WEIGHT.to_string()]), + )?; + + // contract: we are storing "weights" on edges; + let edges = match self.weights_strategy { + WeightsStrategy::None => { + edges_checkpointer + .push_pre_sorted(ctx, "edges", symmetrized_edges.clone(), EDGE_SRC) + .await? + } + WeightsStrategy::PPMI => { + // checkpointing is responsibility of the "ppmi" function; + // think about it as a bad design but as is :) + ppmi( + &symmetrized_edges, + EDGE_WEIGHT, + &mut edges_checkpointer, + ctx, + ) + .await? + } + }; + + let min_w = unsafe_first_f64( + &edges + .clone() + .aggregate(vec![], vec![agg_min(col(EDGE_WEIGHT))])?, + 0, + ) + .await?; + if min_w < 0.0f64 { + return Err(DataFusionError::Plan( + "weights can be negative in PIC!".to_string(), + )); + } + + let mut rng = StdRng::seed_from_u64(self.seed); + + // state: v0 + weighted degree; + // weights: f64 for simplicity; + // eignevector: f32 for better and lightweight KMeans + let mut state = { + let s = edges.clone().aggregate( + vec![col(EDGE_SRC).alias(VERTEX_ID)], + vec![sum(col(EDGE_WEIGHT)).alias("out_deg")], + )?; + + let v0 = match self.init_strategy { + // we approximate Z by n / 2 + InitStrategy::Random => { + let mut r_a = rng.random::(); + while r_a == 0 { + r_a = rng.random::(); + } + let r_b = rng.random::(); + s.with_column( + "v0", + abs(finite_axpb(lit(r_a), col(VERTEX_ID), lit(r_b))) + .div(lit(9_223_372_036_854_776_000.0f64)) // 2^63 + .mul(lit(2.0 / (n as f64))), + )? + } + InitStrategy::DegreeBased => { + // cast to f64 to avoid lost in precision; + // we are not checking here corner cases like graph is empty. + let n = unsafe_first_f64( + &s.clone().aggregate(vec![], vec![sum(col("out_deg"))])?, + 0, + ) + .await?; + + s.with_column("v0", cast(col("out_deg").div(lit(n)), DataType::Float64))? + } + }; + + state_checkpointer + .push_pre_sorted( + ctx, + "state_0", + v0.select(vec![col(VERTEX_ID), col("out_deg"), col("v0")])?, + VERTEX_ID, + ) + .await? + }; + + let mut old_delta = + unsafe_first_f64(&state.clone().aggregate(vec![], vec![sum(col("v0"))])?, 0).await?; + let mut converged = false; + let mut iteration = 0usize; + + while !converged && (iteration < self.max_iterations) { + let triplets = edges.clone().join( + state.clone(), + JoinType::Left, + &vec![EDGE_SRC], + &[VERTEX_ID], + None, + )?; + let msgs = triplets.aggregate( + vec![col(EDGE_DST)], + vec![ + sum(cast(col(format!("v{}", iteration)), DataType::Float64) + .mul(col(EDGE_WEIGHT))) + .alias("msg"), + ], + )?; + iteration += 1; + + state = state_checkpointer + .push_pre_sorted( + ctx, + &format!("state_{}", iteration), + state + .clone() + .join( + msgs, + JoinType::Inner, + &vec![VERTEX_ID], + &vec![EDGE_DST], + None, + )? + .select(vec![ + col(VERTEX_ID), + col("out_deg"), + cast(col("msg").div(col("out_deg")), DataType::Float32) + .alias(format!("v{}", iteration)), + abs(cast(col("msg").div(col("out_deg")), DataType::Float32) + .sub(col(format!("v{}", iteration - 1)))) + .alias("delta"), + ])?, + VERTEX_ID, + ) + .await?; + + states.push(state.clone()); + + // Relative (mass-normalized) delta: sum|v_t − v_{t−1}| / sum(v_t). + let v_col = format!("v{iteration}"); + let stats = state.clone().aggregate( + vec![], + vec![ + sum(cast(col("delta"), DataType::Float64)).alias("__d"), + sum(cast(col(&v_col), DataType::Float64)).alias("__m"), + ], + )?; + let raw_delta = unsafe_first_f64(&stats.clone(), 0).await?; + let mass = unsafe_first_f64(&stats, 1).await?; + let new_delta = if mass.is_finite() && mass > 0.0 { + raw_delta / mass + } else { + raw_delta + }; + + log::info!( + "iteration {iteration} completed: relative delta={new_delta:.3e} (mass={mass:.3})" + ); + + // at least 1 iterataion we need; + // otherwise a risk of something strange due to random init. + if iteration > 1 { + if (old_delta - new_delta).abs() < tol { + if new_delta > 1e-3 { + log::warn!( + "PIC stopped with a still-large delta ({new_delta:.3e}): \ + period-2 oscillation (bipartite graph?) or slow mixing; \ + the embedding is truncated" + ); + } + converged = true; + log::info!("converged at iteration {iteration}"); + } + } + old_delta = new_delta; + } + + if !converged { + if old_delta <= 1e-4 { + log::info!( + "PIC stopped on the iteration budget ({}) with the iterate \ + essentially settled (last relative delta {old_delta:.3e} ≤ 1e-4); \ + raise --tol to {tol:.0e} if you want the exact criterion to fire", + self.max_iterations + ); + } else { + log::warn!( + "PIC stopped on the iteration budget ({}) with the iterate \ + still moving (last relative delta {old_delta:.3e}): slow mixing; \ + the embedding is intentionally truncated — try a lower budget, \ + or --embedding trajectory", + self.max_iterations + ); + } + } + + // clean up; + edges_checkpointer.purge(ctx).await?; + + // reconstruction/embedding: what KMeans consumes. + let avg_v = unsafe_first_f64( + &state.clone().aggregate( + vec![], + vec![avg(cast(col(format!("v{}", iteration)), DataType::Float64))], + )?, + 0, + ) + .await?; + let scale = if avg_v.is_finite() && avg_v > 0.0 { + 1.0 / avg_v + } else { + 1.0 + }; + + let e = match self.embedding_mode { + EmbeddingMode::LastIterate => state.clone().select(vec![ + col(VERTEX_ID), + cast( + make_array(vec![col(format!("v{}", iteration)).mul(lit(scale))]), + DataType::List(Arc::new(Field::new("el", DataType::Float32, true))), + ) + .alias("embedding"), + ])?, + EmbeddingMode::FullTrajectory => { + let mut assembled = states[0].clone().select(vec![col(VERTEX_ID), col("v1")])?; // explicitly prune projection; + + for (t, snap) in states.iter().enumerate().skip(1) { + let v_col = format!("v{}", t + 1); // states[t] holds v_{t+1} + assembled = assembled + .join_on( + snap.clone() + .select(vec![col(VERTEX_ID).alias("__vid"), col(&v_col)])?, + JoinType::Inner, + vec![col(VERTEX_ID).eq(col("__vid"))], + )? + .drop_columns(&["__vid"])?; + } + + assembled.select(vec![ + col(VERTEX_ID), + cast( + make_array( + (1..=iteration) + .map(|t| col(format!("v{t}")).mul(lit(scale))) + .collect::>(), + ), + DataType::List(Arc::new(Field::new("el", DataType::Float32, true))), + ) + .alias("embedding"), + ])? + } + }; + + let embedding = state_checkpointer.push(ctx, "embedding", e).await?; + state_checkpointer.evict_all_but_latest_n(ctx, 1).await?; + + let kmeans = KMeansBuilder::new(&embedding, "embedding") + .seed(rng.next_u64()) + .k_values(&self.k) + .metric(DistanceMetric::L2); + + log::info!("run KMeans on resulted embedding..."); + let rr = kmeans.run().await?; + let num_kmeans_iters = rr.num_iterations; + + log::info!("KMeans converged after {num_kmeans_iters} iterations"); + + let mut final_columns = vec![col(VERTEX_ID), col("embedding")]; + + let mut ks: Vec = Vec::new(); + + for &kk in &self.k { + if !ks.contains(&kk) { + ks.push(kk); + } + } // mirror KMeansBuilder's dedup + for (kk, run) in ks.iter().zip(&rr.runs) { + final_columns.push( + kmeans_assign_expr( + col("embedding"), + run.k, + rr.d, + run.centers.clone(), + DistanceMetric::L2, + ) // run.k = k_eff + .alias(format!("center_{kk}")), + ); + } + + let final_df = embedding.select(final_columns)?; + final_df + .write_parquet(output, DataFrameWriteOptions::new(), None) + .await?; + log::info!("result was written into {output}"); + + state_checkpointer.purge(ctx).await?; + + Ok(rr) + } +} + +impl GraphFrame { + /// Create a new Power Iteration Clustering builder. + pub fn pic(&self) -> PICBuilder { + PICBuilder::new(self.clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::utils::collect_to_i64; + use std::fs; + use std::path::PathBuf; + use std::process::id; + use std::sync::atomic::{AtomicU64, Ordering}; + use url::Url; + + static COUNTER: AtomicU64 = AtomicU64::new(0); + + /// Returns a unique directory under `std::env::temp_dir()` for this test run, creating it. + /// Combining the PID with a process-wide counter guarantees uniqueness across parallel + /// `cargo test` invocations and concurrent tests. + fn unique_temp_dir(label: &str) -> PathBuf { + let n = COUNTER.fetch_add(1, Ordering::SeqCst); + let dir = std::env::temp_dir().join(format!("graphframes_pic_test_{}_{n}_{label}", id())); + fs::create_dir_all(&dir).expect("failed to create unique temp dir"); + dir + } + + /// RAII guard that recursively removes the temp directory when dropped, so tests stay + /// self-contained without depending on the `tempfile` crate. + struct TempGuard(PathBuf); + impl Drop for TempGuard { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } + } + + /// Builds a `SessionContext`, an object_store checkpoint `Path`, an output `file://` URI, + /// and a `TempGuard` that cleans up on drop. The parent temp dir contains two non-overlapping + /// siblings — `checkpoints/` and `output/` — so `validate_output` is satisfied. + fn setup(label: &str) -> Result<(SessionContext, Path, String, TempGuard)> { + let parent = unique_temp_dir(label); + let checkpoint_root = parent.join("checkpoints"); + let output_root = parent.join("output"); + fs::create_dir_all(&checkpoint_root).expect("failed to create checkpoint dir"); + fs::create_dir_all(&output_root).expect("failed to create output dir"); + + let checkpoint_dir = Path::from_filesystem_path(&checkpoint_root) + .expect("checkpoint dir must be convertible to object_store path"); + let output_uri = Url::from_directory_path(&output_root) + .expect("output dir must be convertible to file:// URL") + .to_string(); + + let ctx = SessionContext::new(); + Ok((ctx, checkpoint_dir, output_uri, TempGuard(parent))) + } + + /// Builds a small unweighted `GraphFrame` from vertex ids and `(src, dst)` edges. + fn create_graph(vertices: Vec, edges: Vec<(i64, i64)>) -> Result { + let vertices_df = dataframe!(VERTEX_ID => vertices)?; + let (srcs, dsts): (Vec, Vec) = edges.into_iter().unzip(); + let edges_df = dataframe!(EDGE_SRC => srcs, EDGE_DST => dsts)?; + Ok(GraphFrame { + vertices: vertices_df, + edges: edges_df, + }) + } + + /// Builds a weighted `GraphFrame`; weights are `f64` (the PIC builder contract). + fn create_weighted_graph( + vertices: Vec, + edges: Vec<(i64, i64, f64)>, + ) -> Result { + let vertices_df = dataframe!(VERTEX_ID => vertices)?; + let (srcs, dsts, ws): (Vec, Vec, Vec) = edges.into_iter().fold( + (Vec::new(), Vec::new(), Vec::new()), + |mut acc, (s, d, w)| { + acc.0.push(s); + acc.1.push(d); + acc.2.push(w); + acc + }, + ); + let edges_df = dataframe!(EDGE_SRC => srcs, EDGE_DST => dsts, EDGE_WEIGHT => ws)?; + Ok(GraphFrame { + vertices: vertices_df, + edges: edges_df, + }) + } + + /// Reads back the written PIC output, sorted by id, cached. + async fn read_output(ctx: &SessionContext, output_uri: &str) -> Result { + Ok(ctx + .read_parquet(output_uri, ParquetReadOptions::default()) + .await? + .sort(vec![col(VERTEX_ID).sort(true, false)])? + .cache() + .await?) + } + + /// Two asymmetric cliques joined by a bridge — the smallest graph where PIC + /// has a real signal (two non-equivalent orbits). The symmetric two-triangle + /// version is *not* usable: automorphic clusters are indistinguishable by + /// any spectral method and the embedding gap is zero by construction. + fn two_cliques() -> Result { + create_graph( + vec![1, 2, 3, 4, 5, 6, 7], + vec![ + // 4-clique A = {1,2,3,4} + (1, 2), + (1, 3), + (1, 4), + (2, 3), + (2, 4), + (3, 4), + // bridge + (4, 5), + // 3-clique B = {5,6,7} + (5, 6), + (5, 7), + (6, 7), + ], + ) + } + + /// Smoke test: degree-init PIC runs end-to-end, writes non-empty parquet + /// with exactly the edge-touching vertices, the embedding column plus one + /// assignment column per requested K. + #[tokio::test] + async fn test_pic_run_writes_output() -> Result<()> { + let graph = two_cliques()?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("run_writes_output")?; + let kmeans = PICBuilder::new(graph) + .set_k(2) + .set_max_iterations(10) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + assert!(kmeans.num_iterations >= 1, "KMeans must run at least once"); + // Paper mode (default): the embedding is the final iterate only. + assert_eq!(kmeans.d, 1, "LastIterate embedding must be 1-wide"); + + let out = read_output(&ctx, &output_uri).await?; + assert_eq!( + out.clone().count().await?, + 7, + "all 7 edge-touching vertices must be in the output" + ); + let schema = out.schema(); + assert_eq!(schema.fields().len(), 1 + 1 + 1, "id, embedding, center_2"); + assert!(schema.has_column_with_unqualified_name(VERTEX_ID)); + assert!(schema.has_column_with_unqualified_name("embedding")); + assert!(schema.has_column_with_unqualified_name("center_2")); + + // Assembly pin: every embedding list is exactly `kmeans.d` long + // (1 in the default LastIterate mode). + let widths = out + .clone() + .select(vec![ + cast( + datafusion::functions_nested::expr_fn::array_length(col("embedding")), + DataType::Int64, + ) + .alias("w"), + ])? + .cache() + .await?; + let widths = collect_to_i64(&widths, 0).await?; + assert!( + widths.iter().all(|&w| w as usize == kmeans.d), + "embedding widths {widths:?} must all equal d = {}", + kmeans.d + ); + Ok(()) + } + + /// Extended mode: `FullTrajectory` embeds the iterate history `[v1..vm]`; + /// the embedding width must equal the executed iteration count (within + /// the max_iterations = 10 cap), for every row — pinning the assembly + /// join-chain wiring (`states[t]` holds `v_{t+1}`) and the `make_array` + /// argument count. + #[tokio::test] + async fn test_pic_full_trajectory_mode() -> Result<()> { + let graph = two_cliques()?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("full_trajectory")?; + let kmeans = PICBuilder::new(graph) + .set_k(2) + .set_embedding_mode(EmbeddingMode::FullTrajectory) + .set_max_iterations(10) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + assert!( + kmeans.d >= 2 && kmeans.d <= 10, + "history embedding width should equal the executed iterations, got {}", + kmeans.d + ); + + let out = read_output(&ctx, &output_uri).await?; + assert_eq!(out.clone().count().await?, 7); + let widths = out + .clone() + .select(vec![ + cast( + datafusion::functions_nested::expr_fn::array_length(col("embedding")), + DataType::Int64, + ) + .alias("w"), + ])? + .cache() + .await?; + let widths = collect_to_i64(&widths, 0).await?; + assert!( + widths.iter().all(|&w| w as usize == kmeans.d), + "embedding widths {widths:?} must all equal d = {}", + kmeans.d + ); + Ok(()) + } + + /// A vertex with no incident edge is excluded from the output by contract: + /// its embedding would be a convention (zero row), not a measurement. + #[tokio::test] + async fn test_pic_isolated_vertex_excluded() -> Result<()> { + let graph = create_graph(vec![1, 2, 3, 9], vec![(1, 2), (2, 3)])?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("isolated_vertex")?; + PICBuilder::new(graph) + .set_k(2) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let out = read_output(&ctx, &output_uri).await?; + let ids = collect_to_i64(&out, 0).await?; + assert_eq!(ids, vec![1, 2, 3], "isolated vertex 9 must be excluded"); + Ok(()) + } + + /// Self-loops are dropped by `symmetrize`; the loop count and the degree + /// must be unaffected by them. + #[tokio::test] + async fn test_pic_self_loops_ignored() -> Result<()> { + let graph = create_graph(vec![1, 2, 3], vec![(1, 2), (2, 3), (1, 1), (2, 2)])?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("self_loops")?; + PICBuilder::new(graph) + .set_k(2) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let out = read_output(&ctx, &output_uri).await?; + assert_eq!(out.clone().count().await?, 3); + Ok(()) + } + + /// `set_max_iterations(1)` (and 0) must be rejected: the acceleration + /// criterion needs at least two iterations, and the snapshot assembly + /// indexes `states[0]`. + #[tokio::test] + async fn test_pic_rejects_too_small_max_iterations() -> Result<()> { + let graph = two_cliques()?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("too_small_max_iter")?; + let res = PICBuilder::new(graph) + .set_k(2) + .set_max_iterations(1) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await; + assert!(res.is_err(), "max_iterations < 2 must be rejected"); + Ok(()) + } + + /// Input stored in one direction only: symmetrization must produce the + /// same affinity as the both-directions spelling (same embedding within + /// tolerance on this fixture). + #[tokio::test] + async fn test_pic_symmetrization_from_single_direction() -> Result<()> { + let graph = two_cliques()?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("symmetrize_one_dir")?; + PICBuilder::new(graph) + .set_k(2) + .set_init_strategy(InitStrategy::DegreeBased) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let out = read_output(&ctx, &output_uri).await?; + // All 7 vertices present; assignment column readable as i64 via cast. + let assignments = out + .clone() + .select(vec![cast(col("center_2"), DataType::Int64)])? + .sort(vec![col(VERTEX_ID).sort(true, false)])? + .cache() + .await?; + let labels = collect_to_i64(&assignments, 0).await?; + assert_eq!(labels.len(), 7); + Ok(()) + } + + /// Multiple K values produce one `center_k` column per requested K in + /// request order, and KMeans clamping (k_eff < k) must not panic — + /// the run is matched positionally, not by k value. + #[tokio::test] + async fn test_pic_multiple_k_produces_one_column_per_k() -> Result<()> { + // 4 distinct embedding values at most on this small graph; k=8 must + // clamp without panicking and still produce a `center_8` column. + let graph = two_cliques()?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("multiple_k")?; + let kmeans = PICBuilder::new(graph) + .set_multiple_k(vec![2, 3, 8]) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + assert_eq!(kmeans.runs.len(), 3, "one run per requested K"); + + let out = read_output(&ctx, &output_uri).await?; + let schema = out.schema(); + for name in ["center_2", "center_3", "center_8"] { + assert!( + schema.has_column_with_unqualified_name(name), + "missing {name} in output" + ); + } + Ok(()) + } + + /// Negative weights are rejected with an error (not a panic). + #[tokio::test] + async fn test_pic_negative_weights_rejected() -> Result<()> { + let graph = + create_weighted_graph(vec![1, 2, 3], vec![(1, 2, 1.0), (2, 3, -0.5), (3, 1, 1.0)])?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("negative_weights")?; + let res = PICBuilder::new(graph) + .set_k(2) + .set_edge_weight_col(EDGE_WEIGHT) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await; + match res { + Err(e) => { + let msg = e.to_string(); + assert!( + msg.contains("negative"), + "expected a negative-weight error, got: {msg}" + ); + } + Ok(_) => panic!("negative weights must be rejected"), + } + Ok(()) + } + + /// A non-existent weight column is rejected at plan time. + #[tokio::test] + async fn test_pic_missing_weight_column_rejected() -> Result<()> { + let graph = two_cliques()?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("missing_weight_col")?; + let res = PICBuilder::new(graph) + .set_k(2) + .set_edge_weight_col("no_such_col") + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await; + assert!(res.is_err(), "missing weight column must be rejected"); + Ok(()) + } + + /// Weighted run with a positive-weight column: same shape as unweighted, + /// same vertex coverage. + #[tokio::test] + async fn test_pic_weighted_run_writes_output() -> Result<()> { + let graph = create_weighted_graph( + vec![1, 2, 3, 4, 5, 6, 7], + vec![ + (1, 2, 1.0), + (1, 3, 1.0), + (1, 4, 1.0), + (2, 3, 1.0), + (2, 4, 1.0), + (3, 4, 1.0), + (4, 5, 0.1), + (5, 6, 1.0), + (5, 7, 1.0), + (6, 7, 1.0), + ], + )?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("weighted_run")?; + PICBuilder::new(graph) + .set_k(2) + .set_edge_weight_col(EDGE_WEIGHT) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let out = read_output(&ctx, &output_uri).await?; + assert_eq!(out.clone().count().await?, 7); + Ok(()) + } + + /// Random init: completes end-to-end with the analytic normalization + /// (v0 scaled by 2/n instead of an exact sum). Determinism is pinned by + /// `test_pic_deterministic_given_seed` (the seed is fixed by the builder). + #[tokio::test] + async fn test_pic_random_init_completes() -> Result<()> { + let graph = two_cliques()?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("random_init")?; + PICBuilder::new(graph) + .set_k(2) + .set_init_strategy(InitStrategy::Random) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let out = read_output(&ctx, &output_uri).await?; + assert_eq!(out.clone().count().await?, 7); + // embedding is a non-empty list column; read it back to force errors + let _ = out.clone().select(vec![col("embedding")])?.count().await?; + Ok(()) + } + + /// A star graph is period-2 after degree init: the acceleration criterion + /// stops early with the oscillation warning; the run must still complete + /// and write output (truncated embedding). + #[tokio::test] + async fn test_pic_star_graph_oscillation_completes() -> Result<()> { + // star with 4 leaves, ids 1..5 + let graph = create_graph(vec![1, 2, 3, 4, 5], vec![(1, 2), (1, 3), (1, 4), (1, 5)])?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("star_oscillation")?; + PICBuilder::new(graph) + .set_k(2) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let out = read_output(&ctx, &output_uri).await?; + assert_eq!(out.clone().count().await?, 5); + Ok(()) + } + + /// The PPMI transform: hand-computed reference on a 3-edge symmetric graph. + /// Edges (one direction stored): 1-2 w=3, 1-3 w=1, 2-3 w=2. + /// Symmetrized row sums: d1=4, d2=5, d3=3; W=12. + /// ppmi(1,2) = max(0, ln(3*12/(4*5))) = ln(1.8) > 0 -> kept + /// ppmi(1,3) = max(0, ln(1*12/(4*3))) = ln(1.0) = 0 -> dropped + /// ppmi(2,3) = max(0, ln(2*12/(5*3))) = ln(1.6) > 0 -> kept + /// so the transformed graph has 2 undirected edges: 1-2 and 2-3. + #[tokio::test] + async fn test_pic_ppmi_drops_unassociated_edges() -> Result<()> { + let graph = + create_weighted_graph(vec![1, 2, 3], vec![(1, 2, 3.0), (1, 3, 1.0), (2, 3, 2.0)])?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("ppmi_drop")?; + PICBuilder::new(graph) + .set_k(2) + .set_edge_weight_col(EDGE_WEIGHT) + .set_weights_strategy(WeightsStrategy::PPMI) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + // Vertex 3 keeps edge 2-3 (kept), vertex 1 keeps 1-2: all 3 vertices + // remain in the output; the drop of edge 1-3 is observable only + // through the embedding, pinned by the converged iteration count + // budget below. + let out = read_output(&ctx, &output_uri).await?; + let ids = collect_to_i64(&out, 0).await?; + assert_eq!(ids, vec![1, 2, 3]); + Ok(()) + } + + /// Empty edge set: no vertex is embedded; the run must not panic. + #[tokio::test] + async fn test_pic_empty_edges() -> Result<()> { + let graph = create_graph(vec![1, 2], vec![])?; + let (ctx, checkpoint_dir, output_uri, _guard) = setup("empty_edges")?; + let res = PICBuilder::new(graph) + .set_k(2) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await; + // Either a clear error or an empty output is acceptable; a panic is not. + match res { + Err(e) => { + let msg = e.to_string(); + assert!(!msg.is_empty(), "error must carry a message"); + } + Ok(_) => { + let out = read_output(&ctx, &output_uri).await?; + assert_eq!(out.clone().count().await?, 0, "output must be empty"); + } + } + Ok(()) + } + + /// Same seed -> bitwise identical output: the whole pipeline (init draws, + /// KMeans seeding, order of joins) is deterministic. + #[tokio::test] + async fn test_pic_deterministic_given_seed() -> Result<()> { + async fn run( + graph: GraphFrame, + ctx: &SessionContext, + checkpoint_dir: Path, + output_uri: &str, + ) -> Result { + PICBuilder::new(graph) + .set_k(2) + .set_checkpoint_dir(checkpoint_dir) + .run(ctx, output_uri) + .await?; + read_output(ctx, output_uri).await + } + + let (ctx_a, ck_a, out_a, _g) = setup("deterministic_a")?; + let a = run(two_cliques()?, &ctx_a, ck_a, &out_a).await?; + + let (ctx_b, ck_b, out_b, _g) = setup("deterministic_b")?; + let b = run(two_cliques()?, &ctx_b, ck_b, &out_b).await?; + let va = a + .clone() + .select(vec![cast(col("center_2"), DataType::Int64)])? + .cache() + .await?; + let vb = b + .select(vec![cast(col("center_2"), DataType::Int64)])? + .cache() + .await?; + let la = collect_to_i64(&va, 0).await?; + let lb = collect_to_i64(&vb, 0).await?; + assert_eq!(la, lb, "same seed must produce identical assignments"); + Ok(()) + } + + /// `set_seed` must exist and flow into both init and KMeans: two different + /// seeds may differ, the same seed must not (already covered above); here + /// we only pin the API. + #[test] + fn test_pic_builder_setters_exist() { + let graph = two_cliques().expect("graph"); + let b = PICBuilder::new(graph) + .set_k(3) + .set_multiple_k(vec![2, 4]) + .set_init_strategy(InitStrategy::Random) + .set_weights_strategy(WeightsStrategy::None) + .set_max_iterations(10); + assert!(format!("{b:?}").len() > 0); + } +} diff --git a/src/algorithm/subgraph/maximal_independent_set.rs b/src/algorithm/subgraph/maximal_independent_set.rs index e8edabf..f21330e 100644 --- a/src/algorithm/subgraph/maximal_independent_set.rs +++ b/src/algorithm/subgraph/maximal_independent_set.rs @@ -179,6 +179,7 @@ impl<'a> MISBuilder<'a> { .clone() .select_columns(&[EDGE_SRC, EDGE_DST])?, true, + None, )?, ) .await?; diff --git a/src/expressions.rs b/src/expressions.rs index 78a35c2..90344d6 100644 --- a/src/expressions.rs +++ b/src/expressions.rs @@ -2,9 +2,17 @@ mod common; mod finite_axpb; mod hll; mod kcore_merge; +mod kmeans_assign; +mod kmeans_step; +mod linalg; mod most_common; +pub(crate) use common::as_f32_list_like; 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 use kmeans_assign::kmeans_assign_expr; +pub(crate) use kmeans_assign::kmeans_cost_expr; +pub(crate) use kmeans_step::kmeans_step_expr; +pub(crate) use linalg::{cosine_distance_expr, l2_distance_expr, l2_norm_expr}; pub(crate) use most_common::most_common_expr; diff --git a/src/expressions/common.rs b/src/expressions/common.rs index fead671..5b87cae 100644 --- a/src/expressions/common.rs +++ b/src/expressions/common.rs @@ -1,6 +1,8 @@ use datafusion::arrow::array::{ - Array, ArrayRef, BinaryArray, BinaryViewArray, Int32Array, Int64Array, + Array, ArrayRef, BinaryArray, BinaryViewArray, FixedSizeListArray, Float32Array, Float64Array, + Int32Array, Int64Array, ListArray, }; +use datafusion::arrow::datatypes::DataType; use datafusion::error::{DataFusionError, Result}; /// Helper for other UDFs: vertexId and edgeSrc / edgeDst are all i64 @@ -72,6 +74,122 @@ impl<'a> BinaryLike<'a> { } } +/// Read-only accessor over either a [`FixedSizeListArray`] or a [`ListArray`] +/// whose elements are `Float32`. +/// +/// Fixed-width float vectors are typically materialized as `FixedSizeList`, +/// while a parquet round-trip surfaces them as `List`, so every +/// vector-consuming UDF must accept both representations. +pub(crate) enum F32ListLike<'a> { + Fixed(&'a FixedSizeListArray), + View(&'a ListArray), +} + +impl<'a> F32ListLike<'a> { + pub(crate) fn len(&self) -> usize { + match self { + F32ListLike::Fixed(a) => a.len(), + F32ListLike::View(a) => a.len(), + } + } + + pub(crate) fn null_count(&self) -> usize { + match self { + F32ListLike::Fixed(a) => a.null_count(), + F32ListLike::View(a) => a.null_count(), + } + } + + pub(crate) fn is_null(&self, i: usize) -> bool { + match self { + F32ListLike::Fixed(a) => a.is_null(i), + F32ListLike::View(a) => a.is_null(i), + } + } + + /// Zero-copy `f32` slice for row `i`. + pub(crate) fn value(&self, i: usize) -> &'a [f32] { + match self { + F32ListLike::Fixed(a) => { + let child = a + .values() + .as_any() + .downcast_ref::() + .expect("F32ListLike child validated as Float32"); + let size = a.value_length() as usize; + &child.values()[i * size..(i + 1) * size] + } + F32ListLike::View(a) => { + let child = a + .values() + .as_any() + .downcast_ref::() + .expect("F32ListLike child validated as Float32"); + let offsets = a.value_offsets(); + &child.values()[offsets[i] as usize..offsets[i + 1] as usize] + } + } + } +} + +/// Read-only accessor over either a [`FixedSizeListArray`] or a [`ListArray`] +/// whose elements are `Float64`. +/// +/// Fixed-width float vectors are typically materialized as `FixedSizeList`, +/// while a parquet round-trip surfaces them as `List`, so every +/// vector-consuming UDF must accept both representations. +pub(crate) enum F64ListLike<'a> { + Fixed(&'a FixedSizeListArray), + View(&'a ListArray), +} + +impl<'a> F64ListLike<'a> { + pub(crate) fn len(&self) -> usize { + match self { + F64ListLike::Fixed(a) => a.len(), + F64ListLike::View(a) => a.len(), + } + } + + pub(crate) fn null_count(&self) -> usize { + match self { + F64ListLike::Fixed(a) => a.null_count(), + F64ListLike::View(a) => a.null_count(), + } + } + + pub(crate) fn is_null(&self, i: usize) -> bool { + match self { + F64ListLike::Fixed(a) => a.is_null(i), + F64ListLike::View(a) => a.is_null(i), + } + } + + /// Zero-copy `f64` slice for row `i`. + pub(crate) fn value(&self, i: usize) -> &'a [f64] { + match self { + F64ListLike::Fixed(a) => { + let child = a + .values() + .as_any() + .downcast_ref::() + .expect("F64ListLike child validated as Float64"); + let size = a.value_length() as usize; + &child.values()[i * size..(i + 1) * size] + } + F64ListLike::View(a) => { + let child = a + .values() + .as_any() + .downcast_ref::() + .expect("F64ListLike child validated as Float64"); + let offsets = a.value_offsets(); + &child.values()[offsets[i] as usize..offsets[i + 1] as usize] + } + } + } +} + /// Helper for binary vectors pub(crate) fn as_binary_like<'a>( array: &'a ArrayRef, @@ -89,3 +207,47 @@ pub(crate) fn as_binary_like<'a>( array.data_type() ))) } + +/// Helper for f32 vectors +pub(crate) fn as_f32_list_like<'a>( + array: &'a ArrayRef, + fname: &str, + label: &str, +) -> Result> { + if let Some(a) = array.as_any().downcast_ref::() { + if a.values().data_type() == &DataType::Float32 { + return Ok(F32ListLike::Fixed(a)); + } + } + if let Some(a) = array.as_any().downcast_ref::() { + if a.values().data_type() == &DataType::Float32 { + return Ok(F32ListLike::View(a)); + } + } + Err(DataFusionError::Plan(format!( + "{fname} {label} argument must be FixedSizeList(Float32) or List(Float32), got: {:?}", + array.data_type() + ))) +} + +/// Helper for f64 vectors +pub(crate) fn as_f64_list_like<'a>( + array: &'a ArrayRef, + fname: &str, + label: &str, +) -> Result> { + if let Some(a) = array.as_any().downcast_ref::() { + if a.values().data_type() == &DataType::Float64 { + return Ok(F64ListLike::Fixed(a)); + } + } + if let Some(a) = array.as_any().downcast_ref::() { + if a.values().data_type() == &DataType::Float64 { + return Ok(F64ListLike::View(a)); + } + } + Err(DataFusionError::Plan(format!( + "{fname} {label} argument must be FixedSizeList(Float64) or List(Float64), got: {:?}", + array.data_type() + ))) +} diff --git a/src/expressions/kmeans_assign.rs b/src/expressions/kmeans_assign.rs new file mode 100644 index 0000000..b3acd31 --- /dev/null +++ b/src/expressions/kmeans_assign.rs @@ -0,0 +1,306 @@ +//! Final cluster assignment (`k_means_assign`) and nearest-center cost +//! (`k_means_cost`) scalar UDFs for K-Means. + +use datafusion::arrow::array::{ArrayRef, Float64Array, Int32Array}; +use datafusion::arrow::datatypes::DataType; +use datafusion::error::Result; +use datafusion::logical_expr::{ + ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, +}; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use crate::expressions::common::as_f32_list_like; +use crate::ml::{DistanceMetric, nearest_center, nearest_centers}; + +/// Shared construction of the 1-argument signature: +/// `features` is either `FixedSizeList(d)` or `List`. +fn features_signature(d: usize) -> Signature { + let field = || { + Arc::new(datafusion::arrow::datatypes::Field::new( + "el", + DataType::Float32, + false, + )) + }; + Signature::uniform( + 1, + vec![ + DataType::FixedSizeList(field(), d as i32), + DataType::List(field()), + ], + Volatility::Immutable, + ) +} + +/// `Eq`/`Hash` consistent over the float centers: comparison and hashing are +/// both bit-based (`f32::to_bits`), so equal instances always hash equal. +/// `signature` is a pure function of `d` and therefore not compared. +macro_rules! impl_kmeans_udf_eq_hash { + ($t:ty) => { + impl PartialEq for $t { + fn eq(&self, other: &Self) -> bool { + self.k == other.k + && self.d == other.d + && self.metric == other.metric + && self.centers.len() == other.centers.len() + && self + .centers + .iter() + .zip(&other.centers) + .all(|(a, b)| a.to_bits() == b.to_bits()) + } + } + + impl Eq for $t {} + + impl Hash for $t { + fn hash(&self, state: &mut H) { + self.k.hash(state); + self.d.hash(state); + self.metric.hash(state); + self.centers.len().hash(state); + for c in &self.centers { + c.to_bits().hash(state); + } + } + } + }; +} + +/// Scalar UDF `k_means_assign(features) -> Int32`: index of the nearest +/// center per row (ties break to the first center). +#[derive(Debug)] +pub(crate) struct KMeansAssign { + signature: Signature, + k: usize, + d: usize, + centers: Vec, + metric: DistanceMetric, +} + +impl KMeansAssign { + pub(crate) fn new(k: usize, d: usize, centers: Vec, metric: DistanceMetric) -> Self { + Self { + signature: features_signature(d), + k, + d, + centers, + metric, + } + } +} + +impl_kmeans_udf_eq_hash!(KMeansAssign); + +impl ScalarUDFImpl for KMeansAssign { + fn name(&self) -> &str { + "k_means_assign" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Int32) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + let v = as_f32_list_like(&arrays[0], "k_means_assign", "first")?; + let result: Int32Array = (0..v.len()) + .map(|i| { + Some(nearest_center(v.value(i), &self.centers, self.k, self.d, self.metric) as i32) + }) + .collect(); + Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef)) + } +} + +/// Scalar UDF `k_means_cost(features) -> Float64`: distance to the nearest +/// center per row (squared L2 for the L2 metric — the D^2 cost used by the +/// k-means|| sampling; cosine distance for the Cosine metric). +#[derive(Debug)] +pub(crate) struct KMeansCost { + signature: Signature, + k: usize, + d: usize, + centers: Vec, + metric: DistanceMetric, +} + +impl KMeansCost { + pub(crate) fn new(k: usize, d: usize, centers: Vec, metric: DistanceMetric) -> Self { + Self { + signature: features_signature(d), + k, + d, + centers, + metric, + } + } +} + +impl_kmeans_udf_eq_hash!(KMeansCost); + +impl ScalarUDFImpl for KMeansCost { + fn name(&self) -> &str { + "k_means_cost" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Float64) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + let v = as_f32_list_like(&arrays[0], "k_means_cost", "first")?; + let result: Float64Array = (0..v.len()) + .map(|i| { + let (_, dist) = + nearest_centers(v.value(i), &self.centers, self.k, self.d, self.metric); + Some(dist as f64) + }) + .collect(); + Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef)) + } +} + +/// Builds an [`Expr`] assigning each row to its nearest center. +pub fn kmeans_assign_expr( + features: Expr, + k: usize, + d: usize, + centers: Vec, + metric: DistanceMetric, +) -> Expr { + ScalarUDF::from(KMeansAssign::new(k, d, centers, metric)).call(vec![features]) +} + +/// Builds an [`Expr`] computing the distance from each row to its nearest +/// center. +pub(crate) fn kmeans_cost_expr( + features: Expr, + k: usize, + d: usize, + centers: Vec, + metric: DistanceMetric, +) -> Expr { + ScalarUDF::from(KMeansCost::new(k, d, centers, metric)).call(vec![features]) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::{Array, FixedSizeListArray, Float32Array, RecordBatch}; + use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::common::Result; + use datafusion::prelude::*; + + fn centers_2x2() -> Vec { + vec![10.0, 10.0, 0.0, 0.0] + } + + fn features_table(rows: &[[f32; 2]]) -> Result { + let flat: Vec = rows.iter().flatten().copied().collect(); + let schema = Schema::new(vec![Field::new( + "feat", + DataType::FixedSizeList(Arc::new(Field::new("el", DataType::Float32, false)), 2), + false, + )]); + let fsl = FixedSizeListArray::try_new( + Arc::new(Field::new("el", DataType::Float32, false)), + 2, + Arc::new(Float32Array::from(flat)), + None, + )?; + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(fsl) as ArrayRef])?; + let ctx = SessionContext::new(); + Ok(ctx.read_batch(batch)?) + } + + #[tokio::test] + async fn assign_picks_nearest_center_per_row() -> Result<()> { + let df = features_table(&[[11.0, 11.0], [1.0, 1.0], [4.0, 4.0]])?; + let out = df + .clone() + .select(vec![ + kmeans_assign_expr(col("feat"), 2, 2, centers_2x2(), DistanceMetric::L2) + .alias("cluster"), + ])? + .collect() + .await?; + let clusters = out[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(clusters.len(), 3); + assert_eq!(clusters.value(0), 0); + assert_eq!(clusters.value(1), 1); + // (4,4): squared L2 to (10,10) = 72, to (0,0) = 32 -> cluster 1. + assert_eq!(clusters.value(2), 1); + Ok(()) + } + + #[tokio::test] + async fn cost_returns_squared_l2_to_nearest_center() -> Result<()> { + let df = features_table(&[[11.0, 11.0], [1.0, 1.0]])?; + let out = df + .clone() + .select(vec![ + kmeans_cost_expr(col("feat"), 2, 2, centers_2x2(), DistanceMetric::L2) + .alias("cost"), + ])? + .collect() + .await?; + let costs = out[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(costs.len(), 2); + assert!((costs.value(0) - 2.0).abs() < 1e-6, "{}", costs.value(0)); + assert!((costs.value(1) - 2.0).abs() < 1e-6, "{}", costs.value(1)); + Ok(()) + } + + #[test] + fn builders_reference_the_udf_names() { + let a = kmeans_assign_expr(col("feat"), 2, 2, centers_2x2(), DistanceMetric::L2); + assert!(format!("{a}").contains("k_means_assign")); + let c = kmeans_cost_expr(col("feat"), 2, 2, centers_2x2(), DistanceMetric::L2); + assert!(format!("{c}").contains("k_means_cost")); + } + + #[test] + fn eq_and_hash_agree_on_signed_zero() { + let a = KMeansAssign::new(1, 2, vec![0.0f32, 0.0], DistanceMetric::L2); + let b = KMeansAssign::new(1, 2, vec![-0.0f32, 0.0], DistanceMetric::L2); + if a == b { + let ha = { + let mut h = std::collections::hash_map::DefaultHasher::new(); + a.hash(&mut h); + h.finish() + }; + let hb = { + let mut h = std::collections::hash_map::DefaultHasher::new(); + b.hash(&mut h); + h.finish() + }; + assert_eq!(ha, hb, "equal objects must hash equal"); + } + } + + #[test] + fn eq_respects_centers_length() { + let a = KMeansCost::new(1, 2, vec![1.0f32], DistanceMetric::L2); + let b = KMeansCost::new(1, 2, vec![1.0f32, 2.0], DistanceMetric::L2); + assert_ne!(a, b); + } +} diff --git a/src/expressions/kmeans_step.rs b/src/expressions/kmeans_step.rs new file mode 100644 index 0000000..b92a894 --- /dev/null +++ b/src/expressions/kmeans_step.rs @@ -0,0 +1,420 @@ +use datafusion::error::Result; +use datafusion::logical_expr::function::AccumulatorArgs; +use datafusion::logical_expr::{ + Accumulator, AggregateUDF, AggregateUDFImpl, Expr, Signature, Volatility, +}; + +use std::hash::{Hash, Hasher}; +use std::sync::Arc; + +use datafusion::arrow::array::{ArrayRef, FixedSizeListArray, Float64Array}; +use datafusion::arrow::datatypes::{DataType, Field}; +use datafusion::scalar::ScalarValue; + +use crate::expressions::common::{as_f32_list_like, as_f64_list_like}; +use crate::ml::{DistanceMetric, nearest_centers}; + +#[derive(Debug)] +pub(crate) struct KMeansStepAccumulator { + k: usize, + d: usize, + centers: Vec, + state: Vec, + metric: DistanceMetric, +} + +impl KMeansStepAccumulator { + pub(crate) fn new(k: usize, d: usize, centers: Vec, metric: DistanceMetric) -> Self { + Self { + k: k, + d: d, + centers: centers, + state: vec![0.0; k * d + k], + metric: metric, + } + } + + pub(crate) fn default(k: usize, d: usize, centers: Vec) -> Self { + Self { + k: k, + d: d, + centers: centers, + state: vec![0.0; k * d + k], + metric: DistanceMetric::L2, + } + } + + fn to_scalar(&self) -> ScalarValue { + ScalarValue::FixedSizeList(Arc::new( + FixedSizeListArray::try_new( + Arc::new(Field::new("el", DataType::Float64, false)), + (self.k * self.d + self.k) as i32, + Arc::new(Float64Array::from(self.state.clone())), + None, + ) + .expect("valid fixed size list"), + )) + } +} + +impl Accumulator for KMeansStepAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + debug_assert_eq!(values[0].null_count(), 0usize); + let v = as_f32_list_like(&values[0], "k_means_step", "first")?; + + // no nulls are assumed in feature (embeddings) + for i in 0..v.len() { + let vv = v.value(i); + let (cluster, _) = nearest_centers(vv, &self.centers, self.k, self.d, self.metric); + self.state[self.k * self.d + cluster] += 1.0f64; + + for t in 0..self.d { + self.state[cluster * self.d + t] += vv[t] as f64; + } + } + + Ok(()) + } + + fn state(&mut self) -> Result> { + Ok(vec![self.to_scalar()]) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + let s = as_f64_list_like(&states[0], "k_means_step", "state")?; + + let n = self.k * self.d + self.k; + // non-nulls semantic by contract + for i in 0..s.len() { + let ss = s.value(i); + for t in 0..n { + self.state[t] += ss[t]; + } + } + + Ok(()) + } + + fn size(&self) -> usize { + // 2 usize values + 1 enum + 2x vec struct + k*d of f32 + (k*d + k) of f64 + size_of::() * 2 + + size_of::() + + (self.k * self.d) * size_of::() + + (self.k * self.d + self.k) * size_of::() + + size_of::>() + + size_of::>() + } + + fn evaluate(&mut self) -> Result { + Ok(self.to_scalar()) + } +} + +#[derive(Debug)] +pub(crate) struct KMeansStep { + signature: Signature, + k: usize, + d: usize, + centers: Vec, + metric: DistanceMetric, +} + +impl KMeansStep { + pub(crate) fn new(k: usize, d: usize, centers: Vec, metric: DistanceMetric) -> Self { + Self { + signature: Signature::uniform( + 1, + vec![ + DataType::FixedSizeList( + Arc::new(Field::new("el", DataType::Float32, false)), + d as i32, + ), + DataType::List(Arc::new(Field::new("el", DataType::Float32, false))), + ], + Volatility::Immutable, + ), + k: k, + d: d, + centers: centers, + metric: metric, + } + } + + pub(crate) fn default(k: usize, d: usize, centers: Vec) -> Self { + Self::new(k, d, centers, DistanceMetric::L2) + } +} + +impl PartialEq for KMeansStep { + fn eq(&self, other: &Self) -> bool { + // note: signature cannot differ because it is not exposed anyhow + // to the caller + self.k == other.k + && self.d == other.d + && self.metric == other.metric + && self.centers.len() == other.centers.len() + && self + .centers + .iter() + .zip(&other.centers) + .all(|(a, b)| a.to_bits() == b.to_bits()) + } +} + +impl Eq for KMeansStep {} + +impl Hash for KMeansStep { + fn hash(&self, state: &mut H) { + self.k.hash(state); + self.d.hash(state); + self.metric.hash(state); + self.centers.len().hash(state); + for c in &self.centers { + c.to_bits().hash(state); + } + } +} + +impl AggregateUDFImpl for KMeansStep { + fn name(&self) -> &str { + "k_means_step" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::FixedSizeList( + Arc::new(Field::new("el", DataType::Float64, false)), + (self.k * self.d + self.k) as i32, + )) + } + + fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result> { + Ok(Box::new(KMeansStepAccumulator::new( + self.k, + self.d, + self.centers.clone(), + self.metric, + ))) + } +} + +/// Builds an [`Expr`] applying `k_means_step` to `features`. +/// +/// `k`, `d`, `centers` and `metric` are captured in the UDF instance at +/// construction time. The caller must build a fresh expression +/// per iteration. +pub(crate) fn kmeans_step_expr( + features: Expr, + k: usize, + d: usize, + centers: Vec, + metric: DistanceMetric, +) -> Expr { + AggregateUDF::from(KMeansStep::new(k, d, centers, metric)).call(vec![features]) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::{ + Array, ArrayRef, FixedSizeListArray, Float32Array, Float64Array, RecordBatch, + }; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::Result; + use datafusion::datasource::MemTable; + use datafusion::prelude::*; + use datafusion::scalar::ScalarValue; + use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; + use std::sync::Arc; + + fn centers_2x2() -> Vec { + vec![10.0, 10.0, 0.0, 0.0] + } + + /// `FixedSizeList` column (d=2) from flat rows. + fn feat_batch(rows: &[[f32; 2]]) -> ArrayRef { + let flat: Vec = rows.iter().flatten().copied().collect(); + Arc::new( + FixedSizeListArray::try_new( + Arc::new(Field::new("el", DataType::Float32, false)), + 2, + Arc::new(Float32Array::from(flat)), + None, + ) + .unwrap(), + ) + } + + /// Reads `[k*d sums, k counts]` out of the accumulator state regardless of + /// the current state representation (List today, FixedSizeList per the + /// declared return type once the representation is aligned). + fn state_vec(acc: &mut KMeansStepAccumulator) -> Vec { + let st = acc.state().unwrap(); + match &st[0] { + ScalarValue::FixedSizeList(arr) => { + let binding = arr.value(0); + let v = binding.as_any().downcast_ref::().unwrap(); + (0..v.len()).map(|i| v.value(i)).collect() + } + ScalarValue::List(la) => { + let binding = la.value(0); + let v = binding.as_any().downcast_ref::().unwrap(); + (0..v.len()).map(|i| v.value(i)).collect() + } + other => panic!("unexpected state: {other:?}"), + } + } + + /// State as an [`ArrayRef`] suitable for `merge_batch`. + fn state_array(acc: &mut KMeansStepAccumulator) -> ArrayRef { + let st = acc.state().unwrap(); + match &st[0] { + ScalarValue::FixedSizeList(arr) => arr.clone() as ArrayRef, + ScalarValue::List(la) => la.clone() as ArrayRef, + other => panic!("unexpected state: {other:?}"), + } + } + + #[test] + fn sums_land_in_assigned_center_block() { + // rows (11,11) -> c0, (1,1) -> c1 against c0=(10,10), c1=(0,0) + let mut acc = KMeansStepAccumulator::new(2, 2, centers_2x2(), DistanceMetric::L2); + acc.update_batch(&vec![feat_batch(&[[11.0, 11.0], [1.0, 1.0]])]) + .unwrap(); + assert_eq!( + state_vec(&mut acc), + vec![11.0, 11.0, 1.0, 1.0, 1.0, 1.0], + "per-cluster sums + counts" + ); + } + + #[test] + fn merge_batch_adds_partials() { + let mut a = KMeansStepAccumulator::new(2, 2, centers_2x2(), DistanceMetric::L2); + a.update_batch(&vec![feat_batch(&[[11.0, 11.0]])]).unwrap(); + let mut b = KMeansStepAccumulator::new(2, 2, centers_2x2(), DistanceMetric::L2); + b.update_batch(&vec![feat_batch(&[[1.0, 1.0]])]).unwrap(); + a.merge_batch(&[state_array(&mut b)]).unwrap(); + assert_eq!( + state_vec(&mut a), + vec![11.0, 11.0, 1.0, 1.0, 1.0, 1.0], + "merged partials equal the single-batch result" + ); + } + + #[test] + fn state_shape_matches_declared_return_type() { + // return_type declares FixedSizeList(Float64, k*d+k); state must match. + let mut acc = KMeansStepAccumulator::new(2, 2, centers_2x2(), DistanceMetric::L2); + acc.update_batch(&vec![feat_batch(&[[1.0, 0.0]])]).unwrap(); + let st = acc.state().unwrap(); + match &st[0] { + ScalarValue::FixedSizeList(arr) => { + assert_eq!(arr.len(), 1); + assert_eq!(arr.value_length() as usize, 6); + } + other => panic!("state must be ScalarValue::FixedSizeList, got {other:?}"), + } + } + + fn hash_of(x: &KMeansStep) -> u64 { + let mut h = DefaultHasher::new(); + x.hash(&mut h); + h.finish() + } + + #[test] + fn eq_respects_centers_length() { + let a = KMeansStep::new(1, 2, vec![1.0f32], DistanceMetric::L2); + let b = KMeansStep::new(1, 2, vec![1.0f32, 2.0], DistanceMetric::L2); + assert_ne!(a, b, "prefix-equal centers of different length must differ"); + } + + #[test] + fn eq_and_hash_agree_on_signed_zero() { + // Numeric equality says +0.0 == -0.0; the bit-level hash differs. + // Whatever side equality lands on, the Eq/Hash contract must hold. + let a = KMeansStep::new(1, 2, vec![0.0f32, 0.0], DistanceMetric::L2); + let b = KMeansStep::new(1, 2, vec![-0.0f32, 0.0], DistanceMetric::L2); + if a == b { + assert_eq!(hash_of(&a), hash_of(&b), "equal objects must hash equal"); + } + } + + #[test] + fn builder_produces_kmeans_step_expr() { + let expr = kmeans_step_expr(col("feat"), 2, 2, centers_2x2(), DistanceMetric::L2); + let s = format!("{expr}"); + assert!(s.contains("k_means_step"), "expr display: {s}"); + } + + async fn register_features( + ctx: &SessionContext, + rows: &[[f32; 2]], + partitions: usize, + ) -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "feat", + DataType::FixedSizeList(Arc::new(Field::new("el", DataType::Float32, false)), 2), + false, + )])); + let chunks = rows.chunks(rows.len().div_ceil(partitions)); + let batches: Vec> = chunks + .map(|chunk| { + let flat: Vec = chunk.iter().flatten().copied().collect(); + let fsl = FixedSizeListArray::try_new( + Arc::new(Field::new("el", DataType::Float32, false)), + 2, + Arc::new(Float32Array::from(flat)), + None, + ) + .unwrap(); + vec![RecordBatch::try_new(schema.clone(), vec![Arc::new(fsl) as ArrayRef]).unwrap()] + }) + .collect(); + let table = MemTable::try_new(schema, batches)?; + ctx.register_table("t", Arc::new(table))?; + Ok(()) + } + + async fn run_step(ctx: &SessionContext, centers: Vec) -> Result> { + let udf = AggregateUDF::from(KMeansStep::new(2, 2, centers, DistanceMetric::L2)); + let out = ctx + .table("t") + .await? + .aggregate(vec![], vec![udf.call(vec![col("feat")]).alias("r")])? + .collect() + .await?; + assert_eq!(out[0].num_rows(), 1); + let arr = out[0].column(0); + let fsl = arr + .as_any() + .downcast_ref::() + .expect("result must be a FixedSizeList"); + let binding = fsl.value(0); + let v = binding.as_any().downcast_ref::().unwrap(); + Ok((0..v.len()).map(|i| v.value(i)).collect()) + } + + #[tokio::test] + async fn e2e_aggregate_single_partition() -> Result<()> { + let ctx = SessionContext::new(); + register_features(&ctx, &[[11.0, 11.0], [1.0, 1.0]], 1).await?; + let out = run_step(&ctx, centers_2x2()).await?; + assert_eq!(out, vec![11.0, 11.0, 1.0, 1.0, 1.0, 1.0]); + Ok(()) + } + + #[tokio::test] + async fn e2e_aggregate_two_partitions() -> Result<()> { + let ctx = SessionContext::new_with_config(SessionConfig::new().with_target_partitions(2)); + register_features(&ctx, &[[11.0, 11.0], [1.0, 1.0]], 2).await?; + let out = run_step(&ctx, centers_2x2()).await?; + assert_eq!(out, vec![11.0, 11.0, 1.0, 1.0, 1.0, 1.0]); + Ok(()) + } +} diff --git a/src/expressions/linalg.rs b/src/expressions/linalg.rs new file mode 100644 index 0000000..8eaabde --- /dev/null +++ b/src/expressions/linalg.rs @@ -0,0 +1,392 @@ +//! Scalar UDFs over `f32` vectors: `l2_norm`, `l2_distance`, +//! `cosine_distance`. +//! +//! This module holds the DataFusion wrappers only: SIMD-logic free; the +//! kernels live in [`crate::ml::linalg`]. Contract: +//! vector rows are non-null; the two-argument UDFs require both rows +//! to have the same length. + +use std::sync::Arc; + +use datafusion::arrow::array::{ArrayRef, Float64Array}; +use datafusion::arrow::datatypes::DataType; +use datafusion::common::plan_err; +use datafusion::error::Result; +use datafusion::logical_expr::{ + ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, +}; + +use crate::expressions::common::as_f32_list_like; +use crate::ml::{cosine_distance, l2_distance, l2_norm}; + +/// Both arguments must be same-sized `f32` vectors. +fn validate_vector_args(arg_types: &[DataType], arity: usize, fname: &str) -> Result<()> { + if arg_types.len() != arity { + return plan_err!( + "{fname} expects {arity} argument(s), got {}", + arg_types.len() + ); + } + for (i, t) in arg_types.iter().enumerate() { + let ok = match t { + DataType::FixedSizeList(f, _) => f.data_type() == &DataType::Float32, + DataType::List(f) => f.data_type() == &DataType::Float32, + _ => false, + }; + if !ok { + return plan_err!( + "{fname} argument {i} must be FixedSizeList or List, got {t:?}" + ); + } + } + Ok(()) +} + +/// Scalar UDF `l2_norm(v) -> Float64`: true L2 norm of a vector. +#[derive(Debug, PartialEq, Eq, Hash)] +pub(crate) struct L2NormUDF { + signature: Signature, +} + +impl L2NormUDF { + pub(crate) fn new() -> Self { + Self { + signature: Signature::any(1, Volatility::Immutable), + } + } +} + +impl Default for L2NormUDF { + fn default() -> Self { + Self::new() + } +} + +impl ScalarUDFImpl for L2NormUDF { + fn name(&self) -> &str { + "l2_norm" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + validate_vector_args(arg_types, 1, "l2_norm")?; + Ok(DataType::Float64) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + let v = as_f32_list_like(&arrays[0], "l2_norm", "first")?; + let result: Float64Array = (0..v.len()) + .map(|i| { + let row = v.value(i); + Some(l2_norm(row, row.len()) as f64) + }) + .collect(); + Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef)) + } +} + +/// Scalar UDF `l2_distance(v1, v2) -> Float64`: true L2 distance between two +/// vectors. +#[derive(Debug, PartialEq, Eq, Hash)] +pub(crate) struct L2DistanceUDF { + signature: Signature, +} + +impl L2DistanceUDF { + pub(crate) fn new() -> Self { + Self { + signature: Signature::any(2, Volatility::Immutable), + } + } +} + +impl Default for L2DistanceUDF { + fn default() -> Self { + Self::new() + } +} + +impl ScalarUDFImpl for L2DistanceUDF { + fn name(&self) -> &str { + "l2_distance" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + validate_vector_args(arg_types, 2, "l2_distance")?; + Ok(DataType::Float64) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + let v1 = as_f32_list_like(&arrays[0], "l2_distance", "first")?; + let v2 = as_f32_list_like(&arrays[1], "l2_distance", "second")?; + let len = v1.len().max(v2.len()); + let mut values = Vec::with_capacity(len); + for i in 0..len { + let a = v1.value(i % v1.len()); + let b = v2.value(i % v2.len()); + if a.len() != b.len() { + return plan_err!( + "l2_distance vectors must have the same length, got {} and {}", + a.len(), + b.len() + ); + } + values.push(Some(l2_distance(a, b, a.len()).sqrt() as f64)); + } + let result: Float64Array = values.into_iter().collect(); + Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef)) + } +} + +/// Scalar UDF `cosine_distance(v1, v2) -> Float64`: cosine distance between +/// two vectors (`1 - cosine similarity`; zero-norm rows score `0.0`); +/// +/// scikit-learn semantics +#[derive(Debug, PartialEq, Eq, Hash)] +pub(crate) struct CosineDistanceUDF { + signature: Signature, +} + +impl CosineDistanceUDF { + pub(crate) fn new() -> Self { + Self { + signature: Signature::any(2, Volatility::Immutable), + } + } +} + +impl Default for CosineDistanceUDF { + fn default() -> Self { + Self::new() + } +} + +impl ScalarUDFImpl for CosineDistanceUDF { + fn name(&self) -> &str { + "cosine_distance" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + validate_vector_args(arg_types, 2, "cosine_distance")?; + Ok(DataType::Float64) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + let v1 = as_f32_list_like(&arrays[0], "cosine_distance", "first")?; + let v2 = as_f32_list_like(&arrays[1], "cosine_distance", "second")?; + let len = v1.len().max(v2.len()); + let mut values = Vec::with_capacity(len); + for i in 0..len { + let a = v1.value(i % v1.len()); + let b = v2.value(i % v2.len()); + if a.len() != b.len() { + return plan_err!( + "cosine_distance vectors must have the same length, got {} and {}", + a.len(), + b.len() + ); + } + values.push(Some(cosine_distance( + a, + b, + a.len(), + l2_norm(a, a.len()), + l2_norm(b, b.len()), + ) as f64)); + } + let result: Float64Array = values.into_iter().collect(); + Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef)) + } +} + +/// Builds an [`Expr`] computing the L2 norm of `v`. +pub(crate) fn l2_norm_expr(v: Expr) -> Expr { + ScalarUDF::from(L2NormUDF::new()).call(vec![v]) +} + +/// Builds an [`Expr`] computing the true L2 distance between `v1` and `v2`. +pub(crate) fn l2_distance_expr(v1: Expr, v2: Expr) -> Expr { + ScalarUDF::from(L2DistanceUDF::new()).call(vec![v1, v2]) +} + +/// Builds an [`Expr`] computing the cosine distance between `v1` and `v2`. +pub(crate) fn cosine_distance_expr(v1: Expr, v2: Expr) -> Expr { + ScalarUDF::from(CosineDistanceUDF::new()).call(vec![v1, v2]) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::{Array, FixedSizeListArray, Float32Array, RecordBatch}; + use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::common::Result; + use datafusion::prelude::{DataFrame, SessionContext, col, lit}; + + // ---------------- UDFs ---------------- + + fn two_col_table(rows: &[[f32; 2]], rows2: &[[f32; 2]]) -> Result { + let mk = |rows: &[[f32; 2]], name: &str| -> Result<(Arc, ArrayRef)> { + let flat: Vec = rows.iter().flatten().copied().collect(); + let fsl = FixedSizeListArray::try_new( + Arc::new(Field::new("el", DataType::Float32, false)), + 2, + Arc::new(Float32Array::from(flat)), + None, + )?; + let schema = Schema::new(vec![Field::new( + name, + DataType::FixedSizeList(Arc::new(Field::new("el", DataType::Float32, false)), 2), + false, + )]); + Ok((Arc::new(schema), Arc::new(fsl) as ArrayRef)) + }; + let (s1, a1) = mk(rows, "v1")?; + let (s2, a2) = mk(rows2, "v2")?; + let schema = Schema::new(vec![s1.field(0).clone(), s2.field(0).clone()]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![a1, a2])?; + let ctx = SessionContext::new(); + Ok(ctx.read_batch(batch)?) + } + + #[tokio::test] + async fn udf_l2_norm_returns_true_norm() -> Result<()> { + let df = two_col_table(&[[3.0, 4.0], [0.0, 0.0]], &[[0.0, 0.0], [0.0, 0.0]])?; + let out = df + .clone() + .select(vec![l2_norm_expr(col("v1")).alias("n")])? + .collect() + .await?; + let n = out[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!((n.value(0) - 5.0).abs() < 1e-5, "{}", n.value(0)); + assert_eq!(n.value(1), 0.0); + Ok(()) + } + + #[tokio::test] + async fn udf_l2_distance_returns_true_distance() -> Result<()> { + let df = two_col_table(&[[0.0, 0.0], [1.0, 1.0]], &[[3.0, 4.0], [1.0, 1.0]])?; + let out = df + .clone() + .select(vec![l2_distance_expr(col("v1"), col("v2")).alias("d")])? + .collect() + .await?; + let d = out[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!((d.value(0) - 5.0).abs() < 1e-5, "{}", d.value(0)); // 3-4-5 + assert!((d.value(1) - 0.0).abs() < 1e-6, "{}", d.value(1)); + Ok(()) + } + + #[tokio::test] + async fn udf_cosine_distance_canonical_values() -> Result<()> { + let df = two_col_table( + &[[1.0, 0.0], [1.0, 2.0], [0.0, 0.0]], + &[[1.0, 0.0], [2.0, 4.0], [3.0, 4.0]], + )?; + let out = df + .clone() + .select(vec![cosine_distance_expr(col("v1"), col("v2")).alias("d")])? + .collect() + .await?; + let d = out[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!((d.value(0) - 0.0).abs() < 1e-6, "{}", d.value(0)); // parallel + assert!((d.value(1) - 0.0).abs() < 1e-6, "{}", d.value(1)); // collinear, scaled + assert!((d.value(2) - 0.0).abs() < 1e-6, "{}", d.value(2)); // zero vector policy + Ok(()) + } + + #[tokio::test] + async fn udf_l2_distance_broadcasts_literal_vector() -> Result<()> { + // One column plus a literal vector: the literal broadcasts per row. + let df = two_col_table(&[[0.0, 0.0], [6.0, 8.0]], &[[0.0, 0.0], [0.0, 0.0]])?; + let lit_vec = + datafusion::scalar::ScalarValue::FixedSizeList(Arc::new(FixedSizeListArray::try_new( + Arc::new(Field::new("el", DataType::Float32, false)), + 2, + Arc::new(Float32Array::from(vec![0.0f32, 0.0])), + None, + )?)); + let out = df + .clone() + .select(vec![l2_distance_expr(col("v1"), lit(lit_vec)).alias("d")])? + .collect() + .await?; + let d = out[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(d.len(), 2); + assert!((d.value(0) - 0.0).abs() < 1e-6, "{}", d.value(0)); + assert!((d.value(1) - 10.0).abs() < 1e-5, "{}", d.value(1)); // 6-8-10 + Ok(()) + } + + #[tokio::test] + async fn udf_l2_distance_rejects_mismatched_lengths() -> Result<()> { + // v1: one row of length 3 (List); v2: one row of length 2 (List). + use datafusion::arrow::array::ListArray; + use datafusion::arrow::datatypes::Float32Type; + let v1 = ListArray::from_iter_primitive::(vec![Some(vec![ + Some(1.0f32), + Some(2.0), + Some(3.0), + ])]); + let v2 = ListArray::from_iter_primitive::(vec![Some(vec![ + Some(1.0f32), + Some(2.0), + ])]); + let v1_field = Field::new("v1", v1.data_type().clone(), false); + let v2_field = Field::new("v2", v2.data_type().clone(), false); + let schema = Schema::new(vec![v1_field, v2_field]); + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![Arc::new(v1) as ArrayRef, Arc::new(v2) as ArrayRef], + )?; + let ctx = SessionContext::new(); + let result = ctx + .read_batch(batch)? + .select(vec![l2_distance_expr(col("v1"), col("v2")).alias("d")])? + .collect() + .await; + assert!( + result.is_err(), + "mismatched vector lengths must surface as an error" + ); + Ok(()) + } + + #[test] + fn builders_reference_udf_names() { + assert!(format!("{}", l2_norm_expr(col("v"))).contains("l2_norm")); + assert!(format!("{}", l2_distance_expr(col("a"), col("b"))).contains("l2_distance")); + assert!( + format!("{}", cosine_distance_expr(col("a"), col("b"))).contains("cosine_distance") + ); + } +} diff --git a/src/lib.rs b/src/lib.rs index 8ed0bbd..bd69e7b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,8 +1,14 @@ mod algorithm; mod expressions; mod memory; +mod ml; mod utils; +pub use algorithm::community::power_iteration_clustering::{ + EmbeddingMode, InitStrategy, PICBuilder, WeightsStrategy, +}; +pub use expressions::kmeans_assign_expr; +pub use ml::{DistanceMetric, KMeansBuilder, KMeansResult, KMeansRun}; pub use utils::GraphFramesConfig; use datafusion::arrow::datatypes::DataType; @@ -19,6 +25,8 @@ pub const EDGE_SRC: &str = "src"; pub const EDGE_DST: &str = "dst"; /// Column names for the edge column in triplet representation. pub const EDGE_COL: &str = "edge"; +/// Column names for edge weights +pub const EDGE_WEIGHT: &str = "weight"; /// Column names for the source vertex in triplet representation. pub const SRC_VERTEX: &str = "src_vertex"; /// Column names for the destination vertex in triplet representation. diff --git a/src/main.rs b/src/main.rs index 1e7f638..c8bdc6a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -2,6 +2,7 @@ #[global_allocator] static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc; +use datafusion::dataframe::DataFrameWriteOptions; use datafusion::error::{DataFusionError, Result}; use datafusion::execution::memory_pool::FairSpillPool; use datafusion::execution::runtime_env::RuntimeEnvBuilder; @@ -9,6 +10,9 @@ use datafusion::execution::session_state::SessionStateBuilder; use datafusion::object_store::path::Path as ObjectPath; use datafusion::prelude::*; use graphframes_rs::GraphFramesConfig; +use graphframes_rs::{ + DistanceMetric, EmbeddingMode, InitStrategy, KMeansBuilder, WeightsStrategy, kmeans_assign_expr, +}; use graphframes_rs::{EDGE_DST, EDGE_SRC, GraphFrame, VERTEX_ID}; use std::path::{Path, PathBuf}; use std::sync::Arc; @@ -22,6 +26,30 @@ enum Format { Json, } +#[derive(Debug, Clone, Copy, ValueEnum)] +enum PicInit { + /// Degree vector d/Σd — the paper's recommended (and MLlib's "degree") init. + Degree, + /// Seeded pseudo-random v0, L1-normalized in expectation. + Random, +} + +#[derive(Debug, Clone, Copy, ValueEnum)] +enum PicWeights { + /// Use the weight column (or unit weights) as-is. + None, + /// Positive pointwise mutual information: max(0, ln(w·W/(d_i·d_j))). + Ppmi, +} + +#[derive(Debug, Clone, Copy, ValueEnum)] +enum PicEmbedding { + /// Paper mode (default): cluster on the final iterate only. + Last, + /// Extended mode: cluster on the full iterate history [v1..vm]. + Trajectory, +} + #[derive(Args, Debug)] struct CommonArgs { /// Path (or URI) to the vertices file or directory. @@ -191,6 +219,45 @@ enum Command { to_landmarks: bool, }, + /// Power Iteration Clustering (Lin & Cohen 2010; Spark MLlib-inspired). + Pic { + #[command(flatten)] + common: CommonArgs, + + /// Number of clusters; one cluster column per K (comma-separated). + #[arg(long, value_delimiter = ',', default_values_t = vec![2usize])] + k: Vec, + + /// Maximum power iterations. PIC is *truncated* power iteration: the + /// cluster signal lives in early iterates, so on slow-mixing graphs a + /// lower budget gives a *better* embedding (paper average: 13). + #[arg(long, default_value_t = 20)] + max_iter: usize, + + /// Convergence threshold on the (mass-normalized) acceleration. + /// Flat across graph sizes; smaller = more iterations. + #[arg(long, default_value_t = 1e-5)] + tol: f64, + + /// Initial vector for the power iteration. + #[arg(long, value_enum, default_value_t = PicInit::Degree)] + init: PicInit, + + /// Edge weight transform. + #[arg(long, value_enum, default_value_t = PicWeights::None)] + weights: PicWeights, + + /// What the embedding column holds (paper clusters the final iterate). + #[arg(long, value_enum, default_value_t = PicEmbedding::Last)] + embedding: PicEmbedding, + }, + + /// Raw (graph-free) machine-learning algorithms. + Mllib { + #[command(subcommand)] + cmd: MllibCommand, + }, + /// Classical Label Propagation (CDLP). ClassicalLp { #[command(flatten)] @@ -210,6 +277,82 @@ enum Command { }, } +#[derive(Debug, Clone, ValueEnum)] +enum KmeansArgsMetric { + L2, + Cosine, +} + +#[derive(Subcommand, Debug)] +enum MllibCommand { + /// K-Means (k-means|| init, Lloyd iterations) over a feature column. + /// + /// The vertices file must contain an Int64 `id` column and a + /// Float32 feature column of the shape `List` or + /// `FixedSizeList`; no edges are read. + Kmeans { + /// Path (or URI) to the features (vertices) file or directory. + #[arg(long)] + vertices: String, + + /// Output directory as a `file://` URI. + #[arg(long)] + output: String, + + /// Input file format for `--vertices`. + #[arg(long, value_enum, default_value_t = Format::Parquet)] + format: Format, + + /// Name of the vertex-id column in the input; renamed to `id`. + #[arg(long, default_value = "id")] + id_col_name: String, + + /// Name of the feature column: List / FixedSizeList. + #[arg(long)] + feature_col: String, + + /// Number of clusters; one cluster column per K (comma-separated). + #[arg(long, value_delimiter = ',', default_values_t = vec![2usize])] + k: Vec, + + /// Distance metric. + #[arg(long, value_enum, default_value_t = KmeansArgsMetric::L2)] + metric: KmeansArgsMetric, + + /// Maximum Lloyd iterations. + #[arg(long, default_value_t = 20)] + max_iter: usize, + + /// Convergence tolerance on the center shift. + #[arg(long, default_value_t = 1e-4)] + tol: f64, + + /// k-means|| initialization steps. + #[arg(long, default_value_t = 2)] + init_steps: usize, + + /// Seed for the (deterministic) sampling and Lloyd loop. + #[arg(long, default_value_t = 42)] + seed: u64, + + /// DataFusion spill-pool memory limit, e.g. `4G` or `512M`. + #[arg(long, env = "GRAPHFRAMES_MAX_MEMORY", default_value = "4G")] + max_memory: String, + + /// Parallelism (= DataFusion `target_partitions`). + #[arg(long, env = "GRAPHFRAMES_NUM_WORKERS", default_value_t = 2)] + num_workers: usize, + + /// Base working directory for spill files. + #[arg(long, env = "GRAPHFRAMES_WORKDIR", default_value = "gf_workdir")] + checkpoint_dir: String, + + /// Upper bound on the total size of DataFusion's temporary spill dir. + #[arg(long, env = "GRAPHFRAMES_MAX_TEMP_FILE", default_value = "200G")] + max_temp_file: String, + }, +} + #[derive(Parser, Debug)] #[command( name = "graphframes", @@ -458,6 +601,126 @@ async fn main() -> Result<()> { .run(&ctx, &common.output, false) .await?; } + Command::Pic { + common, + k, + max_iter, + tol, + init, + weights, + embedding, + } => { + let (ctx, g, ckpt) = setup(&common).await?; + let mut b = g + .pic() + .set_multiple_k(k) + .set_max_iterations(max_iter) + .set_tol(tol) + .set_checkpoint_dir(ckpt); + if common.weighted { + b = b.set_edge_weight_col(&common.weight_col_name); + } + b = match init { + PicInit::Degree => b.set_init_strategy(InitStrategy::DegreeBased), + PicInit::Random => b.set_init_strategy(InitStrategy::Random), + }; + b = match weights { + PicWeights::None => b.set_weights_strategy(WeightsStrategy::None), + PicWeights::Ppmi => b.set_weights_strategy(WeightsStrategy::PPMI), + }; + b = match embedding { + PicEmbedding::Last => b.set_embedding_mode(EmbeddingMode::LastIterate), + PicEmbedding::Trajectory => b.set_embedding_mode(EmbeddingMode::FullTrajectory), + }; + let res = b.run(&ctx, &common.output).await?; + log::info!( + "PIC embedding dim = {}, KMeans Lloyd iterations = {}", + res.d, + res.num_iterations + ); + } + Command::Mllib { cmd } => match cmd { + MllibCommand::Kmeans { + vertices, + output, + format, + id_col_name, + feature_col, + k, + metric, + max_iter, + tol, + init_steps, + seed, + max_memory, + num_workers, + checkpoint_dir, + max_temp_file, + } => { + let work = ensure_dir(&checkpoint_dir)?; + let ctx = build_context(&work, &max_memory, num_workers, &max_temp_file)?; + + let raw = read_data(&ctx, &vertices, format).await?; + let features = + raw.select(vec![col(&id_col_name).alias(VERTEX_ID), col(&feature_col)])?; + + let metric = match metric { + KmeansArgsMetric::L2 => DistanceMetric::L2, + KmeansArgsMetric::Cosine => DistanceMetric::Cosine, + }; + + // Deduplicate K values preserving order (KMeansBuilder does the + // same), so runs zip 1:1 with the requested list. + let mut ks: Vec = Vec::new(); + for &kk in &k { + if !ks.contains(&kk) { + ks.push(kk); + } + } + + let res = KMeansBuilder::new(&features, &feature_col) + .k_values(&ks) + .metric(metric) + .max_iter(max_iter) + .tol(tol) + .init_steps(init_steps) + .seed(seed) + .run() + .await?; + log::info!( + "k-means finished after {} Lloyd iterations, d = {}", + res.num_iterations, + res.d + ); + + // One cluster column per requested K (named by the requested + // K; computed with the effective centers, see KMeansRun.k). + let mut columns = vec![col(VERTEX_ID)]; + for (kk, run) in ks.iter().zip(&res.runs) { + log::info!( + "k = {} (effective {}), total metric = {}", + kk, + run.k, + run.total_metric + ); + columns.push( + kmeans_assign_expr( + col(&feature_col), + run.k, + res.d, + run.centers.clone(), + metric, + ) + .alias(format!("cluster_{kk}")), + ); + } + features + .select(columns)? + .write_parquet(&output, DataFrameWriteOptions::new(), None) + .await?; + log::info!("result was written into {output}"); + } + }, Command::ClassicalLp { common, max_iter, diff --git a/src/ml.rs b/src/ml.rs new file mode 100644 index 0000000..6def91a --- /dev/null +++ b/src/ml.rs @@ -0,0 +1,8 @@ +mod distance; +mod kmeans; +mod linalg; + +pub use distance::DistanceMetric; +pub(crate) use distance::{nearest_center, nearest_centers}; +pub use kmeans::{KMeansBuilder, KMeansResult, KMeansRun}; +pub(crate) use linalg::{cosine_distance, l2_distance, l2_norm}; diff --git a/src/ml/distance.rs b/src/ml/distance.rs new file mode 100644 index 0000000..4914636 --- /dev/null +++ b/src/ml/distance.rs @@ -0,0 +1,135 @@ +use crate::ml::linalg::{cosine_distance, l2_distance, l2_norm}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum DistanceMetric { + L2, + Cosine, +} + +/// Nearest center to `feat` among `k` flat `centers` (`k * d` values). +/// +/// Returns the center index (ties break to the first center) and the distance +/// under the metric: **squared** L2 for [`DistanceMetric::L2`] (the D^2 cost +/// the k-means|| sampling consumes) and cosine distance for +/// [`DistanceMetric::Cosine`]. +pub(crate) fn nearest_centers( + feat: &[f32], + centers: &[f32], + k: usize, + d: usize, + metric: DistanceMetric, +) -> (usize, f32) { + let mut best_cluster = 0usize; + let mut best_dist = f32::MAX; + + match metric { + DistanceMetric::L2 => { + for t in 0..k { + let dist = l2_distance(feat, ¢ers[(t * d)..(t * d + d)], d); + if dist < best_dist { + best_dist = dist; + best_cluster = t; + } + } + + (best_cluster, best_dist) + } + DistanceMetric::Cosine => { + let mut c_norms = vec![0.0f32; k]; + for i in 0..k { + c_norms[i] = l2_norm(¢ers[(i * d)..(i * d + d)], d); + } + let x_norm = l2_norm(feat, d); + + for t in 0..k { + let dist = + cosine_distance(feat, ¢ers[(t * d)..(t * d + d)], d, x_norm, c_norms[t]); + if dist < best_dist { + best_dist = dist; + best_cluster = t; + } + } + + (best_cluster, best_dist) + } + } +} + +/// Index of the center nearest to `feat` (ties break to the first center). +/// +/// Thin wrapper over [`nearest_centers`] for the assignment scalar UDF, which +/// does not need the distance value. +pub(crate) fn nearest_center( + feat: &[f32], + centers: &[f32], + k: usize, + d: usize, + metric: DistanceMetric, +) -> usize { + nearest_centers(feat, centers, k, d, metric).0 +} + +#[cfg(test)] +mod tests { + use super::*; + + fn assert_close(actual: f32, expected: f32, what: &str) { + let scale = 1.0 + actual.abs() + expected.abs(); + assert!( + (actual - expected).abs() <= 1e-3 * scale, + "{what}: expected {expected}, got {actual}" + ); + } + + #[test] + fn test_nearest_centers_l2_picks_argmin() { + let feat = vec![0.0f32, 0.0]; + let centers = vec![1.0f32, 1.0, 10.0, 0.0, 5.0, 5.0]; + let (best, dist) = nearest_centers(&feat, ¢ers, 3, 2, DistanceMetric::L2); + assert_eq!(best, 0); + assert_close(dist, 2.0, "squared L2 to center 0"); + } + + #[test] + fn test_nearest_centers_metrics_disagree_on_purpose() { + // Unit vectors keep cosine exact: c0 is orthogonal (cos dist 1.0), c1 + // is parallel (cos dist 0.0) but far in L2. L2 must pick 0, cosine 1. + let feat = vec![1.0f32, 0.0]; + let centers = vec![0.0f32, 1.0, 1000.0, 0.0]; + let (l2_best, _) = nearest_centers(&feat, ¢ers, 2, 2, DistanceMetric::L2); + let (cos_best, _) = nearest_centers(&feat, ¢ers, 2, 2, DistanceMetric::Cosine); + assert_eq!(l2_best, 0); + assert_eq!(cos_best, 1); + } + + #[test] + fn test_nearest_centers_tie_breaks_to_first() { + let feat = vec![0.0f32, 0.0]; + let centers = vec![1.0f32, 0.0, 1.0, 0.0]; + let (best, dist) = nearest_centers(&feat, ¢ers, 2, 2, DistanceMetric::L2); + assert_eq!(best, 0); + assert_close(dist, 1.0, "squared L2 to the first center"); + } + + #[test] + fn test_nearest_centers_single_center() { + let feat = vec![3.0f32, 4.0]; + let centers = vec![0.0f32, 0.0]; + let (best, dist) = nearest_centers(&feat, ¢ers, 1, 2, DistanceMetric::L2); + assert_eq!(best, 0); + assert_close(dist, 25.0, "3-4-5 triangle, squared"); + } + + #[test] + fn test_nearest_center_returns_index_only() { + let feat = vec![0.0f32, 0.0]; + let centers = vec![5.0f32, 5.0, 1.0, 0.0]; + assert_eq!(nearest_center(&feat, ¢ers, 2, 2, DistanceMetric::L2), 1); + // Tie: both centers at squared distance 1 -> first one wins. + let tie_centers = vec![1.0f32, 0.0, -1.0, 0.0]; + assert_eq!( + nearest_center(&feat, &tie_centers, 2, 2, DistanceMetric::L2), + 0 + ); + } +} diff --git a/src/ml/kmeans.rs b/src/ml/kmeans.rs new file mode 100644 index 0000000..a030b17 --- /dev/null +++ b/src/ml/kmeans.rs @@ -0,0 +1,1117 @@ +//! Out-of-core K-Means driver. +//! +//! The run itself writes nothing: every iteration is one streaming scan of +//! the features (no caching, no checkpoints) and the result — iteration +//! count, per-k total metric and centers — is returned to the caller. +//! +//! Initialization follows +//! [Spark MLlib's k-means initializationf](https://github.com/apache/spark/blob/dc46b1479450e8e656bce703d831e3aada95ea5c/mllib/src/main/scala/org/apache/spark/mllib/clustering/KMeans.scala#L387-L459): +//! one hash-selected first center, `init_steps` rounds of D^2-weighted +//! candidate sampling (via the engine's per-row `random()`, since +//! `DataFrame::sample` does not exist), then a weighted local k-means++ over +//! the candidates. Only the candidate sampling is non-deterministic across +//! runs; everything seeded in the driver is derived from `seed`. +//! +//! Bahmani et al., Scalable K-Means++, VLDB 2012 +//! +//! +//! Contract: the feature column holds non-null, same-sized `f32` vectors +//! (`FixedSizeList` or `List`). + +use datafusion::arrow::array::{Array, FixedSizeListArray, Float64Array, Int32Array, Int64Array}; +use datafusion::arrow::datatypes::DataType; +use datafusion::common::{Result, plan_err}; +use datafusion::functions::math::random; +use datafusion::functions_aggregate::count::count; +use datafusion::functions_aggregate::sum::sum; +use datafusion::logical_expr::Expr; +use datafusion::prelude::*; +use rand::rngs::StdRng; + +/// The engine's per-row `random()` as an [`Expr`] (it is exposed as a bare +/// `ScalarUDF`, not an expression function). +fn random_expr() -> Expr { + random().call(vec![]) +} +use rand::{Rng, SeedableRng}; + +use crate::VERTEX_ID; +use crate::expressions::{ + as_f32_list_like, finite_axpb, kmeans_assign_expr, kmeans_cost_expr, kmeans_step_expr, + l2_norm_expr, +}; +use crate::ml::{DistanceMetric, nearest_centers}; + +/// `power(base, exp)` as an [`Expr`] (exposed as a bare `ScalarUDF`). +fn power_expr(base: Expr, exp: Expr) -> Expr { + datafusion::functions::math::power().call(vec![base, exp]) +} + +/// Result of a single K (effective K: when the data has fewer distinct points +/// than requested, K is clamped down, mirroring Spark returning fewer +/// distinct candidate centers). +#[derive(Debug, Clone)] +pub struct KMeansRun { + /// Effective number of centers. + pub k: usize, + /// Sum over all rows of the distance to the assigned center + /// (squared L2 — i.e. inertia — for the L2 metric, cosine distance sum + /// for the Cosine metric). + pub total_metric: f64, + /// Flat `k * d` center coordinates. + pub centers: Vec, +} + +/// Result of a [`KMeansBuilder`] run. +#[derive(Debug, Clone)] +pub struct KMeansResult { + /// Lloyd iterations executed (shared by all K values; the loop stops when + /// every K converged or `max_iter` was reached). + pub num_iterations: usize, + /// Feature dimension. + pub d: usize, + /// One run per requested K, in request order. + pub runs: Vec, +} + +struct RunState { + k_requested: usize, + k_eff: usize, + centers: Vec, + delta: f64, + /// `sum_c ||S_c||^2 / n_c` over the last Lloyd step's cluster sums; + l2_term: f64, +} + +/// Builder for an out-of-core K-Means run over a feature column. +pub struct KMeansBuilder<'a> { + features: &'a DataFrame, + feature_col: String, + ks: Vec, + metric: DistanceMetric, + max_iter: usize, + tol: f64, + init_steps: usize, + seed: u64, +} + +impl<'a> KMeansBuilder<'a> { + pub fn new(features: &'a DataFrame, feature_col: impl Into) -> Self { + Self { + features, + feature_col: feature_col.into(), + ks: Vec::new(), + metric: DistanceMetric::L2, + max_iter: 20, + tol: 1e-4, + init_steps: 2, + seed: 42, + } + } + + /// Requests a single K. + pub fn k(mut self, k: usize) -> Self { + self.ks = vec![k]; + self + } + + /// Requests multiple K values; all of them share one features scan per + /// Lloyd iteration. + pub fn k_values(mut self, ks: &[usize]) -> Self { + self.ks = ks.to_vec(); + self + } + + pub fn metric(mut self, metric: DistanceMetric) -> Self { + self.metric = metric; + self + } + + pub fn max_iter(mut self, max_iter: usize) -> Self { + self.max_iter = max_iter; + self + } + + pub fn tol(mut self, tol: f64) -> Self { + self.tol = tol; + self + } + + /// k-means|| sampling rounds (Spark's `initializationSteps`). + pub fn init_steps(mut self, init_steps: usize) -> Self { + self.init_steps = init_steps; + self + } + + pub fn seed(mut self, seed: u64) -> Self { + self.seed = seed; + self + } + + pub async fn run(&self) -> Result { + if self.ks.is_empty() { + return plan_err!("k-means requires at least one K value"); + } + if self.ks.iter().any(|&k| k == 0) { + return plan_err!("k-means K values must be >= 1"); + } + + // Deduplicate K values, preserving request order. + let mut ks: Vec = Vec::new(); + for &k in &self.ks { + if !ks.contains(&k) { + ks.push(k); + } + } + let d = self.detect_d().await?; + + let mut rng = StdRng::seed_from_u64(self.seed); + let centers_per_k = self.init_centers_multi(&ks, d, &mut rng).await?; + let mut states: Vec = ks + .iter() + .zip(centers_per_k) + .map(|(&k, centers)| RunState { + k_requested: k, + k_eff: centers.len() / d, + centers, + delta: f64::MAX, + l2_term: 0.0, + }) + .collect(); + + let mut iteration = 0usize; + let mut t_sqnorm: Option = None; + while iteration < self.max_iter { + iteration += 1; + let t = self.lloyd_step(&mut states, d).await?; + if t_sqnorm.is_none() { + t_sqnorm = t; + } + if states.iter().all(|s| s.delta < self.tol) { + break; + } + } + + // Total metric per K. + let metrics = if self.metric == DistanceMetric::L2 && iteration > 0 { + match t_sqnorm { + Some(t) => states.iter().map(|s| t - s.l2_term).collect(), + None => self.total_metrics(&states, d).await?, + } + } else { + self.total_metrics(&states, d).await? + }; + + Ok(KMeansResult { + num_iterations: iteration, + d, + runs: states + .into_iter() + .zip(metrics) + .map(|(s, metric)| KMeansRun { + k: s.k_eff, + total_metric: metric, + centers: s.centers, + }) + .collect(), + }) + } + + /// Feature dimension from the column type (`FixedSizeList`) or from the + /// first row (`List`). + async fn detect_d(&self) -> Result { + let field = self + .features + .schema() + .field_with_name(None, &self.feature_col)?; + match field.data_type() { + DataType::FixedSizeList(f, size) => { + if f.data_type() != &DataType::Float32 { + return plan_err!( + "k-means feature column must hold Float32, got {:?}", + f.data_type() + ); + } + Ok(*size as usize) + } + DataType::List(f) => { + if f.data_type() != &DataType::Float32 { + return plan_err!( + "k-means feature column must hold Float32, got {:?}", + f.data_type() + ); + } + let df = self + .features + .clone() + .select_columns(&[&self.feature_col])? + .limit(0, Some(1))?; + let (flat, d) = collect_feature_flat(&df, &self.feature_col).await?; + if flat.is_empty() { + return plan_err!("k-means requires at least one feature row"); + } + Ok(d) + } + other => plan_err!( + "k-means feature column must be FixedSizeList or List, got {other:?}" + ), + } + } + + /// The first feature row: sorted by a seeded affine hash of the vertex + /// id (deterministic given the seed; the frame always carries an Int64 + /// `id` column). + async fn first_center(&self, r_a: i64, r_b: i64) -> Result> { + let sorted = self.features.clone().sort(vec![ + finite_axpb(lit(r_a), col(VERTEX_ID), lit(r_b)).sort(true, true), + ])?; + let df = sorted + .select_columns(&[&self.feature_col])? + .limit(0, Some(1))?; + let (flat, _) = collect_feature_flat(&df, &self.feature_col).await?; + if flat.is_empty() { + return plan_err!("k-means requires at least one feature row"); + } + Ok(flat) + } + + /// Spark MLlib `initKMeansParallel`, run for **all K values with every + /// features scan shared**: one shared first center, then per sampling + /// round one shared cost-sum scan and one shared sampling scan, one + /// shared weighting scan — the number of init passes is independent of + /// how many Ks are requested (per-K work is extra compute in the same + /// scans, not extra scans). + async fn init_centers_multi( + &self, + ks: &[usize], + d: usize, + rng: &mut StdRng, + ) -> Result>> { + // Shared first center (one scan): per-K diversity comes from the + // per-K sampling thresholds afterwards. + let mut r_a = rng.random::(); + while r_a == 0 { + r_a = rng.random::(); + } + let r_b = rng.random::(); + let first = self.first_center(r_a, r_b).await?; + let mut candidates: Vec>> = ks.iter().map(|_| vec![first.clone()]).collect(); + // A K becomes inactive once all its cost mass sits on its centers. + let mut active: Vec = ks.iter().map(|_| true).collect(); + + for _ in 0..self.init_steps { + if !active.iter().any(|a| *a) { + break; + } + + // One scan: per-K cost sums under each K's accumulated centers. + let cost_exprs: Vec = candidates + .iter() + .enumerate() + .map(|(i, c)| { + sum(kmeans_cost_expr( + col(&self.feature_col), + c.len(), + d, + flatten(c), + self.metric, + )) + .alias(format!("__cs{i}")) + }) + .collect(); + let batches = self + .features + .clone() + .aggregate(vec![], cost_exprs)? + .collect() + .await?; + let mut cost_sums = vec![0.0f64; ks.len()]; + for (i, s) in cost_sums.iter_mut().enumerate() { + *s = batches[0] + .column(i) + .as_any() + .downcast_ref::() + .expect("sum of Float64") + .value(0); + if *s <= 0.0 { + active[i] = false; + } + } + if !active.iter().any(|a| *a) { + break; + } + + // Thresholds as per-row expressions; the gate is their sum. + let thresholds: Vec = ks + .iter() + .enumerate() + .map(|(i, &k)| { + if active[i] { + lit(2.0 * (k as f64)) + * kmeans_cost_expr( + col(&self.feature_col), + candidates[i].len(), + d, + flatten(&candidates[i]), + self.metric, + ) + / lit(cost_sums[i]) + } else { + lit(0.0f64) + } + }) + .collect(); + let gate = thresholds + .iter() + .fold(lit(0.0f64), |acc, t| acc + t.clone()); + let mut selection: Vec = vec![col(&self.feature_col)]; + for (i, t) in thresholds.iter().enumerate() { + selection.push( + random_expr() + .lt(t.clone() / gate.clone()) + .alias(format!("__s{i}")), + ); + } + let df = self + .features + .clone() + .filter(random_expr().lt(gate))? + .select(selection)?; + let batches = df.collect().await?; + let mut added = 0usize; + for b in &batches { + let v = as_f32_list_like(b.column(0), "k_means", &self.feature_col)?; + let flags: Vec = (0..ks.len()) + .map(|i| { + b.column(1 + i) + .as_any() + .downcast_ref::() + .expect("boolean flag") + .clone() + }) + .collect(); + for row in 0..v.len() { + for (i, flag) in flags.iter().enumerate() { + if active[i] && !flag.is_null(row) && flag.value(row) { + candidates[i].push(v.value(row).to_vec()); + added += 1; + } + } + } + } + if added == 0 { + break; + } + } + + // Distinct candidates per K; Ks with few enough candidates are + // clamped right away. + let distincts: Vec>> = candidates + .iter() + .map(|cs| { + let mut distinct: Vec> = Vec::new(); + for c in cs { + if !distinct.iter().any(|e| e == c) { + distinct.push(c.clone()); + } + } + distinct + }) + .collect(); + let needs_weights: Vec = ks + .iter() + .zip(&distincts) + .map(|(&k, ds)| ds.len() > k) + .collect(); + + // One shared weighting scan: group by all per-K assignment columns at + // once and decompose the combination counts per K in the driver. + // Clamped Ks contribute a constant column (their weights are unused). + let group_exprs: Vec = ks + .iter() + .enumerate() + .map(|(i, _)| { + if needs_weights[i] { + kmeans_assign_expr( + col(&self.feature_col), + distincts[i].len(), + d, + flatten(&distincts[i]), + self.metric, + ) + .alias(format!("__a{i}")) + } else { + lit(0i32).alias(format!("__a{i}")) + } + }) + .collect(); + let batches = self + .features + .clone() + .aggregate(group_exprs, vec![count(lit(1)).alias("__w")])? + .collect() + .await?; + let mut weights: Vec> = ks + .iter() + .zip(&distincts) + .map(|(_, ds)| vec![0.0f64; ds.len()]) + .collect(); + for b in &batches { + let w = b + .column(ks.len()) + .as_any() + .downcast_ref::() + .expect("count"); + for row in 0..w.len() { + if w.is_null(row) { + continue; + } + for (i, _) in ks.iter().enumerate() { + if !needs_weights[i] { + continue; + } + let a = b + .column(i) + .as_any() + .downcast_ref::() + .expect("assignment") + .value(row); + weights[i][a as usize] += w.value(row) as f64; + } + } + } + + // Local weighted k-means++ per K (in-memory; candidate sets are tiny). + let mut result = Vec::with_capacity(ks.len()); + for (i, &k) in ks.iter().enumerate() { + if distincts[i].len() <= k { + result.push(flatten(&distincts[i])); + } else { + result.push(flatten(&local_kmeans_pp( + rng, + &distincts[i], + &weights[i], + k, + d, + 30, + self.metric, + ))); + } + } + Ok(result) + } + + /// One Lloyd iteration for all K values in a single features scan. + /// + /// Also fuses `sum(||x||^2)` (one norm per row — a fraction of the step's + /// `k*d` distance work) so the final L2 metric needs no extra scan; + /// returns that total (constant across iterations). + async fn lloyd_step(&self, states: &mut [RunState], d: usize) -> Result> { + let (uniq, idx) = dedup_states(states); + let mut exprs: Vec = uniq + .iter() + .enumerate() + .map(|(i, (k, c))| { + kmeans_step_expr(col(&self.feature_col), *k, d, c.clone(), self.metric) + .alias(format!("__km{i}")) + }) + .collect(); + exprs.push( + sum(power_expr( + l2_norm_expr(col(&self.feature_col)), + lit(2.0f64), + )) + .alias("__sq"), + ); + let batches = self + .features + .clone() + .aggregate(vec![], exprs)? + .collect() + .await?; + if batches.is_empty() || batches[0].num_rows() == 0 { + return plan_err!("k-means step aggregate produced no rows"); + } + let sq_col = batches[0].column(uniq.len()); + let t_sqnorm = sq_col + .as_any() + .downcast_ref::() + .expect("sum of Float64") + .value(0); + + for (s, &ui) in states.iter_mut().zip(&idx) { + let (_, centers) = &uniq[ui]; + let arr = batches[0].column(ui); + let fsl = arr + .as_any() + .downcast_ref::() + .ok_or_else(|| { + datafusion::common::DataFusionError::Execution( + "k-means step result must be a FixedSizeList".to_string(), + ) + })?; + let binding = fsl.value(0); + let vals = binding + .as_any() + .downcast_ref::() + .expect("k_means_step returns Float64"); + let k = s.k_eff; + let expected = k * d + k; + if vals.len() != expected { + return plan_err!( + "k-means step returned {expected} values expected, got {}", + vals.len() + ); + } + let mut new_centers = centers.clone(); + let mut delta = 0.0f64; + let mut l2_term = 0.0f64; + for c in 0..k { + let cnt = vals.value(k * d + c); + if cnt > 0.0 { + let mut sq_sum = 0.0f64; + for t in 0..d { + let sv = vals.value(c * d + t); + sq_sum += sv * sv; + let nv = (sv / cnt) as f32; + delta = delta.max((nv as f64 - centers[c * d + t] as f64).abs()); + new_centers[c * d + t] = nv; + } + l2_term += sq_sum / cnt; + } else { + log::warn!( + "k-means: empty cluster {c} (k={}) kept its previous center", + s.k_requested + ); + } + } + s.centers = new_centers; + s.delta = delta; + s.l2_term = l2_term; + } + Ok(Some(t_sqnorm)) + } + + /// Per-K total metric in one scan: `sum(k_means_cost(feat))` per distinct + /// final state. + async fn total_metrics(&self, states: &[RunState], d: usize) -> Result> { + let (uniq, idx) = dedup_states(states); + let exprs: Vec = uniq + .iter() + .enumerate() + .map(|(i, (k, c))| { + sum(kmeans_cost_expr( + col(&self.feature_col), + *k, + d, + c.clone(), + self.metric, + )) + .alias(format!("__mt{i}")) + }) + .collect(); + let batches = self + .features + .clone() + .aggregate(vec![], exprs)? + .collect() + .await?; + if batches.is_empty() || batches[0].num_rows() == 0 { + return plan_err!("k-means metric aggregate produced no rows"); + } + let mut metrics = vec![0.0f64; states.len()]; + for (m, &ui) in metrics.iter_mut().zip(&idx) { + let arr = batches[0].column(ui); + *m = arr + .as_any() + .downcast_ref::() + .expect("sum of Float64") + .value(0); + } + Ok(metrics) + } +} + +/// Collects the (single) selected feature column as a flat buffer plus the +/// row dimension `d` (0 rows -> empty buffer, `d == 0`). +async fn collect_feature_flat(df: &DataFrame, col_name: &str) -> Result<(Vec, usize)> { + let batches = df.clone().collect().await?; + let mut flat: Vec = Vec::new(); + let mut d = 0usize; + for b in &batches { + let v = as_f32_list_like(b.column(0), "k_means", col_name)?; + for i in 0..v.len() { + let row = v.value(i); + if d == 0 { + d = row.len(); + } + flat.extend_from_slice(row); + } + } + Ok((flat, d)) +} + +fn flatten(centers: &[Vec]) -> Vec { + centers.concat() +} + +/// Groups identical (k_eff, centers) states so the aggregate does not carry +/// duplicate expressions (which DataFusion would collapse into one column). +/// Returns the distinct states and, per input state, its index into them. +fn dedup_states(states: &[RunState]) -> (Vec<(usize, Vec)>, Vec) { + let mut uniq: Vec<(usize, Vec)> = Vec::new(); + let mut idx: Vec = Vec::with_capacity(states.len()); + for s in states { + match uniq.iter().position(|(k, c)| { + *k == s.k_eff + && c.len() == s.centers.len() + && c.iter() + .zip(&s.centers) + .all(|(a, b)| a.to_bits() == b.to_bits()) + }) { + Some(i) => idx.push(i), + None => { + uniq.push((s.k_eff, s.centers.clone())); + idx.push(uniq.len() - 1); + } + } + } + (uniq, idx) +} + +/// Weighted pick: returns an index with probability proportional to `weights`. +fn pick_weighted(rng: &mut StdRng, weights: &[f64]) -> usize { + let total: f64 = weights.iter().sum(); + if total <= 0.0 { + return rng.random_range(0..weights.len()); + } + let mut r = rng.random::() * total; + for (i, w) in weights.iter().enumerate() { + r -= w; + if r <= 0.0 { + return i; + } + } + weights.len() - 1 +} + +/// Spark MLlib `LocalKMeans.kMeansPlusPlus`: weighted k-means++ seeding +/// followed by weighted Lloyd iterations over the (small) candidate set. +fn local_kmeans_pp( + rng: &mut StdRng, + points: &[Vec], + weights: &[f64], + k: usize, + d: usize, + max_iter: usize, + metric: DistanceMetric, +) -> Vec> { + let n = points.len(); + + // Weighted k-means++ seeding: first center with probability proportional + // to its weight, each next center with probability proportional to + // weight * D^2 under the current centers. + let mut centers: Vec> = vec![points[pick_weighted(rng, weights)].clone()]; + while centers.len() < k { + let flat = flatten(¢ers); + let mut costs = vec![0.0f64; n]; + let mut total = 0.0; + for (i, p) in points.iter().enumerate() { + let (_, dist) = nearest_centers(p, &flat, centers.len(), d, metric); + costs[i] = dist as f64 * weights[i]; + total += costs[i]; + } + if total <= 0.0 { + // Every remaining point coincides with a center: fall back to a + // weighted pick. + centers.push(points[pick_weighted(rng, weights)].clone()); + continue; + } + let mut r = rng.random::() * total; + let mut pick = n - 1; + for (i, c) in costs.iter().enumerate() { + r -= c; + if r <= 0.0 { + pick = i; + break; + } + } + centers.push(points[pick].clone()); + } + + // Weighted Lloyd iterations. + let mut assignment = vec![usize::MAX; n]; + for _ in 0..max_iter { + let flat = flatten(¢ers); + let mut sums = vec![0.0f64; k * d]; + let mut cnts = vec![0.0f64; k]; + let mut changed = false; + for (i, p) in points.iter().enumerate() { + let (c, _) = nearest_centers(p, &flat, k, d, metric); + if c != assignment[i] { + changed = true; + assignment[i] = c; + } + for t in 0..d { + sums[c * d + t] += p[t] as f64 * weights[i]; + } + cnts[c] += weights[i]; + } + for c in 0..k { + if cnts[c] > 0.0 { + for t in 0..d { + centers[c][t] = (sums[c * d + t] / cnts[c]) as f32; + } + } + } + if !changed { + break; + } + } + centers +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::{ + ArrayRef, FixedSizeListArray, Float32Array, Int64Array, RecordBatch, + }; + use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::datasource::MemTable; + use std::sync::Arc; + use std::time::Instant; + + fn lcg(seed: &mut u32) -> f32 { + *seed = seed.wrapping_mul(1664525).wrapping_add(1013904223); + ((*seed >> 8) % 1001) as f32 / 100.0 - 5.0 + } + + fn blobs(seed: u32, n_per: usize, means: &[[f32; 2]], noise_amp: f32) -> Vec<[f32; 2]> { + let mut s = seed; + let mut rows = Vec::new(); + for m in means { + for _ in 0..n_per { + let nx = lcg(&mut s) * noise_amp; + let ny = lcg(&mut s) * noise_amp; + rows.push([m[0] + nx, m[1] + ny]); + } + } + rows + } + + async fn features_df(ctx: &SessionContext, rows: &[[f32; 2]]) -> Result { + let schema = Schema::new(vec![ + Field::new(VERTEX_ID, DataType::Int64, false), + Field::new( + "feat", + DataType::FixedSizeList(Arc::new(Field::new("el", DataType::Float32, false)), 2), + false, + ), + ]); + let flat: Vec = rows.iter().flatten().copied().collect(); + let fsl = FixedSizeListArray::try_new( + Arc::new(Field::new("el", DataType::Float32, false)), + 2, + Arc::new(Float32Array::from(flat)), + None, + )?; + let ids: Int64Array = (0..rows.len() as i64).collect(); + let batch = RecordBatch::try_new( + Arc::new(schema.clone()), + vec![Arc::new(ids) as ArrayRef, Arc::new(fsl) as ArrayRef], + )?; + let table = MemTable::try_new(Arc::new(schema), vec![vec![batch]])?; + ctx.register_table("km_t", Arc::new(table))?; + Ok(ctx.table("km_t").await?) + } + + fn center_pairs(run: &KMeansRun) -> Vec<(f32, f32)> { + let mut cs: Vec<(f32, f32)> = run.centers.chunks(2).map(|c| (c[0], c[1])).collect(); + cs.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + cs + } + + #[tokio::test] + async fn recovers_three_blobs() -> Result<()> { + let ctx = SessionContext::new(); + let rows = blobs(7, 30, &[[0.0, 0.0], [20.0, 20.0], [-20.0, 20.0]], 1.0); + let df = features_df(&ctx, &rows).await?; + let res = KMeansBuilder::new(&df, "feat").k(3).run().await?; + assert_eq!(res.d, 2); + assert_eq!(res.runs.len(), 1); + let run = &res.runs[0]; + assert_eq!(run.k, 3); + assert!(run.total_metric < 4000.0, "metric {}", run.total_metric); + let want = [(-20.0f32, 20.0f32), (0.0, 0.0), (20.0, 20.0)]; + for (got, w) in center_pairs(run).iter().zip(want) { + assert!( + (got.0 - w.0).abs() < 2.0 && (got.1 - w.1).abs() < 2.0, + "center {got:?} vs {w:?}" + ); + } + assert!( + res.num_iterations < 20, + "well-separated blobs must converge early, took {}", + res.num_iterations + ); + Ok(()) + } + + #[tokio::test] + async fn multi_k_returns_run_per_k() -> Result<()> { + let ctx = SessionContext::new(); + let rows = blobs(11, 30, &[[0.0, 0.0], [20.0, 20.0], [-20.0, 20.0]], 1.0); + let df = features_df(&ctx, &rows).await?; + // Duplicate K values are deduplicated. + let res = KMeansBuilder::new(&df, "feat") + .k_values(&[1, 1, 3]) + .run() + .await?; + assert_eq!(res.runs.len(), 2); + assert_eq!(res.runs[0].k, 1); + assert_eq!(res.runs[1].k, 3); + // The single-cluster center is the global mean (0, 40/3). + let c = center_pairs(&res.runs[0]); + assert_eq!(c.len(), 1); + assert!((c[0].0 - 0.0).abs() < 1.0, "{:?}", c[0]); + assert!((c[0].1 - 40.0 / 3.0).abs() < 1.0, "{:?}", c[0]); + Ok(()) + } + + #[tokio::test] + async fn max_iter_is_respected() -> Result<()> { + let ctx = SessionContext::new(); + let rows = blobs(13, 30, &[[0.0, 0.0], [20.0, 20.0], [-20.0, 20.0]], 1.0); + let df = features_df(&ctx, &rows).await?; + let res = KMeansBuilder::new(&df, "feat") + .k(3) + .max_iter(1) + .run() + .await?; + assert_eq!(res.num_iterations, 1); + Ok(()) + } + + #[tokio::test] + async fn k1_converges_to_the_mean() -> Result<()> { + let ctx = SessionContext::new(); + let rows = [[1.0f32, 2.0], [3.0, 4.0], [5.0, 6.0]]; + let df = features_df(&ctx, &rows).await?; + let res = KMeansBuilder::new(&df, "feat") + .seed(9) + .tol(1e-6) + .k(1) + .run() + .await?; + let run = &res.runs[0]; + assert_eq!(run.k, 1); + assert!((run.centers[0] - 3.0).abs() < 1e-3, "{:?}", run.centers); + assert!((run.centers[1] - 4.0).abs() < 1e-3, "{:?}", run.centers); + assert!( + (run.total_metric - 16.0).abs() < 1e-3, + "metric {}", + run.total_metric + ); + Ok(()) + } + + #[tokio::test] + async fn identical_points_clamp_k_and_share_one_expression() -> Result<()> { + // Fewer distinct points than K: K is clamped down (Spark returns the + // distinct candidates), and both runs end up identical, which must not + // confuse the multi-K aggregate (duplicate expressions collapse). + let ctx = SessionContext::new(); + let rows = vec![[1.0f32, 2.0]; 6]; + let df = features_df(&ctx, &rows).await?; + let res = KMeansBuilder::new(&df, "feat") + .k_values(&[2, 3]) + .run() + .await?; + assert_eq!(res.runs.len(), 2); + for run in &res.runs { + assert_eq!(run.k, 1); + // The L2 metric is algebraic (T - sum ||S_c||^2/n_c); `T` carries + // a ~1e-7-relative sqrt-then-square round trip, so identical + // points yield ~0, not bitwise 0. + assert!(run.total_metric.abs() < 1e-4, "{}", run.total_metric); + assert_eq!(run.centers, vec![1.0, 2.0]); + } + Ok(()) + } + + #[tokio::test] + async fn empty_features_errors() -> Result<()> { + let ctx = SessionContext::new(); + let df = features_df(&ctx, &[]).await?; + let res = KMeansBuilder::new(&df, "feat").k(2).run().await; + assert!(res.is_err(), "empty features must error"); + Ok(()) + } + + #[tokio::test] + async fn single_init_step_still_recovers_blobs() -> Result<()> { + let ctx = SessionContext::new(); + let rows = blobs(19, 30, &[[0.0, 0.0], [20.0, 20.0], [-20.0, 20.0]], 1.0); + let df = features_df(&ctx, &rows).await?; + let res = KMeansBuilder::new(&df, "feat") + .k(3) + .init_steps(1) + .run() + .await?; + let run = &res.runs[0]; + assert_eq!(run.k, 3); + let want = [(-20.0f32, 20.0f32), (0.0, 0.0), (20.0, 20.0)]; + for (got, w) in center_pairs(run).iter().zip(want) { + assert!((got.0 - w.0).abs() < 2.0 && (got.1 - w.1).abs() < 2.0); + } + Ok(()) + } + + /// The fused (algebraic) L2 metric must equal an explicit + /// `sum(k_means_cost)` scan under the *returned* centers. + #[tokio::test] + async fn fused_l2_metric_matches_explicit_scan() -> Result<()> { + let ctx = SessionContext::new(); + let rows = blobs(23, 40, &[[0.0, 0.0], [20.0, 20.0], [-20.0, 20.0]], 2.0); + let df = features_df(&ctx, &rows).await?; + let res = KMeansBuilder::new(&df, "feat").k(3).run().await?; + let run = &res.runs[0]; + let manual = df + .clone() + .aggregate( + vec![], + vec![ + sum(kmeans_cost_expr( + col("feat"), + run.k, + res.d, + run.centers.clone(), + DistanceMetric::L2, + )) + .alias("m"), + ], + )? + .collect() + .await?; + let explicit = manual[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + let scale = 1.0 + explicit.abs(); + assert!( + (run.total_metric - explicit).abs() <= 1e-3 * scale, + "fused {} vs explicit {explicit}", + run.total_metric + ); + Ok(()) + } + + /// The cosine fallback (final scan) is by construction the explicit + /// definition — pin it. + #[tokio::test] + async fn cosine_metric_matches_explicit_scan() -> Result<()> { + let ctx = SessionContext::new(); + let rows = blobs(29, 30, &[[10.0, 0.0], [-5.0, 8.66], [-5.0, -8.66]], 0.1); + let df = features_df(&ctx, &rows).await?; + let res = KMeansBuilder::new(&df, "feat") + .k(3) + .metric(DistanceMetric::Cosine) + .run() + .await?; + let run = &res.runs[0]; + let manual = df + .clone() + .aggregate( + vec![], + vec![ + sum(kmeans_cost_expr( + col("feat"), + run.k, + res.d, + run.centers.clone(), + DistanceMetric::Cosine, + )) + .alias("m"), + ], + )? + .collect() + .await?; + let explicit = manual[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + let scale = 1.0 + explicit.abs(); + assert!( + (run.total_metric - explicit).abs() <= 1e-9 * scale, + "cosine fallback {} vs explicit {explicit}", + run.total_metric + ); + Ok(()) + } + + /// Multi-K init produces a valid run per K and, after Lloyd, every + /// K >= 3 covers all three blobs. (The scan *sharing* itself — one + /// cost-sum scan, one sampling scan and one weighting scan per round + /// regardless of how many Ks — is a property of the query plan; it is + /// measured by the `kmeans_at_scale` bench, not assertable by wall clock + /// at unit-test scale where per-query planning overhead dominates.) + #[tokio::test] + async fn multi_k_init_shared_scans() -> Result<()> { + let ctx = SessionContext::new(); + let rows = blobs(31, 40, &[[0.0, 0.0], [20.0, 20.0], [-20.0, 20.0]], 1.0); + let df = features_df(&ctx, &rows).await?; + + let multi0 = KMeansBuilder::new(&df, "feat") + .k_values(&[3, 5, 7]) + .max_iter(0) + .run() + .await?; + assert_eq!(multi0.runs.len(), 3); + + // Quality after Lloyd: every K >= 3 must cover all three blobs (an + // existence check per blob — larger K may split blobs into several + // centers, so positional zipping would mismatch). + let multi = KMeansBuilder::new(&df, "feat") + .k_values(&[3, 5, 7]) + .run() + .await?; + assert_eq!(multi.runs.len(), 3); + let want = [(-20.0f32, 20.0f32), (0.0, 0.0), (20.0, 20.0)]; + for run in &multi.runs { + assert!(run.k >= 3, "k={}", run.k); + for w in want { + let covered = center_pairs(run) + .iter() + .any(|c| (c.0 - w.0).abs() < 3.0 && (c.1 - w.1).abs() < 3.0); + assert!( + covered, + "k={} misses blob {w:?}: {:?}", + run.k, + center_pairs(run) + ); + } + } + Ok(()) + } + + #[tokio::test] + async fn cosine_metric_separates_directions() -> Result<()> { + let ctx = SessionContext::new(); + let rows = blobs(17, 30, &[[10.0, 0.0], [-5.0, 8.66], [-5.0, -8.66]], 0.1); + let df = features_df(&ctx, &rows).await?; + let res = KMeansBuilder::new(&df, "feat") + .k(3) + .metric(DistanceMetric::Cosine) + .run() + .await?; + let run = &res.runs[0]; + assert_eq!(run.k, 3); + assert!(run.total_metric < 60.0, "metric {}", run.total_metric); + Ok(()) + } +} diff --git a/src/ml/linalg.rs b/src/ml/linalg.rs new file mode 100644 index 0000000..ac17125 --- /dev/null +++ b/src/ml/linalg.rs @@ -0,0 +1,221 @@ +//! SIMD linear-algebra kernels over `f32` vectors. +//! +//! All manual SIMD code for the ml module lives here; this module is +//! DataFusion-free on purpose — the scalar-UDF wrappers live in +//! [`crate::expressions::linalg`]. +//! +//! Contract: vectors are non-null, same-sized slices; for the two-argument +//! kernels both slices have the same length. + +use std::ops::Add; + +use wide::f32x8; + +/// Squared L2 distance `||x - c||^2`. +/// +/// Kept squared (no `sqrt`) because the only hot-loop caller, the K-Means +/// nearest-center search, compares distances and `sqrt` is monotone. +pub(crate) fn l2_distance(x: &[f32], c: &[f32], d: usize) -> f32 { + let mut acc = f32x8::splat(0.0); + let mut t = 0; + while t + 8 <= d { + let xv = f32x8::from(&x[t..t + 8]); + let cv = f32x8::from(&c[t..t + 8]); + acc = (xv - cv).mul_add(xv - cv, acc); + t += 8; + } + + let mut dist = acc.reduce_add(); + + while t < d { + let xi = x[t]; + let ci = c[t]; + dist += (xi - ci) * (xi - ci); + t += 1; + } + + dist +} + +/// True L2 norm `sqrt(||x||^2)`. +pub(crate) fn l2_norm(x: &[f32], d: usize) -> f32 { + let mut acc = f32x8::splat(0.0); + let mut t = 0; + while t + 8 <= d { + let xv = f32x8::from(&x[t..t + 8]); + acc = xv.mul_add(xv, acc); + t += 8; + } + + let mut l2 = acc.reduce_add(); + + while t < d { + let xi = x[t]; + l2 += xi * xi; + t += 1; + } + + l2.sqrt() +} + +/// Cosine distance `1 - (x . c) / (|x| |c|)`. +/// +/// Zero-norm vectors score `0.0`: scikit-learn semantics +pub(crate) fn cosine_distance(x: &[f32], c: &[f32], d: usize, x_norm: f32, c_norm: f32) -> f32 { + let denom = x_norm * c_norm; + + if denom == 0f32 { + return 0f32; + } + + let mut mult = f32x8::splat(0.0); + let mut t = 0; + + while t + 8 <= d { + let xv = f32x8::from(&x[t..t + 8]); + let cv = f32x8::from(&c[t..t + 8]); + + mult = mult.add(xv * cv); + t += 8; + } + + let mut numerator = mult.reduce_add(); + + while t < d { + let xi = x[t]; + let ci = c[t]; + numerator += xi * ci; + t += 1; + } + + 1f32 - numerator / denom +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Naive scalar squared-L2 reference. + fn l2_ref(x: &[f32], c: &[f32]) -> f32 { + x.iter().zip(c).map(|(a, b)| (a - b) * (a - b)).sum() + } + + /// True L2 norm reference: `sqrt(sum of squares)`. + fn norm_ref(x: &[f32]) -> f32 { + x.iter().map(|a| a * a).sum::().sqrt() + } + + /// Naive scalar cosine-distance reference: `1 - dot / (|x| |c|)`. + fn cosine_ref(x: &[f32], c: &[f32]) -> f32 { + let dot: f32 = x.iter().zip(c).map(|(a, b)| a * b).sum(); + let norm_x = norm_ref(x); + let norm_c = norm_ref(c); + if norm_x == 0.0 || norm_c == 0.0 { + 0.0 + } else { + 1.0 - dot / (norm_x * norm_c) + } + } + + fn assert_close(actual: f32, expected: f32, what: &str) { + let scale = 1.0 + actual.abs() + expected.abs(); + assert!( + (actual - expected).abs() <= 1e-3 * scale, + "{what}: expected {expected}, got {actual}" + ); + } + + /// Deterministic pseudo-random vector pair with mixed signs and magnitudes. + fn vec_pair(d: usize, seed: u32) -> (Vec, Vec) { + let mut state = seed; + let mut next = move || { + state = state.wrapping_mul(1664525).wrapping_add(1013904223); + ((state >> 8) % 2001) as f32 / 100.0 - 10.0 + }; + let x = (0..d).map(|_| next()).collect(); + let c = (0..d).map(|_| next()).collect(); + (x, c) + } + + // ---------------- kernels ---------------- + + #[test] + fn test_l2_distance_matches_scalar_reference() { + for d in [1usize, 3, 7, 8, 9, 15, 16, 17, 64, 100] { + let (x, c) = vec_pair(d, 42 + d as u32); + assert_close( + l2_distance(&x, &c, d), + l2_ref(&x, &c), + &format!("l2_distance, d={d}"), + ); + } + } + + #[test] + fn test_l2_distance_identical_vectors_is_zero() { + let x = vec![1.0f32, -2.0, 3.5, 0.0]; + assert_eq!(l2_distance(&x, &x, x.len()), 0.0); + } + + #[test] + fn test_l2_norm_matches_scalar_reference() { + for d in [1usize, 3, 7, 8, 9, 15, 16, 17, 64, 100] { + let (x, _) = vec_pair(d, 42 + d as u32); + assert_close(l2_norm(&x, d), norm_ref(&x), &format!("l2_norm, d={d}")); + } + // Zero vector. + let zero = vec![0.0f32; 17]; + assert_eq!(l2_norm(&zero, 17), 0.0); + // 3-4-5 triangle. + let t = vec![3.0f32, 4.0]; + assert_close(l2_norm(&t, 2), 5.0, "3-4-5"); + } + + #[test] + fn test_cosine_distance_matches_scalar_reference() { + for d in [1usize, 3, 7, 8, 9, 15, 16, 17, 64, 100] { + let (x, c) = vec_pair(d, 7 + d as u32); + assert_close( + cosine_distance(&x, &c, d, norm_ref(&x), norm_ref(&c)), + cosine_ref(&x, &c), + &format!("cosine_distance, d={d}"), + ); + } + } + + #[test] + fn test_cosine_distance_scaled_parallel_vectors_is_zero() { + // Collinear vectors must have cosine distance 0 regardless of scale. + let x = vec![1.0f32, 2.0, 3.0]; + let c = vec![2.0f32, 4.0, 6.0]; + assert_close( + cosine_distance(&x, &c, 3, norm_ref(&x), norm_ref(&c)), + 0.0, + "scaled parallel", + ); + } + + #[test] + fn test_cosine_distance_unit_vector_canonical_values() { + let e1 = vec![1.0f32, 0.0]; + let e2 = vec![0.0f32, 1.0]; + assert_close(cosine_distance(&e1, &e1, 2, 1.0, 1.0), 0.0, "parallel"); + assert_close(cosine_distance(&e1, &e2, 2, 1.0, 1.0), 1.0, "orthogonal"); + assert_close( + cosine_distance(&e1, &vec![-1.0, 0.0], 2, 1.0, 1.0), + 2.0, + "opposite", + ); + } + + #[test] + fn test_cosine_distance_zero_vector_policy() { + // Pinned decision: a zero vector scores as "identical to everything". + let zero = vec![0.0f32, 0.0, 0.0]; + let v = vec![1.0f32, 2.0, 3.0]; + let nv = norm_ref(&v); + assert_eq!(cosine_distance(&zero, &v, 3, 0.0, nv), 0.0); + assert_eq!(cosine_distance(&v, &zero, 3, nv, 0.0), 0.0); + assert_eq!(cosine_distance(&zero, &zero, 3, 0.0, 0.0), 0.0); + } +} diff --git a/src/utils/graph_utils.rs b/src/utils/graph_utils.rs index 1a235e6..2776f57 100644 --- a/src/utils/graph_utils.rs +++ b/src/utils/graph_utils.rs @@ -5,18 +5,39 @@ use datafusion::{error::Result, prelude::DataFrame}; /// Prepares the edge set: drops self-loops, symmetrizes /// (adds the reverse of every edge), and deduplicates (optionally). The result is the /// undirected simple graph the algorithm operates on. -pub(crate) fn symmetrize(edges: &DataFrame, do_distinct: bool) -> Result { +pub(crate) fn symmetrize( + edges: &DataFrame, + do_distinct: bool, + attr_cols: Option>, +) -> Result { // some algorithms to the aggregation over edges and are tolerant // to duplicates; // // that is the only reason this function is // a) packgage-private // b) has second argument - let no_loops = edges.clone().filter(col(EDGE_SRC).not_eq(col(EDGE_DST)))?; - let reversed = no_loops.clone().select(vec![ - col(EDGE_DST).alias(EDGE_SRC), - col(EDGE_SRC).alias(EDGE_DST), - ])?; + let (forward_cols, backward_cols) = match attr_cols { + Some(cols) => { + let mut fcc = vec![col(EDGE_SRC), col(EDGE_DST)]; + let mut bcc = vec![col(EDGE_DST).alias(EDGE_SRC), col(EDGE_SRC).alias(EDGE_DST)]; + + let attrs: Vec = cols.iter().map(|c| col(c)).collect(); + + fcc.extend(attrs.clone()); + bcc.extend(attrs.clone()); + + (fcc, bcc) + } + None => ( + vec![col(EDGE_SRC), col(EDGE_DST)], + vec![col(EDGE_DST).alias(EDGE_SRC), col(EDGE_SRC).alias(EDGE_DST)], + ), + }; + let no_loops = edges + .clone() + .filter(col(EDGE_SRC).not_eq(col(EDGE_DST)))? + .select(forward_cols)?; + let reversed = no_loops.clone().select(backward_cols)?; let res = if do_distinct { no_loops.union(reversed)?.distinct()?