fix: Fix nullability of logical InSubquery expression

datafusion/expr/src/expr_schema.rs

@@ -366,7 +366,12 @@ impl ExprSchemable for Expr {
             | Expr::IsNotUnknown(_)
             | Expr::Exists { .. } => Ok(false),
             Expr::SetComparison(_) => Ok(true),
-            Expr::InSubquery(InSubquery { expr, .. }) => expr.nullable(input_schema),
+            Expr::InSubquery(InSubquery { expr, subquery, .. }) => {
+                let expr_nullable = expr.nullable(input_schema)?;
+                let subquery_nullable = subquery.subquery.schema().field(0).is_nullable();
+
+                Ok(expr_nullable | subquery_nullable)
+            }
             Expr::ScalarSubquery(subquery) => {
                 Ok(subquery.subquery.schema().field(0).is_nullable())
             }
@@ -796,8 +801,13 @@ mod tests {
     use std::collections::HashMap;
 
     use super::*;
-    use crate::{and, col, lit, not, or, out_ref_col_with_metadata, when};
+    use crate::logical_plan::builder::LogicalTableSource;
+    use crate::{
+        LogicalPlanBuilder, and, col, in_subquery, lit, not, or,
+        out_ref_col_with_metadata, when,
+    };
 
+    use arrow::datatypes::Schema;
     use datafusion_common::{DFSchema, assert_or_internal_err};
 
     macro_rules! test_is_expr_nullable {
@@ -1192,6 +1202,64 @@ mod tests {
         }
     }
 
+    /// A scan of `t`, whose single column `a` has the given nullability.
+    fn scan_t(a_nullable: bool) -> LogicalPlanBuilder {
+        let schema = Schema::new(vec![Field::new("a", DataType::Int32, a_nullable)]);
+        let source = Arc::new(LogicalTableSource::new(Arc::new(schema)));
+        LogicalPlanBuilder::scan("t", source, None).unwrap()
+    }
+
+    #[test]
+    fn in_subquery_nullability() {
+        // `x IN (SELECT a FROM t)` evaluates to NULL when `x` is NULL, and when `x`
+        // matches no row while `a` contains a NULL. So it is nullable exactly when
+        // either the compared expression or the subquery's output column is.
+        let cases = [
+            (false, false, false),
+            (false, true, true),
+            (true, false, true),
+            (true, true, true),
+        ];
+
+        for (x_nullable, a_nullable, expected) in cases {
+            let subquery = scan_t(a_nullable)
+                .project(vec![col("a")])
+                .unwrap()
+                .build()
+                .unwrap();
+            let expr = in_subquery(col("x"), Arc::new(subquery));
+            let schema = MockExprSchema::new().with_nullable(x_nullable);
+
+            assert_eq!(expr.nullable(&schema).unwrap(), expected);
+        }
+    }
+
+    #[test]
+    fn in_subquery_nullability_uses_subquery_output_schema() {
+        // `DISTINCT` carries no expressions of its own, but its output column is still
+        // nullable, so the `IN` expression must be nullable too.
+        let subquery = scan_t(true)
+            .project(vec![col("a")])
+            .unwrap()
+            .distinct()
+            .unwrap()
+            .build()
+            .unwrap();
+        let expr = in_subquery(col("x"), Arc::new(subquery));
+        assert!(expr.nullable(&MockExprSchema::new()).unwrap());
+
+        // A computed projection's expressions reference `t.a`, which does not appear in
+        // the subquery's output schema, so nullability must be read off that schema's
+        // single column rather than by resolving the projection's expressions against it.
+        let subquery = scan_t(false)
+            .project(vec![col("a") + lit(1)])
+            .unwrap()
+            .build()
+            .unwrap();
+        let expr = in_subquery(col("x"), Arc::new(subquery));
+        assert!(!expr.nullable(&MockExprSchema::new()).unwrap());
+    }
+
     #[test]
     fn test_scalar_variable() {
         let mut meta = HashMap::new();

datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs

@@ -2541,6 +2541,36 @@ mod tests {
         assert_eq!(simplify(expr_b), expected_b);
     }
 
+    /// `c3_non_null IN (SELECT a FROM t)`, where `a` has the given nullability.
+    fn in_subquery_expr(a_nullable: bool) -> Expr {
+        let schema = Schema::new(vec![Field::new("a", DataType::Int64, a_nullable)]);
+        let source = Arc::new(LogicalTableSource::new(Arc::new(schema)));
+        let subquery = LogicalPlanBuilder::scan("t", source, None)
+            .unwrap()
+            .project(vec![col("a")])
+            .unwrap()
+            .build()
+            .unwrap();
+
+        in_subquery(col("c3_non_null"), Arc::new(subquery))
+    }
+
+    #[test]
+    fn test_simplify_eq_not_self_in_subquery() {
+        // `expr_a`: even though `c3_non_null` is non-nullable, the `IN` evaluates to NULL
+        // when `c3_non_null` matches no row and the subquery's `a` contains a NULL. So the
+        // expression is nullable and `A = A` must not fold to `true`.
+        let expr_a = in_subquery_expr(true);
+        let expected_a = expr_a.clone().is_not_null().or(lit_bool_null());
+
+        // `expr_b`: neither side can be NULL, so the `IN` is non-nullable and `A = A` is true.
+        let expr_b = in_subquery_expr(false);
+        let expected_b = lit(true);
+
+        assert_eq!(simplify(expr_a.clone().eq(expr_a)), expected_a);
+        assert_eq!(simplify(expr_b.clone().eq(expr_b)), expected_b);
+    }
+
     #[test]
     fn test_simplify_or_true() {
         let expr_a = col("c2").or(lit(true));