Skip to content

Commit 9449729

Browse files
mergify[bot]mergify[bot]
authored
Merge pull request #10439 from RinChanNOWWW/catch
feat: do not throw an error if it appears in a short-circuited branch.
2 parents f992dde + 534d194 commit 9449729

File tree

11 files changed

+551
-91
lines changed

11 files changed

+551
-91
lines changed

‎src/query/expression/src/evaluator.rs‎

Lines changed: 176 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,11 @@
1313
// limitations under the License.
1414

1515
use std::collections::HashMap;
16+
use std::ops::Not;
1617

1718
use common_arrow::arrow::bitmap;
19+
use common_arrow::arrow::bitmap::Bitmap;
20+
use common_arrow::arrow::bitmap::MutableBitmap;
1821
use common_exception::ErrorCode;
1922
use common_exception::Result;
2023
use common_exception::Span;
@@ -43,6 +46,9 @@ use crate::FunctionContext;
4346
use crate::FunctionDomain;
4447
use crate::FunctionRegistry;
4548

49+
pub(crate) type EvalError = (Span, Value<AnyType>, Bitmap, String);
50+
pub(crate) type EvalResult<T> = std::result::Result<T, EvalError>;
51+
4652
pub struct Evaluator<'a> {
4753
input_columns: &'a DataBlock,
4854
func_ctx: FunctionContext,
@@ -113,7 +119,15 @@ impl<'a> Evaluator<'a> {
113119
new_evaluator.run(expr)
114120
}
115121

122+
#[inline]
116123
pub fn run(&self, expr: &Expr) -> Result<Value<AnyType>> {
124+
match self.run_impl(expr) {
125+
Ok(v) => Ok(v),
126+
Err((span, _, _, msg)) => Err(ErrorCode::Internal(msg).set_span(span)),
127+
}
128+
}
129+
130+
pub fn run_impl(&self, expr: &Expr) -> EvalResult<Value<AnyType>> {
117131
#[cfg(debug_assertions)]
118132
self.check_expr(expr);
119133

@@ -129,8 +143,8 @@ impl<'a> Evaluator<'a> {
129143
} => {
130144
let cols = args
131145
.iter()
132-
.map(|expr| self.run(expr))
133-
.collect::<Result<Vec<_>>>()?;
146+
.map(|expr| self.run_impl(expr))
147+
.collect::<EvalResult<Vec<_>>>()?;
134148
assert!(
135149
cols.iter()
136150
.filter_map(|val| match val {
@@ -146,9 +160,10 @@ impl<'a> Evaluator<'a> {
146160
validity: None,
147161
errors: None,
148162
tz: self.func_ctx.tz,
163+
already_rendered: false,
149164
};
150165
let result = (function.eval)(cols_ref.as_slice(), &mut ctx);
151-
ctx.render_error(*span, &cols, &function.signature.name)?;
166+
ctx.render_error(*span, &cols, &result, &function.signature.name)?;
152167
Ok(result)
153168
}
154169
Expr::Cast {
@@ -157,13 +172,50 @@ impl<'a> Evaluator<'a> {
157172
expr,
158173
dest_type,
159174
} => {
160-
let value = self.run(expr)?;
175+
let value = self.run_impl(expr)?;
161176
if *is_try {
162177
self.run_try_cast(*span, expr.data_type(), dest_type, value)
163178
} else {
164179
self.run_cast(*span, expr.data_type(), dest_type, value)
165180
}
166181
}
182+
Expr::Catch {
183+
expr, data_type, ..
184+
} => Ok(match self.run_impl(expr) {
185+
Ok(value) => match value {
186+
Value::Scalar(scalar) => {
187+
Value::Scalar(Scalar::Tuple(vec![scalar, Scalar::Null]))
188+
}
189+
Value::Column(col) => {
190+
let len = col.len();
191+
Value::Column(Column::Tuple {
192+
fields: vec![col, Column::Null { len }],
193+
len,
194+
})
195+
}
196+
},
197+
Err((_, value, bitmap, err)) => {
198+
let num_rows = self.input_columns.num_rows();
199+
let inner_type = data_type.as_tuple().unwrap().first().unwrap();
200+
let value_col = value.convert_to_full_column(inner_type, num_rows);
201+
let err_scalar = Scalar::String(err.into_bytes());
202+
let err_col =
203+
ColumnBuilder::repeat(&err_scalar.as_ref(), num_rows, &DataType::String)
204+
.build();
205+
206+
// If the element in the bitmap is true,
207+
// it means the corresponding row is error.
208+
let bitmap = bitmap.not();
209+
let err_col = Column::Nullable(Box::new(NullableColumn {
210+
column: err_col,
211+
validity: bitmap,
212+
}));
213+
Value::Column(Column::Tuple {
214+
fields: vec![value_col, err_col],
215+
len: num_rows,
216+
})
217+
}
218+
}),
167219
};
168220

169221
#[cfg(debug_assertions)]
@@ -203,7 +255,7 @@ impl<'a> Evaluator<'a> {
203255
src_type: &DataType,
204256
dest_type: &DataType,
205257
value: Value<AnyType>,
206-
) -> Result<Value<AnyType>> {
258+
) -> EvalResult<Value<AnyType>> {
207259
if src_type == dest_type {
208260
return Ok(value);
209261
}
@@ -244,17 +296,27 @@ impl<'a> Evaluator<'a> {
244296
other => unreachable!("source: {}", other),
245297
},
246298
(DataType::Nullable(inner_src_ty), _) => match value {
247-
Value::Scalar(Scalar::Null) => Err(ErrorCode::Internal(format!(
248-
"unable to cast type `{src_type}` to type `{dest_type}`"
249-
))
250-
.set_span(span)),
299+
Value::Scalar(Scalar::Null) => {
300+
let mut bitmap = MutableBitmap::new();
301+
bitmap.extend_constant(self.input_columns.num_rows(), false);
302+
Err((
303+
span,
304+
Value::Scalar(Scalar::default_value(dest_type)),
305+
bitmap.into(),
306+
format!("unable to cast type `{src_type}` to type `{dest_type}`"),
307+
))
308+
}
251309
Value::Scalar(_) => self.run_cast(span, inner_src_ty, dest_type, value),
252310
Value::Column(Column::Nullable(col)) => {
253311
if col.validity.unset_bits() > 0 {
254-
return Err(ErrorCode::Internal(format!(
255-
"unable to cast `NULL` to type `{dest_type}`"
256-
))
257-
.set_span(span));
312+
let mut bitmap = MutableBitmap::new();
313+
bitmap.extend_constant(self.input_columns.num_rows(), false);
314+
return Err((
315+
span,
316+
Value::Scalar(Scalar::default_value(dest_type)),
317+
bitmap.into(),
318+
format!("unable to cast `NULL` to type `{dest_type}`"),
319+
));
258320
}
259321
let column = self
260322
.run_cast(span, inner_src_ty, dest_type, Value::Column(col.column))?
@@ -361,7 +423,7 @@ impl<'a> Evaluator<'a> {
361423
self.run_cast(span, src_ty, dest_ty, Value::Scalar(field))
362424
.map(|val| val.into_scalar().unwrap())
363425
})
364-
.collect::<Result<Vec<_>>>()?;
426+
.collect::<EvalResult<Vec<_>>>()?;
365427
Ok(Value::Scalar(Scalar::Tuple(new_fields)))
366428
}
367429
Value::Column(Column::Tuple { fields, len }) => {
@@ -373,7 +435,7 @@ impl<'a> Evaluator<'a> {
373435
self.run_cast(span, src_ty, dest_ty, Value::Column(field))
374436
.map(|val| val.into_column().unwrap())
375437
})
376-
.collect::<Result<_>>()?;
438+
.collect::<EvalResult<_>>()?;
377439
Ok(Value::Column(Column::Tuple {
378440
fields: new_fields,
379441
len,
@@ -383,10 +445,16 @@ impl<'a> Evaluator<'a> {
383445
}
384446
}
385447

386-
_ => Err(ErrorCode::Internal(format!(
387-
"unable to cast type `{src_type}` to type `{dest_type}`"
388-
))
389-
.set_span(span)),
448+
_ => {
449+
let mut bitmap = MutableBitmap::new();
450+
bitmap.extend_constant(self.input_columns.num_rows(), false);
451+
Err((
452+
span,
453+
Value::Scalar(Scalar::default_value(dest_type)),
454+
bitmap.into(),
455+
format!("unable to cast type `{src_type}` to type `{dest_type}`"),
456+
))
457+
}
390458
}
391459
}
392460

@@ -396,7 +464,7 @@ impl<'a> Evaluator<'a> {
396464
src_type: &DataType,
397465
dest_type: &DataType,
398466
value: Value<AnyType>,
399-
) -> Result<Value<AnyType>> {
467+
) -> EvalResult<Value<AnyType>> {
400468
if src_type == dest_type {
401469
return Ok(value);
402470
}
@@ -542,7 +610,7 @@ impl<'a> Evaluator<'a> {
542610
.into_scalar()
543611
.unwrap())
544612
})
545-
.collect::<Result<_>>()?;
613+
.collect::<EvalResult<_>>()?;
546614
Ok(Value::Scalar(Scalar::Tuple(new_fields)))
547615
}
548616
Value::Column(Column::Tuple { fields, len }) => {
@@ -556,7 +624,7 @@ impl<'a> Evaluator<'a> {
556624
.into_column()
557625
.unwrap())
558626
})
559-
.collect::<Result<_>>()?;
627+
.collect::<EvalResult<_>>()?;
560628
let new_col = Column::Tuple {
561629
fields: new_fields,
562630
len,
@@ -567,10 +635,16 @@ impl<'a> Evaluator<'a> {
567635
}
568636
}
569637

570-
_ => Err(ErrorCode::Internal(format!(
571-
"unable to cast type `{src_type}` to type `{dest_type}`"
572-
))
573-
.set_span(span)),
638+
_ => {
639+
let mut bitmap = MutableBitmap::new();
640+
bitmap.extend_constant(self.input_columns.num_rows(), false);
641+
Err((
642+
span,
643+
Value::Scalar(Scalar::default_value(dest_type)),
644+
bitmap.into(),
645+
format!("unable to cast type `{src_type}` to type `{dest_type}`"),
646+
))
647+
}
574648
}
575649
}
576650

@@ -581,7 +655,7 @@ impl<'a> Evaluator<'a> {
581655
dest_type: &DataType,
582656
value: Value<AnyType>,
583657
cast_fn: &str,
584-
) -> Result<Option<Value<AnyType>>> {
658+
) -> EvalResult<Option<Value<AnyType>>> {
585659
let expr = Expr::ColumnRef {
586660
span,
587661
id: 0,
@@ -616,7 +690,7 @@ impl<'a> Evaluator<'a> {
616690
num_rows,
617691
);
618692
let evaluator = Evaluator::new(&block, self.func_ctx, self.fn_registry);
619-
Ok(Some(evaluator.run(&cast_expr)?))
693+
Ok(Some(evaluator.run_impl(&cast_expr)?))
620694
}
621695
}
622696

@@ -769,6 +843,80 @@ impl<'a, Index: ColumnIndex> ConstantFolder<'a, Index> {
769843
new_domain,
770844
)
771845
}
846+
Expr::Catch {
847+
span,
848+
data_type,
849+
expr,
850+
} => {
851+
let (inner_expr, inner_domain) = self.fold_once(expr);
852+
let inner_type = &data_type.as_tuple().unwrap()[0];
853+
854+
if inner_expr.as_constant().is_some() {
855+
return (
856+
Expr::Catch {
857+
span: *span,
858+
data_type: data_type.clone(),
859+
expr: Box::new(inner_expr),
860+
},
861+
inner_domain.map(|d| {
862+
Domain::Tuple(vec![
863+
d,
864+
Domain::Nullable(NullableDomain {
865+
has_null: true,
866+
value: None,
867+
}),
868+
])
869+
}),
870+
);
871+
}
872+
873+
match inner_domain {
874+
// If the inner domain is not None,
875+
// it means the inner expr will not throw an error.
876+
Some(d) => match d.as_singleton() {
877+
Some(scalar) => (
878+
Expr::Catch {
879+
span: *span,
880+
data_type: data_type.clone(),
881+
expr: Box::new(Expr::Constant {
882+
span: *span,
883+
scalar,
884+
data_type: inner_type.clone(),
885+
}),
886+
},
887+
Some(Domain::Tuple(vec![
888+
d,
889+
Domain::Nullable(NullableDomain {
890+
has_null: true,
891+
value: None,
892+
}),
893+
])),
894+
),
895+
None => (
896+
Expr::Catch {
897+
span: *span,
898+
data_type: data_type.clone(),
899+
expr: Box::new(inner_expr),
900+
},
901+
Some(Domain::Tuple(vec![
902+
d,
903+
Domain::Nullable(NullableDomain {
904+
has_null: true,
905+
value: None,
906+
}),
907+
])),
908+
),
909+
},
910+
None => (
911+
Expr::Catch {
912+
span: *span,
913+
data_type: data_type.clone(),
914+
expr: Box::new(inner_expr),
915+
},
916+
None,
917+
),
918+
}
919+
}
772920
Expr::FunctionCall {
773921
span,
774922
id,

0 commit comments

Comments
 (0)