From 369865035a1574d7df62c63230c0fbc2068fd6b4 Mon Sep 17 00:00:00 2001 From: HairstonE Date: Tue, 19 May 2026 11:33:43 -0400 Subject: [PATCH 1/2] Infer placeholder type from subquery, unit tests, and .slt --- datafusion/expr/src/expr.rs | 113 ++++++++++++++++++ .../sqllogictest/test_files/prepare.slt | 32 +++++ 2 files changed, 145 insertions(+) diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index 36ef6cf1f5ba9..0dc58b34a7117 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -2143,6 +2143,17 @@ impl Expr { rewrite_placeholder(item, expr.as_ref(), schema)?; } } + Expr::InSubquery(InSubquery { + expr, + subquery, + negated: _, + }) => { + let subquery_schema = subquery.subquery.schema(); + let column = Expr::Column(Column::new_unqualified( + subquery_schema.fields()[0].name().clone(), + )); + rewrite_placeholder(expr.as_mut(), &column, subquery_schema)?; + } Expr::Like(Like { expr, pattern, .. }) | Expr::SimilarTo(Like { expr, pattern, .. }) => { rewrite_placeholder(pattern.as_mut(), expr.as_ref(), schema)?; @@ -3764,6 +3775,108 @@ mod test { } } + #[test] + fn infer_placeholder_in_subquery() { + // WHERE $1 IN (SELECT a FROM t) + let subquery_field = Field::new("a", DataType::Int32, false); + let subquery_schema = Arc::new( + DFSchema::from_unqualified_fields( + vec![subquery_field].into(), + Default::default(), + ) + .unwrap(), + ); + let subquery = Subquery { + subquery: Arc::new(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: subquery_schema, + })), + outer_ref_columns: vec![], + spans: Spans::new(), + }; + + let in_subquery = Expr::InSubquery(InSubquery { + expr: Box::new(Expr::Placeholder(Placeholder { + id: "$1".to_string(), + field: None, + })), + subquery, + negated: false, + }); + + let outer_schema = DFSchema::empty(); + let (inferred_expr, contains_placeholder) = + in_subquery.infer_placeholder_types(&outer_schema).unwrap(); + + assert!(contains_placeholder); + + match inferred_expr { + Expr::InSubquery(in_subquery) => match *in_subquery.expr { + Expr::Placeholder(placeholder) => { + let inferred = placeholder.field.expect("placeholder field"); + assert_eq!(inferred.data_type(), &DataType::Int32); + assert!(inferred.is_nullable()); + } + _ => panic!("Expected Placeholder expression in InSubquery"), + }, + _ => panic!("Expected InSubquery expression"), + } + } + + #[test] + fn infer_placeholder_not_in_subquery() { + // WHERE $1 NOT IN (SELECT a FROM t) + let subquery_field = Field::new("a", DataType::Int32, false); + let subquery_schema = Arc::new( + DFSchema::from_unqualified_fields( + vec![subquery_field].into(), + Default::default(), + ) + .unwrap(), + ); + let subquery = Subquery { + subquery: Arc::new(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: subquery_schema, + })), + outer_ref_columns: vec![], + spans: Spans::new(), + }; + + let not_in_subquery = Expr::InSubquery(InSubquery { + expr: Box::new(Expr::Placeholder(Placeholder { + id: "$1".to_string(), + field: None, + })), + subquery, + negated: true, + }); + + let outer_schema = DFSchema::empty(); + let (inferred_expr, contains_placeholder) = not_in_subquery + .infer_placeholder_types(&outer_schema) + .unwrap(); + + assert!(contains_placeholder); + + match inferred_expr { + Expr::InSubquery(in_subquery) => { + assert!(in_subquery.negated, "negated flag must be preserved"); + match *in_subquery.expr { + Expr::Placeholder(placeholder) => { + let inferred = placeholder.field.expect("placeholder field"); + assert_eq!(inferred.data_type(), &DataType::Int32); + assert!(inferred.is_nullable()); + } + _ => { + panic!("Expected Placeholder expression in InSubquery") + } + } + } + _ => panic!("Expected InSubquery expression"), + } + } + #[test] fn infer_placeholder_like_and_similar_to() { // name LIKE $1 diff --git a/datafusion/sqllogictest/test_files/prepare.slt b/datafusion/sqllogictest/test_files/prepare.slt index 8e8b1cd8e6ad0..f3225d4acc85d 100644 --- a/datafusion/sqllogictest/test_files/prepare.slt +++ b/datafusion/sqllogictest/test_files/prepare.slt @@ -88,6 +88,38 @@ EXECUTE my_plan('j%'); statement ok DEALLOCATE my_plan +# Allow prepare $1 IN (subquery) +statement ok +PREPARE my_plan AS SELECT id FROM person WHERE $1 IN (SELECT age FROM person); + +query I rowsort +EXECUTE my_plan(20); +---- +1 + +query I rowsort +EXECUTE my_plan(99); +---- + +statement ok +DEALLOCATE my_plan + +# Allow prepare $1 NOT IN (subquery) +statement ok +PREPARE my_plan AS SELECT id FROM person WHERE $1 NOT IN (SELECT age FROM person); + +query I rowsort +EXECUTE my_plan(99); +---- +1 + +query I rowsort +EXECUTE my_plan(20); +---- + +statement ok +DEALLOCATE my_plan + # Check for missing parameters statement ok PREPARE my_plan AS SELECT * FROM person WHERE id < $1; From 02cfd75c3657556e4a2a7b00a9dd2142385c9f58 Mon Sep 17 00:00:00 2001 From: HairstonE Date: Thu, 21 May 2026 17:22:27 -0400 Subject: [PATCH 2/2] Guarantee single-column subquery for InSubquery placeholder --- datafusion/expr/src/expr.rs | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index 0dc58b34a7117..08ec76c7689c3 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -2149,10 +2149,25 @@ impl Expr { negated: _, }) => { let subquery_schema = subquery.subquery.schema(); - let column = Expr::Column(Column::new_unqualified( - subquery_schema.fields()[0].name().clone(), - )); - rewrite_placeholder(expr.as_mut(), &column, subquery_schema)?; + match &subquery_schema.fields()[..] { + [subquery_field] => { + let column = Expr::Column(Column::new_unqualified( + subquery_field.name().clone(), + )); + rewrite_placeholder( + expr.as_mut(), + &column, + subquery_schema, + )?; + } + _ => { + return plan_err!( + "InSubquery should only return one column, but found {}: {}", + subquery_schema.fields().len(), + subquery_schema.field_names().join(", ") + ); + } + } } Expr::Like(Like { expr, pattern, .. }) | Expr::SimilarTo(Like { expr, pattern, .. }) => {