Skip to content

Commit 0476efd

Browse files
sundy-lisundy-li
authored
feat(query): support scalar subquery arguments in table functions (#19213)
* feat(query): support scalar subquery arguments in table functions Enable table functions like generate_series and range to accept scalar subqueries as arguments. For example: SELECT * FROM generate_series( (SELECT count() FROM numbers(10))::int, (SELECT count() FROM numbers(39))::int ); The implementation modifies bind_table_args to evaluate scalar subqueries by executing them using the subquery executor when constant folding fails. * fix: address review comments - Rename bind_table_args_with_subquery_executor to bind_table_args - Return Scalar::Null for empty subquery results instead of error - Update bind_obfuscate.rs to use new signature * update * fix: correct variable name block -> blocks and handle Result type
1 parent d7a05f0 commit 0476efd

File tree

4 files changed

+182
-28
lines changed

4 files changed

+182
-28
lines changed

‎src/query/sql/src/planner/binder/bind_table_reference/bind_obfuscate.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ impl Binder {
5959
self.metadata.clone(),
6060
&[],
6161
);
62-
let mut named_args = bind_table_args(&mut scalar_binder, &[], named_params)?.named;
62+
let mut named_args = bind_table_args(&mut scalar_binder, &[], named_params, &None)?.named;
6363
let seed = match named_args.remove("seed") {
6464
Some(v) => u64_value(&v).ok_or(ErrorCode::BadArguments("invalid seed"))?,
6565
None => {

‎src/query/sql/src/planner/binder/bind_table_reference/bind_table_function.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,12 @@ impl Binder {
136136
self.metadata.clone(),
137137
&[],
138138
);
139-
let table_args = bind_table_args(&mut scalar_binder, params, named_params)?;
139+
let table_args = bind_table_args(
140+
&mut scalar_binder,
141+
params,
142+
named_params,
143+
&self.subquery_executor,
144+
)?;
140145

141146
let tenant = self.ctx.get_tenant();
142147
let udtf_result = databend_common_base::runtime::block_on(async {

‎src/query/sql/src/planner/binder/table_args.rs‎

Lines changed: 122 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
// limitations under the License.
1414

1515
use std::collections::HashMap;
16+
use std::sync::Arc;
1617

1718
use databend_common_ast::ast::Expr;
1819
use databend_common_ast::ast::Identifier;
@@ -21,59 +22,156 @@ use databend_common_exception::ErrorCode;
2122
use databend_common_exception::Result;
2223
use databend_common_expression::Constant;
2324
use databend_common_expression::ConstantFolder;
25+
use databend_common_expression::DataBlock;
2426
use databend_common_expression::Scalar;
2527
use databend_common_functions::BUILTIN_FUNCTIONS;
2628

2729
use crate::ScalarBinder;
2830
use crate::ScalarExpr;
31+
use crate::planner::QueryExecutor;
2932
use crate::plans::ConstantExpr;
3033

34+
/// Check if an AST expression contains a subquery
35+
fn contains_subquery(expr: &Expr) -> bool {
36+
match expr {
37+
Expr::Subquery { .. } => true,
38+
Expr::InSubquery { .. } => true,
39+
Expr::Exists { .. } => true,
40+
Expr::Cast { expr, .. } => contains_subquery(expr),
41+
Expr::TryCast { expr, .. } => contains_subquery(expr),
42+
Expr::FunctionCall { func, .. } => func.args.iter().any(contains_subquery),
43+
Expr::BinaryOp { left, right, .. } => contains_subquery(left) || contains_subquery(right),
44+
Expr::UnaryOp { expr, .. } => contains_subquery(expr),
45+
Expr::IsNull { expr, .. } => contains_subquery(expr),
46+
Expr::IsDistinctFrom { left, right, .. } => {
47+
contains_subquery(left) || contains_subquery(right)
48+
}
49+
Expr::InList { expr, list, .. } => {
50+
contains_subquery(expr) || list.iter().any(contains_subquery)
51+
}
52+
Expr::Between {
53+
expr, low, high, ..
54+
} => contains_subquery(expr) || contains_subquery(low) || contains_subquery(high),
55+
Expr::Case {
56+
operand,
57+
conditions,
58+
results,
59+
else_result,
60+
..
61+
} => {
62+
operand.as_ref().is_some_and(|e| contains_subquery(e))
63+
|| conditions.iter().any(contains_subquery)
64+
|| results.iter().any(contains_subquery)
65+
|| else_result.as_ref().is_some_and(|e| contains_subquery(e))
66+
}
67+
Expr::MapAccess { expr, .. } => contains_subquery(expr),
68+
Expr::Array { exprs, .. } => exprs.iter().any(contains_subquery),
69+
Expr::Tuple { exprs, .. } => exprs.iter().any(contains_subquery),
70+
_ => false,
71+
}
72+
}
73+
74+
/// Execute a subquery and extract a single scalar value from the result
75+
fn execute_subquery_for_scalar(
76+
subquery_executor: &Arc<dyn QueryExecutor>,
77+
expr: &Expr,
78+
) -> Result<Scalar> {
79+
let sql = expr.to_string();
80+
let data_blocks = databend_common_base::runtime::block_on(async {
81+
subquery_executor
82+
.execute_query_with_sql_string(&format!("SELECT {}", sql))
83+
.await
84+
})?;
85+
86+
if data_blocks.is_empty() {
87+
return Ok(Scalar::Null);
88+
}
89+
90+
let block = DataBlock::concat(&data_blocks)?;
91+
92+
if block.num_columns() != 1 {
93+
return Err(ErrorCode::SemanticError(
94+
"Scalar subquery in table function argument must return exactly one column".to_string(),
95+
));
96+
}
97+
98+
if block.num_rows() > 1 {
99+
return Err(ErrorCode::SemanticError(
100+
"Scalar subquery in table function argument returned more than one row".to_string(),
101+
));
102+
}
103+
104+
if block.num_rows() == 0 {
105+
return Ok(Scalar::Null);
106+
}
107+
108+
let column = &block.columns()[0];
109+
let col_value = column.value();
110+
Ok(col_value.index(0).unwrap().to_owned())
111+
}
112+
113+
/// Try to fold an expression to a constant scalar value.
114+
/// If constant folding fails and the expression contains a subquery,
115+
/// execute the subquery and return the result.
116+
fn try_fold_to_scalar(
117+
scalar: ScalarExpr,
118+
ast_expr: &Expr,
119+
scalar_binder: &ScalarBinder<'_>,
120+
subquery_executor: &Option<Arc<dyn QueryExecutor>>,
121+
) -> Result<Scalar> {
122+
let expr = scalar.as_expr()?;
123+
let (expr, _) = ConstantFolder::fold(&expr, &scalar_binder.get_func_ctx()?, &BUILTIN_FUNCTIONS);
124+
125+
match expr.into_constant() {
126+
Ok(Constant { scalar, .. }) => Ok(scalar),
127+
Err(_) => {
128+
if contains_subquery(ast_expr) {
129+
if let Some(executor) = subquery_executor {
130+
return execute_subquery_for_scalar(executor, ast_expr);
131+
}
132+
return Err(ErrorCode::Internal(
133+
"Subquery executor is not available for evaluating table function arguments"
134+
.to_string(),
135+
));
136+
}
137+
Err(ErrorCode::Unimplemented(format!(
138+
"Unsupported table argument type: {:?}",
139+
scalar
140+
)))
141+
}
142+
}
143+
}
144+
31145
pub fn bind_table_args(
32146
scalar_binder: &mut ScalarBinder<'_>,
33147
params: &[Expr],
34148
named_params: &[(Identifier, Expr)],
149+
subquery_executor: &Option<Arc<dyn QueryExecutor>>,
35150
) -> Result<TableArgs> {
36151
let mut args = Vec::with_capacity(params.len());
37152
for arg in params.iter() {
38-
args.push(scalar_binder.bind(arg)?.0);
153+
args.push((scalar_binder.bind(arg)?.0, arg));
39154
}
40155

41156
let mut named_args = Vec::with_capacity(named_params.len());
42157
for (name, arg) in named_params.iter() {
43-
named_args.push((name.clone(), scalar_binder.bind(arg)?.0));
158+
named_args.push((name.clone(), scalar_binder.bind(arg)?.0, arg));
44159
}
45160

46161
let positioned_args = args
47162
.into_iter()
48-
.map(|scalar| {
49-
let expr = scalar.as_expr()?;
50-
let (expr, _) =
51-
ConstantFolder::fold(&expr, &scalar_binder.get_func_ctx()?, &BUILTIN_FUNCTIONS);
52-
match expr.into_constant() {
53-
Ok(Constant { scalar, .. }) => Ok(scalar),
54-
Err(_) => Err(ErrorCode::Unimplemented(format!(
55-
"Unsupported table argument type: {:?}",
56-
scalar
57-
))),
58-
}
163+
.map(|(scalar, ast_expr)| {
164+
try_fold_to_scalar(scalar, ast_expr, scalar_binder, subquery_executor)
59165
})
60166
.collect::<Result<Vec<_>>>()?;
61167

62168
let named_args: HashMap<String, Scalar> = named_args
63169
.into_iter()
64-
.map(|(name, scalar)| match scalar {
170+
.map(|(name, scalar, ast_expr)| match scalar {
65171
ScalarExpr::ConstantExpr(ConstantExpr { value, .. }) => Ok((name.name.clone(), value)),
66172
_ => {
67-
let expr = scalar.as_expr()?;
68-
let (expr, _) =
69-
ConstantFolder::fold(&expr, &scalar_binder.get_func_ctx()?, &BUILTIN_FUNCTIONS);
70-
match expr.into_constant() {
71-
Ok(Constant { scalar, .. }) => Ok((name.name.clone(), scalar)),
72-
Err(_) => Err(ErrorCode::Unimplemented(format!(
73-
"Unsupported table argument type: {:?}",
74-
scalar
75-
))),
76-
}
173+
let value = try_fold_to_scalar(scalar, ast_expr, scalar_binder, subquery_executor)?;
174+
Ok((name.name.clone(), value))
77175
}
78176
})
79177
.collect::<Result<HashMap<_, _>>>()?;

‎tests/sqllogictests/suites/query/functions/02_0063_function_generate_series.test‎

Lines changed: 53 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,13 +35,64 @@ select * from generate_series('2021-03-26'::date,'2021-03-27'::date,1);
3535
2021-03-26
3636
2021-03-27
3737

38-
query T
38+
query I
3939
select * from range(1, 3)
4040
----
4141
1
4242
2
4343

44-
query T
44+
query I
4545
select max(`range`) from range(1, 10000)
4646
----
4747
9999
48+
49+
query I
50+
select generate_series as install_date from generate_series((select count() from numbers(10))::int, (select count() from numbers(39))::int);
51+
----
52+
10
53+
11
54+
12
55+
13
56+
14
57+
15
58+
16
59+
17
60+
18
61+
19
62+
20
63+
21
64+
22
65+
23
66+
24
67+
25
68+
26
69+
27
70+
28
71+
29
72+
30
73+
31
74+
32
75+
33
76+
34
77+
35
78+
36
79+
37
80+
38
81+
39
82+
83+
query I
84+
select * from generate_series((select 1), (select 5), 1);
85+
----
86+
1
87+
2
88+
3
89+
4
90+
5
91+
92+
query I
93+
select * from range((select 3)::int, (select 7)::int);
94+
----
95+
3
96+
4
97+
5
98+
6

0 commit comments

Comments
 (0)