From af409ec29d2a98062d084ff1221efdc3e00c4364 Mon Sep 17 00:00:00 2001
From: "Heres, Daniel" <danielheres@gmail.com>
Date: Sun, 1 Aug 2021 11:43:22 +0200
Subject: [PATCH 1/3] Simplify inlist

---
 .../src/optimizer/simplify_expressions.rs     | 29 ++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/datafusion/src/optimizer/simplify_expressions.rs b/datafusion/src/optimizer/simplify_expressions.rs
index f629eaf95b5f4..1516c2c9401b2 100644
--- a/datafusion/src/optimizer/simplify_expressions.rs
+++ b/datafusion/src/optimizer/simplify_expressions.rs
@@ -251,6 +251,22 @@ fn simplify(expr: &Expr) -> Expr {
             op: *op,
             right: Box::new(simplify(right)),
         },
+        Expr::InList {
+            expr,
+            list,
+            negated,
+        } if list.len() == 2 => {
+            let e1 = (**expr).clone().eq(list[0].clone());
+            let e2 = (**expr).clone().eq(list[1].clone());
+
+            let exp = e1.or(e2);
+
+            if *negated {
+                exp.not()
+            } else {
+                exp
+            }
+        }
         _ => expr.clone(),
     }
 }
@@ -293,7 +309,9 @@ impl SimplifyExpressions {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::logical_plan::{and, binary_expr, col, lit, Expr, LogicalPlanBuilder};
+    use crate::logical_plan::{
+        and, binary_expr, col, in_list, lit, Expr, LogicalPlanBuilder,
+    };
     use crate::test::*;
 
     fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
@@ -396,6 +414,15 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn simplify_inlist() -> Result<()> {
+        let expr = in_list(col("c"), vec![lit(1), lit(2)], false);
+        let expected = col("c").eq(lit(1)).or(col("c").eq(lit(2)));
+
+        assert_eq!(simplify(&expr), expected);
+        Ok(())
+    }
+
     #[test]
     fn test_simplify_simple_and() -> Result<()> {
         // (c > 5) AND (c > 5)

From 608c54cacc4cdb76212e014d671670a9d1e4aba4 Mon Sep 17 00:00:00 2001
From: "Heres, Daniel" <danielheres@gmail.com>
Date: Sun, 1 Aug 2021 11:45:26 +0200
Subject: [PATCH 2/3] Single expr

---
 datafusion/src/optimizer/simplify_expressions.rs | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/datafusion/src/optimizer/simplify_expressions.rs b/datafusion/src/optimizer/simplify_expressions.rs
index 1516c2c9401b2..8e06ddbc188c0 100644
--- a/datafusion/src/optimizer/simplify_expressions.rs
+++ b/datafusion/src/optimizer/simplify_expressions.rs
@@ -251,6 +251,18 @@ fn simplify(expr: &Expr) -> Expr {
             op: *op,
             right: Box::new(simplify(right)),
         },
+        Expr::InList {
+            expr,
+            list,
+            negated,
+        } if list.len() == 1 => {
+            let exp = (**expr).clone().eq(list[0].clone());
+            if *negated {
+                exp.not()
+            } else {
+                exp
+            }
+        }
         Expr::InList {
             expr,
             list,

From 5c57b26d30b11158a25c866109217cb373ac5191 Mon Sep 17 00:00:00 2001
From: "Heres, Daniel" <danielheres@gmail.com>
Date: Sun, 1 Aug 2021 18:40:24 +0200
Subject: [PATCH 3/3] Make more generic, add tests

---
 .../src/optimizer/simplify_expressions.rs     | 49 +++++++++++--------
 1 file changed, 29 insertions(+), 20 deletions(-)

diff --git a/datafusion/src/optimizer/simplify_expressions.rs b/datafusion/src/optimizer/simplify_expressions.rs
index 8e06ddbc188c0..7c9e976d1d0bb 100644
--- a/datafusion/src/optimizer/simplify_expressions.rs
+++ b/datafusion/src/optimizer/simplify_expressions.rs
@@ -255,28 +255,19 @@ fn simplify(expr: &Expr) -> Expr {
             expr,
             list,
             negated,
-        } if list.len() == 1 => {
-            let exp = (**expr).clone().eq(list[0].clone());
+        } if list.len() >= 1 && list.len() <= 2 => {
             if *negated {
-                exp.not()
+                list.iter()
+                    .skip(1)
+                    .fold((**expr).clone().not_eq(list[0].clone()), |acc, e| {
+                        acc.and((**expr).clone().not_eq(e.clone()))
+                    })
             } else {
-                exp
-            }
-        }
-        Expr::InList {
-            expr,
-            list,
-            negated,
-        } if list.len() == 2 => {
-            let e1 = (**expr).clone().eq(list[0].clone());
-            let e2 = (**expr).clone().eq(list[1].clone());
-
-            let exp = e1.or(e2);
-
-            if *negated {
-                exp.not()
-            } else {
-                exp
+                list.iter()
+                    .skip(1)
+                    .fold((**expr).clone().eq(list[0].clone()), |acc, e| {
+                        acc.or((**expr).clone().eq(e.clone()))
+                    })
             }
         }
         _ => expr.clone(),
@@ -435,6 +426,24 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn simplify_inlist_negated() -> Result<()> {
+        let expr = in_list(col("c"), vec![lit(1), lit(2)], true);
+        let expected = col("c").not_eq(lit(1)).and(col("c").not_eq(lit(2)));
+
+        assert_eq!(simplify(&expr), expected);
+        Ok(())
+    }
+
+    #[test]
+    fn simplify_inlist_single() -> Result<()> {
+        let expr = in_list(col("c"), vec![lit(1)], false);
+        let expected = col("c").eq(lit(1));
+
+        assert_eq!(simplify(&expr), expected);
+        Ok(())
+    }
+
     #[test]
     fn test_simplify_simple_and() -> Result<()> {
         // (c > 5) AND (c > 5)