diff --git a/src/query/expression/src/evaluator.rs b/src/query/expression/src/evaluator.rs index ce9ef55f382d9..45af35f3fbd34 100644 --- a/src/query/expression/src/evaluator.rs +++ b/src/query/expression/src/evaluator.rs @@ -40,13 +40,28 @@ use crate::type_check::format_function_argument_mismatch_hint; use crate::type_check::get_simple_cast_function; use crate::types::BooleanType; use crate::types::DataType; +use crate::types::Decimal64Type; +use crate::types::Decimal128Type; +use crate::types::Decimal256Type; use crate::types::DecimalColumn; +use crate::types::DecimalDataKind; use crate::types::DecimalDataType; use crate::types::F32; +use crate::types::Float32Type; +use crate::types::Float64Type; +use crate::types::Int8Type; +use crate::types::Int16Type; +use crate::types::Int32Type; +use crate::types::Int64Type; use crate::types::NullableType; +use crate::types::NumberDataType; use crate::types::NumberScalar; use crate::types::ReturnType; use crate::types::StringType; +use crate::types::UInt8Type; +use crate::types::UInt16Type; +use crate::types::UInt32Type; +use crate::types::UInt64Type; use crate::types::ValueType; use crate::types::VariantType; use crate::types::VectorDataType; @@ -101,6 +116,131 @@ impl<'a> EvaluateOptions<'a> { } } +fn select_binary_if_typed( + flag: &Bitmap, + then_result: &Value, + else_result: &Value, + data_type: &DataType, + len: usize, +) -> Result { + let then_result = then_result.try_downcast::()?; + let else_result = else_result.try_downcast::()?; + let mut output_builder = ColumnBuilder::with_capacity(data_type, len); + + { + let mut output = T::downcast_builder(&mut output_builder); + match (&then_result, &else_result) { + (Value::Scalar(then_value), Value::Scalar(else_value)) => { + let then_value = T::to_scalar_ref(then_value); + let else_value = T::to_scalar_ref(else_value); + for take_then in flag.iter() { + let value = if take_then { + then_value.clone() + } else { + else_value.clone() + }; + T::push_item_mut(&mut output, value); + } + } + (Value::Scalar(then_value), Value::Column(else_column)) => { + let then_value = T::to_scalar_ref(then_value); + for (row, take_then) in flag.iter().enumerate() { + let value = if take_then { + then_value.clone() + } else { + unsafe { T::index_column_unchecked(else_column, row) } + }; + T::push_item_mut(&mut output, value); + } + } + (Value::Column(then_column), Value::Scalar(else_value)) => { + let else_value = T::to_scalar_ref(else_value); + for (row, take_then) in flag.iter().enumerate() { + let value = if take_then { + unsafe { T::index_column_unchecked(then_column, row) } + } else { + else_value.clone() + }; + T::push_item_mut(&mut output, value); + } + } + (Value::Column(then_column), Value::Column(else_column)) => { + for (row, take_then) in flag.iter().enumerate() { + let value = if take_then { + unsafe { T::index_column_unchecked(then_column, row) } + } else { + unsafe { T::index_column_unchecked(else_column, row) } + }; + T::push_item_mut(&mut output, value); + } + } + } + } + + Ok(output_builder.build()) +} + +fn select_binary_numeric_if( + flag: &Bitmap, + then_result: &Value, + else_result: &Value, + data_type: &DataType, + len: usize, +) -> Result> { + macro_rules! select_plain { + ($type:ty) => { + select_binary_if_typed::<$type>(flag, then_result, else_result, data_type, len) + }; + } + macro_rules! select_nullable { + ($type:ty) => { + select_binary_if_typed::>( + flag, + then_result, + else_result, + data_type, + len, + ) + }; + } + macro_rules! select_number { + ($number_type:expr, $select:ident) => { + match $number_type { + NumberDataType::UInt8 => $select!(UInt8Type), + NumberDataType::UInt16 => $select!(UInt16Type), + NumberDataType::UInt32 => $select!(UInt32Type), + NumberDataType::UInt64 => $select!(UInt64Type), + NumberDataType::Int8 => $select!(Int8Type), + NumberDataType::Int16 => $select!(Int16Type), + NumberDataType::Int32 => $select!(Int32Type), + NumberDataType::Int64 => $select!(Int64Type), + NumberDataType::Float32 => $select!(Float32Type), + NumberDataType::Float64 => $select!(Float64Type), + } + }; + } + + let column = match data_type { + DataType::Number(number_type) => select_number!(number_type, select_plain)?, + DataType::Decimal(size) => match size.data_kind() { + DecimalDataKind::Decimal64 => select_plain!(Decimal64Type)?, + DecimalDataKind::Decimal128 => select_plain!(Decimal128Type)?, + DecimalDataKind::Decimal256 => select_plain!(Decimal256Type)?, + }, + DataType::Nullable(inner) => match inner.as_ref() { + DataType::Number(number_type) => select_number!(number_type, select_nullable)?, + DataType::Decimal(size) => match size.data_kind() { + DecimalDataKind::Decimal64 => select_nullable!(Decimal64Type)?, + DecimalDataKind::Decimal128 => select_nullable!(Decimal128Type)?, + DecimalDataKind::Decimal256 => select_nullable!(Decimal256Type)?, + }, + _ => return Ok(None), + }, + _ => return Ok(None), + }; + Ok(Some(column)) +} + pub struct Evaluator<'a> { data_block: &'a DataBlock, func_ctx: &'a FunctionContext, @@ -1641,6 +1781,13 @@ impl<'a> Evaluator<'a> { .all_equal() ); + if let (Some(len), [flag], [then_result]) = (len, flags.as_slice(), results.as_slice()) + && let Some(column) = + select_binary_numeric_if(flag, then_result, &else_result, &generics[0], len)? + { + return Ok(Value::Column(column)); + } + // Pick the results from the result branches depending on the condition. let mut output_builder = ColumnBuilder::with_capacity(&generics[0], len.unwrap_or(1)); for row_idx in 0..(len.unwrap_or(1)) { diff --git a/src/query/functions/benches/bench.rs b/src/query/functions/benches/bench.rs index 37109aa2b7097..29cf741b1ae93 100644 --- a/src/query/functions/benches/bench.rs +++ b/src/query/functions/benches/bench.rs @@ -16,6 +16,59 @@ fn main() { divan::main(); } +#[divan::bench_group(sample_count = 30)] +mod numeric_if { + use databend_common_expression::BlockEntry; + use databend_common_expression::DataBlock; + use databend_common_expression::Evaluator; + use databend_common_expression::FromData; + use databend_common_expression::FunctionContext; + use databend_common_expression::type_check; + use databend_common_expression::types::DataType; + use databend_common_expression::types::NumberDataType; + use databend_common_expression::types::number::Int64Type; + use databend_common_expression_test_support as parser; + use databend_common_functions::BUILTIN_FUNCTIONS; + use divan::counter::ItemsCount; + + fn bench_expr(bencher: divan::Bencher, sql: &str, rows: usize) { + let raw_expr = parser::parse_raw_expr( + sql, + &[("x", DataType::Number(NumberDataType::Int64))], + &BUILTIN_FUNCTIONS, + ); + let expr = type_check::check(&raw_expr, &BUILTIN_FUNCTIONS).unwrap(); + let values = (0..rows).map(|value| value as i64).collect::>(); + let block = DataBlock::new(vec![BlockEntry::Column(Int64Type::from_data(values))], rows); + let func_ctx = FunctionContext::default(); + let evaluator = Evaluator::new(&block, &func_ctx, &BUILTIN_FUNCTIONS); + + bencher + .counter(ItemsCount::new(rows)) + .bench(|| evaluator.run(&expr).unwrap()); + } + + // 65,536 is the default max_block_size; 1,048,576 is a stress size. + #[divan::bench(args = [65_536, 1_048_576])] + fn if_contains(bencher: divan::Bencher, rows: usize) { + bench_expr( + bencher, + "if(contains([1, 3, 5, 7, 11, 13, 17, 19], x), 0, x % 1000003)", + rows, + ); + } + + #[divan::bench(args = [65_536, 1_048_576])] + fn if_modulo(bencher: divan::Bencher, rows: usize) { + bench_expr(bencher, "if(x % 1000003 = 0, 0, x % 1000003)", rows); + } + + #[divan::bench(args = [65_536, 1_048_576])] + fn if_boolean(bencher: divan::Bencher, rows: usize) { + bench_expr(bencher, "if(x % 2 = 0, x > 10, x < 10)", rows); + } +} + // bench fastest │ slowest │ median │ mean │ samples │ iters // ╰─ dummy │ │ │ │ │ // ├─ check │ │ │ │ │ diff --git a/src/query/functions/tests/it/scalars/control.rs b/src/query/functions/tests/it/scalars/control.rs index 6b626d578cb5a..a3c419f70ddc5 100644 --- a/src/query/functions/tests/it/scalars/control.rs +++ b/src/query/functions/tests/it/scalars/control.rs @@ -63,6 +63,36 @@ fn test_if(file: &mut impl Write) { Int64Type::from_data_with_validity(vec![5i64, 6, 7, 8], vec![true, true, false, false]), ), ]); + // Cover the non-nullable numeric fast path for a binary `if`. + run_ast(file, "if(cond_a, expr_true, expr_else)", &[ + ( + "cond_a", + BooleanType::from_data(vec![true, false, true, false]), + ), + ("expr_true", Int64Type::from_data(vec![1i64, 2, 3, 4])), + ("expr_else", Int64Type::from_data(vec![5i64, 6, 7, 8])), + ]); + // Cover the decimal fast path while preserving precision and scale. + run_ast(file, "if(cond_a, expr_true, expr_else)", &[ + ( + "cond_a", + BooleanType::from_data(vec![true, false, true, false]), + ), + ( + "expr_true", + Decimal64Type::from_data_with_size( + [100i64, 200, 300, 400], + Some(DecimalSize::new_unchecked(15, 2)), + ), + ), + ( + "expr_else", + Decimal64Type::from_data_with_size( + [500i64, 600, 700, 800], + Some(DecimalSize::new_unchecked(15, 2)), + ), + ), + ]); run_ast(file, "if(cond_a, expr_a, cond_b, expr_b, expr_else)", &[ ( "cond_a", diff --git a/src/query/functions/tests/it/scalars/testdata/control.txt b/src/query/functions/tests/it/scalars/testdata/control.txt index 698deda407045..87b5b0f841dc8 100644 --- a/src/query/functions/tests/it/scalars/testdata/control.txt +++ b/src/query/functions/tests/it/scalars/testdata/control.txt @@ -137,6 +137,56 @@ evaluation (internal): +-----------+--------------------------------------------------------------------------------+ +ast : if(cond_a, expr_true, expr_else) +raw expr : if(cond_a::Boolean, expr_true::Int64, expr_else::Int64) +checked expr : if(CAST(cond_a AS Boolean NULL), expr_true, expr_else) +evaluation: ++--------+---------------+-----------+-----------+---------+ +| | cond_a | expr_true | expr_else | Output | ++--------+---------------+-----------+-----------+---------+ +| Type | Boolean | Int64 | Int64 | Int64 | +| Domain | {FALSE, TRUE} | {1..=4} | {5..=8} | {1..=8} | +| Row 0 | true | 1 | 5 | 1 | +| Row 1 | false | 2 | 6 | 6 | +| Row 2 | true | 3 | 7 | 3 | +| Row 3 | false | 4 | 8 | 8 | ++--------+---------------+-----------+-----------+---------+ +evaluation (internal): ++-----------+-------------------------------+ +| Column | Data | ++-----------+-------------------------------+ +| cond_a | Column(Boolean([0b____0101])) | +| expr_true | Column(Int64([1, 2, 3, 4])) | +| expr_else | Column(Int64([5, 6, 7, 8])) | +| Output | Int64([1, 6, 3, 8]) | ++-----------+-------------------------------+ + + +ast : if(cond_a, expr_true, expr_else) +raw expr : if(cond_a::Boolean, expr_true::Decimal(15, 2), expr_else::Decimal(15, 2)) +checked expr : if(CAST(cond_a AS Boolean NULL), expr_true, expr_else) +evaluation: ++--------+---------------+----------------+----------------+----------------+ +| | cond_a | expr_true | expr_else | Output | ++--------+---------------+----------------+----------------+----------------+ +| Type | Boolean | Decimal(15, 2) | Decimal(15, 2) | Decimal(15, 2) | +| Domain | {FALSE, TRUE} | {1.00..=4.00} | {5.00..=8.00} | {1.00..=8.00} | +| Row 0 | true | 1.00 | 5.00 | 1.00 | +| Row 1 | false | 2.00 | 6.00 | 6.00 | +| Row 2 | true | 3.00 | 7.00 | 3.00 | +| Row 3 | false | 4.00 | 8.00 | 8.00 | ++--------+---------------+----------------+----------------+----------------+ +evaluation (internal): ++-----------+---------------------------------------------+ +| Column | Data | ++-----------+---------------------------------------------+ +| cond_a | Column(Boolean([0b____0101])) | +| expr_true | Column(Decimal64([1.00, 2.00, 3.00, 4.00])) | +| expr_else | Column(Decimal64([5.00, 6.00, 7.00, 8.00])) | +| Output | Decimal64([1.00, 6.00, 3.00, 8.00]) | ++-----------+---------------------------------------------+ + + ast : if(cond_a, expr_a, cond_b, expr_b, expr_else) raw expr : if(cond_a::Boolean, expr_a::Int64, cond_b::Boolean NULL, expr_b::Int64, expr_else::Int64 NULL) checked expr : if(CAST(cond_a AS Boolean NULL), CAST(expr_a AS Int64 NULL), cond_b, CAST(expr_b AS Int64 NULL), expr_else)