feat: remove nulls from in list

GH-8688
apache · Dec 31, 2023 · 665ca39 · 665ca39
1 parent 7fcfb91
commit 665ca39
Showing 1 changed file with 34 additions and 0 deletions.
diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs
@@ -555,6 +555,19 @@ impl<'a, S: SimplifyInfo> TreeNodeRewriter for Simplifier<'a, S> {
                     )
                 }
             }
+
+            // expr in (null, x, y, z) --> expr in (x, y, z)
+            // expr not in (null, x, y, z) --> expr not in (x, y, z)
+            Expr::InList(InList {
+                expr,
+                list,
+                negated,
+            }) if !is_null(&expr) => Expr::InList(InList {
+                expr,
+                list: list.into_iter().filter(|expr| !is_null(expr)).collect(),
+                negated,
+            }),
+
             //
             // Rules for NotEq
             //
@@ -3116,6 +3129,27 @@ mod tests {
             lit(true)
         );
 
+        // c1 in (null, ...)  -> c1 in (...)
+        assert_eq!(
+            simplify(in_list(
+                col("c1"),
+                vec![lit_bool_null(), lit(1), lit(2), lit(3), lit(4)],
+                true
+            )),
+            in_list(col("c1"), vec![lit(1), lit(2), lit(3), lit(4)], true)
+        );
+
+
+        // c1 not in (null, ...)  -> c1 not in (...)
+        assert_eq!(
+            simplify(in_list(
+                col("c1"),
+                vec![lit_bool_null(), lit(1), lit(2), lit(3), lit(4)],
+                false
+            )),
+            in_list(col("c1"), vec![lit(1), lit(2), lit(3), lit(4)], false)
+        );
+
         assert_eq!(
             simplify(in_list(col("c1"), vec![lit(1)], false)),
             col("c1").eq(lit(1))