diff --git a/rust/lance-graph/src/datafusion_planner.rs b/rust/lance-graph/src/datafusion_planner.rs index 3db581252..e1cf5d85c 100644 --- a/rust/lance-graph/src/datafusion_planner.rs +++ b/rust/lance-graph/src/datafusion_planner.rs @@ -52,6 +52,18 @@ impl DataFusionPlanner { catalog: Some(catalog), } } + + /// Helper to convert DataFusion builder errors into GraphError::PlanError with context + fn plan_error( + &self, + context: &str, + error: E, + ) -> crate::error::GraphError { + crate::error::GraphError::PlanError { + message: format!("{}: {}", context, error), + location: snafu::Location::new(file!(), line!(), column!()), + } + } } // ============================================================================ @@ -348,11 +360,11 @@ impl DataFusionPlanner { LogicalOperator::Filter { input, predicate } => { let input_plan = self.build_operator(ctx, input)?; let expr = self.to_df_boolean_expr(predicate); - Ok(LogicalPlanBuilder::from(input_plan) + LogicalPlanBuilder::from(input_plan) .filter(expr) - .unwrap() + .map_err(|e| self.plan_error("Failed to build filter", e))? .build() - .unwrap()) + .map_err(|e| self.plan_error("Failed to build plan", e)) } LogicalOperator::Project { input, projections } => { let input_plan = self.build_operator(ctx, input)?; @@ -368,19 +380,19 @@ impl DataFusionPlanner { } }) .collect(); - Ok(LogicalPlanBuilder::from(input_plan) + LogicalPlanBuilder::from(input_plan) .project(exprs) - .unwrap() + .map_err(|e| self.plan_error("Failed to build projection", e))? .build() - .unwrap()) + .map_err(|e| self.plan_error("Failed to build plan", e)) } LogicalOperator::Distinct { input } => { let input_plan = self.build_operator(ctx, input)?; - Ok(LogicalPlanBuilder::from(input_plan) + LogicalPlanBuilder::from(input_plan) .distinct() - .unwrap() + .map_err(|e| self.plan_error("Failed to build distinct", e))? .build() - .unwrap()) + .map_err(|e| self.plan_error("Failed to build plan", e)) } LogicalOperator::Sort { input, sort_items } => { use datafusion::logical_expr::SortExpr; @@ -401,27 +413,27 @@ impl DataFusionPlanner { }) .collect(); - Ok(LogicalPlanBuilder::from(input_plan) + LogicalPlanBuilder::from(input_plan) .sort(sort_exprs) - .unwrap() + .map_err(|e| self.plan_error("Failed to build sort", e))? .build() - .unwrap()) + .map_err(|e| self.plan_error("Failed to build plan", e)) } LogicalOperator::Limit { input, count } => { let input_plan = self.build_operator(ctx, input)?; - Ok(LogicalPlanBuilder::from(input_plan) + LogicalPlanBuilder::from(input_plan) .limit(0, Some((*count) as usize)) - .unwrap() + .map_err(|e| self.plan_error("Failed to build limit", e))? .build() - .unwrap()) + .map_err(|e| self.plan_error("Failed to build plan", e)) } LogicalOperator::Offset { input, offset } => { let input_plan = self.build_operator(ctx, input)?; - Ok(LogicalPlanBuilder::from(input_plan) + LogicalPlanBuilder::from(input_plan) .limit((*offset) as usize, None) - .unwrap() + .map_err(|e| self.plan_error("Failed to build offset", e))? .build() - .unwrap()) + .map_err(|e| self.plan_error("Failed to build plan", e)) } LogicalOperator::Expand { input, @@ -486,10 +498,13 @@ impl DataFusionPlanner { ) -> Result { // Try to use catalog if available if let Some(cat) = &self.catalog { + // Catalog exists - check if label is registered if let Some(source) = cat.node_source(label) { // Get schema before moving source let schema = source.schema(); - let mut builder = LogicalPlanBuilder::scan(label, source, None).unwrap(); + let mut builder = LogicalPlanBuilder::scan(label, source, None).map_err(|e| { + self.plan_error(&format!("Failed to scan node source '{}'", label), e) + })?; // Apply property filters using unqualified names (before aliasing) for (k, v) in properties.iter() { @@ -500,7 +515,9 @@ impl DataFusionPlanner { op: Operator::Eq, right: Box::new(lit_expr), }); - builder = builder.filter(filter_expr).unwrap(); + builder = builder.filter(filter_expr).map_err(|e| { + self.plan_error(&format!("Failed to apply filter on property '{}'", k), e) + })?; } // Create qualified column aliases: variable__property @@ -514,27 +531,37 @@ impl DataFusionPlanner { .collect(); // Add projection with qualified aliases - builder = builder.project(qualified_exprs).unwrap(); + builder = builder + .project(qualified_exprs) + .map_err(|e| self.plan_error("Failed to project qualified columns", e))?; - return Ok(builder.build().unwrap()); + return builder + .build() + .map_err(|e| self.plan_error("Failed to build scan plan", e)); + } else { + // Catalog exists but label not found - fail fast + return Err(crate::error::GraphError::ConfigError { + message: format!( + "Node label '{}' not found in catalog. \ + Ensure the label is registered in your GraphConfig with .with_node_label()", + label + ), + location: snafu::Location::new(file!(), line!(), column!()), + }); } } - // Fallback: create a simple table reference that DataFusion can resolve at execution time + // No catalog attached - create empty source fallback for flexibility + // This allows planners created with DataFusionPlanner::new() to work + // without requiring a catalog, though they won't have actual data sources let empty_source = Arc::new(crate::source_catalog::SimpleTableSource::empty()); let builder = LogicalPlanBuilder::scan(label, empty_source, None).map_err(|e| { - crate::error::GraphError::PlanError { - message: format!("Failed to create table scan for {}: {}", label, e), - location: snafu::Location::new(file!(), line!(), column!()), - } + self.plan_error(&format!("Failed to create table scan for '{}'", label), e) })?; builder .build() - .map_err(|e| crate::error::GraphError::PlanError { - message: format!("Failed to build table scan for {}: {}", label, e), - location: snafu::Location::new(file!(), line!(), column!()), - }) + .map_err(|e| self.plan_error("Failed to build scan plan", e)) } /// Build a relationship expansion (graph traversal) as a series of joins @@ -621,8 +648,13 @@ impl DataFusionPlanner { relationship_properties: &HashMap, ) -> Result { let rel_schema = rel_source.schema(); - let mut rel_builder = - LogicalPlanBuilder::scan(&rel_instance.rel_type, rel_source, None).unwrap(); + let mut rel_builder = LogicalPlanBuilder::scan(&rel_instance.rel_type, rel_source, None) + .map_err(|e| { + self.plan_error( + &format!("Failed to scan relationship '{}'", rel_instance.rel_type), + e, + ) + })?; // Apply relationship property filters (e.g., -[r {since: 2020}]->) for (k, v) in relationship_properties.iter() { @@ -632,7 +664,12 @@ impl DataFusionPlanner { op: Operator::Eq, right: Box::new(lit_expr), }); - rel_builder = rel_builder.filter(filter_expr).unwrap(); + rel_builder = rel_builder.filter(filter_expr).map_err(|e| { + self.plan_error( + &format!("Failed to apply relationship filter on '{}'", k), + e, + ) + })?; } // Use unique alias from rel_instance to avoid column conflicts @@ -645,11 +682,11 @@ impl DataFusionPlanner { }) .collect(); - Ok(rel_builder + rel_builder .project(rel_qualified_exprs) - .unwrap() + .map_err(|e| self.plan_error("Failed to project relationship columns", e))? .build() - .unwrap()) + .map_err(|e| self.plan_error("Failed to build relationship scan", e)) } /// Join source node plan with relationship scan @@ -669,14 +706,14 @@ impl DataFusionPlanner { let qualified_left_key = format!("{}__{}", params.source_variable, params.node_id_field); let qualified_right_key = format!("{}__{}", params.rel_qualifier, right_key); - Ok(LogicalPlanBuilder::from(left_plan) + LogicalPlanBuilder::from(left_plan) .join( rel_scan, JoinType::Inner, (vec![qualified_left_key], vec![qualified_right_key]), None, ) - .unwrap()) + .map_err(|e| self.plan_error("Failed to join source to relationship", e)) } /// Join relationship with target node scan @@ -694,17 +731,23 @@ impl DataFusionPlanner { .get(params.target_variable) .cloned() else { - return Ok(builder.build().unwrap()); + return builder + .build() + .map_err(|e| self.plan_error("Failed to build plan (no target label)", e)); }; let Some(target_source) = cat.node_source(&target_label) else { - return Ok(builder.build().unwrap()); + return builder + .build() + .map_err(|e| self.plan_error("Failed to build plan (no target source)", e)); }; // Create target node scan with qualified column aliases and property filters let target_schema = target_source.schema(); - let mut target_builder = - LogicalPlanBuilder::scan(&target_label, target_source, None).unwrap(); + let mut target_builder = LogicalPlanBuilder::scan(&target_label, target_source, None) + .map_err(|e| { + self.plan_error(&format!("Failed to scan target node '{}'", target_label), e) + })?; // Apply target property filters (e.g., (b {age: 30})) for (k, v) in params.target_properties.iter() { @@ -714,7 +757,9 @@ impl DataFusionPlanner { op: Operator::Eq, right: Box::new(lit_expr), }); - target_builder = target_builder.filter(filter_expr).unwrap(); + target_builder = target_builder.filter(filter_expr).map_err(|e| { + self.plan_error(&format!("Failed to apply target filter on '{}'", k), e) + })?; } let target_qualified_exprs: Vec = target_schema @@ -728,9 +773,9 @@ impl DataFusionPlanner { let target_scan = target_builder .project(target_qualified_exprs) - .unwrap() + .map_err(|e| self.plan_error("Failed to project target columns", e))? .build() - .unwrap(); + .map_err(|e| self.plan_error("Failed to build target scan", e))?; // Determine target join keys let target_key = match params.direction { @@ -750,9 +795,11 @@ impl DataFusionPlanner { (vec![qualified_rel_target_key], vec![qualified_target_key]), None, ) - .unwrap(); + .map_err(|e| self.plan_error("Failed to join relationship to target", e))?; - Ok(builder.build().unwrap()) + builder + .build() + .map_err(|e| self.plan_error("Failed to build final join plan", e)) } /// Get the expected qualified column names for variable-length path results @@ -1049,19 +1096,53 @@ impl DataFusionPlanner { ctx: &PlanningContext, target_variable: &str, ) -> Result<(String, &crate::config::NodeMapping)> { - let target_label = ctx - .analysis - .var_to_label - .get(target_variable) - .or_else(|| self.config.node_mappings.keys().next()) - .ok_or_else(|| crate::error::GraphError::ConfigError { + // Try to get label from analysis first + let target_label = if let Some(label) = ctx.analysis.var_to_label.get(target_variable) { + label.clone() + } else if target_variable.starts_with("_temp_") { + // For temporary variables in multi-hop paths (e.g., "_temp_a_1" or "_temp_foo_bar_1"), + // infer the label from the source variable by extracting the base name + // Format: _temp_{source}_{hop_index} + // Note: source can contain underscores, so we reconstruct it from all parts + // between the _temp prefix and the final hop index + let parts: Vec<&str> = target_variable.split('_').collect(); + if parts.len() >= 4 { + // parts[0] = "", parts[1] = "temp", parts[2..len-1] = source variable parts, parts[len-1] = hop index + let source_var = parts[2..parts.len() - 1].join("_"); + ctx.analysis + .var_to_label + .get(&source_var) + .ok_or_else(|| crate::error::GraphError::ConfigError { + message: format!( + "Cannot infer label for temporary variable '{}' \ + from source variable '{}'", + target_variable, source_var + ), + location: snafu::Location::new(file!(), line!(), column!()), + })? + .clone() + } else { + return Err(crate::error::GraphError::ConfigError { + message: format!( + "Invalid temporary variable format: '{}'. \ + Expected format: _temp_{{source}}_{{index}}", + target_variable + ), + location: snafu::Location::new(file!(), line!(), column!()), + }); + } + } else { + // Not in analysis and not a temp variable - this is an error + return Err(crate::error::GraphError::ConfigError { message: format!( - "Cannot infer target label for variable: {}", + "Cannot determine target node label for variable '{}'. \ + This variable was not found in the query analysis. \ + Ensure the query properly defines this node variable.", target_variable ), location: snafu::Location::new(file!(), line!(), column!()), - })? - .clone(); + }); + }; let node_map = self .config @@ -2921,4 +3002,271 @@ mod tests { assert!(s.contains("full_name")); assert!(s.contains("p__age")); } + + #[test] + fn test_temp_variable_with_underscores_in_source() { + // Test that temporary variables work correctly when source variable contains underscores + let cfg = crate::config::GraphConfig::builder() + .with_node_label("Person", "id") + .with_relationship("KNOWS", "src_person_id", "dst_person_id") + .build() + .unwrap(); + let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); + + // Create a scan with a variable name containing underscores + let scan = LogicalOperator::ScanByLabel { + variable: "foo_bar".to_string(), // Variable with underscores + label: "Person".to_string(), + properties: Default::default(), + }; + + let var_expand = LogicalOperator::VariableLengthExpand { + input: Box::new(scan), + source_variable: "foo_bar".to_string(), // Will generate _temp_foo_bar_1 + target_variable: "baz".to_string(), + relationship_types: vec!["KNOWS".to_string()], + direction: crate::ast::RelationshipDirection::Outgoing, + min_length: Some(2), + max_length: Some(2), + relationship_variable: None, + target_properties: Default::default(), + }; + + let result = planner.plan(&var_expand); + + // Should succeed - the temp variable parsing should handle underscores correctly + assert!( + result.is_ok(), + "Should handle source variables with underscores: {:?}", + result.err() + ); + } + + // ======================================================================== + // Failure Scenario Tests + // ======================================================================== + + #[test] + fn test_scan_missing_node_label_with_catalog_fails_fast() { + // Test that when a catalog is attached, scanning a non-existent label fails fast + // This catches typos and configuration issues at planning time + let cfg = crate::config::GraphConfig::builder() + .with_node_label("Person", "id") + .build() + .unwrap(); + let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); + + let scan = LogicalOperator::ScanByLabel { + variable: "x".to_string(), + label: "NonExistentLabel".to_string(), // This label doesn't exist in catalog + properties: Default::default(), + }; + + let result = planner.plan(&scan); + + // Should return ConfigError with helpful message + assert!( + result.is_err(), + "Should fail when catalog exists but label is missing" + ); + match result { + Err(crate::error::GraphError::ConfigError { message, .. }) => { + assert!( + message.contains("NonExistentLabel"), + "Error should mention the missing label" + ); + assert!( + message.contains("not found"), + "Error should indicate label not found" + ); + } + _ => panic!("Expected ConfigError for missing node label"), + } + } + + #[test] + fn test_scan_without_catalog_uses_empty_source() { + // Test that when no catalog is attached, scanning creates an empty source fallback + // This allows DataFusionPlanner::new() to work without requiring a catalog + let cfg = crate::config::GraphConfig::builder() + .with_node_label("Person", "id") + .build() + .unwrap(); + let planner = DataFusionPlanner::new(cfg); // No catalog attached + + let scan = LogicalOperator::ScanByLabel { + variable: "x".to_string(), + label: "AnyLabel".to_string(), // Any label works without catalog + properties: Default::default(), + }; + + let result = planner.plan(&scan); + + // Should succeed with empty source fallback + assert!( + result.is_ok(), + "Should succeed with empty source when no catalog attached" + ); + } + + #[test] + fn test_expand_with_missing_relationship() { + // Test that expanding with non-existent relationship type handles gracefully + let cfg = crate::config::GraphConfig::builder() + .with_node_label("Person", "id") + .with_relationship("KNOWS", "src_id", "dst_id") + .build() + .unwrap(); + let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); + + let scan = LogicalOperator::ScanByLabel { + variable: "a".to_string(), + label: "Person".to_string(), + properties: Default::default(), + }; + + let expand = LogicalOperator::Expand { + input: Box::new(scan), + source_variable: "a".to_string(), + target_variable: "b".to_string(), + target_label: "Person".to_string(), + relationship_types: vec!["NONEXISTENT_REL".to_string()], // Doesn't exist + direction: crate::ast::RelationshipDirection::Outgoing, + relationship_variable: None, + properties: Default::default(), + target_properties: Default::default(), + }; + + let result = planner.plan(&expand); + + // Should handle gracefully - either error or empty result + // The key is no panic + match result { + Ok(_) => {} // Graceful handling + Err(e) => { + // Should be a PlanError + assert!(matches!(e, crate::error::GraphError::PlanError { .. })); + } + } + } + + #[test] + fn test_filter_preserves_error_context() { + // Test that filter errors include helpful context + use crate::ast::{BooleanExpression, PropertyRef, ValueExpression}; + + let cfg = crate::config::GraphConfig::builder() + .with_node_label("Person", "id") + .build() + .unwrap(); + let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); + + let scan = LogicalOperator::ScanByLabel { + variable: "p".to_string(), + label: "Person".to_string(), + properties: Default::default(), + }; + + // Create a filter with a property reference + let filter = LogicalOperator::Filter { + input: Box::new(scan), + predicate: BooleanExpression::Comparison { + left: ValueExpression::Property(PropertyRef { + variable: "p".to_string(), + property: "age".to_string(), + }), + operator: crate::ast::ComparisonOperator::GreaterThan, + right: ValueExpression::Literal(crate::ast::PropertyValue::Integer(30)), + }, + }; + + let result = planner.plan(&filter); + + // Should succeed - this tests that valid filters work + assert!(result.is_ok(), "Valid filter should succeed"); + } + + #[test] + fn test_variable_length_with_invalid_range() { + // Test that invalid variable-length ranges are caught + let cfg = crate::config::GraphConfig::builder() + .with_node_label("Person", "id") + .with_relationship("KNOWS", "src_id", "dst_id") + .build() + .unwrap(); + let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); + + let scan = LogicalOperator::ScanByLabel { + variable: "a".to_string(), + label: "Person".to_string(), + properties: Default::default(), + }; + + let var_expand = LogicalOperator::VariableLengthExpand { + input: Box::new(scan), + source_variable: "a".to_string(), + target_variable: "b".to_string(), + relationship_types: vec!["KNOWS".to_string()], + direction: crate::ast::RelationshipDirection::Outgoing, + min_length: Some(5), // min > max + max_length: Some(2), + relationship_variable: None, + target_properties: Default::default(), + }; + + let result = planner.plan(&var_expand); + + // Should return InvalidPattern error + assert!(result.is_err(), "Invalid range should return error"); + match result { + Err(crate::error::GraphError::InvalidPattern { message, .. }) => { + assert!(message.contains("min"), "Error should mention min"); + assert!(message.contains("max"), "Error should mention max"); + } + _ => panic!("Expected InvalidPattern error"), + } + } + + #[test] + fn test_variable_length_exceeds_max_hops() { + // Test that exceeding MAX_VARIABLE_LENGTH_HOPS is caught + let cfg = crate::config::GraphConfig::builder() + .with_node_label("Person", "id") + .with_relationship("KNOWS", "src_id", "dst_id") + .build() + .unwrap(); + let planner = DataFusionPlanner::with_catalog(cfg, make_catalog()); + + let scan = LogicalOperator::ScanByLabel { + variable: "a".to_string(), + label: "Person".to_string(), + properties: Default::default(), + }; + + let var_expand = LogicalOperator::VariableLengthExpand { + input: Box::new(scan), + source_variable: "a".to_string(), + target_variable: "b".to_string(), + relationship_types: vec!["KNOWS".to_string()], + direction: crate::ast::RelationshipDirection::Outgoing, + min_length: Some(1), + max_length: Some(100), // Way too high + relationship_variable: None, + target_properties: Default::default(), + }; + + let result = planner.plan(&var_expand); + + // Should return UnsupportedFeature error + assert!(result.is_err(), "Exceeding max hops should return error"); + match result { + Err(crate::error::GraphError::UnsupportedFeature { feature, .. }) => { + assert!( + feature.contains("Variable-length"), + "Error should mention variable-length" + ); + } + _ => panic!("Expected UnsupportedFeature error"), + } + } } diff --git a/rust/lance-graph/tests/integration_datafusion_pipeline.rs b/rust/lance-graph/tests/datafusion_pipeline.rs similarity index 93% rename from rust/lance-graph/tests/integration_datafusion_pipeline.rs rename to rust/lance-graph/tests/datafusion_pipeline.rs index bdd5251bc..5a220170a 100644 --- a/rust/lance-graph/tests/integration_datafusion_pipeline.rs +++ b/rust/lance-graph/tests/datafusion_pipeline.rs @@ -687,6 +687,149 @@ async fn test_datafusion_one_hop_filtered_source_age_strict() { assert_eq!(set, expected); } +#[tokio::test] +async fn test_datafusion_one_hop_with_city_filter() { + let config = create_graph_config(); + let person_batch = create_person_dataset(); + let knows_batch = create_knows_dataset(); + + // Query: Filter targets by city using inline property filter + // Tests inline property filter instead of WHERE clause + let query = + CypherQuery::new("MATCH (a:Person)-[:KNOWS]->(b:Person {city: 'Seattle'}) RETURN b.name") + .unwrap() + .with_config(config); + + let mut datasets = HashMap::new(); + datasets.insert("Person".to_string(), person_batch); + datasets.insert("KNOWS".to_string(), knows_batch); + + let out = query.execute_datafusion(datasets).await.unwrap(); + + // Only Eve has city = 'Seattle' and is reachable (David->Eve) + assert_eq!(out.num_rows(), 1); + + let names = out + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(names.value(0), "Eve"); +} + +#[tokio::test] +async fn test_datafusion_one_hop_multiple_properties() { + let config = create_graph_config(); + let person_batch = create_person_dataset(); + let knows_batch = create_knows_dataset(); + + // Query: Return multiple properties from both source and target + let query = CypherQuery::new( + "MATCH (a:Person)-[:KNOWS]->(b:Person) \ + RETURN a.name, a.age, b.name, b.age", + ) + .unwrap() + .with_config(config); + + let mut datasets = HashMap::new(); + datasets.insert("Person".to_string(), person_batch); + datasets.insert("KNOWS".to_string(), knows_batch); + + let out = query.execute_datafusion(datasets).await.unwrap(); + + assert_eq!(out.num_columns(), 4); + assert_eq!(out.num_rows(), 5); + + let a_names = out + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let a_ages = out.column(1).as_any().downcast_ref::().unwrap(); + let b_names = out + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let b_ages = out.column(3).as_any().downcast_ref::().unwrap(); + + // Verify at least one row has correct data + let mut found_alice_bob = false; + for i in 0..out.num_rows() { + if a_names.value(i) == "Alice" && b_names.value(i) == "Bob" { + assert_eq!(a_ages.value(i), 25); + assert_eq!(b_ages.value(i), 35); + found_alice_bob = true; + } + } + assert!(found_alice_bob); +} + +#[tokio::test] +async fn test_datafusion_one_hop_return_relationship_properties() { + let config = create_graph_config(); + let person_batch = create_person_dataset(); + let knows_batch = create_knows_dataset(); + + // Query: Return both node and relationship properties in projection + // This validates qualified relationship columns and aliasing + let query = CypherQuery::new( + "MATCH (a:Person)-[r:KNOWS]->(b:Person) \ + RETURN a.name, r.since_year, b.name \ + ORDER BY a.name, b.name", + ) + .unwrap() + .with_config(config); + + let mut datasets = HashMap::new(); + datasets.insert("Person".to_string(), person_batch); + datasets.insert("KNOWS".to_string(), knows_batch); + + let out = query.execute_datafusion(datasets).await.unwrap(); + + // Should return 3 columns: a.name, r.since_year, b.name + assert_eq!(out.num_columns(), 3); + // Should return 5 edges + assert_eq!(out.num_rows(), 5); + + let a_names = out + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let since_years = out.column(1).as_any().downcast_ref::().unwrap(); + let b_names = out + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + + // Verify first row: Alice -> Bob (2020) + assert_eq!(a_names.value(0), "Alice"); + assert_eq!(since_years.value(0), 2020); + assert_eq!(b_names.value(0), "Bob"); + + // Verify second row: Alice -> Charlie (2018) + assert_eq!(a_names.value(1), "Alice"); + assert_eq!(since_years.value(1), 2018); + assert_eq!(b_names.value(1), "Charlie"); + + // Verify third row: Bob -> Charlie (2019) + assert_eq!(a_names.value(2), "Bob"); + assert_eq!(since_years.value(2), 2019); + assert_eq!(b_names.value(2), "Charlie"); + + // Verify fourth row: Charlie -> David (2021) + assert_eq!(a_names.value(3), "Charlie"); + assert_eq!(since_years.value(3), 2021); + assert_eq!(b_names.value(3), "David"); + + // Verify fifth row: David -> Eve (NULL since_year) + assert_eq!(a_names.value(4), "David"); + assert!(since_years.is_null(4)); // NULL value + assert_eq!(b_names.value(4), "Eve"); +} + // ============================================================================ // Two-Hop Path Query Tests // ============================================================================ @@ -941,40 +1084,6 @@ async fn test_datafusion_two_hop_no_results() { assert_eq!(out.num_rows(), 0); } -#[tokio::test] -async fn test_datafusion_varlength_projection_correctness() { - // Test that variable-length path projection correctly handles qualified column names - // and doesn't accidentally include intermediate node columns - let out = execute_test_query( - "MATCH (a:Person {name: 'Alice'})-[:KNOWS*1..2]->(b:Person) RETURN b.name", - ) - .await; - - // Alice can reach: Bob (1-hop), Charlie (1-hop and 2-hop via Bob), David (2-hop via Charlie) - // Total: 4 results (Bob, Charlie, Charlie, David) - assert_eq!(out.num_rows(), 4); - - // Verify schema only contains source and target columns, not intermediate nodes - let schema = out.schema(); - let column_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - - // Should only have b__ prefixed columns (target), no intermediate node columns - for name in &column_names { - assert!( - name.starts_with("b__"), - "Unexpected column in variable-length result: {}", - name - ); - // Ensure no double-qualified names like "b__intermediate__prop" - let remainder = &name[3..]; // Skip "b__" - assert!( - !remainder.contains("__"), - "Column name contains nested qualifiers: {}", - name - ); - } -} - // ============================================================================ // Complex Query Tests (Advanced Filtering & Multi-Condition) // ============================================================================ @@ -1084,17 +1193,20 @@ async fn test_datafusion_two_hop_return_relationship_properties() { } #[tokio::test] -async fn test_datafusion_one_hop_with_city_filter() { +async fn test_datafusion_two_hop_return_both_relationship_properties() { let config = create_graph_config(); let person_batch = create_person_dataset(); let knows_batch = create_knows_dataset(); - // Query: Filter targets by city using inline property filter - // Tests inline property filter instead of WHERE clause - let query = - CypherQuery::new("MATCH (a:Person)-[:KNOWS]->(b:Person {city: 'Seattle'}) RETURN b.name") - .unwrap() - .with_config(config); + // Query: Return properties from both relationships in a two-hop path + // This validates qualified relationship columns for r1 and r2, and proper aliasing + let query = CypherQuery::new( + "MATCH (a:Person)-[r1:KNOWS]->(b:Person)-[r2:KNOWS]->(c:Person) \ + RETURN a.name, r1.since_year, b.name, r2.since_year, c.name \ + ORDER BY a.name, c.name", + ) + .unwrap() + .with_config(config); let mut datasets = HashMap::new(); datasets.insert("Person".to_string(), person_batch); @@ -1102,15 +1214,56 @@ async fn test_datafusion_one_hop_with_city_filter() { let out = query.execute_datafusion(datasets).await.unwrap(); - // Only Eve has city = 'Seattle' and is reachable (David->Eve) - assert_eq!(out.num_rows(), 1); + // Should return 5 columns: a.name, r1.since_year, b.name, r2.since_year, c.name + assert_eq!(out.num_columns(), 5); + // Should return 4 two-hop paths + assert_eq!(out.num_rows(), 4); - let names = out + let a_names = out .column(0) .as_any() .downcast_ref::() .unwrap(); - assert_eq!(names.value(0), "Eve"); + let r1_years = out.column(1).as_any().downcast_ref::().unwrap(); + let b_names = out + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let r2_years = out.column(3).as_any().downcast_ref::().unwrap(); + let c_names = out + .column(4) + .as_any() + .downcast_ref::() + .unwrap(); + + // Verify first path: Alice -[2020]-> Bob -[2019]-> Charlie + assert_eq!(a_names.value(0), "Alice"); + assert_eq!(r1_years.value(0), 2020); + assert_eq!(b_names.value(0), "Bob"); + assert_eq!(r2_years.value(0), 2019); + assert_eq!(c_names.value(0), "Charlie"); + + // Verify second path: Alice -[2018]-> Charlie -[2021]-> David + assert_eq!(a_names.value(1), "Alice"); + assert_eq!(r1_years.value(1), 2018); + assert_eq!(b_names.value(1), "Charlie"); + assert_eq!(r2_years.value(1), 2021); + assert_eq!(c_names.value(1), "David"); + + // Verify third path: Bob -[2019]-> Charlie -[2021]-> David + assert_eq!(a_names.value(2), "Bob"); + assert_eq!(r1_years.value(2), 2019); + assert_eq!(b_names.value(2), "Charlie"); + assert_eq!(r2_years.value(2), 2021); + assert_eq!(c_names.value(2), "David"); + + // Verify fourth path: Charlie -[2021]-> David -[NULL]-> Eve + assert_eq!(a_names.value(3), "Charlie"); + assert_eq!(r1_years.value(3), 2021); + assert_eq!(b_names.value(3), "David"); + assert!(r2_years.is_null(3)); // NULL value for David -> Eve + assert_eq!(c_names.value(3), "Eve"); } #[tokio::test] @@ -1230,51 +1383,37 @@ async fn test_datafusion_two_hop_same_intermediate_node() { } #[tokio::test] -async fn test_datafusion_one_hop_multiple_properties() { - let config = create_graph_config(); - let person_batch = create_person_dataset(); - let knows_batch = create_knows_dataset(); - - // Query: Return multiple properties from both source and target - let query = CypherQuery::new( - "MATCH (a:Person)-[:KNOWS]->(b:Person) \ - RETURN a.name, a.age, b.name, b.age", +async fn test_datafusion_varlength_projection_correctness() { + // Test that variable-length path projection correctly handles qualified column names + // and doesn't accidentally include intermediate node columns + let out = execute_test_query( + "MATCH (a:Person {name: 'Alice'})-[:KNOWS*1..2]->(b:Person) RETURN b.name", ) - .unwrap() - .with_config(config); - - let mut datasets = HashMap::new(); - datasets.insert("Person".to_string(), person_batch); - datasets.insert("KNOWS".to_string(), knows_batch); - - let out = query.execute_datafusion(datasets).await.unwrap(); + .await; - assert_eq!(out.num_columns(), 4); - assert_eq!(out.num_rows(), 5); + // Alice can reach: Bob (1-hop), Charlie (1-hop and 2-hop via Bob), David (2-hop via Charlie) + // Total: 4 results (Bob, Charlie, Charlie, David) + assert_eq!(out.num_rows(), 4); - let a_names = out - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - let a_ages = out.column(1).as_any().downcast_ref::().unwrap(); - let b_names = out - .column(2) - .as_any() - .downcast_ref::() - .unwrap(); - let b_ages = out.column(3).as_any().downcast_ref::().unwrap(); + // Verify schema only contains source and target columns, not intermediate nodes + let schema = out.schema(); + let column_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); - // Verify at least one row has correct data - let mut found_alice_bob = false; - for i in 0..out.num_rows() { - if a_names.value(i) == "Alice" && b_names.value(i) == "Bob" { - assert_eq!(a_ages.value(i), 25); - assert_eq!(b_ages.value(i), 35); - found_alice_bob = true; - } + // Should only have b__ prefixed columns (target), no intermediate node columns + for name in &column_names { + assert!( + name.starts_with("b__"), + "Unexpected column in variable-length result: {}", + name + ); + // Ensure no double-qualified names like "b__intermediate__prop" + let remainder = &name[3..]; // Skip "b__" + assert!( + !remainder.contains("__"), + "Column name contains nested qualifiers: {}", + name + ); } - assert!(found_alice_bob); } #[tokio::test]