1use crate::storage::query::ast::{
14 CompareOp, FieldRef, Filter as AstFilter, JoinQuery, Projection, QueryExpr,
15};
16use crate::storage::query::sql_lowering::{
17 effective_graph_filter, effective_join_filter, effective_table_filter, effective_vector_filter,
18};
19use crate::storage::schema::Value;
20
21#[derive(Debug, Clone, Default)]
23pub struct RewriteContext {
24 pub property_cache: Vec<CachedProperty>,
26 pub errors: Vec<String>,
28 pub warnings: Vec<String>,
30 pub stats: RewriteStats,
32}
33
34#[derive(Debug, Clone)]
36pub struct CachedProperty {
37 pub source: String,
39 pub property: String,
41 pub cached_value: Option<String>,
43}
44
45#[derive(Debug, Clone, Default)]
47pub struct RewriteStats {
48 pub filters_simplified: u32,
50 pub predicates_pushed: u32,
52 pub properties_cached: u32,
54 pub expressions_normalized: u32,
56}
57
58pub trait RewriteRule: Send + Sync {
60 fn name(&self) -> &str;
62
63 fn apply(&self, query: QueryExpr, ctx: &mut RewriteContext) -> QueryExpr;
65
66 fn is_applicable(&self, query: &QueryExpr) -> bool;
68}
69
70pub struct QueryRewriter {
72 rules: Vec<Box<dyn RewriteRule>>,
74 max_iterations: usize,
76}
77
78impl QueryRewriter {
79 pub fn new() -> Self {
81 let rules: Vec<Box<dyn RewriteRule>> = vec![
82 Box::new(NormalizeRule),
83 Box::new(SimplifyFiltersRule),
84 Box::new(PushdownPredicatesRule),
85 Box::new(EliminateDeadCodeRule),
86 Box::new(FoldConstantsRule),
87 ];
88
89 Self {
90 rules,
91 max_iterations: 10,
92 }
93 }
94
95 pub fn add_rule(&mut self, rule: Box<dyn RewriteRule>) {
97 self.rules.push(rule);
98 }
99
100 pub fn rewrite(&self, query: QueryExpr) -> QueryExpr {
102 let mut ctx = RewriteContext::default();
103 self.rewrite_with_context(query, &mut ctx)
104 }
105
106 pub fn rewrite_with_context(
108 &self,
109 mut query: QueryExpr,
110 ctx: &mut RewriteContext,
111 ) -> QueryExpr {
112 for _iteration in 0..self.max_iterations {
114 let original = format!("{:?}", query);
115
116 for rule in &self.rules {
117 if rule.is_applicable(&query) {
118 query = rule.apply(query, ctx);
119 }
120 }
121
122 if format!("{:?}", query) == original {
124 break;
125 }
126 }
127
128 query
129 }
130}
131
132impl Default for QueryRewriter {
133 fn default() -> Self {
134 Self::new()
135 }
136}
137
138struct NormalizeRule;
144
145impl RewriteRule for NormalizeRule {
146 fn name(&self) -> &str {
147 "Normalize"
148 }
149
150 fn apply(&self, query: QueryExpr, ctx: &mut RewriteContext) -> QueryExpr {
151 match query {
152 QueryExpr::Table(mut tq) => {
153 tq.columns.sort_by(|a, b| {
155 let a_name = projection_name(a);
156 let b_name = projection_name(b);
157 a_name.cmp(&b_name)
158 });
159 ctx.stats.expressions_normalized += 1;
160 QueryExpr::Table(tq)
161 }
162 QueryExpr::Graph(gq) => {
163 QueryExpr::Graph(gq)
165 }
166 QueryExpr::Join(jq) => {
167 let left = self.apply(*jq.left, ctx);
169 let right = self.apply(*jq.right, ctx);
170 QueryExpr::Join(JoinQuery {
171 left: Box::new(left),
172 right: Box::new(right),
173 ..jq
174 })
175 }
176 QueryExpr::Path(pq) => QueryExpr::Path(pq),
177 QueryExpr::Vector(vq) => {
178 QueryExpr::Vector(vq)
180 }
181 QueryExpr::Hybrid(mut hq) => {
182 hq.structured = Box::new(self.apply(*hq.structured, ctx));
184 QueryExpr::Hybrid(hq)
185 }
186 other @ (QueryExpr::Insert(_)
188 | QueryExpr::Update(_)
189 | QueryExpr::Delete(_)
190 | QueryExpr::CreateTable(_)
191 | QueryExpr::CreateCollection(_)
192 | QueryExpr::CreateVector(_)
193 | QueryExpr::DropTable(_)
194 | QueryExpr::DropGraph(_)
195 | QueryExpr::DropVector(_)
196 | QueryExpr::DropDocument(_)
197 | QueryExpr::DropKv(_)
198 | QueryExpr::DropCollection(_)
199 | QueryExpr::Truncate(_)
200 | QueryExpr::AlterTable(_)
201 | QueryExpr::GraphCommand(_)
202 | QueryExpr::SearchCommand(_)
203 | QueryExpr::CreateIndex(_)
204 | QueryExpr::DropIndex(_)
205 | QueryExpr::ProbabilisticCommand(_)
206 | QueryExpr::Ask(_)
207 | QueryExpr::SetConfig { .. }
208 | QueryExpr::ShowConfig { .. }
209 | QueryExpr::SetSecret { .. }
210 | QueryExpr::DeleteSecret { .. }
211 | QueryExpr::ShowSecrets { .. }
212 | QueryExpr::SetTenant(_)
213 | QueryExpr::ShowTenant
214 | QueryExpr::CreateTimeSeries(_)
215 | QueryExpr::DropTimeSeries(_)
216 | QueryExpr::CreateQueue(_)
217 | QueryExpr::AlterQueue(_)
218 | QueryExpr::DropQueue(_)
219 | QueryExpr::QueueSelect(_)
220 | QueryExpr::QueueCommand(_)
221 | QueryExpr::KvCommand(_)
222 | QueryExpr::ConfigCommand(_)
223 | QueryExpr::CreateTree(_)
224 | QueryExpr::DropTree(_)
225 | QueryExpr::TreeCommand(_)
226 | QueryExpr::ExplainAlter(_)
227 | QueryExpr::TransactionControl(_)
228 | QueryExpr::MaintenanceCommand(_)
229 | QueryExpr::CreateSchema(_)
230 | QueryExpr::DropSchema(_)
231 | QueryExpr::CreateSequence(_)
232 | QueryExpr::DropSequence(_)
233 | QueryExpr::CopyFrom(_)
234 | QueryExpr::CreateView(_)
235 | QueryExpr::DropView(_)
236 | QueryExpr::RefreshMaterializedView(_)
237 | QueryExpr::CreatePolicy(_)
238 | QueryExpr::DropPolicy(_)
239 | QueryExpr::CreateServer(_)
240 | QueryExpr::DropServer(_)
241 | QueryExpr::CreateForeignTable(_)
242 | QueryExpr::DropForeignTable(_)
243 | QueryExpr::Grant(_)
244 | QueryExpr::Revoke(_)
245 | QueryExpr::AlterUser(_)
246 | QueryExpr::CreateIamPolicy { .. }
247 | QueryExpr::DropIamPolicy { .. }
248 | QueryExpr::AttachPolicy { .. }
249 | QueryExpr::DetachPolicy { .. }
250 | QueryExpr::ShowPolicies { .. }
251 | QueryExpr::ShowEffectivePermissions { .. }
252 | QueryExpr::SimulatePolicy { .. }
253 | QueryExpr::CreateMigration(_)
254 | QueryExpr::ApplyMigration(_)
255 | QueryExpr::RollbackMigration(_)
256 | QueryExpr::ExplainMigration(_)
257 | QueryExpr::EventsBackfill(_)
258 | QueryExpr::EventsBackfillStatus { .. }) => other,
259 }
260 }
261
262 fn is_applicable(&self, _query: &QueryExpr) -> bool {
263 true
264 }
265}
266
267struct SimplifyFiltersRule;
269
270impl RewriteRule for SimplifyFiltersRule {
271 fn name(&self) -> &str {
272 "SimplifyFilters"
273 }
274
275 fn apply(&self, query: QueryExpr, ctx: &mut RewriteContext) -> QueryExpr {
276 match query {
277 QueryExpr::Table(mut tq) => {
278 if let Some(filter) = effective_table_filter(&tq) {
279 tq.filter = Some(simplify_filter(filter, ctx));
280 }
281 QueryExpr::Table(tq)
282 }
283 QueryExpr::Graph(mut gq) => {
284 if let Some(filter) = effective_graph_filter(&gq) {
285 gq.filter = Some(simplify_filter(filter, ctx));
286 }
287 QueryExpr::Graph(gq)
288 }
289 QueryExpr::Join(mut jq) => {
290 let join_filter = effective_join_filter(&jq);
291 let left = self.apply(*jq.left, ctx);
292 let right = self.apply(*jq.right, ctx);
293 if let Some(filter) = join_filter {
294 jq.filter = Some(simplify_filter(filter, ctx));
295 }
296 jq.left = Box::new(left);
297 jq.right = Box::new(right);
298 QueryExpr::Join(jq)
299 }
300 QueryExpr::Path(pq) => QueryExpr::Path(pq),
301 QueryExpr::Vector(vq) => {
302 QueryExpr::Vector(vq)
305 }
306 QueryExpr::Hybrid(mut hq) => {
307 hq.structured = Box::new(self.apply(*hq.structured, ctx));
309 QueryExpr::Hybrid(hq)
310 }
311 other @ (QueryExpr::Insert(_)
313 | QueryExpr::Update(_)
314 | QueryExpr::Delete(_)
315 | QueryExpr::CreateTable(_)
316 | QueryExpr::CreateCollection(_)
317 | QueryExpr::CreateVector(_)
318 | QueryExpr::DropTable(_)
319 | QueryExpr::DropGraph(_)
320 | QueryExpr::DropVector(_)
321 | QueryExpr::DropDocument(_)
322 | QueryExpr::DropKv(_)
323 | QueryExpr::DropCollection(_)
324 | QueryExpr::Truncate(_)
325 | QueryExpr::AlterTable(_)
326 | QueryExpr::GraphCommand(_)
327 | QueryExpr::SearchCommand(_)
328 | QueryExpr::CreateIndex(_)
329 | QueryExpr::DropIndex(_)
330 | QueryExpr::ProbabilisticCommand(_)
331 | QueryExpr::Ask(_)
332 | QueryExpr::SetConfig { .. }
333 | QueryExpr::ShowConfig { .. }
334 | QueryExpr::SetSecret { .. }
335 | QueryExpr::DeleteSecret { .. }
336 | QueryExpr::ShowSecrets { .. }
337 | QueryExpr::SetTenant(_)
338 | QueryExpr::ShowTenant
339 | QueryExpr::CreateTimeSeries(_)
340 | QueryExpr::DropTimeSeries(_)
341 | QueryExpr::CreateQueue(_)
342 | QueryExpr::AlterQueue(_)
343 | QueryExpr::DropQueue(_)
344 | QueryExpr::QueueSelect(_)
345 | QueryExpr::QueueCommand(_)
346 | QueryExpr::KvCommand(_)
347 | QueryExpr::ConfigCommand(_)
348 | QueryExpr::CreateTree(_)
349 | QueryExpr::DropTree(_)
350 | QueryExpr::TreeCommand(_)
351 | QueryExpr::ExplainAlter(_)
352 | QueryExpr::TransactionControl(_)
353 | QueryExpr::MaintenanceCommand(_)
354 | QueryExpr::CreateSchema(_)
355 | QueryExpr::DropSchema(_)
356 | QueryExpr::CreateSequence(_)
357 | QueryExpr::DropSequence(_)
358 | QueryExpr::CopyFrom(_)
359 | QueryExpr::CreateView(_)
360 | QueryExpr::DropView(_)
361 | QueryExpr::RefreshMaterializedView(_)
362 | QueryExpr::CreatePolicy(_)
363 | QueryExpr::DropPolicy(_)
364 | QueryExpr::CreateServer(_)
365 | QueryExpr::DropServer(_)
366 | QueryExpr::CreateForeignTable(_)
367 | QueryExpr::DropForeignTable(_)
368 | QueryExpr::Grant(_)
369 | QueryExpr::Revoke(_)
370 | QueryExpr::AlterUser(_)
371 | QueryExpr::CreateIamPolicy { .. }
372 | QueryExpr::DropIamPolicy { .. }
373 | QueryExpr::AttachPolicy { .. }
374 | QueryExpr::DetachPolicy { .. }
375 | QueryExpr::ShowPolicies { .. }
376 | QueryExpr::ShowEffectivePermissions { .. }
377 | QueryExpr::SimulatePolicy { .. }
378 | QueryExpr::CreateMigration(_)
379 | QueryExpr::ApplyMigration(_)
380 | QueryExpr::RollbackMigration(_)
381 | QueryExpr::ExplainMigration(_)
382 | QueryExpr::EventsBackfill(_)
383 | QueryExpr::EventsBackfillStatus { .. }) => other,
384 }
385 }
386
387 fn is_applicable(&self, query: &QueryExpr) -> bool {
388 match query {
389 QueryExpr::Table(tq) => effective_table_filter(tq).is_some(),
390 QueryExpr::Graph(gq) => effective_graph_filter(gq).is_some(),
391 QueryExpr::Join(_) => true,
392 QueryExpr::Path(_) => false,
393 QueryExpr::Vector(vq) => effective_vector_filter(vq).is_some(),
394 QueryExpr::Hybrid(_) => true, QueryExpr::Insert(_)
397 | QueryExpr::Update(_)
398 | QueryExpr::Delete(_)
399 | QueryExpr::CreateTable(_)
400 | QueryExpr::CreateCollection(_)
401 | QueryExpr::CreateVector(_)
402 | QueryExpr::DropTable(_)
403 | QueryExpr::DropGraph(_)
404 | QueryExpr::DropVector(_)
405 | QueryExpr::DropDocument(_)
406 | QueryExpr::DropKv(_)
407 | QueryExpr::DropCollection(_)
408 | QueryExpr::Truncate(_)
409 | QueryExpr::AlterTable(_)
410 | QueryExpr::GraphCommand(_)
411 | QueryExpr::SearchCommand(_)
412 | QueryExpr::CreateIndex(_)
413 | QueryExpr::DropIndex(_)
414 | QueryExpr::ProbabilisticCommand(_)
415 | QueryExpr::Ask(_)
416 | QueryExpr::SetConfig { .. }
417 | QueryExpr::ShowConfig { .. }
418 | QueryExpr::SetSecret { .. }
419 | QueryExpr::DeleteSecret { .. }
420 | QueryExpr::ShowSecrets { .. }
421 | QueryExpr::SetTenant(_)
422 | QueryExpr::ShowTenant
423 | QueryExpr::CreateTimeSeries(_)
424 | QueryExpr::DropTimeSeries(_)
425 | QueryExpr::CreateQueue(_)
426 | QueryExpr::AlterQueue(_)
427 | QueryExpr::DropQueue(_)
428 | QueryExpr::QueueSelect(_)
429 | QueryExpr::QueueCommand(_)
430 | QueryExpr::KvCommand(_)
431 | QueryExpr::ConfigCommand(_)
432 | QueryExpr::CreateTree(_)
433 | QueryExpr::DropTree(_)
434 | QueryExpr::TreeCommand(_)
435 | QueryExpr::ExplainAlter(_)
436 | QueryExpr::TransactionControl(_)
437 | QueryExpr::MaintenanceCommand(_)
438 | QueryExpr::CreateSchema(_)
439 | QueryExpr::DropSchema(_)
440 | QueryExpr::CreateSequence(_)
441 | QueryExpr::DropSequence(_)
442 | QueryExpr::CopyFrom(_)
443 | QueryExpr::CreateView(_)
444 | QueryExpr::DropView(_)
445 | QueryExpr::RefreshMaterializedView(_)
446 | QueryExpr::CreatePolicy(_)
447 | QueryExpr::DropPolicy(_)
448 | QueryExpr::CreateServer(_)
449 | QueryExpr::DropServer(_)
450 | QueryExpr::CreateForeignTable(_)
451 | QueryExpr::DropForeignTable(_)
452 | QueryExpr::Grant(_)
453 | QueryExpr::Revoke(_)
454 | QueryExpr::AlterUser(_)
455 | QueryExpr::CreateIamPolicy { .. }
456 | QueryExpr::DropIamPolicy { .. }
457 | QueryExpr::AttachPolicy { .. }
458 | QueryExpr::DetachPolicy { .. }
459 | QueryExpr::ShowPolicies { .. }
460 | QueryExpr::ShowEffectivePermissions { .. }
461 | QueryExpr::SimulatePolicy { .. }
462 | QueryExpr::CreateMigration(_)
463 | QueryExpr::ApplyMigration(_)
464 | QueryExpr::RollbackMigration(_)
465 | QueryExpr::ExplainMigration(_)
466 | QueryExpr::EventsBackfill(_)
467 | QueryExpr::EventsBackfillStatus { .. } => false,
468 }
469 }
470}
471
472struct PushdownPredicatesRule;
474
475impl RewriteRule for PushdownPredicatesRule {
476 fn name(&self) -> &str {
477 "PushdownPredicates"
478 }
479
480 fn apply(&self, query: QueryExpr, ctx: &mut RewriteContext) -> QueryExpr {
481 match query {
482 QueryExpr::Join(mut jq) => {
483 jq.left = Box::new(self.apply(*jq.left, ctx));
489 jq.right = Box::new(self.apply(*jq.right, ctx));
490 ctx.stats.predicates_pushed += 1;
491 QueryExpr::Join(jq)
492 }
493 other => other,
494 }
495 }
496
497 fn is_applicable(&self, query: &QueryExpr) -> bool {
498 matches!(query, QueryExpr::Join(_))
499 }
500}
501
502struct EliminateDeadCodeRule;
504
505impl RewriteRule for EliminateDeadCodeRule {
506 fn name(&self) -> &str {
507 "EliminateDeadCode"
508 }
509
510 fn apply(&self, query: QueryExpr, _ctx: &mut RewriteContext) -> QueryExpr {
511 match query {
512 QueryExpr::Table(mut tq) => {
513 if let Some(filter) = effective_table_filter(&tq).as_ref() {
515 if is_always_true(filter) {
516 tq.filter = None;
517 }
518 }
519 QueryExpr::Table(tq)
520 }
521 other => other,
522 }
523 }
524
525 fn is_applicable(&self, query: &QueryExpr) -> bool {
526 matches!(query, QueryExpr::Table(_))
527 }
528}
529
530struct FoldConstantsRule;
532
533impl RewriteRule for FoldConstantsRule {
534 fn name(&self) -> &str {
535 "FoldConstants"
536 }
537
538 fn apply(&self, query: QueryExpr, _ctx: &mut RewriteContext) -> QueryExpr {
539 query
542 }
543
544 fn is_applicable(&self, _query: &QueryExpr) -> bool {
545 true
546 }
547}
548
549fn projection_name(proj: &Projection) -> String {
554 match proj {
555 Projection::All => "*".to_string(),
556 Projection::Column(name) => name.clone(),
557 Projection::Alias(_, alias) => alias.clone(),
558 Projection::Function(name, _) => name
559 .split_once(':')
560 .map(|(_, alias)| alias.to_string())
561 .unwrap_or_else(|| name.clone()),
562 Projection::Expression(expr, alias) => {
563 alias.clone().unwrap_or_else(|| format!("{:?}", expr))
564 }
565 Projection::Field(field, alias) => alias.clone().unwrap_or_else(|| format!("{:?}", field)),
566 Projection::Window { name, alias, .. } => alias.clone().unwrap_or_else(|| name.clone()),
567 }
568}
569
570fn simplify_filter(filter: AstFilter, ctx: &mut RewriteContext) -> AstFilter {
571 match filter {
572 AstFilter::And(left, right) => {
573 let left = simplify_filter(*left, ctx);
574 let right = simplify_filter(*right, ctx);
575
576 if is_always_true(&left) {
578 ctx.stats.filters_simplified += 1;
579 return right;
580 }
581 if is_always_true(&right) {
582 ctx.stats.filters_simplified += 1;
583 return left;
584 }
585
586 if is_always_false(&left) || is_always_false(&right) {
588 ctx.stats.filters_simplified += 1;
589 return AstFilter::Compare {
590 field: FieldRef::TableColumn {
591 table: String::new(),
592 column: "1".to_string(),
593 },
594 op: CompareOp::Eq,
595 value: Value::Integer(0),
596 };
597 }
598
599 AstFilter::And(Box::new(left), Box::new(right))
600 }
601 AstFilter::Or(left, right) => {
602 let left = simplify_filter(*left, ctx);
603 let right = simplify_filter(*right, ctx);
604
605 if is_always_false(&left) {
607 ctx.stats.filters_simplified += 1;
608 return right;
609 }
610 if is_always_false(&right) {
611 ctx.stats.filters_simplified += 1;
612 return left;
613 }
614
615 if is_always_true(&left) || is_always_true(&right) {
617 ctx.stats.filters_simplified += 1;
618 return AstFilter::Compare {
619 field: FieldRef::TableColumn {
620 table: String::new(),
621 column: "1".to_string(),
622 },
623 op: CompareOp::Eq,
624 value: Value::Integer(1),
625 };
626 }
627
628 AstFilter::Or(Box::new(left), Box::new(right))
629 }
630 AstFilter::Not(inner) => {
631 let inner = simplify_filter(*inner, ctx);
632
633 if let AstFilter::Not(double_inner) = inner {
635 ctx.stats.filters_simplified += 1;
636 return *double_inner;
637 }
638
639 AstFilter::Not(Box::new(inner))
640 }
641 other => other,
642 }
643}
644
645fn is_always_true(filter: &AstFilter) -> bool {
646 match filter {
647 AstFilter::Compare { field, op, value } => {
648 matches!(field, FieldRef::TableColumn { column, .. } if column == "1")
650 && matches!(op, CompareOp::Eq)
651 && matches!(value, Value::Integer(1))
652 }
653 _ => false,
654 }
655}
656
657fn is_always_false(filter: &AstFilter) -> bool {
658 match filter {
659 AstFilter::Compare { field, op, value } => {
660 matches!(field, FieldRef::TableColumn { column, .. } if column == "1")
662 && matches!(op, CompareOp::Eq)
663 && matches!(value, Value::Integer(0))
664 }
665 _ => false,
666 }
667}
668
669#[cfg(test)]
670mod tests {
671 use super::*;
672
673 fn make_field(name: &str) -> FieldRef {
674 FieldRef::TableColumn {
675 table: String::new(),
676 column: name.to_string(),
677 }
678 }
679
680 #[test]
681 fn test_simplify_and_with_true() {
682 let mut ctx = RewriteContext::default();
683
684 let filter = AstFilter::And(
685 Box::new(AstFilter::Compare {
686 field: make_field("1"),
687 op: CompareOp::Eq,
688 value: Value::Integer(1),
689 }),
690 Box::new(AstFilter::Compare {
691 field: make_field("x"),
692 op: CompareOp::Eq,
693 value: Value::Integer(5),
694 }),
695 );
696
697 let simplified = simplify_filter(filter, &mut ctx);
698
699 match simplified {
700 AstFilter::Compare { field, .. } => {
701 assert!(matches!(field, FieldRef::TableColumn { column, .. } if column == "x"));
702 }
703 _ => panic!("Expected Compare filter"),
704 }
705 }
706
707 #[test]
708 fn test_simplify_double_not() {
709 let mut ctx = RewriteContext::default();
710
711 let filter = AstFilter::Not(Box::new(AstFilter::Not(Box::new(AstFilter::Compare {
712 field: make_field("x"),
713 op: CompareOp::Eq,
714 value: Value::Integer(5),
715 }))));
716
717 let simplified = simplify_filter(filter, &mut ctx);
718
719 match simplified {
720 AstFilter::Compare { field, .. } => {
721 assert!(matches!(field, FieldRef::TableColumn { column, .. } if column == "x"));
722 }
723 _ => panic!("Expected Compare filter"),
724 }
725 }
726}