Skip to main content

reddb_server/storage/query/planner/
shape.rs

1use crate::storage::engine::vector_metadata::{MetadataFilter, MetadataValue};
2use crate::storage::query::ast::{
3    Expr, FusionStrategy, GraphPattern, GraphQuery, HybridQuery, JoinQuery, NodePattern,
4    NodeSelector, OrderByClause, PathQuery, Projection, PropertyFilter, QueryExpr, SelectItem,
5    TableQuery, TableSource, VectorQuery, VectorSource,
6};
7use crate::storage::query::sql_lowering::{
8    expr_to_filter, filter_to_expr, projection_from_literal,
9};
10use crate::storage::schema::Value;
11
12const PROJECTION_PARAM_PREFIX: &str = "__shape_projection_param__:";
13const STRING_PARAM_PREFIX: &str = "__shape_string_param__:";
14const VALUE_PARAM_PREFIX: &str = "__shape_value_param__:";
15const ROW_SELECTOR_TABLE_PREFIX: &str = "__shape_row_selector__:";
16const METADATA_VALUE_PARAM_PREFIX: &str = "__shape_metadata_value_param__:";
17const VECTOR_TEXT_PARAM_PREFIX: &str = "__shape_vector_text_param__:";
18const VECTOR_REF_ID_PREFIX: &str = "__shape_vector_ref_id__:";
19const FLOAT32_PARAM_BITS_BASE: u32 = 0x7fc0_0000;
20const FLOAT64_PARAM_BITS_BASE: u64 = 0x7ff8_0000_0000_0000;
21const U32_PARAM_BASE: u32 = 0xfff0_0000;
22
23#[derive(Debug, Clone)]
24pub struct ParameterizedQuery {
25    pub shape: QueryExpr,
26    pub parameter_count: usize,
27}
28
29pub fn parameterize_query_expr(expr: &QueryExpr) -> Option<ParameterizedQuery> {
30    let mut next_index = 0usize;
31    let shape = parameterize_query_expr_inner(expr, &mut next_index)?;
32    Some(ParameterizedQuery {
33        shape,
34        parameter_count: next_index,
35    })
36}
37
38pub fn bind_parameterized_query(
39    expr: &QueryExpr,
40    binds: &[Value],
41    parameter_count: usize,
42) -> Option<QueryExpr> {
43    if binds.len() != parameter_count {
44        return None;
45    }
46    bind_query_expr_inner(expr, binds)
47}
48
49fn parameterize_query_expr_inner(expr: &QueryExpr, next_index: &mut usize) -> Option<QueryExpr> {
50    match expr {
51        QueryExpr::Table(query) => Some(QueryExpr::Table(parameterize_table_query(
52            query, next_index,
53        )?)),
54        QueryExpr::Join(query) => {
55            Some(QueryExpr::Join(parameterize_join_query(query, next_index)?))
56        }
57        QueryExpr::Graph(query) => Some(QueryExpr::Graph(parameterize_graph_query(
58            query, next_index,
59        )?)),
60        QueryExpr::Path(query) => {
61            Some(QueryExpr::Path(parameterize_path_query(query, next_index)?))
62        }
63        QueryExpr::Vector(query) => Some(QueryExpr::Vector(parameterize_vector_query(
64            query, next_index,
65        )?)),
66        QueryExpr::Hybrid(query) => Some(QueryExpr::Hybrid(parameterize_hybrid_query(
67            query, next_index,
68        )?)),
69        _ => None,
70    }
71}
72
73fn parameterize_table_query(query: &TableQuery, next_index: &mut usize) -> Option<TableQuery> {
74    let source = match &query.source {
75        Some(TableSource::Name(name)) => Some(TableSource::Name(name.clone())),
76        Some(TableSource::Subquery(inner)) => Some(TableSource::Subquery(Box::new(
77            parameterize_query_expr_inner(inner, next_index)?,
78        ))),
79        None => None,
80    };
81
82    let select_items = query
83        .select_items
84        .iter()
85        .map(|item| parameterize_select_item(item, next_index))
86        .collect::<Option<Vec<_>>>()?;
87
88    let where_expr = query
89        .where_expr
90        .as_ref()
91        .map(|expr| parameterize_expr(expr, next_index))
92        .or_else(|| {
93            query
94                .filter
95                .as_ref()
96                .map(|filter| parameterize_expr(&filter_to_expr(filter), next_index))
97        });
98
99    let group_by_exprs = if !query.group_by_exprs.is_empty() {
100        query
101            .group_by_exprs
102            .iter()
103            .map(|expr| parameterize_expr(expr, next_index))
104            .collect()
105    } else {
106        Vec::new()
107    };
108
109    let having_expr = query
110        .having_expr
111        .as_ref()
112        .map(|expr| parameterize_expr(expr, next_index))
113        .or_else(|| {
114            query
115                .having
116                .as_ref()
117                .map(|filter| parameterize_expr(&filter_to_expr(filter), next_index))
118        });
119
120    let order_by = query
121        .order_by
122        .iter()
123        .map(|clause| parameterize_order_by(clause, next_index))
124        .collect::<Option<Vec<_>>>()?;
125
126    Some(TableQuery {
127        table: query.table.clone(),
128        source,
129        alias: query.alias.clone(),
130        select_items,
131        columns: Vec::new(),
132        where_expr,
133        filter: None,
134        group_by_exprs,
135        group_by: Vec::new(),
136        having_expr,
137        having: None,
138        order_by,
139        limit: query.limit,
140        offset: query.offset,
141        expand: query.expand.clone(),
142        as_of: query.as_of.clone(),
143    })
144}
145
146fn parameterize_select_item(item: &SelectItem, next_index: &mut usize) -> Option<SelectItem> {
147    match item {
148        SelectItem::Wildcard => Some(SelectItem::Wildcard),
149        SelectItem::Expr { expr, alias } => Some(SelectItem::Expr {
150            expr: parameterize_expr(expr, next_index),
151            alias: alias.clone(),
152        }),
153    }
154}
155
156fn parameterize_order_by(clause: &OrderByClause, next_index: &mut usize) -> Option<OrderByClause> {
157    Some(OrderByClause {
158        field: clause.field.clone(),
159        expr: clause
160            .expr
161            .as_ref()
162            .map(|expr| parameterize_expr(expr, next_index)),
163        ascending: clause.ascending,
164        nulls_first: clause.nulls_first,
165    })
166}
167
168fn parameterize_expr(expr: &Expr, next_index: &mut usize) -> Expr {
169    match expr {
170        Expr::Literal { value, span } => {
171            let index = *next_index;
172            *next_index += 1;
173            let _ = value;
174            Expr::Parameter { index, span: *span }
175        }
176        Expr::Column { .. } | Expr::Parameter { .. } => expr.clone(),
177        Expr::BinaryOp { op, lhs, rhs, span } => Expr::BinaryOp {
178            op: *op,
179            lhs: Box::new(parameterize_expr(lhs, next_index)),
180            rhs: Box::new(parameterize_expr(rhs, next_index)),
181            span: *span,
182        },
183        Expr::UnaryOp { op, operand, span } => Expr::UnaryOp {
184            op: *op,
185            operand: Box::new(parameterize_expr(operand, next_index)),
186            span: *span,
187        },
188        Expr::Cast {
189            inner,
190            target,
191            span,
192        } => Expr::Cast {
193            inner: Box::new(parameterize_expr(inner, next_index)),
194            target: *target,
195            span: *span,
196        },
197        Expr::FunctionCall { name, args, span } => Expr::FunctionCall {
198            name: name.clone(),
199            args: args
200                .iter()
201                .map(|arg| parameterize_expr(arg, next_index))
202                .collect(),
203            span: *span,
204        },
205        Expr::Case {
206            branches,
207            else_,
208            span,
209        } => Expr::Case {
210            branches: branches
211                .iter()
212                .map(|(cond, value)| {
213                    (
214                        parameterize_expr(cond, next_index),
215                        parameterize_expr(value, next_index),
216                    )
217                })
218                .collect(),
219            else_: else_
220                .as_ref()
221                .map(|expr| Box::new(parameterize_expr(expr, next_index))),
222            span: *span,
223        },
224        Expr::IsNull {
225            operand,
226            negated,
227            span,
228        } => Expr::IsNull {
229            operand: Box::new(parameterize_expr(operand, next_index)),
230            negated: *negated,
231            span: *span,
232        },
233        Expr::InList {
234            target,
235            values,
236            negated,
237            span,
238        } => Expr::InList {
239            target: Box::new(parameterize_expr(target, next_index)),
240            values: values
241                .iter()
242                .map(|value| parameterize_expr(value, next_index))
243                .collect(),
244            negated: *negated,
245            span: *span,
246        },
247        Expr::Between {
248            target,
249            low,
250            high,
251            negated,
252            span,
253        } => Expr::Between {
254            target: Box::new(parameterize_expr(target, next_index)),
255            low: Box::new(parameterize_expr(low, next_index)),
256            high: Box::new(parameterize_expr(high, next_index)),
257            negated: *negated,
258            span: *span,
259        },
260    }
261}
262
263fn bind_query_expr_inner(expr: &QueryExpr, binds: &[Value]) -> Option<QueryExpr> {
264    match expr {
265        QueryExpr::Table(query) => Some(QueryExpr::Table(bind_table_query(query, binds)?)),
266        QueryExpr::Join(query) => Some(QueryExpr::Join(bind_join_query(query, binds)?)),
267        QueryExpr::Graph(query) => Some(QueryExpr::Graph(bind_graph_query(query, binds)?)),
268        QueryExpr::Path(query) => Some(QueryExpr::Path(bind_path_query(query, binds)?)),
269        QueryExpr::Vector(query) => Some(QueryExpr::Vector(bind_vector_query(query, binds)?)),
270        QueryExpr::Hybrid(query) => Some(QueryExpr::Hybrid(bind_hybrid_query(query, binds)?)),
271        _ => None,
272    }
273}
274
275fn parameterize_vector_query(query: &VectorQuery, next_index: &mut usize) -> Option<VectorQuery> {
276    Some(VectorQuery {
277        alias: query.alias.clone(),
278        collection: query.collection.clone(),
279        query_vector: parameterize_vector_source(&query.query_vector, next_index)?,
280        k: query.k,
281        filter: query
282            .filter
283            .as_ref()
284            .map(|filter| parameterize_metadata_filter(filter, next_index)),
285        metric: query.metric,
286        include_vectors: query.include_vectors,
287        include_metadata: query.include_metadata,
288        threshold: query
289            .threshold
290            .map(|_| encode_f32_placeholder(allocate_param_index(next_index))),
291    })
292}
293
294fn bind_vector_query(query: &VectorQuery, binds: &[Value]) -> Option<VectorQuery> {
295    Some(VectorQuery {
296        alias: query.alias.clone(),
297        collection: query.collection.clone(),
298        query_vector: bind_vector_source(&query.query_vector, binds)?,
299        k: query.k,
300        filter: query
301            .filter
302            .as_ref()
303            .and_then(|filter| bind_metadata_filter(filter, binds)),
304        metric: query.metric,
305        include_vectors: query.include_vectors,
306        include_metadata: query.include_metadata,
307        threshold: query
308            .threshold
309            .and_then(|value| bind_placeholder_f32(value, binds)),
310    })
311}
312
313fn parameterize_hybrid_query(query: &HybridQuery, next_index: &mut usize) -> Option<HybridQuery> {
314    Some(HybridQuery {
315        alias: query.alias.clone(),
316        structured: Box::new(parameterize_query_expr_inner(
317            &query.structured,
318            next_index,
319        )?),
320        vector: parameterize_vector_query(&query.vector, next_index)?,
321        fusion: parameterize_fusion_strategy(&query.fusion, next_index),
322        limit: query.limit,
323    })
324}
325
326fn bind_hybrid_query(query: &HybridQuery, binds: &[Value]) -> Option<HybridQuery> {
327    Some(HybridQuery {
328        alias: query.alias.clone(),
329        structured: Box::new(bind_query_expr_inner(&query.structured, binds)?),
330        vector: bind_vector_query(&query.vector, binds)?,
331        fusion: bind_fusion_strategy(&query.fusion, binds)?,
332        limit: query.limit,
333    })
334}
335
336fn parameterize_vector_source(
337    source: &VectorSource,
338    next_index: &mut usize,
339) -> Option<VectorSource> {
340    match source {
341        VectorSource::Literal(values) => Some(VectorSource::Literal(
342            values
343                .iter()
344                .map(|_| encode_f32_placeholder(allocate_param_index(next_index)))
345                .collect(),
346        )),
347        VectorSource::Text(_) => Some(VectorSource::Text(format!(
348            "{VECTOR_TEXT_PARAM_PREFIX}{}",
349            allocate_param_index(next_index)
350        ))),
351        VectorSource::Reference { collection, .. } => Some(VectorSource::Reference {
352            collection: format!(
353                "{VECTOR_REF_ID_PREFIX}{}:{collection}",
354                allocate_param_index(next_index)
355            ),
356            vector_id: 0,
357        }),
358        VectorSource::Subquery(expr) => Some(VectorSource::Subquery(Box::new(
359            parameterize_query_expr_inner(expr, next_index)?,
360        ))),
361    }
362}
363
364fn bind_vector_source(source: &VectorSource, binds: &[Value]) -> Option<VectorSource> {
365    match source {
366        VectorSource::Literal(values) => Some(VectorSource::Literal(
367            values
368                .iter()
369                .map(|value| bind_placeholder_f32(*value, binds))
370                .collect::<Option<Vec<_>>>()?,
371        )),
372        VectorSource::Text(text) => {
373            if let Some(index) = parse_placeholder_index(text, VECTOR_TEXT_PARAM_PREFIX) {
374                Some(VectorSource::Text(bind_value_to_string(binds.get(index)?)?))
375            } else {
376                Some(VectorSource::Text(text.clone()))
377            }
378        }
379        VectorSource::Reference {
380            collection,
381            vector_id,
382        } => {
383            if let Some((index, original_collection)) =
384                parse_prefixed_index_with_suffix(collection, VECTOR_REF_ID_PREFIX)
385            {
386                Some(VectorSource::Reference {
387                    collection: original_collection.to_string(),
388                    vector_id: bind_value_to_u64(binds.get(index)?)?,
389                })
390            } else {
391                Some(VectorSource::Reference {
392                    collection: collection.clone(),
393                    vector_id: *vector_id,
394                })
395            }
396        }
397        VectorSource::Subquery(expr) => Some(VectorSource::Subquery(Box::new(
398            bind_query_expr_inner(expr, binds)?,
399        ))),
400    }
401}
402
403fn parameterize_fusion_strategy(fusion: &FusionStrategy, next_index: &mut usize) -> FusionStrategy {
404    match fusion {
405        FusionStrategy::Rerank { .. } => FusionStrategy::Rerank {
406            weight: encode_f32_placeholder(allocate_param_index(next_index)),
407        },
408        FusionStrategy::FilterThenSearch => FusionStrategy::FilterThenSearch,
409        FusionStrategy::SearchThenFilter => FusionStrategy::SearchThenFilter,
410        FusionStrategy::RRF { .. } => FusionStrategy::RRF {
411            k: encode_u32_placeholder(allocate_param_index(next_index)),
412        },
413        FusionStrategy::Intersection => FusionStrategy::Intersection,
414        FusionStrategy::Union { .. } => FusionStrategy::Union {
415            structured_weight: encode_f32_placeholder(allocate_param_index(next_index)),
416            vector_weight: encode_f32_placeholder(allocate_param_index(next_index)),
417        },
418    }
419}
420
421fn bind_fusion_strategy(fusion: &FusionStrategy, binds: &[Value]) -> Option<FusionStrategy> {
422    match fusion {
423        FusionStrategy::Rerank { weight } => Some(FusionStrategy::Rerank {
424            weight: bind_placeholder_f32(*weight, binds)?,
425        }),
426        FusionStrategy::FilterThenSearch => Some(FusionStrategy::FilterThenSearch),
427        FusionStrategy::SearchThenFilter => Some(FusionStrategy::SearchThenFilter),
428        FusionStrategy::RRF { k } => Some(FusionStrategy::RRF {
429            k: bind_placeholder_u32(*k, binds)?,
430        }),
431        FusionStrategy::Intersection => Some(FusionStrategy::Intersection),
432        FusionStrategy::Union {
433            structured_weight,
434            vector_weight,
435        } => Some(FusionStrategy::Union {
436            structured_weight: bind_placeholder_f32(*structured_weight, binds)?,
437            vector_weight: bind_placeholder_f32(*vector_weight, binds)?,
438        }),
439    }
440}
441
442fn parameterize_metadata_filter(filter: &MetadataFilter, next_index: &mut usize) -> MetadataFilter {
443    match filter {
444        MetadataFilter::Eq(key, value) => {
445            MetadataFilter::Eq(key.clone(), parameterize_metadata_value(value, next_index))
446        }
447        MetadataFilter::Ne(key, value) => {
448            MetadataFilter::Ne(key.clone(), parameterize_metadata_value(value, next_index))
449        }
450        MetadataFilter::Gt(key, value) => {
451            MetadataFilter::Gt(key.clone(), parameterize_metadata_value(value, next_index))
452        }
453        MetadataFilter::Gte(key, value) => {
454            MetadataFilter::Gte(key.clone(), parameterize_metadata_value(value, next_index))
455        }
456        MetadataFilter::Lt(key, value) => {
457            MetadataFilter::Lt(key.clone(), parameterize_metadata_value(value, next_index))
458        }
459        MetadataFilter::Lte(key, value) => {
460            MetadataFilter::Lte(key.clone(), parameterize_metadata_value(value, next_index))
461        }
462        MetadataFilter::In(key, values) => MetadataFilter::In(
463            key.clone(),
464            values
465                .iter()
466                .map(|value| parameterize_metadata_value(value, next_index))
467                .collect(),
468        ),
469        MetadataFilter::NotIn(key, values) => MetadataFilter::NotIn(
470            key.clone(),
471            values
472                .iter()
473                .map(|value| parameterize_metadata_value(value, next_index))
474                .collect(),
475        ),
476        MetadataFilter::Contains(_, _) => MetadataFilter::Contains(
477            match filter {
478                MetadataFilter::Contains(key, _) => key.clone(),
479                _ => unreachable!(),
480            },
481            format!("{STRING_PARAM_PREFIX}{}", allocate_param_index(next_index)),
482        ),
483        MetadataFilter::StartsWith(_, _) => MetadataFilter::StartsWith(
484            match filter {
485                MetadataFilter::StartsWith(key, _) => key.clone(),
486                _ => unreachable!(),
487            },
488            format!("{STRING_PARAM_PREFIX}{}", allocate_param_index(next_index)),
489        ),
490        MetadataFilter::EndsWith(_, _) => MetadataFilter::EndsWith(
491            match filter {
492                MetadataFilter::EndsWith(key, _) => key.clone(),
493                _ => unreachable!(),
494            },
495            format!("{STRING_PARAM_PREFIX}{}", allocate_param_index(next_index)),
496        ),
497        MetadataFilter::Exists(key) => MetadataFilter::Exists(key.clone()),
498        MetadataFilter::NotExists(key) => MetadataFilter::NotExists(key.clone()),
499        MetadataFilter::And(filters) => MetadataFilter::And(
500            filters
501                .iter()
502                .map(|filter| parameterize_metadata_filter(filter, next_index))
503                .collect(),
504        ),
505        MetadataFilter::Or(filters) => MetadataFilter::Or(
506            filters
507                .iter()
508                .map(|filter| parameterize_metadata_filter(filter, next_index))
509                .collect(),
510        ),
511        MetadataFilter::Not(inner) => {
512            MetadataFilter::Not(Box::new(parameterize_metadata_filter(inner, next_index)))
513        }
514    }
515}
516
517fn bind_metadata_filter(filter: &MetadataFilter, binds: &[Value]) -> Option<MetadataFilter> {
518    match filter {
519        MetadataFilter::Eq(key, value) => Some(MetadataFilter::Eq(
520            key.clone(),
521            bind_metadata_value(value, binds)?,
522        )),
523        MetadataFilter::Ne(key, value) => Some(MetadataFilter::Ne(
524            key.clone(),
525            bind_metadata_value(value, binds)?,
526        )),
527        MetadataFilter::Gt(key, value) => Some(MetadataFilter::Gt(
528            key.clone(),
529            bind_metadata_value(value, binds)?,
530        )),
531        MetadataFilter::Gte(key, value) => Some(MetadataFilter::Gte(
532            key.clone(),
533            bind_metadata_value(value, binds)?,
534        )),
535        MetadataFilter::Lt(key, value) => Some(MetadataFilter::Lt(
536            key.clone(),
537            bind_metadata_value(value, binds)?,
538        )),
539        MetadataFilter::Lte(key, value) => Some(MetadataFilter::Lte(
540            key.clone(),
541            bind_metadata_value(value, binds)?,
542        )),
543        MetadataFilter::In(key, values) => Some(MetadataFilter::In(
544            key.clone(),
545            values
546                .iter()
547                .map(|value| bind_metadata_value(value, binds))
548                .collect::<Option<Vec<_>>>()?,
549        )),
550        MetadataFilter::NotIn(key, values) => Some(MetadataFilter::NotIn(
551            key.clone(),
552            values
553                .iter()
554                .map(|value| bind_metadata_value(value, binds))
555                .collect::<Option<Vec<_>>>()?,
556        )),
557        MetadataFilter::Contains(key, value) => Some(MetadataFilter::Contains(
558            key.clone(),
559            bind_placeholder_string(value, binds)?.unwrap_or_default(),
560        )),
561        MetadataFilter::StartsWith(key, value) => Some(MetadataFilter::StartsWith(
562            key.clone(),
563            bind_placeholder_string(value, binds)?.unwrap_or_default(),
564        )),
565        MetadataFilter::EndsWith(key, value) => Some(MetadataFilter::EndsWith(
566            key.clone(),
567            bind_placeholder_string(value, binds)?.unwrap_or_default(),
568        )),
569        MetadataFilter::Exists(key) => Some(MetadataFilter::Exists(key.clone())),
570        MetadataFilter::NotExists(key) => Some(MetadataFilter::NotExists(key.clone())),
571        MetadataFilter::And(filters) => Some(MetadataFilter::And(
572            filters
573                .iter()
574                .map(|filter| bind_metadata_filter(filter, binds))
575                .collect::<Option<Vec<_>>>()?,
576        )),
577        MetadataFilter::Or(filters) => Some(MetadataFilter::Or(
578            filters
579                .iter()
580                .map(|filter| bind_metadata_filter(filter, binds))
581                .collect::<Option<Vec<_>>>()?,
582        )),
583        MetadataFilter::Not(inner) => Some(MetadataFilter::Not(Box::new(bind_metadata_filter(
584            inner, binds,
585        )?))),
586    }
587}
588
589fn parameterize_metadata_value(_value: &MetadataValue, next_index: &mut usize) -> MetadataValue {
590    MetadataValue::String(format!(
591        "{METADATA_VALUE_PARAM_PREFIX}{}",
592        allocate_param_index(next_index)
593    ))
594}
595
596fn bind_metadata_value(value: &MetadataValue, binds: &[Value]) -> Option<MetadataValue> {
597    match value {
598        MetadataValue::String(text) => {
599            if let Some(index) = parse_placeholder_index(text, METADATA_VALUE_PARAM_PREFIX) {
600                Some(bind_value_to_metadata_value(binds.get(index)?)?)
601            } else {
602                Some(MetadataValue::String(text.clone()))
603            }
604        }
605        other => Some(other.clone()),
606    }
607}
608
609fn parameterize_join_query(query: &JoinQuery, next_index: &mut usize) -> Option<JoinQuery> {
610    Some(JoinQuery {
611        left: Box::new(parameterize_query_expr_inner(&query.left, next_index)?),
612        right: Box::new(parameterize_query_expr_inner(&query.right, next_index)?),
613        join_type: query.join_type,
614        on: query.on.clone(),
615        filter: query
616            .filter
617            .as_ref()
618            .map(|filter| parameterize_filter(filter, next_index)),
619        order_by: query
620            .order_by
621            .iter()
622            .map(|clause| parameterize_order_by(clause, next_index))
623            .collect::<Option<Vec<_>>>()?,
624        limit: query.limit,
625        offset: query.offset,
626        return_items: query
627            .return_items
628            .iter()
629            .map(|item| parameterize_select_item(item, next_index))
630            .collect::<Option<Vec<_>>>()?,
631        return_: query
632            .return_
633            .iter()
634            .map(|projection| parameterize_projection(projection, next_index))
635            .collect::<Option<Vec<_>>>()?,
636    })
637}
638
639fn bind_join_query(query: &JoinQuery, binds: &[Value]) -> Option<JoinQuery> {
640    Some(JoinQuery {
641        left: Box::new(bind_query_expr_inner(&query.left, binds)?),
642        right: Box::new(bind_query_expr_inner(&query.right, binds)?),
643        join_type: query.join_type,
644        on: query.on.clone(),
645        filter: query
646            .filter
647            .as_ref()
648            .and_then(|filter| bind_filter(filter, binds)),
649        order_by: query
650            .order_by
651            .iter()
652            .map(|clause| bind_order_by(clause, binds))
653            .collect::<Option<Vec<_>>>()?,
654        limit: query.limit,
655        offset: query.offset,
656        return_items: query
657            .return_items
658            .iter()
659            .map(|item| bind_select_item(item, binds))
660            .collect::<Option<Vec<_>>>()?,
661        return_: query
662            .return_
663            .iter()
664            .map(|projection| bind_projection(projection, binds))
665            .collect::<Option<Vec<_>>>()?,
666    })
667}
668
669fn parameterize_graph_query(query: &GraphQuery, next_index: &mut usize) -> Option<GraphQuery> {
670    Some(GraphQuery {
671        alias: query.alias.clone(),
672        pattern: parameterize_graph_pattern(&query.pattern, next_index),
673        filter: query
674            .filter
675            .as_ref()
676            .map(|filter| parameterize_filter(filter, next_index)),
677        return_: query
678            .return_
679            .iter()
680            .map(|projection| parameterize_projection(projection, next_index))
681            .collect::<Option<Vec<_>>>()?,
682    })
683}
684
685fn bind_graph_query(query: &GraphQuery, binds: &[Value]) -> Option<GraphQuery> {
686    Some(GraphQuery {
687        alias: query.alias.clone(),
688        pattern: bind_graph_pattern(&query.pattern, binds)?,
689        filter: query
690            .filter
691            .as_ref()
692            .and_then(|filter| bind_filter(filter, binds)),
693        return_: query
694            .return_
695            .iter()
696            .map(|projection| bind_projection(projection, binds))
697            .collect::<Option<Vec<_>>>()?,
698    })
699}
700
701fn parameterize_path_query(query: &PathQuery, next_index: &mut usize) -> Option<PathQuery> {
702    Some(PathQuery {
703        alias: query.alias.clone(),
704        from: parameterize_node_selector(&query.from, next_index),
705        to: parameterize_node_selector(&query.to, next_index),
706        via: query.via.clone(),
707        max_length: query.max_length,
708        filter: query
709            .filter
710            .as_ref()
711            .map(|filter| parameterize_filter(filter, next_index)),
712        return_: query
713            .return_
714            .iter()
715            .map(|projection| parameterize_projection(projection, next_index))
716            .collect::<Option<Vec<_>>>()?,
717    })
718}
719
720fn bind_path_query(query: &PathQuery, binds: &[Value]) -> Option<PathQuery> {
721    Some(PathQuery {
722        alias: query.alias.clone(),
723        from: bind_node_selector(&query.from, binds)?,
724        to: bind_node_selector(&query.to, binds)?,
725        via: query.via.clone(),
726        max_length: query.max_length,
727        filter: query
728            .filter
729            .as_ref()
730            .and_then(|filter| bind_filter(filter, binds)),
731        return_: query
732            .return_
733            .iter()
734            .map(|projection| bind_projection(projection, binds))
735            .collect::<Option<Vec<_>>>()?,
736    })
737}
738
739fn parameterize_filter(
740    filter: &crate::storage::query::ast::Filter,
741    next_index: &mut usize,
742) -> crate::storage::query::ast::Filter {
743    expr_to_filter(&parameterize_expr(&filter_to_expr(filter), next_index))
744}
745
746fn bind_filter(
747    filter: &crate::storage::query::ast::Filter,
748    binds: &[Value],
749) -> Option<crate::storage::query::ast::Filter> {
750    Some(expr_to_filter(&bind_expr(&filter_to_expr(filter), binds)?))
751}
752
753fn parameterize_projection(projection: &Projection, next_index: &mut usize) -> Option<Projection> {
754    match projection {
755        Projection::All => Some(Projection::All),
756        Projection::Column(column) => {
757            Some(parameterize_projection_column(column, None, next_index))
758        }
759        Projection::Alias(column, alias) => Some(parameterize_projection_column(
760            column,
761            Some(alias.as_str()),
762            next_index,
763        )),
764        Projection::Function(name, args) => Some(Projection::Function(
765            name.clone(),
766            args.iter()
767                .map(|arg| parameterize_projection(arg, next_index))
768                .collect::<Option<Vec<_>>>()?,
769        )),
770        Projection::Expression(filter, alias) => Some(Projection::Expression(
771            Box::new(parameterize_filter(filter, next_index)),
772            alias.clone(),
773        )),
774        Projection::Field(field, alias) => Some(Projection::Field(field.clone(), alias.clone())),
775    }
776}
777
778fn bind_projection(projection: &Projection, binds: &[Value]) -> Option<Projection> {
779    match projection {
780        Projection::All => Some(Projection::All),
781        Projection::Column(column) => bind_projection_column(column, None, binds),
782        Projection::Alias(column, alias) => {
783            bind_projection_column(column, Some(alias.as_str()), binds)
784        }
785        Projection::Function(name, args) => Some(Projection::Function(
786            name.clone(),
787            args.iter()
788                .map(|arg| bind_projection(arg, binds))
789                .collect::<Option<Vec<_>>>()?,
790        )),
791        Projection::Expression(filter, alias) => Some(Projection::Expression(
792            Box::new(bind_filter(filter, binds)?),
793            alias.clone(),
794        )),
795        Projection::Field(field, alias) => Some(Projection::Field(field.clone(), alias.clone())),
796    }
797}
798
799fn parameterize_projection_column(
800    column: &str,
801    alias: Option<&str>,
802    next_index: &mut usize,
803) -> Projection {
804    if column.starts_with("LIT:") {
805        let index = *next_index;
806        *next_index += 1;
807        let placeholder = format!("{PROJECTION_PARAM_PREFIX}{index}");
808        if let Some(alias) = alias {
809            Projection::Alias(placeholder, alias.to_string())
810        } else {
811            Projection::Column(placeholder)
812        }
813    } else if let Some(alias) = alias {
814        Projection::Alias(column.to_string(), alias.to_string())
815    } else {
816        Projection::Column(column.to_string())
817    }
818}
819
820fn bind_projection_column(
821    column: &str,
822    alias: Option<&str>,
823    binds: &[Value],
824) -> Option<Projection> {
825    if let Some(index) = parse_placeholder_index(column, PROJECTION_PARAM_PREFIX) {
826        let projection = projection_from_literal(binds.get(index)?)?;
827        Some(attach_projection_alias(projection, alias))
828    } else if let Some(alias) = alias {
829        Some(Projection::Alias(column.to_string(), alias.to_string()))
830    } else {
831        Some(Projection::Column(column.to_string()))
832    }
833}
834
835fn parameterize_graph_pattern(pattern: &GraphPattern, next_index: &mut usize) -> GraphPattern {
836    GraphPattern {
837        nodes: pattern
838            .nodes
839            .iter()
840            .map(|node| parameterize_node_pattern(node, next_index))
841            .collect(),
842        edges: pattern.edges.clone(),
843    }
844}
845
846fn bind_graph_pattern(pattern: &GraphPattern, binds: &[Value]) -> Option<GraphPattern> {
847    Some(GraphPattern {
848        nodes: pattern
849            .nodes
850            .iter()
851            .map(|node| bind_node_pattern(node, binds))
852            .collect::<Option<Vec<_>>>()?,
853        edges: pattern.edges.clone(),
854    })
855}
856
857fn parameterize_node_pattern(node: &NodePattern, next_index: &mut usize) -> NodePattern {
858    NodePattern {
859        alias: node.alias.clone(),
860        node_label: node.node_label.clone(),
861        properties: node
862            .properties
863            .iter()
864            .map(|property| parameterize_property_filter(property, next_index))
865            .collect(),
866    }
867}
868
869fn bind_node_pattern(node: &NodePattern, binds: &[Value]) -> Option<NodePattern> {
870    Some(NodePattern {
871        alias: node.alias.clone(),
872        node_label: node.node_label.clone(),
873        properties: node
874            .properties
875            .iter()
876            .map(|property| bind_property_filter(property, binds))
877            .collect::<Option<Vec<_>>>()?,
878    })
879}
880
881fn parameterize_property_filter(filter: &PropertyFilter, next_index: &mut usize) -> PropertyFilter {
882    PropertyFilter {
883        name: filter.name.clone(),
884        op: filter.op,
885        value: parameterize_value_placeholder(next_index),
886    }
887}
888
889fn bind_property_filter(filter: &PropertyFilter, binds: &[Value]) -> Option<PropertyFilter> {
890    Some(PropertyFilter {
891        name: filter.name.clone(),
892        op: filter.op,
893        value: bind_value_placeholder(&filter.value, binds)?,
894    })
895}
896
897fn parameterize_node_selector(selector: &NodeSelector, next_index: &mut usize) -> NodeSelector {
898    match selector {
899        NodeSelector::ById(_) => {
900            let index = *next_index;
901            *next_index += 1;
902            NodeSelector::ById(format!("{STRING_PARAM_PREFIX}{index}"))
903        }
904        NodeSelector::ByType { node_label, filter } => NodeSelector::ByType {
905            node_label: node_label.clone(),
906            filter: filter
907                .as_ref()
908                .map(|filter| parameterize_property_filter(filter, next_index)),
909        },
910        NodeSelector::ByRow { table, .. } => {
911            let index = *next_index;
912            *next_index += 1;
913            NodeSelector::ByRow {
914                table: format!("{ROW_SELECTOR_TABLE_PREFIX}{index}:{table}"),
915                row_id: 0,
916            }
917        }
918    }
919}
920
921fn bind_node_selector(selector: &NodeSelector, binds: &[Value]) -> Option<NodeSelector> {
922    match selector {
923        NodeSelector::ById(id) => {
924            if let Some(index) = parse_placeholder_index(id, STRING_PARAM_PREFIX) {
925                Some(NodeSelector::ById(bind_value_to_string(binds.get(index)?)?))
926            } else {
927                Some(NodeSelector::ById(id.clone()))
928            }
929        }
930        NodeSelector::ByType { node_label, filter } => Some(NodeSelector::ByType {
931            node_label: node_label.clone(),
932            filter: filter
933                .as_ref()
934                .and_then(|filter| bind_property_filter(filter, binds)),
935        }),
936        NodeSelector::ByRow { table, row_id } => {
937            if let Some((index, original_table)) = parse_row_selector_placeholder(table) {
938                Some(NodeSelector::ByRow {
939                    table: original_table.to_string(),
940                    row_id: bind_value_to_u64(binds.get(index)?)?,
941                })
942            } else {
943                Some(NodeSelector::ByRow {
944                    table: table.clone(),
945                    row_id: *row_id,
946                })
947            }
948        }
949    }
950}
951
952fn parameterize_value_placeholder(next_index: &mut usize) -> Value {
953    let index = *next_index;
954    *next_index += 1;
955    Value::text(format!("{VALUE_PARAM_PREFIX}{index}"))
956}
957
958fn bind_value_placeholder(value: &Value, binds: &[Value]) -> Option<Value> {
959    match value {
960        Value::Text(text) => {
961            if let Some(index) = parse_placeholder_index(text, VALUE_PARAM_PREFIX) {
962                binds.get(index).cloned()
963            } else {
964                Some(value.clone())
965            }
966        }
967        _ => Some(value.clone()),
968    }
969}
970
971fn attach_projection_alias(projection: Projection, alias: Option<&str>) -> Projection {
972    let Some(alias) = alias else {
973        return projection;
974    };
975    match projection {
976        Projection::Field(field, _) => Projection::Field(field, Some(alias.to_string())),
977        Projection::Expression(filter, _) => {
978            Projection::Expression(filter, Some(alias.to_string()))
979        }
980        Projection::Function(name, args) => {
981            if name.contains(':') {
982                Projection::Function(name, args)
983            } else {
984                Projection::Function(format!("{name}:{alias}"), args)
985            }
986        }
987        Projection::Column(column) => Projection::Alias(column, alias.to_string()),
988        Projection::Alias(column, _) => Projection::Alias(column, alias.to_string()),
989        Projection::All => Projection::All,
990    }
991}
992
993fn bind_value_to_string(value: &Value) -> Option<String> {
994    match value {
995        Value::Null => None,
996        _ => Some(value.to_string()),
997    }
998}
999
1000fn bind_placeholder_string(value: &str, binds: &[Value]) -> Option<Option<String>> {
1001    if let Some(index) = parse_placeholder_index(value, STRING_PARAM_PREFIX) {
1002        Some(bind_value_to_string(binds.get(index)?))
1003    } else {
1004        Some(Some(value.to_string()))
1005    }
1006}
1007
1008fn bind_value_to_u64(value: &Value) -> Option<u64> {
1009    match value {
1010        Value::UnsignedInteger(value) => Some(*value),
1011        Value::Integer(value) if *value >= 0 => Some(*value as u64),
1012        Value::BigInt(value) if *value >= 0 => Some(*value as u64),
1013        Value::Text(value) => value.parse().ok(),
1014        _ => None,
1015    }
1016}
1017
1018fn parse_placeholder_index(value: &str, prefix: &str) -> Option<usize> {
1019    value.strip_prefix(prefix)?.parse().ok()
1020}
1021
1022fn parse_prefixed_index_with_suffix<'a>(value: &'a str, prefix: &str) -> Option<(usize, &'a str)> {
1023    let rest = value.strip_prefix(prefix)?;
1024    let (index, suffix) = rest.split_once(':')?;
1025    Some((index.parse().ok()?, suffix))
1026}
1027
1028fn parse_row_selector_placeholder(value: &str) -> Option<(usize, &str)> {
1029    let rest = value.strip_prefix(ROW_SELECTOR_TABLE_PREFIX)?;
1030    let (index, table) = rest.split_once(':')?;
1031    Some((index.parse().ok()?, table))
1032}
1033
1034fn allocate_param_index(next_index: &mut usize) -> usize {
1035    let index = *next_index;
1036    *next_index += 1;
1037    index
1038}
1039
1040fn encode_f32_placeholder(index: usize) -> f32 {
1041    f32::from_bits(FLOAT32_PARAM_BITS_BASE | (index as u32 & 0x003f_ffff))
1042}
1043
1044fn decode_f32_placeholder(value: f32) -> Option<usize> {
1045    let bits = value.to_bits();
1046    if bits & FLOAT32_PARAM_BITS_BASE == FLOAT32_PARAM_BITS_BASE {
1047        Some((bits & 0x003f_ffff) as usize)
1048    } else {
1049        None
1050    }
1051}
1052
1053fn bind_placeholder_f32(value: f32, binds: &[Value]) -> Option<f32> {
1054    if let Some(index) = decode_f32_placeholder(value) {
1055        bind_value_to_f32(binds.get(index)?)
1056    } else {
1057        Some(value)
1058    }
1059}
1060
1061fn encode_u32_placeholder(index: usize) -> u32 {
1062    U32_PARAM_BASE | (index as u32 & 0x000f_ffff)
1063}
1064
1065fn decode_u32_placeholder(value: u32) -> Option<usize> {
1066    if value & 0xfff0_0000 == U32_PARAM_BASE {
1067        Some((value & 0x000f_ffff) as usize)
1068    } else {
1069        None
1070    }
1071}
1072
1073fn bind_placeholder_u32(value: u32, binds: &[Value]) -> Option<u32> {
1074    if let Some(index) = decode_u32_placeholder(value) {
1075        bind_value_to_u64(binds.get(index)?).and_then(|value| u32::try_from(value).ok())
1076    } else {
1077        Some(value)
1078    }
1079}
1080
1081fn bind_value_to_f32(value: &Value) -> Option<f32> {
1082    match value {
1083        Value::Float(value) => Some(*value as f32),
1084        Value::Integer(value) => Some(*value as f32),
1085        Value::UnsignedInteger(value) => Some(*value as f32),
1086        Value::BigInt(value) => Some(*value as f32),
1087        Value::Text(value) => value.parse().ok(),
1088        _ => None,
1089    }
1090}
1091
1092fn bind_value_to_metadata_value(value: &Value) -> Option<MetadataValue> {
1093    match value {
1094        Value::Text(value) => Some(MetadataValue::String(value.to_string())),
1095        Value::Integer(value) => Some(MetadataValue::Integer(*value)),
1096        Value::UnsignedInteger(value) => i64::try_from(*value).ok().map(MetadataValue::Integer),
1097        Value::BigInt(value) => Some(MetadataValue::Integer(*value)),
1098        Value::Float(value) => Some(MetadataValue::Float(*value)),
1099        Value::Boolean(value) => Some(MetadataValue::Bool(*value)),
1100        Value::Null => Some(MetadataValue::Null),
1101        _ => None,
1102    }
1103}
1104
1105fn bind_table_query(query: &TableQuery, binds: &[Value]) -> Option<TableQuery> {
1106    let source = match &query.source {
1107        Some(TableSource::Name(name)) => Some(TableSource::Name(name.clone())),
1108        Some(TableSource::Subquery(inner)) => Some(TableSource::Subquery(Box::new(
1109            bind_query_expr_inner(inner, binds)?,
1110        ))),
1111        None => None,
1112    };
1113
1114    Some(TableQuery {
1115        table: query.table.clone(),
1116        source,
1117        alias: query.alias.clone(),
1118        select_items: query
1119            .select_items
1120            .iter()
1121            .map(|item| bind_select_item(item, binds))
1122            .collect::<Option<Vec<_>>>()?,
1123        columns: Vec::new(),
1124        where_expr: query
1125            .where_expr
1126            .as_ref()
1127            .and_then(|expr| bind_expr(expr, binds)),
1128        filter: None,
1129        group_by_exprs: query
1130            .group_by_exprs
1131            .iter()
1132            .map(|expr| bind_expr(expr, binds))
1133            .collect::<Option<Vec<_>>>()?,
1134        group_by: Vec::new(),
1135        having_expr: query
1136            .having_expr
1137            .as_ref()
1138            .and_then(|expr| bind_expr(expr, binds)),
1139        having: None,
1140        order_by: query
1141            .order_by
1142            .iter()
1143            .map(|clause| bind_order_by(clause, binds))
1144            .collect::<Option<Vec<_>>>()?,
1145        limit: query.limit,
1146        offset: query.offset,
1147        expand: query.expand.clone(),
1148        as_of: query.as_of.clone(),
1149    })
1150}
1151
1152fn bind_select_item(item: &SelectItem, binds: &[Value]) -> Option<SelectItem> {
1153    match item {
1154        SelectItem::Wildcard => Some(SelectItem::Wildcard),
1155        SelectItem::Expr { expr, alias } => Some(SelectItem::Expr {
1156            expr: bind_expr(expr, binds)?,
1157            alias: alias.clone(),
1158        }),
1159    }
1160}
1161
1162fn bind_order_by(clause: &OrderByClause, binds: &[Value]) -> Option<OrderByClause> {
1163    Some(OrderByClause {
1164        field: clause.field.clone(),
1165        expr: clause.expr.as_ref().and_then(|expr| bind_expr(expr, binds)),
1166        ascending: clause.ascending,
1167        nulls_first: clause.nulls_first,
1168    })
1169}
1170
1171fn bind_expr(expr: &Expr, binds: &[Value]) -> Option<Expr> {
1172    match expr {
1173        Expr::Literal { .. } | Expr::Column { .. } => Some(expr.clone()),
1174        Expr::Parameter { index, span } => Some(Expr::Literal {
1175            value: binds.get(*index)?.clone(),
1176            span: *span,
1177        }),
1178        Expr::BinaryOp { op, lhs, rhs, span } => Some(Expr::BinaryOp {
1179            op: *op,
1180            lhs: Box::new(bind_expr(lhs, binds)?),
1181            rhs: Box::new(bind_expr(rhs, binds)?),
1182            span: *span,
1183        }),
1184        Expr::UnaryOp { op, operand, span } => Some(Expr::UnaryOp {
1185            op: *op,
1186            operand: Box::new(bind_expr(operand, binds)?),
1187            span: *span,
1188        }),
1189        Expr::Cast {
1190            inner,
1191            target,
1192            span,
1193        } => Some(Expr::Cast {
1194            inner: Box::new(bind_expr(inner, binds)?),
1195            target: *target,
1196            span: *span,
1197        }),
1198        Expr::FunctionCall { name, args, span } => Some(Expr::FunctionCall {
1199            name: name.clone(),
1200            args: args
1201                .iter()
1202                .map(|arg| bind_expr(arg, binds))
1203                .collect::<Option<Vec<_>>>()?,
1204            span: *span,
1205        }),
1206        Expr::Case {
1207            branches,
1208            else_,
1209            span,
1210        } => Some(Expr::Case {
1211            branches: branches
1212                .iter()
1213                .map(|(cond, value)| Some((bind_expr(cond, binds)?, bind_expr(value, binds)?)))
1214                .collect::<Option<Vec<_>>>()?,
1215            else_: else_
1216                .as_ref()
1217                .and_then(|expr| bind_expr(expr, binds).map(Box::new)),
1218            span: *span,
1219        }),
1220        Expr::IsNull {
1221            operand,
1222            negated,
1223            span,
1224        } => Some(Expr::IsNull {
1225            operand: Box::new(bind_expr(operand, binds)?),
1226            negated: *negated,
1227            span: *span,
1228        }),
1229        Expr::InList {
1230            target,
1231            values,
1232            negated,
1233            span,
1234        } => Some(Expr::InList {
1235            target: Box::new(bind_expr(target, binds)?),
1236            values: values
1237                .iter()
1238                .map(|value| bind_expr(value, binds))
1239                .collect::<Option<Vec<_>>>()?,
1240            negated: *negated,
1241            span: *span,
1242        }),
1243        Expr::Between {
1244            target,
1245            low,
1246            high,
1247            negated,
1248            span,
1249        } => Some(Expr::Between {
1250            target: Box::new(bind_expr(target, binds)?),
1251            low: Box::new(bind_expr(low, binds)?),
1252            high: Box::new(bind_expr(high, binds)?),
1253            negated: *negated,
1254            span: *span,
1255        }),
1256    }
1257}
1258
1259#[cfg(test)]
1260mod tests {
1261    use super::*;
1262    use crate::storage::query::ast::{BinOp, FieldRef, SelectItem, TableQuery};
1263
1264    #[test]
1265    fn table_shape_round_trips_with_new_binds() {
1266        let query = QueryExpr::Table(TableQuery {
1267            table: "users".to_string(),
1268            source: None,
1269            alias: None,
1270            select_items: vec![SelectItem::Expr {
1271                expr: Expr::Column {
1272                    field: FieldRef::TableColumn {
1273                        table: String::new(),
1274                        column: "name".to_string(),
1275                    },
1276                    span: crate::storage::query::ast::Span::synthetic(),
1277                },
1278                alias: None,
1279            }],
1280            columns: Vec::new(),
1281            where_expr: Some(Expr::BinaryOp {
1282                op: BinOp::Eq,
1283                lhs: Box::new(Expr::Column {
1284                    field: FieldRef::TableColumn {
1285                        table: String::new(),
1286                        column: "age".to_string(),
1287                    },
1288                    span: crate::storage::query::ast::Span::synthetic(),
1289                }),
1290                rhs: Box::new(Expr::Literal {
1291                    value: Value::Integer(18),
1292                    span: crate::storage::query::ast::Span::synthetic(),
1293                }),
1294                span: crate::storage::query::ast::Span::synthetic(),
1295            }),
1296            filter: None,
1297            group_by_exprs: Vec::new(),
1298            group_by: Vec::new(),
1299            having_expr: None,
1300            having: None,
1301            order_by: Vec::new(),
1302            limit: None,
1303            offset: None,
1304            expand: None,
1305            as_of: None,
1306        });
1307
1308        let prepared = parameterize_query_expr(&query).unwrap();
1309        assert_eq!(prepared.parameter_count, 1);
1310
1311        let rebound = bind_parameterized_query(
1312            &prepared.shape,
1313            &[Value::Integer(42)],
1314            prepared.parameter_count,
1315        )
1316        .unwrap();
1317
1318        let QueryExpr::Table(bound_table) = rebound else {
1319            panic!("expected table query");
1320        };
1321        match bound_table.where_expr.unwrap() {
1322            Expr::BinaryOp { rhs, .. } => match *rhs {
1323                Expr::Literal { value, .. } => assert_eq!(value, Value::Integer(42)),
1324                other => panic!("expected rebound literal, got {other:?}"),
1325            },
1326            other => panic!("expected binary op, got {other:?}"),
1327        }
1328    }
1329}