From a650ba6f0ddf2b747b9d7bbfcbe61110b4d8736e Mon Sep 17 00:00:00 2001 From: sundy-li <543950155@qq.com> Date: Sun, 17 Mar 2024 03:57:08 +0800 Subject: [PATCH 1/4] feat(query): add error_or function --- src/query/expression/src/evaluator.rs | 149 ++++++++++-------- src/query/expression/src/filter/selector.rs | 25 ++- src/query/expression/src/function.rs | 4 + src/query/functions/src/scalars/control.rs | 12 ++ .../functions/src/scalars/decimal/cast.rs | 1 + .../it/scalars/testdata/function_list.txt | 1 + .../sql/src/planner/semantic/type_check.rs | 76 ++++++++- 7 files changed, 192 insertions(+), 76 deletions(-) diff --git a/src/query/expression/src/evaluator.rs b/src/query/expression/src/evaluator.rs index 74d8b4b8f7b3..3c6aba3212db 100644 --- a/src/query/expression/src/evaluator.rs +++ b/src/query/expression/src/evaluator.rs @@ -51,6 +51,28 @@ use crate::FunctionEval; use crate::FunctionRegistry; use crate::RemoteExpr; +#[derive(Default)] +pub struct EvaluateOptions<'a> { + pub selection: Option<&'a [u32]>, + pub suppress_error: bool, +} + +impl<'a> EvaluateOptions<'a> { + pub fn new(selection: Option<&'a [u32]>) -> EvaluateOptions<'a> { + Self { + suppress_error: false, + selection, + } + } + + pub fn with_suppress_error(&mut self, suppress_error: bool) -> Self { + Self { + suppress_error, + selection: self.selection, + } + } +} + pub struct Evaluator<'a> { data_block: &'a DataBlock, func_ctx: &'a FunctionContext, @@ -98,7 +120,7 @@ impl<'a> Evaluator<'a> { } pub fn run(&self, expr: &Expr) -> Result> { - self.partial_run(expr, None, None) + self.partial_run(expr, None, &mut EvaluateOptions::default()) } /// Run an expression partially, only the rows that are valid in the validity bitmap @@ -107,7 +129,7 @@ impl<'a> Evaluator<'a> { &self, expr: &Expr, validity: Option, - selection: Option<&[u32]>, + options: &mut EvaluateOptions, ) -> Result> { debug_assert!( validity.is_none() || validity.as_ref().unwrap().len() == self.data_block.num_rows() @@ -125,18 +147,11 @@ impl<'a> Evaluator<'a> { expr, dest_type, } => { - let value = self.partial_run(expr, validity.clone(), selection)?; + let value = self.partial_run(expr, validity.clone(), options)?; if *is_try { self.run_try_cast(*span, expr.data_type(), dest_type, value) } else { - self.run_cast( - *span, - expr.data_type(), - dest_type, - value, - validity, - selection, - ) + self.run_cast(*span, expr.data_type(), dest_type, value, validity, options) } } Expr::FunctionCall { @@ -144,14 +159,12 @@ impl<'a> Evaluator<'a> { args, generics, .. - } if function.signature.name == "if" => { - self.eval_if(args, generics, validity, selection) - } + } if function.signature.name == "if" => self.eval_if(args, generics, validity, options), Expr::FunctionCall { function, args, .. } if function.signature.name == "and_filters" => { - self.eval_and_filters(args, validity, selection) + self.eval_and_filters(args, validity, options) } Expr::FunctionCall { @@ -162,10 +175,18 @@ impl<'a> Evaluator<'a> { generics, .. } => { + let suppress_error = function.signature.name == "is_not_error"; let args = args .iter() - .map(|expr| self.partial_run(expr, validity.clone(), selection)) + .map(|expr| { + self.partial_run( + expr, + validity.clone(), + &mut options.with_suppress_error(suppress_error), + ) + }) .collect::>>()?; + assert!( args.iter() .filter_map(|val| match val { @@ -181,6 +202,7 @@ impl<'a> Evaluator<'a> { validity, errors: None, func_ctx: self.func_ctx, + suppress_error: options.suppress_error, }; let (_, eval) = function.eval.as_scalar().unwrap(); let result = (eval)(cols_ref.as_slice(), &mut ctx); @@ -189,7 +211,7 @@ impl<'a> Evaluator<'a> { id.params(), &args, &function.signature.name, - selection, + options.selection, )?; Ok(result) } @@ -202,7 +224,7 @@ impl<'a> Evaluator<'a> { } => { let args = args .iter() - .map(|expr| self.partial_run(expr, validity.clone(), selection)) + .map(|expr| self.partial_run(expr, validity.clone(), options)) .collect::>>()?; assert!( args.iter() @@ -266,7 +288,7 @@ impl<'a> Evaluator<'a> { dest_type: &DataType, value: Value, validity: Option, - selection: Option<&[u32]>, + options: &mut EvaluateOptions, ) -> Result> { if src_type == dest_type { return Ok(value); @@ -280,7 +302,7 @@ impl<'a> Evaluator<'a> { value.clone(), &cast_fn, validity.clone(), - selection, + options, )? { return Ok(new_value); } @@ -300,14 +322,9 @@ impl<'a> Evaluator<'a> { }, (DataType::Nullable(inner_src_ty), DataType::Nullable(inner_dest_ty)) => match value { Value::Scalar(Scalar::Null) => Ok(Value::Scalar(Scalar::Null)), - Value::Scalar(_) => self.run_cast( - span, - inner_src_ty, - inner_dest_ty, - value, - validity, - selection, - ), + Value::Scalar(_) => { + self.run_cast(span, inner_src_ty, inner_dest_ty, value, validity, options) + } Value::Column(Column::Nullable(col)) => { let validity = validity .map(|validity| (&validity) & (&col.validity)) @@ -319,7 +336,7 @@ impl<'a> Evaluator<'a> { inner_dest_ty, Value::Column(col.column), Some(validity.clone()), - selection, + options, )? .into_column() .unwrap(); @@ -345,7 +362,7 @@ impl<'a> Evaluator<'a> { } } Value::Scalar(_) => { - self.run_cast(span, inner_src_ty, dest_type, value, validity, selection) + self.run_cast(span, inner_src_ty, dest_type, value, validity, options) } Value::Column(Column::Nullable(col)) => { let has_valid_nulls = validity @@ -367,7 +384,7 @@ impl<'a> Evaluator<'a> { dest_type, Value::Column(col.column), validity, - selection, + options, )? .into_column() .unwrap(); @@ -382,7 +399,7 @@ impl<'a> Evaluator<'a> { inner_dest_ty, Value::Scalar(scalar), validity, - selection, + options, ), Value::Column(col) => { let column = self @@ -392,7 +409,7 @@ impl<'a> Evaluator<'a> { inner_dest_ty, Value::Column(col), validity, - selection, + options, )? .into_column() .unwrap(); @@ -429,7 +446,7 @@ impl<'a> Evaluator<'a> { inner_dest_ty, Value::Column(array), validity, - selection, + options, )? .into_column() .unwrap(); @@ -453,7 +470,7 @@ impl<'a> Evaluator<'a> { inner_dest_ty, Value::Column(col.values), validity, - selection, + options, )? .into_column() .unwrap(); @@ -490,7 +507,7 @@ impl<'a> Evaluator<'a> { inner_dest_ty, Value::Column(array), validity, - selection, + options, )? .into_column() .unwrap(); @@ -514,7 +531,7 @@ impl<'a> Evaluator<'a> { inner_dest_ty, Value::Column(col.values), validity, - selection, + options, )? .into_column() .unwrap(); @@ -541,7 +558,7 @@ impl<'a> Evaluator<'a> { dest_ty, Value::Scalar(field), validity.clone(), - selection, + options, ) .map(|val| val.into_scalar().unwrap()) }) @@ -560,7 +577,7 @@ impl<'a> Evaluator<'a> { dest_ty, Value::Column(field), validity.clone(), - selection, + options, ) .map(|val| val.into_column().unwrap()) }) @@ -601,7 +618,7 @@ impl<'a> Evaluator<'a> { value.clone(), &cast_fn, None, - None, + &mut EvaluateOptions::default(), ) { return Ok(new_value); } @@ -777,7 +794,7 @@ impl<'a> Evaluator<'a> { value: Value, cast_fn: &str, validity: Option, - selection: Option<&[u32]>, + options: &mut EvaluateOptions, ) -> Result>> { let expr = Expr::ColumnRef { span, @@ -813,9 +830,7 @@ impl<'a> Evaluator<'a> { }); let block = DataBlock::new(vec![BlockEntry::new(src_type.clone(), value)], num_rows); let evaluator = Evaluator::new(&block, self.func_ctx, self.fn_registry); - Ok(Some( - evaluator.partial_run(&cast_expr, validity, selection)?, - )) + Ok(Some(evaluator.partial_run(&cast_expr, validity, options)?)) } // `if` is a special builtin function that could partially evaluate its arguments @@ -827,7 +842,7 @@ impl<'a> Evaluator<'a> { args: &[Expr], generics: &[DataType], validity: Option, - selection: Option<&[u32]>, + options: &mut EvaluateOptions, ) -> Result> { if args.len() < 3 && args.len() % 2 == 0 { unreachable!() @@ -849,7 +864,7 @@ impl<'a> Evaluator<'a> { let mut flags = Vec::new(); let mut results = Vec::new(); for cond_idx in (0..args.len() - 1).step_by(2) { - let cond = self.partial_run(&args[cond_idx], Some(validity.clone()), selection)?; + let cond = self.partial_run(&args[cond_idx], Some(validity.clone()), options)?; match cond.try_downcast::>().unwrap() { Value::Scalar(None | Some(false)) => { results.push(Value::Scalar(Scalar::default_value(&generics[0]))); @@ -859,7 +874,7 @@ impl<'a> Evaluator<'a> { results.push(self.partial_run( &args[cond_idx + 1], Some(validity.clone()), - selection, + options, )?); validity = Bitmap::new_constant(false, num_rows); flags.push(Bitmap::new_constant(true, len.unwrap_or(1))); @@ -870,7 +885,7 @@ impl<'a> Evaluator<'a> { results.push(self.partial_run( &args[cond_idx + 1], Some((&validity) & (&flag)), - selection, + options, )?); validity = (&validity) & (&flag.not()); flags.push(flag); @@ -878,7 +893,7 @@ impl<'a> Evaluator<'a> { }; conds.push(cond); } - let else_result = self.partial_run(&args[args.len() - 1], Some(validity), selection)?; + let else_result = self.partial_run(&args[args.len() - 1], Some(validity), options)?; // Assert that all the arguments have the same length. assert!( @@ -916,12 +931,12 @@ impl<'a> Evaluator<'a> { &self, args: &[Expr], mut validity: Option, - selection: Option<&[u32]>, + options: &mut EvaluateOptions, ) -> Result> { assert!(args.len() >= 2); for arg in args { - let cond = self.partial_run(arg, validity.clone(), selection)?; + let cond = self.partial_run(arg, validity.clone(), options)?; match &cond { Value::Scalar(Scalar::Null | Scalar::Boolean(false)) => { return Ok(Value::Scalar(Scalar::Boolean(false))); @@ -987,6 +1002,7 @@ impl<'a> Evaluator<'a> { validity: None, errors: None, func_ctx: self.func_ctx, + suppress_error: false, }; let result = (eval)(&cols_ref, &mut ctx, max_nums_per_row); ctx.render_error(*span, id.params(), &args, &function.signature.name, None)?; @@ -1004,6 +1020,7 @@ impl<'a> Evaluator<'a> { return Ok(Scalar::Null); } let mut arg0 = column.index(0).unwrap().to_owned(); + let mut eval_options = EvaluateOptions::default(); for i in 1..column.len() { let arg1 = column.index(i).unwrap().to_owned(); let entries = { @@ -1016,7 +1033,14 @@ impl<'a> Evaluator<'a> { let evaluator = Evaluator::new(&block, self.func_ctx, self.fn_registry); let result = evaluator.run(expr)?; arg0 = self - .run_cast(None, expr.data_type(), &col_type, result, None, None)? + .run_cast( + None, + expr.data_type(), + &col_type, + result, + None, + &mut eval_options, + )? .as_scalar() .unwrap() .clone(); @@ -1151,11 +1175,11 @@ impl<'a> Evaluator<'a> { pub fn get_children( &self, args: &[Expr], - selection: Option<&[u32]>, + options: &mut EvaluateOptions, ) -> Result, DataType)>> { let children = args .iter() - .map(|expr| self.get_select_child(expr, selection)) + .map(|expr| self.get_select_child(expr, options)) .collect::>>()?; assert!( children @@ -1186,7 +1210,7 @@ impl<'a> Evaluator<'a> { pub fn get_select_child( &self, expr: &Expr, - selection: Option<&[u32]>, + options: &mut EvaluateOptions, ) -> Result<(Value, DataType)> { #[cfg(debug_assertions)] self.check_expr(expr); @@ -1206,7 +1230,7 @@ impl<'a> Evaluator<'a> { expr, dest_type, } => { - let value = self.get_select_child(expr, selection)?.0; + let value = self.get_select_child(expr, options)?.0; if *is_try { Ok(( self.run_try_cast(*span, expr.data_type(), dest_type, value)?, @@ -1214,7 +1238,7 @@ impl<'a> Evaluator<'a> { )) } else { Ok(( - self.run_cast(*span, expr.data_type(), dest_type, value, None, selection)?, + self.run_cast(*span, expr.data_type(), dest_type, value, None, options)?, dest_type.clone(), )) } @@ -1227,7 +1251,7 @@ impl<'a> Evaluator<'a> { } if function.signature.name == "if" => { let return_type = self.remove_generics_data_type(generics, &function.signature.return_type); - Ok((self.eval_if(args, generics, None, selection)?, return_type)) + Ok((self.eval_if(args, generics, None, options)?, return_type)) } Expr::FunctionCall { @@ -1238,7 +1262,7 @@ impl<'a> Evaluator<'a> { } if function.signature.name == "and_filters" => { let return_type = self.remove_generics_data_type(generics, &function.signature.return_type); - Ok((self.eval_and_filters(args, None, selection)?, return_type)) + Ok((self.eval_and_filters(args, None, options)?, return_type)) } Expr::FunctionCall { @@ -1251,7 +1275,7 @@ impl<'a> Evaluator<'a> { } => { let args = args .iter() - .map(|expr| self.get_select_child(expr, selection)) + .map(|expr| self.get_select_child(expr, options)) .collect::>>()?; assert!( args.iter() @@ -1272,6 +1296,7 @@ impl<'a> Evaluator<'a> { validity: None, errors: None, func_ctx: self.func_ctx, + suppress_error: options.suppress_error, }; let (_, eval) = function.eval.as_scalar().unwrap(); let result = (eval)(cols_ref.as_slice(), &mut ctx); @@ -1281,7 +1306,7 @@ impl<'a> Evaluator<'a> { id.params(), &args, &function.signature.name, - selection, + options.selection, )?; let return_type = self.remove_generics_data_type(generics, &function.signature.return_type); @@ -1296,7 +1321,7 @@ impl<'a> Evaluator<'a> { } => { let args = args .iter() - .map(|expr| self.partial_run(expr, None, None)) + .map(|expr| self.partial_run(expr, None, &mut EvaluateOptions::default())) .collect::>>()?; assert!( args.iter() diff --git a/src/query/expression/src/filter/selector.rs b/src/query/expression/src/filter/selector.rs index 782333d16583..54c8f801d073 100644 --- a/src/query/expression/src/filter/selector.rs +++ b/src/query/expression/src/filter/selector.rs @@ -25,6 +25,7 @@ use crate::filter::SelectExpr; use crate::filter::SelectOp; use crate::types::DataType; use crate::EvalContext; +use crate::EvaluateOptions; use crate::Evaluator; use crate::Expr; use crate::Scalar; @@ -282,7 +283,8 @@ impl<'a> Selector<'a> { *mutable_false_idx + count, &select_strategy, ); - let children = self.evaluator.get_children(exprs, selection)?; + let mut eval_options = EvaluateOptions::new(selection); + let children = self.evaluator.get_children(exprs, &mut eval_options)?; let (left_value, left_data_type) = children[0].clone(); let (right_value, right_data_type) = children[1].clone(); let left_data_type = self @@ -332,7 +334,11 @@ impl<'a> Selector<'a> { *mutable_false_idx + count, &select_strategy, ); - let result = self.evaluator.eval_if(args, generics, None, selection)?; + let mut eval_options = EvaluateOptions::new(selection); + + let result = self + .evaluator + .eval_if(args, generics, None, &mut eval_options)?; let data_type = self .evaluator .remove_generics_data_type(generics, &function.signature.return_type); @@ -366,9 +372,12 @@ impl<'a> Selector<'a> { *mutable_false_idx + count, &select_strategy, ); + let mut eval_options = EvaluateOptions::new(selection) + .with_suppress_error(function.signature.name == "is_not_error"); + let args = args .iter() - .map(|expr| self.evaluator.partial_run(expr, None, selection)) + .map(|expr| self.evaluator.partial_run(expr, None, &mut eval_options)) .collect::>>()?; assert!( args.iter() @@ -385,6 +394,7 @@ impl<'a> Selector<'a> { validity: None, errors: None, func_ctx: self.evaluator.func_ctx(), + suppress_error: eval_options.suppress_error, }; let (_, eval) = function.eval.as_scalar().unwrap(); let result = (eval)(cols_ref.as_slice(), &mut ctx); @@ -422,7 +432,8 @@ impl<'a> Selector<'a> { *mutable_false_idx + count, &select_strategy, ); - let value = self.evaluator.get_select_child(expr, selection)?.0; + let mut eval_options = EvaluateOptions::new(selection); + let value = self.evaluator.get_select_child(expr, &mut eval_options)?.0; let result = if *is_try { self.evaluator .run_try_cast(*span, expr.data_type(), dest_type, value)? @@ -433,7 +444,7 @@ impl<'a> Selector<'a> { dest_type, value, None, - selection, + &mut eval_options, )? }; self.select_value( @@ -461,9 +472,11 @@ impl<'a> Selector<'a> { *mutable_false_idx + count, &select_strategy, ); + let mut eval_options = EvaluateOptions::new(selection); + let args = args .iter() - .map(|expr| self.evaluator.partial_run(expr, None, selection)) + .map(|expr| self.evaluator.partial_run(expr, None, &mut eval_options)) .collect::>>()?; assert!( args.iter() diff --git a/src/query/expression/src/function.rs b/src/query/expression/src/function.rs index 26dc101ed8d5..f97ed5279bff 100755 --- a/src/query/expression/src/function.rs +++ b/src/query/expression/src/function.rs @@ -119,6 +119,7 @@ pub struct EvalContext<'a> { /// default value in nullable's inner column. pub validity: Option, pub errors: Option<(MutableBitmap, String)>, + pub suppress_error: bool, } /// `FunctionID` is a unique identifier for a function in the registry. It's used to @@ -564,6 +565,9 @@ impl<'a> EvalContext<'a> { func_name: &str, selection: Option<&[u32]>, ) -> Result<()> { + if self.suppress_error { + return Ok(()); + } match &self.errors { Some((valids, error)) => { let first_error_row = if let Some(selection) = selection { diff --git a/src/query/functions/src/scalars/control.rs b/src/query/functions/src/scalars/control.rs index 457cb15729e8..3b401576d045 100644 --- a/src/query/functions/src/scalars/control.rs +++ b/src/query/functions/src/scalars/control.rs @@ -130,4 +130,16 @@ pub fn register(registry: &mut FunctionRegistry) { ValueRef::Scalar(Some(_)) => Value::Scalar(true), }, ); + + registry.register_1_arg_core::, BooleanType, _, _>( + "is_not_error", + |_, _| FunctionDomain::Full, + |arg, ctx| match ctx.errors.take() { + Some((bitmap, _)) => match arg { + ValueRef::Column(_) => Value::Column(bitmap.into()), + ValueRef::Scalar(_) => Value::Scalar(false), + }, + None => Value::Scalar(false), + }, + ); } diff --git a/src/query/functions/src/scalars/decimal/cast.rs b/src/query/functions/src/scalars/decimal/cast.rs index 7b6cb6768e3f..9550df85517c 100644 --- a/src/query/functions/src/scalars/decimal/cast.rs +++ b/src/query/functions/src/scalars/decimal/cast.rs @@ -436,6 +436,7 @@ fn convert_to_decimal_domain( func_ctx, validity: None, errors: None, + suppress_error: false, }; let dest_size = dest_type.size(); let res = convert_to_decimal(&value.as_ref(), &mut ctx, &from_type, dest_type); diff --git a/src/query/functions/tests/it/scalars/testdata/function_list.txt b/src/query/functions/tests/it/scalars/testdata/function_list.txt index 6dbd297207ad..8382816ba9ad 100644 --- a/src/query/functions/tests/it/scalars/testdata/function_list.txt +++ b/src/query/functions/tests/it/scalars/testdata/function_list.txt @@ -2182,6 +2182,7 @@ Functions overloads: 1 is_float(Variant NULL) :: Boolean NULL 0 is_integer(Variant) :: Boolean 1 is_integer(Variant NULL) :: Boolean NULL +0 is_not_error(T0) :: Boolean 0 is_not_null(NULL) :: Boolean 1 is_not_null(T0 NULL) :: Boolean 0 is_null_value(Variant) :: Boolean diff --git a/src/query/sql/src/planner/semantic/type_check.rs b/src/query/sql/src/planner/semantic/type_check.rs index c691f9a450e2..edd8140bf214 100644 --- a/src/query/sql/src/planner/semantic/type_check.rs +++ b/src/query/sql/src/planner/semantic/type_check.rs @@ -2453,6 +2453,8 @@ impl<'a> TypeChecker<'a> { "nvl", "nvl2", "is_null", + "is_error", + "error_or", "coalesce", "last_query_id", "array_sort", @@ -2597,6 +2599,59 @@ impl<'a> TypeChecker<'a> { .await, ) } + ("is_error", &[arg_x]) => { + // Rewrite is_error(x) to not(is_not_error(x)) + Some( + self.resolve_unary_op(span, &UnaryOperator::Not, &Expr::FunctionCall { + span, + func: ASTFunctionCall { + distinct: false, + name: Identifier { + name: "is_not_error".to_string(), + quote: None, + span, + }, + args: vec![arg_x.clone()], + params: vec![], + window: None, + lambda: None, + }, + }) + .await, + ) + } + ("error_or", args) => { + // error_or(arg0, arg1, ..., argN) is essentially + // if(is_not_error(arg0), arg0, is_not_error(arg1), arg1, ..., argN) + let mut new_args = Vec::with_capacity(args.len() * 2 + 1); + + for arg in args.iter() { + let is_not_error = Expr::FunctionCall { + span, + func: ASTFunctionCall { + distinct: false, + name: Identifier { + name: "is_not_error".to_string(), + quote: None, + span, + }, + args: vec![(*arg).clone()], + params: vec![], + window: None, + lambda: None, + }, + }; + new_args.push(is_not_error); + new_args.push((*arg).clone()); + } + new_args.push(Expr::Literal { + span, + lit: Literal::Null, + }); + + let args_ref: Vec<&Expr> = new_args.iter().collect(); + Some(self.resolve_function(span, "if", vec![], &args_ref).await) + } ("coalesce", args) => { // coalesce(arg0, arg1, ..., argN) is essentially // if(is_not_null(arg0), assume_not_null(arg0), is_not_null(arg1), assume_not_null(arg1), ..., argN) @@ -2641,14 +2696,19 @@ impl<'a> TypeChecker<'a> { span, lit: Literal::Null, }); - new_args.push(Expr::Literal { - span, - lit: Literal::Null, - }); - new_args.push(Expr::Literal { - span, - lit: Literal::Null, - }); + + // coalesce(all_null) => null + if new_args.len() == 1 { + new_args.push(Expr::Literal { + span, + lit: Literal::Null, + }); + new_args.push(Expr::Literal { + span, + lit: Literal::Null, + }); + } + let args_ref: Vec<&Expr> = new_args.iter().collect(); Some(self.resolve_function(span, "if", vec![], &args_ref).await) } From b68c5d8c3938416d514344d12cea924e7c4c0ea7 Mon Sep 17 00:00:00 2001 From: sundy-li <543950155@qq.com> Date: Sun, 17 Mar 2024 06:07:21 +0800 Subject: [PATCH 2/4] feat(query): add error_or function --- src/query/expression/src/evaluator.rs | 49 +++++++++++++++++----- src/query/functions/src/scalars/control.rs | 4 +- src/query/sql/src/evaluator/cse.rs | 1 + 3 files changed, 41 insertions(+), 13 deletions(-) diff --git a/src/query/expression/src/evaluator.rs b/src/query/expression/src/evaluator.rs index 3c6aba3212db..ae388ccc0b6f 100644 --- a/src/query/expression/src/evaluator.rs +++ b/src/query/expression/src/evaluator.rs @@ -55,6 +55,7 @@ use crate::RemoteExpr; pub struct EvaluateOptions<'a> { pub selection: Option<&'a [u32]>, pub suppress_error: bool, + pub errors: Option<(MutableBitmap, String)>, } impl<'a> EvaluateOptions<'a> { @@ -62,6 +63,7 @@ impl<'a> EvaluateOptions<'a> { Self { suppress_error: false, selection, + errors: None, } } @@ -69,6 +71,7 @@ impl<'a> EvaluateOptions<'a> { Self { suppress_error, selection: self.selection, + errors: None, } } } @@ -175,16 +178,12 @@ impl<'a> Evaluator<'a> { generics, .. } => { - let suppress_error = function.signature.name == "is_not_error"; + let child_suppress_error = function.signature.name == "is_not_error"; + let mut child_option = options.with_suppress_error(child_suppress_error); + let args = args .iter() - .map(|expr| { - self.partial_run( - expr, - validity.clone(), - &mut options.with_suppress_error(suppress_error), - ) - }) + .map(|expr| self.partial_run(expr, validity.clone(), &mut child_option)) .collect::>>()?; assert!( @@ -196,14 +195,21 @@ impl<'a> Evaluator<'a> { .all_equal() ); let cols_ref = args.iter().map(Value::as_ref).collect::>(); + + let errors = if !child_suppress_error { + None + } else { + child_option.errors.take() + }; let mut ctx = EvalContext { generics, num_rows: self.data_block.num_rows(), validity, - errors: None, + errors, func_ctx: self.func_ctx, suppress_error: options.suppress_error, }; + let (_, eval) = function.eval.as_scalar().unwrap(); let result = (eval)(cols_ref.as_slice(), &mut ctx); ctx.render_error( @@ -213,6 +219,12 @@ impl<'a> Evaluator<'a> { &function.signature.name, options.selection, )?; + + // inject errors into options, parent will handle it + if options.suppress_error { + options.errors = ctx.errors.take(); + } + Ok(result) } Expr::LambdaFunctionCall { @@ -1273,9 +1285,11 @@ impl<'a> Evaluator<'a> { generics, .. } => { + let child_suppress_error = function.signature.name == "is_not_error"; + let mut child_option = options.with_suppress_error(child_suppress_error); let args = args .iter() - .map(|expr| self.get_select_child(expr, options)) + .map(|expr| self.get_select_child(expr, &mut child_option)) .collect::>>()?; assert!( args.iter() @@ -1290,17 +1304,24 @@ impl<'a> Evaluator<'a> { .iter() .map(|(val, _)| Value::as_ref(val)) .collect::>(); + + let errors = if !child_suppress_error { + None + } else { + child_option.errors.take() + }; let mut ctx = EvalContext { generics, num_rows: self.data_block.num_rows(), validity: None, - errors: None, + errors, func_ctx: self.func_ctx, suppress_error: options.suppress_error, }; let (_, eval) = function.eval.as_scalar().unwrap(); let result = (eval)(cols_ref.as_slice(), &mut ctx); let args = args.into_iter().map(|(val, _)| val).collect::>(); + ctx.render_error( *span, id.params(), @@ -1308,6 +1329,12 @@ impl<'a> Evaluator<'a> { &function.signature.name, options.selection, )?; + + // inject errors into options, parent will handle it + if options.suppress_error { + options.errors = ctx.errors.take(); + } + let return_type = self.remove_generics_data_type(generics, &function.signature.return_type); Ok((result, return_type)) diff --git a/src/query/functions/src/scalars/control.rs b/src/query/functions/src/scalars/control.rs index 3b401576d045..a10e5d9f23fc 100644 --- a/src/query/functions/src/scalars/control.rs +++ b/src/query/functions/src/scalars/control.rs @@ -137,9 +137,9 @@ pub fn register(registry: &mut FunctionRegistry) { |arg, ctx| match ctx.errors.take() { Some((bitmap, _)) => match arg { ValueRef::Column(_) => Value::Column(bitmap.into()), - ValueRef::Scalar(_) => Value::Scalar(false), + ValueRef::Scalar(_) => Value::Scalar(bitmap.get(0)), }, - None => Value::Scalar(false), + None => Value::Scalar(true), }, ); } diff --git a/src/query/sql/src/evaluator/cse.rs b/src/query/sql/src/evaluator/cse.rs index c4bc673f19eb..717b0ceb8cf6 100644 --- a/src/query/sql/src/evaluator/cse.rs +++ b/src/query/sql/src/evaluator/cse.rs @@ -124,6 +124,7 @@ pub fn apply_cse( fn count_expressions(expr: &Expr, counter: &mut HashMap) { match expr { Expr::FunctionCall { function, .. } if function.signature.name == "if" => {} + Expr::FunctionCall { function, .. } if function.signature.name == "is_not_error" => {} Expr::FunctionCall { args, .. } | Expr::LambdaFunctionCall { args, .. } => { let entry = counter.entry(expr.clone()).or_insert(0); *entry += 1; From 43ec33924d4794558d4ba705be9863604583eda6 Mon Sep 17 00:00:00 2001 From: sundy-li <543950155@qq.com> Date: Sun, 17 Mar 2024 06:11:30 +0800 Subject: [PATCH 3/4] feat(query): add error_or function --- .../suites/query/02_function/02_0072_error.test | 9 +++++++++ 1 file changed, 9 insertions(+) create mode 100644 tests/sqllogictests/suites/query/02_function/02_0072_error.test diff --git a/tests/sqllogictests/suites/query/02_function/02_0072_error.test b/tests/sqllogictests/suites/query/02_function/02_0072_error.test new file mode 100644 index 000000000000..3ced6fb64544 --- /dev/null +++ b/tests/sqllogictests/suites/query/02_function/02_0072_error.test @@ -0,0 +1,9 @@ +query BBB +select is_error(from_base64('aj')), is_not_error(from_base64('ac')), is_error(3); +---- +1 0 0 + +query T +select error_or(from_base64('aak') , from_base64('aaj'), from_base64('MzQz')); +---- +333433 From 94e5992eb4b42aec2cb538496a2b0e6ee191e8e7 Mon Sep 17 00:00:00 2001 From: sundy-li <543950155@qq.com> Date: Sun, 17 Mar 2024 06:37:48 +0800 Subject: [PATCH 4/4] feat(query): add error_or function --- .../suites/mode/standalone/explain/push_down_filter.test | 4 ++-- .../mode/standalone/explain_native/push_down_filter.test | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/sqllogictests/suites/mode/standalone/explain/push_down_filter.test b/tests/sqllogictests/suites/mode/standalone/explain/push_down_filter.test index 71a345670f3d..eb1489bd6376 100644 --- a/tests/sqllogictests/suites/mode/standalone/explain/push_down_filter.test +++ b/tests/sqllogictests/suites/mode/standalone/explain/push_down_filter.test @@ -145,7 +145,7 @@ AggregateFinal ├── estimated rows: 0.00 ├── EvalScalar │ ├── output columns: [t.id (#0), de (#8)] - │ ├── expressions: [if(CAST(is_not_null(sum(tb.de) (#7)) AS Boolean NULL), CAST(assume_not_null(sum(tb.de) (#7)) AS Int64 NULL), true, 0, NULL, NULL, NULL)] + │ ├── expressions: [if(CAST(is_not_null(sum(tb.de) (#7)) AS Boolean NULL), CAST(assume_not_null(sum(tb.de) (#7)) AS Int64 NULL), true, 0, NULL)] │ ├── estimated rows: 0.00 │ └── AggregateFinal │ ├── output columns: [sum(tb.de) (#7), t.id (#0)] @@ -174,7 +174,7 @@ AggregateFinal │ │ ├── estimated rows: 0.00 │ │ └── EvalScalar │ │ ├── output columns: [t2.sid (#1), sum_arg_0 (#4)] - │ │ ├── expressions: [if(CAST(is_not_null(t3.val (#2)) AS Boolean NULL), CAST(assume_not_null(t3.val (#2)) AS Int32 NULL), true, 0, NULL, NULL, NULL)] + │ │ ├── expressions: [if(CAST(is_not_null(t3.val (#2)) AS Boolean NULL), CAST(assume_not_null(t3.val (#2)) AS Int32 NULL), true, 0, NULL)] │ │ ├── estimated rows: 0.00 │ │ └── TableScan │ │ ├── table: default.default.t2 diff --git a/tests/sqllogictests/suites/mode/standalone/explain_native/push_down_filter.test b/tests/sqllogictests/suites/mode/standalone/explain_native/push_down_filter.test index f43c65eff041..67e237a0657c 100644 --- a/tests/sqllogictests/suites/mode/standalone/explain_native/push_down_filter.test +++ b/tests/sqllogictests/suites/mode/standalone/explain_native/push_down_filter.test @@ -137,7 +137,7 @@ AggregateFinal ├── estimated rows: 0.00 ├── EvalScalar │ ├── output columns: [t.id (#0), de (#8)] - │ ├── expressions: [if(CAST(is_not_null(sum(tb.de) (#7)) AS Boolean NULL), CAST(assume_not_null(sum(tb.de) (#7)) AS Int64 NULL), true, 0, NULL, NULL, NULL)] + │ ├── expressions: [if(CAST(is_not_null(sum(tb.de) (#7)) AS Boolean NULL), CAST(assume_not_null(sum(tb.de) (#7)) AS Int64 NULL), true, 0, NULL)] │ ├── estimated rows: 0.00 │ └── AggregateFinal │ ├── output columns: [sum(tb.de) (#7), t.id (#0)] @@ -166,7 +166,7 @@ AggregateFinal │ │ ├── estimated rows: 0.00 │ │ └── EvalScalar │ │ ├── output columns: [t2.sid (#1), sum_arg_0 (#4)] - │ │ ├── expressions: [if(CAST(is_not_null(t3.val (#2)) AS Boolean NULL), CAST(assume_not_null(t3.val (#2)) AS Int32 NULL), true, 0, NULL, NULL, NULL)] + │ │ ├── expressions: [if(CAST(is_not_null(t3.val (#2)) AS Boolean NULL), CAST(assume_not_null(t3.val (#2)) AS Int32 NULL), true, 0, NULL)] │ │ ├── estimated rows: 0.00 │ │ └── TableScan │ │ ├── table: default.default.t2