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/README.md b/README.md index 79e682c..5c0e8c2 100644 --- a/README.md +++ b/README.md @@ -97,6 +97,8 @@ Subcommands: | `-q`, `--queries ` | Queries file (see format below). | | `-a`, `--algo ` | Algorithm(s). Repeat to run several. | | `-b`, `--base-iri []` | Optional `BASE ` to prepend to each query. Bare `--base-iri` uses `http://example.org/`. | +| `-p`, `--rpqmatrix-optimizer ` | RPQMatrix optimizer: `none`, `join`, `metaac`, `mnc`, or `hybrid`. | + `query` adds `-o, --output ` to write JSON. diff --git a/pathrex-sys/build.rs b/pathrex-sys/build.rs index 7b7ce96..c99354a 100644 --- a/pathrex-sys/build.rs +++ b/pathrex-sys/build.rs @@ -52,6 +52,7 @@ fn main() { let lagraph_src = manifest_dir.join(LAGRAPH_REL_PATH); assert_lagraph_submodule_present(&lagraph_src); + watch_lagraph_sources(&lagraph_src); let graphblas_src = fetch_graphblas(&out_dir); let graphblas_install = build_graphblas_static(&graphblas_src); @@ -77,6 +78,44 @@ fn assert_lagraph_submodule_present(lagraph_src: &Path) { } } +fn watch_lagraph_sources(lagraph_src: &Path) { + println!( + "cargo:rerun-if-changed={}", + lagraph_src.join("CMakeLists.txt").display() + ); + + for rel_dir in ["Config", "cmake_modules", "include", "src", "experimental"] { + let dir = lagraph_src.join(rel_dir); + if dir.exists() { + watch_files_with_extensions(&dir, &["c", "h", "cmake", "in", "txt"]); + } + } +} + +fn watch_files_with_extensions(dir: &Path, extensions: &[&str]) { + let entries = std::fs::read_dir(dir) + .unwrap_or_else(|e| panic!("failed to read directory {}: {e}", dir.display())); + + for entry in entries { + let path = entry + .unwrap_or_else(|e| panic!("failed to read entry in {}: {e}", dir.display())) + .path(); + + if path.is_dir() { + watch_files_with_extensions(&path, extensions); + continue; + } + + let Some(ext) = path.extension().and_then(|ext| ext.to_str()) else { + continue; + }; + + if extensions.contains(&ext) { + println!("cargo:rerun-if-changed={}", path.display()); + } + } +} + /// Clone SuiteSparse:GraphBLAS at [`GRAPHBLAS_TAG`] into /// `$OUT_DIR/graphblas-src/`. Returns the path to the source tree. /// @@ -292,6 +331,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") @@ -304,6 +344,7 @@ fn regenerate_bindings(graphblas_install: &Path) { .allowlist_function("GrB_vxm") .allowlist_item("LAGRAPH_MSG_LEN") .allowlist_item("RPQMatrixOp") + .allowlist_item("RPQMatrixStorage") .allowlist_type("RPQMatrixPlan") .allowlist_type("LAGraph_Graph") .allowlist_type("LAGraph_Kind") @@ -317,7 +358,13 @@ fn regenerate_bindings(graphblas_install: &Path) { .allowlist_function("LAGraph_Cached_AT") .allowlist_function("LAGraph_MMRead") .allowlist_function("LAGraph_RPQMatrix") + .allowlist_function("LAGraph_RPQMatrix_SetGlobalStorageOrientation") + .allowlist_function("LAGraph_RPQMatrix_SetStorageOrientation") + .allowlist_function("LAGraph_RPQMatrix_DupWithStorageOrientation") .allowlist_function("LAGraph_RPQMatrix_reduce") + .allowlist_function("LAGraph_RPQMatrix_reduce_count_vector") + .allowlist_function("LAGraph_RPQMatrix_extended_count_vectors") + .allowlist_function("LAGraph_RPQMatrix_count_vector_.*") .allowlist_function("LAGraph_DestroyRpqMatrixPlan") .allowlist_function("LAGraph_RPQMatrix_label") .allowlist_function("LAGraph_RPQMatrix_Free") diff --git a/pathrex-sys/deps/LAGraph b/pathrex-sys/deps/LAGraph index bc00497..4229d83 160000 --- a/pathrex-sys/deps/LAGraph +++ b/pathrex-sys/deps/LAGraph @@ -1 +1 @@ -Subproject commit bc004979161db08389f52e2eff1e84e5cac42a64 +Subproject commit 4229d83bccb8fbaaeb5f17d7a3e5dd2bcdf1248a diff --git a/pathrex-sys/src/lagraph_sys_generated.rs b/pathrex-sys/src/lagraph_sys_generated.rs index 1a9188f..8252dc1 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; } @@ -320,6 +323,28 @@ pub struct RPQMatrixPlan { pub mat: GrB_Matrix, pub res_mat: GrB_Matrix, } +#[repr(u32)] +#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)] +pub enum RPQMatrixStorage { + RPQ_MATRIX_STORAGE_CSC = 1, + RPQ_MATRIX_STORAGE_CSR = 2, +} +unsafe extern "C" { + pub fn LAGraph_RPQMatrix_SetGlobalStorageOrientation(storage: RPQMatrixStorage) -> GrB_Info; +} +unsafe extern "C" { + pub fn LAGraph_RPQMatrix_SetStorageOrientation( + mat: GrB_Matrix, + storage: RPQMatrixStorage, + ) -> GrB_Info; +} +unsafe extern "C" { + pub fn LAGraph_RPQMatrix_DupWithStorageOrientation( + dst: *mut GrB_Matrix, + src: GrB_Matrix, + storage: RPQMatrixStorage, + ) -> GrB_Info; +} unsafe extern "C" { pub fn LAGraph_RPQMatrix( nnz: *mut GrB_Index, @@ -348,3 +373,57 @@ unsafe extern "C" { reduce_type: u8, ) -> GrB_Info; } +unsafe extern "C" { + pub fn LAGraph_RPQMatrix_reduce_count_vector( + res: *mut GrB_Vector, + mat: GrB_Matrix, + reduce_type: u8, + ) -> GrB_Info; +} +unsafe extern "C" { + pub fn LAGraph_RPQMatrix_extended_count_vectors( + row_extended: *mut GrB_Vector, + col_extended: *mut GrB_Vector, + mat: GrB_Matrix, + row_counts: GrB_Vector, + col_counts: GrB_Vector, + ) -> GrB_Info; +} +unsafe extern "C" { + pub fn LAGraph_RPQMatrix_count_vector_dot( + res: *mut f64, + lhs: GrB_Vector, + rhs: GrB_Vector, + ) -> GrB_Info; +} +unsafe extern "C" { + pub fn LAGraph_RPQMatrix_count_vector_mnc_matmul_nnz( + res: *mut f64, + lhs_rows: GrB_Vector, + lhs_cols: GrB_Vector, + rhs_rows: GrB_Vector, + rhs_cols: GrB_Vector, + lhs_col_extended: GrB_Vector, + rhs_row_extended: GrB_Vector, + ) -> GrB_Info; +} +unsafe extern "C" { + pub fn LAGraph_RPQMatrix_count_vector_sum(res: *mut f64, vector: GrB_Vector) -> GrB_Info; +} +unsafe extern "C" { + pub fn LAGraph_RPQMatrix_count_vector_scale( + res: *mut GrB_Vector, + vector: GrB_Vector, + scale: f64, + cap: f64, + ) -> GrB_Info; +} +unsafe extern "C" { + pub fn LAGraph_RPQMatrix_count_vector_mnc_add( + res: *mut GrB_Vector, + lhs: GrB_Vector, + rhs: GrB_Vector, + lambda: f64, + cap: f64, + ) -> GrB_Info; +} diff --git a/pathrex/Cargo.toml b/pathrex/Cargo.toml index 6401f1e..13b0a84 100644 --- a/pathrex/Cargo.toml +++ b/pathrex/Cargo.toml @@ -45,6 +45,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/bin/pathrex.rs b/pathrex/src/bin/pathrex.rs index 69b508a..5cc0b50 100644 --- a/pathrex/src/bin/pathrex.rs +++ b/pathrex/src/bin/pathrex.rs @@ -16,7 +16,9 @@ //! 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 \ +//! --bench-mode criterion \ +//! --rpqmatrix-optimizer join \ //! --output results.json //! ``` @@ -27,9 +29,12 @@ use chrono::Utc; use clap::Parser; use thiserror::Error; -use pathrex::cli::args::{BenchArgs, Cli, Commands, QueryArgs}; +use pathrex::cli::args::{ + 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}; @@ -55,6 +60,8 @@ enum MainError { #[source] source: std::io::Error, }, + #[error("invalid arguments: {0}")] + InvalidArgs(String), } fn main() { @@ -78,6 +85,75 @@ 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 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(), @@ -87,12 +163,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 +205,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(), @@ -147,48 +226,73 @@ 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)?; - 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, - 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( &common.graph, &common.queries, &common.algo, - path, + common.rpqmatrix_optimizer, + bench_config, + None, )) } } fn run_bench_cmd(args: BenchArgs) -> Result<(), MainError> { let 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!(); @@ -223,11 +327,18 @@ 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, - 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, }; @@ -238,10 +349,19 @@ fn run_bench_cmd(args: BenchArgs) -> Result<(), MainError> { path: args.output.clone(), source: e, })?; + let samples_path = output + .write_samples_to_file(Path::new(&args.output)) + .map_err(|e| MainError::Output { + path: args.output.clone(), + source: e, + })?; eprintln!(); eprintln!("=== Done ==="); eprintln!("Results written to: {}", args.output); + if let Some(path) = samples_path { + eprintln!("Run samples written to: {}", path.display()); + } if let Some(dir) = &args.criterion_dir { eprintln!("Criterion data in: {dir}") } diff --git a/pathrex/src/cli/args.rs b/pathrex/src/cli/args.rs index d95f478..4b7ea84 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( @@ -27,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), } @@ -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. @@ -85,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). @@ -106,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)] @@ -156,6 +235,55 @@ 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 join of source matrices. + Join, + /// MetaAC independence estimator. + #[value(name = "metaac")] + MetaAc, + /// Matrix nonzero count estimator. + Mnc, + /// Join work model with MetaAC result estimates. + Hybrid, + /// 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::Join => write!(f, "join"), + RpqMatrixOptimizer::MetaAc => write!(f, "metaac"), + RpqMatrixOptimizer::Mnc => write!(f, "mnc"), + RpqMatrixOptimizer::Hybrid => write!(f, "hybrid"), + RpqMatrixOptimizer::None => write!(f, "none"), + } + } +} + +impl From for OptimizationStrategy { + fn from(value: RpqMatrixOptimizer) -> Self { + match value { + RpqMatrixOptimizer::None => OptimizationStrategy::NoOpt, + RpqMatrixOptimizer::Join => OptimizationStrategy::Join, + RpqMatrixOptimizer::MetaAc => OptimizationStrategy::MetaAc, + RpqMatrixOptimizer::Mnc => OptimizationStrategy::Mnc, + RpqMatrixOptimizer::Hybrid => OptimizationStrategy::Hybrid, + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -174,4 +302,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/estimates.rs b/pathrex/src/cli/bench/estimates.rs index 1cf72cb..c6459b7 100644 --- a/pathrex/src/cli/bench/estimates.rs +++ b/pathrex/src/cli/bench/estimates.rs @@ -87,7 +87,11 @@ pub fn read_timing_stats( pub fn read_algo_timing(criterion_dir: &Path, group: &str) -> Result { let total = read_timing_stats(criterion_dir, group, "eval_total")?; let ffi_only = read_timing_stats(criterion_dir, group, "eval_ffi_only")?; - Ok(AlgoTiming { total, ffi_only }) + Ok(AlgoTiming { + total, + ffi_only, + samples: None, + }) } #[cfg(test)] diff --git a/pathrex/src/cli/bench/runner.rs b/pathrex/src/cli/bench/runner.rs index 4b5a54a..815e90d 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, AlgoTimingSamples, 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() @@ -65,10 +65,13 @@ fn run_benchmark_group( query_index: usize, ) -> Result, RpqError> where - E: Evaluator + Copy, + E: Evaluator, 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,88 @@ 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, + 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 result = black_box(prepared.execute()?); + ffi_samples.push(elapsed_ns(start)); + drop(result); + } + + Ok(( + result_count.unwrap_or(0), + AlgoTiming { + total: timing_stats(&total_samples), + ffi_only: timing_stats(&ffi_samples), + samples: Some(AlgoTimingSamples { + total_ns: total_samples, + ffi_only_ns: ffi_samples, + }), + }, + )) +} + /// Run the bench loop for every query in `queries` for one evaluator. pub fn run_bench_for_evaluator( args: &BenchArgs, @@ -111,7 +204,7 @@ pub fn run_bench_for_evaluator( checkpointer: &mut Checkpointer, ) -> Result, BenchError> where - E: Evaluator + Copy, + E: Evaluator + Clone, E::Result: ResultCount, { let mut results = Vec::with_capacity(queries.len()); @@ -150,9 +243,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.clone(), query, graph) + .map(|(count, timing)| Ok((Some(count), timing))), + BenchMode::Criterion => { + run_benchmark_group(args, algo_name, evaluator.clone(), 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 f59f3bb..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; +use super::args::{Algo, BenchArgs, BenchMode, RpqMatrixOptimizer}; /// Persistent checkpoint state written to disk as JSON. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -25,9 +25,59 @@ pub struct Checkpoint { pub graph_path: String, pub queries_file: String, 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 { @@ -37,12 +87,20 @@ 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, + bench_config: BenchRunConfig, + ) -> Self { Self { version: 1, graph_path: graph_path.to_string(), queries_file: queries_file.to_string(), algorithms: algorithms.to_vec(), + rpqmatrix_optimizer, + bench_config, completed: Vec::new(), } } @@ -65,6 +123,8 @@ impl Checkpoint { graph_path: &str, queries_file: &str, algorithms: &[Algo], + rpqmatrix_optimizer: RpqMatrixOptimizer, + bench_config: &BenchRunConfig, ) -> Result<(), CheckpointError> { if self.graph_path != graph_path { return Err(CheckpointError::Mismatch(format!( @@ -86,6 +146,18 @@ 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 + ))); + } + if &self.bench_config != bench_config { + return Err(CheckpointError::Mismatch(format!( + "bench_config: checkpoint has {:?}, current is {:?}", + self.bench_config, bench_config + ))); + } Ok(()) } @@ -93,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) @@ -140,21 +220,37 @@ impl Checkpoint { /// Runtime owner for a [`Checkpoint`] paired with its on-disk path. pub struct Checkpointer { inner: Checkpoint, - path: PathBuf, + path: Option, } 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, + bench_config: BenchRunConfig, + path: Option, + ) -> Self { Self { - inner: Checkpoint::new(graph_path, queries_file, algorithms), + 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. @@ -184,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/dispatch.rs b/pathrex/src/cli/dispatch.rs index 2701460..07f57a5 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; @@ -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, 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, + 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..de30e86 100644 --- a/pathrex/src/cli/output.rs +++ b/pathrex/src/cli/output.rs @@ -2,7 +2,7 @@ use std::collections::HashMap; use std::fs; -use std::path::Path; +use std::path::{Path, PathBuf}; use serde::Serialize; @@ -60,6 +60,14 @@ impl AlgoResult { pub struct AlgoTiming { pub total: TimingStats, pub ffi_only: TimingStats, + #[serde(skip)] + pub samples: Option, +} + +#[derive(Debug, Serialize)] +pub struct AlgoTimingSamples { + pub total_ns: Vec, + pub ffi_only_ns: Vec, } /// Timing statistics extracted from criterion estimates. @@ -91,6 +99,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, @@ -100,7 +109,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) } } @@ -116,22 +125,126 @@ 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, 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) + } + + pub fn write_samples_to_file(&self, path: &Path) -> Result, std::io::Error> { + let Some(samples) = BenchSamplesOutput::from_bench_output(self) else { + return Ok(None); + }; + let samples_path = samples_path_for(path); + let json = serde_json::to_string_pretty(&samples).map_err(std::io::Error::other)?; + write_json_to_file(&samples_path, json)?; + Ok(Some(samples_path)) } } +#[derive(Debug, Serialize)] +pub struct BenchSamplesOutput<'a> { + pub metadata: &'a BenchMetadata, + pub results: Vec>, +} + +#[derive(Debug, Serialize)] +pub struct QuerySamples<'a> { + pub query_index: usize, + pub query_id: &'a str, + pub query_text: &'a str, + pub algorithms: HashMap<&'a str, AlgoSamples<'a>>, +} + +#[derive(Debug, Serialize)] +pub struct AlgoSamples<'a> { + #[serde(skip_serializing_if = "Option::is_none")] + pub result_count: Option, + pub total_ns: &'a [f64], + pub ffi_only_ns: &'a [f64], +} + +impl<'a> BenchSamplesOutput<'a> { + pub fn from_bench_output(output: &'a BenchOutput) -> Option { + let mut results = Vec::new(); + + for query in &output.results { + let mut algorithms = HashMap::new(); + for (algo, result) in &query.algorithms { + let Some(timing) = &result.timing else { + continue; + }; + let Some(samples) = &timing.samples else { + continue; + }; + algorithms.insert( + algo.as_str(), + AlgoSamples { + result_count: result.result_count, + total_ns: &samples.total_ns, + ffi_only_ns: &samples.ffi_only_ns, + }, + ); + } + + if !algorithms.is_empty() { + results.push(QuerySamples { + query_index: query.query_index, + query_id: &query.query_id, + query_text: &query.query_text, + algorithms, + }); + } + } + + (!results.is_empty()).then_some(Self { + metadata: &output.metadata, + results, + }) + } +} + +fn samples_path_for(path: &Path) -> PathBuf { + let file_name = path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("bench_results.json"); + let samples_name = file_name + .strip_suffix(".json") + .map(|stem| format!("{stem}.runs.json")) + .unwrap_or_else(|| format!("{file_name}.runs.json")); + + path.with_file_name(samples_name) +} + +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::*; @@ -153,6 +266,7 @@ mod tests { stddev_ns: 0.0, iterations: 10, }, + samples: None, }), ); @@ -176,4 +290,96 @@ 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()); + } + + #[test] + fn bench_output_writes_samples_next_to_output() { + let dir = tempfile::tempdir().expect("tempdir"); + let output_path = dir.path().join("bench.json"); + let output = BenchOutput { + metadata: BenchMetadata { + 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, + bench_mode: "fixed".into(), + runs: Some(2), + warm_up_runs: Some(0), + sample_size: None, + warm_up_secs: None, + measurement_secs: None, + }, + results: vec![QueryResult { + query_index: 0, + query_id: "q0".into(), + query_text: "query".into(), + algorithms: HashMap::from([( + "rpqmatrix".into(), + AlgoResult::ok( + Some(1), + Some(AlgoTiming { + total: TimingStats { + mean_ns: 15.0, + median_ns: 15.0, + stddev_ns: 5.0, + iterations: 2, + }, + ffi_only: TimingStats { + mean_ns: 4.0, + median_ns: 4.0, + stddev_ns: 1.0, + iterations: 2, + }, + samples: Some(AlgoTimingSamples { + total_ns: vec![10.0, 20.0], + ffi_only_ns: vec![3.0, 5.0], + }), + }), + ), + )]), + }], + }; + + let samples_path = output + .write_samples_to_file(&output_path) + .expect("write samples") + .expect("samples path"); + let samples_json = fs::read_to_string(samples_path).expect("read samples"); + let value: serde_json::Value = serde_json::from_str(&samples_json).expect("json"); + + assert_eq!( + value["results"][0]["algorithms"]["rpqmatrix"]["total_ns"][0], + 10.0 + ); + assert_eq!( + value["results"][0]["algorithms"]["rpqmatrix"]["ffi_only_ns"][1], + 5.0 + ); + } } diff --git a/pathrex/src/cli/query.rs b/pathrex/src/cli/query.rs index 152883c..5bbad9a 100644 --- a/pathrex/src/cli/query.rs +++ b/pathrex/src/cli/query.rs @@ -20,7 +20,7 @@ pub fn run_query_for_evaluator( queries: &[LoadedQuery], ) -> Vec where - E: Evaluator + Copy, + E: Evaluator, E::Result: ResultCount, { let mut results = Vec::with_capacity(queries.len()); diff --git a/pathrex/src/graph/inmemory.rs b/pathrex/src/graph/inmemory.rs index 9121101..746da6b 100644 --- a/pathrex/src/graph/inmemory.rs +++ b/pathrex/src/graph/inmemory.rs @@ -11,8 +11,8 @@ use crate::{ }; use super::{ - Backend, Edge, GraphBuilder, GraphDecomposition, GraphError, LagraphGraph, ThreadScope, - compute_outer_inner, load_mm_file, + Backend, Edge, GraphBuilder, GraphDecomposition, GraphError, LagraphGraph, MatrixStorage, + ThreadScope, compute_outer_inner, load_mm_file, }; /// Marker type for the in-memory GraphBLAS-backed backend. @@ -42,6 +42,8 @@ pub struct InMemoryBuilder { next_id: usize, label_buffers: HashMap>, prebuilt: HashMap, + prebuilt_csc: HashMap, + metadata: HashMap, } impl InMemoryBuilder { @@ -52,6 +54,8 @@ impl InMemoryBuilder { next_id: 0, label_buffers: HashMap::new(), prebuilt: HashMap::new(), + prebuilt_csc: HashMap::new(), + metadata: HashMap::new(), } } @@ -111,6 +115,22 @@ impl InMemoryBuilder { ) { self.prebuilt.extend(iter); } + + /// Bulk-install pre-wrapped CSC `(label, LagraphGraph)` pairs. + pub(crate) fn extend_prebuilt_csc>( + &mut self, + iter: I, + ) { + self.prebuilt_csc.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 { @@ -128,10 +148,15 @@ impl GraphBuilder for InMemoryBuilder { let mut graphs: HashMap> = HashMap::with_capacity(self.label_buffers.len() + self.prebuilt.len()); + let mut graphs_csc: HashMap> = + HashMap::with_capacity(self.prebuilt_csc.len()); for (label, lg) in self.prebuilt { graphs.insert(label, Arc::new(lg)); } + for (label, lg) in self.prebuilt_csc { + graphs_csc.insert(label, Arc::new(lg)); + } let label_buffers: Vec<(String, Vec<(usize, usize)>)> = self.label_buffers.into_iter().collect(); @@ -162,11 +187,14 @@ 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, + graphs_csc, + metadata: GraphMetadata { + label_to_data: self.metadata, + }, }) } } @@ -176,6 +204,24 @@ pub struct InMemoryGraph { node_to_id: HashMap, id_to_node: HashMap, graphs: HashMap>, + graphs_csc: 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 { + pub dimension: usize, + pub nonzero_rows: usize, + pub nonzero_cols: usize, + pub nvals: usize, } impl GraphDecomposition for InMemoryGraph { @@ -186,6 +232,19 @@ impl GraphDecomposition for InMemoryGraph { .ok_or_else(|| GraphError::LabelNotFound(label.to_owned())) } + fn get_graph_with_storage( + &self, + label: &str, + storage: MatrixStorage, + ) -> Result, GraphError> { + if storage == MatrixStorage::Csc { + if let Some(graph) = self.graphs_csc.get(label) { + return Ok(Arc::clone(graph)); + } + } + self.get_graph(label) + } + fn get_node_id(&self, string_id: &str) -> Option { self.node_to_id.get(string_id).copied() } @@ -197,6 +256,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 { @@ -204,6 +267,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 { @@ -244,22 +310,52 @@ impl GraphSource for MatrixMarket { let _scope = ThreadScope::enter(outer, inner)?; let mm_dir = self.dir.clone(); - let loaded: Vec<(String, LagraphGraph)> = edge_by_idx - .into_par_iter() - .map( - |(idx, label)| -> Result<(String, LagraphGraph), 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)) - }, - ) - .collect::, GraphError>>()?; - - builder.extend_prebuilt(loaded); + let loaded: Vec<(String, LagraphGraph, LagraphGraph, MatrixMetadata)> = + edge_by_idx + .into_par_iter() + .map( + |(idx, label)| -> Result< + (String, LagraphGraph, LagraphGraph, MatrixMetadata), + GraphError, + > { + let path = mm_dir.join(format!("{}.txt", idx)); + let matrix = load_mm_file(&path)?; + matrix.set_storage_orientation(MatrixStorage::Csr)?; + let csc_matrix = matrix.dup_with_storage_orientation(MatrixStorage::Csc)?; + let lg = LagraphGraph::from_matrix( + matrix, + LAGraph_Kind::LAGraph_ADJACENCY_DIRECTED, + )?; + let lg_csc = LagraphGraph::from_matrix( + csc_matrix, + LAGraph_Kind::LAGraph_ADJACENCY_DIRECTED, + )?; + 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, lg_csc, metadata)) + }, + ) + .collect::, GraphError>>()?; + + let mut loaded_graphs = vec![]; + let mut loaded_graphs_csc = vec![]; + let mut loaded_metadata = vec![]; + for (_i, (name, graph, graph_csc, meta)) in loaded.into_iter().enumerate() { + loaded_graphs.push((name.clone(), graph)); + loaded_graphs_csc.push((name.clone(), graph_csc)); + loaded_metadata.push((name, meta)); + } + builder.extend_prebuilt(loaded_graphs); + builder.extend_prebuilt_csc(loaded_graphs_csc); + builder.extend_metadata(loaded_metadata); Ok(builder) } diff --git a/pathrex/src/graph/mod.rs b/pathrex/src/graph/mod.rs index 0922459..3cc7574 100644 --- a/pathrex/src/graph/mod.rs +++ b/pathrex/src/graph/mod.rs @@ -4,12 +4,15 @@ pub mod inmemory; pub mod wrappers; pub use inmemory::{InMemory, InMemoryBuilder, InMemoryGraph}; -pub use wrappers::{GraphblasMatrix, GraphblasVector, LagraphGraph, load_mm_file}; -pub(crate) use wrappers::{ThreadScope, compute_outer_inner, ensure_grb_init}; +pub use wrappers::{GraphblasMatrix, GraphblasVector, LagraphGraph, MatrixStorage, load_mm_file}; +pub(crate) use wrappers::{ + ThreadScope, compute_outer_inner, ensure_grb_init, set_global_matrix_storage_hint, +}; use std::marker::PhantomData; use std::sync::Arc; +use crate::graph::inmemory::GraphMetadata; use crate::lagraph_sys::GrB_Info; use thiserror::Error; @@ -77,12 +80,25 @@ pub trait GraphDecomposition { /// Returns the [`LagraphGraph`] for `label`. fn get_graph(&self, label: &str) -> Result, GraphError>; + /// Returns the [`LagraphGraph`] for `label` in a preferred storage orientation. + /// Backends that do not maintain multiple orientations may return their default graph. + fn get_graph_with_storage( + &self, + label: &str, + _storage: MatrixStorage, + ) -> Result, GraphError> { + self.get_graph(label) + } + /// Translates a string ID to a contiguous matrix index. fn get_node_id(&self, string_id: &str) -> Option; /// 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 e97cfc5..959bbb7 100644 --- a/pathrex/src/graph/wrappers.rs +++ b/pathrex/src/graph/wrappers.rs @@ -9,13 +9,43 @@ use std::ffi::CString; use std::fs::File; use std::os::fd::IntoRawFd; use std::path::Path; -use std::sync::Once; +use std::sync::{ + Once, + atomic::{AtomicU8, Ordering}, +}; -use crate::{grb_ok, la_ok, lagraph_sys::*}; +use crate::{ + graph::wrappers::ReduceType::{ByCols, ByRows}, + grb_ok, la_ok, + lagraph_sys::*, +}; use super::GraphError; static GRB_INIT: Once = Once::new(); +static GLOBAL_MATRIX_STORAGE_HINT: AtomicU8 = AtomicU8::new(0); + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum MatrixStorage { + Csc, + Csr, +} + +impl MatrixStorage { + fn as_hint_code(self) -> u8 { + match self { + MatrixStorage::Csc => 1, + MatrixStorage::Csr => 2, + } + } + + fn as_rpq_storage(self) -> RPQMatrixStorage { + match self { + MatrixStorage::Csc => RPQMatrixStorage::RPQ_MATRIX_STORAGE_CSC, + MatrixStorage::Csr => RPQMatrixStorage::RPQ_MATRIX_STORAGE_CSR, + } + } +} pub(crate) fn ensure_grb_init() -> Result<(), GraphError> { let mut result = Ok(()); @@ -25,6 +55,22 @@ pub(crate) fn ensure_grb_init() -> Result<(), GraphError> { result } +pub(crate) fn set_global_matrix_storage_hint(storage: MatrixStorage) -> Result<(), GraphError> { + let hint_code = storage.as_hint_code(); + if GLOBAL_MATRIX_STORAGE_HINT.load(Ordering::Acquire) == hint_code { + return Ok(()); + } + + ensure_grb_init()?; + unsafe { + grb_ok!(LAGraph_RPQMatrix_SetGlobalStorageOrientation( + storage.as_rpq_storage(), + ))? + }; + GLOBAL_MATRIX_STORAGE_HINT.store(hint_code, Ordering::Release); + Ok(()) +} + /// Compute a balanced `(outer, inner)` split for LAGraph's two-level threading. /// /// `outer` is the number of user-level concurrent tasks (rayon workers); @@ -69,12 +115,20 @@ impl Drop for ThreadScope { } } +pub enum ReduceType { + ByRows, + ByCols, +} #[derive(Debug)] pub struct LagraphGraph { pub(crate) inner: LAGraph_Graph, } impl LagraphGraph { + pub(crate) fn matrix(&self) -> GrB_Matrix { + unsafe { (*self.inner).A } + } + /// Build a `LagraphGraph` from an RAII-wrapped [`GraphblasMatrix`]. /// /// On success, ownership of the underlying `GrB_Matrix` is transferred @@ -147,6 +201,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 +222,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 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 mut res: GrB_Index = 0; + unsafe { LAGraph_RPQMatrix_reduce(&mut res, matrix, ByCols as u8) }; + Ok(res as usize) + } } impl Drop for LagraphGraph { @@ -239,6 +318,28 @@ impl GraphblasMatrix { pub fn from_raw(raw: GrB_Matrix) -> Self { Self { inner: raw } } + + pub fn set_storage_orientation(&self, storage: MatrixStorage) -> Result<(), GraphError> { + unsafe { + grb_ok!(LAGraph_RPQMatrix_SetStorageOrientation( + self.inner, + storage.as_rpq_storage(), + ))? + }; + Ok(()) + } + + pub fn dup_with_storage_orientation(&self, storage: MatrixStorage) -> Result { + let mut raw: GrB_Matrix = std::ptr::null_mut(); + unsafe { + grb_ok!(LAGraph_RPQMatrix_DupWithStorageOrientation( + &mut raw, + self.inner, + storage.as_rpq_storage(), + ))? + }; + Ok(Self { inner: raw }) + } } impl Drop for GraphblasMatrix { diff --git a/pathrex/src/rpq/rpqmatrix/cost.rs b/pathrex/src/rpq/rpqmatrix/cost.rs new file mode 100644 index 0000000..63d3329 --- /dev/null +++ b/pathrex/src/rpq/rpqmatrix/cost.rs @@ -0,0 +1,1238 @@ +use std::{cmp::Ordering, collections::HashMap}; + +use egg::{CostFunction, Id}; + +use super::{ + plan::RpqPlan, + stats::{CountVector, LabelCountVectors}, +}; + + + +#[derive(Clone, Debug, PartialEq)] +pub struct JoinCost { + pub score: f64, + pub nnz: f64, + pub nnz_r: f64, + pub nnz_c: f64, +} + +impl Eq for JoinCost {} + +impl PartialOrd for JoinCost { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for JoinCost { + 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) + } +} + +// Approach based on formula for SQL JOIN operation +// Got from https://github.com/chernishev/Database-Engines-Course/tree/master/Lecture%203 +pub struct JoinCostFn { + pub n: f64, + pub star_penalty: f64, + pub lr_multiplier: f64, +} + +// 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 JoinCostFn { + type Cost = JoinCost; + + fn cost(&mut self, enode: &RpqPlan, mut costs: C) -> Self::Cost + where + C: FnMut(Id) -> Self::Cost, + { + match enode { + RpqPlan::NamedVertex(_name) => JoinCost { + score: 0.0, + nnz: 1 as f64, + nnz_r: 1 as f64, + nnz_c: 1 as f64, + }, + + RpqPlan::Label(meta) => JoinCost { + 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; + + // TODO: this can exceed `n^2` when child estimates are already loose. + let nnz_est = ca.nnz * cb.nnz / (self.n * self.n); + + JoinCost { + 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); + + // 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; + + 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); + + JoinCost { + score, + nnz: nnz_est, + nnz_r: nnz_r_est, + nnz_c: nnz_c_est, + } + } + + 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; + + JoinCost { + 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); + + // 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; + let score = ca.score + cb.score + op_cost; + + let nnz_est = self.lr_multiplier * ca.nnz * cb.nnz / (self.n * self.n); + + JoinCost { + 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); + + // 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; + let score = ca.score + cb.score + op_cost; + + let nnz_est = self.lr_multiplier * ca.nnz * cb.nnz / (self.n * self.n); + + JoinCost { + 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 + } + } + } + } +} + +// 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() +// } +// } + +#[derive(Clone, Debug)] +pub(super) struct MetaAcCost { + score: f64, + nnz: f64, + nnz_r: f64, + nnz_c: f64, +} + +impl PartialEq for MetaAcCost { + fn eq(&self, other: &Self) -> bool { + self.cmp(other).is_eq() + } +} +impl Eq for MetaAcCost {} +impl PartialOrd for MetaAcCost { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} +impl Ord for MetaAcCost { + fn cmp(&self, other: &Self) -> Ordering { + self.score + .total_cmp(&other.score) + .then(self.nnz.total_cmp(&other.nnz)) + .then(self.nnz_r.total_cmp(&other.nnz_r)) + .then(self.nnz_c.total_cmp(&other.nnz_c)) + } +} + +#[derive(Clone, Debug)] +pub(super) struct HybridCost { + score: f64, + nnz: f64, + nnz_r: f64, + nnz_c: f64, +} + +impl PartialEq for HybridCost { + fn eq(&self, other: &Self) -> bool { + self.cmp(other).is_eq() + } +} +impl Eq for HybridCost {} +impl PartialOrd for HybridCost { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} +impl Ord for HybridCost { + fn cmp(&self, other: &Self) -> Ordering { + self.score + .total_cmp(&other.score) + .then(self.nnz.total_cmp(&other.nnz)) + .then(self.nnz_r.total_cmp(&other.nnz_r)) + .then(self.nnz_c.total_cmp(&other.nnz_c)) + } +} + +// Naive metadata estimator +// got from 2.1 of paper https://mboehm7.github.io/resources/sigmod2019.pdf +pub(super) struct MetaAcCostFn { + pub n: f64, + pub star_penalty: f64, + pub lr_multiplier: f64, +} + +fn metaac_matmul_nnz(lhs_nnz: f64, rhs_nnz: f64, n: f64) -> f64 { + let output_cells = (n * n).max(1.0); + let p = ((lhs_nnz / output_cells) * (rhs_nnz / output_cells)).clamp(0.0, 1.0); + (output_cells * (1.0 - (1.0 - p).powf(n))).clamp(0.0, output_cells) +} + +// Due to Join cost fun estimates only join operation +// and don't even estimate nnz of result matrices, I combine +// it with metaac +pub(super) struct HybridCostFn { + pub n: f64, + pub star_penalty: f64, + pub lr_multiplier: f64, +} + +impl CostFunction for MetaAcCostFn { + type Cost = MetaAcCost; + + fn cost(&mut self, enode: &RpqPlan, mut costs: C) -> Self::Cost + where + C: FnMut(Id) -> Self::Cost, + { + match enode { + RpqPlan::NamedVertex(_name) => MetaAcCost { + score: 0.0, + nnz: 1.0, + nnz_r: 1.0, + nnz_c: 1.0, + }, + RpqPlan::Label(meta) => MetaAcCost { + 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 op_cost = ca.nnz * cb.nnz / self.n.max(1.0); + let score = ca.score + cb.score + op_cost; + let nnz_est = metaac_matmul_nnz(ca.nnz, cb.nnz, self.n); + MetaAcCost { + score, + nnz: nnz_est, + nnz_r: ca.nnz_r.min(self.n), + nnz_c: cb.nnz_c.min(self.n), + } + } + RpqPlan::Alt([a, b]) => { + let ca = costs(*a); + let cb = costs(*b); + let overlap = ca.nnz * cb.nnz / (self.n * self.n).max(1.0); + let nnz_est = (ca.nnz + cb.nnz - overlap).clamp(0.0, self.n * self.n); + let score = ca.score + cb.score + nnz_est; + let nnz_r_est = (ca.nnz_r + cb.nnz_r - ca.nnz_r * cb.nnz_r / self.n.max(1.0)) + .clamp(0.0, self.n); + let nnz_c_est = (ca.nnz_c + cb.nnz_c - ca.nnz_c * cb.nnz_c / self.n.max(1.0)) + .clamp(0.0, self.n); + MetaAcCost { + 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); + MetaAcCost { + score: ca.score + penalty, + 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 op_cost = self.lr_multiplier * ca.nnz * cb.nnz / denom; + let score = ca.score + cb.score + op_cost; + let nnz_est = self.lr_multiplier * ca.nnz * cb.nnz / (self.n * self.n).max(1.0); + MetaAcCost { + score, + nnz: nnz_est, + nnz_r: ca.nnz_r.min(self.n), + nnz_c: cb.nnz_c.min(self.n), + } + } + 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 op_cost = self.lr_multiplier * ca.nnz * cb.nnz / denom; + let score = ca.score + cb.score + op_cost; + let nnz_est = self.lr_multiplier * ca.nnz * cb.nnz / (self.n * self.n).max(1.0); + MetaAcCost { + score, + nnz: nnz_est, + nnz_r: ca.nnz_r.min(self.n), + nnz_c: cb.nnz_c.min(self.n), + } + } + } + } +} + +impl CostFunction for HybridCostFn { + type Cost = HybridCost; + + fn cost(&mut self, enode: &RpqPlan, mut costs: C) -> Self::Cost + where + C: FnMut(Id) -> Self::Cost, + { + match enode { + RpqPlan::NamedVertex(_name) => HybridCost { + score: 0.0, + nnz: 1.0, + nnz_r: 1.0, + nnz_c: 1.0, + }, + RpqPlan::Label(meta) => HybridCost { + 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 universe = (self.n * self.n).max(1.0); + let p = ((ca.nnz / universe) * (cb.nnz / universe)).clamp(0.0, 1.0); + let nnz_est = (universe * (1.0 - (1.0 - p).powf(self.n))).clamp(0.0, universe); + HybridCost { + score, + nnz: nnz_est, + nnz_r: ca.nnz_r.min(self.n), + nnz_c: cb.nnz_c.min(self.n), + } + } + RpqPlan::Alt([a, b]) => { + let ca = costs(*a); + let cb = costs(*b); + let overlap = ca.nnz * cb.nnz / (self.n * self.n).max(1.0); + let nnz_est = (ca.nnz + cb.nnz - overlap).clamp(0.0, self.n * self.n); + let score = ca.score + cb.score + nnz_est; + let nnz_r_est = (ca.nnz_r + cb.nnz_r - ca.nnz_r * cb.nnz_r / self.n.max(1.0)) + .clamp(0.0, self.n); + let nnz_c_est = (ca.nnz_c + cb.nnz_c - ca.nnz_c * cb.nnz_c / self.n.max(1.0)) + .clamp(0.0, self.n); + HybridCost { + 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); + HybridCost { + score: ca.score + penalty, + 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 op_cost = self.lr_multiplier * ca.nnz * cb.nnz / denom; + let score = ca.score + cb.score + op_cost; + let nnz_est = self.lr_multiplier * ca.nnz * cb.nnz / (self.n * self.n).max(1.0); + HybridCost { + score, + nnz: nnz_est, + nnz_r: ca.nnz_r.min(self.n), + nnz_c: cb.nnz_c.min(self.n), + } + } + 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 op_cost = self.lr_multiplier * ca.nnz * cb.nnz / denom; + let score = ca.score + cb.score + op_cost; + let nnz_est = self.lr_multiplier * ca.nnz * cb.nnz / (self.n * self.n).max(1.0); + HybridCost { + score, + nnz: nnz_est, + nnz_r: ca.nnz_r.min(self.n), + nnz_c: cb.nnz_c.min(self.n), + } + } + } + } +} + +#[derive(Clone, Debug)] +pub(super) struct MncCost { + score: f64, + nnz: f64, + nnz_r: f64, + nnz_c: f64, + row_counts: Option, + col_counts: Option, + row_extended: Option, + col_extended: Option, +} + +impl PartialEq for MncCost { + fn eq(&self, other: &Self) -> bool { + self.cmp(other).is_eq() + } +} +impl Eq for MncCost {} +impl PartialOrd for MncCost { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} +impl Ord for MncCost { + fn cmp(&self, other: &Self) -> Ordering { + self.score + .total_cmp(&other.score) + .then(self.nnz.total_cmp(&other.nnz)) + .then(self.nnz_r.total_cmp(&other.nnz_r)) + .then(self.nnz_c.total_cmp(&other.nnz_c)) + } +} + +pub(super) struct MncCostFn { + n: f64, + star_penalty: f64, + lr_multiplier: f64, + labels: HashMap, + vertices: HashMap, + dot_cache: HashMap<(usize, usize), f64>, + matmul_cache: HashMap<(usize, usize, usize, usize, Option, Option), f64>, + scale_cache: HashMap<(usize, u64, u64), CountVector>, + add_cache: HashMap<(usize, usize, u64, u64), CountVector>, +} + +impl MncCostFn { + pub(super) fn new( + n: f64, + labels: HashMap, + vertices: HashMap, + ) -> Self { + Self { + n, + star_penalty: 50.0, + lr_multiplier: 5.0, + labels, + vertices, + dot_cache: HashMap::new(), + matmul_cache: HashMap::new(), + scale_cache: HashMap::new(), + add_cache: HashMap::new(), + } + } + + fn dot(&mut self, a: &CountVector, b: &CountVector) -> Option { + let key = (a.cache_key(), b.cache_key()); + if let Some(value) = self.dot_cache.get(&key) { + return Some(*value); + } + let value = a.dot(b)?; + self.dot_cache.insert(key, value); + Some(value) + } + + fn matmul_nnz(&mut self, a: &MncCost, b: &MncCost) -> Option { + let (ar, ac, br, bc) = ( + a.row_counts.as_ref()?, + a.col_counts.as_ref()?, + b.row_counts.as_ref()?, + b.col_counts.as_ref()?, + ); + let key = ( + ar.cache_key(), + ac.cache_key(), + br.cache_key(), + bc.cache_key(), + a.col_extended.as_ref().map(CountVector::cache_key), + b.row_extended.as_ref().map(CountVector::cache_key), + ); + if let Some(value) = self.matmul_cache.get(&key) { + return Some(*value); + } + let value = CountVector::mnc_matmul_nnz( + ar, + ac, + br, + bc, + a.col_extended.as_ref(), + b.row_extended.as_ref(), + )?; + self.matmul_cache.insert(key, value); + Some(value) + } + + fn scale(&mut self, vector: &CountVector, factor: f64) -> Option { + let key = (vector.cache_key(), factor.to_bits(), self.n.to_bits()); + if let Some(value) = self.scale_cache.get(&key) { + return Some(value.clone()); + } + let value = vector.scale(factor, self.n)?; + self.scale_cache.insert(key, value.clone()); + Some(value) + } + + fn add(&mut self, a: &CountVector, b: &CountVector, lambda: f64) -> Option { + let (x, y) = ( + a.cache_key().min(b.cache_key()), + a.cache_key().max(b.cache_key()), + ); + let key = (x, y, lambda.to_bits(), self.n.to_bits()); + if let Some(value) = self.add_cache.get(&key) { + return Some(value.clone()); + } + let value = a.mnc_add(b, lambda, self.n)?; + self.add_cache.insert(key, value.clone()); + Some(value) + } + + fn seq(&mut self, a: MncCost, b: MncCost) -> MncCost { + let op_cost = match (&a.col_counts, &b.row_counts) { + (Some(ac), Some(br)) => self.dot(ac, br), + _ => None, + } + .unwrap_or_else(|| a.nnz * b.nnz / self.n.max(1.0)); + let nnz = self + .matmul_nnz(&a, &b) + .unwrap_or_else(|| metaac_matmul_nnz(a.nnz, b.nnz, self.n)); + let row_counts = a + .row_counts + .as_ref() + .and_then(|v| self.scale(v, nnz / a.nnz.max(1.0))); + let col_counts = b + .col_counts + .as_ref() + .and_then(|v| self.scale(v, nnz / b.nnz.max(1.0))); + MncCost { + score: a.score + b.score + op_cost, + nnz, + nnz_r: row_counts + .as_ref() + .map_or(a.nnz_r.min(self.n), CountVector::nonzero_count), + nnz_c: col_counts + .as_ref() + .map_or(b.nnz_c.min(self.n), CountVector::nonzero_count), + row_counts, + col_counts, + row_extended: None, + col_extended: None, + } + } + + fn alt(&mut self, a: MncCost, b: MncCost) -> MncCost { + let overlap = a.nnz * b.nnz / (self.n * self.n).max(1.0); + let fallback_nnz = (a.nnz + b.nnz - overlap).clamp(0.0, self.n * self.n); + let fallback_nnz_r = + (a.nnz_r + b.nnz_r - a.nnz_r * b.nnz_r / self.n.max(1.0)).clamp(0.0, self.n); + let fallback_nnz_c = + (a.nnz_c + b.nnz_c - a.nnz_c * b.nnz_c / self.n.max(1.0)).clamp(0.0, self.n); + let denominator = (a.nnz * b.nnz).max(1.0); + let lambda_cols = match (&a.col_counts, &b.col_counts) { + (Some(ac), Some(bc)) => self.dot(ac, bc).unwrap_or(0.0) / denominator, + _ => 0.0, + } + .clamp(0.0, 1.0); + let lambda_rows = match (&a.row_counts, &b.row_counts) { + (Some(ar), Some(br)) => self.dot(ar, br).unwrap_or(0.0) / denominator, + _ => 0.0, + } + .clamp(0.0, 1.0); + let row_counts = match (&a.row_counts, &b.row_counts) { + (Some(ar), Some(br)) => self.add(ar, br, lambda_cols), + _ => None, + }; + let col_counts = match (&a.col_counts, &b.col_counts) { + (Some(ac), Some(bc)) => self.add(ac, bc, lambda_rows), + _ => None, + }; + let nnz = match (&row_counts, &col_counts) { + (Some(rows), Some(cols)) => (rows.sum() + cols.sum()) / 2.0, + (Some(rows), None) => rows.sum(), + (None, Some(cols)) => cols.sum(), + (None, None) => fallback_nnz, + } + .clamp(0.0, self.n * self.n); + MncCost { + score: a.score + b.score + nnz, + nnz, + nnz_r: row_counts + .as_ref() + .map_or(fallback_nnz_r, CountVector::nonzero_count), + nnz_c: col_counts + .as_ref() + .map_or(fallback_nnz_c, CountVector::nonzero_count), + row_counts, + col_counts, + row_extended: None, + col_extended: None, + } + } +} + +impl CostFunction for MncCostFn { + type Cost = MncCost; + + fn cost(&mut self, enode: &RpqPlan, mut costs: C) -> Self::Cost + where + C: FnMut(Id) -> Self::Cost, + { + match enode { + RpqPlan::NamedVertex(name) => { + let counts = self.vertices.get(name); + let row_counts = counts.map(|v| v.row_counts.clone()); + let col_counts = counts.map(|v| v.col_counts.clone()); + let row_extended = counts.map(|v| v.row_extended.clone()); + let col_extended = counts.map(|v| v.col_extended.clone()); + MncCost { + score: 0.0, + nnz: 1.0, + nnz_r: row_counts.as_ref().map_or(1.0, CountVector::nonzero_count), + nnz_c: col_counts.as_ref().map_or(1.0, CountVector::nonzero_count), + row_counts, + col_counts, + row_extended, + col_extended, + } + } + RpqPlan::Label(meta) => { + let counts = self.labels.get(&meta.name); + let row_counts = counts.map(|v| v.row_counts.clone()); + let col_counts = counts.map(|v| v.col_counts.clone()); + let row_extended = counts.map(|v| v.row_extended.clone()); + let col_extended = counts.map(|v| v.col_extended.clone()); + MncCost { + score: 0.0, + nnz: meta.nvals as f64, + nnz_r: row_counts + .as_ref() + .map_or(meta.nonzero_rows as f64, CountVector::nonzero_count), + nnz_c: col_counts + .as_ref() + .map_or(meta.nonzero_cols as f64, CountVector::nonzero_count), + row_counts, + col_counts, + row_extended, + col_extended, + } + } + RpqPlan::Seq([a, b]) => { + let ca = costs(*a); + let cb = costs(*b); + self.seq(ca, cb) + } + RpqPlan::Alt([a, b]) => { + let ca = costs(*a); + let cb = costs(*b); + self.alt(ca, cb) + } + RpqPlan::Star([a]) => { + let ca = costs(*a); + let penalty = self.star_penalty * ca.nnz.max(1.0); + MncCost { + score: ca.score + penalty, + nnz: self.n * self.n, + nnz_r: self.n, + nnz_c: self.n, + row_counts: None, + col_counts: None, + row_extended: None, + col_extended: None, + } + } + 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 op_cost = self.lr_multiplier * ca.nnz * cb.nnz / denom; + let nnz_est = self.lr_multiplier * ca.nnz * cb.nnz / (self.n * self.n).max(1.0); + MncCost { + score: ca.score + cb.score + op_cost, + nnz: nnz_est, + nnz_r: ca.nnz_r.min(self.n), + nnz_c: cb.nnz_c.min(self.n), + row_counts: None, + col_counts: None, + row_extended: None, + col_extended: None, + } + } + 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 op_cost = self.lr_multiplier * ca.nnz * cb.nnz / denom; + let nnz_est = self.lr_multiplier * ca.nnz * cb.nnz / (self.n * self.n).max(1.0); + MncCost { + score: ca.score + cb.score + op_cost, + nnz: nnz_est, + nnz_r: ca.nnz_r.min(self.n), + nnz_c: cb.nnz_c.min(self.n), + row_counts: None, + col_counts: None, + row_extended: None, + col_extended: None, + } + } + } + } +} + +#[cfg(test)] +mod tests { + use egg::RecExpr; + + use crate::rpq::rpqmatrix::{ + optimize::optimize_expr_join, + plan::{LabelMeta, RpqPlan}, + }; + use crate::{graph::GraphDecomposition, utils::build_graph}; + + use super::*; + + #[test] + fn independent_costs_preserve_seq_estimates() { + let a = Id::from(0); + let b = Id::from(1); + let left = RpqPlan::Label(LabelMeta { + name: "left".to_string(), + nvals: 10, + nonzero_rows: 2, + nonzero_cols: 4, + }); + let right = RpqPlan::Label(LabelMeta { + name: "right".to_string(), + nvals: 20, + nonzero_rows: 5, + nonzero_cols: 6, + }); + let seq = RpqPlan::Seq([a, b]); + + let mut metaac = MetaAcCostFn { + n: 10.0, + star_penalty: 50.0, + lr_multiplier: 5.0, + }; + let ma = metaac.cost(&left, |_| unreachable!()); + let mb = metaac.cost(&right, |_| unreachable!()); + let meta_cost = metaac.cost(&seq, |id| if id == a { ma.clone() } else { mb.clone() }); + let expected_nnz = 100.0 * (1.0 - (1.0_f64 - 0.02).powf(10.0)); + assert!((meta_cost.nnz - expected_nnz).abs() < 1e-10); + assert_eq!(meta_cost.score, 20.0); + assert_eq!((meta_cost.nnz_r, meta_cost.nnz_c), (2.0, 6.0)); + + let mut hybrid = HybridCostFn { + n: 10.0, + star_penalty: 50.0, + lr_multiplier: 5.0, + }; + let ha = hybrid.cost(&left, |_| unreachable!()); + let hb = hybrid.cost(&right, |_| unreachable!()); + let hybrid_cost = hybrid.cost(&seq, |id| if id == a { ha.clone() } else { hb.clone() }); + assert!((hybrid_cost.nnz - expected_nnz).abs() < 1e-10); + assert!((hybrid_cost.score - 200.0 / 6.0).abs() < 1e-10); + + let mut mnc = MncCostFn::new(10.0, HashMap::new(), HashMap::new()); + let ca = mnc.cost(&left, |_| unreachable!()); + let cb = mnc.cost(&right, |_| unreachable!()); + let cost = mnc.cost(&seq, |id| if id == a { ca.clone() } else { cb.clone() }); + assert_eq!(cost.score, meta_cost.score); + assert!((cost.nnz - meta_cost.nnz).abs() < 1e-10); + assert_eq!((cost.nnz_r, cost.nnz_c), (2.0, 6.0)); + } + + #[test] + fn mnc_seq_uses_leaf_extensions_without_propagating_them() { + let graph = build_graph(&[ + ("u", "a", "A"), + ("v", "b", "A"), + ("v", "c", "A"), + ("a", "x", "B"), + ("b", "x", "B"), + ("c", "y", "B"), + ]); + let labels = ["A", "B"] + .into_iter() + .map(|name| { + ( + name.to_string(), + LabelCountVectors::from_matrix(graph.get_graph(name).unwrap().matrix()) + .unwrap(), + ) + }) + .collect(); + let mut mnc = MncCostFn::new(7.0, labels, HashMap::new()); + let a = mnc.cost( + &RpqPlan::Label(LabelMeta { + name: "A".to_string(), + nvals: 3, + nonzero_rows: 2, + nonzero_cols: 3, + }), + |_| unreachable!(), + ); + let b = mnc.cost( + &RpqPlan::Label(LabelMeta { + name: "B".to_string(), + nvals: 3, + nonzero_rows: 3, + nonzero_cols: 2, + }), + |_| unreachable!(), + ); + let result = mnc.seq(a, b); + assert_eq!(result.nnz, 3.0); + assert!(result.row_extended.is_none()); + assert!(result.col_extended.is_none()); + } + + fn assert_finite_nonnegative(cost: &JoinCost) { + 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: &JoinCost, b: Id, cb: &JoinCost) -> JoinCost { + 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: &JoinCost) -> JoinCost { + 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 = JoinCost { + score: 10.0, + nnz: 20.0, + nnz_r: 30.0, + nnz_c: 40.0, + }; + + assert!( + JoinCost { + score: 9.0, + nnz: 100.0, + nnz_r: 100.0, + nnz_c: 100.0, + } < base + ); + assert!( + JoinCost { + score: 10.0, + nnz: 19.0, + nnz_r: 100.0, + nnz_c: 100.0, + } < base + ); + assert!( + JoinCost { + score: 10.0, + nnz: 20.0, + nnz_r: 29.0, + nnz_c: 100.0, + } < base + ); + assert!( + JoinCost { + score: 10.0, + nnz: 20.0, + nnz_r: 30.0, + nnz_c: 39.0, + } < base + ); + } + + #[test] + fn join_cost_base_nodes_use_vertex_and_label_metadata() { + let mut cost_fn = JoinCostFn { + 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, + JoinCost { + 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, + JoinCost { + score: 0.0, + nnz: 17.0, + nnz_r: 5.0, + nnz_c: 9.0, + } + ); + } + #[test] + fn join_cost_seq_correctness() { + let mut cost_fn = JoinCostFn { + n: 100.0, + star_penalty: 50.0, + lr_multiplier: 5.0, + }; + let a = Id::from(0); + let b = Id::from(1); + + let ca = JoinCost { + score: 2.0, + nnz: 20.0, + nnz_r: 4.0, + nnz_c: 8.0, + }; + + let cb = JoinCost { + 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, + JoinCost { + score: 65.0, + nnz: 0.06, + nnz_r: 4.0, + nnz_c: 10.0, + } + ); + } + #[test] + fn join_cost_seq_correctness_zero_denom() { + let mut cost_fn = JoinCostFn { + n: 100.0, + star_penalty: 50.0, + lr_multiplier: 5.0, + }; + let a = Id::from(0); + let b = Id::from(1); + + let ca = JoinCost { + score: 2.0, + nnz: 20.0, + nnz_r: 0.0, + nnz_c: 4.0, + }; + + let cb = JoinCost { + 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, + JoinCost { + score: 605.0, + nnz: 0.06, + nnz_r: 0.0, + nnz_c: 0.0, + } + ); + } + #[test] + fn join_cost_alt_with_zero_children_stays_finite_nonnegative() { + let mut cost_fn = JoinCostFn { + n: 100.0, + star_penalty: 50.0, + lr_multiplier: 5.0, + }; + let a = Id::from(0); + let b = Id::from(1); + let ca = JoinCost { + score: 0.0, + nnz: 0.0, + nnz_r: 0.0, + nnz_c: 0.0, + }; + let cb = JoinCost { + 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, + JoinCost { + score: 33.0, + nnz: 30.0, + nnz_r: 0.0, + nnz_c: 10.0, + } + ); + } + + #[test] + fn join_cost_star_with_zero_nnz_uses_min_penalty() { + let mut cost_fn = JoinCostFn { + n: 100.0, + star_penalty: 50.0, + lr_multiplier: 5.0, + }; + let a = Id::from(0); + let ca = JoinCost { + 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, + JoinCost { + score: 57.0, + nnz: 10_000.0, + nnz_r: 100.0, + nnz_c: 100.0, + } + ); + } + + #[test] + fn join_cost_lstar_and_rstar_zero_denom_stay_finite_nonnegative() { + let mut cost_fn = JoinCostFn { + n: 100.0, + star_penalty: 50.0, + lr_multiplier: 5.0, + }; + let a = Id::from(0); + let b = Id::from(1); + let ca = JoinCost { + score: 2.0, + nnz: 20.0, + nnz_r: 0.0, + nnz_c: 4.0, + }; + let cb = JoinCost { + 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 = JoinCost { + score: 3005.0, + nnz: 0.3, + nnz_r: 0.0, + nnz_c: 0.0, + }; + + assert_eq!(lstar, expected); + assert_eq!(rstar, expected); + } + #[test] + fn join_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_join(expr, 100); + let root = opt.as_ref().last().expect("optimized expr is non-empty"); + + assert!(matches!(root, RpqPlan::LStar(_))); + } + #[test] + fn join_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_join(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/src/rpq/rpqmatrix/eval.rs b/pathrex/src/rpq/rpqmatrix/eval.rs new file mode 100644 index 0000000..5777db8 --- /dev/null +++ b/pathrex/src/rpq/rpqmatrix/eval.rs @@ -0,0 +1,178 @@ +use super::expr::{materialize_with_storage, query_to_expr}; +use super::optimize::{ + OptimizationStrategy, OptimizerCache, optimize_expr_hybrid, optimize_expr_join, + optimize_expr_metaac, optimize_expr_mnc, +}; +use super::result::{PreparedRpqMatrix, RpqMatrixResult}; +/// RPQ evaluator backed by `LAGraph_RPQMatrix`. +use crate::eval::Evaluator; +use crate::graph::{GraphDecomposition, MatrixStorage, set_global_matrix_storage_hint}; +use crate::rpq::{Endpoint, RpqError, RpqQuery}; +use std::sync::{Arc, Mutex}; + +#[derive(Clone)] +pub struct RpqMatrixEvaluator { + optimizer: OptimizationStrategy, + cache: Arc>, +} + +impl RpqMatrixEvaluator { + pub fn unoptimized() -> Self { + return RpqMatrixEvaluator { + optimizer: OptimizationStrategy::NoOpt, + cache: Arc::default(), + }; + } + pub fn optimized(opt: OptimizationStrategy) -> Self { + return RpqMatrixEvaluator { + optimizer: opt, + cache: Arc::default(), + }; + } +} + +impl Default for RpqMatrixEvaluator { + fn default() -> Self { + RpqMatrixEvaluator::unoptimized() + } +} + +fn storage_for_query(query: &RpqQuery) -> MatrixStorage { + match (&query.subject, &query.object) { + (Endpoint::Variable(_), Endpoint::Named(_)) => MatrixStorage::Csc, + _ => MatrixStorage::Csr, + } +} + +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 storage = storage_for_query(query); + set_global_matrix_storage_hint(storage)?; + + let mut expr = query_to_expr(query, graph)?; + match self.optimizer { + OptimizationStrategy::NoOpt => {} + OptimizationStrategy::Join => { + expr = optimize_expr_join(expr, graph.num_nodes()); + } + OptimizationStrategy::MetaAc => expr = optimize_expr_metaac(expr, graph.num_nodes()), + OptimizationStrategy::Mnc => expr = optimize_expr_mnc(expr, graph, &self.cache)?, + OptimizationStrategy::Hybrid => expr = optimize_expr_hybrid(expr, graph.num_nodes()), + OptimizationStrategy::RandomOpt + | OptimizationStrategy::Simple + | OptimizationStrategy::Wander => todo!(), + } + + let (plans, owned_matrices) = materialize_with_storage(&expr, graph, storage)?; + + Ok(PreparedRpqMatrix { + plans, + owned_matrices, + storage, + }) + } +} + +#[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"); + } + + #[test] + fn optimizers_preserve_results() { + let graph = build_graph(&[ + ("A", "B", "p"), + ("B", "C", "p"), + ("B", "C", "q"), + ("A", "D", "q"), + ]); + let queries = [ + ( + RpqQuery { + subject: Endpoint::Named("A".into()), + path: PathExpr::Sequence( + Box::new(PathExpr::Label("p".into())), + Box::new(PathExpr::Label("q".into())), + ), + object: Endpoint::Variable("y".into()), + }, + 1, + ), + ( + RpqQuery { + subject: Endpoint::Named("A".into()), + path: PathExpr::ZeroOrMore(Box::new(PathExpr::Label("p".into()))), + object: Endpoint::Variable("y".into()), + }, + 3, + ), + ( + RpqQuery { + subject: Endpoint::Variable("x".into()), + path: PathExpr::ZeroOrMore(Box::new(PathExpr::Label("p".into()))), + object: Endpoint::Named("C".into()), + }, + 3, + ), + ]; + + for optimizer in [ + OptimizationStrategy::MetaAc, + OptimizationStrategy::Mnc, + OptimizationStrategy::Hybrid, + ] { + let evaluator = RpqMatrixEvaluator::optimized(optimizer); + for (query, expected) in &queries { + for _ in 0..2 { + let result = evaluator + .evaluate(query, &graph) + .expect("optimized evaluation"); + assert_eq!( + result.nnz, *expected, + "optimizer={optimizer:?}, query={query:?}" + ); + } + } + } + } +} diff --git a/pathrex/src/rpq/rpqmatrix.rs b/pathrex/src/rpq/rpqmatrix/expr.rs similarity index 51% rename from pathrex/src/rpq/rpqmatrix.rs rename to pathrex/src/rpq/rpqmatrix/expr.rs index 1462abc..796ddc8 100644 --- a/pathrex/src/rpq/rpqmatrix.rs +++ b/pathrex/src/rpq/rpqmatrix/expr.rs @@ -1,50 +1,61 @@ -//! Plan-based RPQ evaluation using `LAGraph_RPQMatrix`. - use std::ptr::null_mut; -use egg::{Id, RecExpr, define_language}; +use egg::{Id, RecExpr}; -use crate::eval::{Evaluator, PreparedEvaluator, ResultCount}; -use crate::graph::{GraphDecomposition, GraphError, GraphblasMatrix}; +use super::plan::{LabelMeta, RpqPlan}; +use crate::graph::{GraphDecomposition, MatrixStorage}; +use crate::grb_ok; use crate::lagraph_sys::*; use crate::rpq::{Endpoint, PathExpr, RpqError, RpqQuery}; -use crate::{grb_ok, la_ok}; - -const RPQMATRIX_REDUCE_BY_COL: u8 = 1; -define_language! { - pub enum RpqPlan { - Label(String), - NamedVertex(String), - "/" = Seq([Id; 2]), - "|" = Alt([Id; 2]), - "*" = Star([Id; 1]), +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) -> Result { +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 +68,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, @@ -87,9 +101,10 @@ pub fn query_to_expr(query: &RpqQuery) -> Result, RpqError> { /// /// Returns the plan array and a list of owned diagonal matrices that must be /// freed after evaluation. -pub fn materialize( +pub fn materialize_with_storage( expr: &RecExpr, graph: &G, + storage: MatrixStorage, ) -> Result<(Vec, Vec), RpqError> { let null_plan = RPQMatrixPlan { op: RPQMatrixOp::RPQ_MATRIX_OP_LABEL, @@ -105,7 +120,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_with_storage(&label.name, storage)?; let mat = unsafe { (*lg.inner).A }; RPQMatrixPlan { op: RPQMatrixOp::RPQ_MATRIX_OP_LABEL, @@ -164,135 +179,24 @@ pub fn materialize( 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, - RPQMATRIX_REDUCE_BY_COL, - ))? - }; - 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)?; - 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; + 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(), + }, - #[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()), + 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(), + }, }; - 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"); - } + Ok((plans, owned_matrices)) } diff --git a/pathrex/src/rpq/rpqmatrix/mod.rs b/pathrex/src/rpq/rpqmatrix/mod.rs new file mode 100644 index 0000000..22c3996 --- /dev/null +++ b/pathrex/src/rpq/rpqmatrix/mod.rs @@ -0,0 +1,13 @@ +//! Plan-based RPQ evaluation using `LAGraph_RPQMatrix`. + +mod cost; +pub mod eval; +mod expr; +mod optimize; +mod plan; +pub mod result; +mod stats; + +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..6c5d159 --- /dev/null +++ b/pathrex/src/rpq/rpqmatrix/optimize.rs @@ -0,0 +1,206 @@ +use egg::{Extractor, RecExpr, Runner}; +use std::{ + collections::{HashMap, HashSet}, + sync::{LazyLock, Mutex}, +}; + +use super::cost::{HybridCostFn, JoinCostFn, MetaAcCostFn, MncCostFn}; +use super::plan::{RpqPlan, make_rules}; +use super::stats::LabelCountVectors; +use crate::{graph::GraphDecomposition, rpq::RpqError}; + +static RULES: LazyLock>> = LazyLock::new(make_rules); + +#[derive(Clone, Copy, Debug)] +pub enum OptimizationStrategy { + NoOpt, + Join, + MetaAc, + Mnc, + Hybrid, + 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 +} + +struct CachedLabel { + matrix_id: usize, + counts: LabelCountVectors, +} + +struct CachedVertex { + vertex: usize, + n: usize, + counts: LabelCountVectors, +} + +#[derive(Default)] +pub(super) struct OptimizerCache { + labels: HashMap, + vertices: HashMap, +} + +fn runner(expr: &RecExpr) -> Runner { + Runner::default() + .with_explanations_disabled() + .with_expr(expr) + .run(&*RULES) +} + +fn extract_with>( + expr: &RecExpr, + cost: C, +) -> RecExpr { + let runner = runner(expr); + Extractor::new(&runner.egraph, cost) + .find_best(runner.roots[0]) + .1 +} + +pub(super) fn optimize_expr_join(expr: RecExpr, graph_size: usize) -> RecExpr { + extract_with( + &expr, + JoinCostFn { + n: graph_size as f64, + star_penalty: 50.0, + lr_multiplier: 5.0, + }, + ) +} + +pub(super) fn optimize_expr_metaac(expr: RecExpr, n: usize) -> RecExpr { + extract_with( + &expr, + MetaAcCostFn { + n: n as f64, + star_penalty: 50.0, + lr_multiplier: 5.0, + }, + ) +} + +pub(super) fn optimize_expr_mnc( + expr: RecExpr, + graph: &G, + cache: &Mutex, +) -> Result, RpqError> { + let labels = cached_label_data(&labels(&expr), graph, cache)?; + let vertices = vertices(&expr, graph)?; + let vertex_counts = vertex_counts(&vertices, graph.num_nodes(), cache)?; + Ok(extract_with( + &expr, + MncCostFn::new(graph.num_nodes() as f64, labels, vertex_counts), + )) +} + +pub(super) fn optimize_expr_hybrid(expr: RecExpr, n: usize) -> RecExpr { + extract_with( + &expr, + HybridCostFn { + n: n as f64, + star_penalty: 50.0, + lr_multiplier: 5.0, + }, + ) +} + +fn labels(expr: &RecExpr) -> HashSet { + expr.as_ref() + .iter() + .filter_map(|node| match node { + RpqPlan::Label(meta) => Some(meta.name.clone()), + _ => None, + }) + .collect() +} + +fn cached_label_data( + names: &HashSet, + graph: &G, + cache: &Mutex, +) -> Result, RpqError> { + let mut cache = cache.lock().expect("RPQ optimizer cache poisoned"); + let mut counts = HashMap::with_capacity(names.len()); + for name in names { + let matrix = graph.get_graph(name)?.matrix(); + let matrix_id = matrix as usize; + let stale = cache + .labels + .get(name) + .is_none_or(|entry| entry.matrix_id != matrix_id); + if stale { + let vectors = LabelCountVectors::from_matrix(matrix) + .ok_or_else(|| RpqError::UnsupportedPath("unable to build count vectors".into()))?; + cache.labels.insert( + name.clone(), + CachedLabel { + matrix_id, + counts: vectors, + }, + ); + } + counts.insert( + name.clone(), + cache + .labels + .get(name) + .expect("cached label was inserted") + .counts + .clone(), + ); + } + Ok(counts) +} + +fn vertices( + expr: &RecExpr, + graph: &G, +) -> Result, RpqError> { + expr.as_ref() + .iter() + .filter_map(|node| match node { + RpqPlan::NamedVertex(name) => Some(name), + _ => None, + }) + .map(|name| { + graph + .get_node_id(name) + .map(|id| (name.clone(), id)) + .ok_or_else(|| RpqError::VertexNotFound(name.clone())) + }) + .collect() +} + +fn vertex_counts( + vertices: &HashMap, + n: usize, + cache: &Mutex, +) -> Result, RpqError> { + let mut cache = cache.lock().expect("RPQ optimizer cache poisoned"); + vertices + .iter() + .map(|(name, &vertex)| { + let stale = cache + .vertices + .get(name) + .is_none_or(|entry| entry.vertex != vertex || entry.n != n); + if stale { + let counts = LabelCountVectors::from_vertex(vertex, n).ok_or_else(|| { + RpqError::UnsupportedPath("unable to build endpoint statistics".into()) + })?; + cache + .vertices + .insert(name.clone(), CachedVertex { vertex, n, counts }); + } + Ok(( + name.clone(), + cache + .vertices + .get(name) + .expect("vertex was cached") + .counts + .clone(), + )) + }) + .collect() +} diff --git a/pathrex/src/rpq/rpqmatrix/plan.rs b/pathrex/src/rpq/rpqmatrix/plan.rs new file mode 100644 index 0000000..7a39527 --- /dev/null +++ b/pathrex/src/rpq/rpqmatrix/plan.rs @@ -0,0 +1,56 @@ +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!("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..72d9215 --- /dev/null +++ b/pathrex/src/rpq/rpqmatrix/result.rs @@ -0,0 +1,81 @@ +use std::ptr::null_mut; + +use crate::eval::{PreparedEvaluator, ResultCount}; +use crate::graph::wrappers::ReduceType::ByCols; +use crate::graph::{GraphError, GraphblasMatrix, MatrixStorage, set_global_matrix_storage_hint}; +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, + pub(super) storage: MatrixStorage, +} + +impl PreparedEvaluator for PreparedRpqMatrix { + type Result = RpqMatrixResult; + type Error = RpqError; + + fn execute(&mut self) -> Result { + set_global_matrix_storage_hint(self.storage)?; + + 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/src/rpq/rpqmatrix/stats.rs b/pathrex/src/rpq/rpqmatrix/stats.rs new file mode 100644 index 0000000..3420a9c --- /dev/null +++ b/pathrex/src/rpq/rpqmatrix/stats.rs @@ -0,0 +1,345 @@ +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; + +use crate::lagraph_sys::{ + GrB_Info, GrB_Matrix, GrB_Vector, GrB_Vector_free, GrB_Vector_nvals, LAGraph_RPQMatrix_Free, + LAGraph_RPQMatrix_count_vector_dot, LAGraph_RPQMatrix_count_vector_mnc_add, + LAGraph_RPQMatrix_count_vector_mnc_matmul_nnz, LAGraph_RPQMatrix_count_vector_scale, + LAGraph_RPQMatrix_count_vector_sum, LAGraph_RPQMatrix_extended_count_vectors, + LAGraph_RPQMatrix_label, LAGraph_RPQMatrix_reduce_count_vector, +}; + +#[derive(Debug)] +struct CountVectorHandle(GrB_Vector); + +static NEXT_VECTOR_ID: AtomicUsize = AtomicUsize::new(1); + +impl Drop for CountVectorHandle { + fn drop(&mut self) { + if !self.0.is_null() { + unsafe { GrB_Vector_free(&mut self.0) }; + } + } +} + +unsafe impl Send for CountVectorHandle {} +unsafe impl Sync for CountVectorHandle {} + +#[derive(Clone, Debug)] +pub(super) struct CountVector { + handle: Arc, + nvals: usize, + sum: f64, + id: usize, +} + +impl CountVector { + fn from_matrix(matrix: GrB_Matrix, by_columns: bool) -> Option { + let mut vector = std::ptr::null_mut(); + let code = unsafe { + LAGraph_RPQMatrix_reduce_count_vector(&mut vector, matrix, u8::from(by_columns)) + }; + (code == GrB_Info::GrB_SUCCESS) + .then(|| Self::from_owned(vector)) + .flatten() + } + + fn from_owned(mut vector: GrB_Vector) -> Option { + let mut nvals = 0; + let mut sum = 0.0; + let ok = unsafe { + GrB_Vector_nvals(&mut nvals, vector) == GrB_Info::GrB_SUCCESS + && LAGraph_RPQMatrix_count_vector_sum(&mut sum, vector) == GrB_Info::GrB_SUCCESS + }; + if ok { + Some(Self { + handle: Arc::new(CountVectorHandle(vector)), + nvals: nvals as usize, + sum, + id: NEXT_VECTOR_ID.fetch_add(1, Ordering::Relaxed), + }) + } else { + unsafe { GrB_Vector_free(&mut vector) }; + None + } + } + + fn raw(&self) -> GrB_Vector { + self.handle.0 + } + + pub(super) fn cache_key(&self) -> usize { + self.id + } + + pub(super) fn sum(&self) -> f64 { + self.sum + } + + pub(super) fn dot(&self, other: &Self) -> Option { + let mut result = 0.0; + let code = + unsafe { LAGraph_RPQMatrix_count_vector_dot(&mut result, self.raw(), other.raw()) }; + (code == GrB_Info::GrB_SUCCESS).then_some(result) + } + + pub(super) fn mnc_matmul_nnz( + lhs_rows: &Self, + lhs_cols: &Self, + rhs_rows: &Self, + rhs_cols: &Self, + lhs_col_extended: Option<&Self>, + rhs_row_extended: Option<&Self>, + ) -> Option { + let mut result = 0.0; + let code = unsafe { + LAGraph_RPQMatrix_count_vector_mnc_matmul_nnz( + &mut result, + lhs_rows.raw(), + lhs_cols.raw(), + rhs_rows.raw(), + rhs_cols.raw(), + lhs_col_extended.map_or(std::ptr::null_mut(), Self::raw), + rhs_row_extended.map_or(std::ptr::null_mut(), Self::raw), + ) + }; + (code == GrB_Info::GrB_SUCCESS).then_some(result) + } + + pub(super) fn mnc_add(&self, other: &Self, lambda: f64, cap: f64) -> Option { + let mut result = std::ptr::null_mut(); + let code = unsafe { + LAGraph_RPQMatrix_count_vector_mnc_add( + &mut result, + self.raw(), + other.raw(), + lambda, + cap, + ) + }; + (code == GrB_Info::GrB_SUCCESS) + .then(|| Self::from_owned(result)) + .flatten() + } + + pub(super) fn scale(&self, scale: f64, cap: f64) -> Option { + let mut result = std::ptr::null_mut(); + let code = + unsafe { LAGraph_RPQMatrix_count_vector_scale(&mut result, self.raw(), scale, cap) }; + (code == GrB_Info::GrB_SUCCESS) + .then(|| Self::from_owned(result)) + .flatten() + } + + #[allow(dead_code)] + pub(super) fn nonzero_count(&self) -> f64 { + self.nvals as f64 + } +} + +#[derive(Clone, Debug)] +pub(super) struct LabelCountVectors { + pub row_counts: CountVector, + pub col_counts: CountVector, + pub row_extended: CountVector, + pub col_extended: CountVector, +} + +impl LabelCountVectors { + pub(super) fn from_matrix(matrix: GrB_Matrix) -> Option { + let row_counts = CountVector::from_matrix(matrix, false)?; + let col_counts = CountVector::from_matrix(matrix, true)?; + let mut row_extended = std::ptr::null_mut(); + let mut col_extended = std::ptr::null_mut(); + let code = unsafe { + LAGraph_RPQMatrix_extended_count_vectors( + &mut row_extended, + &mut col_extended, + matrix, + row_counts.raw(), + col_counts.raw(), + ) + }; + if code != GrB_Info::GrB_SUCCESS { + return None; + } + let row_extended = CountVector::from_owned(row_extended); + let col_extended = CountVector::from_owned(col_extended); + Some(Self { + row_counts, + col_counts, + row_extended: row_extended?, + col_extended: col_extended?, + }) + } + + pub(super) fn from_vertex(vertex: usize, n: usize) -> Option { + let mut matrix = std::ptr::null_mut(); + let create = unsafe { LAGraph_RPQMatrix_label(&mut matrix, vertex as _, n as _, n as _) }; + if create != GrB_Info::GrB_SUCCESS { + return None; + } + let counts = Self::from_matrix(matrix); + let free = unsafe { LAGraph_RPQMatrix_Free(&mut matrix) }; + (free == GrB_Info::GrB_SUCCESS).then_some(counts).flatten() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{graph::GraphDecomposition, utils::build_graph}; + + #[test] + fn graphblas_apply_zero_retains_explicit_zero_entries() { + use crate::lagraph_sys::{ + GrB_BinaryOp, GrB_Descriptor, GrB_Type, GrB_Vector_new, GrB_Vector_nvals, + }; + + unsafe extern "C" { + static mut GrB_FP64: GrB_Type; + static mut GrB_TIMES_FP64: GrB_BinaryOp; + fn GrB_Vector_apply_BinaryOp1st_FP64( + output: GrB_Vector, + mask: GrB_Vector, + accum: GrB_BinaryOp, + op: GrB_BinaryOp, + scalar: f64, + input: GrB_Vector, + descriptor: GrB_Descriptor, + ) -> GrB_Info; + } + + let graph = build_graph(&[("A", "B", "p"), ("B", "C", "p")]); + let counts = + LabelCountVectors::from_matrix(graph.get_graph("p").unwrap().matrix()).unwrap(); + let mut raw = std::ptr::null_mut(); + unsafe { + assert_eq!(GrB_Vector_new(&mut raw, GrB_FP64, 3), GrB_Info::GrB_SUCCESS); + assert_eq!( + GrB_Vector_apply_BinaryOp1st_FP64( + raw, + std::ptr::null_mut(), + std::ptr::null_mut(), + GrB_TIMES_FP64, + 0.0, + counts.row_counts.raw(), + std::ptr::null_mut(), + ), + GrB_Info::GrB_SUCCESS, + ); + let mut stored = 0; + assert_eq!(GrB_Vector_nvals(&mut stored, raw), GrB_Info::GrB_SUCCESS); + assert_eq!(stored, 2, "GraphBLAS should retain the two explicit zeros"); + } + let result = CountVector::from_owned(raw).unwrap(); + assert_eq!(result.sum(), 0.0); + assert_eq!(result.nonzero_count(), 2.0); + } + + #[test] + fn mnc_vector_operations_use_graphblas_counts() { + let graph = build_graph(&[("A", "B", "p"), ("B", "C", "p")]); + let matrix = graph.get_graph("p").unwrap().matrix(); + let counts = LabelCountVectors::from_matrix(matrix).unwrap(); + assert_eq!(counts.row_counts.sum(), 2.0); + assert_eq!(counts.col_counts.sum(), 2.0); + assert_eq!(counts.col_counts.dot(&counts.row_counts), Some(1.0)); + assert_eq!( + CountVector::mnc_matmul_nnz( + &counts.row_counts, + &counts.col_counts, + &counts.row_counts, + &counts.col_counts, + Some(&counts.col_extended), + Some(&counts.row_extended), + ), + Some(1.0), + ); + let union = counts + .row_counts + .mnc_add(&counts.row_counts, 0.5, 3.0) + .unwrap(); + assert_eq!(union.sum(), 3.0); + assert_eq!(union.nonzero_count(), 2.0); + let fractional_union = counts + .row_counts + .mnc_add(&counts.row_counts, 0.5, 3.0) + .unwrap(); + assert_eq!(fractional_union.sum(), 3.0); + let empty = counts.row_counts.scale(0.0, 3.0).unwrap(); + assert_eq!(empty.sum(), 0.0); + assert_eq!(empty.nonzero_count(), 0.0); + let fractional = counts.row_counts.scale(0.4, 3.0).unwrap(); + assert_eq!(fractional.sum(), 0.8); + assert_eq!(fractional.nonzero_count(), 2.0); + } + + #[test] + fn mnc_extension_vectors_account_for_singleton_rows_and_columns() { + let graph = build_graph(&[ + ("u", "a", "A"), + ("v", "b", "A"), + ("v", "c", "A"), + ("a", "x", "B"), + ("b", "x", "B"), + ("c", "y", "B"), + ]); + let a = LabelCountVectors::from_matrix(graph.get_graph("A").unwrap().matrix()).unwrap(); + let b = LabelCountVectors::from_matrix(graph.get_graph("B").unwrap().matrix()).unwrap(); + assert_eq!(a.row_counts.sum(), 3.0); + assert_eq!(a.col_counts.sum(), 3.0); + assert_eq!(b.row_counts.sum(), 3.0); + assert_eq!(b.col_counts.sum(), 3.0); + assert_eq!(a.row_counts.nonzero_count(), 2.0); + assert_eq!(a.col_counts.nonzero_count(), 3.0); + assert_eq!(b.row_counts.nonzero_count(), 3.0); + assert_eq!(b.col_counts.nonzero_count(), 2.0); + assert_eq!(a.row_extended.nonzero_count(), 2.0); + assert_eq!(a.col_extended.nonzero_count(), 1.0); + assert_eq!(b.row_extended.nonzero_count(), 1.0); + assert_eq!(b.col_extended.nonzero_count(), 2.0); + assert_eq!(a.row_extended.sum(), 3.0); + assert_eq!(a.col_extended.sum(), 1.0); + assert_eq!(b.row_extended.sum(), 1.0); + assert_eq!(b.col_extended.sum(), 3.0); + + let estimate = |lhs_ext, rhs_ext| { + CountVector::mnc_matmul_nnz( + &a.row_counts, + &a.col_counts, + &b.row_counts, + &b.col_counts, + lhs_ext, + rhs_ext, + ) + .unwrap() + }; + assert!((estimate(None, None) - 2.3125).abs() < 1e-10); + let lhs_only = estimate(Some(&a.col_extended), None); + let rhs_only = estimate(None, Some(&b.row_extended)); + assert!((lhs_only - 2.5).abs() < 1e-10, "lhs-only: {lhs_only}"); + assert!((rhs_only - 2.5).abs() < 1e-10, "rhs-only: {rhs_only}"); + assert_eq!(estimate(Some(&a.col_extended), Some(&b.row_extended)), 3.0); + } + + #[test] + fn fractional_counts_do_not_trigger_the_exact_singleton_case() { + let graph = build_graph(&[("A", "A", "p"), ("B", "B", "p")]); + let counts = + LabelCountVectors::from_matrix(graph.get_graph("p").unwrap().matrix()).unwrap(); + let fractional = counts.row_counts.scale(0.4, 2.0).unwrap(); + let estimate = CountVector::mnc_matmul_nnz( + &fractional, + &fractional, + &fractional, + &fractional, + None, + None, + ) + .unwrap(); + assert!((estimate - 0.3136).abs() < 1e-10, "estimate: {estimate}"); + } +} 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`]. 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/rpqmatrix_tests.rs b/pathrex/tests/rpqmatrix_tests.rs index 3f84d80..064f9af 100644 --- a/pathrex/tests/rpqmatrix_tests.rs +++ b/pathrex/tests/rpqmatrix_tests.rs @@ -3,10 +3,13 @@ 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::{RpqMatrixEvaluator, RpqMatrixResult}; +use pathrex::rpq::rpqmatrix::OptimizationStrategy::{Hybrid, Join, MetaAc, Mnc}; +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; @@ -67,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); @@ -79,7 +82,6 @@ fn run_la_n_egg_case(case_name: &str) { ); let graph = &*LA_N_EGG_GRAPH; - let evaluator = RpqMatrixEvaluator; for (i, (query, expected_nnz)) in queries.iter().zip(expected.iter()).enumerate() { let result = evaluator.evaluate(query, graph).unwrap_or_else(|e| { @@ -95,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_join(case_name: &str) { + run_la_n_egg_case_with_evaluator(case_name, RpqMatrixEvaluator::optimized(Join)); +} + fn label(s: &str) -> PathExpr { PathExpr::Label(s.to_string()) } @@ -123,12 +133,37 @@ fn matrix_entry_set(result: &RpqMatrixResult, row: GrB_Index, col: GrB_Index) -> } } +// TODO: made it reusable for different optimizers +fn evaluate_default_and_join( + graph: &InMemoryGraph, + query: &RpqQuery, +) -> (RpqMatrixResult, RpqMatrixResult) { + let default_result = RpqMatrixEvaluator::default() + .evaluate(query, graph) + .expect("default evaluator should succeed"); + let optimized_result = RpqMatrixEvaluator::optimized(Join) + .evaluate(query, graph) + .expect("join 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] 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 +177,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 +197,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 +217,12 @@ 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 +237,9 @@ 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 +251,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 +273,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 +300,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 +332,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 +362,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 +376,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 +389,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 +404,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 +418,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 +443,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 +482,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 +509,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"))); @@ -494,3 +535,128 @@ fn test_la_n_egg_any_con() { fn test_la_n_egg_con_any() { run_la_n_egg_case("con-any"); } + +#[test] +fn test_la_n_egg_any_any_join_optimizer() { + run_la_n_egg_case_join("any-any"); +} + +#[test] +fn test_la_n_egg_any_con_join_optimizer() { + run_la_n_egg_case_join("con-any"); +} + +#[test] +fn test_la_n_egg_cases_with_core_optimizers() { + for optimizer in [MetaAc, Mnc, Hybrid] { + let evaluator = RpqMatrixEvaluator::optimized(optimizer); + for case in ["any-any", "any-con", "con-any"] { + run_la_n_egg_case_with_evaluator(case, evaluator.clone()); + } + } +} + +#[test] +fn test_join_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_join(&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_join_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_join(&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_join_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_join(&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" + ); + } +} 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