1use itertools::Itertools;
17use mz_ore::id_gen::IdGen;
18use mz_ore::stack::{CheckedRecursion, RecursionGuard};
19use mz_repr::namespaces::{MZ_CATALOG_SCHEMA, MZ_UNSAFE_SCHEMA, PG_CATALOG_SCHEMA};
20use mz_sql_parser::ast::visit_mut::{self, VisitMut, VisitMutNode};
21use mz_sql_parser::ast::{
22 Expr, Function, FunctionArgs, HomogenizingFunction, Ident, IsExprConstruct, Op, OrderByExpr,
23 Query, Select, SelectItem, TableAlias, TableFactor, TableWithJoins, Value, WindowSpec,
24};
25use mz_sql_parser::ident;
26
27use crate::names::{Aug, PartialItemName, ResolvedDataType, ResolvedItemName};
28use crate::plan::{PlanError, StatementContext};
29use crate::{ORDINALITY_COL_NAME, normalize};
30
31pub(crate) fn transform<N>(scx: &StatementContext, node: &mut N) -> Result<(), PlanError>
32where
33 N: for<'a> VisitMutNode<'a, Aug>,
34{
35 let mut func_rewriter = FuncRewriter::new(scx);
36 node.visit_mut(&mut func_rewriter);
37 func_rewriter.status?;
38
39 let mut desugarer = Desugarer::new(scx);
40 node.visit_mut(&mut desugarer);
41 desugarer.status
42}
43
44struct FuncRewriter<'a> {
65 scx: &'a StatementContext<'a>,
66 status: Result<(), PlanError>,
67 rewriting_table_factor: bool,
68}
69
70impl<'a> FuncRewriter<'a> {
71 fn new(scx: &'a StatementContext<'a>) -> FuncRewriter<'a> {
72 FuncRewriter {
73 scx,
74 status: Ok(()),
75 rewriting_table_factor: false,
76 }
77 }
78
79 fn resolve_known_valid_data_type(&self, name: &PartialItemName) -> ResolvedDataType {
80 let item = self
81 .scx
82 .catalog
83 .resolve_type(name)
84 .expect("data type known to be valid");
85 let full_name = self.scx.catalog.resolve_full_name(item.name());
86 ResolvedDataType::Named {
87 id: item.id(),
88 qualifiers: item.name().qualifiers.clone(),
89 full_name,
90 modifiers: vec![],
91 print_id: true,
92 }
93 }
94
95 fn int32_data_type(&self) -> ResolvedDataType {
96 self.resolve_known_valid_data_type(&PartialItemName {
97 database: None,
98 schema: Some(PG_CATALOG_SCHEMA.into()),
99 item: "int4".into(),
100 })
101 }
102
103 fn plan_divide(lhs: Expr<Aug>, rhs: Expr<Aug>) -> Expr<Aug> {
106 lhs.divide(Expr::Case {
107 operand: None,
108 conditions: vec![rhs.clone().equals(Expr::number("0"))],
109 results: vec![Expr::null()],
110 else_result: Some(Box::new(rhs)),
111 })
112 }
113
114 fn plan_agg(
115 &mut self,
116 name: ResolvedItemName,
117 expr: Expr<Aug>,
118 order_by: Vec<OrderByExpr<Aug>>,
119 filter: Option<Box<Expr<Aug>>>,
120 distinct: bool,
121 over: Option<WindowSpec<Aug>>,
122 ) -> Expr<Aug> {
123 if self.rewriting_table_factor && self.status.is_ok() {
124 self.status = Err(PlanError::Unstructured(
125 "aggregate functions are not supported in functions in FROM".to_string(),
126 ))
127 }
128 Expr::Function(Function {
129 name,
130 args: FunctionArgs::Args {
131 args: vec![expr],
132 order_by,
133 },
134 filter,
135 over,
136 distinct,
137 })
138 }
139
140 fn plan_avg(
141 &mut self,
142 expr: Expr<Aug>,
143 filter: Option<Box<Expr<Aug>>>,
144 distinct: bool,
145 over: Option<WindowSpec<Aug>>,
146 ) -> Expr<Aug> {
147 let sum = self
148 .plan_agg(
149 self.scx
150 .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
151 expr.clone(),
152 vec![],
153 filter.clone(),
154 distinct,
155 over.clone(),
156 )
157 .call_unary(
158 self.scx
159 .dangerous_resolve_name(vec![MZ_UNSAFE_SCHEMA, "mz_avg_promotion"]),
160 );
161 let count = self.plan_agg(
162 self.scx
163 .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "count"]),
164 expr,
165 vec![],
166 filter,
167 distinct,
168 over,
169 );
170 Self::plan_divide(sum, count)
171 }
172
173 fn plan_avg_internal_v1(
175 &mut self,
176 expr: Expr<Aug>,
177 filter: Option<Box<Expr<Aug>>>,
178 distinct: bool,
179 over: Option<WindowSpec<Aug>>,
180 ) -> Expr<Aug> {
181 let sum = self
182 .plan_agg(
183 self.scx
184 .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
185 expr.clone(),
186 vec![],
187 filter.clone(),
188 distinct,
189 over.clone(),
190 )
191 .call_unary(
192 self.scx
193 .dangerous_resolve_name(vec![MZ_UNSAFE_SCHEMA, "mz_avg_promotion_internal_v1"]),
194 );
195 let count = self.plan_agg(
196 self.scx
197 .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "count"]),
198 expr,
199 vec![],
200 filter,
201 distinct,
202 over,
203 );
204 Self::plan_divide(sum, count)
205 }
206
207 fn plan_variance(
208 &mut self,
209 expr: Expr<Aug>,
210 filter: Option<Box<Expr<Aug>>>,
211 distinct: bool,
212 sample: bool,
213 over: Option<WindowSpec<Aug>>,
214 ) -> Expr<Aug> {
215 let expr = expr.call_unary(
230 self.scx
231 .dangerous_resolve_name(vec![MZ_UNSAFE_SCHEMA, "mz_avg_promotion"]),
232 );
233 let expr_squared = expr.clone().multiply(expr.clone());
234 let sum_squares = if distinct {
235 let case_squared = |condition| Expr::Case {
252 operand: None,
253 conditions: vec![condition],
254 results: vec![expr_squared.clone()],
255 else_result: None,
256 };
257 let sum_squares_nonneg = self.plan_agg(
258 self.scx
259 .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
260 case_squared(expr.clone().gt_eq(Expr::number("0"))),
261 vec![],
262 filter.clone(),
263 distinct,
264 over.clone(),
265 );
266 let sum_squares_neg = self.plan_agg(
267 self.scx
268 .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
269 case_squared(expr.clone().lt(Expr::number("0"))),
270 vec![],
271 filter.clone(),
272 distinct,
273 over.clone(),
274 );
275 let coalesce_zero = |sum| Expr::HomogenizingFunction {
276 function: HomogenizingFunction::Coalesce,
277 exprs: vec![sum, Expr::number("0")],
278 };
279 coalesce_zero(sum_squares_nonneg).binop(Op::bare("+"), coalesce_zero(sum_squares_neg))
280 } else {
281 self.plan_agg(
282 self.scx
283 .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
284 expr_squared,
285 vec![],
286 filter.clone(),
287 distinct,
288 over.clone(),
289 )
290 };
291 let sum = self.plan_agg(
292 self.scx
293 .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
294 expr.clone(),
295 vec![],
296 filter.clone(),
297 distinct,
298 over.clone(),
299 );
300 let sum_squared = sum.clone().multiply(sum);
301 let count = self.plan_agg(
302 self.scx
303 .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "count"]),
304 expr,
305 vec![],
306 filter,
307 distinct,
308 over,
309 );
310 let result = Self::plan_divide(
311 sum_squares.minus(Self::plan_divide(sum_squared, count.clone())),
312 if sample {
313 count.minus(Expr::number("1"))
314 } else {
315 count
316 },
317 );
318 let result_is_null = Expr::IsExpr {
333 expr: Box::new(result.clone()),
334 construct: IsExprConstruct::Null,
335 negated: false,
336 };
337 Expr::Case {
338 operand: None,
339 conditions: vec![result_is_null],
340 results: vec![Expr::Value(Value::Null)],
341 else_result: Some(Box::new(Expr::HomogenizingFunction {
342 function: HomogenizingFunction::Greatest,
343 exprs: vec![result, Expr::number("0")],
344 })),
345 }
346 }
347
348 fn plan_stddev(
349 &mut self,
350 expr: Expr<Aug>,
351 filter: Option<Box<Expr<Aug>>>,
352 distinct: bool,
353 sample: bool,
354 over: Option<WindowSpec<Aug>>,
355 ) -> Expr<Aug> {
356 self.plan_variance(expr, filter, distinct, sample, over)
357 .call_unary(
358 self.scx
359 .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sqrt"]),
360 )
361 }
362
363 fn plan_bool_and(
364 &mut self,
365 expr: Expr<Aug>,
366 filter: Option<Box<Expr<Aug>>>,
367 distinct: bool,
368 over: Option<WindowSpec<Aug>>,
369 ) -> Expr<Aug> {
370 let sum = self.plan_agg(
382 self.scx
383 .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
384 expr.negate().cast(self.int32_data_type()),
385 vec![],
386 filter,
387 distinct,
388 over,
389 );
390 sum.equals(Expr::Value(Value::Number(0.to_string())))
391 }
392
393 fn plan_bool_or(
394 &mut self,
395 expr: Expr<Aug>,
396 filter: Option<Box<Expr<Aug>>>,
397 distinct: bool,
398 over: Option<WindowSpec<Aug>>,
399 ) -> Expr<Aug> {
400 let sum = self.plan_agg(
412 self.scx
413 .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "sum"]),
414 expr.or(Expr::Value(Value::Boolean(false)))
415 .cast(self.int32_data_type()),
416 vec![],
417 filter,
418 distinct,
419 over,
420 );
421 sum.gt(Expr::Value(Value::Number(0.to_string())))
422 }
423
424 fn rewrite_function(&mut self, func: &Function<Aug>) -> Option<(Ident, Expr<Aug>)> {
425 if let Function {
426 name,
427 args: FunctionArgs::Args { args, order_by: _ },
428 filter,
429 distinct,
430 over,
431 } = func
432 {
433 let pg_catalog_id = self
434 .scx
435 .catalog
436 .resolve_schema(None, PG_CATALOG_SCHEMA)
437 .expect("pg_catalog schema exists")
438 .id();
439 let mz_catalog_id = self
440 .scx
441 .catalog
442 .resolve_schema(None, MZ_CATALOG_SCHEMA)
443 .expect("mz_catalog schema exists")
444 .id();
445 let name = match name {
446 ResolvedItemName::Item {
447 qualifiers,
448 full_name,
449 ..
450 } => {
451 if ![*pg_catalog_id, *mz_catalog_id].contains(&qualifiers.schema_spec) {
452 return None;
453 }
454 full_name.item.clone()
455 }
456 _ => unreachable!(),
457 };
458
459 let filter = filter.clone();
460 let distinct = *distinct;
461 let over = over.clone();
462 let expr = if args.len() == 1 {
463 let arg = args[0].clone();
464 match name.as_str() {
465 "avg_internal_v1" => self.plan_avg_internal_v1(arg, filter, distinct, over),
466 "avg" => self.plan_avg(arg, filter, distinct, over),
467 "variance" | "var_samp" => {
468 self.plan_variance(arg, filter, distinct, true, over)
469 }
470 "var_pop" => self.plan_variance(arg, filter, distinct, false, over),
471 "stddev" | "stddev_samp" => self.plan_stddev(arg, filter, distinct, true, over),
472 "stddev_pop" => self.plan_stddev(arg, filter, distinct, false, over),
473 "bool_and" => self.plan_bool_and(arg, filter, distinct, over),
474 "bool_or" => self.plan_bool_or(arg, filter, distinct, over),
475 _ => return None,
476 }
477 } else if args.len() == 2 {
478 let (lhs, rhs) = (args[0].clone(), args[1].clone());
479 match name.as_str() {
480 "mod" => lhs.modulo(rhs),
481 "pow" => Expr::call(
482 self.scx
483 .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, "power"]),
484 vec![lhs, rhs],
485 ),
486 _ => return None,
487 }
488 } else {
489 return None;
490 };
491 Some((Ident::new_unchecked(name), expr))
492 } else {
493 None
494 }
495 }
496
497 fn rewrite_expr(&mut self, expr: &Expr<Aug>) -> Option<(Ident, Expr<Aug>)> {
498 match expr {
499 Expr::Function(function) => self.rewrite_function(function),
500 Expr::Identifier(ident) if ident.len() == 1 => {
504 let ident = normalize::ident(ident[0].clone());
505 let fn_ident = match ident.as_str() {
506 "current_role" => Some("current_user"),
507 "current_schema" | "current_timestamp" | "current_user" | "session_user"
508 | "current_catalog" => Some(ident.as_str()),
509 _ => None,
510 };
511 match fn_ident {
512 None => None,
513 Some(fn_ident) => {
514 let expr = Expr::call_nullary(
515 self.scx
516 .dangerous_resolve_name(vec![PG_CATALOG_SCHEMA, fn_ident]),
517 );
518 Some((Ident::new_unchecked(ident), expr))
519 }
520 }
521 }
522 _ => None,
523 }
524 }
525}
526
527impl<'ast> VisitMut<'ast, Aug> for FuncRewriter<'_> {
528 fn visit_select_item_mut(&mut self, item: &'ast mut SelectItem<Aug>) {
529 if let SelectItem::Expr { expr, alias: None } = item {
530 visit_mut::visit_expr_mut(self, expr);
531 if let Some((alias, expr)) = self.rewrite_expr(expr) {
532 *item = SelectItem::Expr {
533 expr,
534 alias: Some(alias),
535 };
536 }
537 } else {
538 visit_mut::visit_select_item_mut(self, item);
539 }
540 }
541
542 fn visit_table_with_joins_mut(&mut self, item: &'ast mut TableWithJoins<Aug>) {
543 visit_mut::visit_table_with_joins_mut(self, item);
544 match &mut item.relation {
545 TableFactor::Function {
546 function,
547 alias,
548 with_ordinality,
549 } => {
550 self.rewriting_table_factor = true;
551 if let Some((ident, expr)) = self.rewrite_function(function) {
554 let mut select = Select::default().project(SelectItem::Expr {
555 expr,
556 alias: Some(match &alias {
557 Some(TableAlias { name, columns, .. }) => {
558 columns.get(0).unwrap_or(name).clone()
559 }
560 None => ident,
561 }),
562 });
563
564 if *with_ordinality {
565 select = select.project(SelectItem::Expr {
566 expr: Expr::Value(Value::Number("1".into())),
567 alias: Some(ident!(ORDINALITY_COL_NAME)),
568 });
569 }
570
571 item.relation = TableFactor::Derived {
572 lateral: false,
573 subquery: Box::new(Query {
574 ctes: mz_sql_parser::ast::CteBlock::Simple(vec![]),
575 body: mz_sql_parser::ast::SetExpr::Select(Box::new(select)),
576 order_by: vec![],
577 limit: None,
578 offset: None,
579 }),
580 alias: alias.clone(),
581 }
582 }
583 self.rewriting_table_factor = false;
584 }
585 _ => {}
586 }
587 }
588
589 fn visit_expr_mut(&mut self, expr: &'ast mut Expr<Aug>) {
590 visit_mut::visit_expr_mut(self, expr);
591 if let Some((_name, new_expr)) = self.rewrite_expr(expr) {
592 *expr = new_expr;
593 }
594 }
595}
596
597struct Desugarer<'a> {
602 scx: &'a StatementContext<'a>,
603 status: Result<(), PlanError>,
604 id_gen: IdGen,
605 recursion_guard: RecursionGuard,
606}
607
608impl<'a> CheckedRecursion for Desugarer<'a> {
609 fn recursion_guard(&self) -> &RecursionGuard {
610 &self.recursion_guard
611 }
612}
613
614impl<'a, 'ast> VisitMut<'ast, Aug> for Desugarer<'a> {
615 fn visit_expr_mut(&mut self, expr: &'ast mut Expr<Aug>) {
616 self.visit_internal(Self::visit_expr_mut_internal, expr);
617 }
618}
619
620impl<'a> Desugarer<'a> {
621 fn visit_internal<F, X>(&mut self, f: F, x: X)
622 where
623 F: Fn(&mut Self, X) -> Result<(), PlanError>,
624 {
625 if self.status.is_ok() {
626 let status = self.checked_recur_mut(|d| f(d, x));
629 if self.status.is_ok() {
630 self.status = status;
631 }
632 }
633 }
634
635 fn new(scx: &'a StatementContext) -> Desugarer<'a> {
636 Desugarer {
637 scx,
638 status: Ok(()),
639 id_gen: Default::default(),
640 recursion_guard: RecursionGuard::with_limit(1024), }
642 }
643
644 fn visit_expr_mut_internal(&mut self, expr: &mut Expr<Aug>) -> Result<(), PlanError> {
645 while let Expr::Nested(e) = expr {
647 *expr = e.take();
648 }
649
650 if let Expr::Between {
653 expr: e,
654 low,
655 high,
656 negated,
657 } = expr
658 {
659 if *negated {
660 *expr = Expr::lt(*e.clone(), low.take()).or(e.take().gt(high.take()));
661 } else {
662 *expr = e.clone().gt_eq(low.take()).and(e.take().lt_eq(high.take()));
663 }
664 }
665
666 if let Expr::InList {
676 expr: e,
677 list,
678 negated,
679 } = expr
680 {
681 if let Expr::Row { .. } = &**e {
682 if *negated {
683 *expr = list
684 .drain(..)
685 .map(|r| e.clone().not_equals(r))
686 .reduce(|e1, e2| e1.and(e2))
687 .expect("list known to contain at least one element");
688 } else {
689 *expr = list
690 .drain(..)
691 .map(|r| e.clone().equals(r))
692 .reduce(|e1, e2| e1.or(e2))
693 .expect("list known to contain at least one element");
694 }
695 }
696 }
697
698 if let Expr::InSubquery {
701 expr: e,
702 subquery,
703 negated,
704 } = expr
705 {
706 if *negated {
707 *expr = Expr::AllSubquery {
708 left: Box::new(e.take()),
709 op: Op::bare("<>"),
710 right: Box::new(subquery.take()),
711 };
712 } else {
713 *expr = Expr::AnySubquery {
714 left: Box::new(e.take()),
715 op: Op::bare("="),
716 right: Box::new(subquery.take()),
717 };
718 }
719 }
720
721 if let Expr::AnyExpr { left, op, right } | Expr::AllExpr { left, op, right } = expr {
727 let binding = ident!("elem");
728
729 let subquery = Query::select(
730 Select::default()
731 .from(TableWithJoins {
732 relation: TableFactor::Function {
733 function: Function {
734 name: self
735 .scx
736 .dangerous_resolve_name(vec![MZ_CATALOG_SCHEMA, "unnest"]),
737 args: FunctionArgs::args(vec![right.take()]),
738 filter: None,
739 over: None,
740 distinct: false,
741 },
742 alias: Some(TableAlias {
743 name: ident!("_"),
744 columns: vec![binding.clone()],
745 strict: true,
746 }),
747 with_ordinality: false,
748 },
749 joins: vec![],
750 })
751 .project(SelectItem::Expr {
752 expr: Expr::Identifier(vec![binding]),
753 alias: None,
754 }),
755 );
756
757 let left = Box::new(left.take());
758
759 let op = op.clone();
760
761 *expr = match expr {
762 Expr::AnyExpr { .. } => Expr::AnySubquery {
763 left,
764 op,
765 right: Box::new(subquery),
766 },
767 Expr::AllExpr { .. } => Expr::AllSubquery {
768 left,
769 op,
770 right: Box::new(subquery),
771 },
772 _ => unreachable!(),
773 };
774 }
775
776 if let Expr::AnySubquery { left, op, right } | Expr::AllSubquery { left, op, right } = expr
782 {
783 let left = match &mut **left {
784 Expr::Row { .. } => left.take(),
785 _ => Expr::Row {
786 exprs: vec![left.take()],
787 },
788 };
789
790 let arity = match &left {
791 Expr::Row { exprs } => exprs.len(),
792 _ => unreachable!(),
793 };
794
795 let bindings: Vec<_> = (0..arity)
796 .map(|col| {
799 let unique_id = self.id_gen.allocate_id();
800 Ident::new_unchecked(format!("right_col{col}_{unique_id}"))
801 })
802 .collect();
803
804 let subquery_unique_id = self.id_gen.allocate_id();
805 let subquery_name = Ident::new_unchecked(format!("subquery{subquery_unique_id}"));
807 let select = Select::default()
808 .from(TableWithJoins::subquery(
809 right.take(),
810 TableAlias {
811 name: subquery_name,
812 columns: bindings.clone(),
813 strict: true,
814 },
815 ))
816 .project(SelectItem::Expr {
817 expr: left
818 .binop(
819 op.clone(),
820 Expr::Row {
821 exprs: bindings
822 .into_iter()
823 .map(|b| Expr::Identifier(vec![b]))
824 .collect(),
825 },
826 )
827 .call_unary(self.scx.dangerous_resolve_name(match expr {
828 Expr::AnySubquery { .. } => vec![MZ_UNSAFE_SCHEMA, "mz_any"],
829 Expr::AllSubquery { .. } => vec![MZ_UNSAFE_SCHEMA, "mz_all"],
830 _ => unreachable!(),
831 })),
832 alias: None,
833 });
834
835 *expr = Expr::Subquery(Box::new(Query::select(select)));
836 }
837
838 if let Expr::Op {
854 op,
855 expr1: left,
856 expr2: Some(right),
857 } = expr
858 {
859 if let (Expr::Row { exprs: left }, Expr::Row { exprs: right }) =
860 (&mut **left, &mut **right)
861 {
862 if matches!(normalize::op(op)?, "=" | "<>" | "<" | "<=" | ">" | ">=") {
863 if left.len() != right.len() {
864 sql_bail!("unequal number of entries in row expressions");
865 }
866 if left.is_empty() {
867 assert!(right.is_empty());
868 sql_bail!("cannot compare rows of zero length");
869 }
870 }
871 match normalize::op(op)? {
872 "=" | "<>" => {
873 let mut pairs = left.iter_mut().zip_eq(right);
874 let mut new = pairs
875 .next()
876 .map(|(l, r)| l.take().equals(r.take()))
877 .expect("cannot compare rows of zero length");
878 for (l, r) in pairs {
879 new = l.take().equals(r.take()).and(new);
880 }
881 if normalize::op(op)? == "<>" {
882 new = new.negate();
883 }
884 *expr = new;
885 }
886 "<" | "<=" | ">" | ">=" => {
887 let strict_op = match normalize::op(op)? {
888 "<" | "<=" => "<",
889 ">" | ">=" => ">",
890 _ => unreachable!(),
891 };
892 let (l, r) = (left.last_mut().unwrap(), right.last_mut().unwrap());
893 let mut new = l.take().binop(op.clone(), r.take());
894 for (l, r) in left
895 .iter_mut()
896 .rev()
897 .zip_eq(right.into_iter().rev())
898 .skip(1)
899 {
900 new = l
901 .clone()
902 .binop(Op::bare(strict_op), r.clone())
903 .or(l.take().equals(r.take()).and(new));
904 }
905 *expr = new;
906 }
907 _ if left.len() == 1 && right.len() == 1 => {
908 let left = left.remove(0);
909 let right = right.remove(0);
910 *expr = left.binop(op.clone(), right);
911 }
912 _ => (),
913 }
914 }
915 }
916
917 visit_mut::visit_expr_mut(self, expr);
918 Ok(())
919 }
920}