lance_csr/datafusion_planner/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! DataFusion-based physical planner for graph queries
5//!
6//! Translates graph logical plans into DataFusion logical plans using a two-phase approach:
7//!
8//! ## Phase 1: Analysis
9//! - Assigns unique IDs to relationship instances to avoid column conflicts
10//! - Collects variable-to-label mappings and required datasets
11//!
12//! ## Phase 2: Plan Building
13//! - Nodes -> Table scans, Relationships -> Linking tables, Traversals -> Joins
14//! - Variable-length paths (`*1..3`) use unrolling: generate fixed-length plans + UNION
15//! - All columns qualified as `{variable}__{column}` to avoid ambiguity
16
17pub mod analysis;
18mod builder;
19mod config_helpers;
20mod cypher_udf;
21pub(crate) mod expression;
22mod join_ops;
23mod list_folds;
24mod math_functions;
25mod scalar_functions;
26mod scan_ops;
27mod udf;
28pub mod vector_ops;
29
30#[cfg(test)]
31mod test_fixtures;
32
33// Re-export public types
34pub use analysis::{PlanningContext, QueryAnalysis, RelationshipInstance};
35pub(crate) use scalar_functions::sql_expr;
36
37use crate::config::GraphConfig;
38use crate::error::Result;
39use crate::logical_plan::LogicalOperator;
40use datafusion::logical_expr::LogicalPlan;
41use lance_csr_catalog::GraphSourceCatalog;
42use std::sync::Arc;
43
44/// Planner abstraction for graph-to-physical planning
45pub trait GraphPhysicalPlanner {
46    fn plan(&self, logical_plan: &LogicalOperator) -> Result<LogicalPlan>;
47}
48
49/// DataFusion-based physical planner
50pub struct DataFusionPlanner {
51    pub(crate) config: GraphConfig,
52    pub(crate) catalog: Option<Arc<dyn GraphSourceCatalog>>,
53    pub(crate) adjacency: Option<crate::csr_expand::AdjacencyCatalog>,
54    pub(crate) parameter_types: std::collections::HashMap<String, arrow_schema::DataType>,
55}
56
57impl DataFusionPlanner {
58    pub fn new(config: GraphConfig) -> Self {
59        Self {
60            config,
61            catalog: None,
62            adjacency: None,
63            parameter_types: Default::default(),
64        }
65    }
66
67    pub fn with_catalog(config: GraphConfig, catalog: Arc<dyn GraphSourceCatalog>) -> Self {
68        Self {
69            config,
70            catalog: Some(catalog),
71            adjacency: None,
72            parameter_types: Default::default(),
73        }
74    }
75
76    /// Helper to convert DataFusion builder errors into GraphError::PlanError with context
77    pub(crate) fn plan_error<E: std::fmt::Display>(
78        &self,
79        context: &str,
80        error: E,
81    ) -> crate::error::GraphError {
82        crate::error::GraphError::PlanError {
83            message: format!("{}: {}", context, error),
84            location: snafu::Location::new(file!(), line!(), column!()),
85        }
86    }
87}
88
89impl GraphPhysicalPlanner for DataFusionPlanner {
90    fn plan(&self, logical_plan: &LogicalOperator) -> Result<LogicalPlan> {
91        // Phase 1: Analyze query structure
92        let analysis = analysis::analyze(logical_plan)?;
93
94        // Phase 2: Build execution plan with context
95        let mut ctx = PlanningContext::new(&analysis);
96        self.build_operator(&mut ctx, logical_plan)
97    }
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103    use crate::ast::{
104        BooleanExpression, ComparisonOperator, PropertyRef, PropertyValue, RelationshipDirection,
105        ValueExpression,
106    };
107    use crate::logical_plan::LogicalOperator;
108    use test_fixtures::{make_catalog, person_config, person_knows_config, person_scan};
109
110    #[test]
111    fn test_filter_preserves_error_context() {
112        // Test that filter errors include helpful context
113        let planner = DataFusionPlanner::with_catalog(person_config(), make_catalog());
114
115        let scan = person_scan("p");
116
117        // Create a filter with a property reference
118        let filter = LogicalOperator::Filter {
119            input: Box::new(scan),
120            predicate: BooleanExpression::Comparison {
121                left: ValueExpression::Property(PropertyRef {
122                    variable: "p".to_string(),
123                    property: "age".to_string(),
124                }),
125                operator: ComparisonOperator::GreaterThan,
126                right: ValueExpression::Literal(PropertyValue::Integer(30)),
127            },
128        };
129
130        let result = planner.plan(&filter);
131
132        // Should succeed - this tests that valid filters work
133        assert!(result.is_ok(), "Valid filter should succeed");
134    }
135
136    #[test]
137    fn test_exists_on_relationship_property_is_qualified() {
138        // Test that EXISTS on relationship properties uses qualified column names
139        let planner = DataFusionPlanner::with_catalog(person_knows_config(), make_catalog());
140
141        let scan_a = person_scan("a");
142        let expand = LogicalOperator::Expand {
143            input: Box::new(scan_a),
144            source_variable: "a".to_string(),
145            target_variable: "b".to_string(),
146            target_label: "Person".to_string(),
147            relationship_types: vec!["KNOWS".to_string()],
148            direction: RelationshipDirection::Outgoing,
149            relationship_variable: Some("r".to_string()),
150            properties: Default::default(),
151            target_properties: Default::default(),
152        };
153        let pred = BooleanExpression::Exists(PropertyRef {
154            variable: "r".into(),
155            property: "src_person_id".into(),
156        });
157        let filter = LogicalOperator::Filter {
158            input: Box::new(expand),
159            predicate: pred,
160        };
161        let df_plan = planner.plan(&filter).unwrap();
162        let s = format!("{:?}", df_plan);
163        assert!(s.contains("Filter"), "missing Filter: {}", s);
164        assert!(
165            s.contains("r__src_person_id") || s.contains("IsNotNull"),
166            "missing qualified rel column or IsNotNull in filter: {}",
167            s
168        );
169    }
170
171    #[test]
172    fn test_in_list_on_relationship_property_is_qualified() {
173        // Test that IN lists on relationship properties use qualified column names
174        let planner = DataFusionPlanner::with_catalog(person_knows_config(), make_catalog());
175
176        let scan_a = person_scan("a");
177        let expand = LogicalOperator::Expand {
178            input: Box::new(scan_a),
179            source_variable: "a".to_string(),
180            target_variable: "b".to_string(),
181            target_label: "Person".to_string(),
182            relationship_types: vec!["KNOWS".to_string()],
183            direction: RelationshipDirection::Outgoing,
184            relationship_variable: Some("r".to_string()),
185            properties: Default::default(),
186            target_properties: Default::default(),
187        };
188        let filter = LogicalOperator::Filter {
189            input: Box::new(expand),
190            predicate: BooleanExpression::In {
191                expression: ValueExpression::Property(PropertyRef {
192                    variable: "r".into(),
193                    property: "src_person_id".into(),
194                }),
195                list: vec![
196                    ValueExpression::Literal(PropertyValue::Integer(1)),
197                    ValueExpression::Literal(PropertyValue::Integer(2)),
198                ],
199            },
200        };
201        let df_plan = planner.plan(&filter).unwrap();
202        let s = format!("{:?}", df_plan);
203        assert!(s.contains("Filter"), "missing Filter: {}", s);
204        assert!(
205            s.contains("r__src_person_id"),
206            "missing qualified rel column in IN list filter: {}",
207            s
208        );
209    }
210
211    #[test]
212    fn test_exists_and_in_on_node_props_materialized() {
213        // Test that EXISTS and IN expressions on node properties work correctly
214        let planner = DataFusionPlanner::with_catalog(person_config(), make_catalog());
215
216        let scan_a = person_scan("a");
217        let pred = BooleanExpression::And(
218            Box::new(BooleanExpression::Exists(PropertyRef {
219                variable: "a".into(),
220                property: "name".into(),
221            })),
222            Box::new(BooleanExpression::In {
223                expression: ValueExpression::Property(PropertyRef {
224                    variable: "a".into(),
225                    property: "age".into(),
226                }),
227                list: vec![
228                    ValueExpression::Literal(PropertyValue::Integer(20)),
229                    ValueExpression::Literal(PropertyValue::Integer(30)),
230                ],
231            }),
232        );
233        let filter = LogicalOperator::Filter {
234            input: Box::new(scan_a),
235            predicate: pred,
236        };
237        let df_plan = planner.plan(&filter).unwrap();
238        let s = format!("{:?}", df_plan);
239        assert!(s.contains("Filter"), "missing Filter: {}", s);
240        assert!(
241            s.contains("a__name") || s.contains("IsNotNull"),
242            "missing EXISTS on a__name: {}",
243            s
244        );
245        assert!(
246            s.contains("a__age") || s.contains("age"),
247            "missing IN on a.age: {}",
248            s
249        );
250    }
251}
lance_csr/datafusion_planner/mod.rs

lance_csr/datafusion_planner/
mod.rs
