1use crate::ast::{CompareOp, FieldRef, Filter as AstFilter, JoinQuery, Projection, QueryExpr};
14use crate::sql_lowering::{
15 effective_graph_filter, effective_join_filter, effective_table_filter, effective_vector_filter,
16};
17use reddb_types::Value;
18
19#[derive(Debug, Clone, Default)]
21pub struct RewriteContext {
22 pub property_cache: Vec<CachedProperty>,
24 pub errors: Vec<String>,
26 pub warnings: Vec<String>,
28 pub stats: RewriteStats,
30}
31
32#[derive(Debug, Clone)]
34pub struct CachedProperty {
35 pub source: String,
37 pub property: String,
39 pub cached_value: Option<String>,
41}
42
43#[derive(Debug, Clone, Default)]
45pub struct RewriteStats {
46 pub filters_simplified: u32,
48 pub predicates_pushed: u32,
50 pub properties_cached: u32,
52 pub expressions_normalized: u32,
54}
55
56pub trait RewriteRule: Send + Sync {
58 fn name(&self) -> &str;
60
61 fn apply(&self, query: QueryExpr, ctx: &mut RewriteContext) -> QueryExpr;
63
64 fn is_applicable(&self, query: &QueryExpr) -> bool;
66}
67
68pub struct QueryRewriter {
70 rules: Vec<Box<dyn RewriteRule>>,
72 max_iterations: usize,
74}
75
76impl QueryRewriter {
77 pub fn new() -> Self {
79 let rules: Vec<Box<dyn RewriteRule>> = vec![
80 Box::new(NormalizeRule),
81 Box::new(SimplifyFiltersRule),
82 Box::new(PushdownPredicatesRule),
83 Box::new(EliminateDeadCodeRule),
84 Box::new(FoldConstantsRule),
85 ];
86
87 Self {
88 rules,
89 max_iterations: 10,
90 }
91 }
92
93 pub fn add_rule(&mut self, rule: Box<dyn RewriteRule>) {
95 self.rules.push(rule);
96 }
97
98 pub fn rewrite(&self, query: QueryExpr) -> QueryExpr {
100 let mut ctx = RewriteContext::default();
101 self.rewrite_with_context(query, &mut ctx)
102 }
103
104 pub fn rewrite_with_context(
106 &self,
107 mut query: QueryExpr,
108 ctx: &mut RewriteContext,
109 ) -> QueryExpr {
110 for _iteration in 0..self.max_iterations {
112 let original = format!("{:?}", query);
113
114 for rule in &self.rules {
115 if rule.is_applicable(&query) {
116 query = rule.apply(query, ctx);
117 }
118 }
119
120 if format!("{:?}", query) == original {
122 break;
123 }
124 }
125
126 query
127 }
128}
129
130impl Default for QueryRewriter {
131 fn default() -> Self {
132 Self::new()
133 }
134}
135
136struct NormalizeRule;
142
143impl RewriteRule for NormalizeRule {
144 fn name(&self) -> &str {
145 "Normalize"
146 }
147
148 fn apply(&self, query: QueryExpr, ctx: &mut RewriteContext) -> QueryExpr {
149 match query {
150 QueryExpr::Table(mut tq) => {
151 tq.columns.sort_by(|a, b| {
153 let a_name = projection_name(a);
154 let b_name = projection_name(b);
155 a_name.cmp(&b_name)
156 });
157 ctx.stats.expressions_normalized += 1;
158 QueryExpr::Table(tq)
159 }
160 QueryExpr::Graph(gq) => {
161 QueryExpr::Graph(gq)
163 }
164 QueryExpr::Join(jq) => {
165 let left = self.apply(*jq.left, ctx);
167 let right = self.apply(*jq.right, ctx);
168 QueryExpr::Join(JoinQuery {
169 left: Box::new(left),
170 right: Box::new(right),
171 ..jq
172 })
173 }
174 QueryExpr::Path(pq) => QueryExpr::Path(pq),
175 QueryExpr::Vector(vq) => {
176 QueryExpr::Vector(vq)
178 }
179 QueryExpr::Hybrid(mut hq) => {
180 hq.structured = Box::new(self.apply(*hq.structured, ctx));
182 QueryExpr::Hybrid(hq)
183 }
184 other @ (QueryExpr::Insert(_)
186 | QueryExpr::Update(_)
187 | QueryExpr::Delete(_)
188 | QueryExpr::CreateTable(_)
189 | QueryExpr::CreateCollection(_)
190 | QueryExpr::CreateVector(_)
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::CreateVcsRef(_)
200 | QueryExpr::DropVcsRef(_)
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::VcsCommand(_)
233 | QueryExpr::CreateSchema(_)
234 | QueryExpr::DropSchema(_)
235 | QueryExpr::CreateSequence(_)
236 | QueryExpr::DropSequence(_)
237 | QueryExpr::CopyFrom(_)
238 | QueryExpr::CreateView(_)
239 | QueryExpr::DropView(_)
240 | QueryExpr::RefreshMaterializedView(_)
241 | QueryExpr::CreatePolicy(_)
242 | QueryExpr::DropPolicy(_)
243 | QueryExpr::CreateServer(_)
244 | QueryExpr::DropServer(_)
245 | QueryExpr::CreateForeignTable(_)
246 | QueryExpr::DropForeignTable(_)
247 | QueryExpr::Grant(_)
248 | QueryExpr::Revoke(_)
249 | QueryExpr::AlterUser(_)
250 | QueryExpr::CreateUser(_)
251 | QueryExpr::CreateIamPolicy { .. }
252 | QueryExpr::DropIamPolicy { .. }
253 | QueryExpr::AttachPolicy { .. }
254 | QueryExpr::DetachPolicy { .. }
255 | QueryExpr::ShowPolicies { .. }
256 | QueryExpr::ShowEffectivePermissions { .. }
257 | QueryExpr::RankOf(_)
258 | QueryExpr::ApproxRankOf(_)
259 | QueryExpr::RankRange(_)
260 | QueryExpr::SimulatePolicy { .. }
261 | QueryExpr::LintPolicy { .. }
262 | QueryExpr::MigratePolicyMode { .. }
263 | QueryExpr::CreateMigration(_)
264 | QueryExpr::ApplyMigration(_)
265 | QueryExpr::RollbackMigration(_)
266 | QueryExpr::ExplainMigration(_)
267 | QueryExpr::EventsBackfill(_)
268 | QueryExpr::EventsBackfillStatus { .. }) => other,
269 }
270 }
271
272 fn is_applicable(&self, _query: &QueryExpr) -> bool {
273 true
274 }
275}
276
277struct SimplifyFiltersRule;
279
280impl RewriteRule for SimplifyFiltersRule {
281 fn name(&self) -> &str {
282 "SimplifyFilters"
283 }
284
285 fn apply(&self, query: QueryExpr, ctx: &mut RewriteContext) -> QueryExpr {
286 match query {
287 QueryExpr::Table(mut tq) => {
288 if let Some(filter) = effective_table_filter(&tq) {
289 tq.filter = Some(simplify_filter(filter, ctx));
290 }
291 QueryExpr::Table(tq)
292 }
293 QueryExpr::Graph(mut gq) => {
294 if let Some(filter) = effective_graph_filter(&gq) {
295 gq.filter = Some(simplify_filter(filter, ctx));
296 }
297 QueryExpr::Graph(gq)
298 }
299 QueryExpr::Join(mut jq) => {
300 let join_filter = effective_join_filter(&jq);
301 let left = self.apply(*jq.left, ctx);
302 let right = self.apply(*jq.right, ctx);
303 if let Some(filter) = join_filter {
304 jq.filter = Some(simplify_filter(filter, ctx));
305 }
306 jq.left = Box::new(left);
307 jq.right = Box::new(right);
308 QueryExpr::Join(jq)
309 }
310 QueryExpr::Path(pq) => QueryExpr::Path(pq),
311 QueryExpr::Vector(vq) => {
312 QueryExpr::Vector(vq)
315 }
316 QueryExpr::Hybrid(mut hq) => {
317 hq.structured = Box::new(self.apply(*hq.structured, ctx));
319 QueryExpr::Hybrid(hq)
320 }
321 other @ (QueryExpr::Insert(_)
323 | QueryExpr::Update(_)
324 | QueryExpr::Delete(_)
325 | QueryExpr::CreateTable(_)
326 | QueryExpr::CreateCollection(_)
327 | QueryExpr::CreateVector(_)
328 | QueryExpr::DropTable(_)
329 | QueryExpr::DropGraph(_)
330 | QueryExpr::DropVector(_)
331 | QueryExpr::DropDocument(_)
332 | QueryExpr::DropKv(_)
333 | QueryExpr::DropCollection(_)
334 | QueryExpr::Truncate(_)
335 | QueryExpr::AlterTable(_)
336 | QueryExpr::CreateVcsRef(_)
337 | QueryExpr::DropVcsRef(_)
338 | QueryExpr::GraphCommand(_)
339 | QueryExpr::SearchCommand(_)
340 | QueryExpr::CreateIndex(_)
341 | QueryExpr::DropIndex(_)
342 | QueryExpr::ProbabilisticCommand(_)
343 | QueryExpr::Ask(_)
344 | QueryExpr::SetConfig { .. }
345 | QueryExpr::ShowConfig { .. }
346 | QueryExpr::SetSecret { .. }
347 | QueryExpr::DeleteSecret { .. }
348 | QueryExpr::ShowSecrets { .. }
349 | QueryExpr::SetTenant(_)
350 | QueryExpr::ShowTenant
351 | QueryExpr::CreateTimeSeries(_)
352 | QueryExpr::CreateMetric(_)
353 | QueryExpr::AlterMetric(_)
354 | QueryExpr::CreateSlo(_)
355 | QueryExpr::DropTimeSeries(_)
356 | QueryExpr::CreateQueue(_)
357 | QueryExpr::AlterQueue(_)
358 | QueryExpr::DropQueue(_)
359 | QueryExpr::QueueSelect(_)
360 | QueryExpr::QueueCommand(_)
361 | QueryExpr::KvCommand(_)
362 | QueryExpr::ConfigCommand(_)
363 | QueryExpr::CreateTree(_)
364 | QueryExpr::DropTree(_)
365 | QueryExpr::TreeCommand(_)
366 | QueryExpr::ExplainAlter(_)
367 | QueryExpr::TransactionControl(_)
368 | QueryExpr::MaintenanceCommand(_)
369 | QueryExpr::VcsCommand(_)
370 | QueryExpr::CreateSchema(_)
371 | QueryExpr::DropSchema(_)
372 | QueryExpr::CreateSequence(_)
373 | QueryExpr::DropSequence(_)
374 | QueryExpr::CopyFrom(_)
375 | QueryExpr::CreateView(_)
376 | QueryExpr::DropView(_)
377 | QueryExpr::RefreshMaterializedView(_)
378 | QueryExpr::CreatePolicy(_)
379 | QueryExpr::DropPolicy(_)
380 | QueryExpr::CreateServer(_)
381 | QueryExpr::DropServer(_)
382 | QueryExpr::CreateForeignTable(_)
383 | QueryExpr::DropForeignTable(_)
384 | QueryExpr::Grant(_)
385 | QueryExpr::Revoke(_)
386 | QueryExpr::AlterUser(_)
387 | QueryExpr::CreateUser(_)
388 | QueryExpr::CreateIamPolicy { .. }
389 | QueryExpr::DropIamPolicy { .. }
390 | QueryExpr::AttachPolicy { .. }
391 | QueryExpr::DetachPolicy { .. }
392 | QueryExpr::ShowPolicies { .. }
393 | QueryExpr::ShowEffectivePermissions { .. }
394 | QueryExpr::RankOf(_)
395 | QueryExpr::ApproxRankOf(_)
396 | QueryExpr::RankRange(_)
397 | QueryExpr::SimulatePolicy { .. }
398 | QueryExpr::LintPolicy { .. }
399 | QueryExpr::MigratePolicyMode { .. }
400 | QueryExpr::CreateMigration(_)
401 | QueryExpr::ApplyMigration(_)
402 | QueryExpr::RollbackMigration(_)
403 | QueryExpr::ExplainMigration(_)
404 | QueryExpr::EventsBackfill(_)
405 | QueryExpr::EventsBackfillStatus { .. }) => other,
406 }
407 }
408
409 fn is_applicable(&self, query: &QueryExpr) -> bool {
410 match query {
411 QueryExpr::Table(tq) => effective_table_filter(tq).is_some(),
412 QueryExpr::Graph(gq) => effective_graph_filter(gq).is_some(),
413 QueryExpr::Join(_) => true,
414 QueryExpr::Path(_) => false,
415 QueryExpr::Vector(vq) => effective_vector_filter(vq).is_some(),
416 QueryExpr::Hybrid(_) => true, QueryExpr::Insert(_)
419 | QueryExpr::Update(_)
420 | QueryExpr::Delete(_)
421 | QueryExpr::CreateTable(_)
422 | QueryExpr::CreateCollection(_)
423 | QueryExpr::CreateVector(_)
424 | QueryExpr::DropTable(_)
425 | QueryExpr::DropGraph(_)
426 | QueryExpr::DropVector(_)
427 | QueryExpr::DropDocument(_)
428 | QueryExpr::DropKv(_)
429 | QueryExpr::DropCollection(_)
430 | QueryExpr::Truncate(_)
431 | QueryExpr::AlterTable(_)
432 | QueryExpr::CreateVcsRef(_)
433 | QueryExpr::DropVcsRef(_)
434 | QueryExpr::GraphCommand(_)
435 | QueryExpr::SearchCommand(_)
436 | QueryExpr::CreateIndex(_)
437 | QueryExpr::DropIndex(_)
438 | QueryExpr::ProbabilisticCommand(_)
439 | QueryExpr::Ask(_)
440 | QueryExpr::SetConfig { .. }
441 | QueryExpr::ShowConfig { .. }
442 | QueryExpr::SetSecret { .. }
443 | QueryExpr::DeleteSecret { .. }
444 | QueryExpr::ShowSecrets { .. }
445 | QueryExpr::SetTenant(_)
446 | QueryExpr::ShowTenant
447 | QueryExpr::CreateTimeSeries(_)
448 | QueryExpr::CreateMetric(_)
449 | QueryExpr::AlterMetric(_)
450 | QueryExpr::CreateSlo(_)
451 | QueryExpr::DropTimeSeries(_)
452 | QueryExpr::CreateQueue(_)
453 | QueryExpr::AlterQueue(_)
454 | QueryExpr::DropQueue(_)
455 | QueryExpr::QueueSelect(_)
456 | QueryExpr::QueueCommand(_)
457 | QueryExpr::KvCommand(_)
458 | QueryExpr::ConfigCommand(_)
459 | QueryExpr::CreateTree(_)
460 | QueryExpr::DropTree(_)
461 | QueryExpr::TreeCommand(_)
462 | QueryExpr::ExplainAlter(_)
463 | QueryExpr::TransactionControl(_)
464 | QueryExpr::MaintenanceCommand(_)
465 | QueryExpr::VcsCommand(_)
466 | QueryExpr::CreateSchema(_)
467 | QueryExpr::DropSchema(_)
468 | QueryExpr::CreateSequence(_)
469 | QueryExpr::DropSequence(_)
470 | QueryExpr::CopyFrom(_)
471 | QueryExpr::CreateView(_)
472 | QueryExpr::DropView(_)
473 | QueryExpr::RefreshMaterializedView(_)
474 | QueryExpr::CreatePolicy(_)
475 | QueryExpr::DropPolicy(_)
476 | QueryExpr::CreateServer(_)
477 | QueryExpr::DropServer(_)
478 | QueryExpr::CreateForeignTable(_)
479 | QueryExpr::DropForeignTable(_)
480 | QueryExpr::Grant(_)
481 | QueryExpr::Revoke(_)
482 | QueryExpr::AlterUser(_)
483 | QueryExpr::CreateUser(_)
484 | QueryExpr::CreateIamPolicy { .. }
485 | QueryExpr::DropIamPolicy { .. }
486 | QueryExpr::AttachPolicy { .. }
487 | QueryExpr::DetachPolicy { .. }
488 | QueryExpr::ShowPolicies { .. }
489 | QueryExpr::ShowEffectivePermissions { .. }
490 | QueryExpr::RankOf(_)
491 | QueryExpr::ApproxRankOf(_)
492 | QueryExpr::RankRange(_)
493 | QueryExpr::SimulatePolicy { .. }
494 | QueryExpr::LintPolicy { .. }
495 | QueryExpr::MigratePolicyMode { .. }
496 | QueryExpr::CreateMigration(_)
497 | QueryExpr::ApplyMigration(_)
498 | QueryExpr::RollbackMigration(_)
499 | QueryExpr::ExplainMigration(_)
500 | QueryExpr::EventsBackfill(_)
501 | QueryExpr::EventsBackfillStatus { .. } => false,
502 }
503 }
504}
505
506struct PushdownPredicatesRule;
508
509impl RewriteRule for PushdownPredicatesRule {
510 fn name(&self) -> &str {
511 "PushdownPredicates"
512 }
513
514 fn apply(&self, query: QueryExpr, ctx: &mut RewriteContext) -> QueryExpr {
515 match query {
516 QueryExpr::Join(mut jq) => {
517 jq.left = Box::new(self.apply(*jq.left, ctx));
523 jq.right = Box::new(self.apply(*jq.right, ctx));
524 ctx.stats.predicates_pushed += 1;
525 QueryExpr::Join(jq)
526 }
527 other => other,
528 }
529 }
530
531 fn is_applicable(&self, query: &QueryExpr) -> bool {
532 matches!(query, QueryExpr::Join(_))
533 }
534}
535
536struct EliminateDeadCodeRule;
538
539impl RewriteRule for EliminateDeadCodeRule {
540 fn name(&self) -> &str {
541 "EliminateDeadCode"
542 }
543
544 fn apply(&self, query: QueryExpr, _ctx: &mut RewriteContext) -> QueryExpr {
545 match query {
546 QueryExpr::Table(mut tq) => {
547 if let Some(filter) = effective_table_filter(&tq).as_ref() {
549 if is_always_true(filter) {
550 tq.filter = None;
551 }
552 }
553 QueryExpr::Table(tq)
554 }
555 other => other,
556 }
557 }
558
559 fn is_applicable(&self, query: &QueryExpr) -> bool {
560 matches!(query, QueryExpr::Table(_))
561 }
562}
563
564struct FoldConstantsRule;
566
567impl RewriteRule for FoldConstantsRule {
568 fn name(&self) -> &str {
569 "FoldConstants"
570 }
571
572 fn apply(&self, query: QueryExpr, _ctx: &mut RewriteContext) -> QueryExpr {
573 query
576 }
577
578 fn is_applicable(&self, _query: &QueryExpr) -> bool {
579 true
580 }
581}
582
583fn projection_name(proj: &Projection) -> String {
588 match proj {
589 Projection::All => "*".to_string(),
590 Projection::Column(name) => name.clone(),
591 Projection::Alias(_, alias) => alias.clone(),
592 Projection::Function(name, _) => name
593 .split_once(':')
594 .map(|(_, alias)| alias.to_string())
595 .unwrap_or_else(|| name.clone()),
596 Projection::Expression(expr, alias) => {
597 alias.clone().unwrap_or_else(|| format!("{:?}", expr))
598 }
599 Projection::Field(field, alias) => alias.clone().unwrap_or_else(|| format!("{:?}", field)),
600 Projection::Window { name, alias, .. } => alias.clone().unwrap_or_else(|| name.clone()),
601 }
602}
603
604fn simplify_filter(filter: AstFilter, ctx: &mut RewriteContext) -> AstFilter {
605 match filter {
606 AstFilter::And(left, right) => {
607 let left = simplify_filter(*left, ctx);
608 let right = simplify_filter(*right, ctx);
609
610 if is_always_true(&left) {
612 ctx.stats.filters_simplified += 1;
613 return right;
614 }
615 if is_always_true(&right) {
616 ctx.stats.filters_simplified += 1;
617 return left;
618 }
619
620 if is_always_false(&left) || is_always_false(&right) {
622 ctx.stats.filters_simplified += 1;
623 return AstFilter::Compare {
624 field: FieldRef::TableColumn {
625 table: String::new(),
626 column: "1".to_string(),
627 },
628 op: CompareOp::Eq,
629 value: Value::Integer(0),
630 };
631 }
632
633 AstFilter::And(Box::new(left), Box::new(right))
634 }
635 AstFilter::Or(left, right) => {
636 let left = simplify_filter(*left, ctx);
637 let right = simplify_filter(*right, ctx);
638
639 if is_always_false(&left) {
641 ctx.stats.filters_simplified += 1;
642 return right;
643 }
644 if is_always_false(&right) {
645 ctx.stats.filters_simplified += 1;
646 return left;
647 }
648
649 if is_always_true(&left) || is_always_true(&right) {
651 ctx.stats.filters_simplified += 1;
652 return AstFilter::Compare {
653 field: FieldRef::TableColumn {
654 table: String::new(),
655 column: "1".to_string(),
656 },
657 op: CompareOp::Eq,
658 value: Value::Integer(1),
659 };
660 }
661
662 AstFilter::Or(Box::new(left), Box::new(right))
663 }
664 AstFilter::Not(inner) => {
665 let inner = simplify_filter(*inner, ctx);
666
667 if let AstFilter::Not(double_inner) = inner {
669 ctx.stats.filters_simplified += 1;
670 return *double_inner;
671 }
672
673 AstFilter::Not(Box::new(inner))
674 }
675 other => other,
676 }
677}
678
679fn is_always_true(filter: &AstFilter) -> bool {
680 match filter {
681 AstFilter::Compare { field, op, value } => {
682 matches!(field, FieldRef::TableColumn { column, .. } if column == "1")
684 && matches!(op, CompareOp::Eq)
685 && matches!(value, Value::Integer(1))
686 }
687 _ => false,
688 }
689}
690
691fn is_always_false(filter: &AstFilter) -> bool {
692 match filter {
693 AstFilter::Compare { field, op, value } => {
694 matches!(field, FieldRef::TableColumn { column, .. } if column == "1")
696 && matches!(op, CompareOp::Eq)
697 && matches!(value, Value::Integer(0))
698 }
699 _ => false,
700 }
701}
702
703#[cfg(test)]
704mod tests {
705 use super::*;
706 use crate::ast::{JoinCondition, TableQuery, WindowSpec};
707
708 fn make_field(name: &str) -> FieldRef {
709 FieldRef::TableColumn {
710 table: String::new(),
711 column: name.to_string(),
712 }
713 }
714
715 #[test]
716 fn test_simplify_and_with_true() {
717 let mut ctx = RewriteContext::default();
718
719 let filter = AstFilter::And(
720 Box::new(AstFilter::Compare {
721 field: make_field("1"),
722 op: CompareOp::Eq,
723 value: Value::Integer(1),
724 }),
725 Box::new(AstFilter::Compare {
726 field: make_field("x"),
727 op: CompareOp::Eq,
728 value: Value::Integer(5),
729 }),
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
742 #[test]
743 fn test_simplify_double_not() {
744 let mut ctx = RewriteContext::default();
745
746 let filter = AstFilter::Not(Box::new(AstFilter::Not(Box::new(AstFilter::Compare {
747 field: make_field("x"),
748 op: CompareOp::Eq,
749 value: Value::Integer(5),
750 }))));
751
752 let simplified = simplify_filter(filter, &mut ctx);
753
754 match simplified {
755 AstFilter::Compare { field, .. } => {
756 assert!(matches!(field, FieldRef::TableColumn { column, .. } if column == "x"));
757 }
758 _ => panic!("Expected Compare filter"),
759 }
760 }
761
762 #[test]
763 fn projection_name_uses_visible_output_name_for_all_projection_shapes() {
764 assert_eq!(projection_name(&Projection::All), "*");
765 assert_eq!(
766 projection_name(&Projection::Column("raw".to_string())),
767 "raw"
768 );
769 assert_eq!(
770 projection_name(&Projection::Alias("raw".to_string(), "alias".to_string())),
771 "alias"
772 );
773 assert_eq!(
774 projection_name(&Projection::Function(
775 "LOWER:display".to_string(),
776 Vec::new()
777 )),
778 "display"
779 );
780 assert_eq!(
781 projection_name(&Projection::Expression(
782 Box::new(AstFilter::Compare {
783 field: make_field("x"),
784 op: CompareOp::Eq,
785 value: Value::Integer(1),
786 }),
787 Some("expr_alias".to_string()),
788 )),
789 "expr_alias"
790 );
791 assert_eq!(
792 projection_name(&Projection::Field(
793 FieldRef::node_prop("n", "name"),
794 Some("node_name".to_string()),
795 )),
796 "node_name"
797 );
798 assert_eq!(
799 projection_name(&Projection::Window {
800 name: "ROW_NUMBER".to_string(),
801 args: Vec::new(),
802 window: Box::new(WindowSpec::default()),
803 alias: Some("rn".to_string()),
804 }),
805 "rn"
806 );
807 }
808
809 #[test]
810 fn normalize_rule_sorts_table_columns_by_output_name() {
811 let mut table = TableQuery::new("users");
812 table.columns = vec![
813 Projection::Column("z".to_string()),
814 Projection::Function("LOWER:a_alias".to_string(), Vec::new()),
815 Projection::Alias("name".to_string(), "m".to_string()),
816 ];
817
818 let mut ctx = RewriteContext::default();
819 let normalized = NormalizeRule.apply(QueryExpr::Table(table), &mut ctx);
820 let QueryExpr::Table(table) = normalized else {
821 panic!("expected table query");
822 };
823
824 assert_eq!(ctx.stats.expressions_normalized, 1);
825 assert_eq!(
826 table
827 .columns
828 .iter()
829 .map(projection_name)
830 .collect::<Vec<_>>(),
831 vec!["a_alias", "m", "z"]
832 );
833 }
834
835 #[test]
836 fn simplify_filter_covers_or_true_false_and_not_paths() {
837 let mut ctx = RewriteContext::default();
838 let truth = AstFilter::Compare {
839 field: make_field("1"),
840 op: CompareOp::Eq,
841 value: Value::Integer(1),
842 };
843 let falsehood = AstFilter::Compare {
844 field: make_field("1"),
845 op: CompareOp::Eq,
846 value: Value::Integer(0),
847 };
848 let predicate = AstFilter::Compare {
849 field: make_field("x"),
850 op: CompareOp::Eq,
851 value: Value::Integer(5),
852 };
853
854 assert_eq!(
855 simplify_filter(
856 AstFilter::Or(Box::new(falsehood.clone()), Box::new(predicate.clone())),
857 &mut ctx,
858 ),
859 predicate
860 );
861 assert!(is_always_true(&simplify_filter(
862 AstFilter::Or(Box::new(truth), Box::new(falsehood.clone())),
863 &mut ctx,
864 )));
865 assert!(is_always_false(&simplify_filter(
866 AstFilter::And(
867 Box::new(falsehood),
868 Box::new(AstFilter::IsNotNull(make_field("x")))
869 ),
870 &mut ctx,
871 )));
872 assert!(ctx.stats.filters_simplified >= 3);
873 }
874
875 #[test]
876 fn query_rewriter_runs_rules_until_fixed_point_and_exposes_context() {
877 let mut table = TableQuery::new("users");
878 table.filter = Some(AstFilter::And(
879 Box::new(AstFilter::Compare {
880 field: make_field("1"),
881 op: CompareOp::Eq,
882 value: Value::Integer(1),
883 }),
884 Box::new(AstFilter::Compare {
885 field: make_field("age"),
886 op: CompareOp::Ge,
887 value: Value::Integer(18),
888 }),
889 ));
890
891 let mut ctx = RewriteContext::default();
892 let rewritten =
893 QueryRewriter::default().rewrite_with_context(QueryExpr::Table(table), &mut ctx);
894 let QueryExpr::Table(table) = rewritten else {
895 panic!("expected table query");
896 };
897
898 assert!(matches!(
899 table.filter,
900 Some(AstFilter::Compare {
901 field: FieldRef::TableColumn { column, .. },
902 op: CompareOp::Ge,
903 value: Value::Integer(18),
904 }) if column == "age"
905 ));
906 assert!(ctx.stats.filters_simplified >= 1);
907 }
908
909 #[test]
910 fn query_rewriter_recurses_into_join_children_and_tracks_pushdown() {
911 let mut left = TableQuery::new("users");
912 left.columns = vec![
913 Projection::Column("z".to_string()),
914 Projection::Column("a".to_string()),
915 ];
916 left.filter = Some(AstFilter::And(
917 Box::new(AstFilter::Compare {
918 field: make_field("1"),
919 op: CompareOp::Eq,
920 value: Value::Integer(1),
921 }),
922 Box::new(AstFilter::Compare {
923 field: make_field("age"),
924 op: CompareOp::Ge,
925 value: Value::Integer(18),
926 }),
927 ));
928
929 let join = JoinQuery::new(
930 QueryExpr::Table(left),
931 QueryExpr::Table(TableQuery::new("orders")),
932 JoinCondition::new(make_field("id"), make_field("user_id")),
933 );
934
935 let mut ctx = RewriteContext::default();
936 let rewritten =
937 QueryRewriter::default().rewrite_with_context(QueryExpr::Join(join), &mut ctx);
938 let QueryExpr::Join(join) = rewritten else {
939 panic!("expected join query");
940 };
941 let QueryExpr::Table(left) = join.left.as_ref() else {
942 panic!("expected table on left side");
943 };
944
945 assert_eq!(
946 left.columns.iter().map(projection_name).collect::<Vec<_>>(),
947 vec!["a", "z"]
948 );
949 assert!(matches!(
950 &left.filter,
951 Some(AstFilter::Compare {
952 field: FieldRef::TableColumn { ref column, .. },
953 op: CompareOp::Ge,
954 value: Value::Integer(18),
955 }) if column == "age"
956 ));
957 assert!(ctx.stats.predicates_pushed >= 1);
958 }
959
960 #[test]
961 fn query_rewriter_eliminates_always_true_table_filters() {
962 let mut table = TableQuery::new("users");
963 table.filter = Some(AstFilter::Compare {
964 field: make_field("1"),
965 op: CompareOp::Eq,
966 value: Value::Integer(1),
967 });
968
969 let rewritten = QueryRewriter::default().rewrite(QueryExpr::Table(table));
970 let QueryExpr::Table(table) = rewritten else {
971 panic!("expected table query");
972 };
973
974 assert!(table.filter.is_none());
975 }
976
977 struct CountingRule;
978
979 impl RewriteRule for CountingRule {
980 fn name(&self) -> &str {
981 "CountingRule"
982 }
983
984 fn apply(&self, query: QueryExpr, ctx: &mut RewriteContext) -> QueryExpr {
985 ctx.warnings.push(self.name().to_string());
986 query
987 }
988
989 fn is_applicable(&self, query: &QueryExpr) -> bool {
990 matches!(query, QueryExpr::Table(_))
991 }
992 }
993
994 #[test]
995 fn custom_rules_can_be_added_after_defaults() {
996 let mut rewriter = QueryRewriter::new();
997 rewriter.add_rule(Box::new(CountingRule));
998
999 let mut ctx = RewriteContext::default();
1000 let rewritten =
1001 rewriter.rewrite_with_context(QueryExpr::Table(TableQuery::new("users")), &mut ctx);
1002
1003 assert!(matches!(rewritten, QueryExpr::Table(_)));
1004 assert!(ctx.warnings.iter().any(|warning| warning == "CountingRule"));
1005 }
1006}