From 95d7319885bbe8c3d48669a3e7f92623e9a60345 Mon Sep 17 00:00:00 2001 From: Rodion Suvorov Date: Wed, 24 Jun 2026 15:33:50 +0300 Subject: [PATCH 1/8] feat: collect metada from mm format --- .gitignore | 2 + pathrex-sys/build.rs | 1 + pathrex-sys/src/lagraph_sys_generated.rs | 3 ++ pathrex/src/cli/dispatch.rs | 2 +- pathrex/src/graph/inmemory.rs | 51 +++++++++++++++++++++--- pathrex/src/graph/wrappers.rs | 35 +++++++++++++++- 6 files changed, 87 insertions(+), 7 deletions(-) diff --git a/.gitignore b/.gitignore index 0728338..c7c7c3d 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ +.vscode +Cargo.lock # Generated by Cargo # will have compiled files and executables debug diff --git a/pathrex-sys/build.rs b/pathrex-sys/build.rs index 7b7ce96..b4cfeb8 100644 --- a/pathrex-sys/build.rs +++ b/pathrex-sys/build.rs @@ -292,6 +292,7 @@ fn regenerate_bindings(graphblas_install: &Path) { .allowlist_item("GrB_Info") .allowlist_function("GrB_Matrix_new") .allowlist_function("GrB_Matrix_nvals") + .allowlist_function("GrB_Matrix_nrows") .allowlist_function("GrB_Matrix_dup") .allowlist_function("GrB_Matrix_free") .allowlist_function("GrB_Matrix_extractElement_BOOL") diff --git a/pathrex-sys/src/lagraph_sys_generated.rs b/pathrex-sys/src/lagraph_sys_generated.rs index 1a9188f..b0bcca7 100644 --- a/pathrex-sys/src/lagraph_sys_generated.rs +++ b/pathrex-sys/src/lagraph_sys_generated.rs @@ -158,6 +158,9 @@ unsafe extern "C" { unsafe extern "C" { pub fn GrB_Matrix_dup(C: *mut GrB_Matrix, A: GrB_Matrix) -> GrB_Info; } +unsafe extern "C" { + pub fn GrB_Matrix_nrows(nrows: *mut GrB_Index, A: GrB_Matrix) -> GrB_Info; +} unsafe extern "C" { pub fn GrB_Matrix_nvals(nvals: *mut GrB_Index, A: GrB_Matrix) -> GrB_Info; } diff --git a/pathrex/src/cli/dispatch.rs b/pathrex/src/cli/dispatch.rs index 2701460..cdcdfa2 100644 --- a/pathrex/src/cli/dispatch.rs +++ b/pathrex/src/cli/dispatch.rs @@ -1,5 +1,4 @@ //! Typed dispatch from CLI algorithm choices to concrete evaluators. - use crate::cli::args::{Algo, BenchArgs, QueryArgs}; use crate::cli::bench::error::BenchError; use crate::cli::bench::runner::run_bench_for_evaluator; @@ -8,6 +7,7 @@ use crate::cli::loader::LoadedQuery; use crate::cli::output::QueryResult; use crate::cli::query::run_query_for_evaluator; use crate::graph::InMemoryGraph; + use crate::rpq::nfarpq::NfaRpqEvaluator; use crate::rpq::rpqmatrix::RpqMatrixEvaluator; diff --git a/pathrex/src/graph/inmemory.rs b/pathrex/src/graph/inmemory.rs index 9121101..f2e90d8 100644 --- a/pathrex/src/graph/inmemory.rs +++ b/pathrex/src/graph/inmemory.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use std::{collections::HashMap, io::Read}; +use libc::uinput_setup; use rayon::prelude::*; use crate::formats::mm::{apply_base_iri, parse_index_map}; @@ -42,6 +43,7 @@ pub struct InMemoryBuilder { next_id: usize, label_buffers: HashMap>, prebuilt: HashMap, + metadata: HashMap, } impl InMemoryBuilder { @@ -52,6 +54,7 @@ impl InMemoryBuilder { next_id: 0, label_buffers: HashMap::new(), prebuilt: HashMap::new(), + metadata: HashMap::new(), } } @@ -111,6 +114,14 @@ impl InMemoryBuilder { ) { self.prebuilt.extend(iter); } + + /// Bulk-install pre-wrapped `(label, MatrixMetadata)` pairs into `metadata`. + pub(crate) fn extend_metadata>( + &mut self, + iter: I, + ) { + self.metadata.extend(iter); + } } impl GraphBuilder for InMemoryBuilder { @@ -162,11 +173,13 @@ impl GraphBuilder for InMemoryBuilder { for (label, lg) in built { graphs.insert(label, Arc::new(lg)); } - Ok(InMemoryGraph { node_to_id: self.node_to_id, id_to_node: self.id_to_node, graphs, + metadata: GraphMetadata { + label_to_data: self.metadata, + }, }) } } @@ -176,6 +189,17 @@ pub struct InMemoryGraph { node_to_id: HashMap, id_to_node: HashMap, graphs: HashMap>, + metadata: GraphMetadata, +} + +pub struct GraphMetadata { + label_to_data: HashMap, +} +pub struct MatrixMetadata { + dimension: usize, + nonzero_rows: usize, + nonzero_cols: usize, + nvals: usize, } impl GraphDecomposition for InMemoryGraph { @@ -244,22 +268,39 @@ impl GraphSource for MatrixMarket { let _scope = ThreadScope::enter(outer, inner)?; let mm_dir = self.dir.clone(); - let loaded: Vec<(String, LagraphGraph)> = edge_by_idx + let loaded: Vec<(String, LagraphGraph, MatrixMetadata)> = edge_by_idx .into_par_iter() .map( - |(idx, label)| -> Result<(String, LagraphGraph), GraphError> { + |(idx, label)| -> Result<(String, LagraphGraph, MatrixMetadata), GraphError> { let path = mm_dir.join(format!("{}.txt", idx)); let matrix = load_mm_file(&path)?; let lg = LagraphGraph::from_matrix( matrix, LAGraph_Kind::LAGraph_ADJACENCY_DIRECTED, )?; - Ok((label, lg)) + let dimension = lg.dimension()?; + let nonzero_rows = lg.nonzero_rows()?; + let nonzero_cols = lg.nonzero_cols()?; + let nvals = lg.nvals()?; + let metadata = MatrixMetadata { + dimension: dimension as usize, + nonzero_rows: nonzero_rows, + nonzero_cols: nonzero_cols, + nvals: nvals as usize, + }; + Ok((label, lg, metadata)) }, ) .collect::, GraphError>>()?; - builder.extend_prebuilt(loaded); + let mut loaded_graphs = vec![]; + let mut loaded_metadata = vec![]; + for (_i, (name, graph, meta)) in loaded.into_iter().enumerate() { + loaded_graphs.push((name.clone(), graph)); + loaded_metadata.push((name, meta)); + } + builder.extend_prebuilt(loaded_graphs); + builder.extend_metadata(loaded_metadata); Ok(builder) } diff --git a/pathrex/src/graph/wrappers.rs b/pathrex/src/graph/wrappers.rs index e97cfc5..6d5d5c0 100644 --- a/pathrex/src/graph/wrappers.rs +++ b/pathrex/src/graph/wrappers.rs @@ -11,7 +11,11 @@ use std::os::fd::IntoRawFd; use std::path::Path; use std::sync::Once; -use crate::{grb_ok, la_ok, lagraph_sys::*}; +use crate::{ + graph::wrappers::ReduceType::{ByCols, ByRows}, + grb_ok, la_ok, + lagraph_sys::*, +}; use super::GraphError; @@ -69,6 +73,10 @@ impl Drop for ThreadScope { } } +enum ReduceType { + ByRows, + ByCols, +} #[derive(Debug)] pub struct LagraphGraph { pub(crate) inner: LAGraph_Graph, @@ -147,6 +155,17 @@ impl LagraphGraph { unsafe { la_ok!(LAGraph_CheckGraph(self.inner)) } } + /// Number of rows and cols in the underlying adjacency matrix. + pub fn dimension(&self) -> Result { + if self.inner.is_null() { + return Ok(0); + } + let matrix: GrB_Matrix = unsafe { (*self.inner).A }; + let mut dimension: GrB_Index = 0; + unsafe { grb_ok!(GrB_Matrix_nrows(&mut dimension, matrix))? }; + Ok(dimension) + } + /// Number of stored (non-zero) values in the underlying adjacency matrix. pub fn nvals(&self) -> Result { if self.inner.is_null() { @@ -157,6 +176,20 @@ impl LagraphGraph { unsafe { grb_ok!(GrB_Matrix_nvals(&mut nvals, matrix))? }; Ok(nvals) } + + pub fn nonzero_cols(&self) -> Result { + let matrix: GrB_Matrix = unsafe { (*self.inner).A }; + let res: usize = 0; + unsafe { LAGraph_RPQMatrix_reduce(&mut (res as u64), matrix, ByRows as u8) }; + Ok(res) + } + + pub fn nonzero_rows(&self) -> Result { + let matrix: GrB_Matrix = unsafe { (*self.inner).A }; + let res: usize = 0; + unsafe { LAGraph_RPQMatrix_reduce(&mut (res as u64), matrix, ByCols as u8) }; + Ok(res) + } } impl Drop for LagraphGraph { From bf06e96845ca0ab95e2bba4bf21f90ada0dcb136 Mon Sep 17 00:00:00 2001 From: Rodion Suvorov Date: Tue, 7 Jul 2026 20:19:33 +0300 Subject: [PATCH 2/8] feat: [wip] add optimizer based on metadata --- pathrex/Cargo.toml | 1 + pathrex/src/graph/inmemory.rs | 27 +++- pathrex/src/graph/mod.rs | 4 + pathrex/src/graph/wrappers.rs | 2 +- pathrex/src/rpq/rpqmatrix.rs | 293 ++++++++++++++++++++++++++++++++-- pathrex/src/utils.rs | 3 + 6 files changed, 307 insertions(+), 23 deletions(-) diff --git a/pathrex/Cargo.toml b/pathrex/Cargo.toml index 6401f1e..b439e61 100644 --- a/pathrex/Cargo.toml +++ b/pathrex/Cargo.toml @@ -38,6 +38,7 @@ serde_json = { version = "1", optional = true } chrono = { version = "0.4", features = ["serde"], optional = true } criterion = { version = "0.5", optional = true } tempfile = { version = "3", optional = true } +rand = "0.10.2" [features] default = [] diff --git a/pathrex/src/graph/inmemory.rs b/pathrex/src/graph/inmemory.rs index f2e90d8..043e536 100644 --- a/pathrex/src/graph/inmemory.rs +++ b/pathrex/src/graph/inmemory.rs @@ -1,7 +1,6 @@ use std::sync::Arc; use std::{collections::HashMap, io::Read}; -use libc::uinput_setup; use rayon::prelude::*; use crate::formats::mm::{apply_base_iri, parse_index_map}; @@ -191,15 +190,21 @@ pub struct InMemoryGraph { graphs: HashMap>, metadata: GraphMetadata, } - pub struct GraphMetadata { label_to_data: HashMap, } + +impl GraphMetadata { + pub fn matrix(&self, label: &str) -> Option<&MatrixMetadata> { + self.label_to_data.get(label) + } +} + pub struct MatrixMetadata { - dimension: usize, - nonzero_rows: usize, - nonzero_cols: usize, - nvals: usize, + pub dimension: usize, + pub nonzero_rows: usize, + pub nonzero_cols: usize, + pub nvals: usize, } impl GraphDecomposition for InMemoryGraph { @@ -221,6 +226,10 @@ impl GraphDecomposition for InMemoryGraph { fn num_nodes(&self) -> usize { self.id_to_node.len() } + + fn get_metadata(&self) -> Option<&GraphMetadata> { + Some(&self.metadata) + } } impl InMemoryGraph { @@ -228,6 +237,9 @@ impl InMemoryGraph { pub fn num_labels(&self) -> usize { self.graphs.len() } + pub fn metadata(&self, label: &str) -> Option<&MatrixMetadata> { + self.metadata.label_to_data.get(label) + } } impl GraphSource for Csv { @@ -442,4 +454,7 @@ mod tests { assert!(graph.get_graph("http://example.org/knows").is_ok()); assert!(graph.get_graph("http://example.org/likes").is_ok()); } + + #[test] + fn test_metadata_from_mm() {} } diff --git a/pathrex/src/graph/mod.rs b/pathrex/src/graph/mod.rs index 0922459..62e2a2c 100644 --- a/pathrex/src/graph/mod.rs +++ b/pathrex/src/graph/mod.rs @@ -10,6 +10,7 @@ pub(crate) use wrappers::{ThreadScope, compute_outer_inner, ensure_grb_init}; use std::marker::PhantomData; use std::sync::Arc; +use crate::graph::inmemory::GraphMetadata; use crate::lagraph_sys::GrB_Info; use thiserror::Error; @@ -83,6 +84,9 @@ pub trait GraphDecomposition { /// Translates a matrix index back to a string ID. fn get_node_name(&self, mapped_id: usize) -> Option; fn num_nodes(&self) -> usize; + fn get_metadata(&self) -> Option<&GraphMetadata> { + None + } } /// Associates a backend marker type with a concrete [`GraphBuilder`] and diff --git a/pathrex/src/graph/wrappers.rs b/pathrex/src/graph/wrappers.rs index 6d5d5c0..c684d60 100644 --- a/pathrex/src/graph/wrappers.rs +++ b/pathrex/src/graph/wrappers.rs @@ -73,7 +73,7 @@ impl Drop for ThreadScope { } } -enum ReduceType { +pub enum ReduceType { ByRows, ByCols, } diff --git a/pathrex/src/rpq/rpqmatrix.rs b/pathrex/src/rpq/rpqmatrix.rs index 1462abc..d66fa57 100644 --- a/pathrex/src/rpq/rpqmatrix.rs +++ b/pathrex/src/rpq/rpqmatrix.rs @@ -1,50 +1,292 @@ //! Plan-based RPQ evaluation using `LAGraph_RPQMatrix`. use std::ptr::null_mut; +use std::{cmp::Ordering, fmt::Display, str::FromStr}; -use egg::{Id, RecExpr, define_language}; +use egg::{CostFunction, Id, RecExpr, define_language, rewrite}; use crate::eval::{Evaluator, PreparedEvaluator, ResultCount}; +use crate::graph::wrappers::ReduceType::ByCols; use crate::graph::{GraphDecomposition, GraphError, GraphblasMatrix}; use crate::lagraph_sys::*; use crate::rpq::{Endpoint, PathExpr, RpqError, RpqQuery}; use crate::{grb_ok, la_ok}; -const RPQMATRIX_REDUCE_BY_COL: u8 = 1; +#[derive(Clone, Hash, Ord, Eq, PartialEq, PartialOrd, Debug)] +struct LabelMeta { + pub name: String, + pub nvals: usize, + pub nonzero_rows: usize, + pub nonzero_cols: usize, +} + +impl FromStr for LabelMeta { + type Err = ::Err; + // This is needed for the builtin egg parser. Only used in tests. + fn from_str(s: &str) -> Result { + Ok(LabelMeta { + name: "-".to_string(), + nvals: s.parse()?, + nonzero_rows: s.parse()?, + nonzero_cols: s.parse()?, + }) + } +} + +impl Display for LabelMeta { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "({}, {})", self.name, self.nvals) + } +} define_language! { pub enum RpqPlan { - Label(String), + Label(LabelMeta), NamedVertex(String), "/" = Seq([Id; 2]), "|" = Alt([Id; 2]), "*" = Star([Id; 1]), + "l*" = LStar([Id; 2]), + "*r" = RStar([Id; 2]), + } +} + +pub fn make_rules() -> Vec> { + vec![ + rewrite!("assoc-sec-1"; "(/ ?a (/ ?b ?c))" => "(/ (/ ?a ?b) ?c)"), + rewrite!("assoc-sec-2"; "(/ (/ ?a ?b) ?c)" => "(/ ?a (/ ?b ?c))"), + rewrite!("commute-alt"; "(| ?a ?b)" => "(| ?b ?a)"), + rewrite!("assoc-alt"; "(| ?a (| ?b ?c))" => "(| (| ?a ?b) ?c)"), + rewrite!("distribute-1"; "(/ ?a (| ?b ?c))" => "(| (/ ?a ?b) (/ ?a ?c))"), + rewrite!("distribute-2"; "(/ (| ?a ?b) ?c)" => "(| (/ ?a ?c) (/ ?b ?c))"), + rewrite!("distribute-3"; "(| (/ ?a ?b) (/ ?a ?c))" => "(/ ?a (| ?b ?c))"), + rewrite!("distribute-4"; "(| (/ ?a ?c) (/ ?b ?c))" => "(/ (| ?a ?b) ?c)"), + rewrite!("build-lstar"; "(/ (* ?a) ?b)" => "(l* ?a ?b)"), + rewrite!("build-rstar"; "(/ ?a (* ?b))" => "(*r ?a ?b)"), + ] +} + +pub struct RandomCostFn; +impl CostFunction for RandomCostFn { + type Cost = f64; + fn cost(&mut self, _enode: &RpqPlan, _costs: C) -> Self::Cost + where + C: FnMut(Id) -> Self::Cost, + { + rand::random() + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct CardCost { + pub score: f64, + pub nnz: f64, + pub nnz_r: f64, + pub nnz_c: f64, +} + +impl Eq for CardCost {} + +impl PartialOrd for CardCost { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for CardCost { + fn cmp(&self, other: &Self) -> Ordering { + match self.score.total_cmp(&other.score) { + Ordering::Equal => {} + ord => return ord, + } + match self.nnz.total_cmp(&other.nnz) { + Ordering::Equal => {} + ord => return ord, + } + match self.nnz_r.total_cmp(&other.nnz_r) { + Ordering::Equal => {} + ord => return ord, + } + self.nnz_c.total_cmp(&other.nnz_c) } } -fn to_expr_aux(path: &PathExpr, expr: &mut RecExpr) -> Result { +pub struct CardinalityCostFn { + pub n: f64, + pub star_penalty: f64, + pub lr_multiplier: f64, +} + +// TODO: check value intervals +impl CostFunction for CardinalityCostFn { + type Cost = CardCost; + + fn cost(&mut self, enode: &RpqPlan, mut costs: C) -> Self::Cost + where + C: FnMut(Id) -> Self::Cost, + { + match enode { + RpqPlan::NamedVertex(_name) => CardCost { + score: 0.0, + nnz: 1 as f64, + nnz_r: 1 as f64, + nnz_c: 1 as f64, + }, + RpqPlan::Label(meta) => CardCost { + score: 0.0, + nnz: meta.nvals as f64, + nnz_r: meta.nonzero_rows as f64, + nnz_c: meta.nonzero_cols as f64, + }, + + RpqPlan::Seq([a, b]) => { + let ca = costs(*a); + let cb = costs(*b); + + let denom = ca.nnz_r.max(cb.nnz_c).max(1.0); + let op_cost = (ca.nnz * cb.nnz) / denom; + let score = ca.score + cb.score + op_cost; + + let nnz_est = ca.nnz * cb.nnz / (self.n * self.n); + + CardCost { + score, + nnz: nnz_est, + nnz_r: ca.nnz_r.min(self.n), // TODO: better reduce estimators + nnz_c: cb.nnz_c.min(self.n), // TODO: better reduce estimators + } + } + + RpqPlan::Alt([a, b]) => { + let ca = costs(*a); + let cb = costs(*b); + + let overlap = (ca.nnz * cb.nnz) / (self.n * self.n); + let op_cost = ca.nnz + cb.nnz - overlap; + let score = ca.score + cb.score + op_cost; + + let nnz_est = (ca.nnz + cb.nnz - overlap).min(self.n * self.n).max(0.0); + + let nnz_r_est = (ca.nnz_r + cb.nnz_r - (ca.nnz_r * cb.nnz_r) / self.n) + .min(self.n) + .max(0.0); + + let nnz_c_est = (ca.nnz_c + cb.nnz_c - (ca.nnz_c * cb.nnz_c) / self.n) + .min(self.n) + .max(0.0); + + CardCost { + score, + nnz: nnz_est, + nnz_r: nnz_r_est, + nnz_c: nnz_c_est, + } + } + + RpqPlan::Star([a]) => { + let ca = costs(*a); + + let penalty = self.star_penalty * ca.nnz.max(1.0); + let score = ca.score + penalty; + + CardCost { + score, + nnz: self.n * self.n, + nnz_r: self.n, + nnz_c: self.n, + } + } + + RpqPlan::LStar([a, b]) => { + let ca = costs(*a); + let cb = costs(*b); + + let denom = ca.nnz_r.max(cb.nnz_c).max(1.0); + let base = (ca.nnz * cb.nnz) / denom; + let op_cost = self.lr_multiplier * base; + let score = ca.score + cb.score + op_cost; + + let nnz_est = self.lr_multiplier * ca.nnz * cb.nnz / (self.n * self.n); + + CardCost { + score, + nnz: nnz_est, + nnz_r: ca.nnz_r.min(self.n), // TODO: better reduce estimators + nnz_c: cb.nnz_c.min(self.n), // TODO: better reduce estimators + } + } + + RpqPlan::RStar([a, b]) => { + let ca = costs(*a); + let cb = costs(*b); + + let denom = ca.nnz_r.max(cb.nnz_c).max(1.0); + let base = (ca.nnz * cb.nnz) / denom; + + let op_cost = self.lr_multiplier * base; + let score = ca.score + cb.score + op_cost; + + let nnz_est = self.lr_multiplier * ca.nnz * cb.nnz / (self.n * self.n); + + CardCost { + score, + nnz: nnz_est, + nnz_r: ca.nnz_r.min(self.n), // TODO: better reduce estimators + nnz_c: cb.nnz_c.min(self.n), // TODO: better reduce estimators + } + } + } + } +} + +fn label_meta(label: &str, graph: &G) -> Result { + if let Some(metadata) = graph.get_metadata().and_then(|m| m.matrix(label)) { + return Ok(LabelMeta { + name: label.to_owned(), + nvals: metadata.nvals, + nonzero_rows: metadata.nonzero_rows, + nonzero_cols: metadata.nonzero_cols, + }); + } + + // TODO: maybe create optimized (for mm format) and nonoptimized (for other formats) plans + let lg = graph.get_graph(label)?; + let nvals = lg.nvals()? as usize; + Ok(LabelMeta { + name: label.to_owned(), + nvals, + nonzero_rows: nvals, + nonzero_cols: nvals, + }) +} + +fn to_expr_aux( + path: &PathExpr, + expr: &mut RecExpr, + graph: &G, +) -> Result { match path { - PathExpr::Label(label) => Ok(expr.add(RpqPlan::Label(label.clone()))), + PathExpr::Label(label) => Ok(expr.add(RpqPlan::Label(label_meta(label, graph)?))), PathExpr::Sequence(lhs, rhs) => { - let l = to_expr_aux(lhs, expr)?; - let r = to_expr_aux(rhs, expr)?; + let l = to_expr_aux(lhs, expr, graph)?; + let r = to_expr_aux(rhs, expr, graph)?; Ok(expr.add(RpqPlan::Seq([l, r]))) } PathExpr::Alternative(lhs, rhs) => { - let l = to_expr_aux(lhs, expr)?; - let r = to_expr_aux(rhs, expr)?; + let l = to_expr_aux(lhs, expr, graph)?; + let r = to_expr_aux(rhs, expr, graph)?; Ok(expr.add(RpqPlan::Alt([l, r]))) } PathExpr::ZeroOrMore(inner) => { - let i = to_expr_aux(inner, expr)?; + let i = to_expr_aux(inner, expr, graph)?; Ok(expr.add(RpqPlan::Star([i]))) } PathExpr::OneOrMore(inner) => { - let e = to_expr_aux(inner, expr)?; + let e = to_expr_aux(inner, expr, graph)?; let s = expr.add(RpqPlan::Star([e])); Ok(expr.add(RpqPlan::Seq([e, s]))) } @@ -57,9 +299,12 @@ fn to_expr_aux(path: &PathExpr, expr: &mut RecExpr) -> Result`]. -pub fn query_to_expr(query: &RpqQuery) -> Result, RpqError> { +pub fn query_to_expr( + query: &RpqQuery, + graph: &G, +) -> Result, RpqError> { let mut expr = RecExpr::default(); - let path_root = to_expr_aux(&query.path, &mut expr)?; + let path_root = to_expr_aux(&query.path, &mut expr, graph)?; let _root = match (&query.subject, &query.object) { (Endpoint::Variable(_), Endpoint::Variable(_)) => path_root, @@ -105,7 +350,7 @@ pub fn materialize( for (id, node) in expr.as_ref().iter().enumerate() { plans[id] = match node { RpqPlan::Label(label) => { - let lg = graph.get_graph(label)?; + let lg = graph.get_graph(&label.name)?; let mat = unsafe { (*lg.inner).A }; RPQMatrixPlan { op: RPQMatrixOp::RPQ_MATRIX_OP_LABEL, @@ -164,6 +409,22 @@ pub fn materialize( mat: null_mut(), res_mat: null_mut(), }, + + RpqPlan::LStar([l, r]) => RPQMatrixPlan { + op: RPQMatrixOp::RPQ_MATRIX_OP_KLEENE_L, + lhs: unsafe { plans.as_mut_ptr().add(usize::from(*l)) }, + rhs: unsafe { plans.as_mut_ptr().add(usize::from(*r)) }, + mat: null_mut(), + res_mat: null_mut(), + }, + + RpqPlan::RStar([l, r]) => RPQMatrixPlan { + op: RPQMatrixOp::RPQ_MATRIX_OP_KLEENE_R, + lhs: unsafe { plans.as_mut_ptr().add(usize::from(*l)) }, + rhs: unsafe { plans.as_mut_ptr().add(usize::from(*r)) }, + mat: null_mut(), + res_mat: null_mut(), + }, }; } @@ -186,7 +447,7 @@ impl RpqMatrixResult { grb_ok!(LAGraph_RPQMatrix_reduce( &mut count, self.matrix.inner, - RPQMATRIX_REDUCE_BY_COL, + ByCols as u8, ))? }; Ok(count as u64) @@ -254,7 +515,7 @@ impl Evaluator for RpqMatrixEvaluator { query: &RpqQuery, graph: &G, ) -> Result { - let expr = query_to_expr(query)?; + let expr = query_to_expr(query, graph)?; let (plans, owned_matrices) = materialize(&expr, graph)?; Ok(PreparedRpqMatrix { diff --git a/pathrex/src/utils.rs b/pathrex/src/utils.rs index 30477fb..7b9ea4b 100644 --- a/pathrex/src/utils.rs +++ b/pathrex/src/utils.rs @@ -26,6 +26,9 @@ impl GraphDecomposition for CountOutput { fn num_nodes(&self) -> usize { self.0 } + fn get_metadata(&self) -> Option<&inmemory::GraphMetadata> { + None + } } /// A minimal [`GraphBuilder`] that counts pushed edges and produces a [`CountOutput`]. From 632c3c18e280116380392ec71b81a0a7cce72446 Mon Sep 17 00:00:00 2001 From: Rodion Suvorov Date: Wed, 15 Jul 2026 18:16:10 +0300 Subject: [PATCH 3/8] feat: create rpqmatrix submodule --- pathrex/src/cli/dispatch.rs | 6 +- pathrex/src/rpq/rpqmatrix.rs | 559 -------------------------- pathrex/src/rpq/rpqmatrix/cost.rs | 178 ++++++++ pathrex/src/rpq/rpqmatrix/eval.rs | 97 +++++ pathrex/src/rpq/rpqmatrix/expr.rs | 205 ++++++++++ pathrex/src/rpq/rpqmatrix/mod.rs | 12 + pathrex/src/rpq/rpqmatrix/optimize.rs | 35 ++ pathrex/src/rpq/rpqmatrix/plan.rs | 57 +++ pathrex/src/rpq/rpqmatrix/result.rs | 78 ++++ pathrex/tests/rpqmatrix_tests.rs | 41 +- 10 files changed, 686 insertions(+), 582 deletions(-) delete mode 100644 pathrex/src/rpq/rpqmatrix.rs create mode 100644 pathrex/src/rpq/rpqmatrix/cost.rs create mode 100644 pathrex/src/rpq/rpqmatrix/eval.rs create mode 100644 pathrex/src/rpq/rpqmatrix/expr.rs create mode 100644 pathrex/src/rpq/rpqmatrix/mod.rs create mode 100644 pathrex/src/rpq/rpqmatrix/optimize.rs create mode 100644 pathrex/src/rpq/rpqmatrix/plan.rs create mode 100644 pathrex/src/rpq/rpqmatrix/result.rs diff --git a/pathrex/src/cli/dispatch.rs b/pathrex/src/cli/dispatch.rs index cdcdfa2..cd19ee6 100644 --- a/pathrex/src/cli/dispatch.rs +++ b/pathrex/src/cli/dispatch.rs @@ -9,7 +9,7 @@ use crate::cli::query::run_query_for_evaluator; use crate::graph::InMemoryGraph; use crate::rpq::nfarpq::NfaRpqEvaluator; -use crate::rpq::rpqmatrix::RpqMatrixEvaluator; +use crate::rpq::rpqmatrix::eval::RpqMatrixEvaluator; fn merge_results(all: &mut Vec, per_algo: Vec) { for result in per_algo { @@ -35,7 +35,7 @@ pub fn dispatch_query( let name = algo.to_string(); let per_algo = match algo { Algo::NfaRpq => run_query_for_evaluator(&name, NfaRpqEvaluator, graph, queries), - Algo::Rpqmatrix => run_query_for_evaluator(&name, RpqMatrixEvaluator, graph, queries), + Algo::Rpqmatrix => run_query_for_evaluator(&name, RpqMatrixEvaluator::default(), graph, queries), }; merge_results(&mut all, per_algo); } @@ -67,7 +67,7 @@ pub fn dispatch_bench( args, algo, &name, - RpqMatrixEvaluator, + RpqMatrixEvaluator::default(), graph, queries, checkpointer, diff --git a/pathrex/src/rpq/rpqmatrix.rs b/pathrex/src/rpq/rpqmatrix.rs deleted file mode 100644 index d66fa57..0000000 --- a/pathrex/src/rpq/rpqmatrix.rs +++ /dev/null @@ -1,559 +0,0 @@ -//! Plan-based RPQ evaluation using `LAGraph_RPQMatrix`. - -use std::ptr::null_mut; -use std::{cmp::Ordering, fmt::Display, str::FromStr}; - -use egg::{CostFunction, Id, RecExpr, define_language, rewrite}; - -use crate::eval::{Evaluator, PreparedEvaluator, ResultCount}; -use crate::graph::wrappers::ReduceType::ByCols; -use crate::graph::{GraphDecomposition, GraphError, GraphblasMatrix}; -use crate::lagraph_sys::*; -use crate::rpq::{Endpoint, PathExpr, RpqError, RpqQuery}; -use crate::{grb_ok, la_ok}; - -#[derive(Clone, Hash, Ord, Eq, PartialEq, PartialOrd, Debug)] -struct LabelMeta { - pub name: String, - pub nvals: usize, - pub nonzero_rows: usize, - pub nonzero_cols: usize, -} - -impl FromStr for LabelMeta { - type Err = ::Err; - // This is needed for the builtin egg parser. Only used in tests. - fn from_str(s: &str) -> Result { - Ok(LabelMeta { - name: "-".to_string(), - nvals: s.parse()?, - nonzero_rows: s.parse()?, - nonzero_cols: s.parse()?, - }) - } -} - -impl Display for LabelMeta { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "({}, {})", self.name, self.nvals) - } -} - -define_language! { - pub enum RpqPlan { - Label(LabelMeta), - NamedVertex(String), - "/" = Seq([Id; 2]), - "|" = Alt([Id; 2]), - "*" = Star([Id; 1]), - "l*" = LStar([Id; 2]), - "*r" = RStar([Id; 2]), - } -} - -pub fn make_rules() -> Vec> { - vec![ - rewrite!("assoc-sec-1"; "(/ ?a (/ ?b ?c))" => "(/ (/ ?a ?b) ?c)"), - rewrite!("assoc-sec-2"; "(/ (/ ?a ?b) ?c)" => "(/ ?a (/ ?b ?c))"), - rewrite!("commute-alt"; "(| ?a ?b)" => "(| ?b ?a)"), - rewrite!("assoc-alt"; "(| ?a (| ?b ?c))" => "(| (| ?a ?b) ?c)"), - rewrite!("distribute-1"; "(/ ?a (| ?b ?c))" => "(| (/ ?a ?b) (/ ?a ?c))"), - rewrite!("distribute-2"; "(/ (| ?a ?b) ?c)" => "(| (/ ?a ?c) (/ ?b ?c))"), - rewrite!("distribute-3"; "(| (/ ?a ?b) (/ ?a ?c))" => "(/ ?a (| ?b ?c))"), - rewrite!("distribute-4"; "(| (/ ?a ?c) (/ ?b ?c))" => "(/ (| ?a ?b) ?c)"), - rewrite!("build-lstar"; "(/ (* ?a) ?b)" => "(l* ?a ?b)"), - rewrite!("build-rstar"; "(/ ?a (* ?b))" => "(*r ?a ?b)"), - ] -} - -pub struct RandomCostFn; -impl CostFunction for RandomCostFn { - type Cost = f64; - fn cost(&mut self, _enode: &RpqPlan, _costs: C) -> Self::Cost - where - C: FnMut(Id) -> Self::Cost, - { - rand::random() - } -} - -#[derive(Clone, Debug, PartialEq)] -pub struct CardCost { - pub score: f64, - pub nnz: f64, - pub nnz_r: f64, - pub nnz_c: f64, -} - -impl Eq for CardCost {} - -impl PartialOrd for CardCost { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for CardCost { - fn cmp(&self, other: &Self) -> Ordering { - match self.score.total_cmp(&other.score) { - Ordering::Equal => {} - ord => return ord, - } - match self.nnz.total_cmp(&other.nnz) { - Ordering::Equal => {} - ord => return ord, - } - match self.nnz_r.total_cmp(&other.nnz_r) { - Ordering::Equal => {} - ord => return ord, - } - self.nnz_c.total_cmp(&other.nnz_c) - } -} - -pub struct CardinalityCostFn { - pub n: f64, - pub star_penalty: f64, - pub lr_multiplier: f64, -} - -// TODO: check value intervals -impl CostFunction for CardinalityCostFn { - type Cost = CardCost; - - fn cost(&mut self, enode: &RpqPlan, mut costs: C) -> Self::Cost - where - C: FnMut(Id) -> Self::Cost, - { - match enode { - RpqPlan::NamedVertex(_name) => CardCost { - score: 0.0, - nnz: 1 as f64, - nnz_r: 1 as f64, - nnz_c: 1 as f64, - }, - RpqPlan::Label(meta) => CardCost { - score: 0.0, - nnz: meta.nvals as f64, - nnz_r: meta.nonzero_rows as f64, - nnz_c: meta.nonzero_cols as f64, - }, - - RpqPlan::Seq([a, b]) => { - let ca = costs(*a); - let cb = costs(*b); - - let denom = ca.nnz_r.max(cb.nnz_c).max(1.0); - let op_cost = (ca.nnz * cb.nnz) / denom; - let score = ca.score + cb.score + op_cost; - - let nnz_est = ca.nnz * cb.nnz / (self.n * self.n); - - CardCost { - score, - nnz: nnz_est, - nnz_r: ca.nnz_r.min(self.n), // TODO: better reduce estimators - nnz_c: cb.nnz_c.min(self.n), // TODO: better reduce estimators - } - } - - RpqPlan::Alt([a, b]) => { - let ca = costs(*a); - let cb = costs(*b); - - let overlap = (ca.nnz * cb.nnz) / (self.n * self.n); - let op_cost = ca.nnz + cb.nnz - overlap; - let score = ca.score + cb.score + op_cost; - - let nnz_est = (ca.nnz + cb.nnz - overlap).min(self.n * self.n).max(0.0); - - let nnz_r_est = (ca.nnz_r + cb.nnz_r - (ca.nnz_r * cb.nnz_r) / self.n) - .min(self.n) - .max(0.0); - - let nnz_c_est = (ca.nnz_c + cb.nnz_c - (ca.nnz_c * cb.nnz_c) / self.n) - .min(self.n) - .max(0.0); - - CardCost { - score, - nnz: nnz_est, - nnz_r: nnz_r_est, - nnz_c: nnz_c_est, - } - } - - RpqPlan::Star([a]) => { - let ca = costs(*a); - - let penalty = self.star_penalty * ca.nnz.max(1.0); - let score = ca.score + penalty; - - CardCost { - score, - nnz: self.n * self.n, - nnz_r: self.n, - nnz_c: self.n, - } - } - - RpqPlan::LStar([a, b]) => { - let ca = costs(*a); - let cb = costs(*b); - - let denom = ca.nnz_r.max(cb.nnz_c).max(1.0); - let base = (ca.nnz * cb.nnz) / denom; - let op_cost = self.lr_multiplier * base; - let score = ca.score + cb.score + op_cost; - - let nnz_est = self.lr_multiplier * ca.nnz * cb.nnz / (self.n * self.n); - - CardCost { - score, - nnz: nnz_est, - nnz_r: ca.nnz_r.min(self.n), // TODO: better reduce estimators - nnz_c: cb.nnz_c.min(self.n), // TODO: better reduce estimators - } - } - - RpqPlan::RStar([a, b]) => { - let ca = costs(*a); - let cb = costs(*b); - - let denom = ca.nnz_r.max(cb.nnz_c).max(1.0); - let base = (ca.nnz * cb.nnz) / denom; - - let op_cost = self.lr_multiplier * base; - let score = ca.score + cb.score + op_cost; - - let nnz_est = self.lr_multiplier * ca.nnz * cb.nnz / (self.n * self.n); - - CardCost { - score, - nnz: nnz_est, - nnz_r: ca.nnz_r.min(self.n), // TODO: better reduce estimators - nnz_c: cb.nnz_c.min(self.n), // TODO: better reduce estimators - } - } - } - } -} - -fn label_meta(label: &str, graph: &G) -> Result { - if let Some(metadata) = graph.get_metadata().and_then(|m| m.matrix(label)) { - return Ok(LabelMeta { - name: label.to_owned(), - nvals: metadata.nvals, - nonzero_rows: metadata.nonzero_rows, - nonzero_cols: metadata.nonzero_cols, - }); - } - - // TODO: maybe create optimized (for mm format) and nonoptimized (for other formats) plans - let lg = graph.get_graph(label)?; - let nvals = lg.nvals()? as usize; - Ok(LabelMeta { - name: label.to_owned(), - nvals, - nonzero_rows: nvals, - nonzero_cols: nvals, - }) -} - -fn to_expr_aux( - path: &PathExpr, - expr: &mut RecExpr, - graph: &G, -) -> Result { - match path { - PathExpr::Label(label) => Ok(expr.add(RpqPlan::Label(label_meta(label, graph)?))), - - PathExpr::Sequence(lhs, rhs) => { - let l = to_expr_aux(lhs, expr, graph)?; - let r = to_expr_aux(rhs, expr, graph)?; - Ok(expr.add(RpqPlan::Seq([l, r]))) - } - - PathExpr::Alternative(lhs, rhs) => { - let l = to_expr_aux(lhs, expr, graph)?; - let r = to_expr_aux(rhs, expr, graph)?; - Ok(expr.add(RpqPlan::Alt([l, r]))) - } - - PathExpr::ZeroOrMore(inner) => { - let i = to_expr_aux(inner, expr, graph)?; - Ok(expr.add(RpqPlan::Star([i]))) - } - - PathExpr::OneOrMore(inner) => { - let e = to_expr_aux(inner, expr, graph)?; - let s = expr.add(RpqPlan::Star([e])); - Ok(expr.add(RpqPlan::Seq([e, s]))) - } - - PathExpr::ZeroOrOne(_) => Err(RpqError::UnsupportedPath( - "ZeroOrOne (?) is not supported by RPQMatrix".into(), - )), - } -} - -/// Compile a [`RpqQuery`] into -/// [`RecExpr`]. -pub fn query_to_expr( - query: &RpqQuery, - graph: &G, -) -> Result, RpqError> { - let mut expr = RecExpr::default(); - let path_root = to_expr_aux(&query.path, &mut expr, graph)?; - - let _root = match (&query.subject, &query.object) { - (Endpoint::Variable(_), Endpoint::Variable(_)) => path_root, - (Endpoint::Named(name), Endpoint::Variable(_)) => { - let diag = expr.add(RpqPlan::NamedVertex(name.clone())); - expr.add(RpqPlan::Seq([diag, path_root])) - } - (Endpoint::Variable(_), Endpoint::Named(name)) => { - let diag = expr.add(RpqPlan::NamedVertex(name.clone())); - expr.add(RpqPlan::Seq([path_root, diag])) - } - (Endpoint::Named(sub), Endpoint::Named(obj)) => { - let diag_sub = expr.add(RpqPlan::NamedVertex(sub.clone())); - let seq1 = expr.add(RpqPlan::Seq([diag_sub, path_root])); - let diag_obj = expr.add(RpqPlan::NamedVertex(obj.clone())); - expr.add(RpqPlan::Seq([seq1, diag_obj])) - } - }; - - Ok(expr) -} - -/// Convert a [`RecExpr`] into the flat [`RPQMatrixPlan`] array that -/// `LAGraph_RPQMatrix` expects. -/// -/// Returns the plan array and a list of owned diagonal matrices that must be -/// freed after evaluation. -pub fn materialize( - expr: &RecExpr, - graph: &G, -) -> Result<(Vec, Vec), RpqError> { - let null_plan = RPQMatrixPlan { - op: RPQMatrixOp::RPQ_MATRIX_OP_LABEL, - lhs: null_mut(), - rhs: null_mut(), - mat: null_mut(), - res_mat: null_mut(), - }; - let mut plans = vec![null_plan; expr.len()]; - let mut owned_matrices: Vec = Vec::new(); - let n = graph.num_nodes() as GrB_Index; - - for (id, node) in expr.as_ref().iter().enumerate() { - plans[id] = match node { - RpqPlan::Label(label) => { - let lg = graph.get_graph(&label.name)?; - let mat = unsafe { (*lg.inner).A }; - RPQMatrixPlan { - op: RPQMatrixOp::RPQ_MATRIX_OP_LABEL, - lhs: null_mut(), - rhs: null_mut(), - mat, - res_mat: null_mut(), - } - } - - RpqPlan::NamedVertex(name) => { - let vertex_id = graph - .get_node_id(name) - .ok_or_else(|| RpqError::VertexNotFound(name.clone()))? - as GrB_Index; - let mut mat: GrB_Matrix = null_mut(); - unsafe { - crate::graph::ensure_grb_init()?; - grb_ok!(LAGraph_RPQMatrix_label(&mut mat, vertex_id, n, n,))? - }; - if mat.is_null() { - return Err(RpqError::Graph(crate::graph::GraphError::GraphBlas( - GrB_Info::GrB_INVALID_VALUE, - ))); - } - owned_matrices.push(mat); - RPQMatrixPlan { - op: RPQMatrixOp::RPQ_MATRIX_OP_LABEL, - lhs: null_mut(), - rhs: null_mut(), - mat, - res_mat: null_mut(), - } - } - - RpqPlan::Seq([l, r]) => RPQMatrixPlan { - op: RPQMatrixOp::RPQ_MATRIX_OP_CONCAT, - lhs: unsafe { plans.as_mut_ptr().add(usize::from(*l)) }, - rhs: unsafe { plans.as_mut_ptr().add(usize::from(*r)) }, - mat: null_mut(), - res_mat: null_mut(), - }, - - RpqPlan::Alt([l, r]) => RPQMatrixPlan { - op: RPQMatrixOp::RPQ_MATRIX_OP_LOR, - lhs: unsafe { plans.as_mut_ptr().add(usize::from(*l)) }, - rhs: unsafe { plans.as_mut_ptr().add(usize::from(*r)) }, - mat: null_mut(), - res_mat: null_mut(), - }, - - RpqPlan::Star([i]) => RPQMatrixPlan { - op: RPQMatrixOp::RPQ_MATRIX_OP_KLEENE, - lhs: null_mut(), - rhs: unsafe { plans.as_mut_ptr().add(usize::from(*i)) }, - mat: null_mut(), - res_mat: null_mut(), - }, - - RpqPlan::LStar([l, r]) => RPQMatrixPlan { - op: RPQMatrixOp::RPQ_MATRIX_OP_KLEENE_L, - lhs: unsafe { plans.as_mut_ptr().add(usize::from(*l)) }, - rhs: unsafe { plans.as_mut_ptr().add(usize::from(*r)) }, - mat: null_mut(), - res_mat: null_mut(), - }, - - RpqPlan::RStar([l, r]) => RPQMatrixPlan { - op: RPQMatrixOp::RPQ_MATRIX_OP_KLEENE_R, - lhs: unsafe { plans.as_mut_ptr().add(usize::from(*l)) }, - rhs: unsafe { plans.as_mut_ptr().add(usize::from(*r)) }, - mat: null_mut(), - res_mat: null_mut(), - }, - }; - } - - Ok((plans, owned_matrices)) -} - -/// Output of [`RpqMatrixEvaluator`]: full path relation matrix and its nnz. -#[derive(Debug)] -pub struct RpqMatrixResult { - pub nnz: u64, - pub matrix: GraphblasMatrix, -} - -impl RpqMatrixResult { - /// Count distinct reachable target vertices by reducing the path relation - /// matrix to its non-empty columns. - pub fn reachable_target_count(&self) -> Result { - let mut count: GrB_Index = 0; - unsafe { - grb_ok!(LAGraph_RPQMatrix_reduce( - &mut count, - self.matrix.inner, - ByCols as u8, - ))? - }; - Ok(count as u64) - } -} - -impl ResultCount for RpqMatrixResult { - fn result_count(&self) -> Result { - Ok(self.reachable_target_count()? as usize) - } -} - -pub struct PreparedRpqMatrix { - plans: Vec, - owned_matrices: Vec, -} - -impl PreparedEvaluator for PreparedRpqMatrix { - type Result = RpqMatrixResult; - type Error = RpqError; - - fn execute(&mut self) -> Result { - let root_ptr = unsafe { self.plans.as_mut_ptr().add(self.plans.len() - 1) }; - - let mut nnz: GrB_Index = 0; - unsafe { la_ok!(LAGraph_RPQMatrix(&mut nnz, root_ptr))? }; - - let mut matrix_inner: GrB_Matrix = null_mut(); - unsafe { grb_ok!(GrB_Matrix_dup(&mut matrix_inner, (*root_ptr).res_mat))? }; - let matrix = GraphblasMatrix { - inner: matrix_inner, - }; - - unsafe { grb_ok!(LAGraph_DestroyRpqMatrixPlan(root_ptr))? }; - - Ok(RpqMatrixResult { - nnz: nnz as u64, - matrix, - }) - } -} - -impl Drop for PreparedRpqMatrix { - fn drop(&mut self) { - for mat in &mut self.owned_matrices { - unsafe { - LAGraph_RPQMatrix_Free(mat); - } - } - } -} - -/// RPQ evaluator backed by `LAGraph_RPQMatrix`. -#[derive(Clone, Copy)] -pub struct RpqMatrixEvaluator; - -impl Evaluator for RpqMatrixEvaluator { - type Query = RpqQuery; - type Result = RpqMatrixResult; - type Error = RpqError; - type Prepared = PreparedRpqMatrix; - - fn prepare( - &self, - query: &RpqQuery, - graph: &G, - ) -> Result { - let expr = query_to_expr(query, graph)?; - let (plans, owned_matrices) = materialize(&expr, graph)?; - - Ok(PreparedRpqMatrix { - plans, - owned_matrices, - }) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::rpq::{Endpoint, PathExpr, RpqQuery}; - use crate::utils::build_graph; - - #[test] - fn evaluate_single_edge_nnz() { - let graph = build_graph(&[("A", "B", "p")]); - let q = RpqQuery { - subject: Endpoint::Variable("x".into()), - path: PathExpr::Label("p".into()), - object: Endpoint::Variable("y".into()), - }; - let result = RpqMatrixEvaluator.evaluate(&q, &graph).expect("evaluate"); - assert_eq!(result.nnz, 1); - } - - #[test] - fn evaluate_named_subject_no_match_nnz() { - // Graph: A --p--> B - // Query: p ?y -> C has no outgoing p edges, nnz=0 - let graph = build_graph(&[("A", "B", "p"), ("C", "D", "q")]); - let q = RpqQuery { - subject: Endpoint::Named("C".into()), - path: PathExpr::Label("p".into()), - object: Endpoint::Variable("y".into()), - }; - let result = RpqMatrixEvaluator.evaluate(&q, &graph).expect("evaluate"); - assert_eq!(result.nnz, 0, "C has no outgoing p edges"); - } -} diff --git a/pathrex/src/rpq/rpqmatrix/cost.rs b/pathrex/src/rpq/rpqmatrix/cost.rs new file mode 100644 index 0000000..b3bdd67 --- /dev/null +++ b/pathrex/src/rpq/rpqmatrix/cost.rs @@ -0,0 +1,178 @@ +use std::cmp::Ordering; + +use egg::{CostFunction, Id}; + +use super::plan::RpqPlan; + +#[derive(Clone, Debug, PartialEq)] +pub struct CardCost { + pub score: f64, + pub nnz: f64, + pub nnz_r: f64, + pub nnz_c: f64, +} + +impl Eq for CardCost {} + +impl PartialOrd for CardCost { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for CardCost { + fn cmp(&self, other: &Self) -> Ordering { + match self.score.total_cmp(&other.score) { + Ordering::Equal => {} + ord => return ord, + } + match self.nnz.total_cmp(&other.nnz) { + Ordering::Equal => {} + ord => return ord, + } + match self.nnz_r.total_cmp(&other.nnz_r) { + Ordering::Equal => {} + ord => return ord, + } + self.nnz_c.total_cmp(&other.nnz_c) + } +} + +pub struct CardinalityCostFn { + pub n: f64, + pub star_penalty: f64, + pub lr_multiplier: f64, +} + +// TODO: check value intervals +impl CostFunction for CardinalityCostFn { + type Cost = CardCost; + + fn cost(&mut self, enode: &RpqPlan, mut costs: C) -> Self::Cost + where + C: FnMut(Id) -> Self::Cost, + { + match enode { + RpqPlan::NamedVertex(_name) => CardCost { + score: 0.0, + nnz: 1 as f64, + nnz_r: 1 as f64, + nnz_c: 1 as f64, + }, + RpqPlan::Label(meta) => CardCost { + score: 0.0, + nnz: meta.nvals as f64, + nnz_r: meta.nonzero_rows as f64, + nnz_c: meta.nonzero_cols as f64, + }, + + RpqPlan::Seq([a, b]) => { + let ca = costs(*a); + let cb = costs(*b); + + let denom = ca.nnz_r.max(cb.nnz_c).max(1.0); + let op_cost = (ca.nnz * cb.nnz) / denom; + let score = ca.score + cb.score + op_cost; + + let nnz_est = ca.nnz * cb.nnz / (self.n * self.n); + + CardCost { + score, + nnz: nnz_est, + nnz_r: ca.nnz_r.min(self.n), // TODO: better reduce estimators + nnz_c: cb.nnz_c.min(self.n), // TODO: better reduce estimators + } + } + + RpqPlan::Alt([a, b]) => { + let ca = costs(*a); + let cb = costs(*b); + + let overlap = (ca.nnz * cb.nnz) / (self.n * self.n); + let op_cost = ca.nnz + cb.nnz - overlap; + let score = ca.score + cb.score + op_cost; + + let nnz_est = (ca.nnz + cb.nnz - overlap).min(self.n * self.n).max(0.0); + + let nnz_r_est = (ca.nnz_r + cb.nnz_r - (ca.nnz_r * cb.nnz_r) / self.n) + .min(self.n) + .max(0.0); + + let nnz_c_est = (ca.nnz_c + cb.nnz_c - (ca.nnz_c * cb.nnz_c) / self.n) + .min(self.n) + .max(0.0); + + CardCost { + score, + nnz: nnz_est, + nnz_r: nnz_r_est, + nnz_c: nnz_c_est, + } + } + + RpqPlan::Star([a]) => { + let ca = costs(*a); + + let penalty = self.star_penalty * ca.nnz.max(1.0); + let score = ca.score + penalty; + + CardCost { + score, + nnz: self.n * self.n, + nnz_r: self.n, + nnz_c: self.n, + } + } + + RpqPlan::LStar([a, b]) => { + let ca = costs(*a); + let cb = costs(*b); + + let denom = ca.nnz_r.max(cb.nnz_c).max(1.0); + let base = (ca.nnz * cb.nnz) / denom; + let op_cost = self.lr_multiplier * base; + let score = ca.score + cb.score + op_cost; + + let nnz_est = self.lr_multiplier * ca.nnz * cb.nnz / (self.n * self.n); + + CardCost { + score, + nnz: nnz_est, + nnz_r: ca.nnz_r.min(self.n), // TODO: better reduce estimators + nnz_c: cb.nnz_c.min(self.n), // TODO: better reduce estimators + } + } + + RpqPlan::RStar([a, b]) => { + let ca = costs(*a); + let cb = costs(*b); + + let denom = ca.nnz_r.max(cb.nnz_c).max(1.0); + let base = (ca.nnz * cb.nnz) / denom; + + let op_cost = self.lr_multiplier * base; + let score = ca.score + cb.score + op_cost; + + let nnz_est = self.lr_multiplier * ca.nnz * cb.nnz / (self.n * self.n); + + CardCost { + score, + nnz: nnz_est, + nnz_r: ca.nnz_r.min(self.n), // TODO: better reduce estimators + nnz_c: cb.nnz_c.min(self.n), // TODO: better reduce estimators + } + } + } + } +} + +pub struct RandomCostFn; +impl CostFunction for RandomCostFn { + type Cost = f64; + fn cost(&mut self, _enode: &RpqPlan, _costs: C) -> Self::Cost + where + C: FnMut(Id) -> Self::Cost, + { + rand::random() + } +} diff --git a/pathrex/src/rpq/rpqmatrix/eval.rs b/pathrex/src/rpq/rpqmatrix/eval.rs new file mode 100644 index 0000000..a52661e --- /dev/null +++ b/pathrex/src/rpq/rpqmatrix/eval.rs @@ -0,0 +1,97 @@ + +/// RPQ evaluator backed by `LAGraph_RPQMatrix`. +use crate::eval::Evaluator; +use crate::graph::GraphDecomposition; +use super::optimize::{OptimizationStrategy,optimize_expr_cardinality}; +use crate::rpq::{RpqError, RpqQuery}; +use super::result::{RpqMatrixResult,PreparedRpqMatrix}; +use super::expr::{query_to_expr,materialize}; + + + +#[derive(Clone, Copy)] +pub struct RpqMatrixEvaluator { + optimizer: OptimizationStrategy, +} + +impl RpqMatrixEvaluator { + pub fn unoptimized() -> Self { + return RpqMatrixEvaluator { optimizer: OptimizationStrategy::NoOpt }; + } + pub fn optimized(opt: OptimizationStrategy) -> Self { + return RpqMatrixEvaluator { optimizer: opt }; + } +} + +impl Default for RpqMatrixEvaluator { + fn default() -> Self { + RpqMatrixEvaluator::unoptimized() + } +} + +impl Evaluator for RpqMatrixEvaluator { + type Query = RpqQuery; + type Result = RpqMatrixResult; + type Error = RpqError; + type Prepared = PreparedRpqMatrix; + + fn prepare( + &self, + query: &RpqQuery, + graph: &G, + ) -> Result { + let mut expr = query_to_expr(query, graph)?; + match self.optimizer { + OptimizationStrategy::NoOpt => {} + OptimizationStrategy::Cardinality => { + expr = optimize_expr_cardinality(expr, graph.num_nodes()); + } + _ => todo!(), + } + + let (plans, owned_matrices) = materialize(&expr, graph)?; + + Ok(PreparedRpqMatrix { + plans, + owned_matrices, + }) + } +} + + +#[cfg(test)] +mod tests { + use super::*; + use crate::rpq::{Endpoint, PathExpr, RpqQuery}; + use crate::utils::build_graph; + + #[test] + fn evaluate_single_edge_nnz() { + let graph = build_graph(&[("A", "B", "p")]); + let q = RpqQuery { + subject: Endpoint::Variable("x".into()), + path: PathExpr::Label("p".into()), + object: Endpoint::Variable("y".into()), + }; + let result = RpqMatrixEvaluator::default() + .evaluate(&q, &graph) + .expect("evaluate"); + assert_eq!(result.nnz, 1); + } + + #[test] + fn evaluate_named_subject_no_match_nnz() { + // Graph: A --p--> B + // Query: p ?y -> C has no outgoing p edges, nnz=0 + let graph = build_graph(&[("A", "B", "p"), ("C", "D", "q")]); + let q = RpqQuery { + subject: Endpoint::Named("C".into()), + path: PathExpr::Label("p".into()), + object: Endpoint::Variable("y".into()), + }; + let result = RpqMatrixEvaluator::default() + .evaluate(&q, &graph) + .expect("evaluate"); + assert_eq!(result.nnz, 0, "C has no outgoing p edges"); + } +} diff --git a/pathrex/src/rpq/rpqmatrix/expr.rs b/pathrex/src/rpq/rpqmatrix/expr.rs new file mode 100644 index 0000000..558d1a3 --- /dev/null +++ b/pathrex/src/rpq/rpqmatrix/expr.rs @@ -0,0 +1,205 @@ + + +use std::ptr::null_mut; + +use egg::{Id, RecExpr}; + +use crate::graph::GraphDecomposition; +use crate::lagraph_sys::*; +use crate::rpq::{Endpoint, PathExpr, RpqError, RpqQuery}; +use super::plan::{RpqPlan,LabelMeta}; +use crate::grb_ok; + + + +fn label_meta(label: &str, graph: &G) -> Result { + if let Some(metadata) = graph.get_metadata().and_then(|m| m.matrix(label)) { + return Ok(LabelMeta { + name: label.to_owned(), + nvals: metadata.nvals, + nonzero_rows: metadata.nonzero_rows, + nonzero_cols: metadata.nonzero_cols, + }); + } + + // TODO: maybe create optimized (for mm format) and nonoptimized (for other formats) plans + let lg = graph.get_graph(label)?; + let nvals = lg.nvals()? as usize; + Ok(LabelMeta { + name: label.to_owned(), + nvals, + nonzero_rows: nvals, + nonzero_cols: nvals, + }) +} + +fn to_expr_aux( + path: &PathExpr, + expr: &mut RecExpr, + graph: &G, +) -> Result { + match path { + PathExpr::Label(label) => Ok(expr.add(RpqPlan::Label(label_meta(label, graph)?))), + + PathExpr::Sequence(lhs, rhs) => { + let l = to_expr_aux(lhs, expr, graph)?; + let r = to_expr_aux(rhs, expr, graph)?; + Ok(expr.add(RpqPlan::Seq([l, r]))) + } + + PathExpr::Alternative(lhs, rhs) => { + let l = to_expr_aux(lhs, expr, graph)?; + let r = to_expr_aux(rhs, expr, graph)?; + Ok(expr.add(RpqPlan::Alt([l, r]))) + } + + PathExpr::ZeroOrMore(inner) => { + let i = to_expr_aux(inner, expr, graph)?; + Ok(expr.add(RpqPlan::Star([i]))) + } + + PathExpr::OneOrMore(inner) => { + let e = to_expr_aux(inner, expr, graph)?; + let s = expr.add(RpqPlan::Star([e])); + Ok(expr.add(RpqPlan::Seq([e, s]))) + } + + PathExpr::ZeroOrOne(_) => Err(RpqError::UnsupportedPath( + "ZeroOrOne (?) is not supported by RPQMatrix".into(), + )), + } +} + +/// Compile a [`RpqQuery`] into +/// [`RecExpr`]. +pub fn query_to_expr( + query: &RpqQuery, + graph: &G, +) -> Result, RpqError> { + let mut expr = RecExpr::default(); + let path_root = to_expr_aux(&query.path, &mut expr, graph)?; + + let _root = match (&query.subject, &query.object) { + (Endpoint::Variable(_), Endpoint::Variable(_)) => path_root, + (Endpoint::Named(name), Endpoint::Variable(_)) => { + let diag = expr.add(RpqPlan::NamedVertex(name.clone())); + expr.add(RpqPlan::Seq([diag, path_root])) + } + (Endpoint::Variable(_), Endpoint::Named(name)) => { + let diag = expr.add(RpqPlan::NamedVertex(name.clone())); + expr.add(RpqPlan::Seq([path_root, diag])) + } + (Endpoint::Named(sub), Endpoint::Named(obj)) => { + let diag_sub = expr.add(RpqPlan::NamedVertex(sub.clone())); + let seq1 = expr.add(RpqPlan::Seq([diag_sub, path_root])); + let diag_obj = expr.add(RpqPlan::NamedVertex(obj.clone())); + expr.add(RpqPlan::Seq([seq1, diag_obj])) + } + }; + + Ok(expr) +} + +/// Convert a [`RecExpr`] into the flat [`RPQMatrixPlan`] array that +/// `LAGraph_RPQMatrix` expects. +/// +/// Returns the plan array and a list of owned diagonal matrices that must be +/// freed after evaluation. +pub fn materialize( + expr: &RecExpr, + graph: &G, +) -> Result<(Vec, Vec), RpqError> { + let null_plan = RPQMatrixPlan { + op: RPQMatrixOp::RPQ_MATRIX_OP_LABEL, + lhs: null_mut(), + rhs: null_mut(), + mat: null_mut(), + res_mat: null_mut(), + }; + let mut plans = vec![null_plan; expr.len()]; + let mut owned_matrices: Vec = Vec::new(); + let n = graph.num_nodes() as GrB_Index; + + for (id, node) in expr.as_ref().iter().enumerate() { + plans[id] = match node { + RpqPlan::Label(label) => { + let lg = graph.get_graph(&label.name)?; + let mat = unsafe { (*lg.inner).A }; + RPQMatrixPlan { + op: RPQMatrixOp::RPQ_MATRIX_OP_LABEL, + lhs: null_mut(), + rhs: null_mut(), + mat, + res_mat: null_mut(), + } + } + + RpqPlan::NamedVertex(name) => { + let vertex_id = graph + .get_node_id(name) + .ok_or_else(|| RpqError::VertexNotFound(name.clone()))? + as GrB_Index; + let mut mat: GrB_Matrix = null_mut(); + unsafe { + crate::graph::ensure_grb_init()?; + grb_ok!(LAGraph_RPQMatrix_label(&mut mat, vertex_id, n, n,))? + }; + if mat.is_null() { + return Err(RpqError::Graph(crate::graph::GraphError::GraphBlas( + GrB_Info::GrB_INVALID_VALUE, + ))); + } + owned_matrices.push(mat); + RPQMatrixPlan { + op: RPQMatrixOp::RPQ_MATRIX_OP_LABEL, + lhs: null_mut(), + rhs: null_mut(), + mat, + res_mat: null_mut(), + } + } + + RpqPlan::Seq([l, r]) => RPQMatrixPlan { + op: RPQMatrixOp::RPQ_MATRIX_OP_CONCAT, + lhs: unsafe { plans.as_mut_ptr().add(usize::from(*l)) }, + rhs: unsafe { plans.as_mut_ptr().add(usize::from(*r)) }, + mat: null_mut(), + res_mat: null_mut(), + }, + + RpqPlan::Alt([l, r]) => RPQMatrixPlan { + op: RPQMatrixOp::RPQ_MATRIX_OP_LOR, + lhs: unsafe { plans.as_mut_ptr().add(usize::from(*l)) }, + rhs: unsafe { plans.as_mut_ptr().add(usize::from(*r)) }, + mat: null_mut(), + res_mat: null_mut(), + }, + + RpqPlan::Star([i]) => RPQMatrixPlan { + op: RPQMatrixOp::RPQ_MATRIX_OP_KLEENE, + lhs: null_mut(), + rhs: unsafe { plans.as_mut_ptr().add(usize::from(*i)) }, + mat: null_mut(), + res_mat: null_mut(), + }, + + RpqPlan::LStar([l, r]) => RPQMatrixPlan { + op: RPQMatrixOp::RPQ_MATRIX_OP_KLEENE_L, + lhs: unsafe { plans.as_mut_ptr().add(usize::from(*l)) }, + rhs: unsafe { plans.as_mut_ptr().add(usize::from(*r)) }, + mat: null_mut(), + res_mat: null_mut(), + }, + + RpqPlan::RStar([l, r]) => RPQMatrixPlan { + op: RPQMatrixOp::RPQ_MATRIX_OP_KLEENE_R, + lhs: unsafe { plans.as_mut_ptr().add(usize::from(*l)) }, + rhs: unsafe { plans.as_mut_ptr().add(usize::from(*r)) }, + mat: null_mut(), + res_mat: null_mut(), + }, + }; + } + + Ok((plans, owned_matrices)) +} \ No newline at end of file diff --git a/pathrex/src/rpq/rpqmatrix/mod.rs b/pathrex/src/rpq/rpqmatrix/mod.rs new file mode 100644 index 0000000..09b62fa --- /dev/null +++ b/pathrex/src/rpq/rpqmatrix/mod.rs @@ -0,0 +1,12 @@ +//! Plan-based RPQ evaluation using `LAGraph_RPQMatrix`. + +mod plan; +mod cost; +pub mod eval; +mod expr; +mod optimize; +pub mod result; + +pub use eval::RpqMatrixEvaluator; +pub use optimize::OptimizationStrategy; +pub use result::RpqMatrixResult; diff --git a/pathrex/src/rpq/rpqmatrix/optimize.rs b/pathrex/src/rpq/rpqmatrix/optimize.rs new file mode 100644 index 0000000..7070d8d --- /dev/null +++ b/pathrex/src/rpq/rpqmatrix/optimize.rs @@ -0,0 +1,35 @@ +use egg::{Extractor, RecExpr, Runner}; + +use super::cost::CardinalityCostFn; +use super::plan::{RpqPlan, make_rules}; + +#[derive(Clone, Copy)] +pub enum OptimizationStrategy { + NoOpt, + Cardinality, + RandomOpt, // TODO: should be same as random from la-n-egg-rpq: https://github.com/SparseLinearAlgebra/la-n-egg-rpq/blob/main/src/main.rs#L75 + Simple, // TODO + Wander, // TODO +} + +pub(super) fn optimize_expr_cardinality( + expr: RecExpr, + graph_size: usize, +) -> RecExpr { + let rules = make_rules(); + let runner = Runner::default() + .with_explanations_disabled() + .with_expr(&expr) + .run(&rules); + + let extractor = Extractor::new( + &runner.egraph, + CardinalityCostFn { + n: graph_size as f64, + star_penalty: 50.0, + lr_multiplier: 5.0, + }, + ); + let (_, plan) = extractor.find_best(runner.roots[0]); + plan +} diff --git a/pathrex/src/rpq/rpqmatrix/plan.rs b/pathrex/src/rpq/rpqmatrix/plan.rs new file mode 100644 index 0000000..8b12b5c --- /dev/null +++ b/pathrex/src/rpq/rpqmatrix/plan.rs @@ -0,0 +1,57 @@ +use std::{fmt::Display, str::FromStr}; + +use egg::{Id, define_language, rewrite}; + +#[derive(Clone, Hash, Ord, Eq, PartialEq, PartialOrd, Debug)] +pub(super) struct LabelMeta { + pub name: String, + pub nvals: usize, + pub nonzero_rows: usize, + pub nonzero_cols: usize, +} + +impl FromStr for LabelMeta { + type Err = ::Err; + // This is needed for the builtin egg parser. Only used in tests. + fn from_str(s: &str) -> Result { + Ok(LabelMeta { + name: "-".to_string(), + nvals: s.parse()?, + nonzero_rows: s.parse()?, + nonzero_cols: s.parse()?, + }) + } +} + +impl Display for LabelMeta { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "({}, {})", self.name, self.nvals) + } +} + +define_language! { + pub enum RpqPlan { + Label(LabelMeta), + NamedVertex(String), + "/" = Seq([Id; 2]), + "|" = Alt([Id; 2]), + "*" = Star([Id; 1]), + "l*" = LStar([Id; 2]), + "*r" = RStar([Id; 2]), + } +} + +pub(super) fn make_rules() -> Vec> { + vec![ + rewrite!("assoc-sec-1"; "(/ ?a (/ ?b ?c))" => "(/ (/ ?a ?b) ?c)"), + rewrite!("assoc-sec-2"; "(/ (/ ?a ?b) ?c)" => "(/ ?a (/ ?b ?c))"), + rewrite!("commute-alt"; "(| ?a ?b)" => "(| ?b ?a)"), + rewrite!("assoc-alt"; "(| ?a (| ?b ?c))" => "(| (| ?a ?b) ?c)"), + rewrite!("distribute-1"; "(/ ?a (| ?b ?c))" => "(| (/ ?a ?b) (/ ?a ?c))"), + rewrite!("distribute-2"; "(/ (| ?a ?b) ?c)" => "(| (/ ?a ?c) (/ ?b ?c))"), + rewrite!("distribute-3"; "(| (/ ?a ?b) (/ ?a ?c))" => "(/ ?a (| ?b ?c))"), + rewrite!("distribute-4"; "(| (/ ?a ?c) (/ ?b ?c))" => "(/ (| ?a ?b) ?c)"), + rewrite!("build-lstar"; "(/ (* ?a) ?b)" => "(l* ?a ?b)"), + rewrite!("build-rstar"; "(/ ?a (* ?b))" => "(*r ?a ?b)"), + ] +} diff --git a/pathrex/src/rpq/rpqmatrix/result.rs b/pathrex/src/rpq/rpqmatrix/result.rs new file mode 100644 index 0000000..428ae3d --- /dev/null +++ b/pathrex/src/rpq/rpqmatrix/result.rs @@ -0,0 +1,78 @@ +use std::ptr::null_mut; + +use crate::eval::{PreparedEvaluator, ResultCount}; +use crate::graph::wrappers::ReduceType::ByCols; +use crate::graph::{GraphError, GraphblasMatrix}; +use crate::lagraph_sys::*; + +use crate::rpq::RpqError; +use crate::{grb_ok, la_ok}; + +/// Output of [`RpqMatrixEvaluator`]: full path relation matrix and its nnz. +#[derive(Debug)] +pub struct RpqMatrixResult { + pub nnz: u64, + pub matrix: GraphblasMatrix, +} + +impl RpqMatrixResult { + /// Count distinct reachable target vertices by reducing the path relation + /// matrix to its non-empty columns. + pub fn reachable_target_count(&self) -> Result { + let mut count: GrB_Index = 0; + unsafe { + grb_ok!(LAGraph_RPQMatrix_reduce( + &mut count, + self.matrix.inner, + ByCols as u8, + ))? + }; + Ok(count as u64) + } +} + +impl ResultCount for RpqMatrixResult { + fn result_count(&self) -> Result { + Ok(self.reachable_target_count()? as usize) + } +} + +pub struct PreparedRpqMatrix { + pub(super) plans: Vec, + pub(super) owned_matrices: Vec, +} + +impl PreparedEvaluator for PreparedRpqMatrix { + type Result = RpqMatrixResult; + type Error = RpqError; + + fn execute(&mut self) -> Result { + let root_ptr = unsafe { self.plans.as_mut_ptr().add(self.plans.len() - 1) }; + + let mut nnz: GrB_Index = 0; + unsafe { la_ok!(LAGraph_RPQMatrix(&mut nnz, root_ptr))? }; + + let mut matrix_inner: GrB_Matrix = null_mut(); + unsafe { grb_ok!(GrB_Matrix_dup(&mut matrix_inner, (*root_ptr).res_mat))? }; + let matrix = GraphblasMatrix { + inner: matrix_inner, + }; + + unsafe { grb_ok!(LAGraph_DestroyRpqMatrixPlan(root_ptr))? }; + + Ok(RpqMatrixResult { + nnz: nnz as u64, + matrix, + }) + } +} + +impl Drop for PreparedRpqMatrix { + fn drop(&mut self) { + for mat in &mut self.owned_matrices { + unsafe { + LAGraph_RPQMatrix_Free(mat); + } + } + } +} diff --git a/pathrex/tests/rpqmatrix_tests.rs b/pathrex/tests/rpqmatrix_tests.rs index 3f84d80..df888d6 100644 --- a/pathrex/tests/rpqmatrix_tests.rs +++ b/pathrex/tests/rpqmatrix_tests.rs @@ -6,7 +6,8 @@ use std::sync::LazyLock; use pathrex::formats::mm::MatrixMarket; use pathrex::graph::{Graph, GraphDecomposition, GraphError, InMemory, InMemoryGraph}; use pathrex::lagraph_sys::{GrB_Index, GrB_Info, GrB_Matrix_extractElement_BOOL}; -use pathrex::rpq::rpqmatrix::{RpqMatrixEvaluator, RpqMatrixResult}; +use pathrex::rpq::rpqmatrix::eval::{RpqMatrixEvaluator}; +use pathrex::rpq::rpqmatrix::result::{RpqMatrixResult}; use pathrex::rpq::{Endpoint, PathExpr, PreparedRpq, RpqError, RpqEvaluator, RpqQuery}; use pathrex::sparql::parse_rpq; use pathrex::utils::build_graph; @@ -79,7 +80,7 @@ fn run_la_n_egg_case(case_name: &str) { ); let graph = &*LA_N_EGG_GRAPH; - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); for (i, (query, expected_nnz)) in queries.iter().zip(expected.iter()).enumerate() { let result = evaluator.evaluate(query, graph).unwrap_or_else(|e| { @@ -128,7 +129,7 @@ fn matrix_entry_set(result: &RpqMatrixResult, row: GrB_Index, col: GrB_Index) -> #[test] fn test_single_label_variable_variable() { let graph = build_graph(&[("A", "B", "knows"), ("B", "C", "knows")]); - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); let result = evaluator .evaluate(&rq(var("x"), label("knows"), var("y")), &graph) @@ -142,7 +143,7 @@ fn test_single_label_variable_variable() { #[test] fn test_single_label_named_source() { let graph = build_graph(&[("A", "B", "knows"), ("B", "C", "knows")]); - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); let result = evaluator .evaluate(&rq(named_ep("A"), label("knows"), var("y")), &graph) @@ -162,7 +163,7 @@ fn test_single_label_named_source() { #[test] fn test_sequence_path() { let graph = build_graph(&[("A", "B", "knows"), ("B", "C", "likes")]); - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); let path = PathExpr::Sequence(Box::new(label("knows")), Box::new(label("likes"))); @@ -182,8 +183,8 @@ fn prepared_rpqmatrix_execution_matches_evaluate() { var("y"), ); - let direct = RpqMatrixEvaluator.evaluate(&query, &graph).expect("direct"); - let mut prepared = RpqMatrixEvaluator.prepare(&query, &graph).expect("prepare"); + let direct = RpqMatrixEvaluator::default().evaluate(&query, &graph).expect("direct"); + let mut prepared = RpqMatrixEvaluator::default().prepare(&query, &graph).expect("prepare"); let prepared_result = prepared.execute().expect("execute"); assert_eq!(prepared_result.nnz, direct.nnz); @@ -198,7 +199,7 @@ fn prepared_rpqmatrix_execution_can_run_twice() { var("y"), ); - let mut prepared = RpqMatrixEvaluator.prepare(&query, &graph).expect("prepare"); + let mut prepared = RpqMatrixEvaluator::default().prepare(&query, &graph).expect("prepare"); let first = prepared.execute().expect("first"); let second = prepared.execute().expect("second"); @@ -210,7 +211,7 @@ fn prepared_rpqmatrix_execution_can_run_twice() { #[test] fn test_sequence_path_named_source() { let graph = build_graph(&[("A", "B", "knows"), ("B", "C", "likes")]); - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); let path = PathExpr::Sequence(Box::new(label("knows")), Box::new(label("likes"))); @@ -232,7 +233,7 @@ fn test_sequence_path_named_source() { #[test] fn test_alternative_path() { let graph = build_graph(&[("A", "B", "knows"), ("A", "C", "likes")]); - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); let path = PathExpr::Alternative(Box::new(label("knows")), Box::new(label("likes"))); @@ -259,7 +260,7 @@ fn test_alternative_path() { #[test] fn test_zero_or_more_path() { let graph = build_graph(&[("A", "B", "knows"), ("B", "C", "knows")]); - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); let path = PathExpr::ZeroOrMore(Box::new(label("knows"))); @@ -291,7 +292,7 @@ fn test_zero_or_more_path() { #[test] fn test_one_or_more_path() { let graph = build_graph(&[("A", "B", "knows"), ("B", "C", "knows")]); - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); let path = PathExpr::OneOrMore(Box::new(label("knows"))); @@ -321,7 +322,7 @@ fn test_one_or_more_path() { #[test] fn test_zero_or_one_unsupported() { let graph = build_graph(&[("A", "B", "knows")]); - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); let path = PathExpr::ZeroOrOne(Box::new(label("knows"))); let result = evaluator.evaluate(&rq(var("x"), path, var("y")), &graph); @@ -335,7 +336,7 @@ fn test_zero_or_one_unsupported() { #[test] fn test_label_not_found() { let graph = build_graph(&[("A", "B", "knows")]); - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); let result = evaluator.evaluate(&rq(var("x"), label("nonexistent"), var("y")), &graph); @@ -348,7 +349,7 @@ fn test_label_not_found() { #[test] fn test_vertex_not_found() { let graph = build_graph(&[("A", "B", "knows")]); - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); let result = evaluator.evaluate(&rq(named_ep("Z"), label("knows"), var("y")), &graph); @@ -363,7 +364,7 @@ fn test_vertex_not_found() { #[test] fn test_bound_object() { let graph = build_graph(&[("A", "B", "knows"), ("C", "D", "knows")]); - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); let result = evaluator .evaluate(&rq(var("x"), label("knows"), named_ep("B")), &graph) @@ -377,7 +378,7 @@ fn test_bound_object() { #[test] fn test_bound_subject_and_object() { let graph = build_graph(&[("A", "B", "knows"), ("C", "D", "knows")]); - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); let result = evaluator .evaluate(&rq(named_ep("A"), label("knows"), named_ep("B")), &graph) @@ -402,7 +403,7 @@ fn test_cycle_graph_star() { ("B", "C", "knows"), ("C", "A", "knows"), ]); - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); let path = PathExpr::ZeroOrMore(Box::new(label("knows"))); @@ -441,7 +442,7 @@ fn test_complex_path() { ("B", "C", "likes"), ("C", "D", "knows"), ]); - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); // knows / likes* / knows let path = PathExpr::Sequence( @@ -468,7 +469,7 @@ fn test_complex_path() { #[test] fn test_no_matching_path() { let graph = build_graph(&[("A", "B", "knows")]); - let evaluator = RpqMatrixEvaluator; + let evaluator = RpqMatrixEvaluator::default(); let path = PathExpr::Sequence(Box::new(label("knows")), Box::new(label("likes"))); From beb681c9f4ef77f6bede04f3373e3cc145f86552 Mon Sep 17 00:00:00 2001 From: Rodion Suvorov Date: Wed, 15 Jul 2026 19:52:41 +0300 Subject: [PATCH 4/8] feat: support optimizer for rpqmatrix in cli --- pathrex/src/bin/pathrex.rs | 44 +++++++++++++++++++++++++++-- pathrex/src/cli/args.rs | 47 +++++++++++++++++++++++++++++++ pathrex/src/cli/checkpoint.rs | 29 ++++++++++++++++--- pathrex/src/cli/dispatch.rs | 12 +++++--- pathrex/src/cli/loader.rs | 2 +- pathrex/src/cli/output.rs | 2 ++ pathrex/src/rpq/rpqmatrix/eval.rs | 14 ++++----- pathrex/src/rpq/rpqmatrix/expr.rs | 10 ++----- pathrex/src/rpq/rpqmatrix/mod.rs | 2 +- pathrex/tests/rpqmatrix_tests.rs | 16 +++++++---- 10 files changed, 145 insertions(+), 33 deletions(-) diff --git a/pathrex/src/bin/pathrex.rs b/pathrex/src/bin/pathrex.rs index 69b508a..6b4339a 100644 --- a/pathrex/src/bin/pathrex.rs +++ b/pathrex/src/bin/pathrex.rs @@ -16,7 +16,8 @@ //! cargo run --release --bin pathrex --features bench -- bench \ //! --graph tests/testdata/mm_graph \ //! --queries tests/testdata/cases/any-any/queries.txt \ -//! --algo nfa rpqmatrix \ +//! --algo nfarpq rpqmatrix \ +//! --rpqmatrix-optimizer cardinality \ //! --output results.json //! ``` @@ -27,7 +28,9 @@ use chrono::Utc; use clap::Parser; use thiserror::Error; -use pathrex::cli::args::{BenchArgs, Cli, Commands, QueryArgs}; +use pathrex::cli::args::{ + Algo, BenchArgs, Cli, Commands, CommonArgs, GraphFormat, QueryArgs, RpqMatrixOptimizer, +}; use pathrex::cli::bench::BenchError; use pathrex::cli::checkpoint::{Checkpoint, CheckpointError, Checkpointer}; use pathrex::cli::dispatch::{dispatch_bench, dispatch_query}; @@ -55,6 +58,8 @@ enum MainError { #[source] source: std::io::Error, }, + #[error("invalid arguments: {0}")] + InvalidArgs(String), } fn main() { @@ -78,6 +83,26 @@ fn run() -> Result<(), MainError> { } } +fn validate_common_args(common: &CommonArgs) -> Result<(), MainError> { + if common.rpqmatrix_optimizer == RpqMatrixOptimizer::None { + return Ok(()); + } + + if !common.algo.contains(&Algo::Rpqmatrix) { + return Err(MainError::InvalidArgs( + "--rpqmatrix-optimizer can only be used when --algo includes rpqmatrix".to_string(), + )); + } + + if common.format != GraphFormat::Mm { + return Err(MainError::InvalidArgs( + "--rpqmatrix-optimizer can only be used with --format mm".to_string(), + )); + } + + Ok(()) +} + fn load_query_file(path: &str, base_iri: Option<&str>) -> Result, MainError> { load_queries(Path::new(path), base_iri).map_err(|e| MainError::Queries { path: path.to_string(), @@ -87,12 +112,14 @@ fn load_query_file(path: &str, base_iri: Option<&str>) -> Result Result<(), MainError> { let common = &args.common; + validate_common_args(common)?; eprintln!("=== pathrex query ==="); eprintln!("Graph: {}", common.graph); eprintln!("Format: {}", common.format); eprintln!("Queries: {}", common.queries); eprintln!("Algos: {:?}", common.algo); + eprintln!("RPQMatrix optimizer: {}", common.rpqmatrix_optimizer); eprintln!(); eprintln!("[1/2] Loading graph..."); @@ -127,6 +154,7 @@ fn run_query_cmd(args: QueryArgs) -> Result<(), MainError> { graph_path: common.graph.clone(), graph_format: common.format.to_string(), queries_file: common.queries.clone(), + rpqmatrix_optimizer: Some(common.rpqmatrix_optimizer.to_string()), base_iri: common.base_iri.clone(), num_nodes: graph.num_nodes(), num_labels: graph.num_labels(), @@ -152,7 +180,12 @@ fn build_checkpointer(args: &BenchArgs, queries_len: usize) -> Result { - cp.validate(&common.graph, &common.queries, &common.algo)?; + cp.validate( + &common.graph, + &common.queries, + &common.algo, + common.rpqmatrix_optimizer, + )?; let cper = Checkpointer::with_inner(cp, path); eprintln!( " resuming: {}/{} queries fully done", @@ -167,6 +200,7 @@ fn build_checkpointer(args: &BenchArgs, queries_len: usize) -> Result Result Result Result<(), MainError> { let common = &args.common; + validate_common_args(common)?; eprintln!("=== pathrex bench ==="); eprintln!("Graph: {}", common.graph); eprintln!("Format: {}", common.format); eprintln!("Queries: {}", common.queries); eprintln!("Algos: {:?}", common.algo); + eprintln!("RPQMatrix optimizer: {}", common.rpqmatrix_optimizer); eprintln!("Output: {}", args.output); eprintln!(); @@ -223,6 +260,7 @@ fn run_bench_cmd(args: BenchArgs) -> Result<(), MainError> { graph_format: common.format.to_string(), queries_file: common.queries.clone(), base_iri: common.base_iri.clone(), + rpqmatrix_optimizer: Some(common.rpqmatrix_optimizer.to_string()), num_nodes: graph.num_nodes(), num_labels: graph.num_labels(), sample_size: args.sample_size, diff --git a/pathrex/src/cli/args.rs b/pathrex/src/cli/args.rs index d95f478..73e39ad 100644 --- a/pathrex/src/cli/args.rs +++ b/pathrex/src/cli/args.rs @@ -11,6 +11,8 @@ use clap::{Args, Parser, Subcommand, ValueEnum}; +use crate::rpq::rpqmatrix::OptimizationStrategy; + /// Top-level CLI for pathrex. #[derive(Parser, Debug)] #[command( @@ -62,6 +64,15 @@ pub struct CommonArgs { /// Algorithms to use. #[arg(short = 'a', long, value_enum, num_args = 1.., required = true)] pub algo: Vec, + + /// Optimizer type (only for the RPQMatrix algorithm and MatrixMarket source graph). + #[arg( + short = 'p', + long = "rpqmatrix-optimizer", + value_enum, + default_value_t = RpqMatrixOptimizer::None + )] + pub rpqmatrix_optimizer: RpqMatrixOptimizer, } /// Arguments for the `query` subcommand. @@ -156,6 +167,42 @@ impl std::fmt::Display for GraphFormat { } } +/// Optimizer types. +/// Only for RPQMatrix algorithm and MatrixMarket source graph. +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum, serde::Serialize, serde::Deserialize)] +#[value(rename_all = "lowercase")] +#[serde(rename_all = "lowercase")] +pub enum RpqMatrixOptimizer { + /// Optimizer based on cardinality of source matrices. + Cardinality, + /// Without any optimizations. + None, +} + +impl Default for RpqMatrixOptimizer { + fn default() -> Self { + Self::None + } +} + +impl std::fmt::Display for RpqMatrixOptimizer { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + RpqMatrixOptimizer::Cardinality => write!(f, "cardinality"), + RpqMatrixOptimizer::None => write!(f, "none"), + } + } +} + +impl From for OptimizationStrategy { + fn from(value: RpqMatrixOptimizer) -> Self { + match value { + RpqMatrixOptimizer::None => OptimizationStrategy::NoOpt, + RpqMatrixOptimizer::Cardinality => OptimizationStrategy::Cardinality, + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/pathrex/src/cli/checkpoint.rs b/pathrex/src/cli/checkpoint.rs index f59f3bb..2d4f13b 100644 --- a/pathrex/src/cli/checkpoint.rs +++ b/pathrex/src/cli/checkpoint.rs @@ -16,7 +16,7 @@ use std::path::{Path, PathBuf}; use serde::{Deserialize, Serialize}; use thiserror::Error; -use super::args::Algo; +use super::args::{Algo, RpqMatrixOptimizer}; /// Persistent checkpoint state written to disk as JSON. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -25,6 +25,8 @@ pub struct Checkpoint { pub graph_path: String, pub queries_file: String, pub algorithms: Vec, + #[serde(default)] + pub rpqmatrix_optimizer: RpqMatrixOptimizer, pub completed: Vec, } @@ -37,12 +39,18 @@ pub struct QueryCompletion { impl Checkpoint { /// Create a fresh checkpoint for a new benchmark run. - pub fn new(graph_path: &str, queries_file: &str, algorithms: &[Algo]) -> Self { + pub fn new( + graph_path: &str, + queries_file: &str, + algorithms: &[Algo], + rpqmatrix_optimizer: RpqMatrixOptimizer, + ) -> Self { Self { version: 1, graph_path: graph_path.to_string(), queries_file: queries_file.to_string(), algorithms: algorithms.to_vec(), + rpqmatrix_optimizer, completed: Vec::new(), } } @@ -65,6 +73,7 @@ impl Checkpoint { graph_path: &str, queries_file: &str, algorithms: &[Algo], + rpqmatrix_optimizer: RpqMatrixOptimizer, ) -> Result<(), CheckpointError> { if self.graph_path != graph_path { return Err(CheckpointError::Mismatch(format!( @@ -86,6 +95,12 @@ impl Checkpoint { self.algorithms, algorithms ))); } + if self.rpqmatrix_optimizer != rpqmatrix_optimizer { + return Err(CheckpointError::Mismatch(format!( + "rpqmatrix_optimizer: checkpoint has '{}', current is '{}'", + self.rpqmatrix_optimizer, rpqmatrix_optimizer + ))); + } Ok(()) } @@ -145,9 +160,15 @@ pub struct Checkpointer { impl Checkpointer { /// Create a new checkpointer with no completions. - pub fn fresh(graph_path: &str, queries_file: &str, algorithms: &[Algo], path: PathBuf) -> Self { + pub fn fresh( + graph_path: &str, + queries_file: &str, + algorithms: &[Algo], + rpqmatrix_optimizer: RpqMatrixOptimizer, + path: PathBuf, + ) -> Self { Self { - inner: Checkpoint::new(graph_path, queries_file, algorithms), + inner: Checkpoint::new(graph_path, queries_file, algorithms, rpqmatrix_optimizer), path, } } diff --git a/pathrex/src/cli/dispatch.rs b/pathrex/src/cli/dispatch.rs index cd19ee6..07f57a5 100644 --- a/pathrex/src/cli/dispatch.rs +++ b/pathrex/src/cli/dispatch.rs @@ -9,7 +9,7 @@ use crate::cli::query::run_query_for_evaluator; use crate::graph::InMemoryGraph; use crate::rpq::nfarpq::NfaRpqEvaluator; -use crate::rpq::rpqmatrix::eval::RpqMatrixEvaluator; +use crate::rpq::rpqmatrix::RpqMatrixEvaluator; fn merge_results(all: &mut Vec, per_algo: Vec) { for result in per_algo { @@ -30,12 +30,16 @@ pub fn dispatch_query( queries: &[LoadedQuery], ) -> Vec { let mut all = Vec::new(); - for algo in &args.common.algo { let name = algo.to_string(); let per_algo = match algo { Algo::NfaRpq => run_query_for_evaluator(&name, NfaRpqEvaluator, graph, queries), - Algo::Rpqmatrix => run_query_for_evaluator(&name, RpqMatrixEvaluator::default(), graph, queries), + Algo::Rpqmatrix => run_query_for_evaluator( + &name, + RpqMatrixEvaluator::optimized(args.common.rpqmatrix_optimizer.into()), + graph, + queries, + ), }; merge_results(&mut all, per_algo); } @@ -67,7 +71,7 @@ pub fn dispatch_bench( args, algo, &name, - RpqMatrixEvaluator::default(), + RpqMatrixEvaluator::optimized(args.common.rpqmatrix_optimizer.into()), graph, queries, checkpointer, diff --git a/pathrex/src/cli/loader.rs b/pathrex/src/cli/loader.rs index a652043..f99b4fa 100644 --- a/pathrex/src/cli/loader.rs +++ b/pathrex/src/cli/loader.rs @@ -67,7 +67,7 @@ pub fn load_graph( }) } GraphFormat::Rdf => { - let rdf = Rdf::from_path(graph_path).unwrap(); + let rdf = Rdf::from_path(graph_path).unwrap(); // TODO: handle panic Graph::::try_from(rdf).map_err(|e| GraphLoadError::Build { path: graph_path.to_string(), source: e, diff --git a/pathrex/src/cli/output.rs b/pathrex/src/cli/output.rs index 4dbe0a2..c7aa2d1 100644 --- a/pathrex/src/cli/output.rs +++ b/pathrex/src/cli/output.rs @@ -91,6 +91,7 @@ pub struct QueryMetadata { pub graph_path: String, pub graph_format: String, pub queries_file: String, + pub rpqmatrix_optimizer: Option, #[serde(skip_serializing_if = "Option::is_none")] pub base_iri: Option, pub num_nodes: usize, @@ -116,6 +117,7 @@ pub struct BenchMetadata { pub graph_path: String, pub graph_format: String, pub queries_file: String, + pub rpqmatrix_optimizer: Option, #[serde(skip_serializing_if = "Option::is_none")] pub base_iri: Option, pub num_nodes: usize, diff --git a/pathrex/src/rpq/rpqmatrix/eval.rs b/pathrex/src/rpq/rpqmatrix/eval.rs index a52661e..43bc8ea 100644 --- a/pathrex/src/rpq/rpqmatrix/eval.rs +++ b/pathrex/src/rpq/rpqmatrix/eval.rs @@ -1,13 +1,10 @@ - +use super::expr::{materialize, query_to_expr}; +use super::optimize::{OptimizationStrategy, optimize_expr_cardinality}; +use super::result::{PreparedRpqMatrix, RpqMatrixResult}; /// RPQ evaluator backed by `LAGraph_RPQMatrix`. use crate::eval::Evaluator; use crate::graph::GraphDecomposition; -use super::optimize::{OptimizationStrategy,optimize_expr_cardinality}; use crate::rpq::{RpqError, RpqQuery}; -use super::result::{RpqMatrixResult,PreparedRpqMatrix}; -use super::expr::{query_to_expr,materialize}; - - #[derive(Clone, Copy)] pub struct RpqMatrixEvaluator { @@ -16,7 +13,9 @@ pub struct RpqMatrixEvaluator { impl RpqMatrixEvaluator { pub fn unoptimized() -> Self { - return RpqMatrixEvaluator { optimizer: OptimizationStrategy::NoOpt }; + return RpqMatrixEvaluator { + optimizer: OptimizationStrategy::NoOpt, + }; } pub fn optimized(opt: OptimizationStrategy) -> Self { return RpqMatrixEvaluator { optimizer: opt }; @@ -58,7 +57,6 @@ impl Evaluator for RpqMatrixEvaluator { } } - #[cfg(test)] mod tests { use super::*; diff --git a/pathrex/src/rpq/rpqmatrix/expr.rs b/pathrex/src/rpq/rpqmatrix/expr.rs index 558d1a3..2c27c3a 100644 --- a/pathrex/src/rpq/rpqmatrix/expr.rs +++ b/pathrex/src/rpq/rpqmatrix/expr.rs @@ -1,16 +1,12 @@ - - use std::ptr::null_mut; use egg::{Id, RecExpr}; +use super::plan::{LabelMeta, RpqPlan}; use crate::graph::GraphDecomposition; +use crate::grb_ok; use crate::lagraph_sys::*; use crate::rpq::{Endpoint, PathExpr, RpqError, RpqQuery}; -use super::plan::{RpqPlan,LabelMeta}; -use crate::grb_ok; - - fn label_meta(label: &str, graph: &G) -> Result { if let Some(metadata) = graph.get_metadata().and_then(|m| m.matrix(label)) { @@ -202,4 +198,4 @@ pub fn materialize( } Ok((plans, owned_matrices)) -} \ No newline at end of file +} diff --git a/pathrex/src/rpq/rpqmatrix/mod.rs b/pathrex/src/rpq/rpqmatrix/mod.rs index 09b62fa..52dab8c 100644 --- a/pathrex/src/rpq/rpqmatrix/mod.rs +++ b/pathrex/src/rpq/rpqmatrix/mod.rs @@ -1,10 +1,10 @@ //! Plan-based RPQ evaluation using `LAGraph_RPQMatrix`. -mod plan; mod cost; pub mod eval; mod expr; mod optimize; +mod plan; pub mod result; pub use eval::RpqMatrixEvaluator; diff --git a/pathrex/tests/rpqmatrix_tests.rs b/pathrex/tests/rpqmatrix_tests.rs index df888d6..36ec0a5 100644 --- a/pathrex/tests/rpqmatrix_tests.rs +++ b/pathrex/tests/rpqmatrix_tests.rs @@ -6,8 +6,8 @@ use std::sync::LazyLock; use pathrex::formats::mm::MatrixMarket; use pathrex::graph::{Graph, GraphDecomposition, GraphError, InMemory, InMemoryGraph}; use pathrex::lagraph_sys::{GrB_Index, GrB_Info, GrB_Matrix_extractElement_BOOL}; -use pathrex::rpq::rpqmatrix::eval::{RpqMatrixEvaluator}; -use pathrex::rpq::rpqmatrix::result::{RpqMatrixResult}; +use pathrex::rpq::rpqmatrix::eval::RpqMatrixEvaluator; +use pathrex::rpq::rpqmatrix::result::RpqMatrixResult; use pathrex::rpq::{Endpoint, PathExpr, PreparedRpq, RpqError, RpqEvaluator, RpqQuery}; use pathrex::sparql::parse_rpq; use pathrex::utils::build_graph; @@ -183,8 +183,12 @@ fn prepared_rpqmatrix_execution_matches_evaluate() { var("y"), ); - let direct = RpqMatrixEvaluator::default().evaluate(&query, &graph).expect("direct"); - let mut prepared = RpqMatrixEvaluator::default().prepare(&query, &graph).expect("prepare"); + let direct = RpqMatrixEvaluator::default() + .evaluate(&query, &graph) + .expect("direct"); + let mut prepared = RpqMatrixEvaluator::default() + .prepare(&query, &graph) + .expect("prepare"); let prepared_result = prepared.execute().expect("execute"); assert_eq!(prepared_result.nnz, direct.nnz); @@ -199,7 +203,9 @@ fn prepared_rpqmatrix_execution_can_run_twice() { var("y"), ); - let mut prepared = RpqMatrixEvaluator::default().prepare(&query, &graph).expect("prepare"); + let mut prepared = RpqMatrixEvaluator::default() + .prepare(&query, &graph) + .expect("prepare"); let first = prepared.execute().expect("first"); let second = prepared.execute().expect("second"); From 96883d5f0d844827333dceea9874bd5402173c19 Mon Sep 17 00:00:00 2001 From: Rodion Suvorov Date: Sat, 18 Jul 2026 12:14:08 +0300 Subject: [PATCH 5/8] feat: add unit tests for cardinality-based optimizer --- pathrex/Cargo.toml | 1 + pathrex/src/rpq/rpqmatrix/cost.rs | 387 +++++++++++++++++++++++++++++- pathrex/tests/rpqmatrix_tests.rs | 9 + 3 files changed, 388 insertions(+), 9 deletions(-) diff --git a/pathrex/Cargo.toml b/pathrex/Cargo.toml index b439e61..c2ea5ae 100644 --- a/pathrex/Cargo.toml +++ b/pathrex/Cargo.toml @@ -46,6 +46,7 @@ bench = ["clap", "serde", "serde_json", "chrono", "criterion", "tempfile"] [dev-dependencies] tempfile = "3" +expect-test = "1.5.1" [[bin]] name = "pathrex" diff --git a/pathrex/src/rpq/rpqmatrix/cost.rs b/pathrex/src/rpq/rpqmatrix/cost.rs index b3bdd67..c3632bb 100644 --- a/pathrex/src/rpq/rpqmatrix/cost.rs +++ b/pathrex/src/rpq/rpqmatrix/cost.rs @@ -44,7 +44,8 @@ pub struct CardinalityCostFn { pub lr_multiplier: f64, } -// TODO: check value intervals +// TODO: enforce or encode `n > 0`; several estimates divide by `n` or `n^2`. +// TODO: decide whether all estimated cardinalities should be clamped to `[0, n^2]`. impl CostFunction for CardinalityCostFn { type Cost = CardCost; @@ -59,6 +60,7 @@ impl CostFunction for CardinalityCostFn { nnz_r: 1 as f64, nnz_c: 1 as f64, }, + RpqPlan::Label(meta) => CardCost { score: 0.0, nnz: meta.nvals as f64, @@ -74,6 +76,7 @@ impl CostFunction for CardinalityCostFn { let op_cost = (ca.nnz * cb.nnz) / denom; let score = ca.score + cb.score + op_cost; + // TODO: this can exceed `n^2` when child estimates are already loose. let nnz_est = ca.nnz * cb.nnz / (self.n * self.n); CardCost { @@ -88,6 +91,7 @@ impl CostFunction for CardinalityCostFn { let ca = costs(*a); let cb = costs(*b); + // TODO: score uses the raw union estimate; decide if it should be clamped too. let overlap = (ca.nnz * cb.nnz) / (self.n * self.n); let op_cost = ca.nnz + cb.nnz - overlap; let score = ca.score + cb.score + op_cost; @@ -113,6 +117,7 @@ impl CostFunction for CardinalityCostFn { RpqPlan::Star([a]) => { let ca = costs(*a); + // TODO: full dense closure is a conservative upper bound, not a tight estimate. let penalty = self.star_penalty * ca.nnz.max(1.0); let score = ca.score + penalty; @@ -128,6 +133,8 @@ impl CostFunction for CardinalityCostFn { let ca = costs(*a); let cb = costs(*b); + // TODO: LStar/RStar currently reuse Seq-like row/column estimates and do not + // model the closure side directly. let denom = ca.nnz_r.max(cb.nnz_c).max(1.0); let base = (ca.nnz * cb.nnz) / denom; let op_cost = self.lr_multiplier * base; @@ -147,6 +154,8 @@ impl CostFunction for CardinalityCostFn { let ca = costs(*a); let cb = costs(*b); + // TODO: LStar/RStar currently reuse Seq-like row/column estimates and do not + // model the closure side directly. let denom = ca.nnz_r.max(cb.nnz_c).max(1.0); let base = (ca.nnz * cb.nnz) / denom; @@ -166,13 +175,373 @@ impl CostFunction for CardinalityCostFn { } } -pub struct RandomCostFn; -impl CostFunction for RandomCostFn { - type Cost = f64; - fn cost(&mut self, _enode: &RpqPlan, _costs: C) -> Self::Cost - where - C: FnMut(Id) -> Self::Cost, - { - rand::random() +// TODO: random cost fn for evaluating of accuracy of our solution +// pub struct _RandomCostFn; +// impl CostFunction for RandomCostFn { +// type Cost = f64; +// fn cost(&mut self, _enode: &RpqPlan, _costs: C) -> Self::Cost +// where +// C: FnMut(Id) -> Self::Cost, +// { +// rand::random() +// } +// } + +#[cfg(test)] +mod tests { + use egg::RecExpr; + + use crate::rpq::rpqmatrix::{ + optimize::optimize_expr_cardinality, + plan::{LabelMeta, RpqPlan}, + }; + + use super::*; + + fn assert_finite_nonnegative(cost: &CardCost) { + assert!(cost.score.is_finite(), "score must be finite: {cost:?}"); + assert!(cost.nnz.is_finite(), "nnz must be finite: {cost:?}"); + assert!(cost.nnz_r.is_finite(), "nnz_r must be finite: {cost:?}"); + assert!(cost.nnz_c.is_finite(), "nnz_c must be finite: {cost:?}"); + + assert!(cost.score >= 0.0, "score must be non-negative: {cost:?}"); + assert!(cost.nnz >= 0.0, "nnz must be non-negative: {cost:?}"); + assert!(cost.nnz_r >= 0.0, "nnz_r must be non-negative: {cost:?}"); + assert!(cost.nnz_c >= 0.0, "nnz_c must be non-negative: {cost:?}"); + } + + fn child_cost(id: Id, a: Id, ca: &CardCost, b: Id, cb: &CardCost) -> CardCost { + if id == a { + ca.clone() + } else if id == b { + cb.clone() + } else { + panic!("unexpected child id: {id:?}") + } + } + + fn unary_child_cost(id: Id, child: Id, cost: &CardCost) -> CardCost { + if id == child { + cost.clone() + } else { + panic!("unexpected child id: {id:?}") + } + } + + #[test] + fn card_cost_order_uses_score_then_nnz_then_rows_then_cols() { + let base = CardCost { + score: 10.0, + nnz: 20.0, + nnz_r: 30.0, + nnz_c: 40.0, + }; + + assert!( + CardCost { + score: 9.0, + nnz: 100.0, + nnz_r: 100.0, + nnz_c: 100.0, + } < base + ); + assert!( + CardCost { + score: 10.0, + nnz: 19.0, + nnz_r: 100.0, + nnz_c: 100.0, + } < base + ); + assert!( + CardCost { + score: 10.0, + nnz: 20.0, + nnz_r: 29.0, + nnz_c: 100.0, + } < base + ); + assert!( + CardCost { + score: 10.0, + nnz: 20.0, + nnz_r: 30.0, + nnz_c: 39.0, + } < base + ); + } + + #[test] + fn cardinality_cost_base_nodes_use_vertex_and_label_metadata() { + let mut cost_fn = CardinalityCostFn { + n: 100.0, + star_penalty: 50.0, + lr_multiplier: 5.0, + }; + + let named = cost_fn.cost(&RpqPlan::NamedVertex("A".to_string()), |_| { + panic!("NamedVertex must not request child costs") + }); + assert_eq!( + named, + CardCost { + score: 0.0, + nnz: 1.0, + nnz_r: 1.0, + nnz_c: 1.0, + } + ); + + let label = cost_fn.cost( + &RpqPlan::Label(LabelMeta { + name: "knows".to_string(), + nvals: 17, + nonzero_rows: 5, + nonzero_cols: 9, + }), + |_| panic!("Label must not request child costs"), + ); + assert_eq!( + label, + CardCost { + score: 0.0, + nnz: 17.0, + nnz_r: 5.0, + nnz_c: 9.0, + } + ); + } + #[test] + fn cardinality_cost_seq_correctness() { + let mut cost_fn = CardinalityCostFn { + n: 100.0, + star_penalty: 50.0, + lr_multiplier: 5.0, + }; + let a = Id::from(0); + let b = Id::from(1); + + let ca = CardCost { + score: 2.0, + nnz: 20.0, + nnz_r: 4.0, + nnz_c: 8.0, + }; + + let cb = CardCost { + score: 3.0, + nnz: 30.0, + nnz_r: 6.0, + nnz_c: 10.0, + }; + + let seq = cost_fn.cost(&RpqPlan::Seq([a, b]), |id| { + if id == a { + ca.clone() + } else if id == b { + cb.clone() + } else { + panic!("unexpected child id: {id:?}") + } + }); + assert_eq!( + seq, + CardCost { + score: 65.0, + nnz: 0.06, + nnz_r: 4.0, + nnz_c: 10.0, + } + ); + } + #[test] + fn cardinality_cost_seq_correctness_zero_denom() { + let mut cost_fn = CardinalityCostFn { + n: 100.0, + star_penalty: 50.0, + lr_multiplier: 5.0, + }; + let a = Id::from(0); + let b = Id::from(1); + + let ca = CardCost { + score: 2.0, + nnz: 20.0, + nnz_r: 0.0, + nnz_c: 4.0, + }; + + let cb = CardCost { + score: 3.0, + nnz: 30.0, + nnz_r: 4.0, + nnz_c: 0.0, + }; + + let seq = cost_fn.cost(&RpqPlan::Seq([a, b]), |id| { + if id == a { + ca.clone() + } else if id == b { + cb.clone() + } else { + panic!("unexpected child id: {id:?}") + } + }); + assert_eq!( + seq, + CardCost { + score: 605.0, + nnz: 0.06, + nnz_r: 0.0, + nnz_c: 0.0, + } + ); + } + #[test] + fn cardinality_cost_alt_with_zero_children_stays_finite_nonnegative() { + let mut cost_fn = CardinalityCostFn { + n: 100.0, + star_penalty: 50.0, + lr_multiplier: 5.0, + }; + let a = Id::from(0); + let b = Id::from(1); + let ca = CardCost { + score: 0.0, + nnz: 0.0, + nnz_r: 0.0, + nnz_c: 0.0, + }; + let cb = CardCost { + score: 3.0, + nnz: 30.0, + nnz_r: 0.0, + nnz_c: 10.0, + }; + + let alt = cost_fn.cost(&RpqPlan::Alt([a, b]), |id| child_cost(id, a, &ca, b, &cb)); + + assert_finite_nonnegative(&alt); + assert_eq!( + alt, + CardCost { + score: 33.0, + nnz: 30.0, + nnz_r: 0.0, + nnz_c: 10.0, + } + ); + } + + #[test] + fn cardinality_cost_star_with_zero_nnz_uses_min_penalty() { + let mut cost_fn = CardinalityCostFn { + n: 100.0, + star_penalty: 50.0, + lr_multiplier: 5.0, + }; + let a = Id::from(0); + let ca = CardCost { + score: 7.0, + nnz: 0.0, + nnz_r: 0.0, + nnz_c: 0.0, + }; + + let star = cost_fn.cost(&RpqPlan::Star([a]), |id| unary_child_cost(id, a, &ca)); + + assert_finite_nonnegative(&star); + assert_eq!( + star, + CardCost { + score: 57.0, + nnz: 10_000.0, + nnz_r: 100.0, + nnz_c: 100.0, + } + ); + } + + #[test] + fn cardinality_cost_lstar_and_rstar_zero_denom_stay_finite_nonnegative() { + let mut cost_fn = CardinalityCostFn { + n: 100.0, + star_penalty: 50.0, + lr_multiplier: 5.0, + }; + let a = Id::from(0); + let b = Id::from(1); + let ca = CardCost { + score: 2.0, + nnz: 20.0, + nnz_r: 0.0, + nnz_c: 4.0, + }; + let cb = CardCost { + score: 3.0, + nnz: 30.0, + nnz_r: 4.0, + nnz_c: 0.0, + }; + + let lstar = cost_fn.cost(&RpqPlan::LStar([a, b]), |id| child_cost(id, a, &ca, b, &cb)); + let rstar = cost_fn.cost(&RpqPlan::RStar([a, b]), |id| child_cost(id, a, &ca, b, &cb)); + + assert_finite_nonnegative(&lstar); + assert_finite_nonnegative(&rstar); + let expected = CardCost { + score: 3005.0, + nnz: 0.3, + nnz_r: 0.0, + nnz_c: 0.0, + }; + + assert_eq!(lstar, expected); + assert_eq!(rstar, expected); + } + #[test] + fn cardinality_cost_build_lstar() { + let mut expr = RecExpr::default(); + let a = expr.add(RpqPlan::Label(LabelMeta { + name: "knows".to_string(), + nvals: 17, + nonzero_rows: 5, + nonzero_cols: 9, + })); + let b = expr.add(RpqPlan::Label(LabelMeta { + name: "knows".to_string(), + nvals: 17, + nonzero_rows: 5, + nonzero_cols: 9, + })); + let star = expr.add(RpqPlan::Star([a])); + let _seq = expr.add(RpqPlan::Seq([star, b])); + let opt = optimize_expr_cardinality(expr, 100); + let root = opt.as_ref().last().expect("optimized expr is non-empty"); + + assert!(matches!(root, RpqPlan::LStar(_))); + } + #[test] + fn cardinality_cost_build_rstar() { + let mut expr = RecExpr::default(); + let a = expr.add(RpqPlan::Label(LabelMeta { + name: "knows".to_string(), + nvals: 17, + nonzero_rows: 5, + nonzero_cols: 9, + })); + let b = expr.add(RpqPlan::Label(LabelMeta { + name: "knows".to_string(), + nvals: 17, + nonzero_rows: 5, + nonzero_cols: 9, + })); + let star = expr.add(RpqPlan::Star([b])); + let _seq = expr.add(RpqPlan::Seq([a, star])); + let opt = optimize_expr_cardinality(expr, 100); + let root = opt.as_ref().last().expect("optimized expr is non-empty"); + + assert!(matches!(root, RpqPlan::RStar(_))); } + //TODO: maybe cover other rules } diff --git a/pathrex/tests/rpqmatrix_tests.rs b/pathrex/tests/rpqmatrix_tests.rs index 36ec0a5..6ce5d0d 100644 --- a/pathrex/tests/rpqmatrix_tests.rs +++ b/pathrex/tests/rpqmatrix_tests.rs @@ -501,3 +501,12 @@ fn test_la_n_egg_any_con() { fn test_la_n_egg_con_any() { run_la_n_egg_case("con-any"); } + +#[test] +fn test_cardinality_optimizer_give_same_result_unoptimized_way_1() {} + +#[test] +fn test_cardinality_optimizer_give_same_result_unoptimized_way_2() {} + +#[test] +fn test_cardinality_optimizer_give_same_result_unoptimized_way_3() {} From 4cb4e60f4c4d72bf7dda5abc7905a38ddc91a6b6 Mon Sep 17 00:00:00 2001 From: Rodion Suvorov Date: Sat, 18 Jul 2026 12:52:48 +0300 Subject: [PATCH 6/8] feat: add integration tests --- pathrex/tests/rpqmatrix_tests.rs | 150 +++++++++++++++++++++++++++++-- 1 file changed, 145 insertions(+), 5 deletions(-) diff --git a/pathrex/tests/rpqmatrix_tests.rs b/pathrex/tests/rpqmatrix_tests.rs index 6ce5d0d..f73a7e8 100644 --- a/pathrex/tests/rpqmatrix_tests.rs +++ b/pathrex/tests/rpqmatrix_tests.rs @@ -3,9 +3,11 @@ use std::io::{BufRead, BufReader}; use std::path::Path; use std::sync::LazyLock; +use pathrex::eval::ResultCount; use pathrex::formats::mm::MatrixMarket; use pathrex::graph::{Graph, GraphDecomposition, GraphError, InMemory, InMemoryGraph}; use pathrex::lagraph_sys::{GrB_Index, GrB_Info, GrB_Matrix_extractElement_BOOL}; +use pathrex::rpq::rpqmatrix::OptimizationStrategy::Cardinality; use pathrex::rpq::rpqmatrix::eval::RpqMatrixEvaluator; use pathrex::rpq::rpqmatrix::result::RpqMatrixResult; use pathrex::rpq::{Endpoint, PathExpr, PreparedRpq, RpqError, RpqEvaluator, RpqQuery}; @@ -68,7 +70,7 @@ fn load_expected_nnz(case_dir: &Path) -> Vec { .collect() } -fn run_la_n_egg_case(case_name: &str) { +fn run_la_n_egg_case_with_evaluator(case_name: &str, evaluator: RpqMatrixEvaluator) { let case_dir = Path::new(CASES_DIR).join(case_name); let queries = load_queries(&case_dir); let expected = load_expected_nnz(&case_dir); @@ -80,7 +82,6 @@ fn run_la_n_egg_case(case_name: &str) { ); let graph = &*LA_N_EGG_GRAPH; - let evaluator = RpqMatrixEvaluator::default(); for (i, (query, expected_nnz)) in queries.iter().zip(expected.iter()).enumerate() { let result = evaluator.evaluate(query, graph).unwrap_or_else(|e| { @@ -96,6 +97,14 @@ fn run_la_n_egg_case(case_name: &str) { } } +fn run_la_n_egg_case(case_name: &str) { + run_la_n_egg_case_with_evaluator(case_name, RpqMatrixEvaluator::default()); +} + +fn run_la_n_egg_case_cardinality(case_name: &str) { + run_la_n_egg_case_with_evaluator(case_name, RpqMatrixEvaluator::optimized(Cardinality)); +} + fn label(s: &str) -> PathExpr { PathExpr::Label(s.to_string()) } @@ -124,6 +133,31 @@ fn matrix_entry_set(result: &RpqMatrixResult, row: GrB_Index, col: GrB_Index) -> } } +// TODO: made it reusable for different optimizers +fn evaluate_default_and_cardinality( + graph: &InMemoryGraph, + query: &RpqQuery, +) -> (RpqMatrixResult, RpqMatrixResult) { + let default_result = RpqMatrixEvaluator::default() + .evaluate(query, graph) + .expect("default evaluator should succeed"); + let optimized_result = RpqMatrixEvaluator::optimized(Cardinality) + .evaluate(query, graph) + .expect("cardinality optimizer should succeed"); + + assert_eq!( + default_result.nnz, optimized_result.nnz, + "optimized evaluator should preserve result nnz" + ); + assert_eq!( + default_result.result_count().expect("default count"), + optimized_result.result_count().expect("optimized count"), + "optimized evaluator should preserve result count" + ); + + (default_result, optimized_result) +} + /// Graph: A --knows--> B --knows--> C /// Query: ?x ?y #[test] @@ -503,10 +537,116 @@ fn test_la_n_egg_con_any() { } #[test] -fn test_cardinality_optimizer_give_same_result_unoptimized_way_1() {} +fn test_la_n_egg_any_any_cardinality_optimizer() { + run_la_n_egg_case_cardinality("any-any"); +} #[test] -fn test_cardinality_optimizer_give_same_result_unoptimized_way_2() {} +fn test_la_n_egg_any_con_cardinality_optimizer() { + run_la_n_egg_case_cardinality("con-any"); +} #[test] -fn test_cardinality_optimizer_give_same_result_unoptimized_way_3() {} +fn test_cardinality_optimizer_give_same_result_unoptimized_way_1() { + let graph = build_graph(&[ + ("A", "B", "knows"), + ("B", "C", "knows"), + ("C", "D", "likes"), + ("A", "E", "likes"), + ]); + + // knows* / likes can be rewritten to LStar. + let path = PathExpr::Sequence( + Box::new(PathExpr::ZeroOrMore(Box::new(label("knows")))), + Box::new(label("likes")), + ); + let query = rq(named_ep("A"), path, var("y")); + + let (default_result, optimized_result) = evaluate_default_and_cardinality(&graph, &query); + assert_eq!(default_result.nnz, 2); + + let a_id = graph.get_node_id("A").expect("A should exist") as GrB_Index; + let d_id = graph.get_node_id("D").expect("D should exist") as GrB_Index; + let e_id = graph.get_node_id("E").expect("E should exist") as GrB_Index; + + for result in [&default_result, &optimized_result] { + assert!( + matrix_entry_set(result, a_id, d_id), + "D should be reachable via knows*/likes" + ); + assert!( + matrix_entry_set(result, a_id, e_id), + "E should be reachable via zero knows hops then likes" + ); + } +} + +#[test] +fn test_cardinality_optimizer_give_same_result_unoptimized_way_2() { + let graph = build_graph(&[ + ("A", "B", "knows"), + ("B", "C", "likes"), + ("C", "D", "knows"), + ]); + + // knows / likes* / knows + let path = PathExpr::Sequence( + Box::new(PathExpr::Sequence( + Box::new(label("knows")), + Box::new(PathExpr::ZeroOrMore(Box::new(label("likes")))), + )), + Box::new(label("knows")), + ); + + let query = rq(named_ep("A"), path, var("y")); + let (default_result, optimized_result) = evaluate_default_and_cardinality(&graph, &query); + + assert_eq!(default_result.nnz, 1); + let a_id = graph.get_node_id("A").expect("A should exist") as GrB_Index; + let d_id = graph.get_node_id("D").expect("D should exist") as GrB_Index; + assert!( + matrix_entry_set(&default_result, a_id, d_id), + "D should be reachable via knows/likes*/knows" + ); + assert!( + matrix_entry_set(&optimized_result, a_id, d_id), + "D should be reachable via knows/likes*/knows" + ); +} + +#[test] +fn test_cardinality_optimizer_give_same_result_unoptimized_way_3() { + let graph = build_graph(&[ + ("A", "B", "knows"), + ("B", "C", "likes"), + ("B", "D", "hates"), + ]); + + // knows / (likes | hates) can be rewritten by distributivity rules. + let path = PathExpr::Sequence( + Box::new(label("knows")), + Box::new(PathExpr::Alternative( + Box::new(label("likes")), + Box::new(label("hates")), + )), + ); + let query = rq(named_ep("A"), path, var("y")); + + let (default_result, optimized_result) = evaluate_default_and_cardinality(&graph, &query); + assert_eq!(default_result.nnz, 2); + + let a_id = graph.get_node_id("A").expect("A should exist") as GrB_Index; + let c_id = graph.get_node_id("C").expect("C should exist") as GrB_Index; + let d_id = graph.get_node_id("D").expect("D should exist") as GrB_Index; + + for result in [&default_result, &optimized_result] { + assert!( + matrix_entry_set(result, a_id, c_id), + "C should be reachable via knows/likes" + ); + assert!( + matrix_entry_set(result, a_id, d_id), + "D should be reachable via knows/hates" + ); + } +} From 97824238b24f33a7f93a3271f234f11e1a2b889e Mon Sep 17 00:00:00 2001 From: Rodion Suvorov Date: Sat, 18 Jul 2026 13:08:58 +0300 Subject: [PATCH 7/8] feat: add tests for metadata --- pathrex/src/graph/inmemory.rs | 3 --- pathrex/src/graph/wrappers.rs | 12 ++++----- pathrex/tests/mm_tests.rs | 28 ++++++++++++++++++++ pathrex/tests/testdata/mm_small/1.txt | 3 +++ pathrex/tests/testdata/mm_small/edges.txt | 3 +++ pathrex/tests/testdata/mm_small/vertices.txt | 3 +++ 6 files changed, 43 insertions(+), 9 deletions(-) create mode 100644 pathrex/tests/testdata/mm_small/1.txt create mode 100644 pathrex/tests/testdata/mm_small/edges.txt create mode 100644 pathrex/tests/testdata/mm_small/vertices.txt diff --git a/pathrex/src/graph/inmemory.rs b/pathrex/src/graph/inmemory.rs index 043e536..9274a68 100644 --- a/pathrex/src/graph/inmemory.rs +++ b/pathrex/src/graph/inmemory.rs @@ -454,7 +454,4 @@ mod tests { assert!(graph.get_graph("http://example.org/knows").is_ok()); assert!(graph.get_graph("http://example.org/likes").is_ok()); } - - #[test] - fn test_metadata_from_mm() {} } diff --git a/pathrex/src/graph/wrappers.rs b/pathrex/src/graph/wrappers.rs index c684d60..137b7f8 100644 --- a/pathrex/src/graph/wrappers.rs +++ b/pathrex/src/graph/wrappers.rs @@ -179,16 +179,16 @@ impl LagraphGraph { pub fn nonzero_cols(&self) -> Result { let matrix: GrB_Matrix = unsafe { (*self.inner).A }; - let res: usize = 0; - unsafe { LAGraph_RPQMatrix_reduce(&mut (res as u64), matrix, ByRows as u8) }; - Ok(res) + let mut res: GrB_Index = 0; + unsafe { LAGraph_RPQMatrix_reduce(&mut res, matrix, ByRows as u8) }; + Ok(res as usize) } pub fn nonzero_rows(&self) -> Result { let matrix: GrB_Matrix = unsafe { (*self.inner).A }; - let res: usize = 0; - unsafe { LAGraph_RPQMatrix_reduce(&mut (res as u64), matrix, ByCols as u8) }; - Ok(res) + let mut res: GrB_Index = 0; + unsafe { LAGraph_RPQMatrix_reduce(&mut res, matrix, ByCols as u8) }; + Ok(res as usize) } } diff --git a/pathrex/tests/mm_tests.rs b/pathrex/tests/mm_tests.rs index 0998e1b..f914ccd 100644 --- a/pathrex/tests/mm_tests.rs +++ b/pathrex/tests/mm_tests.rs @@ -215,3 +215,31 @@ fn test_mm_graph_empty_label_handling() { let result = graph.get_graph(""); assert!(result.is_err(), "Empty label should not exist in the graph"); } + +#[test] +fn test_mm_graph_correct_metadata() { + let mm = MatrixMarket::from_dir("tests/testdata/mm_small"); + let graph = Graph::::try_from(mm).expect("Failed to load graph"); + + let result = graph + .get_metadata() + .expect("metadata should exist") + .matrix("knows") + .expect("matrix with metadata should exist"); + assert!( + result.dimension == 4, + "dimension of matrix should be calculated correctly" + ); + assert!( + result.nvals == 3, + "nonzero vals of matrix should be calculated correctly" + ); + assert!( + result.nonzero_cols == 2, + "nonzero columns of matrix should be calculated correctly" + ); + assert!( + result.nonzero_rows == 2, + "nonzero rows of matrix should be calculated correctly" + ); +} diff --git a/pathrex/tests/testdata/mm_small/1.txt b/pathrex/tests/testdata/mm_small/1.txt new file mode 100644 index 0000000..cd8561a --- /dev/null +++ b/pathrex/tests/testdata/mm_small/1.txt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:cde9a9330825c73e01fe5fe1c0be5d226cef1b63f314ed410c7a7620c717a623 +size 66 diff --git a/pathrex/tests/testdata/mm_small/edges.txt b/pathrex/tests/testdata/mm_small/edges.txt new file mode 100644 index 0000000..bda8e20 --- /dev/null +++ b/pathrex/tests/testdata/mm_small/edges.txt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b107c484753ed13cc4e1e10c97ebda2e327e3be24d7fc2e84d62075c8f0909f5 +size 9 diff --git a/pathrex/tests/testdata/mm_small/vertices.txt b/pathrex/tests/testdata/mm_small/vertices.txt new file mode 100644 index 0000000..79ef958 --- /dev/null +++ b/pathrex/tests/testdata/mm_small/vertices.txt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d7e475829a22013136702d92adef2253a9fae510fff62dba7975c7062e27c425 +size 23 From 73331c45bfab18a12a3fff07af374003884fce70 Mon Sep 17 00:00:00 2001 From: Rodion Suvorov Date: Sat, 18 Jul 2026 22:12:17 +0300 Subject: [PATCH 8/8] feat: add fixed number of runs benchmark mode --- pathrex/src/bin/pathrex.rs | 143 +++++++++++++---- pathrex/src/cli/args.rs | 149 ++++++++++++++++-- pathrex/src/cli/bench/runner.rs | 125 +++++++++++++-- pathrex/src/cli/checkpoint.rs | 90 ++++++++++- pathrex/src/cli/output.rs | 51 +++++- .../tests/testdata/cases/any-con/queries.txt | 4 +- 6 files changed, 490 insertions(+), 72 deletions(-) diff --git a/pathrex/src/bin/pathrex.rs b/pathrex/src/bin/pathrex.rs index 6b4339a..67ffe58 100644 --- a/pathrex/src/bin/pathrex.rs +++ b/pathrex/src/bin/pathrex.rs @@ -17,6 +17,7 @@ //! --graph tests/testdata/mm_graph \ //! --queries tests/testdata/cases/any-any/queries.txt \ //! --algo nfarpq rpqmatrix \ +//! --bench-mode criterion \ //! --rpqmatrix-optimizer cardinality \ //! --output results.json //! ``` @@ -29,10 +30,11 @@ use clap::Parser; use thiserror::Error; use pathrex::cli::args::{ - Algo, BenchArgs, Cli, Commands, CommonArgs, GraphFormat, QueryArgs, RpqMatrixOptimizer, + Algo, BenchArgs, BenchMode, Cli, Commands, CommonArgs, GraphFormat, QueryArgs, + RpqMatrixOptimizer, }; use pathrex::cli::bench::BenchError; -use pathrex::cli::checkpoint::{Checkpoint, CheckpointError, Checkpointer}; +use pathrex::cli::checkpoint::{BenchRunConfig, Checkpoint, CheckpointError, Checkpointer}; use pathrex::cli::dispatch::{dispatch_bench, dispatch_query}; use pathrex::cli::loader::{GraphLoadError, LoadedQuery, load_graph, load_queries}; use pathrex::cli::output::{BenchMetadata, BenchOutput, QueryMetadata, QueryOutput}; @@ -103,6 +105,55 @@ fn validate_common_args(common: &CommonArgs) -> Result<(), MainError> { Ok(()) } +fn validate_bench_args(args: &BenchArgs) -> Result<(), MainError> { + validate_common_args(&args.common)?; + + if args.resume && args.checkpoint.is_none() { + return Err(MainError::InvalidArgs( + "--resume requires --checkpoint".to_string(), + )); + } + + match args.bench_mode { + BenchMode::Fixed => { + if args.criterion_dir.is_some() + || args.plots + || args.sample_size.is_some() + || args.warm_up.is_some() + || args.measurement.is_some() + { + return Err(MainError::InvalidArgs( + "criterion options can only be used with --bench-mode criterion".to_string(), + )); + } + if args.fixed_runs() == 0 { + return Err(MainError::InvalidArgs( + "--runs must be greater than 0".to_string(), + )); + } + } + BenchMode::Criterion => { + if args.runs.is_some() || args.warm_up_runs.is_some() { + return Err(MainError::InvalidArgs( + "fixed-run options can only be used with --bench-mode fixed".to_string(), + )); + } + if args.plots && args.criterion_dir.is_none() { + return Err(MainError::InvalidArgs( + "--plots requires --criterion-dir".to_string(), + )); + } + if args.criterion_sample_size() < 10 { + return Err(MainError::InvalidArgs( + "--sample-size must be at least 10 for criterion".to_string(), + )); + } + } + } + + Ok(()) +} + fn load_query_file(path: &str, base_iri: Option<&str>) -> Result, MainError> { load_queries(Path::new(path), base_iri).map_err(|e| MainError::Queries { path: path.to_string(), @@ -175,35 +226,49 @@ fn run_query_cmd(args: QueryArgs) -> Result<(), MainError> { fn build_checkpointer(args: &BenchArgs, queries_len: usize) -> Result { let common = &args.common; - let path = PathBuf::from(&args.checkpoint); - - if args.resume { - match Checkpoint::load(&path)? { - Some(cp) => { - cp.validate( - &common.graph, - &common.queries, - &common.algo, - common.rpqmatrix_optimizer, - )?; - let cper = Checkpointer::with_inner(cp, path); - eprintln!( - " resuming: {}/{} queries fully done", - cper.fully_done_count(&common.algo), - queries_len - ); - Ok(cper) - } - None => { - eprintln!(" no checkpoint file found, starting fresh"); - Ok(Checkpointer::fresh( - &common.graph, - &common.queries, - &common.algo, - common.rpqmatrix_optimizer, - path, - )) + let bench_config = BenchRunConfig::from_args(args); + + if let Some(checkpoint) = &args.checkpoint { + let path = PathBuf::from(checkpoint); + if args.resume { + match Checkpoint::load(&path)? { + Some(cp) => { + cp.validate( + &common.graph, + &common.queries, + &common.algo, + common.rpqmatrix_optimizer, + &bench_config, + )?; + let cper = Checkpointer::with_inner(cp, path); + eprintln!( + " resuming: {}/{} queries fully done", + cper.fully_done_count(&common.algo), + queries_len + ); + Ok(cper) + } + None => { + eprintln!(" no checkpoint file found, starting fresh"); + Ok(Checkpointer::fresh( + &common.graph, + &common.queries, + &common.algo, + common.rpqmatrix_optimizer, + bench_config, + Some(path), + )) + } } + } else { + Ok(Checkpointer::fresh( + &common.graph, + &common.queries, + &common.algo, + common.rpqmatrix_optimizer, + bench_config, + Some(path), + )) } } else { Ok(Checkpointer::fresh( @@ -211,20 +276,22 @@ fn build_checkpointer(args: &BenchArgs, queries_len: usize) -> Result Result<(), MainError> { let common = &args.common; - validate_common_args(common)?; + validate_bench_args(&args)?; eprintln!("=== pathrex bench ==="); eprintln!("Graph: {}", common.graph); eprintln!("Format: {}", common.format); eprintln!("Queries: {}", common.queries); eprintln!("Algos: {:?}", common.algo); + eprintln!("Bench mode: {}", args.bench_mode); eprintln!("RPQMatrix optimizer: {}", common.rpqmatrix_optimizer); eprintln!("Output: {}", args.output); eprintln!(); @@ -263,9 +330,15 @@ fn run_bench_cmd(args: BenchArgs) -> Result<(), MainError> { rpqmatrix_optimizer: Some(common.rpqmatrix_optimizer.to_string()), num_nodes: graph.num_nodes(), num_labels: graph.num_labels(), - sample_size: args.sample_size, - warm_up_secs: args.warm_up, - measurement_secs: args.measurement, + bench_mode: args.bench_mode.to_string(), + runs: (args.bench_mode == BenchMode::Fixed).then(|| args.fixed_runs()), + warm_up_runs: (args.bench_mode == BenchMode::Fixed).then(|| args.fixed_warm_up_runs()), + sample_size: (args.bench_mode == BenchMode::Criterion) + .then(|| args.criterion_sample_size()), + warm_up_secs: (args.bench_mode == BenchMode::Criterion) + .then(|| args.criterion_warm_up_secs()), + measurement_secs: (args.bench_mode == BenchMode::Criterion) + .then(|| args.criterion_measurement_secs()), }, results, }; diff --git a/pathrex/src/cli/args.rs b/pathrex/src/cli/args.rs index 73e39ad..cf7b5f8 100644 --- a/pathrex/src/cli/args.rs +++ b/pathrex/src/cli/args.rs @@ -29,7 +29,7 @@ pub struct Cli { pub enum Commands { /// Run queries once and report result counts Query(QueryArgs), - /// Benchmark RPQ evaluators with criterion + /// Benchmark RPQ evaluators Bench(BenchArgs), } @@ -96,14 +96,18 @@ pub struct BenchArgs { #[arg(short = 'o', long, default_value = "bench_results.json")] pub output: String, - /// Checkpoint file path. - #[arg(short = 'c', long, default_value = "bench_checkpoint.json")] - pub checkpoint: String, + /// Optional checkpoint file path. + #[arg(short = 'c', long)] + pub checkpoint: Option, /// Resume from checkpoint, skipping completed queries. #[arg(long)] pub resume: bool, + /// Benchmarking mode. + #[arg(long, value_enum, default_value_t = BenchMode::Fixed)] + pub bench_mode: BenchMode, + /// Directory for criterion output. When omitted, criterion writes into a /// per-group temporary directory that is wiped immediately after each /// benchmark group is parsed (default behavior). @@ -117,16 +121,80 @@ pub struct BenchArgs { pub plots: bool, /// Criterion sample size per benchmark group. - #[arg(long, default_value_t = 10)] - pub sample_size: usize, + #[arg(long)] + pub sample_size: Option, /// Criterion warm-up time in seconds. - #[arg(long, default_value_t = 1)] - pub warm_up: u64, + #[arg(long)] + pub warm_up: Option, /// Criterion measurement time in seconds. - #[arg(long, default_value_t = 5)] - pub measurement: u64, + #[arg(long)] + pub measurement: Option, + + /// Number of warm-up runs. + #[arg(long = "warm-up-runs")] + pub warm_up_runs: Option, + + /// Number of measured runs. + #[arg(long)] + pub runs: Option, +} + +impl BenchArgs { + pub const DEFAULT_FIXED_RUNS: u64 = 10; + pub const DEFAULT_FIXED_WARM_UP_RUNS: u64 = 0; + pub const DEFAULT_CRITERION_SAMPLE_SIZE: usize = 10; + pub const DEFAULT_CRITERION_WARM_UP_SECS: u64 = 1; + pub const DEFAULT_CRITERION_MEASUREMENT_SECS: u64 = 5; + + pub fn fixed_runs(&self) -> u64 { + self.runs.unwrap_or(Self::DEFAULT_FIXED_RUNS) + } + + pub fn fixed_warm_up_runs(&self) -> u64 { + self.warm_up_runs + .unwrap_or(Self::DEFAULT_FIXED_WARM_UP_RUNS) + } + + pub fn criterion_sample_size(&self) -> usize { + self.sample_size + .unwrap_or(Self::DEFAULT_CRITERION_SAMPLE_SIZE) + } + + pub fn criterion_warm_up_secs(&self) -> u64 { + self.warm_up.unwrap_or(Self::DEFAULT_CRITERION_WARM_UP_SECS) + } + + pub fn criterion_measurement_secs(&self) -> u64 { + self.measurement + .unwrap_or(Self::DEFAULT_CRITERION_MEASUREMENT_SECS) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "lowercase")] +#[value(rename_all = "lowercase")] +pub enum BenchMode { + /// Fixed number of runs per query. + Fixed, + /// Criterion time-based benchmark. + Criterion, +} + +impl Default for BenchMode { + fn default() -> Self { + Self::Fixed + } +} + +impl std::fmt::Display for BenchMode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + BenchMode::Fixed => write!(f, "fixed"), + BenchMode::Criterion => write!(f, "criterion"), + } + } } #[derive(Debug, Clone, PartialEq, Eq, Hash, ValueEnum, serde::Serialize, serde::Deserialize)] @@ -221,4 +289,65 @@ mod tests { assert!(result.is_err()); } + + #[test] + fn bench_defaults_to_fixed_runs_without_checkpoint_or_criterion() { + let cli = Cli::parse_from([ + "pathrex", + "bench", + "--graph", + "graph", + "--queries", + "queries", + "--algo", + "rpqmatrix", + ]); + + let Commands::Bench(args) = cli.command else { + panic!("expected bench command"); + }; + + assert_eq!(args.bench_mode, BenchMode::Fixed); + assert_eq!(args.fixed_runs(), BenchArgs::DEFAULT_FIXED_RUNS); + assert_eq!( + args.fixed_warm_up_runs(), + BenchArgs::DEFAULT_FIXED_WARM_UP_RUNS + ); + assert!(args.checkpoint.is_none()); + assert!(args.criterion_dir.is_none()); + assert!(args.sample_size.is_none()); + assert!(args.warm_up.is_none()); + assert!(args.measurement.is_none()); + } + + #[test] + fn criterion_mode_accepts_optional_criterion_settings() { + let cli = Cli::parse_from([ + "pathrex", + "bench", + "--graph", + "graph", + "--queries", + "queries", + "--algo", + "rpqmatrix", + "--bench-mode", + "criterion", + "--sample-size", + "20", + "--warm-up", + "2", + "--measurement", + "7", + ]); + + let Commands::Bench(args) = cli.command else { + panic!("expected bench command"); + }; + + assert_eq!(args.bench_mode, BenchMode::Criterion); + assert_eq!(args.criterion_sample_size(), 20); + assert_eq!(args.criterion_warm_up_secs(), 2); + assert_eq!(args.criterion_measurement_secs(), 7); + } } diff --git a/pathrex/src/cli/bench/runner.rs b/pathrex/src/cli/bench/runner.rs index 4b5a54a..97cf885 100644 --- a/pathrex/src/cli/bench/runner.rs +++ b/pathrex/src/cli/bench/runner.rs @@ -1,15 +1,15 @@ use std::collections::HashMap; use std::path::{Path, PathBuf}; -use std::time::Duration; +use std::time::{Duration, Instant}; -use criterion::{Criterion, black_box}; +use criterion::{BatchSize, Criterion, black_box}; -use crate::cli::args::{Algo, BenchArgs}; +use crate::cli::args::{Algo, BenchArgs, BenchMode}; use crate::cli::bench::error::BenchError; use crate::cli::bench::estimates::read_algo_timing; use crate::cli::checkpoint::Checkpointer; use crate::cli::loader::LoadedQuery; -use crate::cli::output::{AlgoResult, QueryResult}; +use crate::cli::output::{AlgoResult, AlgoTiming, QueryResult, TimingStats}; use crate::eval::{Evaluator, PreparedEvaluator, ResultCount}; use crate::graph::InMemoryGraph; use crate::rpq::{RpqError, RpqQuery}; @@ -41,9 +41,9 @@ impl GroupOutput { pub(crate) fn build_criterion(args: &BenchArgs, output_dir: &Path) -> Criterion { let c = Criterion::default() - .sample_size(args.sample_size) - .warm_up_time(Duration::from_secs(args.warm_up)) - .measurement_time(Duration::from_secs(args.measurement)) + .sample_size(args.criterion_sample_size()) + .warm_up_time(Duration::from_secs(args.criterion_warm_up_secs())) + .measurement_time(Duration::from_secs(args.criterion_measurement_secs())) .output_directory(output_dir); if args.plots { c.with_plots() @@ -68,7 +68,10 @@ where E: Evaluator + Copy, E::Result: ResultCount, { - let mut prepared = evaluator.prepare(query, graph)?; + // Validate preparation once so query/graph errors are reported through the + // normal benchmark error path. The `eval_ffi_only` benchmark below creates a + // fresh prepared state per measured iteration. + let _prepared = evaluator.prepare(query, graph)?; let group = group_name(query_index, algo_name); let output = match GroupOutput::for_group(args) { @@ -89,9 +92,17 @@ where }); g.bench_function("eval_ffi_only", |b| { - b.iter(|| { - let _ = black_box(prepared.execute()); - }); + b.iter_batched( + || { + evaluator + .prepare(query, graph) + .expect("prepare should keep succeeding during benchmark") + }, + |mut prepared| { + let _ = black_box(prepared.execute()); + }, + BatchSize::PerIteration, + ); }); g.finish(); @@ -100,6 +111,83 @@ where Ok(read_algo_timing(&output_path, &group)) } +fn timing_stats(samples_ns: &[f64]) -> TimingStats { + let mut sorted = samples_ns.to_vec(); + sorted.sort_by(f64::total_cmp); + + let len = sorted.len(); + let mean = sorted.iter().sum::() / len as f64; + let median = if len % 2 == 0 { + (sorted[len / 2 - 1] + sorted[len / 2]) / 2.0 + } else { + sorted[len / 2] + }; + let variance = sorted + .iter() + .map(|sample| { + let diff = sample - mean; + diff * diff + }) + .sum::() + / len as f64; + + TimingStats { + mean_ns: mean, + median_ns: median, + stddev_ns: variance.sqrt(), + iterations: len, + } +} + +fn elapsed_ns(start: Instant) -> f64 { + start.elapsed().as_nanos() as f64 +} + +fn run_fixed_group( + args: &BenchArgs, + evaluator: E, + query: &RpqQuery, + graph: &InMemoryGraph, +) -> Result<(usize, AlgoTiming), RpqError> +where + E: Evaluator + Copy, + E::Result: ResultCount, +{ + for _ in 0..args.fixed_warm_up_runs() { + let _ = black_box(evaluator.evaluate(query, graph)?); + } + + let mut total_samples = Vec::with_capacity(args.fixed_runs() as usize); + let mut result_count = None; + for _ in 0..args.fixed_runs() { + let start = Instant::now(); + let result = black_box(evaluator.evaluate(query, graph)?); + total_samples.push(elapsed_ns(start)); + result_count = Some(result.result_count().map_err(RpqError::Graph)?); + } + + for _ in 0..args.fixed_warm_up_runs() { + let mut prepared = evaluator.prepare(query, graph)?; + let _ = black_box(prepared.execute()?); + } + + let mut ffi_samples = Vec::with_capacity(args.fixed_runs() as usize); + for _ in 0..args.fixed_runs() { + let mut prepared = evaluator.prepare(query, graph)?; + let start = Instant::now(); + let _ = black_box(prepared.execute()?); + ffi_samples.push(elapsed_ns(start)); + } + + Ok(( + result_count.unwrap_or(0), + AlgoTiming { + total: timing_stats(&total_samples), + ffi_only: timing_stats(&ffi_samples), + }, + )) +} + /// Run the bench loop for every query in `queries` for one evaluator. pub fn run_bench_for_evaluator( args: &BenchArgs, @@ -150,9 +238,18 @@ where eprintln!("[query #{}] id={}", idx, loaded.id); eprintln!(" [bench] algo={algo_name}"); - match run_benchmark_group(args, algo_name, evaluator, query, graph, idx) { - Ok(Ok(timing)) => { - algorithms.insert(algo_name.to_string(), AlgoResult::ok(None, Some(timing))); + let bench_result = match args.bench_mode { + BenchMode::Fixed => run_fixed_group(args, evaluator, query, graph) + .map(|(count, timing)| Ok((Some(count), timing))), + BenchMode::Criterion => { + run_benchmark_group(args, algo_name, evaluator, query, graph, idx) + .map(|result| result.map(|timing| (None, timing))) + } + }; + + match bench_result { + Ok(Ok((count, timing))) => { + algorithms.insert(algo_name.to_string(), AlgoResult::ok(count, Some(timing))); } Ok(Err(e)) => return Err(e), Err(e) => { diff --git a/pathrex/src/cli/checkpoint.rs b/pathrex/src/cli/checkpoint.rs index 2d4f13b..7de7fc8 100644 --- a/pathrex/src/cli/checkpoint.rs +++ b/pathrex/src/cli/checkpoint.rs @@ -16,7 +16,7 @@ use std::path::{Path, PathBuf}; use serde::{Deserialize, Serialize}; use thiserror::Error; -use super::args::{Algo, RpqMatrixOptimizer}; +use super::args::{Algo, BenchArgs, BenchMode, RpqMatrixOptimizer}; /// Persistent checkpoint state written to disk as JSON. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -27,9 +27,57 @@ pub struct Checkpoint { pub algorithms: Vec, #[serde(default)] pub rpqmatrix_optimizer: RpqMatrixOptimizer, + #[serde(default)] + pub bench_config: BenchRunConfig, pub completed: Vec, } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct BenchRunConfig { + pub bench_mode: BenchMode, + pub runs: Option, + pub warm_up_runs: Option, + pub sample_size: Option, + pub warm_up_secs: Option, + pub measurement_secs: Option, +} + +impl BenchRunConfig { + pub fn from_args(args: &BenchArgs) -> Self { + match args.bench_mode { + BenchMode::Fixed => Self { + bench_mode: args.bench_mode, + runs: Some(args.fixed_runs()), + warm_up_runs: Some(args.fixed_warm_up_runs()), + sample_size: None, + warm_up_secs: None, + measurement_secs: None, + }, + BenchMode::Criterion => Self { + bench_mode: args.bench_mode, + runs: None, + warm_up_runs: None, + sample_size: Some(args.criterion_sample_size()), + warm_up_secs: Some(args.criterion_warm_up_secs()), + measurement_secs: Some(args.criterion_measurement_secs()), + }, + } + } +} + +impl Default for BenchRunConfig { + fn default() -> Self { + Self { + bench_mode: BenchMode::Fixed, + runs: Some(BenchArgs::DEFAULT_FIXED_RUNS), + warm_up_runs: Some(BenchArgs::DEFAULT_FIXED_WARM_UP_RUNS), + sample_size: None, + warm_up_secs: None, + measurement_secs: None, + } + } +} + /// Tracks which algorithms have been completed for a single query. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct QueryCompletion { @@ -44,6 +92,7 @@ impl Checkpoint { queries_file: &str, algorithms: &[Algo], rpqmatrix_optimizer: RpqMatrixOptimizer, + bench_config: BenchRunConfig, ) -> Self { Self { version: 1, @@ -51,6 +100,7 @@ impl Checkpoint { queries_file: queries_file.to_string(), algorithms: algorithms.to_vec(), rpqmatrix_optimizer, + bench_config, completed: Vec::new(), } } @@ -74,6 +124,7 @@ impl Checkpoint { queries_file: &str, algorithms: &[Algo], rpqmatrix_optimizer: RpqMatrixOptimizer, + bench_config: &BenchRunConfig, ) -> Result<(), CheckpointError> { if self.graph_path != graph_path { return Err(CheckpointError::Mismatch(format!( @@ -101,6 +152,12 @@ impl Checkpoint { self.rpqmatrix_optimizer, rpqmatrix_optimizer ))); } + if &self.bench_config != bench_config { + return Err(CheckpointError::Mismatch(format!( + "bench_config: checkpoint has {:?}, current is {:?}", + self.bench_config, bench_config + ))); + } Ok(()) } @@ -108,6 +165,14 @@ impl Checkpoint { pub fn save(&self, path: &Path) -> Result<(), CheckpointError> { let json = serde_json::to_string_pretty(self).map_err(CheckpointError::Serialize)?; + if let Some(parent) = path + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + { + fs::create_dir_all(parent) + .map_err(|e| CheckpointError::Io(parent.display().to_string(), e))?; + } + // Write to a temp file first, then rename for atomicity. let tmp_path = path.with_extension("json.tmp"); fs::write(&tmp_path, &json) @@ -155,7 +220,7 @@ impl Checkpoint { /// Runtime owner for a [`Checkpoint`] paired with its on-disk path. pub struct Checkpointer { inner: Checkpoint, - path: PathBuf, + path: Option, } impl Checkpointer { @@ -165,17 +230,27 @@ impl Checkpointer { queries_file: &str, algorithms: &[Algo], rpqmatrix_optimizer: RpqMatrixOptimizer, - path: PathBuf, + bench_config: BenchRunConfig, + path: Option, ) -> Self { Self { - inner: Checkpoint::new(graph_path, queries_file, algorithms, rpqmatrix_optimizer), + inner: Checkpoint::new( + graph_path, + queries_file, + algorithms, + rpqmatrix_optimizer, + bench_config, + ), path, } } /// Wrap an existing [`Checkpoint`] (e.g. one loaded from disk). pub fn with_inner(inner: Checkpoint, path: PathBuf) -> Self { - Self { inner, path } + Self { + inner, + path: Some(path), + } } /// Number of queries that have *all* requested algorithms done. @@ -205,7 +280,10 @@ impl Checkpointer { algo: &Algo, ) -> Result<(), CheckpointError> { self.inner.mark_algo_done(query_index, algo); - self.inner.save(&self.path) + if let Some(path) = &self.path { + self.inner.save(path)?; + } + Ok(()) } } diff --git a/pathrex/src/cli/output.rs b/pathrex/src/cli/output.rs index c7aa2d1..efcb3ff 100644 --- a/pathrex/src/cli/output.rs +++ b/pathrex/src/cli/output.rs @@ -101,7 +101,7 @@ pub struct QueryMetadata { impl QueryOutput { pub fn write_to_file(&self, path: &Path) -> Result<(), std::io::Error> { let json = serde_json::to_string_pretty(self).map_err(std::io::Error::other)?; - fs::write(path, json) + write_json_to_file(path, json) } } @@ -122,18 +122,36 @@ pub struct BenchMetadata { pub base_iri: Option, pub num_nodes: usize, pub num_labels: usize, - pub sample_size: usize, - pub warm_up_secs: u64, - pub measurement_secs: u64, + pub bench_mode: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub runs: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub warm_up_runs: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub sample_size: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub warm_up_secs: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub measurement_secs: Option, } impl BenchOutput { pub fn write_to_file(&self, path: &Path) -> Result<(), std::io::Error> { let json = serde_json::to_string_pretty(self).map_err(std::io::Error::other)?; - fs::write(path, json) + write_json_to_file(path, json) } } +fn write_json_to_file(path: &Path, json: String) -> Result<(), std::io::Error> { + if let Some(parent) = path + .parent() + .filter(|parent| !parent.as_os_str().is_empty()) + { + fs::create_dir_all(parent)?; + } + fs::write(path, json) +} + #[cfg(test)] mod tests { use super::*; @@ -178,4 +196,27 @@ mod tests { let v = serde_json::to_value(&r).expect("serialize"); assert_eq!(v["status"], "panic"); } + + #[test] + fn query_output_creates_parent_directory() { + let dir = tempfile::tempdir().expect("tempdir"); + let output_path = dir.path().join("nested").join("query.json"); + let output = QueryOutput { + metadata: QueryMetadata { + timestamp: "now".into(), + graph_path: "graph".into(), + graph_format: "mm".into(), + queries_file: "queries".into(), + rpqmatrix_optimizer: Some("none".into()), + base_iri: None, + num_nodes: 0, + num_labels: 0, + }, + results: Vec::new(), + }; + + output.write_to_file(&output_path).expect("write output"); + + assert!(output_path.exists()); + } } diff --git a/pathrex/tests/testdata/cases/any-con/queries.txt b/pathrex/tests/testdata/cases/any-con/queries.txt index d16ebd8..b82088e 100644 --- a/pathrex/tests/testdata/cases/any-con/queries.txt +++ b/pathrex/tests/testdata/cases/any-con/queries.txt @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:c05d3be471b0ba767b1d0257690a34fbbe81c7ef6bfef51c5077447ecfcc1ada -size 1129 +oid sha256:cd0217a7caff9f056c5d1e132a6772c6de11cbddb2f5ac2517182d1e776a42b7 +size 1131