Skip to content

Commit eb45093

Browse files
HairstonEkosiew
HairstonE
and
kosiew
authored
Infer placeholder type from ANY/ALL subquery, unit tests (#22545)
## Which issue does this PR close? Closes #22475. ## Rationale for this change `$1 = ANY (SELECT ...)` and `$1 <> ALL (SELECT ...)` left the placeholder untyped because `infer_placeholder_types` had no arm for `SetComparison`. ## What changes are included in this PR? Adds the `SetComparison` arm to `infer_placeholder_types`, reading the type from the subquery's projected column. Covers all quantifiers (`ANY`, `ALL`) and comparison operators. ## How are these changes tested? Unit tests for `ANY` and `ALL` placeholder inference, plus end-to-end sqllogictests with `PREPARE`/`EXECUTE`. ## Are there any user-facing changes? No. --------- Co-authored-by: kosiew <kosiew@gmail.com>
1 parent 754e113 commit eb45093

2 files changed

Lines changed: 191 additions & 20 deletions

File tree

‎datafusion/expr/src/expr.rs‎

Lines changed: 143 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -2191,26 +2191,23 @@ impl Expr {
21912191
subquery,
21922192
negated: _,
21932193
}) => {
2194-
let subquery_schema = subquery.subquery.schema();
2195-
match &subquery_schema.fields()[..] {
2196-
[subquery_field] => {
2197-
let column = Expr::Column(Column::new_unqualified(
2198-
subquery_field.name().clone(),
2199-
));
2200-
rewrite_placeholder(
2201-
expr.as_mut(),
2202-
&column,
2203-
subquery_schema,
2204-
)?;
2205-
}
2206-
_ => {
2207-
return plan_err!(
2208-
"InSubquery should only return one column, but found {}: {}",
2209-
subquery_schema.fields().len(),
2210-
subquery_schema.field_names().join(", ")
2211-
);
2212-
}
2213-
}
2194+
rewrite_placeholder_from_subquery(
2195+
"InSubquery",
2196+
expr.as_mut(),
2197+
subquery,
2198+
)?;
2199+
}
2200+
Expr::SetComparison(SetComparison {
2201+
expr,
2202+
subquery,
2203+
op: _,
2204+
quantifier: _,
2205+
}) => {
2206+
rewrite_placeholder_from_subquery(
2207+
"SetComparison",
2208+
expr.as_mut(),
2209+
subquery,
2210+
)?;
22142211
}
22152212
Expr::Like(Like { expr, pattern, .. })
22162213
| Expr::SimilarTo(Like { expr, pattern, .. }) => {
@@ -2938,6 +2935,26 @@ macro_rules! expr_vec_fmt {
29382935
.join(", ")
29392936
}};
29402937
}
2938+
/// Infer an untyped placeholder on the left of a single-column subquery predicate from the subquery projection
2939+
fn rewrite_placeholder_from_subquery(
2940+
kind: &str,
2941+
expr: &mut Expr,
2942+
subquery: &Subquery,
2943+
) -> Result<()> {
2944+
let subquery_schema = subquery.subquery.schema();
2945+
match &subquery_schema.fields()[..] {
2946+
[subquery_field] => {
2947+
let column =
2948+
Expr::Column(Column::new_unqualified(subquery_field.name().clone()));
2949+
rewrite_placeholder(expr, &column, subquery_schema)
2950+
}
2951+
_ => plan_err!(
2952+
"{kind} should only return one column, but found {}: {}",
2953+
subquery_schema.fields().len(),
2954+
subquery_schema.field_names().join(", ")
2955+
),
2956+
}
2957+
}
29412958

29422959
struct SchemaDisplay<'a>(&'a Expr);
29432960
impl Display for SchemaDisplay<'_> {
@@ -3946,6 +3963,112 @@ mod test {
39463963
}
39473964
}
39483965

3966+
#[test]
3967+
fn infer_placeholder_set_comparison_any() {
3968+
// WHERE $1 = ANY (SELECT a FROM t) -- parallel to infer_placeholder_in_subquery
3969+
let subquery_field = Field::new("a", DataType::Int32, false);
3970+
let subquery_schema = Arc::new(
3971+
DFSchema::from_unqualified_fields(
3972+
vec![subquery_field].into(),
3973+
Default::default(),
3974+
)
3975+
.unwrap(),
3976+
);
3977+
let subquery = Subquery {
3978+
subquery: Arc::new(LogicalPlan::EmptyRelation(EmptyRelation {
3979+
produce_one_row: false,
3980+
schema: subquery_schema,
3981+
})),
3982+
outer_ref_columns: vec![],
3983+
spans: Spans::new(),
3984+
};
3985+
3986+
let set_cmp = Expr::SetComparison(SetComparison {
3987+
expr: Box::new(Expr::Placeholder(Placeholder {
3988+
id: "$1".to_string(),
3989+
field: None,
3990+
})),
3991+
subquery,
3992+
op: Operator::Eq,
3993+
quantifier: SetQuantifier::Any,
3994+
});
3995+
3996+
let outer_schema = DFSchema::empty();
3997+
let (inferred_expr, contains_placeholder) =
3998+
set_cmp.infer_placeholder_types(&outer_schema).unwrap();
3999+
4000+
assert!(contains_placeholder);
4001+
4002+
match inferred_expr {
4003+
Expr::SetComparison(sc) => {
4004+
assert_eq!(sc.quantifier, SetQuantifier::Any);
4005+
match *sc.expr {
4006+
Expr::Placeholder(p) => {
4007+
let inferred =
4008+
p.field.expect("placeholder field should be Int32");
4009+
assert_eq!(inferred.data_type(), &DataType::Int32);
4010+
assert!(inferred.is_nullable());
4011+
}
4012+
_ => panic!("Expected Placeholder expression in SetComparison"),
4013+
}
4014+
}
4015+
_ => panic!("Expected SetComparison expression"),
4016+
}
4017+
}
4018+
4019+
#[test]
4020+
fn infer_placeholder_set_comparison_all() {
4021+
// WHERE $1 <> ALL (SELECT a FROM t)
4022+
let subquery_field = Field::new("a", DataType::Int32, false);
4023+
let subquery_schema = Arc::new(
4024+
DFSchema::from_unqualified_fields(
4025+
vec![subquery_field].into(),
4026+
Default::default(),
4027+
)
4028+
.unwrap(),
4029+
);
4030+
let subquery = Subquery {
4031+
subquery: Arc::new(LogicalPlan::EmptyRelation(EmptyRelation {
4032+
produce_one_row: false,
4033+
schema: subquery_schema,
4034+
})),
4035+
outer_ref_columns: vec![],
4036+
spans: Spans::new(),
4037+
};
4038+
4039+
let set_cmp = Expr::SetComparison(SetComparison {
4040+
expr: Box::new(Expr::Placeholder(Placeholder {
4041+
id: "$1".to_string(),
4042+
field: None,
4043+
})),
4044+
subquery,
4045+
op: Operator::NotEq,
4046+
quantifier: SetQuantifier::All,
4047+
});
4048+
4049+
let outer_schema = DFSchema::empty();
4050+
let (inferred_expr, contains_placeholder) =
4051+
set_cmp.infer_placeholder_types(&outer_schema).unwrap();
4052+
4053+
assert!(contains_placeholder);
4054+
4055+
match inferred_expr {
4056+
Expr::SetComparison(sc) => {
4057+
assert_eq!(sc.quantifier, SetQuantifier::All);
4058+
match *sc.expr {
4059+
Expr::Placeholder(p) => {
4060+
let inferred =
4061+
p.field.expect("placeholder field should be Int32");
4062+
assert_eq!(inferred.data_type(), &DataType::Int32);
4063+
assert!(inferred.is_nullable());
4064+
}
4065+
_ => panic!("Expected Placeholder expression in SetComparison"),
4066+
}
4067+
}
4068+
_ => panic!("Expected SetComparison expression"),
4069+
}
4070+
}
4071+
39494072
#[test]
39504073
fn infer_placeholder_like_and_similar_to() {
39514074
// name LIKE $1

‎datafusion/sqllogictest/test_files/prepare.slt‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,54 @@ EXECUTE my_plan(20);
139139
statement ok
140140
DEALLOCATE my_plan
141141

142+
# Allow prepare $1 = ANY (subquery)
143+
statement ok
144+
PREPARE my_plan AS SELECT id FROM person WHERE $1 = ANY (SELECT age FROM person);
145+
146+
query I rowsort
147+
EXECUTE my_plan(20);
148+
----
149+
1
150+
151+
query I rowsort
152+
EXECUTE my_plan(99);
153+
----
154+
155+
statement ok
156+
DEALLOCATE my_plan
157+
158+
# Allow prepare $1 <> ALL (subquery)
159+
statement ok
160+
PREPARE my_plan AS SELECT id FROM person WHERE $1 <> ALL (SELECT age FROM person);
161+
162+
query I rowsort
163+
EXECUTE my_plan(99);
164+
----
165+
1
166+
167+
query I rowsort
168+
EXECUTE my_plan(20);
169+
----
170+
171+
statement ok
172+
DEALLOCATE my_plan
173+
174+
# Allow prepare $1 < ALL (subquery)
175+
statement ok
176+
PREPARE my_plan AS SELECT id FROM person WHERE $1 < ALL (SELECT age FROM person);
177+
178+
query I rowsort
179+
EXECUTE my_plan(10);
180+
----
181+
1
182+
183+
query I rowsort
184+
EXECUTE my_plan(50);
185+
----
186+
187+
statement ok
188+
DEALLOCATE my_plan
189+
142190
# Check for missing parameters
143191
statement ok
144192
PREPARE my_plan AS SELECT * FROM person WHERE id < $1;

0 commit comments

Comments
 (0)