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::CreateMetric(_)
216 | QueryExpr::AlterMetric(_)
217 | QueryExpr::CreateSlo(_)
218 | QueryExpr::DropTimeSeries(_)
219 | QueryExpr::CreateQueue(_)
220 | QueryExpr::AlterQueue(_)
221 | QueryExpr::DropQueue(_)
222 | QueryExpr::QueueSelect(_)
223 | QueryExpr::QueueCommand(_)
224 | QueryExpr::KvCommand(_)
225 | QueryExpr::ConfigCommand(_)
226 | QueryExpr::CreateTree(_)
227 | QueryExpr::DropTree(_)
228 | QueryExpr::TreeCommand(_)
229 | QueryExpr::ExplainAlter(_)
230 | QueryExpr::TransactionControl(_)
231 | QueryExpr::MaintenanceCommand(_)
232 | QueryExpr::CreateSchema(_)
233 | QueryExpr::DropSchema(_)
234 | QueryExpr::CreateSequence(_)
235 | QueryExpr::DropSequence(_)
236 | QueryExpr::CopyFrom(_)
237 | QueryExpr::CreateView(_)
238 | QueryExpr::DropView(_)
239 | QueryExpr::RefreshMaterializedView(_)
240 | QueryExpr::CreatePolicy(_)
241 | QueryExpr::DropPolicy(_)
242 | QueryExpr::CreateServer(_)
243 | QueryExpr::DropServer(_)
244 | QueryExpr::CreateForeignTable(_)
245 | QueryExpr::DropForeignTable(_)
246 | QueryExpr::Grant(_)
247 | QueryExpr::Revoke(_)
248 | QueryExpr::AlterUser(_)
249 | QueryExpr::CreateIamPolicy { .. }
250 | QueryExpr::DropIamPolicy { .. }
251 | QueryExpr::AttachPolicy { .. }
252 | QueryExpr::DetachPolicy { .. }
253 | QueryExpr::ShowPolicies { .. }
254 | QueryExpr::ShowEffectivePermissions { .. }
255 | QueryExpr::SimulatePolicy { .. }
256 | QueryExpr::LintPolicy { .. }
257 | QueryExpr::MigratePolicyMode { .. }
258 | QueryExpr::CreateMigration(_)
259 | QueryExpr::ApplyMigration(_)
260 | QueryExpr::RollbackMigration(_)
261 | QueryExpr::ExplainMigration(_)
262 | QueryExpr::EventsBackfill(_)
263 | QueryExpr::EventsBackfillStatus { .. }) => other,
264 }
265 }
266
267 fn is_applicable(&self, _query: &QueryExpr) -> bool {
268 true
269 }
270}
271
272struct SimplifyFiltersRule;
274
275impl RewriteRule for SimplifyFiltersRule {
276 fn name(&self) -> &str {
277 "SimplifyFilters"
278 }
279
280 fn apply(&self, query: QueryExpr, ctx: &mut RewriteContext) -> QueryExpr {
281 match query {
282 QueryExpr::Table(mut tq) => {
283 if let Some(filter) = effective_table_filter(&tq) {
284 tq.filter = Some(simplify_filter(filter, ctx));
285 }
286 QueryExpr::Table(tq)
287 }
288 QueryExpr::Graph(mut gq) => {
289 if let Some(filter) = effective_graph_filter(&gq) {
290 gq.filter = Some(simplify_filter(filter, ctx));
291 }
292 QueryExpr::Graph(gq)
293 }
294 QueryExpr::Join(mut jq) => {
295 let join_filter = effective_join_filter(&jq);
296 let left = self.apply(*jq.left, ctx);
297 let right = self.apply(*jq.right, ctx);
298 if let Some(filter) = join_filter {
299 jq.filter = Some(simplify_filter(filter, ctx));
300 }
301 jq.left = Box::new(left);
302 jq.right = Box::new(right);
303 QueryExpr::Join(jq)
304 }
305 QueryExpr::Path(pq) => QueryExpr::Path(pq),
306 QueryExpr::Vector(vq) => {
307 QueryExpr::Vector(vq)
310 }
311 QueryExpr::Hybrid(mut hq) => {
312 hq.structured = Box::new(self.apply(*hq.structured, ctx));
314 QueryExpr::Hybrid(hq)
315 }
316 other @ (QueryExpr::Insert(_)
318 | QueryExpr::Update(_)
319 | QueryExpr::Delete(_)
320 | QueryExpr::CreateTable(_)
321 | QueryExpr::CreateCollection(_)
322 | QueryExpr::CreateVector(_)
323 | QueryExpr::DropTable(_)
324 | QueryExpr::DropGraph(_)
325 | QueryExpr::DropVector(_)
326 | QueryExpr::DropDocument(_)
327 | QueryExpr::DropKv(_)
328 | QueryExpr::DropCollection(_)
329 | QueryExpr::Truncate(_)
330 | QueryExpr::AlterTable(_)
331 | QueryExpr::GraphCommand(_)
332 | QueryExpr::SearchCommand(_)
333 | QueryExpr::CreateIndex(_)
334 | QueryExpr::DropIndex(_)
335 | QueryExpr::ProbabilisticCommand(_)
336 | QueryExpr::Ask(_)
337 | QueryExpr::SetConfig { .. }
338 | QueryExpr::ShowConfig { .. }
339 | QueryExpr::SetSecret { .. }
340 | QueryExpr::DeleteSecret { .. }
341 | QueryExpr::ShowSecrets { .. }
342 | QueryExpr::SetTenant(_)
343 | QueryExpr::ShowTenant
344 | QueryExpr::CreateTimeSeries(_)
345 | QueryExpr::CreateMetric(_)
346 | QueryExpr::AlterMetric(_)
347 | QueryExpr::CreateSlo(_)
348 | QueryExpr::DropTimeSeries(_)
349 | QueryExpr::CreateQueue(_)
350 | QueryExpr::AlterQueue(_)
351 | QueryExpr::DropQueue(_)
352 | QueryExpr::QueueSelect(_)
353 | QueryExpr::QueueCommand(_)
354 | QueryExpr::KvCommand(_)
355 | QueryExpr::ConfigCommand(_)
356 | QueryExpr::CreateTree(_)
357 | QueryExpr::DropTree(_)
358 | QueryExpr::TreeCommand(_)
359 | QueryExpr::ExplainAlter(_)
360 | QueryExpr::TransactionControl(_)
361 | QueryExpr::MaintenanceCommand(_)
362 | QueryExpr::CreateSchema(_)
363 | QueryExpr::DropSchema(_)
364 | QueryExpr::CreateSequence(_)
365 | QueryExpr::DropSequence(_)
366 | QueryExpr::CopyFrom(_)
367 | QueryExpr::CreateView(_)
368 | QueryExpr::DropView(_)
369 | QueryExpr::RefreshMaterializedView(_)
370 | QueryExpr::CreatePolicy(_)
371 | QueryExpr::DropPolicy(_)
372 | QueryExpr::CreateServer(_)
373 | QueryExpr::DropServer(_)
374 | QueryExpr::CreateForeignTable(_)
375 | QueryExpr::DropForeignTable(_)
376 | QueryExpr::Grant(_)
377 | QueryExpr::Revoke(_)
378 | QueryExpr::AlterUser(_)
379 | QueryExpr::CreateIamPolicy { .. }
380 | QueryExpr::DropIamPolicy { .. }
381 | QueryExpr::AttachPolicy { .. }
382 | QueryExpr::DetachPolicy { .. }
383 | QueryExpr::ShowPolicies { .. }
384 | QueryExpr::ShowEffectivePermissions { .. }
385 | QueryExpr::SimulatePolicy { .. }
386 | QueryExpr::LintPolicy { .. }
387 | QueryExpr::MigratePolicyMode { .. }
388 | QueryExpr::CreateMigration(_)
389 | QueryExpr::ApplyMigration(_)
390 | QueryExpr::RollbackMigration(_)
391 | QueryExpr::ExplainMigration(_)
392 | QueryExpr::EventsBackfill(_)
393 | QueryExpr::EventsBackfillStatus { .. }) => other,
394 }
395 }
396
397 fn is_applicable(&self, query: &QueryExpr) -> bool {
398 match query {
399 QueryExpr::Table(tq) => effective_table_filter(tq).is_some(),
400 QueryExpr::Graph(gq) => effective_graph_filter(gq).is_some(),
401 QueryExpr::Join(_) => true,
402 QueryExpr::Path(_) => false,
403 QueryExpr::Vector(vq) => effective_vector_filter(vq).is_some(),
404 QueryExpr::Hybrid(_) => true, QueryExpr::Insert(_)
407 | QueryExpr::Update(_)
408 | QueryExpr::Delete(_)
409 | QueryExpr::CreateTable(_)
410 | QueryExpr::CreateCollection(_)
411 | QueryExpr::CreateVector(_)
412 | QueryExpr::DropTable(_)
413 | QueryExpr::DropGraph(_)
414 | QueryExpr::DropVector(_)
415 | QueryExpr::DropDocument(_)
416 | QueryExpr::DropKv(_)
417 | QueryExpr::DropCollection(_)
418 | QueryExpr::Truncate(_)
419 | QueryExpr::AlterTable(_)
420 | QueryExpr::GraphCommand(_)
421 | QueryExpr::SearchCommand(_)
422 | QueryExpr::CreateIndex(_)
423 | QueryExpr::DropIndex(_)
424 | QueryExpr::ProbabilisticCommand(_)
425 | QueryExpr::Ask(_)
426 | QueryExpr::SetConfig { .. }
427 | QueryExpr::ShowConfig { .. }
428 | QueryExpr::SetSecret { .. }
429 | QueryExpr::DeleteSecret { .. }
430 | QueryExpr::ShowSecrets { .. }
431 | QueryExpr::SetTenant(_)
432 | QueryExpr::ShowTenant
433 | QueryExpr::CreateTimeSeries(_)
434 | QueryExpr::CreateMetric(_)
435 | QueryExpr::AlterMetric(_)
436 | QueryExpr::CreateSlo(_)
437 | QueryExpr::DropTimeSeries(_)
438 | QueryExpr::CreateQueue(_)
439 | QueryExpr::AlterQueue(_)
440 | QueryExpr::DropQueue(_)
441 | QueryExpr::QueueSelect(_)
442 | QueryExpr::QueueCommand(_)
443 | QueryExpr::KvCommand(_)
444 | QueryExpr::ConfigCommand(_)
445 | QueryExpr::CreateTree(_)
446 | QueryExpr::DropTree(_)
447 | QueryExpr::TreeCommand(_)
448 | QueryExpr::ExplainAlter(_)
449 | QueryExpr::TransactionControl(_)
450 | QueryExpr::MaintenanceCommand(_)
451 | QueryExpr::CreateSchema(_)
452 | QueryExpr::DropSchema(_)
453 | QueryExpr::CreateSequence(_)
454 | QueryExpr::DropSequence(_)
455 | QueryExpr::CopyFrom(_)
456 | QueryExpr::CreateView(_)
457 | QueryExpr::DropView(_)
458 | QueryExpr::RefreshMaterializedView(_)
459 | QueryExpr::CreatePolicy(_)
460 | QueryExpr::DropPolicy(_)
461 | QueryExpr::CreateServer(_)
462 | QueryExpr::DropServer(_)
463 | QueryExpr::CreateForeignTable(_)
464 | QueryExpr::DropForeignTable(_)
465 | QueryExpr::Grant(_)
466 | QueryExpr::Revoke(_)
467 | QueryExpr::AlterUser(_)
468 | QueryExpr::CreateIamPolicy { .. }
469 | QueryExpr::DropIamPolicy { .. }
470 | QueryExpr::AttachPolicy { .. }
471 | QueryExpr::DetachPolicy { .. }
472 | QueryExpr::ShowPolicies { .. }
473 | QueryExpr::ShowEffectivePermissions { .. }
474 | QueryExpr::SimulatePolicy { .. }
475 | QueryExpr::LintPolicy { .. }
476 | QueryExpr::MigratePolicyMode { .. }
477 | QueryExpr::CreateMigration(_)
478 | QueryExpr::ApplyMigration(_)
479 | QueryExpr::RollbackMigration(_)
480 | QueryExpr::ExplainMigration(_)
481 | QueryExpr::EventsBackfill(_)
482 | QueryExpr::EventsBackfillStatus { .. } => false,
483 }
484 }
485}
486
487struct PushdownPredicatesRule;
489
490impl RewriteRule for PushdownPredicatesRule {
491 fn name(&self) -> &str {
492 "PushdownPredicates"
493 }
494
495 fn apply(&self, query: QueryExpr, ctx: &mut RewriteContext) -> QueryExpr {
496 match query {
497 QueryExpr::Join(mut jq) => {
498 jq.left = Box::new(self.apply(*jq.left, ctx));
504 jq.right = Box::new(self.apply(*jq.right, ctx));
505 ctx.stats.predicates_pushed += 1;
506 QueryExpr::Join(jq)
507 }
508 other => other,
509 }
510 }
511
512 fn is_applicable(&self, query: &QueryExpr) -> bool {
513 matches!(query, QueryExpr::Join(_))
514 }
515}
516
517struct EliminateDeadCodeRule;
519
520impl RewriteRule for EliminateDeadCodeRule {
521 fn name(&self) -> &str {
522 "EliminateDeadCode"
523 }
524
525 fn apply(&self, query: QueryExpr, _ctx: &mut RewriteContext) -> QueryExpr {
526 match query {
527 QueryExpr::Table(mut tq) => {
528 if let Some(filter) = effective_table_filter(&tq).as_ref() {
530 if is_always_true(filter) {
531 tq.filter = None;
532 }
533 }
534 QueryExpr::Table(tq)
535 }
536 other => other,
537 }
538 }
539
540 fn is_applicable(&self, query: &QueryExpr) -> bool {
541 matches!(query, QueryExpr::Table(_))
542 }
543}
544
545struct FoldConstantsRule;
547
548impl RewriteRule for FoldConstantsRule {
549 fn name(&self) -> &str {
550 "FoldConstants"
551 }
552
553 fn apply(&self, query: QueryExpr, _ctx: &mut RewriteContext) -> QueryExpr {
554 query
557 }
558
559 fn is_applicable(&self, _query: &QueryExpr) -> bool {
560 true
561 }
562}
563
564fn projection_name(proj: &Projection) -> String {
569 match proj {
570 Projection::All => "*".to_string(),
571 Projection::Column(name) => name.clone(),
572 Projection::Alias(_, alias) => alias.clone(),
573 Projection::Function(name, _) => name
574 .split_once(':')
575 .map(|(_, alias)| alias.to_string())
576 .unwrap_or_else(|| name.clone()),
577 Projection::Expression(expr, alias) => {
578 alias.clone().unwrap_or_else(|| format!("{:?}", expr))
579 }
580 Projection::Field(field, alias) => alias.clone().unwrap_or_else(|| format!("{:?}", field)),
581 Projection::Window { name, alias, .. } => alias.clone().unwrap_or_else(|| name.clone()),
582 }
583}
584
585fn simplify_filter(filter: AstFilter, ctx: &mut RewriteContext) -> AstFilter {
586 match filter {
587 AstFilter::And(left, right) => {
588 let left = simplify_filter(*left, ctx);
589 let right = simplify_filter(*right, ctx);
590
591 if is_always_true(&left) {
593 ctx.stats.filters_simplified += 1;
594 return right;
595 }
596 if is_always_true(&right) {
597 ctx.stats.filters_simplified += 1;
598 return left;
599 }
600
601 if is_always_false(&left) || is_always_false(&right) {
603 ctx.stats.filters_simplified += 1;
604 return AstFilter::Compare {
605 field: FieldRef::TableColumn {
606 table: String::new(),
607 column: "1".to_string(),
608 },
609 op: CompareOp::Eq,
610 value: Value::Integer(0),
611 };
612 }
613
614 AstFilter::And(Box::new(left), Box::new(right))
615 }
616 AstFilter::Or(left, right) => {
617 let left = simplify_filter(*left, ctx);
618 let right = simplify_filter(*right, ctx);
619
620 if is_always_false(&left) {
622 ctx.stats.filters_simplified += 1;
623 return right;
624 }
625 if is_always_false(&right) {
626 ctx.stats.filters_simplified += 1;
627 return left;
628 }
629
630 if is_always_true(&left) || is_always_true(&right) {
632 ctx.stats.filters_simplified += 1;
633 return AstFilter::Compare {
634 field: FieldRef::TableColumn {
635 table: String::new(),
636 column: "1".to_string(),
637 },
638 op: CompareOp::Eq,
639 value: Value::Integer(1),
640 };
641 }
642
643 AstFilter::Or(Box::new(left), Box::new(right))
644 }
645 AstFilter::Not(inner) => {
646 let inner = simplify_filter(*inner, ctx);
647
648 if let AstFilter::Not(double_inner) = inner {
650 ctx.stats.filters_simplified += 1;
651 return *double_inner;
652 }
653
654 AstFilter::Not(Box::new(inner))
655 }
656 other => other,
657 }
658}
659
660fn is_always_true(filter: &AstFilter) -> bool {
661 match filter {
662 AstFilter::Compare { field, op, value } => {
663 matches!(field, FieldRef::TableColumn { column, .. } if column == "1")
665 && matches!(op, CompareOp::Eq)
666 && matches!(value, Value::Integer(1))
667 }
668 _ => false,
669 }
670}
671
672fn is_always_false(filter: &AstFilter) -> bool {
673 match filter {
674 AstFilter::Compare { field, op, value } => {
675 matches!(field, FieldRef::TableColumn { column, .. } if column == "1")
677 && matches!(op, CompareOp::Eq)
678 && matches!(value, Value::Integer(0))
679 }
680 _ => false,
681 }
682}
683
684#[cfg(test)]
685mod tests {
686 use super::*;
687
688 fn make_field(name: &str) -> FieldRef {
689 FieldRef::TableColumn {
690 table: String::new(),
691 column: name.to_string(),
692 }
693 }
694
695 #[test]
696 fn test_simplify_and_with_true() {
697 let mut ctx = RewriteContext::default();
698
699 let filter = AstFilter::And(
700 Box::new(AstFilter::Compare {
701 field: make_field("1"),
702 op: CompareOp::Eq,
703 value: Value::Integer(1),
704 }),
705 Box::new(AstFilter::Compare {
706 field: make_field("x"),
707 op: CompareOp::Eq,
708 value: Value::Integer(5),
709 }),
710 );
711
712 let simplified = simplify_filter(filter, &mut ctx);
713
714 match simplified {
715 AstFilter::Compare { field, .. } => {
716 assert!(matches!(field, FieldRef::TableColumn { column, .. } if column == "x"));
717 }
718 _ => panic!("Expected Compare filter"),
719 }
720 }
721
722 #[test]
723 fn test_simplify_double_not() {
724 let mut ctx = RewriteContext::default();
725
726 let filter = AstFilter::Not(Box::new(AstFilter::Not(Box::new(AstFilter::Compare {
727 field: make_field("x"),
728 op: CompareOp::Eq,
729 value: Value::Integer(5),
730 }))));
731
732 let simplified = simplify_filter(filter, &mut ctx);
733
734 match simplified {
735 AstFilter::Compare { field, .. } => {
736 assert!(matches!(field, FieldRef::TableColumn { column, .. } if column == "x"));
737 }
738 _ => panic!("Expected Compare filter"),
739 }
740 }
741}