diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 039bbad65a660..4d35c07eab661 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/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(); diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs index 39c8541b51b2f..b5d781722c0d0 100644 --- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs +++ b/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));