Attempts to optimize `knowledge_stack` is a pre-allocated vector but is expected not to contain any elements.
(
expr: &mut MirScalarExpr,
column_types: &[ReprColumnType],
column_knowledge: &[DatumKnowledge],
knowledge_stack: &mut Vec<DatumKnowledge>,
)
| 752 | /// |
| 753 | /// `knowledge_stack` is a pre-allocated vector but is expected not to contain any elements. |
| 754 | fn optimize( |
| 755 | expr: &mut MirScalarExpr, |
| 756 | column_types: &[ReprColumnType], |
| 757 | column_knowledge: &[DatumKnowledge], |
| 758 | knowledge_stack: &mut Vec<DatumKnowledge>, |
| 759 | ) -> Result<DatumKnowledge, TransformError> { |
| 760 | // Storage for `DatumKnowledge` being propagated up through the |
| 761 | // `MirScalarExpr`. When a node is visited, pop off as many `DatumKnowledge` |
| 762 | // as the number of children the node has, and then push the |
| 763 | // `DatumKnowledge` corresponding to the node back onto the stack. |
| 764 | // Post-order traversal means that if a node has `n` children, the top `n` |
| 765 | // `DatumKnowledge` in the stack are the `DatumKnowledge` corresponding to |
| 766 | // the children. |
| 767 | assert!(knowledge_stack.is_empty()); |
| 768 | expr.visit_mut_pre_post( |
| 769 | &mut |e| { |
| 770 | if let MirScalarExpr::If { then, els, .. } = e { |
| 771 | Some(vec![then, els]) |
| 772 | } else { |
| 773 | None |
| 774 | } |
| 775 | }, |
| 776 | &mut |e| { |
| 777 | let result = match e { |
| 778 | MirScalarExpr::Column(index, _) => { |
| 779 | let index = *index; |
| 780 | if let DatumKnowledge::Lit { value, typ } = &column_knowledge[index] { |
| 781 | *e = MirScalarExpr::Literal( |
| 782 | value.clone(), |
| 783 | typ.clone().nullable(column_knowledge[index].nullable()), |
| 784 | ); |
| 785 | } |
| 786 | column_knowledge[index].clone() |
| 787 | } |
| 788 | MirScalarExpr::Literal(_, _) | MirScalarExpr::CallUnmaterializable(_) => { |
| 789 | DatumKnowledge::from(&*e) |
| 790 | } |
| 791 | MirScalarExpr::CallUnary { func, expr: _ } => { |
| 792 | let knowledge = knowledge_stack.pop().unwrap(); |
| 793 | if matches!(&knowledge, DatumKnowledge::Lit { .. }) { |
| 794 | e.reduce(column_types); |
| 795 | } else if func == &UnaryFunc::IsNull(func::IsNull) && !knowledge.nullable() { |
| 796 | *e = MirScalarExpr::literal_false(); |
| 797 | }; |
| 798 | DatumKnowledge::from(&*e) |
| 799 | } |
| 800 | MirScalarExpr::CallBinary { |
| 801 | func: _, |
| 802 | expr1: _, |
| 803 | expr2: _, |
| 804 | } => { |
| 805 | let knowledge2 = knowledge_stack.pop().unwrap(); |
| 806 | let knowledge1 = knowledge_stack.pop().unwrap(); |
| 807 | if crate::any![ |
| 808 | matches!(knowledge1, DatumKnowledge::Lit { .. }), |
| 809 | matches!(knowledge2, DatumKnowledge::Lit { .. }), |
| 810 | ] { |
| 811 | e.reduce(column_types); |