linkrs_query/planning/statements/seeks/
variable_prop_index_seek.rs

1//! Variable attribute index lookup strategy
2//!
3//! Index lookup based on variable attributes, used in cases where the value of a runtime variable needs to be determined.
4//!
5//! Applicable scenarios:
6//! - MATCH (v:Person) WHERE v.name = $varName
7//! - MATCH (v:Person) WHERE v.age > $minAge
8//! - Variable binding in parameterized queries
9
10use super::seek_strategy::SeekStrategy;
11use super::seek_strategy_base::{IndexInfo, SeekResult, SeekStrategyContext, SeekStrategyType};
12use crate::storage::StorageReader;
13use linkrs_core::{StorageError, Value};
14
15/// Variable attribute predicate
16#[derive(Debug, Clone, PartialEq)]
17pub struct VariablePropertyPredicate {
18    pub property: String,
19    pub op: VariablePredicateOp,
20    pub variable_name: String,
21}
22
23/// Variable predicate operation types
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum VariablePredicateOp {
26    Eq, // =
27    Ne, // !=
28    Lt, // <
29    Le, // <=
30    Gt, // >
31    Ge, // >=
32    In, // IN
33}
34
35impl VariablePredicateOp {
36    /// Convert to ordinary predicate operations
37    pub fn to_predicate_op(&self) -> super::prop_index_seek::PredicateOp {
38        match self {
39            VariablePredicateOp::Eq => super::prop_index_seek::PredicateOp::Eq,
40            VariablePredicateOp::Ne => super::prop_index_seek::PredicateOp::Ne,
41            VariablePredicateOp::Lt => super::prop_index_seek::PredicateOp::Lt,
42            VariablePredicateOp::Le => super::prop_index_seek::PredicateOp::Le,
43            VariablePredicateOp::Gt => super::prop_index_seek::PredicateOp::Gt,
44            VariablePredicateOp::Ge => super::prop_index_seek::PredicateOp::Ge,
45            VariablePredicateOp::In => super::prop_index_seek::PredicateOp::In,
46        }
47    }
48}
49
50impl std::str::FromStr for VariablePredicateOp {
51    type Err = String;
52
53    fn from_str(s: &str) -> Result<Self, Self::Err> {
54        match s.to_uppercase().as_str() {
55            "=" | "==" => Ok(VariablePredicateOp::Eq),
56            "!=" | "<>" => Ok(VariablePredicateOp::Ne),
57            "<" => Ok(VariablePredicateOp::Lt),
58            "<=" => Ok(VariablePredicateOp::Le),
59            ">" => Ok(VariablePredicateOp::Gt),
60            ">=" => Ok(VariablePredicateOp::Ge),
61            "IN" => Ok(VariablePredicateOp::In),
62            _ => Err(format!("Invalid predicate operator: {}", s)),
63        }
64    }
65}
66
67/// Variable attribute index lookup strategy
68#[derive(Debug, Clone)]
69pub struct VariablePropIndexSeek {
70    predicates: Vec<VariablePropertyPredicate>,
71    variable_values: std::collections::HashMap<String, Value>,
72}
73
74impl VariablePropIndexSeek {
75    pub fn new(predicates: Vec<VariablePropertyPredicate>) -> Self {
76        Self {
77            predicates,
78            variable_values: std::collections::HashMap::new(),
79        }
80    }
81
82    /// Binding variable values
83    pub fn bind_variable(&mut self, name: &str, value: Value) {
84        self.variable_values.insert(name.to_string(), value);
85    }
86
87    /// Batch binding of variable values
88    pub fn bind_variables(&mut self, values: std::collections::HashMap<String, Value>) {
89        self.variable_values.extend(values);
90    }
91
92    /// Check whether all variables have been bound.
93    pub fn all_variables_bound(&self) -> bool {
94        self.predicates
95            .iter()
96            .all(|pred| self.variable_values.contains_key(&pred.variable_name))
97    }
98
99    /// Extract variable attribute predicates from the list of expressions.
100    pub fn extract_predicates(
101        expressions: &[linkrs_core::Expression],
102    ) -> Vec<VariablePropertyPredicate> {
103        let mut predicates = Vec::new();
104
105        for expr in expressions {
106            if let Some(pred) = Self::extract_predicate(expr) {
107                predicates.push(pred);
108            }
109        }
110
111        predicates
112    }
113
114    /// Extracting variable attribute predicates from a single expression
115    fn extract_predicate(expr: &linkrs_core::Expression) -> Option<VariablePropertyPredicate> {
116        use linkrs_core::types::operators::BinaryOperator;
117
118        match expr {
119            linkrs_core::Expression::Binary { op, left, right } => {
120                let op_str = match op {
121                    BinaryOperator::Equal => "=",
122                    BinaryOperator::NotEqual => "!=",
123                    BinaryOperator::LessThan => "<",
124                    BinaryOperator::LessThanOrEqual => "<=",
125                    BinaryOperator::GreaterThan => ">",
126                    BinaryOperator::GreaterThanOrEqual => ">=",
127                    _ => return None,
128                };
129
130                // Try to extract the property name and variable name: v.name = $var
131                if let (Some(prop), Some(var_name)) =
132                    (Self::extract_property(left), Self::extract_variable(right))
133                {
134                    if let Ok(pred_op) = op_str.parse::<VariablePredicateOp>() {
135                        return Some(VariablePropertyPredicate {
136                            property: prop,
137                            op: pred_op,
138                            variable_name: var_name,
139                        });
140                    }
141                }
142
143                // Swap left and right attempts: $var = v.name
144                if let (Some(prop), Some(var_name)) =
145                    (Self::extract_property(right), Self::extract_variable(left))
146                {
147                    let swapped_op = match op_str {
148                        "<" => VariablePredicateOp::Gt,
149                        "<=" => VariablePredicateOp::Ge,
150                        ">" => VariablePredicateOp::Lt,
151                        ">=" => VariablePredicateOp::Le,
152                        _ => op_str.parse::<VariablePredicateOp>().ok()?,
153                    };
154                    return Some(VariablePropertyPredicate {
155                        property: prop,
156                        op: swapped_op,
157                        variable_name: var_name,
158                    });
159                }
160
161                None
162            }
163            _ => None,
164        }
165    }
166
167    /// Extract attribute names from the expression.
168    fn extract_property(expr: &linkrs_core::Expression) -> Option<String> {
169        match expr {
170            linkrs_core::Expression::Property { object, property } => {
171                // Check for node attribute access, e.g. v.name
172                if matches!(object.as_ref(), linkrs_core::Expression::Variable(_)) {
173                    Some(property.clone())
174                } else {
175                    None
176                }
177            }
178            _ => None,
179        }
180    }
181
182    /// Extract variable names from the expression.
183    fn extract_variable(expr: &linkrs_core::Expression) -> Option<String> {
184        match expr {
185            // Variables that start with the symbol $ indicate parameters.
186            linkrs_core::Expression::Variable(name) if name.starts_with('$') => {
187                Some(name[1..].to_string())
188            }
189            _ => None,
190        }
191    }
192
193    /// Find an index that is suitable for the predicate of the variable attribute.
194    fn find_best_index<'a>(
195        &'a self,
196        context: &'a SeekStrategyContext,
197    ) -> Option<(&'a IndexInfo, &'a VariablePropertyPredicate)> {
198        for pred in &self.predicates {
199            if let Some(index) = context.get_index_for_property(&pred.property) {
200                return Some((index, pred));
201            }
202        }
203        None
204    }
205
206    /// Does the evaluated value satisfy the predicate condition?
207    fn value_matches(&self, value: &Value, pred: &VariablePropertyPredicate) -> bool {
208        // Obtain the value of a variable
209        let var_value = match self.variable_values.get(&pred.variable_name) {
210            Some(v) => v,
211            None => return false, // The variable is not bound, so a match cannot be made.
212        };
213
214        match pred.op {
215            VariablePredicateOp::Eq => value == var_value,
216            VariablePredicateOp::Ne => value != var_value,
217            VariablePredicateOp::Lt => Self::compare_values(value, var_value)
218                .map(|c| c < 0)
219                .unwrap_or(false),
220            VariablePredicateOp::Le => Self::compare_values(value, var_value)
221                .map(|c| c <= 0)
222                .unwrap_or(false),
223            VariablePredicateOp::Gt => Self::compare_values(value, var_value)
224                .map(|c| c > 0)
225                .unwrap_or(false),
226            VariablePredicateOp::Ge => Self::compare_values(value, var_value)
227                .map(|c| c >= 0)
228                .unwrap_or(false),
229            VariablePredicateOp::In => {
230                // The IN operation requires that the variable value be a list.
231                matches!(var_value, Value::List(list) if list.contains(value))
232            }
233        }
234    }
235
236    /// Compare two values
237    fn compare_values(left: &Value, right: &Value) -> Option<i32> {
238        match (left, right) {
239            (Value::SmallInt(i1), Value::SmallInt(i2)) => Some(i1.cmp(i2) as i32),
240            (Value::Int(i1), Value::Int(i2)) => Some(i1.cmp(i2) as i32),
241            (Value::BigInt(i1), Value::BigInt(i2)) => Some(i1.cmp(i2) as i32),
242            (Value::Float(f1), Value::Float(f2)) => f1.partial_cmp(f2).map(|c| c as i32),
243            (Value::Double(f1), Value::Double(f2)) => f1.partial_cmp(f2).map(|c| c as i32),
244            (Value::String(s1), Value::String(s2)) => Some(s1.cmp(s2) as i32),
245            _ => None,
246        }
247    }
248}
249
250fn has_distinct_multi_labels(labels: &[String]) -> bool {
251    labels
252        .iter()
253        .collect::<std::collections::HashSet<_>>()
254        .len()
255        > 1
256}
257
258fn vertex_matches_single_label(
259    vertex: &linkrs_core::Vertex,
260    pattern: &super::seek_strategy_base::NodePattern,
261) -> bool {
262    if pattern.labels.is_empty() {
263        return true;
264    }
265    pattern.labels.iter().all(|label| vertex.tag.name == *label)
266}
267
268impl SeekStrategy for VariablePropIndexSeek {
269    fn execute<S: StorageReader>(
270        &self,
271        storage: &S,
272        context: &SeekStrategyContext,
273    ) -> Result<SeekResult, StorageError> {
274        // Check whether the variable has been bound.
275        if !self.all_variables_bound() {
276            return Err(StorageError::invalid_input(
277                "Variable property lookups require all variables to be bound".to_string(),
278            ));
279        }
280
281        let mut vertex_ids = Vec::new();
282        let mut rows_scanned = 0;
283
284        if has_distinct_multi_labels(&context.node_pattern.labels) {
285            return Ok(SeekResult {
286                vertex_ids,
287                strategy_used: SeekStrategyType::VariablePropIndexSeek,
288                rows_scanned,
289            });
290        }
291
292        // Find the best index.
293        if let Some((index_info, primary_pred)) = self.find_best_index(context) {
294            // Retrieve the vertices corresponding to the tags.
295            let space_name = "default"; // In fact, the relevant information should be obtained from the context.
296            let vertices = storage.scan_vertices_by_tag(space_name, &index_info.target_name)?;
297            rows_scanned = vertices.len();
298
299            // Filter the vertices that satisfy all predicates.
300            for vertex in vertices {
301                if !vertex_matches_single_label(&vertex, &context.node_pattern) {
302                    continue;
303                }
304                let mut matches_all = true;
305
306                // Check the subject and verb.
307                if let Some(prop_value) = vertex.property_value(&primary_pred.property) {
308                    if !self.value_matches(&prop_value, primary_pred) {
309                        matches_all = false;
310                    }
311                } else {
312                    matches_all = false;
313                }
314
315                // Check the other predicates.
316                if matches_all {
317                    for pred in &self.predicates {
318                        if pred.property != primary_pred.property {
319                            if let Some(prop_value) = vertex.property_value(&pred.property) {
320                                if !self.value_matches(&prop_value, pred) {
321                                    matches_all = false;
322                                    break;
323                                }
324                            } else {
325                                matches_all = false;
326                                break;
327                            }
328                        }
329                    }
330                }
331
332                if matches_all {
333                    vertex_ids.push(Value::from(*vertex.vid()));
334                }
335            }
336        }
337
338        Ok(SeekResult {
339            vertex_ids,
340            strategy_used: SeekStrategyType::VariablePropIndexSeek,
341            rows_scanned,
342        })
343    }
344
345    fn supports(&self, _context: &SeekStrategyContext) -> bool {
346        // Support is available as long as there is a predicate with variable attributes.
347        !self.predicates.is_empty()
348    }
349}
350
351#[cfg(test)]
352mod tests {
353    use super::*;
354    use linkrs_core::Expression;
355
356    #[test]
357    fn test_variable_predicate_op_from_str() {
358        assert_eq!(
359            "=".parse::<VariablePredicateOp>().ok(),
360            Some(VariablePredicateOp::Eq)
361        );
362        assert_eq!(
363            "<".parse::<VariablePredicateOp>().ok(),
364            Some(VariablePredicateOp::Lt)
365        );
366        assert_eq!(
367            ">=".parse::<VariablePredicateOp>().ok(),
368            Some(VariablePredicateOp::Ge)
369        );
370        assert_eq!(
371            "IN".parse::<VariablePredicateOp>().ok(),
372            Some(VariablePredicateOp::In)
373        );
374        assert!("unknown".parse::<VariablePredicateOp>().is_err());
375    }
376
377    #[test]
378    fn test_extract_variable_predicate() {
379        let expr = Expression::binary(
380            Expression::property(Expression::variable("v"), "name"),
381            linkrs_core::BinaryOperator::Equal,
382            Expression::variable("$varName"),
383        );
384
385        let pred = VariablePropIndexSeek::extract_predicate(&expr);
386        assert!(pred.is_some());
387
388        let pred = pred.expect("Failed to extract predicate");
389        assert_eq!(pred.property, "name");
390        assert_eq!(pred.op, VariablePredicateOp::Eq);
391        assert_eq!(pred.variable_name, "varName");
392    }
393
394    #[test]
395    fn test_variable_binding() {
396        let mut seek = VariablePropIndexSeek::new(vec![VariablePropertyPredicate {
397            property: "age".to_string(),
398            op: VariablePredicateOp::Gt,
399            variable_name: "minAge".to_string(),
400        }]);
401
402        assert!(!seek.all_variables_bound());
403
404        seek.bind_variable("minAge", Value::Int(18));
405
406        assert!(seek.all_variables_bound());
407    }
408}
linkrs_query/planning/statements/seeks/variable_prop_index_seek.rs

linkrs_query/planning/statements/seeks/
variable_prop_index_seek.rs
