diff --git a/crates/lance-graph/src/logical_plan.rs b/crates/lance-graph/src/logical_plan.rs index 9d6c32f8a..4af72e269 100644 --- a/crates/lance-graph/src/logical_plan.rs +++ b/crates/lance-graph/src/logical_plan.rs @@ -429,6 +429,9 @@ impl<'a> LogicalPlanner<'a> { self.variables .insert(target_variable.clone(), target_label.clone()); + let relationship_types = + self.resolve_relationship_types(&segment.relationship.types)?; + // Optimize fixed-length var-length expansions (*1 or *1..1) let next_plan = match segment.relationship.length.as_ref() { Some(length_range) @@ -439,7 +442,7 @@ impl<'a> LogicalPlanner<'a> { source_variable: current_src.clone(), target_variable: target_variable.clone(), target_label: target_label.clone(), - relationship_types: segment.relationship.types.clone(), + relationship_types: relationship_types.clone(), direction: segment.relationship.direction.clone(), relationship_variable: segment.relationship.variable.clone(), properties: segment.relationship.properties.clone(), @@ -450,7 +453,7 @@ impl<'a> LogicalPlanner<'a> { input: Box::new(plan), source_variable: current_src.clone(), target_variable: target_variable.clone(), - relationship_types: segment.relationship.types.clone(), + relationship_types: relationship_types.clone(), direction: segment.relationship.direction.clone(), relationship_variable: segment.relationship.variable.clone(), min_length: length_range.min, @@ -462,7 +465,7 @@ impl<'a> LogicalPlanner<'a> { source_variable: current_src.clone(), target_variable: target_variable.clone(), target_label: target_label.clone(), - relationship_types: segment.relationship.types.clone(), + relationship_types: relationship_types.clone(), direction: segment.relationship.direction.clone(), relationship_variable: segment.relationship.variable.clone(), properties: segment.relationship.properties.clone(), @@ -477,6 +480,35 @@ impl<'a> LogicalPlanner<'a> { Ok(plan) } + /// The relationship types a path segment can traverse. A relationship without a type, as in + /// `-[]->` or `-[*1..2]->`, can follow any relationship type; that is only supported when + /// the graph has a single relationship type. + fn resolve_relationship_types(&self, types: &[String]) -> Result> { + if !types.is_empty() { + return Ok(types.to_vec()); + } + let mut configured: Vec<&str> = self + .config + .relationship_mappings + .values() + .map(|mapping| mapping.relationship_type.as_str()) + .collect(); + if let [only] = configured.as_slice() { + return Ok(vec![only.to_string()]); + } + configured.sort_unstable(); + Err(GraphError::UnsupportedFeature { + feature: format!( + "Relationships without a type when the graph has {} relationship types ({}); \ + specify the type, e.g. -[:{}]->", + configured.len(), + configured.join(", "), + configured.first().copied().unwrap_or("TYPE"), + ), + location: snafu::Location::new(file!(), line!(), column!()), + }) + } + /// Extract the main variable from a logical plan (for chaining) #[allow(clippy::only_used_in_recursion)] fn extract_variable_from_plan(&self, plan: &LogicalOperator) -> Result { diff --git a/crates/lance-graph/tests/test_relationship_types.rs b/crates/lance-graph/tests/test_relationship_types.rs new file mode 100644 index 000000000..cec11142c --- /dev/null +++ b/crates/lance-graph/tests/test_relationship_types.rs @@ -0,0 +1,103 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Relationships without a type, such as `-[]->` and `-[*1..2]->` +//! (https://github.com/lance-format/lance-graph/issues/104). + +use arrow_array::{Int64Array, RecordBatch}; +use arrow_schema::{DataType, Field, Schema}; +use lance_graph::config::GraphConfig; +use lance_graph::{CypherQuery, ExecutionStrategy}; +use std::collections::HashMap; +use std::sync::Arc; + +fn int_batch(columns: &[(&str, Vec)]) -> RecordBatch { + let schema = Arc::new(Schema::new( + columns + .iter() + .map(|(name, _)| Field::new(*name, DataType::Int64, false)) + .collect::>(), + )); + let arrays = columns + .iter() + .map(|(_, values)| Arc::new(Int64Array::from(values.clone())) as _) + .collect(); + RecordBatch::try_new(schema, arrays).unwrap() +} + +/// Person 1 FOLLOWS 2, and 2 FOLLOWS 3. With `with_blocks`, the graph also has a +/// BLOCKS relationship type (1 BLOCKS 3). +async fn run(query: &str, with_blocks: bool) -> lance_graph::Result> { + let mut builder = GraphConfig::builder() + .with_node_label("Person", "id") + .with_relationship("FOLLOWS", "src", "dst"); + let mut datasets = HashMap::from([ + ("Person".to_string(), int_batch(&[("id", vec![1, 2, 3])])), + ( + "FOLLOWS".to_string(), + int_batch(&[("src", vec![1, 2]), ("dst", vec![2, 3])]), + ), + ]); + if with_blocks { + builder = builder.with_relationship("BLOCKS", "src", "dst"); + datasets.insert( + "BLOCKS".to_string(), + int_batch(&[("src", vec![1]), ("dst", vec![3])]), + ); + } + let result = CypherQuery::new(query)? + .with_config(builder.build()?) + .execute(datasets, Some(ExecutionStrategy::DataFusion)) + .await?; + let ids = result + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let mut ids = ids.values().to_vec(); + ids.sort(); + Ok(ids) +} + +#[tokio::test] +async fn test_untyped_relationship_uses_the_only_relationship_type() { + let single_hop = run( + "MATCH (a:Person {id: 1})-[]->(b:Person) RETURN b.id AS id", + false, + ) + .await + .unwrap(); + assert_eq!(single_hop, vec![2]); + + let variable_length = run( + "MATCH (a:Person {id: 1})-[*1..2]->(b:Person) RETURN b.id AS id", + false, + ) + .await + .unwrap(); + assert_eq!(variable_length, vec![2, 3]); +} + +#[tokio::test] +async fn test_untyped_relationship_with_several_relationship_types_is_rejected() { + for query in [ + "MATCH (a:Person {id: 1})-[]->(b:Person) RETURN b.id AS id", + "MATCH (a:Person {id: 1})-[*1..2]->(b:Person) RETURN b.id AS id", + ] { + let error = run(query, true).await.unwrap_err().to_string(); + assert!( + error.contains("Relationships without a type") && error.contains("BLOCKS, FOLLOWS"), + "{query}: {error}" + ); + } + + // A typed relationship still works in the same graph. + let typed = run( + "MATCH (a:Person {id: 1})-[:BLOCKS]->(b:Person) RETURN b.id AS id", + true, + ) + .await + .unwrap(); + assert_eq!(typed, vec![3]); +}