From dbd08cc121fc10051c655eec2e17dfbc2b888f8a Mon Sep 17 00:00:00 2001 From: beinan Date: Fri, 24 Oct 2025 20:22:11 +0000 Subject: [PATCH] Make the knowledge graph support query with where-in clause --- python/python/knowledge_graph/main.py | 42 +++++++++++++-- rust/lance-graph/src/datafusion_planner.rs | 38 +++++++++++++ rust/lance-graph/src/parser.rs | 62 +++++++++++++++++++++- 3 files changed, 136 insertions(+), 6 deletions(-) diff --git a/python/python/knowledge_graph/main.py b/python/python/knowledge_graph/main.py index 1fca1b9f9..addc46650 100644 --- a/python/python/knowledge_graph/main.py +++ b/python/python/knowledge_graph/main.py @@ -397,6 +397,16 @@ def _build_query_prompt( "Use the schema summary to craft queries that directly answer the " "question." ), + ( + " • Use the schema summary and allowed relationship_type values to " + "identify candidate relationship directions and types." + ), + ( + " • When the schema lists relationship_type values and the question " + "does not narrow them down, treat the list as exhaustive and include " + "every value in your filter using OR clauses or " + "WHERE rel.relationship_type IN [...]." + ), ( "Always specify node labels and relationship types in MATCH patterns " "that introduce aliases." @@ -405,12 +415,25 @@ def _build_query_prompt( (" • MATCH (e:Entity) to scan entity rows (name, name_lower, entity_id)."), ( " • MATCH (src:Entity)-[rel:RELATIONSHIP]->(dst:Entity) to traverse " - "relationships (relationship_type column)." + "relationships (relationship_type column); `src` aligns with " + "`source_entity_id` and `dst` with `target_entity_id`." + ), + ( + " • Decide which node should be `src` versus `dst` based on the " + "relationship meaning in the question and schema hints." + ), + ( + " • Map natural language roles (team, person, product, etc.) to the " + "`entity_type` column so queries filter to the expected entities." ), " • Use WHERE e.column = 'value' for node-level filters.", ( " • Filter relationships with WHERE rel.relationship_type = 'VALUE' " - "or by comparing rel.source_entity_id / rel.target_entity_id." + "or by comparing rel.source_entity_id / rel.target_entity_id; when the " + "question does not name a specific relationship type, include every " + "relevant value from the schema summary using OR clauses or " + "WHERE rel.relationship_type IN [...], explicitly note which values " + "you considered, and avoid emitting only a single guessed type." ), ( " • Select columns using the aliases you define, such as e.name or " @@ -421,8 +444,19 @@ def _build_query_prompt( "filter rel.relationship_type instead of [:TYPE]." ), ( - "Example: MATCH (src:Entity)-[rel:RELATIONSHIP]->(dst:Entity) " - f"WHERE rel.relationship_type = '{example_rel_type}' RETURN rel." + "Example: MATCH (part:Entity)-[rel:RELATIONSHIP]->(whole:Entity) " + f"WHERE rel.relationship_type = '{example_rel_type}' " + "RETURN part.name, whole.name." + ), + ( + "Example: MATCH (a:Entity)-[rel:RELATIONSHIP]->(b:Entity) WHERE " + "rel.relationship_type = 'TYPE_A' OR rel.relationship_type = 'TYPE_B' " + "RETURN a.name, b.name." + ), + ( + "Example: MATCH (src:Entity)-[rel:RELATIONSHIP]->(dst:Entity) WHERE " + "rel.relationship_type IN ['TYPE_A', 'TYPE_B', 'TYPE_C'] " + "RETURN src.name, dst.name." ), ( "Example: MATCH (dst:Entity) WHERE dst.name_lower = 'acme corp' " diff --git a/rust/lance-graph/src/datafusion_planner.rs b/rust/lance-graph/src/datafusion_planner.rs index 4b058edd8..c7668ab46 100644 --- a/rust/lance-graph/src/datafusion_planner.rs +++ b/rust/lance-graph/src/datafusion_planner.rs @@ -264,6 +264,15 @@ impl DataFusionPlanner { right: Box::new(r), }) } + BE::In { expression, list } => { + use datafusion::logical_expr::expr::InList as DFInList; + let expr = self.to_df_value_expr(expression); + let list_exprs = list + .iter() + .map(|item| self.to_df_value_expr(item)) + .collect::>(); + Expr::InList(DFInList::new(Box::new(expr), list_exprs, false)) + } BE::And(l, r) => Expr::BinaryExpr(BinaryExpr { left: Box::new(self.to_df_boolean_expr(l)), op: Operator::And, @@ -334,6 +343,35 @@ mod tests { ) } + #[test] + fn test_df_boolean_expr_in_list() { + let cfg = crate::config::GraphConfig::builder().build().unwrap(); + let planner = DataFusionPlanner::new(cfg); + let expr = BooleanExpression::In { + expression: ValueExpression::Property(PropertyRef { + variable: "rel".into(), + property: "relationship_type".into(), + }), + list: vec![ + ValueExpression::Literal(PropertyValue::String("WORKS_FOR".into())), + ValueExpression::Literal(PropertyValue::String("PART_OF".into())), + ], + }; + + if let Expr::InList(in_list) = planner.to_df_boolean_expr(&expr) { + assert!(!in_list.negated); + assert_eq!(in_list.list.len(), 2); + match *in_list.expr { + Expr::Column(ref col_expr) => { + assert_eq!(col_expr.name(), "relationship_type"); + } + other => panic!("Expected column expression, got {:?}", other), + } + } else { + panic!("Expected InList expression"); + } + } + #[test] fn test_df_planner_scan_filter_project() { let scan = LogicalOperator::ScanByLabel { diff --git a/rust/lance-graph/src/parser.rs b/rust/lance-graph/src/parser.rs index e0b8a0b45..dc90f7b46 100644 --- a/rust/lance-graph/src/parser.rs +++ b/rust/lance-graph/src/parser.rs @@ -13,7 +13,7 @@ use nom::{ bytes::complete::{tag, tag_no_case, take_while1}, character::complete::{char, multispace0, multispace1}, combinator::{map, opt, recognize}, - multi::{many0, separated_list0}, + multi::{many0, separated_list0, separated_list1}, sequence::{delimited, pair, preceded, tuple}, IResult, }; @@ -316,6 +316,20 @@ fn comparison_expression(input: &str) -> IResult<&str, BooleanExpression> { let (input, _) = multispace0(input)?; let (input, left) = value_expression(input)?; let (input, _) = multispace0(input)?; + let left_clone = left.clone(); + + if let Ok((input_after_in, (_, _, list))) = + tuple((tag_no_case("IN"), multispace0, value_expression_list))(input) + { + return Ok(( + input_after_in, + BooleanExpression::In { + expression: left, + list, + }, + )); + } + let (input, operator) = comparison_operator(input)?; let (input, _) = multispace0(input)?; let (input, right) = value_expression(input)?; @@ -323,7 +337,7 @@ fn comparison_expression(input: &str) -> IResult<&str, BooleanExpression> { Ok(( input, BooleanExpression::Comparison { - left, + left: left_clone, operator, right, }, @@ -352,6 +366,17 @@ fn value_expression(input: &str) -> IResult<&str, ValueExpression> { ))(input) } +fn value_expression_list(input: &str) -> IResult<&str, Vec> { + delimited( + tuple((char('['), multispace0)), + separated_list1( + tuple((multispace0, char(','), multispace0)), + value_expression, + ), + tuple((multispace0, char(']'))), + )(input) +} + // Parse a property reference: variable.property fn property_reference(input: &str) -> IResult<&str, PropertyRef> { let (input, variable) = identifier(input)?; @@ -726,6 +751,39 @@ mod tests { } } + #[test] + fn test_parse_query_with_in_clause() { + let query = "MATCH (src:Entity)-[rel:RELATIONSHIP]->(dst:Entity) WHERE rel.relationship_type IN ['WORKS_FOR', 'PART_OF'] RETURN src.name"; + let result = parse_cypher_query(query).unwrap(); + + let where_clause = result.where_clause.expect("Expected WHERE clause"); + match where_clause.expression { + BooleanExpression::In { expression, list } => { + match expression { + ValueExpression::Property(prop_ref) => { + assert_eq!(prop_ref.variable, "rel"); + assert_eq!(prop_ref.property, "relationship_type"); + } + _ => panic!("Expected property reference in IN expression"), + } + assert_eq!(list.len(), 2); + match &list[0] { + ValueExpression::Literal(PropertyValue::String(val)) => { + assert_eq!(val, "WORKS_FOR"); + } + _ => panic!("Expected first list item to be a string literal"), + } + match &list[1] { + ValueExpression::Literal(PropertyValue::String(val)) => { + assert_eq!(val, "PART_OF"); + } + _ => panic!("Expected second list item to be a string literal"), + } + } + other => panic!("Expected IN expression, got {:?}", other), + } + } + #[test] fn test_parse_query_with_limit() { let query = "MATCH (n:Person) RETURN n.name LIMIT 10";