Skip to content

Commit e8dad40

Browse files
youngsofunyoungsofun
committed
fix(query): refresh lambda types after planner rewrites
1 parent 9661185 commit e8dad40

14 files changed

Lines changed: 284 additions & 69 deletions

File tree

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -326,14 +326,16 @@ impl Binder {
326326
)
327327
})
328328
.collect::<Result<Vec<_>>>()?;
329-
Ok(ScalarExpr::LambdaFunction(LambdaFunc {
329+
let mut lambda_func = LambdaFunc {
330330
span: lambda_func.span,
331331
func_name: lambda_func.func_name.clone(),
332332
args,
333333
lambda_expr: lambda_func.lambda_expr.clone(),
334334
lambda_display: lambda_func.lambda_display.clone(),
335335
return_type: lambda_func.return_type.clone(),
336-
}))
336+
};
337+
lambda_func.refresh_return_type()?;
338+
Ok(ScalarExpr::LambdaFunction(lambda_func))
337339
}
338340
window @ ScalarExpr::WindowFunction(_) => {
339341
if !rewrite_flags.needs_window_rewrite {

‎src/query/sql/src/planner/optimizer/ir/expr/type_validator.rs‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ use crate::Symbol;
2828
use crate::plans::AggregateFunction;
2929
use crate::plans::BoundColumnRef;
3030
use crate::plans::FunctionCall;
31+
use crate::plans::LambdaFunc;
3132
use crate::plans::Operator;
3233
use crate::plans::RelOperator;
3334
use crate::plans::ScalarExpr;
@@ -268,6 +269,17 @@ impl ScalarTypeValidator<'_> {
268269
Ok(())
269270
}
270271

272+
fn validate_lambda_function(&mut self, function: &LambdaFunc) -> Result<()> {
273+
let inferred = function.infer_return_type()?;
274+
if inferred != *function.return_type {
275+
return Err(ErrorCode::Internal(format!(
276+
"SExpr lambda return type mismatch for {}: stored {:?}, inferred {inferred:?}",
277+
function.func_name, function.return_type
278+
)));
279+
}
280+
Ok(())
281+
}
282+
271283
fn validate_aggregate_function(&mut self, aggregate: &AggregateFunction) -> Result<()> {
272284
let argument_types = aggregate
273285
.args
@@ -320,4 +332,13 @@ impl ScalarExprVisitor<'_> for ScalarTypeValidator<'_> {
320332
}
321333
Ok(())
322334
}
335+
336+
fn visit_lambda_function(&mut self, function: &LambdaFunc) -> Result<()> {
337+
self.validate_lambda_function(function)?;
338+
339+
for argument in &function.args {
340+
self.visit(argument)?;
341+
}
342+
Ok(())
343+
}
323344
}

‎src/query/sql/src/planner/optimizer/optimizers/operator/decorrelate/flatten_scalar.rs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -108,14 +108,16 @@ impl SubqueryDecorrelatorOptimizer {
108108
.iter()
109109
.map(|arg| self.flatten_scalar(arg, correlated_columns, derived_columns))
110110
.collect::<Result<Vec<_>>>()?;
111-
Ok(ScalarExpr::LambdaFunction(LambdaFunc {
111+
let mut lambda = LambdaFunc {
112112
span: lambda.span,
113113
func_name: lambda.func_name.clone(),
114114
args,
115115
lambda_expr: lambda.lambda_expr.clone(),
116116
lambda_display: lambda.lambda_display.clone(),
117117
return_type: lambda.return_type.clone(),
118-
}))
118+
};
119+
lambda.refresh_return_type()?;
120+
Ok(ScalarExpr::LambdaFunction(lambda))
119121
}
120122
ScalarExpr::CastExpr(cast_expr) => {
121123
let scalar =

‎src/query/sql/src/planner/optimizer/optimizers/operator/filter/equivalent_constants_visitor.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ use crate::optimizer::optimizers::operator::filter::remove_trivial_type_cast;
2222
use crate::plans::BoundColumnRef;
2323
use crate::plans::ComparisonOp;
2424
use crate::plans::FunctionCall;
25+
use crate::plans::LambdaFunc;
2526
use crate::plans::VisitorMut;
2627
use crate::plans::walk_expr_mut;
2728

@@ -173,4 +174,14 @@ impl VisitorMut<'_> for EquivalentConstantsVisitorInner {
173174
}
174175
func.refresh_return_type()
175176
}
177+
178+
fn visit_lambda_function(&mut self, lambda: &mut LambdaFunc) -> Result<()> {
179+
for argument in &mut lambda.args {
180+
let mut visitor = EquivalentConstantsVisitorInner::default()
181+
.eq_constants(self.eq_constants.clone())
182+
.left_visit_order(self.left_visit_order);
183+
visitor.visit(argument)?;
184+
}
185+
lambda.refresh_return_type()
186+
}
176187
}

‎src/query/sql/src/planner/optimizer/optimizers/rule/agg_rules/rule_grouping_sets_to_union.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ use crate::plans::CastExpr;
3939
use crate::plans::ConstantExpr;
4040
use crate::plans::EvalScalar;
4141
use crate::plans::FunctionCall;
42+
use crate::plans::LambdaFunc;
4243
use crate::plans::MaterializedCTE;
4344
use crate::plans::MaterializedCTERef;
4445
use crate::plans::RelOp;
@@ -307,4 +308,11 @@ impl VisitorMut<'_> for ReplaceColumnForGroupingSetsVisitor {
307308
}
308309
function.refresh_return_type()
309310
}
311+
312+
fn visit_lambda_function(&mut self, lambda: &mut LambdaFunc) -> Result<()> {
313+
for argument in &mut lambda.args {
314+
self.visit(argument)?;
315+
}
316+
lambda.refresh_return_type()
317+
}
310318
}

‎src/query/sql/src/planner/optimizer/optimizers/rule/agg_rules/rule_hierarchical_grouping_sets.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ use crate::plans::CastExpr;
4444
use crate::plans::ConstantExpr;
4545
use crate::plans::EvalScalar;
4646
use crate::plans::FunctionCall;
47+
use crate::plans::LambdaFunc;
4748
use crate::plans::MaterializedCTE;
4849
use crate::plans::MaterializedCTERef;
4950
use crate::plans::RelOp;
@@ -1031,4 +1032,11 @@ impl VisitorMut<'_> for GroupingSetsNullVisitor {
10311032
}
10321033
function.refresh_return_type()
10331034
}
1035+
1036+
fn visit_lambda_function(&mut self, lambda: &mut LambdaFunc) -> Result<()> {
1037+
for argument in &mut lambda.args {
1038+
self.visit(argument)?;
1039+
}
1040+
lambda.refresh_return_type()
1041+
}
10341042
}

‎src/query/sql/src/planner/optimizer/optimizers/rule/agg_rules/view_rewrite.rs‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,21 @@ impl QueryInfo {
267267
return Err(ErrorCode::Internal("Can't found column from index"));
268268
}
269269
ScalarExpr::ConstantExpr(_) => scalar.clone(),
270+
ScalarExpr::LambdaFunction(lambda) => {
271+
let mut new_args = Vec::with_capacity(lambda.args.len());
272+
for arg in &lambda.args {
273+
let Some(new_arg) =
274+
self.check_output_cols(arg, index_output_cols, new_selection_set)?
275+
else {
276+
return Ok(None);
277+
};
278+
new_args.push(new_arg);
279+
}
280+
let mut new_lambda = lambda.clone();
281+
new_lambda.args = new_args;
282+
new_lambda.refresh_return_type()?;
283+
ScalarExpr::LambdaFunction(new_lambda)
284+
}
270285
ScalarExpr::FunctionCall(func) => {
271286
let mut valid = true;
272287
let mut new_args = Vec::with_capacity(func.arguments.len());

‎src/query/sql/src/planner/optimizer/optimizers/rule/filter_rules/rule_push_down_filter_eval_scalar.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ use crate::optimizer::optimizers::rule::TransformResult;
2525
use crate::plans::EvalScalar;
2626
use crate::plans::Filter;
2727
use crate::plans::FunctionCall;
28+
use crate::plans::LambdaFunc;
2829
use crate::plans::RelOp;
2930
use crate::plans::ScalarExpr;
3031
use crate::plans::ScalarItem;
@@ -80,6 +81,13 @@ impl RulePushDownFilterEvalScalar {
8081
}
8182
function.refresh_return_type()
8283
}
84+
85+
fn visit_lambda_function(&mut self, lambda: &mut LambdaFunc) -> Result<()> {
86+
for argument in &mut lambda.args {
87+
self.visit(argument)?;
88+
}
89+
lambda.refresh_return_type()
90+
}
8391
}
8492

8593
let mut visitor = PredicateVisitor { items };

‎src/query/sql/src/planner/optimizer/optimizers/rule/scalar_rules/rule_merge_eval_scalar.rs‎

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ use crate::optimizer::optimizers::rule::RuleID;
2525
use crate::optimizer::optimizers::rule::TransformResult;
2626
use crate::plans::EvalScalar;
2727
use crate::plans::FunctionCall;
28+
use crate::plans::LambdaFunc;
2829
use crate::plans::RelOp;
2930
use crate::plans::ScalarExpr;
3031
use crate::plans::ScalarItem;
@@ -129,8 +130,8 @@ struct TrivialComposer<'a> {
129130
supported: bool,
130131
}
131132

132-
impl<'a> VisitorMut<'a> for TrivialComposer<'_> {
133-
fn visit(&mut self, expr: &'a mut ScalarExpr) -> Result<()> {
133+
impl VisitorMut<'_> for TrivialComposer<'_> {
134+
fn visit(&mut self, expr: &mut ScalarExpr) -> Result<()> {
134135
if let ScalarExpr::BoundColumnRef(column) = expr
135136
&& let Some(item) = self
136137
.down_items
@@ -150,12 +151,19 @@ impl<'a> VisitorMut<'a> for TrivialComposer<'_> {
150151
walk_expr_mut(self, expr)
151152
}
152153

153-
fn visit_function_call(&mut self, function: &'a mut FunctionCall) -> Result<()> {
154+
fn visit_function_call(&mut self, function: &mut FunctionCall) -> Result<()> {
154155
for argument in &mut function.arguments {
155156
self.visit(argument)?;
156157
}
157158
function.refresh_return_type()
158159
}
160+
161+
fn visit_lambda_function(&mut self, lambda: &mut LambdaFunc) -> Result<()> {
162+
for argument in &mut lambda.args {
163+
self.visit(argument)?;
164+
}
165+
lambda.refresh_return_type()
166+
}
159167
}
160168

161169
impl Rule for RuleMergeEvalScalar {

‎src/query/sql/src/planner/plans/scalar_expr.rs‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1172,6 +1172,98 @@ pub struct LambdaFunc {
11721172
pub return_type: Box<DataType>,
11731173
}
11741174

1175+
impl LambdaFunc {
1176+
pub fn infer_return_type(&self) -> Result<DataType> {
1177+
if self.func_name == "json_path_transform" {
1178+
let [json, path, ..] = self.args.as_slice() else {
1179+
return Err(ErrorCode::Internal(
1180+
"json_path_transform requires json and path arguments",
1181+
));
1182+
};
1183+
let return_type = DataType::Variant;
1184+
return Ok(
1185+
if json.data_type().is_nullable_or_null() || path.data_type().is_nullable_or_null()
1186+
{
1187+
return_type.wrap_nullable()
1188+
} else {
1189+
return_type
1190+
},
1191+
);
1192+
}
1193+
1194+
// Captured columns precede the collection argument.
1195+
let collection_type = self
1196+
.args
1197+
.last()
1198+
.ok_or_else(|| ErrorCode::Internal("lambda function requires a collection argument"))?
1199+
.data_type();
1200+
let is_nullable = collection_type.is_nullable_or_null();
1201+
let collection_type = collection_type.remove_nullable();
1202+
if matches!(
1203+
collection_type,
1204+
DataType::Null | DataType::EmptyArray | DataType::EmptyMap
1205+
) {
1206+
return Ok(collection_type);
1207+
}
1208+
let lambda_type = || {
1209+
self.lambda_expr
1210+
.as_expr(&BUILTIN_FUNCTIONS)
1211+
.data_type()
1212+
.clone()
1213+
};
1214+
1215+
let return_type = match self.func_name.as_str() {
1216+
"array_filter" | "map_filter" => collection_type,
1217+
"array_reduce" => match collection_type {
1218+
DataType::Array(inner_type) => inner_type.wrap_nullable(),
1219+
_ => {
1220+
return Err(ErrorCode::Internal(
1221+
"array_reduce requires an array argument",
1222+
));
1223+
}
1224+
},
1225+
"array_transform" | "array_apply" | "array_map" => {
1226+
DataType::Array(Box::new(lambda_type()))
1227+
}
1228+
"map_transform_keys" | "map_transform_values" => {
1229+
let DataType::Map(box DataType::Tuple(fields)) = collection_type else {
1230+
return Err(ErrorCode::Internal(
1231+
"map lambda function requires a map argument",
1232+
));
1233+
};
1234+
if fields.len() != 2 {
1235+
return Err(ErrorCode::Internal(
1236+
"map lambda function requires key and value fields",
1237+
));
1238+
}
1239+
let fields = if self.func_name == "map_transform_keys" {
1240+
vec![lambda_type(), fields[1].clone()]
1241+
} else {
1242+
vec![fields[0].clone(), lambda_type()]
1243+
};
1244+
DataType::Map(Box::new(DataType::Tuple(fields)))
1245+
}
1246+
_ => {
1247+
return Err(ErrorCode::Internal(format!(
1248+
"unsupported lambda function {}",
1249+
self.func_name
1250+
)));
1251+
}
1252+
};
1253+
1254+
Ok(if is_nullable {
1255+
return_type.wrap_nullable()
1256+
} else {
1257+
return_type
1258+
})
1259+
}
1260+
1261+
pub fn refresh_return_type(&mut self) -> Result<()> {
1262+
self.return_type = Box::new(self.infer_return_type()?);
1263+
Ok(())
1264+
}
1265+
}
1266+
11751267
#[derive(Clone, Educe)]
11761268
#[educe(Debug, PartialEq, Eq, Hash)]
11771269
pub struct FunctionCall {

0 commit comments

Comments
 (0)