diff --git a/README.md b/README.md index e2a8943..f88e314 100644 --- a/README.md +++ b/README.md @@ -47,6 +47,8 @@ Available algorithms: - `shortest-path`: Multi-Source Shortest Path; - `classical-lp`: Classical (Raghavan) Label Propagaion; - `pic`: Power Iteration Clustering; +- `fastrp`: FastRP (Fast Random Projection) vertex embeddings (optional per-iteration L1/L2 degree normalization, optional weighted linear combination of the iterates `H = Σ w_t·H_t`, and optional L2 output normalization); +- `fastrp-clustering`: graph clustering over FastRP embeddings (unit norm) + K-Means; - `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). @@ -110,3 +112,7 @@ TBD ### Subgraphs - **Maximal Independent Set**: _Ghaffari, Mohsen. "An improved distributed algorithm for maximal independent set." Proceedings of the twenty-seventh annual ACM-SIAM symposium on Discrete algorithms. Society for Industrial and Applied Mathematics, 2016._ + +### Embeddings + +- **FastRP**: _Chen, Haochen, et al. "Fast and accurate network embeddings via very sparse random projection." Proceedings of the 28th ACM international conference on information and knowledge management. 2019._ diff --git a/src/algorithm.rs b/src/algorithm.rs index 20d3e35..58d2338 100644 --- a/src/algorithm.rs +++ b/src/algorithm.rs @@ -1,5 +1,6 @@ mod centrality; pub(crate) mod community; mod connectivity; +pub(crate) mod embeddings; mod pregel; mod subgraph; diff --git a/src/algorithm/community.rs b/src/algorithm/community.rs index 9dda77f..25baa54 100644 --- a/src/algorithm/community.rs +++ b/src/algorithm/community.rs @@ -1,2 +1,3 @@ mod classical_lp; +pub(crate) mod fastrp_clustering; pub(crate) mod power_iteration_clustering; diff --git a/src/algorithm/community/fastrp_clustering.rs b/src/algorithm/community/fastrp_clustering.rs new file mode 100644 index 0000000..54d71e0 --- /dev/null +++ b/src/algorithm/community/fastrp_clustering.rs @@ -0,0 +1,492 @@ +//! Graph clustering via FastRP embeddings + K-Means. +//! +//! A thin pipeline on top of the two existing building blocks: build +//! unit-length FastRP embeddings ([`crate::algorithm::embeddings::fastrp`], +//! output normalization forced on) and cluster them with K-Means +//! ([`crate::ml::KMeansBuilder`], k-means|| init, one column per requested +//! `k`). +//! +//! The output parquet contains `[id, embedding, center_{k}...]`: the unit +//! norm embedding plus the assigned center index for every requested `k`. + +use datafusion::{ + dataframe::DataFrameWriteOptions, + error::{DataFusionError, Result}, + execution::object_store::ObjectStoreUrl, + object_store::path::Path, + prelude::*, +}; +use futures::{StreamExt, TryStreamExt}; + +use crate::{ + GraphFrame, VERTEX_ID, + algorithm::embeddings::fastrp::{EMBEDDING, FastRPNormalization}, + expressions::kmeans_assign_expr, + memory::CheckpointConfig, + ml::{DistanceMetric, KMeansBuilder, KMeansResult}, + utils::{GraphFramesConfig, scoped_ctx}, +}; + +/// Builder for the FastRP + K-Means clustering pipeline. +pub struct FastRPClusteringBuilder { + graph: GraphFrame, + dim: usize, + iterations: usize, + seed: u64, + normalization: FastRPNormalization, + iteration_weights: Option>, + k: Vec, + metric: DistanceMetric, + max_iter: usize, + tol: f64, + kmeans_init_steps: usize, + checkpoint_config: CheckpointConfig, +} + +impl FastRPClusteringBuilder { + pub fn new(graph: GraphFrame) -> Self { + FastRPClusteringBuilder { + graph: graph, + dim: 0, + iterations: 4, + seed: 42, + normalization: FastRPNormalization::None, + iteration_weights: None, + k: vec![2], + metric: DistanceMetric::L2, + max_iter: 20, + tol: 1e-4, + kmeans_init_steps: 2, + checkpoint_config: CheckpointConfig::default_local_fs(), + } + } + + /// Embedding dimension `D`. + pub fn set_dim(mut self, v: usize) -> Self { + self.dim = v; + self + } + + /// FastRP propagation iterations `K` (default: 4). + pub fn set_iterations(mut self, v: usize) -> Self { + self.iterations = v; + self + } + + /// Seed for the random projections and the k-means|| init (default: 42). + pub fn set_seed(mut self, v: u64) -> Self { + self.seed = v; + self + } + + /// Per-iteration normalization of the propagated vectors (default: + /// [`FastRPNormalization::None`]). + pub fn set_normalization(mut self, v: FastRPNormalization) -> Self { + self.normalization = v; + self + } + + /// Per-iterate weights of the FastRP linear combination + /// `H = Σ w_t · H_t` over `H_1..H_K` (must have exactly `K` entries; + /// default: all ones). The random init is never part of the combination. + pub fn set_iteration_weights(mut self, v: Vec) -> Self { + self.iteration_weights = Some(v); + self + } + + /// Set a single `k`: number of clusters. + pub fn set_k(mut self, k: usize) -> Self { + self.k = vec![k]; + self + } + + /// Set multiple k values: one `center_{k}` column per K in the output. + pub fn set_multiple_k(mut self, kk: Vec) -> Self { + self.k = kk; + self + } + + /// K-Means distance metric (default: L2; embeddings are unit norm, so + /// L2 on the sphere ranks neighbors like cosine would). + pub fn set_metric(mut self, v: DistanceMetric) -> Self { + self.metric = v; + self + } + + /// Maximum Lloyd iterations (default: 20). + pub fn set_max_iter(mut self, v: usize) -> Self { + self.max_iter = v; + self + } + + /// Convergence tolerance on the center shift (default: 1e-4). + pub fn set_tol(mut self, v: f64) -> Self { + self.tol = v; + self + } + + /// k-means|| initialization steps (default: 2). + pub fn set_kmeans_init_steps(mut self, v: usize) -> Self { + self.kmeans_init_steps = v; + 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 + } + + /// Run the pipeline and write `[id, embedding, center_{k}...]` parquet + /// to `output`. Returns the [`KMeansResult`] of the K-Means stage. + pub async fn run(self, ctx: &SessionContext, output: &str) -> Result { + if self.dim == 0 { + return Err(DataFusionError::Plan( + "FastRP clustering requires a positive embedding dimension: set `set_dim(D)`" + .to_string(), + )); + } + if self.k.is_empty() { + return Err(DataFusionError::Plan( + "FastRP clustering requires at least one k value: set `set_k(k)`".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::Uuid::new_v4().to_string(); + log::info!("start FastRP clustering with ID {run_id}"); + + let run_dir = self.checkpoint_config.dir.clone().join(run_id.clone()); + let embeddings_dir = run_dir.clone().join("embeddings"); + let embeddings_uri = format!( + "{}{}/", + self.checkpoint_config.store_url.as_str(), + embeddings_dir + ); + + // Stage 1: unit-length FastRP embeddings. + self.graph + .fastrp() + .dim(self.dim) + .iterations(self.iterations) + .seed(self.seed) + .normalization(self.normalization) + .iteration_weights(self.iteration_weights.clone().unwrap_or_else(|| { + // mirror the FastRP default: all ones over H_1..H_K + vec![1.0; self.iterations] + })) + .norm_output(true) + .set_checkpoint_dir(run_dir.join("fastrp_checkpoints")) + .with_checkpoint_store(self.checkpoint_config.store_url.clone()) + .run(&ctx, &embeddings_uri) + .await?; + + // Stage 2: K-Means over the embeddings. + let raw = ctx + .read_parquet(&embeddings_uri, ParquetReadOptions::default()) + .await?; + let features = raw.select(vec![col(VERTEX_ID), col(EMBEDDING)])?; + + let kmeans = KMeansBuilder::new(&features, EMBEDDING) + .seed(self.seed) + .k_values(&self.k) + .metric(self.metric) + .max_iter(self.max_iter) + .tol(self.tol) + .init_steps(self.kmeans_init_steps); + + log::info!("run KMeans on the FastRP embedding..."); + let result = kmeans.run().await?; + log::info!( + "KMeans converged after {} iterations", + result.num_iterations + ); + + // Final table: keep the embedding (handy for inspection) and assign + // the center index per requested k. + let mut final_columns = vec![col(VERTEX_ID), col(EMBEDDING)]; + for (kk, run) in self.k.iter().zip(&result.runs) { + final_columns.push( + kmeans_assign_expr( + col(EMBEDDING), + run.k, // effective k + result.d, + run.centers.clone(), + self.metric, + ) + .alias(format!("center_{kk}")), + ); + } + + features + .select(final_columns)? + .write_parquet(output, DataFrameWriteOptions::new(), None) + .await?; + log::info!("result was written into {output}"); + + // clean up the intermediate embeddings + let store = ctx + .runtime_env() + .object_store(&self.checkpoint_config.store_url)?; + let paths = store + .list(Some(&embeddings_dir)) + .map_ok(|m| m.location) + .boxed(); + store.delete_stream(paths).try_collect::>().await?; + + Ok(result) + } +} + +impl GraphFrame { + /// Create a new FastRP + K-Means clustering builder. + pub fn fastrp_clustering(&self) -> FastRPClusteringBuilder { + FastRPClusteringBuilder::new(self.clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + use std::fs; + use std::path::PathBuf; + use std::process::id; + use std::sync::atomic::{AtomicU64, Ordering}; + + use crate::utils::symmetrize; + use crate::{EDGE_DST, EDGE_SRC, VERTEX_ID}; + + static COUNTER: AtomicU64 = AtomicU64::new(0); + + fn unique_temp_dir(label: &str) -> PathBuf { + let n = COUNTER.fetch_add(1, Ordering::SeqCst); + let dir = + std::env::temp_dir().join(format!("graphframes_fastrp_clu_test_{}_{n}_{label}", id())); + fs::create_dir_all(&dir).expect("failed to create unique temp dir"); + dir + } + + struct TempGuard(PathBuf); + + impl Drop for TempGuard { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } + } + + fn setup(label: &str) -> (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("checkpoint dir"); + fs::create_dir_all(&output_root).expect("output dir"); + + let checkpoint_dir = Path::from_filesystem_path(&checkpoint_root).unwrap(); + let output_uri = url::Url::from_directory_path(&output_root) + .unwrap() + .to_string(); + + ( + SessionContext::new(), + checkpoint_dir, + output_uri, + TempGuard(parent), + ) + } + + /// Two 4-cliques {1..4}, {5..8} joined by the single bridge 4-5, + /// symmetrized. + fn two_cliques() -> Result { + let vertices = dataframe!(VERTEX_ID => Vec::::from(vec![1, 2, 3, 4, 5, 6, 7, 8]))?; + let mut edges = Vec::new(); + for (a, b) in [ + (1, 2), + (1, 3), + (2, 3), + (1, 4), + (2, 4), + (3, 4), // clique A + (5, 6), + (5, 7), + (6, 7), + (5, 8), + (6, 8), + (7, 8), // clique B + (4, 5), // bridge + ] { + edges.push(vec![a, b]); + } + let edges_df = dataframe!( + EDGE_SRC => Vec::::from(edges.iter().map(|e| e[0]).collect::>()), + EDGE_DST => Vec::::from(edges.iter().map(|e| e[1]).collect::>()), + )?; + let edges_df = symmetrize(&edges_df, true, None)?; + GraphFrame::try_new(vertices, edges_df) + } + + /// Read `[id, ]` into a map (column looked up by name). + async fn read_clusters(df: DataFrame, col_name: &str) -> Result> { + let mut map = HashMap::new(); + for batch in df.collect().await? { + let schema = batch.schema(); + let id_idx = schema.index_of(VERTEX_ID)?; + let c_idx = schema.index_of(col_name)?; + let ids = batch + .column(id_idx) + .as_any() + .downcast_ref::() + .unwrap(); + let clusters = batch + .column(c_idx) + .as_any() + .downcast_ref::() + .unwrap(); + for i in 0..batch.num_rows() { + map.insert(ids.value(i), clusters.value(i) as i64); + } + } + Ok(map) + } + + #[tokio::test] + async fn fastrp_clustering_separates_cliques() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("cliques"); + let g = two_cliques()?; + + g.fastrp_clustering() + .set_dim(8) + .set_iterations(3) + .set_seed(42) + .set_k(2) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let out = ctx + .read_parquet(&output_uri, ParquetReadOptions::default()) + .await?; + let clusters = read_clusters(out, "center_2").await?; + assert_eq!(clusters.len(), 8); + + // at most the (symmetrized) bridge may cross the cluster boundary + let mut crossing = 0; + for (a, b) in [ + (1, 2), + (1, 3), + (2, 3), + (1, 4), + (2, 4), + (3, 4), + (5, 6), + (5, 7), + (6, 7), + (5, 8), + (6, 8), + (7, 8), + (4, 5), + ] { + if clusters[&a] != clusters[&b] { + crossing += 1; + } + } + assert!( + crossing <= 2, + "the two cliques must land in different clusters, crossing={crossing}" + ); + Ok(()) + } + + #[tokio::test] + async fn fastrp_clustering_multiple_k_columns() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("multi_k"); + let g = two_cliques()?; + + g.fastrp_clustering() + .set_dim(8) + .set_iterations(2) + .set_seed(42) + .set_multiple_k(vec![2, 3]) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let out = ctx + .read_parquet(&output_uri, ParquetReadOptions::default()) + .await?; + let names: Vec = out + .schema() + .fields() + .iter() + .map(|f| f.name().to_string()) + .collect(); + assert!(names.contains(&"center_2".to_string()), "cols={names:?}"); + assert!(names.contains(&"center_3".to_string()), "cols={names:?}"); + assert!(names.contains(&"embedding".to_string()), "cols={names:?}"); + Ok(()) + } + + #[tokio::test] + async fn fastrp_clustering_is_deterministic_for_a_seed() -> Result<()> { + let run = || async { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("det"); + let g = two_cliques()?; + g.fastrp_clustering() + .set_dim(8) + .set_iterations(3) + .set_seed(5) + .set_k(2) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + let out = ctx + .read_parquet(&output_uri, ParquetReadOptions::default()) + .await?; + read_clusters(out, "center_2").await + }; + + let first = run().await?; + let second = run().await?; + assert_eq!(first, second, "same seed must reproduce the clustering"); + Ok(()) + } + + #[tokio::test] + async fn fastrp_clustering_requires_dim_and_k() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("validation"); + let g = two_cliques()?; + + let no_dim = g + .fastrp_clustering() + .set_checkpoint_dir(checkpoint_dir.clone()) + .run(&ctx, &output_uri) + .await; + assert!(no_dim.is_err(), "dim(0) must be rejected"); + + let no_k = g + .fastrp_clustering() + .set_dim(4) + .set_multiple_k(vec![]) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await; + assert!(no_k.is_err(), "empty k must be rejected"); + Ok(()) + } +} diff --git a/src/algorithm/embeddings.rs b/src/algorithm/embeddings.rs new file mode 100644 index 0000000..d00eee3 --- /dev/null +++ b/src/algorithm/embeddings.rs @@ -0,0 +1,3 @@ +//! Graph embedding algorithms. + +pub(crate) mod fastrp; diff --git a/src/algorithm/embeddings/fastrp.rs b/src/algorithm/embeddings/fastrp.rs new file mode 100644 index 0000000..56d209e --- /dev/null +++ b/src/algorithm/embeddings/fastrp.rs @@ -0,0 +1,1043 @@ +//! FastRP (Fast Random Projection) vertex embeddings. +//! +//! A simplified FastRP (Chen & King, CIKM 2019): every vertex starts from a +//! deterministic sparse random projection vector `H_0 ∈ {−1, 0, +1}^D` +//! (see `fastrp_init`), then `K` rounds of message propagation sum the +//! source vectors into every destination: +//! +//! ```text +//! H_{t+1}(dst) = Σ_{(src, dst) ∈ E} H_t(src) +//! ``` +//! +//! Normalization (per iteration): the propagated vector of each source may +//! be divided by its out-degree (linear, `L1`) or by the square +//! root of its out-degree. The out-degree is computed once and +//! attached to the edge table, so the loop needs no extra joins. +//! +//! Deviation from the paper, on purpose: no concatenation of iterates. The +//! final embedding is the weighted linear combination +//! +//! ```text +//! H = Σ_{t=1..K} w_t · H_t +//! ``` +//! +//! (`iteration_weights`, length `K`, default all ones). The random init +//! `H_0` always stays out of the combination — it exists only as the message +//! source of the first iteration. A zero weight drops the corresponding +//! iterate from the pipeline entirely (it is never joined at the end). +//! +//! Unreached vertices keep no state at all during the loop (the aggregate +//! only emits vertices that received a message, which is the sum identity) +//! and are filled with the zero vector in one final streaming pass over the +//! output. + +use datafusion::{ + arrow::datatypes::DataType, + dataframe::DataFrameWriteOptions, + error::{DataFusionError, Result}, + execution::object_store::ObjectStoreUrl, + functions::math::sqrt, + functions_aggregate::count::count, + object_store::path::Path, + prelude::*, +}; + +use crate::{ + EDGE_DST, EDGE_SRC, GraphFrame, VERTEX_ID, + expressions::{ + fastrp_init_expr, l2_norm_expr, vec_scale_expr, vec_sum_expr, vec_weighted_sum_expr, + vec_zero_scalar, + }, + memory::{CheckpointConfig, ParquetCheckpointer}, + utils::{GraphFramesConfig, scoped_ctx}, +}; + +/// Name of the embedding column produced by the algorithm. +pub(crate) const EMBEDDING: &str = "embedding"; + +/// Name of the per-source out-degree column (only present when a per-step +/// normalization is enabled; attached to the edge table so the loop needs +/// no extra joins). +const DEGREE: &str = "__fastrp_degree"; + +/// Name of the temporary L2-norm column of the output pass. +const NORM: &str = "__fastrp_norm"; + +/// Name of the running accumulator column of the combine fold. +const ACC: &str = "__fastrp_acc"; + +/// Name of the accumulator's id column during the combine joins. +const ACC_VID: &str = "__fastrp_acc_vid"; + +/// Per-iteration normalization of the propagated vectors, applied to every +/// message: each source vector is divided by a function of the source +/// out-degree before the group-by sum (`H_{t+1} = S_norm · H_t`). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FastRPNormalization { + /// Plain sum, no normalization (`S = A`). + None, + /// Linear normalization (`S = A · D⁻¹` in the paper): every propagated + /// vector is divided by the out-degree of its source. + L1, + /// Square normalization (`S = A · D^(-1/2)` in the paper): every + /// propagated vector is divided by the square root of the source + /// out-degree. + L2, +} + +/// Builder for FastRP embeddings. +/// +/// Required parameter: `dim` (`D`); `iterations` (`K`) defaults to 4. +pub struct FastRPBuilder<'a> { + graph: &'a GraphFrame, + dim: usize, + iterations: usize, + seed: u64, + normalization: FastRPNormalization, + norm_output: bool, + iteration_weights: Option>, + checkpoint_config: CheckpointConfig, +} + +impl<'a> FastRPBuilder<'a> { + pub fn new(graph: &'a GraphFrame) -> Self { + FastRPBuilder { + graph: graph, + dim: 0, // must be set explicitly + iterations: 4, // a few hops mix the random projections well + seed: 42, + normalization: FastRPNormalization::None, + norm_output: false, + iteration_weights: None, + checkpoint_config: CheckpointConfig::default_local_fs(), + } + } + + /// Embedding dimension `D` (required). + pub fn dim(mut self, v: usize) -> Self { + self.dim = v; + self + } + + /// Number of propagation iterations `K` (default: 4). + pub fn iterations(mut self, v: usize) -> Self { + self.iterations = v; + self + } + + /// Seed for the per-vertex random projections. The init is a pure + /// function of `(seed, id)`, so runs are reproducible. + pub fn seed(mut self, v: u64) -> Self { + self.seed = v; + self + } + + /// Per-iteration normalization of the propagated vectors (default: + /// [`FastRPNormalization::None`]). + pub fn normalization(mut self, v: FastRPNormalization) -> Self { + self.normalization = v; + self + } + + /// L2-normalize the final embeddings to unit length (default: `false`). + /// Zero vectors stay zero, so downstream cosine/K-Means consumers never + /// see NaN. + pub fn norm_output(mut self, v: bool) -> Self { + self.norm_output = v; + self + } + + /// Per-iterate weights of the final linear combination + /// `H = Σ w_t · H_t` over `H_1..H_K`. Must have exactly `K` entries + /// (default: all ones, i.e. sum every iterate). The random init `H_0` + /// never enters the combination. Iterate `t` with `w_t = 0` is dropped + /// from the pipeline entirely (not joined at the end). + pub fn iteration_weights(mut self, v: Vec) -> Self { + self.iteration_weights = Some(v); + 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 + } + + /// Run FastRP and write `[VERTEX_ID, EMBEDDING]` parquet to `output`. + /// + /// The embedding column is a non-null `f32` vector column of length `D` + /// (`FixedSizeList`, surfaced as `List` by a parquet + /// round-trip — all vector consumers in this crate accept both). + pub async fn run(self, ctx: &SessionContext, output: &str) -> Result { + if self.dim == 0 { + return Err(DataFusionError::Plan( + "FastRP requires a positive embedding dimension: set `dim(D)`".to_string(), + )); + } + let weights = match &self.iteration_weights { + Some(w) => w.clone(), + None => vec![1.0; self.iterations], + }; + if weights.len() != self.iterations { + return Err(DataFusionError::Plan(format!( + "iteration_weights must have exactly {} entries (one per iteration, H_1..H_{}), got {}", + self.iterations, + self.iterations, + weights.len() + ))); + } + + 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::Uuid::new_v4().to_string(); + log::info!("start FastRP with ID {run_id}"); + + // Only the id (and the derived embedding) should travel through the + // iterations; attribute columns would inflate every shuffle. + let vertices = self.graph.vertices.clone().select_columns(&[VERTEX_ID])?; + let mut edges = self + .graph + .edges + .clone() + .select_columns(&[EDGE_SRC, EDGE_DST])?; + + // Per-step normalization scales every propagated vector by a factor + // derived from the source out-degree. The degree is computed once and + // attached to the edge table (every edge has a source, so the inner + // join is exact and no degree-zero division can happen). + if self.normalization != FastRPNormalization::None { + let degree_vid = "__fastrp_degree_vid"; + let degrees = edges.clone().aggregate( + vec![col(EDGE_SRC).alias(degree_vid)], + vec![count(col(EDGE_DST)).alias(DEGREE)], + )?; + edges = edges + .join_on( + degrees, + JoinType::Inner, + vec![col(EDGE_SRC).eq(col(degree_vid))], + )? + .select(vec![col(EDGE_SRC), col(EDGE_DST), col(DEGREE)])?; + } + + // Pre-sorted, co-partitioned checkpoints: the edges by `src`, every + // iterate by `id` — the per-iteration join and the final K-way + // combine are sort-merge joins without an extra sort. + let mut edges_checkpointer = ParquetCheckpointer::new( + self.checkpoint_config.store_url.clone(), + self.checkpoint_config + .dir + .clone() + .join(run_id.clone()) + .join("edges"), + ); + let mut states_checkpointer = ParquetCheckpointer::new( + self.checkpoint_config.store_url.clone(), + self.checkpoint_config + .dir + .clone() + .join(run_id.clone()) + .join("states"), + ); + let mut vertices_checkpointer = ParquetCheckpointer::new( + self.checkpoint_config.store_url.clone(), + self.checkpoint_config + .dir + .clone() + .join(run_id.clone()) + .join("vertices"), + ); + + let edges = edges_checkpointer + .push_pre_sorted(&ctx, "edges", edges, EDGE_SRC) + .await?; + + // H_0: the deterministic sparse random projections (message source of + // iteration 1; never part of the final combination). + let init = vertices.clone().select(vec![ + col(VERTEX_ID), + fastrp_init_expr(col(VERTEX_ID), self.dim, self.seed).alias(EMBEDDING), + ])?; + let mut state = states_checkpointer + .push_pre_sorted(&ctx, "state-0", init, VERTEX_ID) + .await?; + + // Per-step normalization scales every propagated vector by a factor + // derived from the source out-degree (which rides on the edge table). + let message = match self.normalization { + FastRPNormalization::None => col(EMBEDDING), + FastRPNormalization::L1 => vec_scale_expr( + col(EMBEDDING), + lit(1.0f64) / cast(col(DEGREE), DataType::Float64), + ), + FastRPNormalization::L2 => vec_scale_expr( + col(EMBEDDING), + lit(1.0f64) / sqrt().call(vec![cast(col(DEGREE), DataType::Float64)]), + ), + }; + + // Propagation: iterate `t` sums the vectors of all sources reached by + // `t - 1` hops. The aggregate emits only vertices that received a + // message (a missing state row *is* the zero vector), so no null + // handling is needed inside the loop. + let mut states: Vec = Vec::with_capacity(self.iterations); + for t in 1..=self.iterations { + let triplets = edges.clone().join_on( + state.clone(), + JoinType::Inner, + vec![col(EDGE_SRC).eq(col(VERTEX_ID))], + )?; + let messages = triplets.select(vec![ + col(EDGE_DST).alias(VERTEX_ID), + message.clone().alias(EMBEDDING), + ])?; + let aggregated = messages.aggregate( + vec![col(VERTEX_ID)], + vec![vec_sum_expr(col(EMBEDDING), self.dim).alias(EMBEDDING)], + )?; + state = states_checkpointer + .push_pre_sorted(&ctx, &format!("state-{t}"), aggregated, VERTEX_ID) + .await?; + states.push(state.clone()); + } + + // Final combination: `H = Σ w_t · H_t` over the iterates with a + // non-zero weight, computed as a *sequential checkpointed fold*: + // + // acc_0 = w_{t0} · H_{t0} + // acc_{i+1} = acc_i + w_t · H_t (checkpointed every step) + // + // At any moment only two iterate columns (the running accumulator and + // the iterate being folded in) are in memory; every intermediate + // accumulator is offloaded to a pre-sorted parquet checkpoint, so the + // combine never materializes the full `K · |V| · D` history. + let base = vertices_checkpointer + .push_pre_sorted(&ctx, "vertices", vertices, VERTEX_ID) + .await?; + + let mut folds_checkpointer = ParquetCheckpointer::new( + self.checkpoint_config.store_url.clone(), + self.checkpoint_config + .dir + .clone() + .join(run_id.clone()) + .join("folds"), + ); + + // iterates that actually participate (non-zero weight), in order + let kept: Vec<(usize, f64)> = weights + .iter() + .enumerate() + .filter(|(_, w)| **w != 0.0) + .map(|(t, w)| (t, *w)) + .collect(); + + let mut fixed = if self.iterations == 0 { + // No propagation at all: the embedding is the raw random init. + state.select(vec![col(VERTEX_ID), col(EMBEDDING)])? + } else if kept.is_empty() { + // Degenerate but valid: all weights are zero. + let fsl_type = DataType::FixedSizeList( + datafusion::arrow::datatypes::Field::new("el", DataType::Float32, false).into(), + self.dim as i32, + ); + let zero = vec_zero_scalar(&fsl_type, self.dim)?; + base.select(vec![col(VERTEX_ID), lit(zero).alias(EMBEDDING)])? + } else { + let zero = vec_zero_scalar( + &DataType::FixedSizeList( + datafusion::arrow::datatypes::Field::new("el", DataType::Float32, false).into(), + self.dim as i32, + ), + self.dim, + )?; + + // Sequential checkpointed fold over the participating iterates: + // + // acc = w_{t0} · H_{t0} + w_{t1} · H_{t1} + // acc = acc + w_t · H_t (one join per iterate, + // checkpointed every step) + // + // Only two iterate columns are ever in memory (the running + // accumulator and the iterate being folded in); every + // intermediate accumulator is offloaded to a pre-sorted parquet + // checkpoint, and all keys are covered by co-partitioned SMJs. + let (t0, w0) = kept[0]; + let mut acc = if kept.len() >= 2 { + let (t1, w1) = kept[1]; + let side_a = states[t0].clone().select(vec![ + col(VERTEX_ID).alias(ACC_VID), + col(EMBEDDING).alias("__fastrp_ha"), + ])?; + let side_b = states[t1].clone().select(vec![ + col(VERTEX_ID).alias("__fastrp_hb_vid"), + col(EMBEDDING).alias("__fastrp_hb"), + ])?; + side_a + .join_on( + side_b, + // full outer: a vertex may be present in one iterate + // and absent (zero contribution) in the other + JoinType::Full, + vec![col(ACC_VID).eq(col("__fastrp_hb_vid"))], + )? + .select(vec![ + coalesce(vec![col(ACC_VID), col("__fastrp_hb_vid")]).alias(ACC_VID), + vec_weighted_sum_expr( + &[(lit(w0), col("__fastrp_ha")), (lit(w1), col("__fastrp_hb"))], + self.dim, + ) + .alias(ACC), + ])? + } else { + // single kept iterate: it is the accumulator (scaled) + states[t0].clone().select(vec![ + col(VERTEX_ID).alias(ACC_VID), + vec_scale_expr(col(EMBEDDING), lit(w0)).alias(ACC), + ])? + }; + folds_checkpointer + .push_pre_sorted(&ctx, "fold-0", acc.clone(), ACC_VID) + .await?; + + // remaining iterates fold in one at a time: acc = acc + w_t · H_t. + // The full outer join grows the key set as iterates reach new + // vertices (an absent iterate row is a zero contribution). + for (t, w) in kept.iter().skip(2) { + let side = states[*t].clone().select(vec![ + col(VERTEX_ID).alias("__fastrp_ht_vid"), + col(EMBEDDING).alias("__fastrp_ht"), + ])?; + acc = acc + .join_on( + side, + // full outer, same reason as the first fold step + JoinType::Full, + vec![col(ACC_VID).eq(col("__fastrp_ht_vid"))], + )? + .select(vec![ + coalesce(vec![col(ACC_VID), col("__fastrp_ht_vid")]).alias(ACC_VID), + vec_weighted_sum_expr( + &[(lit(1.0), col(ACC)), (lit(*w), col("__fastrp_ht"))], + self.dim, + ) + .alias(ACC), + ])?; + acc = folds_checkpointer + .push_pre_sorted(&ctx, "fold", acc.clone(), ACC_VID) + .await?; + // only the latest accumulator checkpoint is needed + folds_checkpointer.evict_all_but_latest_n(&ctx, 1).await?; + } + + // Left-join the accumulator onto the full vertex set: vertices + // that were never reached have no accumulator row, and a missing + // row is the zero vector. + let embedding = when(col(ACC).is_null(), lit(zero)) + .otherwise(col(ACC))? + .alias(EMBEDDING); + base.join_on(acc, JoinType::Left, vec![col(VERTEX_ID).eq(col(ACC_VID))])? + .select(vec![col(VERTEX_ID), embedding])? + }; + + // Optional: L2-normalize the final embeddings to unit length. + // Zero vectors (norm == 0) stay zero, so the pass can not produce + // NaN/Inf and downstream cosine/K-Means consumers are safe. + if self.norm_output { + fixed = fixed + .with_column(NORM, l2_norm_expr(col(EMBEDDING)))? + .select(vec![ + col(VERTEX_ID), + when( + col(NORM).not_eq(lit(0.0f64)), + vec_scale_expr(col(EMBEDDING), lit(1.0f64) / col(NORM)), + ) + .otherwise(col(EMBEDDING))? + .alias(EMBEDDING), + ])?; + } + + fixed + .write_parquet(output, DataFrameWriteOptions::new(), None) + .await?; + log::info!("result was written into {output}"); + + // clean up: all iterate checkpoints live under the run directory + edges_checkpointer.purge(&ctx).await?; + states_checkpointer.purge(&ctx).await?; + vertices_checkpointer.purge(&ctx).await?; + + log::info!( + "FastRP {run_id} finished after {} iterations, output: {output}", + self.iterations + ); + Ok(self.iterations) + } +} + +impl GraphFrame { + /// Create a new FastRP algorithm builder + pub fn fastrp(&self) -> FastRPBuilder<'_> { + FastRPBuilder::new(self) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + use std::fs; + use std::path::PathBuf; + use std::process::id; + use std::sync::atomic::{AtomicU64, Ordering}; + + use crate::expressions::as_f32_list_like; + use crate::ml::fastrp_init_fill; + use crate::utils::symmetrize; + + use datafusion::arrow::array::Int64Array; + + static COUNTER: AtomicU64 = AtomicU64::new(0); + + /// Unique temp dir per test (same pattern as the pregel 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_fastrp_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. + struct TempGuard(PathBuf); + + impl Drop for TempGuard { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.0); + } + } + + fn setup(label: &str) -> (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("checkpoint dir"); + fs::create_dir_all(&output_root).expect("output dir"); + + let checkpoint_dir = Path::from_filesystem_path(&checkpoint_root).unwrap(); + let output_uri = url::Url::from_directory_path(&output_root) + .unwrap() + .to_string(); + + ( + SessionContext::new(), + checkpoint_dir, + output_uri, + TempGuard(parent), + ) + } + + fn create_graph(vertices: Vec, edges: Vec>) -> Result { + let vertices_df = dataframe!(VERTEX_ID => Vec::::from(vertices))?; + let edges_df = dataframe!( + EDGE_SRC => Vec::::from(edges.iter().map(|e| e[0]).collect::>()), + EDGE_DST => Vec::::from(edges.iter().map(|e| e[1]).collect::>()), + )?; + GraphFrame::try_new(vertices_df, edges_df) + } + + /// Read `[id, embedding]` parquet into a map (accepts both vector types). + async fn embeddings_map(df: DataFrame) -> Result>> { + let mut map = HashMap::new(); + for batch in df.collect().await? { + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let vecs = as_f32_list_like(batch.column(1), "test", "embedding")?; + for i in 0..batch.num_rows() { + map.insert(ids.value(i), vecs.value(i).to_vec()); + } + } + Ok(map) + } + + /// The init kernel reference for one vertex. + fn init_ref(id: i64, seed: u64, d: usize) -> Vec { + let mut v = vec![0.0f32; d]; + fastrp_init_fill(id, seed, d, &mut v); + v + } + + #[tokio::test] + async fn fastrp_requires_positive_dim() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("no_dim"); + let g = create_graph(vec![1, 2], vec![vec![1, 2]])?; + let result = g + .fastrp() + .iterations(1) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await; + assert!(result.is_err(), "dim(0) must be rejected"); + Ok(()) + } + + #[tokio::test] + async fn fastrp_k1_on_path_sums_neighbor_inits() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("k1_path"); + let d = 2; + let seed = 42u64; + // path: 1 -> 2 -> 3 (vertex 1 has no in-edges) + let g = create_graph(vec![1, 2, 3], vec![vec![1, 2], vec![2, 3]])?; + + g.fastrp() + .dim(d) + .iterations(1) + .seed(seed) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let got = embeddings_map( + ctx.read_parquet(&output_uri, ParquetReadOptions::default()) + .await?, + ) + .await?; + + assert_eq!(got.len(), 3); + assert_eq!(got[&1], vec![0.0, 0.0], "no in-edges -> zero vector"); + assert_eq!(got[&2], init_ref(1, seed, d)); + assert_eq!(got[&3], init_ref(2, seed, d)); + Ok(()) + } + + #[tokio::test] + async fn fastrp_k0_returns_random_init() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("k0"); + let d = 8; + let seed = 7u64; + let g = create_graph(vec![1, 2], vec![vec![1, 2]])?; + + let iterations = g + .fastrp() + .dim(d) + .iterations(0) + .seed(seed) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + assert_eq!(iterations, 0); + + let got = embeddings_map( + ctx.read_parquet(&output_uri, ParquetReadOptions::default()) + .await?, + ) + .await?; + assert_eq!(got[&1], init_ref(1, seed, d)); + assert_eq!(got[&2], init_ref(2, seed, d)); + Ok(()) + } + + #[tokio::test] + async fn fastrp_k2_propagates_two_hops() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("k2_path"); + let d = 4; + let seed = 9u64; + // path: 1 -> 2 -> 3 + let g = create_graph(vec![1, 2, 3], vec![vec![1, 2], vec![2, 3]])?; + + g.fastrp() + .dim(d) + .iterations(2) + .seed(seed) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let got = embeddings_map( + ctx.read_parquet(&output_uri, ParquetReadOptions::default()) + .await?, + ) + .await?; + // default weights [1, 1]: H(v) = H_1(v) + H_2(v) + // H_1 = {2: init(1), 3: init(2)}; H_2 = {3: init(1)} + assert_eq!( + got[&3], + add(&init_ref(2, seed, d), &init_ref(1, seed, d)), + "vertex 3 sums both iterates" + ); + assert_eq!(got[&2], init_ref(1, seed, d), "H_2(2) is absent -> zero"); + assert_eq!(got[&1], vec![0.0; d]); + Ok(()) + } + + #[tokio::test] + async fn fastrp_weights_can_drop_early_iterates() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("weights_drop"); + let d = 4; + let seed = 9u64; + // path: 1 -> 2 -> 3 + let g = create_graph(vec![1, 2, 3], vec![vec![1, 2], vec![2, 3]])?; + + // w = [0, 1]: only the last iterate survives — the classic + // "last iterate" embedding + g.fastrp() + .dim(d) + .iterations(2) + .seed(seed) + .iteration_weights(vec![0.0, 1.0]) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let got = embeddings_map( + ctx.read_parquet(&output_uri, ParquetReadOptions::default()) + .await?, + ) + .await?; + assert_eq!( + got[&3], + init_ref(1, seed, d), + "two hops reach vertex 1's init" + ); + assert_eq!(got[&2], vec![0.0; d]); + assert_eq!(got[&1], vec![0.0; d]); + Ok(()) + } + + #[tokio::test] + async fn fastrp_weights_scale_the_iterates() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("weights_scale"); + let d = 2; + let seed = 42u64; + // triangle 1 -> 2 -> 3 -> 1: every vertex has an exact expression + // H_1(v) = init(pred(v)), H_2(v) = init(pred(pred(v))) + let g = create_graph(vec![1, 2, 3], vec![vec![1, 2], vec![2, 3], vec![3, 1]])?; + + g.fastrp() + .dim(d) + .iterations(2) + .seed(seed) + .iteration_weights(vec![1.0, 2.0]) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let got = embeddings_map( + ctx.read_parquet(&output_uri, ParquetReadOptions::default()) + .await?, + ) + .await?; + let i1 = init_ref(1, seed, d); + let i2 = init_ref(2, seed, d); + let i3 = init_ref(3, seed, d); + // H_1(1)=init(3), H_1(2)=init(1), H_1(3)=init(2) + // edges into 1: 3->1, into 2: 1->2, into 3: 2->3 + // H_2(1)=H_1(3)=init(2), H_2(2)=H_1(1)=init(1), H_2(3)=H_1(2)=init(1) + // H(1) = 1*H_1(1) + 2*H_2(1) = init(3) + 2*init(2), exact in f32 + assert_eq!(got[&1], add_scaled(&i3, 1.0, &i2, 2.0)); + assert_eq!(got[&2], add_scaled(&i1, 1.0, &i3, 2.0)); + assert_eq!(got[&3], add_scaled(&i2, 1.0, &i1, 2.0)); + Ok(()) + } + + #[tokio::test] + async fn fastrp_combine_folds_match_reference() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("fold3"); + let d = 2; + let seed = 42u64; + // triangle: H_t(v) = init(pred^t(v)) — exact expressions per vertex + let g = create_graph(vec![1, 2, 3], vec![vec![1, 2], vec![2, 3], vec![3, 1]])?; + + g.fastrp() + .dim(d) + .iterations(3) + .seed(seed) + .iteration_weights(vec![0.5, 1.0, 2.0]) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let got = embeddings_map( + ctx.read_parquet(&output_uri, ParquetReadOptions::default()) + .await?, + ) + .await?; + let i1 = init_ref(1, seed, d); + let i2 = init_ref(2, seed, d); + let i3 = init_ref(3, seed, d); + + // H_1(1)=i3, H_2(1)=i2, H_3(1)=i1 => H(1) = 0.5*i3 + 1*i2 + 2*i1 + assert_eq!( + got[&1], + add_scaled(&add_scaled(&i3, 0.5, &i2, 1.0), 1.0, &i1, 2.0) + ); + // H_1(2)=i1, H_2(2)=i3, H_3(2)=i2 => H(2) = 0.5*i1 + 1*i3 + 2*i2 + assert_eq!( + got[&2], + add_scaled(&add_scaled(&i1, 0.5, &i3, 1.0), 1.0, &i2, 2.0) + ); + // H_1(3)=i2, H_2(3)=i1, H_3(3)=i3 => H(3) = 0.5*i2 + 1*i1 + 2*i3 + assert_eq!( + got[&3], + add_scaled(&add_scaled(&i2, 0.5, &i1, 1.0), 1.0, &i3, 2.0) + ); + Ok(()) + } + + #[tokio::test] + async fn fastrp_weights_length_is_validated() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("weights_len"); + let g = create_graph(vec![1, 2], vec![vec![1, 2]])?; + let result = g + .fastrp() + .dim(4) + .iterations(2) + .iteration_weights(vec![1.0]) // 1 entry for K=2 + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await; + assert!(result.is_err(), "weight length must match iterations"); + Ok(()) + } + + /// `a*w_a + b*w_b`, exact in f32 for these integral values. + fn add_scaled(a: &[f32], w_a: f32, b: &[f32], w_b: f32) -> Vec { + a.iter().zip(b).map(|(x, y)| x * w_a + y * w_b).collect() + } + + fn add(a: &[f32], b: &[f32]) -> Vec { + a.iter().zip(b).map(|(x, y)| x + y).collect() + } + + #[tokio::test] + async fn fastrp_symmetrized_graph_is_symmetric() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("undirected"); + let d = 3; + let seed = 5u64; + let base = create_graph(vec![1, 2], vec![vec![1, 2]])?; + let edges = symmetrize(&base.edges, true, None)?; + let g = GraphFrame { + vertices: base.vertices, + edges, + }; + + g.fastrp() + .dim(d) + .iterations(1) + .seed(seed) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let got = embeddings_map( + ctx.read_parquet(&output_uri, ParquetReadOptions::default()) + .await?, + ) + .await?; + assert_eq!(got[&1], init_ref(2, seed, d)); + assert_eq!(got[&2], init_ref(1, seed, d)); + Ok(()) + } + + #[tokio::test] + async fn fastrp_fills_isolated_vertices_with_zero() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("isolated"); + let d = 2; + let seed = 1u64; + // vertex 3 has no edges at all; vertex 1 has no in-edges + let g = create_graph(vec![1, 2, 3], vec![vec![2, 1]])?; + + g.fastrp() + .dim(d) + .iterations(1) + .seed(seed) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let got = embeddings_map( + ctx.read_parquet(&output_uri, ParquetReadOptions::default()) + .await?, + ) + .await?; + assert_eq!(got[&3], vec![0.0; d]); + assert_eq!(got[&2], vec![0.0; d]); + assert_eq!(got[&1], init_ref(2, seed, d)); + Ok(()) + } + + #[tokio::test] + async fn fastrp_is_reproducible_for_a_seed() -> Result<()> { + let d = 6; + let seed = 77u64; + + let build = || async { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("repro"); + let g = create_graph( + vec![1, 2, 3, 4], + vec![vec![1, 2], vec![2, 3], vec![3, 4], vec![1, 4]], + )?; + g.fastrp() + .dim(d) + .iterations(2) + .seed(seed) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + embeddings_map( + ctx.read_parquet(&output_uri, ParquetReadOptions::default()) + .await?, + ) + .await + }; + + let first = build().await?; + let second = build().await?; + assert_eq!(first, second, "same seed must reproduce the embeddings"); + Ok(()) + } + + #[tokio::test] + async fn fastrp_l1_divides_by_source_degree() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("l1"); + let d = 2; + let seed = 42u64; + // star: 1 -> 2, 1 -> 3; out-degree of 1 is 2 + let g = create_graph(vec![1, 2, 3], vec![vec![1, 2], vec![1, 3]])?; + + g.fastrp() + .dim(d) + .iterations(1) + .seed(seed) + .normalization(FastRPNormalization::L1) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let got = embeddings_map( + ctx.read_parquet(&output_uri, ParquetReadOptions::default()) + .await?, + ) + .await?; + // H(2) = H(3) = init(1) / deg(1), with the exact f32 factor + let s = (1.0f64 / 2.0f64) as f32; + let expected: Vec = init_ref(1, seed, d).iter().map(|x| x * s).collect(); + assert_eq!(got[&2], expected); + assert_eq!(got[&3], expected); + assert_eq!(got[&1], vec![0.0; d]); + Ok(()) + } + + #[tokio::test] + async fn fastrp_l2_divides_by_sqrt_source_degree() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("l2"); + let d = 2; + let seed = 42u64; + let g = create_graph(vec![1, 2, 3], vec![vec![1, 2], vec![1, 3]])?; + + g.fastrp() + .dim(d) + .iterations(1) + .seed(seed) + .normalization(FastRPNormalization::L2) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let got = embeddings_map( + ctx.read_parquet(&output_uri, ParquetReadOptions::default()) + .await?, + ) + .await?; + // the expression computes 1 / sqrt(deg) in f64, then scales in f32 + let s = (1.0f64 / (2.0f64).sqrt()) as f32; + let expected: Vec = init_ref(1, seed, d).iter().map(|x| x * s).collect(); + assert_eq!(got[&2], expected); + assert_eq!(got[&3], expected); + Ok(()) + } + + #[tokio::test] + async fn fastrp_norm_output_gives_unit_norms() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("norm_output"); + let d = 4; + let seed = 11u64; + let g = create_graph( + vec![1, 2, 3], + vec![vec![1, 2], vec![2, 3], vec![3, 1], vec![1, 3]], + )?; + + g.fastrp() + .dim(d) + .iterations(2) + .seed(seed) + .norm_output(true) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let got = embeddings_map( + ctx.read_parquet(&output_uri, ParquetReadOptions::default()) + .await?, + ) + .await?; + assert_eq!(got.len(), 3); + for (id, v) in &got { + let norm_sq: f32 = v.iter().map(|x| x * x).sum(); + let norm = norm_sq.sqrt(); + assert!( + (norm - 1.0).abs() < 1e-4 || norm == 0.0, + "vertex {id}: expected unit norm (or zero), got {norm}" + ); + } + Ok(()) + } + + #[tokio::test] + async fn fastrp_norm_output_zero_vectors_stay_zero() -> Result<()> { + let (ctx, checkpoint_dir, output_uri, _guard) = setup("norm_output_zero"); + let d = 3; + let seed = 3u64; + // vertex 1 has no in-edges -> zero embedding after K=1 + let g = create_graph(vec![1, 2], vec![vec![1, 2]])?; + + g.fastrp() + .dim(d) + .iterations(1) + .seed(seed) + .norm_output(true) + .set_checkpoint_dir(checkpoint_dir) + .run(&ctx, &output_uri) + .await?; + + let got = embeddings_map( + ctx.read_parquet(&output_uri, ParquetReadOptions::default()) + .await?, + ) + .await?; + assert_eq!(got[&1], vec![0.0; d], "zero vector must stay exactly zero"); + // and vertex 2 is unit-norm + let norm: f32 = got[&2].iter().map(|x| x * x).sum::().sqrt(); + assert!((norm - 1.0).abs() < 1e-4); + Ok(()) + } +} diff --git a/src/expressions.rs b/src/expressions.rs index 90344d6..a90ae3e 100644 --- a/src/expressions.rs +++ b/src/expressions.rs @@ -1,3 +1,4 @@ +#![allow(dead_code, unused)] // one day I will remove this mod common; mod finite_axpb; mod hll; @@ -14,5 +15,8 @@ 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 linalg::{ + cosine_distance_expr, fastrp_init_expr, l2_distance_expr, l2_norm_expr, vec_scale_expr, + vec_sum_expr, vec_weighted_sum_expr, vec_zero_scalar, +}; pub(crate) use most_common::most_common_expr; diff --git a/src/expressions/common.rs b/src/expressions/common.rs index 5b87cae..d58b701 100644 --- a/src/expressions/common.rs +++ b/src/expressions/common.rs @@ -93,6 +93,22 @@ impl<'a> F32ListLike<'a> { } } + /// Length of a vector row. Only meaningful for a non-empty array: for a + /// `List` it reports row `0`'s length (rows may legally differ, which is + /// the mismatch the vector UDFs reject). + pub(crate) fn value_length(&self) -> usize { + match self { + F32ListLike::Fixed(a) => a.value_length() as usize, + F32ListLike::View(a) => { + if a.len() > 0 { + a.value(0).len() + } else { + 0 + } + } + } + } + pub(crate) fn null_count(&self) -> usize { match self { F32ListLike::Fixed(a) => a.null_count(), diff --git a/src/expressions/linalg.rs b/src/expressions/linalg.rs index 8eaabde..7e780b8 100644 --- a/src/expressions/linalg.rs +++ b/src/expressions/linalg.rs @@ -1,5 +1,6 @@ -//! Scalar UDFs over `f32` vectors: `l2_norm`, `l2_distance`, -//! `cosine_distance`. +//! Scalar UDFs and vector aggregates over `f32` vectors: `l2_norm`, +//! `l2_distance`, `cosine_distance`, the `vec_sum` aggregate and the +//! FastRP `fastrp_init` random-projection init. //! //! This module holds the DataFusion wrappers only: SIMD-logic free; the //! kernels live in [`crate::ml::linalg`]. Contract: @@ -8,16 +9,25 @@ use std::sync::Arc; -use datafusion::arrow::array::{ArrayRef, Float64Array}; -use datafusion::arrow::datatypes::DataType; +use datafusion::arrow::array::{ + Array, ArrayRef, BooleanArray, FixedSizeListArray, Float32Array, Float64Array, ListArray, +}; +use datafusion::arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; +use datafusion::arrow::datatypes::{DataType, Field}; use datafusion::common::plan_err; use datafusion::error::Result; +use datafusion::logical_expr::function::AccumulatorArgs; +use datafusion::logical_expr::groups_accumulator::{EmitTo, GroupsAccumulator}; use datafusion::logical_expr::{ - ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, + Accumulator, AggregateUDF, AggregateUDFImpl, ColumnarValue, Expr, ScalarFunctionArgs, + ScalarUDF, ScalarUDFImpl, Signature, Volatility, }; +use datafusion::scalar::ScalarValue; -use crate::expressions::common::as_f32_list_like; -use crate::ml::{cosine_distance, l2_distance, l2_norm}; +use crate::expressions::common::{as_f32_list_like, downcast_int64}; +use crate::ml::{ + cosine_distance, fastrp_init_fill, l2_distance, l2_norm, vec_add, vec_scale, vec_weighted_sum, +}; /// Both arguments must be same-sized `f32` vectors. fn validate_vector_args(arg_types: &[DataType], arity: usize, fname: &str) -> Result<()> { @@ -96,6 +106,7 @@ pub(crate) struct L2DistanceUDF { signature: Signature, } +#[allow(dead_code)] impl L2DistanceUDF { pub(crate) fn new() -> Self { Self { @@ -156,6 +167,7 @@ pub(crate) struct CosineDistanceUDF { signature: Signature, } +#[allow(dead_code)] impl CosineDistanceUDF { pub(crate) fn new() -> Self { Self { @@ -228,11 +240,698 @@ pub(crate) fn cosine_distance_expr(v1: Expr, v2: Expr) -> Expr { ScalarUDF::from(CosineDistanceUDF::new()).call(vec![v1, v2]) } +/// Builds the `el: Float32` child field shared by all vector columns here. +pub(crate) fn f32_child_field() -> Arc { + Arc::new(Field::new("el", DataType::Float32, false)) +} + +/// Wraps a flat buffer into a non-null `FixedSizeList` array. +fn fsl_array(d: usize, flat: Vec) -> Arc { + debug_assert_eq!( + flat.len() % d, + 0, + "flat vector buffer length must be a multiple of d" + ); + Arc::new( + FixedSizeListArray::try_new( + f32_child_field(), + d as i32, + Arc::new(Float32Array::from(flat)), + None, + ) + .expect("valid fixed size list"), + ) +} + +// ---------------- vec_sum: aggregate UDAF ---------------- +#[derive(Debug, PartialEq, Eq, Hash)] +pub(crate) struct VectorSum { + signature: Signature, + d: usize, +} + +impl VectorSum { + pub(crate) fn new(d: usize) -> Self { + Self { + // Nested-type signatures would pin the child field name and + // nullability, but a parquet round-trip renames/loosens it + // (`el` -> `element`, nullable). Accept any single argument and + // validate the vector shape in `return_type` instead. + signature: Signature::any(1, Volatility::Immutable), + d, + } + } +} + +impl AggregateUDFImpl for VectorSum { + fn name(&self) -> &str { + "vec_sum" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + validate_vector_args(arg_types, 1, "vec_sum")?; + Ok(DataType::FixedSizeList(f32_child_field(), self.d as i32)) + } + + fn accumulator(&self, _args: AccumulatorArgs) -> Result> { + Ok(Box::new(VecSumAccumulator { + inner: VecSumGroupsAccumulator::new(self.d), + })) + } + + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + true + } + + fn create_groups_accumulator( + &self, + _args: AccumulatorArgs, + ) -> Result> { + Ok(Box::new(VecSumGroupsAccumulator::new(self.d))) + } +} + +/// Grouped running sums for [`VectorSum`]: one dense `Vec` holding +/// `total_num_groups * d` floats; group `g` occupies `sums[g*d .. (g+1)*d]`. +/// +/// The buffer is zero-initialized (the sum identity), so a group is never +/// null. New groups arriving with a larger `total_num_groups` extend the +/// buffer; `EmitTo::First(n)` drains the first `n * d` floats so the +/// remaining group indices shift down, exactly as the trait contract +/// requires. +#[derive(Debug)] +struct VecSumGroupsAccumulator { + d: usize, + sums: Vec, +} + +impl VecSumGroupsAccumulator { + fn new(d: usize) -> Self { + Self { + d, + sums: Vec::new(), + } + } + + /// Grow the dense buffer to `total_num_groups` zero vectors. + fn resize(&mut self, total_num_groups: usize) { + let needed = total_num_groups.saturating_mul(self.d); + if self.sums.len() < needed { + self.sums.resize(needed, 0.0f32); + } + } + + /// Add the rows selected by `group_indices` (optionally filtered) into + /// the running group sums. + fn accumulate( + &mut self, + v: &crate::expressions::common::F32ListLike, + group_indices: &[usize], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + use crate::expressions::common::F32ListLike; + + self.resize(total_num_groups); + match v { + // Fast path: FixedSizeList rows are constant-sized, one length + // check is enough for the whole array. + F32ListLike::Fixed(_) => { + if v.value_length() != self.d { + return plan_err!( + "vec_sum: vector size {} does not match the declared dimension {}", + v.value_length(), + self.d + ); + } + } + // List rows are checked per row below (parquet round-trips make + // this the common on-disk representation). + F32ListLike::View(_) => {} + } + + for (row, &g) in group_indices.iter().enumerate() { + if let Some(filter) = opt_filter + && (filter.is_null(row) || !filter.value(row)) + { + continue; + } + let vals = v.value(row); + if vals.len() != self.d { + return plan_err!( + "vec_sum: vector length {} does not match the declared dimension {}", + vals.len(), + self.d + ); + } + let base = g * self.d; + vec_add(&mut self.sums[base..base + self.d], vals); + } + Ok(()) + } + + /// Detach the emitted groups' flat sums, honoring the `EmitTo::First(n)` + /// "shift down" contract. Groups beyond the accumulated tail (defensively, + /// if the hash table allocated more groups than this accumulator ever saw + /// rows for) keep the zero identity. + fn take_state(&mut self, emit_to: EmitTo) -> Vec { + match emit_to { + EmitTo::All => std::mem::take(&mut self.sums), + EmitTo::First(n) => { + let n_floats = n.saturating_mul(self.d); + let mut taken = Vec::with_capacity(n_floats); + if n_floats <= self.sums.len() { + taken.extend(self.sums.drain(0..n_floats)); + } else { + taken.extend(self.sums.drain(0..)); + taken.resize(n_floats, 0.0); + } + taken + } + } + } +} + +impl GroupsAccumulator for VecSumGroupsAccumulator { + fn update_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + debug_assert_eq!(values.len(), 1, "vec_sum takes exactly one argument"); + // Non-null contract: without a filter the whole batch must be + // null-free; with a FILTER clause the excluded rows are allowed to + // be null (they never reach the accumulation loop). + if opt_filter.is_none() { + debug_assert_eq!( + values[0].null_count(), + 0, + "vec_sum input must be non-null by contract" + ); + } + let v = as_f32_list_like(&values[0], "vec_sum", "first")?; + self.accumulate(&v, group_indices, opt_filter, total_num_groups) + } + + fn evaluate(&mut self, emit_to: EmitTo) -> Result { + let flat = self.take_state(emit_to); + Ok(fsl_array(self.d, flat)) + } + + fn state(&mut self, emit_to: EmitTo) -> Result> { + // SUM is self-combinable: the partial state is the (partial) sum. + let flat = self.take_state(emit_to); + Ok(vec![fsl_array(self.d, flat) as ArrayRef]) + } + + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + // State may come back from a parquet spill as List — accept + // both representations like every vector consumer in this crate. + let v = as_f32_list_like(&values[0], "vec_sum", "state")?; + self.accumulate(&v, group_indices, None, total_num_groups) + } + + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + // Partial-aggregation bypass: for SUM the input rows *are* valid + // state rows; filtered-out rows become the zero-vector identity. + let v = as_f32_list_like(&values[0], "vec_sum", "first")?; + let mut flat = Vec::with_capacity(v.len() * self.d); + for row in 0..v.len() { + let keep = opt_filter.map_or(true, |f| !f.is_null(row) && f.value(row)); + if keep { + let vals = v.value(row); + if vals.len() != self.d { + return plan_err!( + "vec_sum: vector length {} does not match the declared dimension {}", + vals.len(), + self.d + ); + } + flat.extend_from_slice(vals); + } else { + flat.extend(std::iter::repeat_n(0.0f32, self.d)); + } + } + Ok(vec![fsl_array(self.d, flat) as ArrayRef]) + } + + fn size(&self) -> usize { + // capacity-based so DataFusion's spill accounting sees the real cost + self.sums.capacity() * size_of::() + 2 * size_of::() + } +} + +/// No-GROUP-BY fallback: a single implicit group over the grouped +/// accumulator. +#[derive(Debug)] +struct VecSumAccumulator { + inner: VecSumGroupsAccumulator, +} + +impl Accumulator for VecSumAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + // every row belongs to the single implicit group 0 + let indices = vec![0usize; values[0].len()]; + self.inner.update_batch(values, &indices, None, 1) + } + + fn state(&mut self) -> Result> { + Ok(vec![self.evaluate()?]) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + let indices = vec![0usize; states[0].len()]; + self.inner.merge_batch(states, &indices, 1) + } + + fn size(&self) -> usize { + self.inner.size() + } + + fn evaluate(&mut self) -> Result { + let arr = self.inner.evaluate(EmitTo::All)?; + ScalarValue::try_from_array(&arr, 0) + } +} + +/// Builds an [`Expr`] summing same-sized `f32` vectors per group. +pub(crate) fn vec_sum_expr(v: Expr, d: usize) -> Expr { + AggregateUDF::from(VectorSum::new(d)).call(vec![v]) +} + +// ---------------- vec_scale: scalar UDF ---------------- + +/// Scalar UDF `vec_scale(v, s) -> v * s`: element-wise product of an `f32` +/// vector with a scalar (Float32 or Float64; the factor may vary per row). +#[derive(Debug, PartialEq, Eq, Hash)] +pub(crate) struct VecScale { + signature: Signature, +} + +impl VecScale { + pub(crate) fn new() -> Self { + Self { + signature: Signature::any(2, Volatility::Immutable), + } + } +} + +impl Default for VecScale { + fn default() -> Self { + Self::new() + } +} + +fn validate_vec_scale_args(arg_types: &[DataType]) -> Result<()> { + if arg_types.len() != 2 { + return plan_err!("vec_scale expects 2 arguments, got {}", arg_types.len()); + } + let is_vector = match &arg_types[0] { + DataType::FixedSizeList(f, _) => f.data_type() == &DataType::Float32, + DataType::List(f) => f.data_type() == &DataType::Float32, + _ => false, + }; + if !is_vector { + return plan_err!( + "vec_scale argument 0 must be FixedSizeList or List, got {:?}", + arg_types[0] + ); + } + match &arg_types[1] { + DataType::Float32 | DataType::Float64 => Ok(()), + other => plan_err!("vec_scale argument 1 must be Float32 or Float64, got {other:?}"), + } +} + +impl ScalarUDFImpl for VecScale { + fn name(&self) -> &str { + "vec_scale" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + validate_vec_scale_args(arg_types)?; + Ok(arg_types[0].clone()) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + let v = as_f32_list_like(&arrays[0], "vec_scale", "first")?; + + // per-row factor (a broadcast literal becomes a constant array) + let factors: Vec = match arrays[1].data_type() { + DataType::Float32 => arrays[1] + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + .collect(), + DataType::Float64 => arrays[1] + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .map(|x| *x as f32) + .collect(), + other => { + return plan_err!( + "vec_scale scalar factor must be Float32 or Float64, got {other:?}" + ); + } + }; + + match &v { + crate::expressions::common::F32ListLike::Fixed(a) => { + let size = a.value_length() as usize; + let mut flat: Vec = Vec::with_capacity(a.len() * size); + let mut validity: Vec = Vec::with_capacity(a.len()); + let mut buf = vec![0.0f32; size]; + for row in 0..a.len() { + if a.is_null(row) { + validity.push(false); + flat.extend(std::iter::repeat_n(0.0f32, size)); + } else { + validity.push(true); + let s = factors[row % factors.len()]; + buf.copy_from_slice(v.value(row)); + vec_scale(&mut buf, s); + flat.extend_from_slice(&buf); + } + } + let out = FixedSizeListArray::try_new( + child_field(&arrays[0])?, + size as i32, + Arc::new(Float32Array::from(flat)), + Some(NullBuffer::from(validity)), + )?; + Ok(ColumnarValue::Array(Arc::new(out))) + } + crate::expressions::common::F32ListLike::View(a) => { + let mut flat: Vec = Vec::with_capacity(a.values().len()); + let mut buf: Vec = Vec::new(); + for row in 0..a.len() { + let len = a.value(row).len(); + if a.is_null(row) { + // masked by the validity buffer; keep offsets valid + flat.extend(std::iter::repeat_n(0.0f32, len)); + } else { + let s = factors[row % factors.len()]; + buf.clear(); + buf.extend_from_slice(v.value(row)); + vec_scale(&mut buf, s); + flat.extend_from_slice(&buf); + } + } + let out = ListArray::new( + child_field(&arrays[0])?, + a.offsets().clone(), + Arc::new(Float32Array::from(flat)), + a.nulls().cloned(), + ); + Ok(ColumnarValue::Array(Arc::new(out))) + } + } + } +} + +/// Child field of a vector column (`List(Float32, f)` / `FixedSizeList(f, d)`). +fn child_field(array: &ArrayRef) -> Result> { + match array.data_type() { + DataType::List(f) => Ok(Arc::clone(f)), + DataType::FixedSizeList(f, _) => Ok(Arc::clone(f)), + other => plan_err!("expected a vector column, got {other:?}"), + } +} + +/// Scalar UDF `vec_weighted_sum(w0, v0, w1, v1, ...) -> FixedSizeList`: +/// the linear combination `Σ w_i · v_i` computed in one fused pass (SIMD +/// kernel in [`crate::ml::linalg`]). +/// +/// The number of (scalar, vector) term pairs is fixed at plan time. A NULL +/// vector row contributes zero — the "never reached" FastRP state — so the +/// combination never propagates NULLs. Vectors must be same-sized `f32` +/// (length `d`, captured at plan time); scalars are Float32 or Float64 and +/// may vary per row. +#[derive(Debug, PartialEq, Eq, Hash)] +pub(crate) struct VecWeightedSum { + signature: Signature, + d: usize, + terms: usize, +} + +impl VecWeightedSum { + pub(crate) fn new(d: usize, terms: usize) -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + d, + terms, + } + } +} + +fn validate_vec_weighted_sum_args(arg_types: &[DataType], terms: usize) -> Result<()> { + if arg_types.len() != 2 * terms { + return plan_err!( + "vec_weighted_sum expects {} arguments ({} (weight, vector) pairs), got {}", + 2 * terms, + terms, + arg_types.len() + ); + } + for i in 0..terms { + let (w, v) = (&arg_types[2 * i], &arg_types[2 * i + 1]); + match w { + DataType::Float32 | DataType::Float64 => {} + other => { + return plan_err!( + "vec_weighted_sum weight must be Float32 or Float64, got {other:?}" + ); + } + } + let is_vector = match v { + DataType::FixedSizeList(f, _) => f.data_type() == &DataType::Float32, + DataType::List(f) => f.data_type() == &DataType::Float32, + _ => false, + }; + if !is_vector { + return plan_err!( + "vec_weighted_sum vector must be FixedSizeList or List, got {v:?}" + ); + } + } + Ok(()) +} + +impl ScalarUDFImpl for VecWeightedSum { + fn name(&self) -> &str { + "vec_weighted_sum" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + validate_vec_weighted_sum_args(arg_types, self.terms)?; + Ok(DataType::FixedSizeList(f32_child_field(), self.d as i32)) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + // user-defined signature: accept the validated types as-is + validate_vec_weighted_sum_args(arg_types, self.terms)?; + Ok(arg_types.to_vec()) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + let n = arrays[0].len(); + + // per-row factors (a broadcast literal becomes a constant array) + let factors: Vec> = (0..self.terms) + .map(|i| match arrays[2 * i].data_type() { + DataType::Float32 => arrays[2 * i] + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + .collect(), + DataType::Float64 => arrays[2 * i] + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .map(|x| *x as f32) + .collect(), + other => unreachable!("validated at plan time: {other:?}"), + }) + .collect(); + + let vecs: Vec = (0..self.terms) + .map(|i| as_f32_list_like(&arrays[2 * i + 1], "vec_weighted_sum", "vector")) + .collect::>()?; + + let mut flat = vec![0.0f32; n * self.d]; + let mut buf = vec![0.0f32; self.d]; + let mut terms: Vec<(f32, &[f32])> = Vec::with_capacity(self.terms); + for row in 0..n { + terms.clear(); + for ti in 0..self.terms { + let v = &vecs[ti]; + if !v.is_null(row) { + let vals = v.value(row); + if vals.len() != self.d { + return plan_err!( + "vec_weighted_sum: vector length {} does not match the declared dimension {}", + vals.len(), + self.d + ); + } + terms.push((factors[ti][row % factors[ti].len()], vals)); + } + } + vec_weighted_sum(&mut buf, &terms); + flat[row * self.d..(row + 1) * self.d].copy_from_slice(&buf); + } + + Ok(ColumnarValue::Array(fsl_array(self.d, flat))) + } +} + +/// Builds an [`Expr`] computing the linear combination `Σ w_i · v_i` over +/// `(weight, vector)` term pairs (vectors may be NULL: a NULL term is zero). +pub(crate) fn vec_weighted_sum_expr(terms: &[(Expr, Expr)], d: usize) -> Expr { + let args: Vec = terms + .iter() + .flat_map(|(w, v)| vec![w.clone(), v.clone()]) + .collect(); + ScalarUDF::from(VecWeightedSum::new(d, terms.len())).call(args) +} + +/// Builds an [`Expr`] scaling an `f32` vector by a scalar (per row). +pub(crate) fn vec_scale_expr(v: Expr, s: Expr) -> Expr { + ScalarUDF::from(VecScale::new()).call(vec![v, s]) +} + +// ---------------- fastrp_init: scalar UDF ---------------- + +/// Scalar UDF `fastrp_init(id) -> FixedSizeList`: the +/// deterministic sparse random projection init of FastRP. +/// +/// `d` and `seed` are captured at plan time; the vector for a row is a pure +/// function of `(id, seed)` (see [`crate::ml::fastrp_init_fill`]), so init is +/// reproducible across re-scans and runs. +#[derive(Debug, PartialEq, Eq, Hash)] +pub(crate) struct FastRPInit { + signature: Signature, + d: usize, + seed: u64, +} + +impl FastRPInit { + pub(crate) fn new(d: usize, seed: u64) -> Self { + Self { + signature: Signature::exact(vec![DataType::Int64], Volatility::Immutable), + d, + seed, + } + } +} + +impl ScalarUDFImpl for FastRPInit { + fn name(&self) -> &str { + "fastrp_init" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + if arg_types.len() != 1 || arg_types[0] != DataType::Int64 { + return plan_err!("fastrp_init expects a single Int64 argument, got {arg_types:?}"); + } + Ok(DataType::FixedSizeList(f32_child_field(), self.d as i32)) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let arrays = ColumnarValue::values_to_arrays(&args.args)?; + let ids = downcast_int64(&arrays[0], "fastrp_init", "first")?; + + let mut flat = Vec::with_capacity(ids.len() * self.d); + let mut buf = vec![0.0f32; self.d]; + for i in 0..ids.len() { + fastrp_init_fill(ids.value(i), self.seed, self.d, &mut buf); + flat.extend_from_slice(&buf); + } + Ok(ColumnarValue::Array(fsl_array(self.d, flat))) + } +} + +/// Builds an [`Expr`] initializing the FastRP random projection vectors. +pub(crate) fn fastrp_init_expr(id: Expr, d: usize, seed: u64) -> Expr { + ScalarUDF::from(FastRPInit::new(d, seed)).call(vec![id]) +} + +/// Zero-vector [`ScalarValue`] matching the concrete vector type of a column +/// (`List` or `FixedSizeList`, including the very same +/// child field). +/// +/// A parquet round-trip turns `FixedSizeList` into `List`, so the literal +/// must adapt to stay type-compatible with `coalesce` / `when`. +pub(crate) fn vec_zero_scalar(vector_type: &DataType, d: usize) -> Result { + let zeros = Arc::new(Float32Array::from(vec![0.0f32; d])); + match vector_type { + DataType::List(f) => { + let arr = ListArray::new( + Arc::clone(f), + OffsetBuffer::new(ScalarBuffer::from(vec![0i32, d as i32])), + zeros, + None, + ); + Ok(ScalarValue::List(Arc::new(arr))) + } + DataType::FixedSizeList(f, size) if *size as usize == d => { + let arr = FixedSizeListArray::try_new(Arc::clone(f), *size, zeros, None)?; + Ok(ScalarValue::FixedSizeList(Arc::new(arr))) + } + other => plan_err!("expected a Float32 vector column of length {d}, got {other:?}"), + } +} + #[cfg(test)] mod tests { use super::*; - use datafusion::arrow::array::{Array, FixedSizeListArray, Float32Array, RecordBatch}; - use datafusion::arrow::datatypes::{Field, Schema}; + use std::collections::HashMap; + + use datafusion::arrow::array::{ + Array, FixedSizeListArray, Float32Array, Int64Array, RecordBatch, + }; + use datafusion::arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::common::Result; use datafusion::prelude::{DataFrame, SessionContext, col, lit}; @@ -389,4 +1088,642 @@ mod tests { format!("{}", cosine_distance_expr(col("a"), col("b"))).contains("cosine_distance") ); } + + // ---------------- vec_sum / fastrp_init ---------------- + + use crate::ml::fastrp_init_fill; + use datafusion::logical_expr::ExprFunctionExt; + + /// `(g: Int64, v: FixedSizeList)` table from `(group, vec)` rows. + fn group_table(rows: &[(i64, &[f32; 3])]) -> Result { + let flat: Vec = rows.iter().flat_map(|(_, v)| v.iter().copied()).collect(); + let groups = Int64Array::from(rows.iter().map(|(g, _)| *g).collect::>()); + let schema = Schema::new(vec![ + Field::new("g", DataType::Int64, false), + Field::new("v", DataType::FixedSizeList(f32_child_field(), 3), false), + ]); + let batch = + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(groups), fsl_array(3, flat)])?; + Ok(SessionContext::new().read_batch(batch)?) + } + + /// Serial per-group reference sums. + fn group_sums_ref(rows: &[(i64, &[f32; 3])]) -> Vec<(i64, [f32; 3])> { + let mut order: Vec = Vec::new(); + let mut sums: HashMap = HashMap::new(); + for (g, v) in rows { + if !sums.contains_key(g) { + order.push(*g); + sums.insert(*g, [0.0; 3]); + } + let s = sums.get_mut(g).unwrap(); + for (t, x) in v.iter().enumerate() { + s[t] += x; + } + } + order.into_iter().map(|g| (g, sums[&g])).collect() + } + + /// Downcast column `idx` into per-row `f32` vectors (any representation). + async fn rows_as_f32(df: DataFrame, idx: usize) -> Result>> { + let batches = df.collect().await?; + let mut out = Vec::new(); + for batch in &batches { + let col = batch.column(idx); + let v = crate::expressions::common::as_f32_list_like(col, "test", "first")?; + for i in 0..v.len() { + out.push(v.value(i).to_vec()); + } + } + Ok(out) + } + + #[tokio::test] + async fn udaf_vec_sum_matches_serial_reference() -> Result<()> { + let rows: Vec<(i64, &[f32; 3])> = vec![ + (2, &[1.0, 1.0, 1.0]), + (1, &[10.0, 0.0, -4.0]), + (2, &[0.5, 0.5, 0.5]), + (1, &[1.0, 2.0, 3.0]), + (2, &[8.25, -0.5, 0.125]), + (0, &[7.0, 7.0, 7.0]), + ]; + let df = group_table(&rows)?; + let out = df.aggregate(vec![col("g")], vec![vec_sum_expr(col("v"), 3)])?; + // group order is not guaranteed: sort by key for comparison + let mut got: Vec<(i64, [f32; 3])> = out + .collect() + .await? + .iter() + .flat_map(|b| { + let g = b.column(0).as_any().downcast_ref::().unwrap(); + let v = crate::expressions::common::as_f32_list_like(b.column(1), "test", "first") + .unwrap(); + (0..b.num_rows()).map(move |i| { + let mut a = [0.0f32; 3]; + a.copy_from_slice(v.value(i)); + (g.value(i), a) + }) + }) + .collect(); + got.sort_by_key(|(g, _)| *g); + + let mut expected = group_sums_ref(&rows); + expected.sort_by_key(|(g, _)| *g); + assert_eq!(got, expected); + Ok(()) + } + + #[tokio::test] + async fn udaf_vec_sum_accepts_list_input() -> Result<()> { + // Simulate a parquet round-trip: List instead of FixedSizeList. + let rows: Vec> = vec![ + vec![1.0, 2.0, 3.0], + vec![0.5, 0.5, 0.5], + vec![-1.0, 0.0, 1.0], + ]; + let values = Arc::new(Float32Array::from( + rows.iter().flatten().copied().collect::>(), + )); + let lists = ListArray::new( + f32_child_field(), + OffsetBuffer::new(ScalarBuffer::from(vec![0i32, 3, 6, 9])), + values, + None, + ); + let groups = Int64Array::from(vec![1i64, 1, 2]); + let schema = Schema::new(vec![ + Field::new("g", DataType::Int64, false), + Field::new("v", DataType::List(f32_child_field()), false), + ]); + let batch = + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(groups), Arc::new(lists)])?; + let df = SessionContext::new().read_batch(batch)?; + + let out = df + .aggregate(vec![col("g")], vec![vec_sum_expr(col("v"), 3)])? + .sort(vec![col("g").sort(true, true)])?; + let got = rows_as_f32(out, 1).await?; + assert_eq!(got, vec![vec![1.5, 2.5, 3.5], vec![-1.0, 0.0, 1.0]]); + Ok(()) + } + + #[tokio::test] + async fn udaf_vec_sum_rejects_mismatched_list_lengths() -> Result<()> { + let values = Arc::new(Float32Array::from(vec![1.0f32, 2.0, 3.0, 4.0])); + let lists = ListArray::new( + f32_child_field(), + OffsetBuffer::new(ScalarBuffer::from(vec![0i32, 3, 4])), + values, + None, + ); + let groups = Int64Array::from(vec![1i64, 1]); + let schema = Schema::new(vec![ + Field::new("g", DataType::Int64, false), + Field::new("v", DataType::List(f32_child_field()), false), + ]); + let batch = + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(groups), Arc::new(lists)])?; + let df = SessionContext::new().read_batch(batch)?; + let result = df + .aggregate(vec![col("g")], vec![vec_sum_expr(col("v"), 3)])? + .collect() + .await; + assert!( + result.is_err(), + "unequal vector lengths must surface as an error" + ); + Ok(()) + } + + #[tokio::test] + async fn udaf_vec_sum_respects_aggregate_filter() -> Result<()> { + // The FastRP pattern: NULL sources are excluded by an aggregate + // FILTER (WHERE v IS NOT NULL) and must not poison the sums. + // The input mimics a parquet round-trip: `List` with a + // nullable child field (not the canonical `el`). + let flat: Vec = vec![1.0, 2.0, 3.0, 10.0, 10.0, 10.0]; + let lists = ListArray::new( + Arc::new(Field::new("element", DataType::Float32, true)), + OffsetBuffer::new(ScalarBuffer::from(vec![0i32, 3, 3, 6])), + Arc::new(Float32Array::from(flat)), + // the middle row is a genuine NULL vector (never-reached source) + Some(NullBuffer::from(vec![true, false, true])), + ); + let groups = Int64Array::from(vec![1i64, 1, 1]); + let schema = Schema::new(vec![ + Field::new("g", DataType::Int64, false), + Field::new( + "v", + DataType::List(Arc::new(Field::new("element", DataType::Float32, true))), + true, + ), + ]); + let batch = + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(groups), Arc::new(lists)])?; + let df = SessionContext::new().read_batch(batch)?; + + let agg = vec_sum_expr(col("v"), 3) + .filter(col("v").is_not_null()) + .build()?; + let out = df.aggregate(vec![col("g")], vec![agg])?; + let got = rows_as_f32(out, 1).await?; + assert_eq!(got, vec![vec![11.0, 12.0, 13.0]]); + Ok(()) + } + + #[tokio::test] + async fn accumulator_single_group_fallback() -> Result<()> { + // No GROUP BY: DataFusion uses the plain Accumulator path. + let rows: Vec<&[f32; 3]> = vec![&[1.0, 1.0, 1.0], &[2.0, 0.5, -3.0], &[4.0, 4.0, 4.0]]; + let df = group_table(&rows.iter().map(|v| (0i64, *v)).collect::>())?; + let out = df.aggregate(vec![], vec![vec_sum_expr(col("v"), 3)])?; + let got = rows_as_f32(out, 0).await?; + assert_eq!(got, vec![vec![7.0, 5.5, 2.0]]); + Ok(()) + } + + // --------- direct GroupsAccumulator trait-level tests --------- + + fn fsl_batch(rows: &[[f32; 3]]) -> ArrayRef { + let flat: Vec = rows.iter().flat_map(|r| r.iter().copied()).collect(); + fsl_array(3, flat) + } + + #[test] + fn groups_accumulator_updates_and_emits_in_group_order() -> Result<()> { + let mut acc = VecSumGroupsAccumulator::new(3); + acc.update_batch( + &[fsl_batch(&[ + [1.0, 1.0, 1.0], + [2.0, 2.0, 2.0], + [10.0, 0.0, 0.0], + ])], + &[0, 1, 0], + None, + 2, + )?; + acc.update_batch( + &[fsl_batch(&[[0.5, 0.5, 0.5], [1.0, 1.0, 1.0]])], + &[1, 0], + None, + 2, + )?; + + let out = acc.evaluate(EmitTo::All)?; + let v = crate::expressions::common::as_f32_list_like(&out, "test", "first")?; + assert_eq!(v.len(), 2); + assert_eq!(v.value(0), &[12.0, 2.0, 2.0]); + assert_eq!(v.value(1), &[2.5, 2.5, 2.5]); + Ok(()) + } + + #[test] + fn groups_accumulator_emit_first_shifts_indices() -> Result<()> { + let mut acc = VecSumGroupsAccumulator::new(3); + acc.update_batch( + &[fsl_batch(&[ + [1.0, 1.0, 1.0], + [2.0, 2.0, 2.0], + [3.0, 3.0, 3.0], + [4.0, 4.0, 4.0], + ])], + &[0, 1, 2, 3], + None, + 4, + )?; + + // Emit groups 0 and 1; indices 2, 3 shift down to 0, 1. + let emitted = acc.evaluate(EmitTo::First(2))?; + let v = crate::expressions::common::as_f32_list_like(&emitted, "test", "first")?; + assert_eq!(v.value(0), &[1.0, 1.0, 1.0]); + assert_eq!(v.value(1), &[2.0, 2.0, 2.0]); + + acc.update_batch( + &[fsl_batch(&[[100.0, 100.0, 100.0], [10.0, 10.0, 10.0]])], + &[0, 1], // old groups 2, 3 + None, + 2, + )?; + let out = acc.evaluate(EmitTo::All)?; + let v = crate::expressions::common::as_f32_list_like(&out, "test", "first")?; + assert_eq!(v.value(0), &[103.0, 103.0, 103.0]); + assert_eq!(v.value(1), &[14.0, 14.0, 14.0]); + Ok(()) + } + + #[test] + fn groups_accumulator_merge_accepts_parquet_list_state() -> Result<()> { + // Phase 1: partial sums. + let mut partial = VecSumGroupsAccumulator::new(3); + partial.update_batch( + &[fsl_batch(&[[1.0, 2.0, 3.0], [0.5, 0.5, 0.5]])], + &[0, 1], + None, + 2, + )?; + let state = partial.state(EmitTo::All)?; + + // Simulate a parquet spill: FixedSizeList state comes back as List. + let state_list = + datafusion::arrow::compute::cast(&state[0], &DataType::List(f32_child_field()))?; + + // Phase 2: merge (as List!) into a fresh accumulator. + let mut final_acc = VecSumGroupsAccumulator::new(3); + final_acc.merge_batch(&[state_list], &[0, 1], 2)?; + final_acc.update_batch( + &[fsl_batch(&[[1.0, 1.0, 1.0], [7.0, 7.0, 7.0]])], + &[0, 1], + None, + 2, + )?; + + let out = final_acc.evaluate(EmitTo::All)?; + let v = crate::expressions::common::as_f32_list_like(&out, "test", "first")?; + assert_eq!(v.value(0), &[2.0, 3.0, 4.0]); + assert_eq!(v.value(1), &[7.5, 7.5, 7.5]); + Ok(()) + } + + #[test] + fn groups_accumulator_respects_opt_filter() -> Result<()> { + let mut acc = VecSumGroupsAccumulator::new(3); + let filter = BooleanArray::from(vec![true, false, true]); + acc.update_batch( + &[fsl_batch(&[ + [1.0, 0.0, 0.0], + [1000.0, 1000.0, 1000.0], + [0.0, 2.0, 0.0], + ])], + &[0, 0, 0], + Some(&filter), + 1, + )?; + let out = acc.evaluate(EmitTo::All)?; + let v = crate::expressions::common::as_f32_list_like(&out, "test", "first")?; + assert_eq!(v.value(0), &[1.0, 2.0, 0.0]); + Ok(()) + } + + #[test] + fn groups_accumulator_convert_to_state_zeroes_filtered_rows() -> Result<()> { + let acc = VecSumGroupsAccumulator::new(3); + let filter = BooleanArray::from(vec![true, false, true]); + let state = acc.convert_to_state( + &[fsl_batch(&[ + [1.0, 1.0, 1.0], + [9.0, 9.0, 9.0], + [3.0, 3.0, 3.0], + ])], + Some(&filter), + )?; + let v = crate::expressions::common::as_f32_list_like(&state[0], "test", "first")?; + assert_eq!(v.value(0), &[1.0, 1.0, 1.0]); + assert_eq!( + v.value(1), + &[0.0, 0.0, 0.0], + "filtered row -> zero identity" + ); + assert_eq!(v.value(2), &[3.0, 3.0, 3.0]); + Ok(()) + } + + // --------- fastrp_init --------- + + #[tokio::test] + async fn udf_fastrp_init_is_deterministic_and_matches_kernel() -> Result<()> { + let ids = Int64Array::from(vec![1i64, 2, 3, -7]); + let schema = Schema::new(vec![Field::new("id", DataType::Int64, false)]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(ids)])?; + let df = SessionContext::new().read_batch(batch)?; + + let got = rows_as_f32( + df.clone() + .select(vec![fastrp_init_expr(col("id"), 4, 42)])?, + 0, + ) + .await?; + + let mut expected = vec![0.0f32; 4]; + for (i, id) in [1i64, 2, 3, -7].iter().enumerate() { + fastrp_init_fill(*id, 42, 4, &mut expected); + assert_eq!(got[i], expected, "fastrp_init({id}) must match the kernel"); + } + + // Same input, different seed: different init. + let other = rows_as_f32(df.select(vec![fastrp_init_expr(col("id"), 4, 43)])?, 0).await?; + assert_ne!(got, other); + Ok(()) + } + + #[test] + fn vec_zero_scalar_matches_column_type() -> Result<()> { + let list_type = DataType::List(f32_child_field()); + let list_zero = vec_zero_scalar(&list_type, 3)?; + assert_eq!(list_zero.data_type(), list_type); + + let fsl_type = DataType::FixedSizeList(f32_child_field(), 3); + let fsl_zero = vec_zero_scalar(&fsl_type, 3)?; + assert_eq!(fsl_zero.data_type(), fsl_type); + + // Each literal holds a single row of d zeros. + match &list_zero { + ScalarValue::List(arr) => { + assert_eq!(arr.len(), 1); + let row = arr.value(0); + let f32s = row.as_any().downcast_ref::().unwrap(); + assert_eq!(f32s.values(), &[0.0f32; 3]); + } + other => panic!("expected List scalar, got {other:?}"), + } + assert!(matches!(fsl_zero, ScalarValue::FixedSizeList(_))); + Ok(()) + } + + #[test] + fn builders_reference_vector_udf_names() { + assert!(format!("{}", vec_sum_expr(col("v"), 4)).contains("vec_sum")); + assert!(format!("{}", fastrp_init_expr(col("id"), 4, 0)).contains("fastrp_init")); + assert!(format!("{}", vec_scale_expr(col("v"), lit(2.0f32))).contains("vec_scale")); + assert!( + format!( + "{}", + vec_weighted_sum_expr(&[(lit(1.0f64), col("a")), (lit(2.0f64), col("b"))], 4) + ) + .contains("vec_weighted_sum") + ); + } + + // ---------------- vec_weighted_sum ---------------- + + #[tokio::test] + async fn udf_vec_weighted_sum_matches_serial_reference() -> Result<()> { + // rows: [v0, v1, NULL]; weights: [1.5, -2.0, 0.5] + let terms: Vec<(Expr, Expr)> = vec![ + (lit(1.5f64), col("h1")), + (lit(-2.0f64), col("h2")), + (lit(0.5f64), col("h3")), + ]; + let d = 3; + let mk = |vals: Vec| { + FixedSizeListArray::try_new( + f32_child_field(), + d as i32, + Arc::new(Float32Array::from(vals)), + None, + ) + .unwrap() + }; + // h1 = [(1,2,3), (0,0,0), NULL] + let h1 = FixedSizeListArray::try_new( + f32_child_field(), + d as i32, + Arc::new(Float32Array::from(vec![ + Some(1.0f32), + Some(2.0), + Some(3.0), + Some(0.0), + Some(0.0), + Some(0.0), + Some(9.0), + Some(9.0), + Some(9.0), + ])), + Some(NullBuffer::from(vec![true, true, false])), + )?; + // h2 = [(0.5,1,-1); 3 rows] + let h2 = mk(vec![0.5, 1.0, -1.0, 0.5, 1.0, -1.0, 0.5, 1.0, -1.0]); + // h3 = [(2,2,2); 3 rows] + let h3 = mk(vec![2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0, 2.0]); + + let schema = Schema::new(vec![ + Field::new( + "h1", + DataType::FixedSizeList(f32_child_field(), d as i32), + true, + ), + Field::new( + "h2", + DataType::FixedSizeList(f32_child_field(), d as i32), + false, + ), + Field::new( + "h3", + DataType::FixedSizeList(f32_child_field(), d as i32), + false, + ), + ]); + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![Arc::new(h1), Arc::new(h2), Arc::new(h3)], + )?; + let df = SessionContext::new().read_batch(batch)?; + + let out = df.select(vec![vec_weighted_sum_expr(&terms, d).alias("c")])?; + let got = rows_as_f32(out, 0).await?; + + // row 0: 1.5*(1,2,3) - 2*(0.5,1,-1) + 0.5*(2,2,2) = (1.5, 2, 7.5) + assert!( + (got[0][0] - 1.5).abs() < 1e-6 + && (got[0][1] - 2.0).abs() < 1e-6 + && (got[0][2] - 7.5).abs() < 1e-4, + "{:?}", + got[0] + ); + // row 1: 1.5*(0,0,0) - 2*(0.5,1,-1) + 0.5*(2,2,2) = (0,-1,3) + assert!( + (got[1][0].abs()) < 1e-6 + && (got[1][1] + 1.0).abs() < 1e-6 + && (got[1][2] - 3.0).abs() < 1e-5, + "{:?}", + got[1] + ); + // row 2: h1 is NULL -> contributes zero, i.e. identical to row 1's + // explicit zero vector: -2*(0.5,1,-1) + 0.5*(2,2,2) = (0,-1,3) + assert_eq!( + got[2], got[1], + "NULL vector must behave like the zero vector" + ); + Ok(()) + } + + #[test] + fn udf_vec_weighted_sum_validates_args() -> Result<()> { + // one (weight, vector) pair -> exactly 2 arguments + let udf = VecWeightedSum::new(2, 1); + assert!( + udf.return_type(&[DataType::Float64, DataType::List(f32_child_field())]) + .is_ok() + ); + // wrong arity + assert!(udf.return_type(&[DataType::Float64]).is_err()); + // non-float scalar + assert!( + udf.return_type(&[DataType::Int64, DataType::List(f32_child_field())]) + .is_err() + ); + // non-vector argument + assert!( + udf.return_type(&[DataType::Float64, DataType::Float32]) + .is_err() + ); + Ok(()) + } + + // ---------------- vec_scale ---------------- + + /// `(v: F32 vector of len `d`)` single-column table. + fn vec_table(rows: Vec>>, d: usize) -> Result { + let flat: Vec> = rows.iter().flatten().copied().collect(); + let fsl = FixedSizeListArray::try_new( + f32_child_field(), + d as i32, + Arc::new(Float32Array::from(flat)), + None, + )?; + let schema = Schema::new(vec![Field::new( + "v", + DataType::FixedSizeList(f32_child_field(), d as i32), + true, + )]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(fsl)])?; + Ok(SessionContext::new().read_batch(batch)?) + } + + #[tokio::test] + async fn udf_vec_scale_multiplies_by_scalar_literal() -> Result<()> { + let df = vec_table(vec![vec![Some(1.0), Some(-2.0), Some(0.5)]], 3)?; + let out = df + .clone() + .select(vec![vec_scale_expr(col("v"), lit(2.0f32)).alias("s")])?; + let got = rows_as_f32(out, 0).await?; + assert_eq!(got, vec![vec![2.0, -4.0, 1.0]]); + + // Float64 factor is accepted too. + let got = rows_as_f32( + df.select(vec![vec_scale_expr(col("v"), lit(0.5f64)).alias("s")])?, + 0, + ) + .await?; + assert_eq!(got, vec![vec![0.5, -1.0, 0.25]]); + Ok(()) + } + + #[tokio::test] + async fn udf_vec_scale_passes_nulls_through() -> Result<()> { + // Mimics the FastRP message path: NULL states stay NULL after scaling. + let flat: Vec> = vec![Some(1.0), Some(2.0), None, None]; + let nullable_child = Arc::new(Field::new("el", DataType::Float32, true)); + let fsl = FixedSizeListArray::try_new( + Arc::clone(&nullable_child), + 2, + Arc::new(Float32Array::from(flat)), + Some(NullBuffer::from(vec![true, false])), + )?; + let schema = Schema::new(vec![Field::new( + "v", + DataType::FixedSizeList(nullable_child, 2), + true, + )]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(fsl)])?; + let df = SessionContext::new().read_batch(batch)?; + + let out = df.select(vec![vec_scale_expr(col("v"), lit(3.0f32)).alias("s")])?; + let batches = out.collect().await?; + let v = + crate::expressions::common::as_f32_list_like(batches[0].column(0), "test", "first")?; + assert!(!v.is_null(0)); + assert_eq!(v.value(0), &[3.0, 6.0]); + assert!(v.is_null(1), "null vector must stay null"); + Ok(()) + } + + #[tokio::test] + async fn udf_vec_scale_with_l2_norm_gives_unit_norm() -> Result<()> { + // The norm_output composition: v / ||v|| has unit L2 norm. + let df = vec_table(vec![vec![Some(3.0), Some(4.0)]], 2)?; + let out = df.select(vec![ + vec_scale_expr(col("v"), lit(1.0f64) / l2_norm_expr(col("v"))).alias("s"), + ])?; + let got = rows_as_f32(out, 0).await?; + assert!( + (got[0][0] - 0.6).abs() < 1e-6 && (got[0][1] - 0.8).abs() < 1e-6, + "expected (0.6, 0.8), got {:?}", + got[0] + ); + let norm = (got[0][0] * got[0][0] + got[0][1] * got[0][1]).sqrt(); + assert!((norm - 1.0).abs() < 1e-5); + Ok(()) + } + + #[tokio::test] + async fn udf_vec_scale_preserves_list_representation() -> Result<()> { + let values = Arc::new(Float32Array::from(vec![1.0f32, 2.0, 3.0, 4.0])); + let lists = ListArray::new( + f32_child_field(), + OffsetBuffer::new(ScalarBuffer::from(vec![0i32, 2, 4])), + values, + None, + ); + let schema = Schema::new(vec![Field::new( + "v", + DataType::List(f32_child_field()), + false, + )]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(lists)])?; + let df = SessionContext::new().read_batch(batch)?; + + let out = df.select(vec![vec_scale_expr(col("v"), lit(10.0f32)).alias("s")])?; + let batches = out.collect().await?; + assert_eq!( + batches[0].column(0).data_type(), + &DataType::List(f32_child_field()), + "List in -> List out" + ); + let v = + crate::expressions::common::as_f32_list_like(batches[0].column(0), "test", "first")?; + assert_eq!(v.value(0), &[10.0, 20.0]); + assert_eq!(v.value(1), &[30.0, 40.0]); + Ok(()) + } } diff --git a/src/lib.rs b/src/lib.rs index bd69e7b..0e10ff6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,9 +4,11 @@ mod memory; mod ml; mod utils; +pub use algorithm::community::fastrp_clustering::FastRPClusteringBuilder; pub use algorithm::community::power_iteration_clustering::{ EmbeddingMode, InitStrategy, PICBuilder, WeightsStrategy, }; +pub use algorithm::embeddings::fastrp::{FastRPBuilder, FastRPNormalization}; pub use expressions::kmeans_assign_expr; pub use ml::{DistanceMetric, KMeansBuilder, KMeansResult, KMeansRun}; pub use utils::GraphFramesConfig; diff --git a/src/main.rs b/src/main.rs index c8bdc6a..fb77da0 100644 --- a/src/main.rs +++ b/src/main.rs @@ -11,7 +11,8 @@ 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, + DistanceMetric, EmbeddingMode, FastRPNormalization, InitStrategy, KMeansBuilder, + WeightsStrategy, kmeans_assign_expr, }; use graphframes_rs::{EDGE_DST, EDGE_SRC, GraphFrame, VERTEX_ID}; use std::path::{Path, PathBuf}; @@ -26,6 +27,16 @@ enum Format { Json, } +#[derive(Debug, Clone, Copy, ValueEnum)] +enum FastrpNormalization { + /// Plain sum, no normalization. + None, + /// Linear normalization: divide by the source out-degree. + L1, + /// Square normalization: divide by the square root of the source out-degree. + L2, +} + #[derive(Debug, Clone, Copy, ValueEnum)] enum PicInit { /// Degree vector d/Σd — the paper's recommended (and MLlib's "degree") init. @@ -258,6 +269,91 @@ enum Command { cmd: MllibCommand, }, + /// FastRP (Fast Random Projection) vertex embeddings. + Fastrp { + #[command(flatten)] + common: CommonArgs, + + /// Embedding dimension D. + #[arg(long)] + dim: usize, + + /// Number of propagation iterations K. The embedding is the last + /// iterate H_K; 0 returns the raw random projections. + #[arg(long, default_value_t = 4)] + iterations: usize, + + /// Seed for the per-vertex sparse random projections. + #[arg(long, default_value_t = 42)] + seed: u64, + + /// Per-iteration normalization of the propagated vectors: + /// `l1` divides each vector by the source out-degree (paper's S⁻¹), + /// `l2` by the square root of it (S^(-1/2)). + #[arg(long, value_enum, default_value_t = FastrpNormalization::None)] + normalization: FastrpNormalization, + + /// L2-normalize the final embeddings to unit length. + #[arg(long)] + norm_output: bool, + + /// Per-iterate weights of the final linear combination + /// H = Σ w_t·H_t over H_1..H_K (must have K entries; default: all + /// ones). The random init is never part of the combination; a zero + /// weight drops the iterate entirely. + #[arg(long, value_delimiter = ',')] + iteration_weights: Option>, + }, + + /// Graph clustering: FastRP embeddings (unit norm) + K-Means. + FastrpClustering { + #[command(flatten)] + common: CommonArgs, + + /// Embedding dimension D. + #[arg(long)] + dim: usize, + + /// FastRP propagation iterations K. + #[arg(long, default_value_t = 4)] + iterations: usize, + + /// Seed for the random projections and the k-means|| init. + #[arg(long, default_value_t = 42)] + seed: u64, + + /// Per-iteration normalization of the propagated vectors + /// (see fastrp). + #[arg(long, value_enum, default_value_t = FastrpNormalization::None)] + normalization: FastrpNormalization, + + /// Per-iterate weights of the FastRP linear combination + /// H = Σ w_t·H_t over H_1..H_K (must have K entries; default: all + /// ones). + #[arg(long, value_delimiter = ',')] + iteration_weights: Option>, + + /// Number of clusters; one center column per K (comma-separated). + #[arg(long, value_delimiter = ',', default_values_t = vec![2usize])] + k: Vec, + + /// K-Means distance metric (embeddings are L2-normalized). + #[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, + }, + /// Classical Label Propagation (CDLP). ClassicalLp { #[command(flatten)] @@ -287,13 +383,13 @@ enum KmeansArgsMetric { enum MllibCommand { /// K-Means (k-means|| init, Lloyd iterations) over a feature column. /// - /// The vertices file must contain an Int64 `id` column and a + /// The features (vertices) file must contain an Int64 `id` column and a /// Float32 feature column of the shape `List` or - /// `FixedSizeList`; no edges are read. + /// `FixedSizeList`. Kmeans { /// Path (or URI) to the features (vertices) file or directory. #[arg(long)] - vertices: String, + features: String, /// Output directory as a `file://` URI. #[arg(long)] @@ -303,7 +399,7 @@ enum MllibCommand { #[arg(long, value_enum, default_value_t = Format::Parquet)] format: Format, - /// Name of the vertex-id column in the input; renamed to `id`. + /// Name of the vertex-id column in the input. #[arg(long, default_value = "id")] id_col_name: String, @@ -641,7 +737,7 @@ async fn main() -> Result<()> { } Command::Mllib { cmd } => match cmd { MllibCommand::Kmeans { - vertices, + features, output, format, id_col_name, @@ -660,7 +756,7 @@ async fn main() -> Result<()> { 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 raw = read_data(&ctx, &features, format).await?; let features = raw.select(vec![col(&id_col_name).alias(VERTEX_ID), col(&feature_col)])?; @@ -721,6 +817,84 @@ async fn main() -> Result<()> { log::info!("result was written into {output}"); } }, + Command::Fastrp { + common, + dim, + iterations, + seed, + normalization, + norm_output, + iteration_weights, + } => { + let (ctx, g, ckpt) = setup(&common).await?; + let normalization = match normalization { + FastrpNormalization::None => FastRPNormalization::None, + FastrpNormalization::L1 => FastRPNormalization::L1, + FastrpNormalization::L2 => FastRPNormalization::L2, + }; + let mut builder = g + .fastrp() + .dim(dim) + .iterations(iterations) + .seed(seed) + .normalization(normalization) + .norm_output(norm_output); + if let Some(w) = iteration_weights { + builder = builder.iteration_weights(w); + } + let iterations = builder + .set_checkpoint_dir(ckpt) + .run(&ctx, &common.output) + .await?; + log::info!("FastRP finished after {iterations} iterations"); + } + Command::FastrpClustering { + common, + dim, + iterations, + seed, + normalization, + iteration_weights, + k, + metric, + max_iter, + tol, + init_steps, + } => { + let (ctx, g, ckpt) = setup(&common).await?; + let normalization = match normalization { + FastrpNormalization::None => FastRPNormalization::None, + FastrpNormalization::L1 => FastRPNormalization::L1, + FastrpNormalization::L2 => FastRPNormalization::L2, + }; + let metric = match metric { + KmeansArgsMetric::L2 => DistanceMetric::L2, + KmeansArgsMetric::Cosine => DistanceMetric::Cosine, + }; + let mut builder = g + .fastrp_clustering() + .set_dim(dim) + .set_iterations(iterations) + .set_seed(seed) + .set_normalization(normalization); + if let Some(w) = iteration_weights { + builder = builder.set_iteration_weights(w); + } + let res = builder + .set_multiple_k(k) + .set_metric(metric) + .set_max_iter(max_iter) + .set_tol(tol) + .set_kmeans_init_steps(init_steps) + .set_checkpoint_dir(ckpt) + .run(&ctx, &common.output) + .await?; + log::info!( + "FastRP clustering: d = {}, KMeans iterations = {}", + res.d, + res.num_iterations + ); + } Command::ClassicalLp { common, max_iter, diff --git a/src/ml.rs b/src/ml.rs index 6def91a..6126aff 100644 --- a/src/ml.rs +++ b/src/ml.rs @@ -5,4 +5,6 @@ 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}; +pub(crate) use linalg::{ + cosine_distance, fastrp_init_fill, l2_distance, l2_norm, vec_add, vec_scale, vec_weighted_sum, +}; diff --git a/src/ml/linalg.rs b/src/ml/linalg.rs index ac17125..46ef1a0 100644 --- a/src/ml/linalg.rs +++ b/src/ml/linalg.rs @@ -1,7 +1,7 @@ -//! SIMD linear-algebra kernels over `f32` vectors. +//! SIMD linear-algebra kernels and pure `f32` vector helpers. //! //! All manual SIMD code for the ml module lives here; this module is -//! DataFusion-free on purpose — the scalar-UDF wrappers live in +//! DataFusion-free on purpose — the scalar-UDF / UDAF wrappers live in //! [`crate::expressions::linalg`]. //! //! Contract: vectors are non-null, same-sized slices; for the two-argument @@ -11,6 +11,133 @@ use std::ops::Add; use wide::f32x8; +/// Accumulate `v` into `acc` in place: `acc[i] += v[i]` for every `i`. +/// +/// Hot path of the `vec_sum` group accumulator: both slices have the same +/// length (the equal-size, non-null vector contract of this module). +pub(crate) fn vec_add(acc: &mut [f32], v: &[f32]) { + debug_assert_eq!( + acc.len(), + v.len(), + "vec_add slices must have the same length" + ); + let d = acc.len().min(v.len()); + + let mut t = 0; + while t + 8 <= d { + let av = f32x8::from(&acc[t..t + 8]); + let vv = f32x8::from(&v[t..t + 8]); + acc[t..t + 8].copy_from_slice((av + vv).as_array_ref()); + t += 8; + } + + while t < d { + acc[t] += v[t]; + t += 1; + } +} + +/// Scale `v` in place: `v[i] *= s` for every `i`. +/// +/// Used by the `vec_scale` scalar UDF (per-step FastRP normalization and the +/// final output normalization). +pub(crate) fn vec_scale(v: &mut [f32], s: f32) { + let d = v.len(); + let sv = f32x8::splat(s); + + let mut t = 0; + while t + 8 <= d { + let vv = f32x8::from(&v[t..t + 8]); + v[t..t + 8].copy_from_slice((vv * sv).as_array_ref()); + t += 8; + } + + while t < d { + v[t] *= s; + t += 1; + } +} + +/// Weighted sum of `terms` into `out`: `out[i] = Σ (w * v[i])`. +/// +/// Fused form of "scale every vector, then add them up": one pass, one +/// output buffer. Hot path of the `vec_weighted_sum` scalar UDF (the final +/// linear combination of FastRP iterates). All vectors have the same length +/// (the equal-size, non-null contract of this module). +pub(crate) fn vec_weighted_sum(out: &mut [f32], terms: &[(f32, &[f32])]) { + debug_assert!( + terms.iter().all(|(_, v)| v.len() == out.len()), + "vec_weighted_sum terms must have the same length as out" + ); + let d = out.len(); + + let mut t = 0; + while t + 8 <= d { + let mut acc = f32x8::splat(0.0); + for (w, v) in terms { + acc = f32x8::from(&v[t..t + 8]).mul_add(f32x8::splat(*w), acc); + } + out[t..t + 8].copy_from_slice(acc.as_array_ref()); + t += 8; + } + + while t < d { + let mut acc = 0.0f32; + for (w, v) in terms { + acc = v[t].mul_add(*w, acc); + } + out[t] = acc; + t += 1; + } +} + +/// Deterministic sparse random projection vector of a single vertex — the +/// `U^(0)` init of FastRP (Chen & King, CIKM 2019). +/// +/// Entries are drawn from `{−1, 0, +1}`: `+1` and `−1` each with probability +/// `1 / (2·√q)`, zero otherwise, where `q` is the largest power of two not +/// exceeding `d` (the paper's `2^d*`). The result is a pure function of +/// `(id, seed)`, so re-scans and re-runs reproduce the same init. +pub(crate) fn fastrp_init_fill(id: i64, seed: u64, d: usize, out: &mut [f32]) { + debug_assert_eq!(out.len(), d, "fastrp_init_fill output must have length d"); + let q = if d.is_power_of_two() { + d + } else { + d.next_power_of_two() >> 1 + }; + let inv = 1.0f32 / (q as f32).sqrt(); // per-sign probability + + let mut state = splitmix64(seed ^ id.cast_unsigned().wrapping_mul(GOLDEN)); + for slot in out.iter_mut().take(d) { + let u = splitmix_unit(&mut state); + *slot = if u < 0.5 * inv { + 1.0 + } else if u < inv { + -1.0 + } else { + 0.0 + }; + } +} + +/// `2^64 / φ`, the SplitMix64 golden-gamma constant. +const GOLDEN: u64 = 0x9E37_79B9_7F4A_7C15; + +/// SplitMix64 next-state mixer (used as a tiny dependency-free RNG). +fn splitmix64(state: u64) -> u64 { + let mut z = state.wrapping_add(GOLDEN); + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) +} + +/// Uniform `f32` in `[0, 1)` from a SplitMix64 stream (53 random bits). +fn splitmix_unit(state: &mut u64) -> f32 { + let z = splitmix64(*state); + *state = z; + (z >> 11) as f32 / (1u64 << 53) as f32 +} + /// Squared L2 distance `||x - c||^2`. /// /// Kept squared (no `sqrt`) because the only hot-loop caller, the K-Means @@ -218,4 +345,145 @@ mod tests { 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); } + + // ---------------- vec_add / fastrp_init ---------------- + + /// Naive scalar `acc += v` reference. + fn vec_add_ref(acc: &mut [f32], v: &[f32]) { + for (a, b) in acc.iter_mut().zip(v) { + *a += *b; + } + } + + #[test] + fn test_vec_add_matches_scalar_reference() { + for d in [1usize, 3, 7, 8, 9, 15, 16, 17, 64, 100] { + let (x, c) = vec_pair(d, 11 + d as u32); + let mut acc = x.clone(); + let mut acc_ref = x.clone(); + vec_add(&mut acc, &c); + vec_add_ref(&mut acc_ref, &c); + assert_eq!(acc, acc_ref, "vec_add, d={d}"); + } + } + + #[test] + fn test_vec_add_repeated_accumulation() { + let d = 20; + let (_, c) = vec_pair(d, 5); + let mut acc = vec![0.0f32; d]; + for _ in 0..7 { + vec_add(&mut acc, &c); + } + let mut expected = vec![0.0f32; d]; + for _ in 0..7 { + vec_add_ref(&mut expected, &c); + } + assert_eq!(acc, expected); + } + + /// Naive scalar `v *= s` reference. + fn vec_scale_ref(v: &mut [f32], s: f32) { + for x in v.iter_mut() { + *x *= s; + } + } + + #[test] + fn test_vec_scale_matches_scalar_reference() { + for d in [1usize, 3, 7, 8, 9, 15, 16, 17, 64, 100] { + let (x, _) = vec_pair(d, 3 + d as u32); + let mut v = x.clone(); + let mut v_ref = x.clone(); + let s = 0.375f32; // exact in f32 + vec_scale(&mut v, s); + vec_scale_ref(&mut v_ref, s); + assert_eq!(v, v_ref, "vec_scale, d={d}"); + } + // identity and zero scale + let mut v = vec![1.0f32, -2.0, 3.5]; + vec_scale(&mut v, 1.0); + assert_eq!(v, vec![1.0, -2.0, 3.5]); + vec_scale(&mut v, 0.0); + assert_eq!(v, vec![0.0; 3]); + } + + #[test] + fn test_vec_weighted_sum_matches_scalar_reference() { + for d in [1usize, 3, 7, 8, 9, 16, 17, 64, 100] { + let (a, _) = vec_pair(d, 21); + let (b, _) = vec_pair(d, 22); + let (c, _) = vec_pair(d, 23); + let terms = vec![ + (0.5f32, a.as_slice()), + (-1.25, b.as_slice()), + (2.0, c.as_slice()), + ]; + + let mut out = vec![0.0f32; d]; + vec_weighted_sum(&mut out, &terms); + + let mut expected = vec![0.0f32; d]; + for (w, v) in &terms { + for i in 0..d { + expected[i] = v[i].mul_add(*w, expected[i]); + } + } + // wide's mul_add may be emulated (not hardware-fused), so the + // last bits can differ from scalar f32::mul_add + for i in 0..d { + assert_close( + out[i], + expected[i], + &format!("vec_weighted_sum, d={d}, i={i}"), + ); + } + } + // no terms at all -> zero vector + let mut out = vec![9.0f32; 5]; + vec_weighted_sum(&mut out, &[]); + assert_eq!(out, vec![0.0; 5]); + } + + #[test] + fn test_fastrp_init_is_deterministic_and_in_domain() { + for d in [1usize, 2, 63, 64, 65, 256] { + let mut a = vec![0.0f32; d]; + let mut b = vec![0.0f32; d]; + fastrp_init_fill(1234, 42, d, &mut a); + fastrp_init_fill(1234, 42, d, &mut b); + assert_eq!(a, b, "same (id, seed) must reproduce the init, d={d}"); + assert!( + a.iter().all(|x| *x == -1.0 || *x == 0.0 || *x == 1.0), + "entries must be sparse ternary, d={d}" + ); + } + } + + #[test] + fn test_fastrp_init_varies_with_id_and_seed() { + let d = 256; + let mut base = vec![0.0f32; d]; + fastrp_init_fill(7, 42, d, &mut base); + + let mut other = vec![0.0f32; d]; + fastrp_init_fill(8, 42, d, &mut other); + assert_ne!(base, other, "different ids must init differently"); + + fastrp_init_fill(7, 43, d, &mut other); + assert_ne!(base, other, "different seeds must init differently"); + } + + #[test] + fn test_fastrp_init_density_is_sparse() { + // density 1/sqrt(q) per sign with q = largest power of two <= d; + // for d=1024 -> q=1024 -> ~3/4 of the entries are zero. Both extremes + // are astronomically unlikely ((3/4)^1024 and (1/4)^1024). + let d = 1024; + let mut v = vec![0.0f32; d]; + fastrp_init_fill(3, 42, d, &mut v); + let non_zero = v.iter().filter(|x| **x != 0.0).count(); + assert!(non_zero > 0, "expected some ±1 entries"); + assert!(non_zero < d, "expected some zero entries"); + } }