Skip to content

Commit 44027bc

Browse files
authored
Fix: Infer placeholder type from subquery (#22436)
## Which issue does this PR close? Closes #15979. ## Rationale for this change `$1 IN (SELECT ...)` left the placeholder untyped because `infer_placeholder_types` had no arm for `InSubquery` and made `get_parameter_types()` return `None` for these placeholders. ## What changes are included in this PR? Adds the `InSubquery` arm to `infer_placeholder_types`, reading the type from the subquery's projected column. Covers both `IN` and `NOT IN`. ## How are these changes tested? Unit tests for `IN` and `NOT IN` placeholder inference, plus end-to-end sqllogictests with `PREPARE`/`EXECUTE`. ## Are there any user-facing changes? No.
1 parent 50d74a7 commit 44027bc

2 files changed

Lines changed: 145 additions & 0 deletions

File tree

datafusion/expr/src/expr.rs

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2186,6 +2186,17 @@ impl Expr {
21862186
rewrite_placeholder(item, expr.as_ref(), schema)?;
21872187
}
21882188
}
2189+
Expr::InSubquery(InSubquery {
2190+
expr,
2191+
subquery,
2192+
negated: _,
2193+
}) => {
2194+
let subquery_schema = subquery.subquery.schema();
2195+
let column = Expr::Column(Column::new_unqualified(
2196+
subquery_schema.fields()[0].name().clone(),
2197+
));
2198+
rewrite_placeholder(expr.as_mut(), &column, subquery_schema)?;
2199+
}
21892200
Expr::Like(Like { expr, pattern, .. })
21902201
| Expr::SimilarTo(Like { expr, pattern, .. }) => {
21912202
rewrite_placeholder(pattern.as_mut(), expr.as_ref(), schema)?;
@@ -3817,6 +3828,108 @@ mod test {
38173828
}
38183829
}
38193830

3831+
#[test]
3832+
fn infer_placeholder_in_subquery() {
3833+
// WHERE $1 IN (SELECT a FROM t)
3834+
let subquery_field = Field::new("a", DataType::Int32, false);
3835+
let subquery_schema = Arc::new(
3836+
DFSchema::from_unqualified_fields(
3837+
vec![subquery_field].into(),
3838+
Default::default(),
3839+
)
3840+
.unwrap(),
3841+
);
3842+
let subquery = Subquery {
3843+
subquery: Arc::new(LogicalPlan::EmptyRelation(EmptyRelation {
3844+
produce_one_row: false,
3845+
schema: subquery_schema,
3846+
})),
3847+
outer_ref_columns: vec![],
3848+
spans: Spans::new(),
3849+
};
3850+
3851+
let in_subquery = Expr::InSubquery(InSubquery {
3852+
expr: Box::new(Expr::Placeholder(Placeholder {
3853+
id: "$1".to_string(),
3854+
field: None,
3855+
})),
3856+
subquery,
3857+
negated: false,
3858+
});
3859+
3860+
let outer_schema = DFSchema::empty();
3861+
let (inferred_expr, contains_placeholder) =
3862+
in_subquery.infer_placeholder_types(&outer_schema).unwrap();
3863+
3864+
assert!(contains_placeholder);
3865+
3866+
match inferred_expr {
3867+
Expr::InSubquery(in_subquery) => match *in_subquery.expr {
3868+
Expr::Placeholder(placeholder) => {
3869+
let inferred = placeholder.field.expect("placeholder field");
3870+
assert_eq!(inferred.data_type(), &DataType::Int32);
3871+
assert!(inferred.is_nullable());
3872+
}
3873+
_ => panic!("Expected Placeholder expression in InSubquery"),
3874+
},
3875+
_ => panic!("Expected InSubquery expression"),
3876+
}
3877+
}
3878+
3879+
#[test]
3880+
fn infer_placeholder_not_in_subquery() {
3881+
// WHERE $1 NOT IN (SELECT a FROM t)
3882+
let subquery_field = Field::new("a", DataType::Int32, false);
3883+
let subquery_schema = Arc::new(
3884+
DFSchema::from_unqualified_fields(
3885+
vec![subquery_field].into(),
3886+
Default::default(),
3887+
)
3888+
.unwrap(),
3889+
);
3890+
let subquery = Subquery {
3891+
subquery: Arc::new(LogicalPlan::EmptyRelation(EmptyRelation {
3892+
produce_one_row: false,
3893+
schema: subquery_schema,
3894+
})),
3895+
outer_ref_columns: vec![],
3896+
spans: Spans::new(),
3897+
};
3898+
3899+
let not_in_subquery = Expr::InSubquery(InSubquery {
3900+
expr: Box::new(Expr::Placeholder(Placeholder {
3901+
id: "$1".to_string(),
3902+
field: None,
3903+
})),
3904+
subquery,
3905+
negated: true,
3906+
});
3907+
3908+
let outer_schema = DFSchema::empty();
3909+
let (inferred_expr, contains_placeholder) = not_in_subquery
3910+
.infer_placeholder_types(&outer_schema)
3911+
.unwrap();
3912+
3913+
assert!(contains_placeholder);
3914+
3915+
match inferred_expr {
3916+
Expr::InSubquery(in_subquery) => {
3917+
assert!(in_subquery.negated, "negated flag must be preserved");
3918+
match *in_subquery.expr {
3919+
Expr::Placeholder(placeholder) => {
3920+
let inferred = placeholder.field.expect("placeholder field");
3921+
assert_eq!(inferred.data_type(), &DataType::Int32);
3922+
assert!(inferred.is_nullable());
3923+
}
3924+
_ => {
3925+
panic!("Expected Placeholder expression in InSubquery")
3926+
}
3927+
}
3928+
}
3929+
_ => panic!("Expected InSubquery expression"),
3930+
}
3931+
}
3932+
38203933
#[test]
38213934
fn infer_placeholder_like_and_similar_to() {
38223935
// name LIKE $1

datafusion/sqllogictest/test_files/prepare.slt

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,38 @@ EXECUTE my_plan('j%');
107107
statement ok
108108
DEALLOCATE my_plan
109109

110+
# Allow prepare $1 IN (subquery)
111+
statement ok
112+
PREPARE my_plan AS SELECT id FROM person WHERE $1 IN (SELECT age FROM person);
113+
114+
query I rowsort
115+
EXECUTE my_plan(20);
116+
----
117+
1
118+
119+
query I rowsort
120+
EXECUTE my_plan(99);
121+
----
122+
123+
statement ok
124+
DEALLOCATE my_plan
125+
126+
# Allow prepare $1 NOT IN (subquery)
127+
statement ok
128+
PREPARE my_plan AS SELECT id FROM person WHERE $1 NOT IN (SELECT age FROM person);
129+
130+
query I rowsort
131+
EXECUTE my_plan(99);
132+
----
133+
1
134+
135+
query I rowsort
136+
EXECUTE my_plan(20);
137+
----
138+
139+
statement ok
140+
DEALLOCATE my_plan
141+
110142
# Check for missing parameters
111143
statement ok
112144
PREPARE my_plan AS SELECT * FROM person WHERE id < $1;

0 commit comments

Comments
 (0)