diff --git a/Cargo.lock b/Cargo.lock index 0f372ba..7c61376 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4110,6 +4110,7 @@ dependencies = [ "arrow-array", "arrow-schema", "criterion", + "datafusion", "futures", "lance", "lance-graph", diff --git a/crates/lance-graph-benches/Cargo.toml b/crates/lance-graph-benches/Cargo.toml index 6d0b613..3881098 100644 --- a/crates/lance-graph-benches/Cargo.toml +++ b/crates/lance-graph-benches/Cargo.toml @@ -10,6 +10,7 @@ lance-graph = { path = "../lance-graph", version = "0.2.0" } lance = "13.0.0" arrow-array = "58" arrow-schema = "58" +datafusion = { version = "54", default-features = false } criterion = { version = "0.5", features = ["async", "async_tokio", "html_reports"] } futures = "0.3" tempfile = "3" diff --git a/crates/lance-graph-benches/README.md b/crates/lance-graph-benches/README.md index ce63a57..fc00241 100644 --- a/crates/lance-graph-benches/README.md +++ b/crates/lance-graph-benches/README.md @@ -26,3 +26,27 @@ cargo bench -p lance-graph-benches --bench graph_execution This crate is not published to crates.io and is excluded from releases. It exists solely for performance testing during development. + + +## Query template reuse + +Measure fresh parsing/planning/execution against reuse of an unoptimized scalar-parameter template on the +same 10,000-row Arrow snapshot. Both paths optimize and build a physical plan for each execution. The harness +checks identical Arrow results, row count and ID checksum before timing; catalog construction is excluded. + +```bash +cargo bench -p lance-graph-benches --bench graph_execution -- query_templates --warm-up-time 1 +``` + +For a stage breakdown on a separate fixed 1,000-row scan, run the optional core diagnostic: + +```bash +cargo test --release -p lance-graph profile_query_planning_stages -- --ignored --nocapture +``` + +It reports elapsed means for 100 iterations of parsing, semantic analysis, graph logical planning, +DataFusion logical planning, optimization, physical planning, execution and cached binding. Physical timing +calls DataFusion's query planner directly so optimization is counted only once. It verifies 990 result rows +and their ID checksum. These local measurements do not replace the two SF1 suites required for performance +submissions; record runtime, hardware, revision and build profile with any reported numbers. The engine's +Python cache also includes dictionary conversion and lookup, which this Rust microbenchmark does not time. diff --git a/crates/lance-graph-benches/benches/graph_execution.rs b/crates/lance-graph-benches/benches/graph_execution.rs index 2b6831b..333ec6e 100644 --- a/crates/lance-graph-benches/benches/graph_execution.rs +++ b/crates/lance-graph-benches/benches/graph_execution.rs @@ -336,5 +336,75 @@ fn bench_cypher_execution(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench_cypher_execution); +// A fixed snapshot isolates template reuse from dataset/catalog construction. +// The result count and checksum are checked outside the timed loops. +fn bench_query_templates(c: &mut Criterion) { + use datafusion::datasource::{DefaultTableSource, MemTable}; + use datafusion::execution::context::SessionContext; + use lance_graph::{GraphSourceCatalog, InMemoryCatalog}; + let rt = tokio::runtime::Runtime::new().unwrap(); + let context = SessionContext::new(); + let batch = make_people_batch(10_000); + let provider = Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch]]).unwrap()); + let catalog: Arc = Arc::new( + InMemoryCatalog::new() + .with_node_source("Person", Arc::new(DefaultTableSource::new(provider))), + ); + let config = GraphConfig::builder() + .node_table("Person", "person_id") + .build() + .unwrap(); + let text = "MATCH (p:Person) WHERE p.age >= $minimum RETURN p.person_id AS id ORDER BY id"; + let fresh = || { + CypherQuery::new(text) + .unwrap() + .with_config(config.clone()) + .with_parameter("minimum", 50) + }; + let mut template = fresh(); + assert!(template.prepare_cached_plan(catalog.clone()).unwrap()); + let expected = rt + .block_on(fresh().execute_with_catalog_and_context(catalog.clone(), context.clone())) + .unwrap(); + let cached = rt + .block_on(template.execute_with_catalog_and_context(catalog.clone(), context.clone())) + .unwrap(); + assert_eq!(cached, expected); + let expected_ids: Vec = (0..10_000).filter(|id| 20 + id % 60 >= 50).collect(); + assert_eq!(cached.num_rows(), expected_ids.len()); + let ids = cached + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!( + ids.values().iter().map(|id| i64::from(*id)).sum::(), + expected_ids.iter().map(|id| i64::from(*id)).sum::() + ); + let mut group = c.benchmark_group("query_templates"); + group.bench_function("cold_parse_plan_execute", |b| { + b.iter(|| { + black_box( + rt.block_on( + fresh().execute_with_catalog_and_context(catalog.clone(), context.clone()), + ) + .unwrap(), + ); + }) + }); + group.bench_function("cached_bind_plan_execute", |b| { + b.iter(|| { + let query = template.clone().with_parameter("minimum", black_box(50)); + black_box( + rt.block_on( + query.execute_with_catalog_and_context(catalog.clone(), context.clone()), + ) + .unwrap(), + ); + }) + }); + group.finish(); +} + +criterion_group!(benches, bench_cypher_execution, bench_query_templates); criterion_main!(benches); diff --git a/crates/lance-graph-python/src/graph.rs b/crates/lance-graph-python/src/graph.rs index 1112d3d..4690d76 100644 --- a/crates/lance-graph-python/src/graph.rs +++ b/crates/lance-graph-python/src/graph.rs @@ -14,8 +14,8 @@ //! Graph query functionality for Lance datasets -use std::collections::HashMap; -use std::sync::{Arc, RwLock}; +use std::collections::{HashMap, VecDeque}; +use std::sync::{Arc, Mutex, RwLock}; use arrow::compute::concat_batches; use arrow::ffi_stream::ArrowArrayStreamReader; @@ -1234,6 +1234,7 @@ enum EngineSource { /// keep the `Arc` they started with, so they finish on their original snapshot. struct EngineState { sources: HashMap, + plans: Mutex>, catalog: Arc, context: SessionContext, /// CSR adjacencies by relationship type; empty unless the engine uses CSR. @@ -1244,6 +1245,33 @@ struct EngineState { key_stats: HashMap, } +/// A small LRU of unoptimized plans. Values are bound after lookup, before optimization. +struct CachedQuery { + text: String, + parameter_types: Vec<(String, u8)>, + query: RustCypherQuery, +} + +// Containers change expression shape; keep them on the uncached path. +fn parameter_types(parameters: &HashMap) -> Option> { + let mut types = parameters + .iter() + .map(|(name, value)| { + let kind = match value { + JsonValue::Null => 0, + JsonValue::Bool(_) => 1, + JsonValue::String(_) => 2, + JsonValue::Number(value) if value.is_i64() => 3, + JsonValue::Number(value) if value.is_f64() => 4, + _ => return None, + }; + Some((name.clone(), kind)) + }) + .collect::>>()?; + types.sort_unstable(); + Some(types) +} + /// Distinct counts for one dataset's node-id and edge-endpoint columns. #[derive(Clone)] struct KeyStatsEntry { @@ -1573,6 +1601,7 @@ fn build_engine_state( } Ok(EngineState { sources, + plans: Mutex::default(), catalog: Arc::new(catalog), context: ctx, csr, @@ -1979,24 +2008,73 @@ impl CypherEngine { fn parse( &self, + py: Python, query: &str, params: Option<&Bound<'_, PyDict>>, state: &EngineState, + use_cache: bool, ) -> PyResult { let mut parameters = HashMap::new(); for (key, value) in params.into_iter().flat_map(|p| p.iter()) { - parameters.insert(key.extract::()?, python_to_json(&value)?); + parameters.insert( + key.extract::()?.to_lowercase(), + python_to_json(&value)?, + ); + } + py.detach(|| self.plan_query(query, parameters, state, use_cache)) + .map_err(graph_error_to_pyerr) + } + + fn plan_query( + &self, + text: &str, + parameters: HashMap, + state: &EngineState, + use_cache: bool, + ) -> Result { + let key = use_cache.then(|| parameter_types(¶meters)).flatten(); + if let Some(key) = &key { + let mut plans = state.plans.lock().unwrap(); + if let Some(index) = plans + .iter() + .position(|entry| entry.text == text && entry.parameter_types == *key) + { + let entry = plans.remove(index).unwrap(); + let query = entry.query.clone().with_parameters(parameters); + plans.push_back(entry); + return Ok(query); + } } - let query = RustCypherQuery::new(query) - .map_err(graph_error_to_pyerr)? + let query = RustCypherQuery::new(text)? .with_parameters(parameters) .with_config(self.config.clone()) .with_fan_out(state.fan_out()); - Ok(if self.use_csr { + let mut query = if self.use_csr { query.with_adjacency(state.adjacency_catalog()) } else { query - }) + }; + if let Some(parameter_types) = key { + // Some operators require literal values during planning. If a template + // cannot be built, ordinary planning preserves their existing behavior. + if query + .prepare_cached_plan(state.catalog.clone()) + .unwrap_or(false) + { + let mut plans = state.plans.lock().unwrap(); + plans + .retain(|entry| entry.text != text || entry.parameter_types != parameter_types); + if plans.len() == 128 { + plans.pop_front(); + } + plans.push_back(CachedQuery { + text: text.to_string(), + parameter_types, + query: query.clone(), + }); + } + } + Ok(query) } } @@ -2072,7 +2150,7 @@ impl CypherEngine { params: Option<&Bound<'_, PyDict>>, ) -> PyResult> { let state = self.snapshot(); - let cypher_query = self.parse(query, params, &state)?; + let cypher_query = self.parse(py, query, params, &state, true)?; let result_batch = RT .block_on( Some(py), @@ -2100,7 +2178,7 @@ impl CypherEngine { use arrow::pyarrow::IntoPyArrow; let state = self.snapshot(); - let cypher_query = self.parse(query, params, &state)?; + let cypher_query = self.parse(py, query, params, &state, true)?; let stream = RT .block_on( Some(py), @@ -2126,7 +2204,7 @@ impl CypherEngine { params: Option<&Bound<'_, PyDict>>, ) -> PyResult { let state = self.snapshot(); - let cypher_query = self.parse(query, params, &state)?; + let cypher_query = self.parse(py, query, params, &state, false)?; RT.block_on( Some(py), cypher_query @@ -2254,7 +2332,7 @@ impl CypherEngine { vector_search: &VectorSearch, ) -> PyResult> { let state = self.snapshot(); - let cypher_query = self.parse(query, None, &state)?; + let cypher_query = self.parse(py, query, None, &state, false)?; let vs = vector_search.inner.clone(); // Execute query to get candidates, then apply vector reranking @@ -2515,41 +2593,113 @@ pub fn register_graph_module(py: Python, parent_module: &Bound<'_, PyModule>) -> mod tests { use super::*; + fn test_state(id: i64) -> EngineState { + let batch = RecordBatch::try_from_iter([( + "id", + Arc::new(arrow_array::Int64Array::from(vec![id])) as arrow_array::ArrayRef, + )]) + .unwrap(); + let provider = + Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch.clone()]]).unwrap()); + let context = SessionContext::new(); + context.register_table("person", provider.clone()).unwrap(); + EngineState { + sources: HashMap::from([("Person".to_string(), EngineSource::Arrow(batch))]), + plans: Mutex::default(), + catalog: Arc::new( + InMemoryCatalog::new() + .with_node_source("Person", Arc::new(DefaultTableSource::new(provider))), + ), + context, + csr: HashMap::new(), + node_index: HashMap::new(), + key_stats: HashMap::new(), + } + } + + #[test] + fn test_plan_cache_reuses_bounds_and_refreshes() { + let engine = CypherEngine { + config: RustGraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(), + state: RwLock::new(Arc::new(test_state(1))), + use_csr: false, + }; + let snapshot = engine.snapshot(); + let text = "MATCH (p:Person) WHERE p.id = $value RETURN p.id"; + for value in [1, 2, 3] { + engine + .plan_query( + text, + HashMap::from([("value".into(), value.into())]), + &snapshot, + true, + ) + .unwrap(); + assert_eq!(snapshot.plans.lock().unwrap().len(), 1); + // The retained template still has its first parameters: later calls + // really used lookup and rebinding rather than replacing the entry. + assert_eq!( + snapshot.plans.lock().unwrap()[0].query.parameters()["value"], + JsonValue::from(1) + ); + } + for value in [JsonValue::Null, JsonValue::from(1.5)] { + engine + .plan_query( + text, + HashMap::from([("value".into(), value)]), + &snapshot, + true, + ) + .unwrap(); + } + assert_eq!(snapshot.plans.lock().unwrap().len(), 3); + engine + .plan_query( + "MATCH (p:Person) WHERE p.id IN $ids RETURN p.id", + HashMap::from([("ids".into(), serde_json::json!([1, 2]))]), + &snapshot, + true, + ) + .unwrap(); + assert_eq!(snapshot.plans.lock().unwrap().len(), 3); + for value in 0..130 { + engine + .plan_query( + &format!("MATCH (p:Person) RETURN p.id AS id LIMIT {value}"), + HashMap::new(), + &snapshot, + true, + ) + .unwrap(); + } + assert_eq!(snapshot.plans.lock().unwrap().len(), 128); + assert!(snapshot + .plans + .lock() + .unwrap() + .iter() + .all(|entry| entry.text != text)); + assert!(engine.try_publish_refresh(&snapshot, test_state(2))); + assert!(engine.snapshot().plans.lock().unwrap().is_empty()); + assert_eq!(snapshot.plans.lock().unwrap().len(), 128); + } + #[test] fn test_refresh_rejects_stale_snapshot() { use arrow_array::Int64Array; use std::sync::mpsc; use std::time::Duration; - fn state(id: i64) -> EngineState { - let batch = RecordBatch::try_from_iter([( - "id", - Arc::new(Int64Array::from(vec![id])) as arrow_array::ArrayRef, - )]) - .unwrap(); - let provider = - Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch.clone()]]).unwrap()); - let context = SessionContext::new(); - context.register_table("person", provider.clone()).unwrap(); - EngineState { - sources: HashMap::from([("Person".to_string(), EngineSource::Arrow(batch))]), - catalog: Arc::new( - InMemoryCatalog::new() - .with_node_source("Person", Arc::new(DefaultTableSource::new(provider))), - ), - context, - csr: HashMap::new(), - node_index: HashMap::new(), - key_stats: HashMap::new(), - } - } - let engine = CypherEngine { config: RustGraphConfig::builder() .node_table("Person", "id") .build() .unwrap(), - state: RwLock::new(Arc::new(state(1))), + state: RwLock::new(Arc::new(test_state(1))), use_csr: false, }; let reader_snapshot = engine.snapshot(); @@ -2561,10 +2711,10 @@ mod tests { let starting_snapshot = engine.snapshot(); started_tx.send(()).unwrap(); published_rx.recv_timeout(Duration::from_secs(5)).unwrap(); - engine.try_publish_refresh(&starting_snapshot, state(2)) + engine.try_publish_refresh(&starting_snapshot, test_state(2)) }); started_rx.recv_timeout(Duration::from_secs(5)).unwrap(); - assert!(engine.try_publish_refresh(&engine.snapshot(), state(3))); + assert!(engine.try_publish_refresh(&engine.snapshot(), test_state(3))); published_tx.send(()).unwrap(); assert!(!older.join().unwrap()); }); diff --git a/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs index 0cf71fc..35e8d00 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/aggregate_ops.rs @@ -96,7 +96,7 @@ impl DataFusionPlanner { let first = aggregates.len(); split_aggregates(&mut rewritten, &mut aggregates); for (column, aggregate) in &aggregates[first..] { - let mut aggregate = super::super::expression::to_df_value_expr(aggregate); + let mut aggregate = self.value_expr(aggregate); // collect() gathers in the order of a preceding `WITH ... ORDER BY`. if let datafusion::logical_expr::Expr::AggregateFunction(f) = &mut aggregate { if f.func.name() == "array_agg" && !f.params.distinct { @@ -105,10 +105,10 @@ impl DataFusionPlanner { } agg_exprs.push(aggregate.alias(column)); } - outer.push(super::super::expression::to_df_value_expr(&rewritten).alias(alias)); + outer.push(self.value_expr(&rewritten).alias(alias)); } else { // Group expressions: use raw expression for grouping, no alias - group_exprs.push(super::super::expression::to_df_value_expr(&p.expression)); + group_exprs.push(self.value_expr(&p.expression)); } } @@ -130,7 +130,7 @@ impl DataFusionPlanner { })); } else if !super::super::expression::contains_aggregate(&p.expression) { // Re-create the expression and apply alias - let expr = super::super::expression::to_df_value_expr(&p.expression); + let expr = self.value_expr(&p.expression); let aliased = if let Some(alias) = &p.alias { expr.alias(alias.to_lowercase()) } else { diff --git a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs index c10b116..9a6c2d3 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/basic_ops.rs @@ -87,7 +87,7 @@ impl DataFusionPlanner { ) -> Result { let input_plan = self.build_operator(ctx, input)?; let predicate = super::super::expression::bind_boolean(predicate, input_plan.schema()); - let expr = super::super::expression::to_df_boolean_expr(&predicate); + let expr = self.boolean_expr(&predicate); LogicalPlanBuilder::from(input_plan) .filter(expr) .map_err(|e| self.plan_error("Failed to build filter", e))? @@ -175,7 +175,7 @@ impl DataFusionPlanner { continue; } let bound = super::super::expression::bind_value(&p.expression, &schema); - let expr = super::super::expression::to_df_value_expr(&bound); + let expr = self.value_expr(&bound); // Apply alias if provided, otherwise use Cypher dot notation // Normalize alias to lowercase for case-insensitive behavior exprs.push(if let Some(alias) = &p.alias { @@ -234,7 +234,7 @@ impl DataFusionPlanner { let bound = super::super::expression::bind_value(&item.expression, input_plan.schema()); SortExpr { - expr: super::super::expression::to_df_value_expr(&bound), + expr: self.value_expr(&bound), asc: matches!(item.direction, crate::ast::SortDirection::Ascending), nulls_first: true, } @@ -289,7 +289,7 @@ impl DataFusionPlanner { // Convert expression to DataFusion Expr let bound = super::super::expression::bind_value(expression, input_plan.schema()); - let df_expr = super::super::expression::to_df_value_expr(&bound); + let df_expr = self.value_expr(&bound); // A list of node ids: unnest under a temporary name, then join each id back to the // node table so the alias carries the node's columns like any matched node. let unnested = match label { diff --git a/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs b/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs index 43b77d2..8858abc 100644 --- a/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/builder/comprehension_ops.rs @@ -16,7 +16,7 @@ use datafusion::logical_expr::{ }; use super::super::analysis::PlanningContext; -use super::super::expression::{bind_boolean, bind_value, to_df_boolean_expr, to_df_value_expr}; +use super::super::expression::{bind_boolean, bind_value}; use crate::ast::{BooleanExpression, ValueExpression}; use crate::case_insensitive::qualify_column; use crate::datafusion_planner::DataFusionPlanner; @@ -54,7 +54,7 @@ impl DataFusionPlanner { .map_err(error("Failed to number comprehension rows"))? .build() .map_err(error("Failed to number comprehension rows"))?; - let list = to_df_value_expr(&bind_value(list, numbered.schema())); + let list = self.value_expr(&bind_value(list, numbered.schema())); let length = int64(coalesce(vec![cardinality(list.clone()), lit(0_u64)])); let last = datafusion::functions::core::expr_fn::greatest(vec![length.clone(), lit(1_i64)]); let positions = gen_series(lit(1_i64), last, lit(1_i64)); @@ -130,9 +130,9 @@ impl DataFusionPlanner { // Padding rows (position past the list's end) are never collected. let mut kept = col(&pos).lt_eq(length); if let Some(keep) = keep { - kept = kept.and(to_df_boolean_expr(&bind_boolean(keep, elements.schema()))); + kept = kept.and(self.boolean_expr(&bind_boolean(keep, elements.schema()))); } - let value = to_df_value_expr(&bind_value(map, elements.schema())); + let value = self.value_expr(&bind_value(map, elements.schema())); let collected = array_agg(value) .filter(kept) .order_by(vec![col(&pos).sort(true, false)]) diff --git a/crates/lance-graph/src/datafusion_planner/expression.rs b/crates/lance-graph/src/datafusion_planner/expression.rs index 18545b6..1ec7894 100644 --- a/crates/lance-graph/src/datafusion_planner/expression.rs +++ b/crates/lance-graph/src/datafusion_planner/expression.rs @@ -19,6 +19,43 @@ use datafusion_functions_aggregate::min_max::max; use datafusion_functions_aggregate::min_max::min; use datafusion_functions_aggregate::sum::sum; +impl super::DataFusionPlanner { + pub(crate) fn value_expr(&self, value: &ValueExpression) -> Expr { + self.type_parameters(to_df_value_expr(value)) + } + + pub(crate) fn boolean_expr(&self, value: &BooleanExpression) -> Expr { + self.type_parameters(to_df_boolean_expr(value)) + } + + fn type_parameters(&self, expression: Expr) -> Expr { + use datafusion::common::tree_node::{Transformed, TreeNode}; + if self.parameter_types.is_empty() { + return expression; + } + expression + .transform_up(|expression| { + if let Expr::Placeholder(mut parameter) = expression { + parameter.field = self + .parameter_types + .get(parameter.id.trim_start_matches('$')) + .map(|data_type| { + std::sync::Arc::new(arrow_schema::Field::new( + "", + data_type.clone(), + true, + )) + }); + Ok(Transformed::yes(Expr::Placeholder(parameter))) + } else { + Ok(Transformed::no(expression)) + } + }) + .expect("typing placeholders does not fail") + .data + } +} + /// Helper function to create LIKE expressions with consistent settings fn create_like_expr(expression: &ValueExpression, pattern: &str, case_insensitive: bool) -> Expr { Expr::Like(datafusion::logical_expr::Like { @@ -131,10 +168,7 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { } VE::Literal(PV::Variable(v)) => col(v.to_lowercase()), VE::Literal(PV::Parameter(name)) => { - panic!( - "Parameter ${} should have been substituted during semantic analysis", - name - ); + datafusion::logical_expr::placeholder(format!("${}", name.to_lowercase())) } VE::Literal(PV::Property(prop)) => { // Create qualified column name: variable__property (lowercase for case-insensitivity) @@ -411,10 +445,7 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr { lit(scalar) } VE::Parameter(name) => { - panic!( - "Parameter ${} should have been substituted during semantic analysis", - name - ); + datafusion::logical_expr::placeholder(format!("${}", name.to_lowercase())) } VE::Boolean(b) => to_df_boolean_expr(b), VE::Map(entries) if entries.is_empty() => lit(datafusion::common::ScalarValue::Struct( diff --git a/crates/lance-graph/src/datafusion_planner/join_ops.rs b/crates/lance-graph/src/datafusion_planner/join_ops.rs index da9ed1c..61282b2 100644 --- a/crates/lance-graph/src/datafusion_planner/join_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/join_ops.rs @@ -59,9 +59,7 @@ impl DataFusionPlanner { // Apply target property filters (e.g., (b {age: 30})) for (k, v) in target_properties.iter() { - let lit_expr = super::expression::to_df_value_expr( - &crate::ast::ValueExpression::Literal(v.clone()), - ); + let lit_expr = self.value_expr(&crate::ast::ValueExpression::Literal(v.clone())); let filter_expr = Expr::BinaryExpr(BinaryExpr { left: Box::new(col(k.to_lowercase())), op: Operator::Eq, diff --git a/crates/lance-graph/src/datafusion_planner/mod.rs b/crates/lance-graph/src/datafusion_planner/mod.rs index 6c8d8ab..8c58f29 100644 --- a/crates/lance-graph/src/datafusion_planner/mod.rs +++ b/crates/lance-graph/src/datafusion_planner/mod.rs @@ -47,6 +47,7 @@ pub struct DataFusionPlanner { pub(crate) config: GraphConfig, pub(crate) catalog: Option>, pub(crate) adjacency: Option, + pub(crate) parameter_types: std::collections::HashMap, } impl DataFusionPlanner { @@ -55,6 +56,7 @@ impl DataFusionPlanner { config, catalog: None, adjacency: None, + parameter_types: Default::default(), } } @@ -63,6 +65,7 @@ impl DataFusionPlanner { config, catalog: Some(catalog), adjacency: None, + parameter_types: Default::default(), } } diff --git a/crates/lance-graph/src/datafusion_planner/scan_ops.rs b/crates/lance-graph/src/datafusion_planner/scan_ops.rs index 33434f8..5a5466f 100644 --- a/crates/lance-graph/src/datafusion_planner/scan_ops.rs +++ b/crates/lance-graph/src/datafusion_planner/scan_ops.rs @@ -44,9 +44,8 @@ impl DataFusionPlanner { let filter_exprs: Vec = properties .iter() .map(|(k, v)| { - let lit_expr = super::expression::to_df_value_expr( - &crate::ast::ValueExpression::Literal(v.clone()), - ); + let lit_expr = + self.value_expr(&crate::ast::ValueExpression::Literal(v.clone())); Expr::BinaryExpr(BinaryExpr { left: Box::new(col(k)), op: Operator::Eq, @@ -140,9 +139,7 @@ impl DataFusionPlanner { // Apply relationship property filters (e.g., -[r {since: 2020}]->) for (k, v) in relationship_properties.iter() { - let lit_expr = super::expression::to_df_value_expr( - &crate::ast::ValueExpression::Literal(v.clone()), - ); + let lit_expr = self.value_expr(&crate::ast::ValueExpression::Literal(v.clone())); let filter_expr = Expr::BinaryExpr(BinaryExpr { left: Box::new(col(k)), op: Operator::Eq, diff --git a/crates/lance-graph/src/query.rs b/crates/lance-graph/src/query.rs index 093248b..d8c3eb1 100644 --- a/crates/lance-graph/src/query.rs +++ b/crates/lance-graph/src/query.rs @@ -129,7 +129,24 @@ pub struct CypherQuery { adjacency: Option, /// Optional mean fan-out per relationship type and direction; see [`FanOut`] fan_out: Option, + /// An unoptimized template belonging to the caller's catalog snapshot. + cached_plan: Option>, } +struct CachedPlan { + graph: crate::logical_plan::LogicalOperator, + datafusion: datafusion::logical_expr::LogicalPlan, + catalog: Arc, + parameter_types: HashMap, +} + +impl std::fmt::Debug for CachedPlan { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CachedPlan") + .field("parameter_types", &self.parameter_types) + .finish_non_exhaustive() + } +} + impl CypherQuery { /// Create a new Cypher query from a query string pub fn new(query: &str) -> Result { @@ -142,11 +159,13 @@ impl CypherQuery { parameters: HashMap::new(), adjacency: None, fan_out: None, + cached_plan: None, }) } /// Set the graph configuration for this query pub fn with_config(mut self, config: GraphConfig) -> Self { + self.cached_plan = None; self.config = Some(config); self } @@ -157,6 +176,7 @@ impl CypherQuery { /// [`crate::csr_expand::CsrQueryPlanner`]. Each adjacency must be built from the /// same relationship data the context serves, or results will be wrong. pub fn with_adjacency(mut self, adjacency: crate::csr_expand::AdjacencyCatalog) -> Self { + self.cached_plan = None; self.adjacency = Some(adjacency); self } @@ -164,6 +184,7 @@ impl CypherQuery { /// Mean fan-out per relationship type and direction, so a pattern that continues /// earlier rows can start from its cheaper bound end (see [`crate::logical_plan::FanOut`]). pub fn with_fan_out(mut self, fan_out: crate::logical_plan::FanOut) -> Self { + self.cached_plan = None; self.fan_out = Some(fan_out); self } @@ -207,6 +228,71 @@ impl CypherQuery { &self.parameters } + /// Build a snapshot-local scalar-parameter template for the Python engine's cache. + /// Binding must precede optimization so pushed CSR filters never retain old values. + #[doc(hidden)] + pub fn prepare_cached_plan( + &mut self, + catalog: Arc, + ) -> Result { + let parameter_counts = self + .ast + .limit + .iter() + .chain(self.ast.skip.iter()) + .chain( + self.ast + .with_clause + .iter() + .filter_map(|clause| clause.limit.as_ref()), + ) + .chain( + self.ast + .further_with_stages + .iter() + .filter_map(|stage| stage.with_clause.limit.as_ref()), + ) + .any(|count| matches!(count, crate::ast::RowCount::Parameter(_))); + if parameter_counts || self.scalar_parameters().is_none() { + return Ok(false); + } + let (graph, datafusion) = self.build_logical_plans(catalog.clone(), true)?; + self.cached_plan = Some(Arc::new(CachedPlan { + graph, + datafusion, + catalog, + parameter_types: self + .scalar_parameters() + .unwrap() + .into_iter() + .map(|(name, value)| (name, value.data_type())) + .collect(), + })); + Ok(true) + } + + fn scalar_parameters(&self) -> Option> { + use datafusion::common::ScalarValue; + self.parameters + .iter() + .map(|(name, value)| { + let value = match value { + serde_json::Value::Null => ScalarValue::Null, + serde_json::Value::Bool(value) => ScalarValue::Boolean(Some(*value)), + serde_json::Value::String(value) => ScalarValue::Utf8(Some(value.clone())), + serde_json::Value::Number(value) if value.is_i64() => { + ScalarValue::Int64(value.as_i64()) + } + serde_json::Value::Number(value) if value.is_f64() => { + ScalarValue::Float64(value.as_f64()) + } + _ => return None, + }; + Some((name.clone(), value)) + }) + .collect() + } + /// Get the required config, returning an error if not set fn require_config(&self) -> Result<&GraphConfig> { self.config.as_ref().ok_or_else(|| GraphError::ConfigError { @@ -876,6 +962,37 @@ impl CypherQuery { ) -> Result<( crate::logical_plan::LogicalOperator, datafusion::logical_expr::LogicalPlan, + )> { + if let Some(template) = &self.cached_plan { + if let Some(parameters) = self.scalar_parameters() { + let same_types = parameters.len() == template.parameter_types.len() + && parameters.iter().all(|(name, value)| { + template.parameter_types.get(name) == Some(&value.data_type()) + }); + if !Arc::ptr_eq(&catalog, &template.catalog) || !same_types { + return self.build_logical_plans(catalog, false); + } + let plan = template + .datafusion + .clone() + .with_param_values(parameters) + .map_err(|error| GraphError::PlanError { + message: format!("Failed to bind query parameters: {error}"), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + return Ok((template.graph.clone(), plan)); + } + } + self.build_logical_plans(catalog, false) + } + + fn build_logical_plans( + &self, + catalog: Arc, + parameterized: bool, + ) -> Result<( + crate::logical_plan::LogicalOperator, + datafusion::logical_expr::LogicalPlan, )> { use crate::datafusion_planner::{DataFusionPlanner, GraphPhysicalPlanner}; use crate::semantic::SemanticAnalyzer; @@ -894,12 +1011,24 @@ impl CypherQuery { // Phase 2: Graph Logical Plan let mut logical_planner = LogicalPlanner::new(config).with_fan_out(self.fan_out.as_ref()); - // Use the transformed AST (with parameters substituted) - let logical_plan = logical_planner.plan(&semantic.ast)?; + let ast = if parameterized { + &self.ast + } else { + &semantic.ast + }; + let logical_plan = logical_planner.plan(ast)?; // Phase 3: DataFusion Logical Plan let mut df_planner = DataFusionPlanner::with_catalog(config.clone(), catalog); df_planner.adjacency = self.adjacency.clone(); + if parameterized { + df_planner.parameter_types = self + .scalar_parameters() + .unwrap_or_default() + .into_iter() + .map(|(name, value)| (name, value.data_type())) + .collect(); + } let df_logical_plan = df_planner.plan(&logical_plan)?; Ok((logical_plan, df_logical_plan)) @@ -1325,6 +1454,7 @@ impl CypherQueryBuilder { parameters: self.parameters, adjacency: None, fan_out: None, + cached_plan: None, }; Ok(query) @@ -1360,6 +1490,266 @@ mod tests { assert!(query.parameters().contains_key("maxage")); } + #[tokio::test] + async fn cached_templates_rebind_scalars_and_preserve_schema() { + use arrow_array::{ArrayRef, Int64Array, StringArray}; + let config = GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + let datasets = HashMap::from([( + "Person".into(), + RecordBatch::try_from_iter([ + ("id", Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef), + ( + "name", + Arc::new(StringArray::from(vec!["Ada", "Bob", "Cam"])) as ArrayRef, + ), + ]) + .unwrap(), + )]); + let cases = [ + ( + "MATCH (p:Person) WHERE p.id >= $value RETURN p.id ORDER BY p.id", + vec![serde_json::json!(1), serde_json::json!(3)], + ), + ( + "MATCH (p:Person) RETURN p.id + $value AS adjusted ORDER BY adjusted", + vec![serde_json::json!(1.5), serde_json::json!(8.5)], + ), + ( + "MATCH (p:Person) WHERE p.name = $value RETURN p.id", + vec![serde_json::json!("Ada"), serde_json::json!("Cam")], + ), + ( + "MATCH (p:Person) RETURN $value AS value", + vec![serde_json::json!(true), serde_json::json!(false)], + ), + ( + "MATCH (p:Person) RETURN $value AS value", + vec![serde_json::Value::Null], + ), + ( + "MATCH (p:Person) RETURN $value", + vec![serde_json::json!(1), serde_json::json!(2)], + ), + ( + "MATCH (p:Person) RETURN p.id + $value", + vec![serde_json::json!(1), serde_json::json!(2)], + ), + ( + "MATCH (p:Person) RETURN p.id AS id, [x IN range(1, $value) | x + p.id] AS values ORDER BY id", + vec![serde_json::json!(1), serde_json::json!(3)], + ), + ]; + for (text, values) in cases { + let mut template = CypherQuery::new(text) + .unwrap() + .with_config(config.clone()) + .with_parameter("value", values[0].clone()); + let (catalog, context) = template + .build_catalog_and_context_from_datasets(datasets.clone()) + .await + .unwrap(); + let catalog: Arc = Arc::new(catalog); + assert!( + template.prepare_cached_plan(catalog.clone()).unwrap(), + "{text}" + ); + for value in values { + let cached = template.clone().with_parameter("value", value.clone()); + let fresh = CypherQuery::new(text) + .unwrap() + .with_config(config.clone()) + .with_parameter("value", value); + let actual = cached + .execute_with_catalog_and_context(catalog.clone(), context.clone()) + .await + .unwrap(); + let expected = fresh + .execute_with_catalog_and_context(catalog.clone(), context.clone()) + .await + .unwrap(); + assert_eq!(actual, expected, "{text}"); + } + } + } + + #[tokio::test] + async fn cached_templates_respect_catalog_and_parameter_types() { + use arrow_array::{ArrayRef, Int64Array}; + let config = GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + let data = |id| { + HashMap::from([( + "Person".into(), + RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from(vec![id])) as ArrayRef, + )]) + .unwrap(), + )]) + }; + let mut query = CypherQuery::new("MATCH (p:Person) WHERE p.id = $value RETURN p.id") + .unwrap() + .with_config(config.clone()) + .with_parameter("value", 1); + let (catalog, _) = query + .build_catalog_and_context_from_datasets(data(1)) + .await + .unwrap(); + assert!(query.prepare_cached_plan(Arc::new(catalog)).unwrap()); + // A different catalog must never serve rows from the prepared snapshot. + assert_eq!(query.execute(data(2), None).await.unwrap().num_rows(), 0); + assert!(query + .clone() + .with_parameter("value", "wrong") + .execute(data(1), None) + .await + .unwrap_err() + .to_string() + .contains("Type mismatch")); + assert!(query.with_config(config).cached_plan.is_none()); + let mut containers = CypherQuery::new("MATCH (p:Person) WHERE p.id IN $ids RETURN p.id") + .unwrap() + .with_parameter("ids", serde_json::json!([1, 2])); + assert!(!containers + .prepare_cached_plan(Arc::new(crate::InMemoryCatalog::new())) + .unwrap()); + } + + #[test] + fn parameterized_row_counts_bypass_templates() { + for text in [ + "MATCH (p:Person) RETURN p.id LIMIT $count", + "MATCH (p:Person) RETURN p.id SKIP $count", + "MATCH (p:Person) WITH p LIMIT $count RETURN p.id", + "MATCH (p:Person) WITH p WITH p LIMIT $count RETURN p.id", + ] { + let mut query = CypherQuery::new(text).unwrap().with_parameter("count", 1); + assert!( + !query + .prepare_cached_plan(Arc::new(crate::InMemoryCatalog::new())) + .unwrap(), + "{text}" + ); + assert!(query.cached_plan.is_none()); + } + } + + // Opt-in diagnostic: reports stage means for a fixed local scan, with no timing + // assertions. Physical planning calls the query planner directly, avoiding the + // extra optimize() performed by SessionState::create_physical_plan(). + #[tokio::test] + #[ignore = "local stage timing; run with --ignored --nocapture"] + async fn profile_query_planning_stages() { + use crate::datafusion_planner::{DataFusionPlanner, GraphPhysicalPlanner}; + use crate::semantic::SemanticAnalyzer; + use arrow_array::{ArrayRef, Int64Array}; + use std::time::Instant; + let text = "MATCH (p:Person) WHERE p.id >= $minimum RETURN p.id AS id ORDER BY id"; + let config = GraphConfig::builder() + .node_table("Person", "id") + .build() + .unwrap(); + let parameters = HashMap::from([("minimum".into(), serde_json::json!(10))]); + let mut template = CypherQuery::new(text) + .unwrap() + .with_config(config.clone()) + .with_parameters(parameters.clone()); + let data = HashMap::from([( + "Person".into(), + RecordBatch::try_from_iter([( + "id", + Arc::new(Int64Array::from_iter_values(0..1000)) as ArrayRef, + )]) + .unwrap(), + )]); + let (catalog, context) = template + .build_catalog_and_context_from_datasets(data) + .await + .unwrap(); + let catalog: Arc = Arc::new(catalog); + template.prepare_cached_plan(catalog.clone()).unwrap(); + let state = context.state(); + let mut totals = [std::time::Duration::ZERO; 8]; + for _ in 0..100 { + let start = Instant::now(); + let ast = parse_cypher_query(text).unwrap(); + totals[0] += start.elapsed(); + let start = Instant::now(); + let semantic = SemanticAnalyzer::new(config.clone()) + .with_catalog(catalog.clone()) + .analyze(&ast, ¶meters) + .unwrap(); + assert!(semantic.errors.is_empty()); + totals[1] += start.elapsed(); + let start = Instant::now(); + let graph = LogicalPlanner::new(&config).plan(&semantic.ast).unwrap(); + totals[2] += start.elapsed(); + let start = Instant::now(); + let logical = DataFusionPlanner::with_catalog(config.clone(), catalog.clone()) + .plan(&graph) + .unwrap(); + totals[3] += start.elapsed(); + let start = Instant::now(); + let optimized = state.optimize(&logical).unwrap(); + totals[4] += start.elapsed(); + let start = Instant::now(); + let physical = state + .query_planner() + .create_physical_plan(&optimized, &state) + .await + .unwrap(); + totals[5] += start.elapsed(); + let start = Instant::now(); + let batches = datafusion::physical_plan::collect(physical, context.task_ctx()) + .await + .unwrap(); + totals[6] += start.elapsed(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + 990 + ); + let checksum: i64 = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + }) + .sum(); + assert_eq!(checksum, (10 + 999) * 990 / 2); + let start = Instant::now(); + let _ = template.create_logical_plans(catalog.clone()).unwrap(); + totals[7] += start.elapsed(); + } + for (name, duration) in [ + "parse", + "semantic", + "graph_logical", + "datafusion_logical", + "optimize", + "physical", + "execute", + "cached_bind", + ] + .into_iter() + .zip(totals) + { + println!( + "{name}: {:.3} us mean (100 iterations, 1000 input rows, 990 result rows)", + duration.as_secs_f64() * 1_000_000.0 / 100.0 + ); + } + } + #[test] fn test_query_builder() { let config = GraphConfig::builder() diff --git a/docs/api-reference.md b/docs/api-reference.md index ee94517..67d99af 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -116,7 +116,11 @@ CypherEngine(config: GraphConfig, datasets: dict, use_csr: bool = False) ``` Registers `datasets` once and reuses the catalog, DataFusion session and (with `use_csr=True`) CSR index -across queries. Safe to share between threads. +across queries. Safe to share between threads. `execute()` and `execute_batches()` reuse up to 128 +unoptimized query templates per snapshot, keyed by exact text and scalar parameter names/types. Values are +bound before optimization on every call. Containers, parameterized `SKIP`/`LIMIT` and unsupported template +shapes use ordinary planning. `explain()` always plans directly; successful `refresh()` clears the cache. +See [Repeated queries](lance.md#repeated-queries) for the scope and limits. | Parameter | Description | | --- | --- | diff --git a/docs/lance.md b/docs/lance.md index 23037af..4e09671 100644 --- a/docs/lance.md +++ b/docs/lance.md @@ -82,6 +82,19 @@ What `refresh()` guarantees: A refresh reopens each dataset by its URI with the same `storage_options` it was opened with, so object-store credentials carry over. +## Repeated queries + +`execute()` and `execute_batches()` keep up to 128 query templates per engine snapshot. Reuse the exact +query text with `$name` parameters to reuse parsing, semantic analysis, graph logical planning and +DataFusion logical planning. Every execution binds its own values before optimization and physical planning; +rows and results are never cached. Query text and parameter names/types identify each template, with nulls +separate from non-null values. Least recently used templates are evicted at the limit. + +Lists, maps and parameterized `SKIP`/`LIMIT` counts use ordinary planning because their shape or value can +change the plan. Queries whose operators cannot produce a reusable template fall back to ordinary planning. `explain()` plans the current values directly. A successful `refresh()` starts an empty cache for the +new snapshot, including a refresh whose datasets did not change; active readers retain their existing plans. +These rules apply with and without CSR. They do not promise a speedup for a particular workload. + ## Streaming large results `execute()` returns one table, held in memory. For big outputs, stream: diff --git a/python/python/tests/test_plan_cache.py b/python/python/tests/test_plan_cache.py new file mode 100644 index 0000000..30aa404 --- /dev/null +++ b/python/python/tests/test_plan_cache.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The lance-csr Authors + +"""Repeated execution must bind current values before filters reach CSR scans.""" + +from concurrent.futures import ThreadPoolExecutor + +import pyarrow as pa +import pytest +from lance_csr import CypherEngine, GraphConfig + + +@pytest.fixture(params=[False, True]) +def engine(request): + config = ( + GraphConfig.builder() + .node_table("Person", "id") + .rel_table("KNOWS", "Person", "Person", "src", "dst") + .build() + ) + return CypherEngine( + config, + { + "Person": pa.table({"id": [1, 2, 3], "age": [20, 30, 40]}), + "KNOWS": pa.table({"src": [1, 1, 2], "dst": [2, 3, 3]}), + }, + use_csr=request.param, + ) + + +def test_cached_target_filters_bind_each_value(engine): + query = ( + "MATCH (a:Person)-[:KNOWS]->(b:Person) " + "WHERE a.id = $start AND b.age >= $minimum RETURN b.id AS id ORDER BY id" + ) + for minimum, ids in [(20, [2, 3]), (35, [3]), (45, []), (20, [2, 3])]: + params = {"start": 1, "minimum": minimum} + expected = [{"id": value} for value in ids] + assert engine.execute(query, params).to_pylist() == expected + assert engine.execute_batches(query, params).read_all().to_pylist() == expected + + +def test_cached_parameter_types_and_container_fallback(engine): + query = "MATCH (p:Person) WHERE p.id = 1 RETURN $value AS value" + for value in [1, 2, None, 1.5, "Ada", "Bob", True, False, [1, 2], [3], {"x": 1}]: + assert engine.execute(query, {"value": value}).to_pylist() == [{"value": value}] + membership = "MATCH (p:Person) WHERE p.id IN $ids RETURN p.id AS id ORDER BY id" + for ids in [[1, 2], [3], []]: + assert engine.execute(membership, {"ids": ids}).to_pylist() == [ + {"id": value} for value in ids + ] + + +def test_unaliased_parameters_keep_existing_column_names(engine): + query = "MATCH (p:Person) WHERE p.id = 1 RETURN $value" + first = engine.execute(query, {"value": 1}) + second = engine.execute(query, {"value": 2}) + assert first.column_names == ["expr"] + assert second.column_names == ["expr"] + assert first.column(0).to_pylist() == [1] + assert second.column(0).to_pylist() == [2] + + +def test_cached_queries_keep_concurrent_parameters_separate(engine): + query = "MATCH (p:Person) WHERE p.id = $id RETURN p.age AS age" + engine.execute(query, {"id": 1}) + + def execute(person_id): + return engine.execute(query, {"id": person_id}).to_pylist() + + ids = [1, 2, 3] * 8 + with ThreadPoolExecutor(max_workers=4) as pool: + results = list(pool.map(execute, ids)) + assert results == [[{"age": 10 * (person_id + 1)}] for person_id in ids] + + +def test_cached_query_still_rejects_missing_or_wrong_types(engine): + query = "MATCH (p:Person) WHERE p.id = $id RETURN p.age AS age" + engine.execute(query, {"id": 1}) + with pytest.raises(ValueError, match="[Pp]arameter"): + engine.execute(query) + with pytest.raises(ValueError, match="Type mismatch"): + engine.execute(query, {"id": "wrong"}) + + +@pytest.mark.parametrize( + "query", + [ + "MATCH (p:Person) RETURN p.id AS id ORDER BY id LIMIT $count", + "MATCH (p:Person) WITH p ORDER BY p.id LIMIT $count RETURN p.id AS id", + "MATCH (p:Person) WITH p WITH p ORDER BY p.id LIMIT $count RETURN p.id AS id", + ], +) +def test_parameterized_limits_replan_each_count(engine, query): + for count in [1, 3, 0, 2]: + expected = [{"id": value} for value in range(1, count + 1)] + assert engine.execute(query, {"count": count}).to_pylist() == expected + assert ( + engine.execute_batches(query, {"count": count}).read_all().to_pylist() + == expected + )