reddb_server/storage/query/planner/
rewriter.rs1use 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.clone(),
559 Projection::Expression(expr, _) => format!("{:?}", expr),
560 Projection::Field(field, alias) => alias.clone().unwrap_or_else(|| format!("{:?}", field)),
561 }
562}
563
564fn simplify_filter(filter: AstFilter, ctx: &mut RewriteContext) -> AstFilter {
565 match filter {
566 AstFilter::And(left, right) => {
567 let left = simplify_filter(*left, ctx);
568 let right = simplify_filter(*right, ctx);
569
570 if is_always_true(&left) {
572 ctx.stats.filters_simplified += 1;
573 return right;
574 }
575 if is_always_true(&right) {
576 ctx.stats.filters_simplified += 1;
577 return left;
578 }
579
580 if is_always_false(&left) || is_always_false(&right) {
582 ctx.stats.filters_simplified += 1;
583 return AstFilter::Compare {
584 field: FieldRef::TableColumn {
585 table: String::new(),
586 column: "1".to_string(),
587 },
588 op: CompareOp::Eq,
589 value: Value::Integer(0),
590 };
591 }
592
593 AstFilter::And(Box::new(left), Box::new(right))
594 }
595 AstFilter::Or(left, right) => {
596 let left = simplify_filter(*left, ctx);
597 let right = simplify_filter(*right, ctx);
598
599 if is_always_false(&left) {
601 ctx.stats.filters_simplified += 1;
602 return right;
603 }
604 if is_always_false(&right) {
605 ctx.stats.filters_simplified += 1;
606 return left;
607 }
608
609 if is_always_true(&left) || is_always_true(&right) {
611 ctx.stats.filters_simplified += 1;
612 return AstFilter::Compare {
613 field: FieldRef::TableColumn {
614 table: String::new(),
615 column: "1".to_string(),
616 },
617 op: CompareOp::Eq,
618 value: Value::Integer(1),
619 };
620 }
621
622 AstFilter::Or(Box::new(left), Box::new(right))
623 }
624 AstFilter::Not(inner) => {
625 let inner = simplify_filter(*inner, ctx);
626
627 if let AstFilter::Not(double_inner) = inner {
629 ctx.stats.filters_simplified += 1;
630 return *double_inner;
631 }
632
633 AstFilter::Not(Box::new(inner))
634 }
635 other => other,
636 }
637}
638
639fn is_always_true(filter: &AstFilter) -> bool {
640 match filter {
641 AstFilter::Compare { field, op, value } => {
642 matches!(field, FieldRef::TableColumn { column, .. } if column == "1")
644 && matches!(op, CompareOp::Eq)
645 && matches!(value, Value::Integer(1))
646 }
647 _ => false,
648 }
649}
650
651fn is_always_false(filter: &AstFilter) -> bool {
652 match filter {
653 AstFilter::Compare { field, op, value } => {
654 matches!(field, FieldRef::TableColumn { column, .. } if column == "1")
656 && matches!(op, CompareOp::Eq)
657 && matches!(value, Value::Integer(0))
658 }
659 _ => false,
660 }
661}
662
663#[cfg(test)]
664mod tests {
665 use super::*;
666
667 fn make_field(name: &str) -> FieldRef {
668 FieldRef::TableColumn {
669 table: String::new(),
670 column: name.to_string(),
671 }
672 }
673
674 #[test]
675 fn test_simplify_and_with_true() {
676 let mut ctx = RewriteContext::default();
677
678 let filter = AstFilter::And(
679 Box::new(AstFilter::Compare {
680 field: make_field("1"),
681 op: CompareOp::Eq,
682 value: Value::Integer(1),
683 }),
684 Box::new(AstFilter::Compare {
685 field: make_field("x"),
686 op: CompareOp::Eq,
687 value: Value::Integer(5),
688 }),
689 );
690
691 let simplified = simplify_filter(filter, &mut ctx);
692
693 match simplified {
694 AstFilter::Compare { field, .. } => {
695 assert!(matches!(field, FieldRef::TableColumn { column, .. } if column == "x"));
696 }
697 _ => panic!("Expected Compare filter"),
698 }
699 }
700
701 #[test]
702 fn test_simplify_double_not() {
703 let mut ctx = RewriteContext::default();
704
705 let filter = AstFilter::Not(Box::new(AstFilter::Not(Box::new(AstFilter::Compare {
706 field: make_field("x"),
707 op: CompareOp::Eq,
708 value: Value::Integer(5),
709 }))));
710
711 let simplified = simplify_filter(filter, &mut ctx);
712
713 match simplified {
714 AstFilter::Compare { field, .. } => {
715 assert!(matches!(field, FieldRef::TableColumn { column, .. } if column == "x"));
716 }
717 _ => panic!("Expected Compare filter"),
718 }
719 }
720}