extract all columns from underlying expr

risingwavelabs · Feb 27, 2024 · fbc1f9a · fbc1f9a
1 parent e9386d1
commit fbc1f9a
Showing 1 changed file with 45 additions and 12 deletions.
diff --git a/src/frontend/src/expr/utils.rs b/src/frontend/src/expr/utils.rs
@@ -623,16 +623,50 @@ fn check_pattern(e1: ExprImpl, e2: ExprImpl) -> (bool, Option<ExprImpl>) {
         func_type == ExprType::IsNull || func_type == ExprType::IsNotNull
     }
 
+    /// Simply extract every possible `InputRef` out from the input `expr`
+    fn extract_column(expr: ExprImpl, columns: &mut Vec<ExprImpl>) {
+        match expr {
+            ExprImpl::FunctionCall(func_call) => {
+                // `IsNotNull( ... )` or `IsNull( ... )` will be ignored
+                if is_null_or_not_null(func_call.func_type()) {
+                    return;
+                }
+                for sub_expr in func_call.inputs().iter() {
+                    extract_column(sub_expr.clone(), columns);
+                }
+            }
+            ExprImpl::InputRef(_) => {
+                columns.push(expr);
+            }
+            _ => (),
+        }
+    }
+
     /// Try wrapping inner expression with `IsNotNull`
-    fn try_wrap_inner_expression(expr: ExprImpl, func_type: ExprType) -> Option<ExprImpl> {
-        if is_null_or_not_null(func_type) {
-            None
-        } else {
-            let Ok(expr) = FunctionCall::new(ExprType::IsNotNull, vec![expr]) else {
+    /// Note: only column (i.e., `InputRef`) will be extracted and connected via `AND`
+    fn try_wrap_inner_expression(expr: ExprImpl) -> Option<ExprImpl> {
+        let mut columns = vec![];
+
+        extract_column(expr, &mut columns);
+        if columns.is_empty() {
+            return None;
+        }
+
+        let mut inputs: Vec<ExprImpl> = vec![];
+        // From [`c1`, `c2`, ... , `cn`] to [`IsNotNull(c1)`, ... , `IsNotNull(cn)`]
+        for column in columns {
+            let Ok(expr) = FunctionCall::new(ExprType::IsNotNull, vec![column]) else {
                 return None;
             };
-            Some(expr.into())
+            inputs.push(expr.into());
         }
+
+        // Connect them with `AND`
+        // i.e., AND [`IsNotNull(c1)`, ... , `IsNotNull(cn)`]
+        let Ok(expr) = FunctionCall::new(ExprType::And, inputs) else {
+            return None;
+        };
+        Some(expr.into())
     }
 
     let ExprImpl::FunctionCall(e1_func) = e1.clone() else {
@@ -644,18 +678,16 @@ fn check_pattern(e1: ExprImpl, e2: ExprImpl) -> (bool, Option<ExprImpl>) {
     if e1_func.func_type() != ExprType::Not && e2_func.func_type() != ExprType::Not {
         return (false, None);
     }
-    println!("e1: {:#?}", e1);
-    println!("e2: {:#?}", e2);
     if e1_func.func_type() != ExprType::Not {
         if e2_func.inputs().len() != 1 {
             return (false, None);
         }
-        (e1 == e2_func.inputs()[0].clone(), None)
+        (e1 == e2_func.inputs()[0].clone(), try_wrap_inner_expression(e1))
     } else {
         if e1_func.inputs().len() != 1 {
             return (false, None);
         }
-        (e2 == e1_func.inputs()[0].clone(), None)
+        (e2 == e1_func.inputs()[0].clone(), try_wrap_inner_expression(e2))
     }
 }
 
@@ -675,10 +707,11 @@ impl ExprRewriter for SimplifyFilterExpressionRewriter {
             return expr;
         }
         let inputs = func_call.inputs();
-        let (optimizable_flag, _) = check_pattern(inputs[0].clone(), inputs[1].clone());
+        let (optimizable_flag, columns) = check_pattern(inputs[0].clone(), inputs[1].clone());
         if optimizable_flag {
             match func_call.func_type() {
-                ExprType::Or => ExprImpl::literal_bool(true),
+                ExprType::Or => if let Some(columns) = columns { columns } else { ExprImpl::literal_bool(true) },
+                // `AND` will always be false, no matter the underlying columns are null or not
                 ExprType::And => ExprImpl::literal_bool(false),
                 _ => expr,
             }