From a3cba88db2507b272339beee659618675c08b66c Mon Sep 17 00:00:00 2001 From: JoshuaTang <1240604020@qq.com> Date: Sun, 2 Nov 2025 11:49:56 -0800 Subject: [PATCH] feat: add an explain method to show generated plans --- rust/lance-graph/src/query.rs | 389 ++++++++++++++---- rust/lance-graph/tests/test_explain_output.rs | 323 +++++++++++++++ 2 files changed, 625 insertions(+), 87 deletions(-) create mode 100644 rust/lance-graph/tests/test_explain_output.rs diff --git a/rust/lance-graph/src/query.rs b/rust/lance-graph/src/query.rs index 76e869ba2..077c3e44f 100644 --- a/rust/lance-graph/src/query.rs +++ b/rust/lance-graph/src/query.rs @@ -84,6 +84,14 @@ impl CypherQuery { &self.parameters } + /// 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 { + message: "Graph configuration is required for query execution".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }) + } + /// Execute using the DataFusion planner with in-memory datasets /// /// # Overview @@ -120,46 +128,12 @@ impl CypherQuery { &self, datasets: HashMap, ) -> Result { - use crate::source_catalog::InMemoryCatalog; - use datafusion::datasource::{DefaultTableSource, MemTable}; - use datafusion::execution::context::SessionContext; use std::sync::Arc; - if datasets.is_empty() { - return Err(GraphError::ConfigError { - message: "No input datasets provided".to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - }); - } - - // Create session context and catalog, register tables in both - let ctx = SessionContext::new(); - let mut catalog: InMemoryCatalog = InMemoryCatalog::new(); - - for (name, batch) in &datasets { - let mem_table = Arc::new( - MemTable::try_new(batch.schema(), vec![vec![batch.clone()]]).map_err(|e| { - GraphError::PlanError { - message: format!("Failed to create MemTable for {}: {}", name, e), - location: snafu::Location::new(file!(), line!(), column!()), - } - })?, - ); - - let table_source = Arc::new(DefaultTableSource::new(mem_table.clone())); - - // Register as both node and relationship source (planner will use whichever is appropriate) - catalog = catalog - .with_node_source(name, table_source.clone()) - .with_relationship_source(name, table_source.clone()); - - // Register in session context for execution (using the same MemTable instance) - ctx.register_table(name, mem_table) - .map_err(|e| GraphError::PlanError { - message: format!("Failed to register table {}: {}", name, e), - location: snafu::Location::new(file!(), line!(), column!()), - })?; - } + // Build catalog and context from datasets + let (catalog, ctx) = self + .build_catalog_and_context_from_datasets(datasets) + .await?; // Delegate to common execution logic self.execute_with_catalog_and_context(Arc::new(catalog), ctx) @@ -222,14 +196,7 @@ impl CypherQuery { use datafusion::datasource::DefaultTableSource; use std::sync::Arc; - // Require a config - let config = self - .config - .as_ref() - .ok_or_else(|| GraphError::ConfigError { - message: "Graph configuration is required for query execution".to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - })?; + let config = self.require_config()?; // Build catalog by querying SessionContext for table providers let mut catalog = InMemoryCatalog::new(); @@ -314,35 +281,12 @@ impl CypherQuery { catalog: std::sync::Arc, ctx: datafusion::execution::context::SessionContext, ) -> Result { - use crate::datafusion_planner::{DataFusionPlanner, GraphPhysicalPlanner}; - use crate::semantic::SemanticAnalyzer; use arrow::compute::concat_batches; - // Require a config for DataFusion execution - let config = self - .config - .as_ref() - .ok_or_else(|| GraphError::ConfigError { - message: "Graph configuration is required for DataFusion execution".to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - })?; + // Create logical plans (phases 1-3) + let (_logical_plan, df_logical_plan) = self.create_logical_plans(catalog)?; - // Phase 1: Semantic Analysis - let mut analyzer = SemanticAnalyzer::new(config.clone()); - analyzer.analyze(&self.ast)?; - - // Phase 2: Logical Planning - let mut logical_planner = LogicalPlanner::new(); - let logical_plan = logical_planner.plan(&self.ast)?; - - // Phase 3: DataFusion Logical Planning - // Convert graph logical plan to DataFusion logical plan - let df_planner = DataFusionPlanner::with_catalog(config.clone(), catalog); - let df_logical_plan = df_planner.plan(&logical_plan)?; - - // Phase 4: Physical Planning and Execution - // DataFusion optimizes the logical plan, creates a physical execution plan, - // and executes it against the pre-configured SessionContext + // Execute the DataFusion plan (phase 4) let df = ctx .execute_logical_plan(df_logical_plan) .await @@ -374,6 +318,289 @@ impl CypherQuery { }) } + /// Explain the query execution plan using in-memory datasets + /// + /// Returns a formatted string showing the query execution plan at different stages: + /// - Graph Logical Plan (graph-specific operators) + /// - DataFusion Logical Plan (optimized relational plan) + /// - DataFusion Physical Plan (execution plan with optimizations) + /// + /// This is useful for understanding query performance, debugging, and optimization. + /// + /// # Arguments + /// * `datasets` - HashMap of table name to RecordBatch (nodes and relationships) + /// + /// # Returns + /// A formatted string containing the execution plan at multiple levels + /// + /// # Errors + /// Returns error if planning fails + /// + /// # Example + /// ```ignore + /// use std::collections::HashMap; + /// use arrow::record_batch::RecordBatch; + /// use lance_graph::query::CypherQuery; + /// + /// // Create in-memory datasets + /// let mut datasets = HashMap::new(); + /// datasets.insert("Person".to_string(), person_batch); + /// datasets.insert("KNOWS".to_string(), knows_batch); + /// + /// let query = CypherQuery::parse("MATCH (p:Person) WHERE p.age > 30 RETURN p.name")? + /// .with_config(config); + /// + /// let plan = query.explain_datafusion(datasets).await?; + /// println!("{}", plan); + /// ``` + pub async fn explain_datafusion( + &self, + datasets: HashMap, + ) -> Result { + use std::sync::Arc; + + // Build catalog and context from datasets + let (catalog, ctx) = self + .build_catalog_and_context_from_datasets(datasets) + .await?; + + // Delegate to the internal explain method + self.explain_internal(Arc::new(catalog), ctx).await + } + + /// Helper to build catalog and context from in-memory datasets + async fn build_catalog_and_context_from_datasets( + &self, + datasets: HashMap, + ) -> Result<( + crate::source_catalog::InMemoryCatalog, + datafusion::execution::context::SessionContext, + )> { + use crate::source_catalog::InMemoryCatalog; + use datafusion::datasource::{DefaultTableSource, MemTable}; + use datafusion::execution::context::SessionContext; + use std::sync::Arc; + + if datasets.is_empty() { + return Err(GraphError::ConfigError { + message: "No input datasets provided".to_string(), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + + // Create session context and catalog + let ctx = SessionContext::new(); + let mut catalog = InMemoryCatalog::new(); + + // Register all datasets as tables + for (name, batch) in &datasets { + let mem_table = Arc::new( + MemTable::try_new(batch.schema(), vec![vec![batch.clone()]]).map_err(|e| { + GraphError::PlanError { + message: format!("Failed to create MemTable for {}: {}", name, e), + location: snafu::Location::new(file!(), line!(), column!()), + } + })?, + ); + + // Register in session context for execution + ctx.register_table(name, mem_table.clone()) + .map_err(|e| GraphError::PlanError { + message: format!("Failed to register table {}: {}", name, e), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + + let table_source = Arc::new(DefaultTableSource::new(mem_table)); + + // Register as both node and relationship source + // The planner will use whichever is appropriate based on the query + catalog = catalog + .with_node_source(name, table_source.clone()) + .with_relationship_source(name, table_source); + } + + Ok((catalog, ctx)) + } + + /// Internal helper to explain the query execution plan with explicit catalog and session context + async fn explain_internal( + &self, + catalog: std::sync::Arc, + ctx: datafusion::execution::context::SessionContext, + ) -> Result { + // Create all plans (phases 1-4) + let (logical_plan, df_logical_plan, physical_plan) = + self.create_plans(catalog, &ctx).await?; + + // Format the explain output + self.format_explain_output(&logical_plan, &df_logical_plan, physical_plan.as_ref()) + } + + /// Helper to create logical plans (graph logical, DataFusion logical) + /// + /// This performs phases 1-3 of query execution (semantic analysis, graph logical planning, + /// DataFusion logical planning) without creating the physical plan. + fn create_logical_plans( + &self, + catalog: std::sync::Arc, + ) -> Result<( + crate::logical_plan::LogicalOperator, + datafusion::logical_expr::LogicalPlan, + )> { + use crate::datafusion_planner::{DataFusionPlanner, GraphPhysicalPlanner}; + use crate::semantic::SemanticAnalyzer; + + let config = self.require_config()?; + + // Phase 1: Semantic Analysis + let mut analyzer = SemanticAnalyzer::new(config.clone()); + analyzer.analyze(&self.ast)?; + + // Phase 2: Graph Logical Plan + let mut logical_planner = LogicalPlanner::new(); + let logical_plan = logical_planner.plan(&self.ast)?; + + // Phase 3: DataFusion Logical Plan + let df_planner = DataFusionPlanner::with_catalog(config.clone(), catalog); + let df_logical_plan = df_planner.plan(&logical_plan)?; + + Ok((logical_plan, df_logical_plan)) + } + + /// Helper to create all plans (graph logical, DataFusion logical, physical) + async fn create_plans( + &self, + catalog: std::sync::Arc, + ctx: &datafusion::execution::context::SessionContext, + ) -> Result<( + crate::logical_plan::LogicalOperator, + datafusion::logical_expr::LogicalPlan, + std::sync::Arc, + )> { + // Phases 1-3: Create logical plans + let (logical_plan, df_logical_plan) = self.create_logical_plans(catalog)?; + + // Phase 4: DataFusion Physical Plan + let df = ctx + .execute_logical_plan(df_logical_plan.clone()) + .await + .map_err(|e| GraphError::ExecutionError { + message: format!("Failed to execute DataFusion plan: {}", e), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + + let physical_plan = + df.create_physical_plan() + .await + .map_err(|e| GraphError::ExecutionError { + message: format!("Failed to create physical plan: {}", e), + location: snafu::Location::new(file!(), line!(), column!()), + })?; + + Ok((logical_plan, df_logical_plan, physical_plan)) + } + + /// Format explain output as a table + fn format_explain_output( + &self, + logical_plan: &crate::logical_plan::LogicalOperator, + df_logical_plan: &datafusion::logical_expr::LogicalPlan, + physical_plan: &dyn datafusion::physical_plan::ExecutionPlan, + ) -> Result { + // Format output with query first, then table + let mut output = String::new(); + + // Show Cypher query before the table + output.push_str("Cypher Query:\n"); + output.push_str(&format!(" {}\n\n", self.query_text)); + + // Build table rows (without the query) + let mut rows = vec![]; + + // Row 1: Graph Logical Plan + let graph_plan_str = format!("{:#?}", logical_plan); + rows.push(("graph_logical_plan", graph_plan_str)); + + // Row 2: DataFusion Logical Plan + let df_logical_str = format!("{}", df_logical_plan.display_indent()); + rows.push(("logical_plan", df_logical_str)); + + // Row 3: DataFusion Physical Plan + let df_physical_str = format!( + "{}", + datafusion::physical_plan::displayable(physical_plan).indent(true) + ); + rows.push(("physical_plan", df_physical_str)); + + // Calculate column widths + let plan_type_width = rows.iter().map(|(t, _)| t.len()).max().unwrap_or(10); + let plan_width = rows + .iter() + .map(|(_, p)| p.lines().map(|l| l.len()).max().unwrap_or(0)) + .max() + .unwrap_or(50); + + // Build table + let separator = format!( + "+{}+{}+", + "-".repeat(plan_type_width + 2), + "-".repeat(plan_width + 2) + ); + + output.push_str(&separator); + output.push('\n'); + + // Header + output.push_str(&format!( + "| {: = plan_content.lines().collect(); + if lines.is_empty() { + output.push_str(&format!( + "| {: p, _ => return Ok(None), }; - let cfg = self.config.as_ref().ok_or_else(|| GraphError::PlanError { - message: "Graph configuration is required for execution".to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - })?; + let cfg = self.require_config()?; // Handle single-segment variable-length paths by unrolling ranges (*1..N, capped) if path.segments.len() == 1 { @@ -835,10 +1053,7 @@ impl CypherQuery { let end_alias = seg.end_node.variable.as_deref().unwrap_or(end_label); // Validate mappings - let cfg = self.config.as_ref().ok_or_else(|| GraphError::PlanError { - message: "Graph configuration is required for execution".to_string(), - location: snafu::Location::new(file!(), line!(), column!()), - })?; + let cfg = self.require_config()?; let start_map = cfg .get_node_mapping(start_label) .ok_or_else(|| GraphError::PlanError { diff --git a/rust/lance-graph/tests/test_explain_output.rs b/rust/lance-graph/tests/test_explain_output.rs new file mode 100644 index 000000000..a112927d3 --- /dev/null +++ b/rust/lance-graph/tests/test_explain_output.rs @@ -0,0 +1,323 @@ +use arrow_array::{Int64Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema}; +use lance_graph::config::GraphConfig; +use lance_graph::query::CypherQuery; +use std::collections::HashMap; +use std::sync::Arc; + +/// Helper to create a Person dataset +fn create_person_dataset() -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), + Field::new("age", DataType::Int64, false), + Field::new("city", DataType::Utf8, false), + ])); + + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5])), + Arc::new(StringArray::from(vec![ + "Alice", "Bob", "Carol", "David", "Eve", + ])), + Arc::new(Int64Array::from(vec![28, 34, 29, 42, 31])), + Arc::new(StringArray::from(vec!["NYC", "SF", "LA", "NYC", "SF"])), + ], + ) + .unwrap() +} + +/// Helper to create a KNOWS relationship dataset +fn create_knows_dataset() -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("src_id", DataType::Int64, false), + Field::new("dst_id", DataType::Int64, false), + Field::new("since", DataType::Int64, false), + ])); + + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![1, 1, 2, 3, 4])), + Arc::new(Int64Array::from(vec![2, 3, 3, 4, 5])), + Arc::new(Int64Array::from(vec![2020, 2019, 2021, 2018, 2022])), + ], + ) + .unwrap() +} + +/// Helper to create a WORKS_AT relationship dataset +fn create_works_at_dataset() -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("person_id", DataType::Int64, false), + Field::new("company_id", DataType::Int64, false), + Field::new("role", DataType::Utf8, false), + ])); + + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5])), + Arc::new(Int64Array::from(vec![10, 10, 20, 20, 30])), + Arc::new(StringArray::from(vec![ + "Engineer", "Manager", "Engineer", "Director", "Engineer", + ])), + ], + ) + .unwrap() +} + +/// Helper to create a Company dataset +fn create_company_dataset() -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), + Field::new("industry", DataType::Utf8, false), + ])); + + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![10, 20, 30])), + Arc::new(StringArray::from(vec!["TechCorp", "DataInc", "AILabs"])), + Arc::new(StringArray::from(vec!["Technology", "Data", "AI"])), + ], + ) + .unwrap() +} + +#[tokio::test] +async fn test_explain_simple_node_scan() { + println!("\n=== EXAMPLE 1: Simple Node Scan with Filter ===\n"); + + let config = GraphConfig::builder() + .with_node_label("Person", "id") + .build() + .unwrap(); + + let query = CypherQuery::new("MATCH (p:Person) WHERE p.age > 30 RETURN p.name, p.age") + .unwrap() + .with_config(config); + + let mut datasets = HashMap::new(); + datasets.insert("Person".to_string(), create_person_dataset()); + + let plan = query.explain_datafusion(datasets).await.unwrap(); + println!("{}", plan); + + assert!(plan.contains("Cypher Query:")); + assert!(plan.contains("MATCH (p:Person)")); + assert!(!plan.contains("| cypher_query")); // Query should not be in table +} + +#[tokio::test] +async fn test_explain_one_hop_relationship() { + println!("\n=== EXAMPLE 2: One-Hop Relationship Traversal ===\n"); + + let config = GraphConfig::builder() + .with_node_label("Person", "id") + .with_relationship("KNOWS", "src_id", "dst_id") + .build() + .unwrap(); + + let query = CypherQuery::new( + "MATCH (a:Person)-[:KNOWS]->(b:Person) \ + WHERE a.age > 25 \ + RETURN a.name, b.name, b.city \ + ORDER BY a.name", + ) + .unwrap() + .with_config(config); + + let mut datasets = HashMap::new(); + datasets.insert("Person".to_string(), create_person_dataset()); + datasets.insert("KNOWS".to_string(), create_knows_dataset()); + + let plan = query.explain_datafusion(datasets).await.unwrap(); + println!("{}", plan); + + assert!(plan.contains("| graph_logical_plan")); + assert!(plan.contains("Expand")); + assert!(plan.contains("HashJoin") || plan.contains("Join")); +} + +#[tokio::test] +async fn test_explain_two_hop_path() { + println!("\n=== EXAMPLE 3: Two-Hop Path (Friends of Friends) ===\n"); + + let config = GraphConfig::builder() + .with_node_label("Person", "id") + .with_relationship("KNOWS", "src_id", "dst_id") + .build() + .unwrap(); + + let query = CypherQuery::new( + "MATCH (a:Person)-[:KNOWS]->(b:Person)-[:KNOWS]->(c:Person) \ + WHERE a.city = 'NYC' AND c.age < 35 \ + RETURN a.name AS source, b.name AS intermediate, c.name AS target \ + ORDER BY source, target \ + LIMIT 10", + ) + .unwrap() + .with_config(config); + + let mut datasets = HashMap::new(); + datasets.insert("Person".to_string(), create_person_dataset()); + datasets.insert("KNOWS".to_string(), create_knows_dataset()); + + let plan = query.explain_datafusion(datasets).await.unwrap(); + println!("{}", plan); + + assert!(plan.contains("| physical_plan")); + // Should have multiple joins for two-hop path + let join_count = plan.matches("Join").count(); + assert!( + join_count >= 2, + "Expected at least 2 joins for two-hop path" + ); +} + +#[tokio::test] +async fn test_explain_multi_relationship_types() { + println!("\n=== EXAMPLE 4: Multiple Relationship Types ===\n"); + + let config = GraphConfig::builder() + .with_node_label("Person", "id") + .with_node_label("Company", "id") + .with_relationship("KNOWS", "src_id", "dst_id") + .with_relationship("WORKS_AT", "person_id", "company_id") + .build() + .unwrap(); + + let query = CypherQuery::new( + "MATCH (p:Person)-[:WORKS_AT]->(c:Company) \ + WHERE c.industry = 'Technology' \ + RETURN p.name, c.name, p.age \ + ORDER BY p.age DESC", + ) + .unwrap() + .with_config(config); + + let mut datasets = HashMap::new(); + datasets.insert("Person".to_string(), create_person_dataset()); + datasets.insert("Company".to_string(), create_company_dataset()); + datasets.insert("WORKS_AT".to_string(), create_works_at_dataset()); + + let plan = query.explain_datafusion(datasets).await.unwrap(); + println!("{}", plan); + + assert!(plan.contains("Company")); + assert!(plan.contains("WORKS_AT")); +} + +#[tokio::test] +async fn test_explain_distinct() { + println!("\n=== EXAMPLE 5: DISTINCT Query ===\n"); + + let config = GraphConfig::builder() + .with_node_label("Person", "id") + .with_relationship("KNOWS", "src_id", "dst_id") + .build() + .unwrap(); + + let query = CypherQuery::new( + "MATCH (p:Person)-[:KNOWS]->(f:Person) \ + RETURN DISTINCT p.city", + ) + .unwrap() + .with_config(config); + + let mut datasets = HashMap::new(); + datasets.insert("Person".to_string(), create_person_dataset()); + datasets.insert("KNOWS".to_string(), create_knows_dataset()); + + let plan = query.explain_datafusion(datasets).await.unwrap(); + println!("{}", plan); + + assert!(plan.contains("Distinct") || plan.contains("DISTINCT")); +} + +#[tokio::test] +async fn test_explain_complex_filter() { + println!("\n=== EXAMPLE 6: Complex Boolean Filter ===\n"); + + let config = GraphConfig::builder() + .with_node_label("Person", "id") + .build() + .unwrap(); + + let query = CypherQuery::new( + "MATCH (p:Person) \ + WHERE (p.age > 30 AND p.city = 'NYC') OR (p.age < 30 AND p.city = 'SF') \ + RETURN p.name, p.age, p.city", + ) + .unwrap() + .with_config(config); + + let mut datasets = HashMap::new(); + datasets.insert("Person".to_string(), create_person_dataset()); + + let plan = query.explain_datafusion(datasets).await.unwrap(); + println!("{}", plan); + + assert!(plan.contains("Filter")); + assert!(plan.contains("OR") || plan.contains("And")); +} + +#[tokio::test] +async fn test_explain_with_skip_and_limit() { + println!("\n=== EXAMPLE 7: Pagination with SKIP and LIMIT ===\n"); + + let config = GraphConfig::builder() + .with_node_label("Person", "id") + .build() + .unwrap(); + + let query = CypherQuery::new( + "MATCH (p:Person) \ + RETURN p.name, p.age \ + ORDER BY p.age DESC \ + SKIP 2 \ + LIMIT 3", + ) + .unwrap() + .with_config(config); + + let mut datasets = HashMap::new(); + datasets.insert("Person".to_string(), create_person_dataset()); + + let plan = query.explain_datafusion(datasets).await.unwrap(); + println!("{}", plan); + + assert!(plan.contains("Limit") || plan.contains("LIMIT")); +} + +#[tokio::test] +async fn test_explain_relationship_properties() { + println!("\n=== EXAMPLE 8: Filter on Relationship Properties ===\n"); + + let config = GraphConfig::builder() + .with_node_label("Person", "id") + .with_relationship("KNOWS", "src_id", "dst_id") + .build() + .unwrap(); + + let query = CypherQuery::new( + "MATCH (a:Person)-[r:KNOWS]->(b:Person) \ + WHERE r.since > 2020 \ + RETURN a.name, b.name, r.since \ + ORDER BY r.since", + ) + .unwrap() + .with_config(config); + + let mut datasets = HashMap::new(); + datasets.insert("Person".to_string(), create_person_dataset()); + datasets.insert("KNOWS".to_string(), create_knows_dataset()); + + let plan = query.explain_datafusion(datasets).await.unwrap(); + println!("{}", plan); + + assert!(plan.contains("since") || plan.contains("Filter")); +}