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