fix(query): optimize numeric if branch selection

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<T: ValueType>(
+    flag: &Bitmap,
+    then_result: &Value<AnyType>,
+    else_result: &Value<AnyType>,
+    data_type: &DataType,
+    len: usize,
+) -> Result<Column> {
+    let then_result = then_result.try_downcast::<T>()?;
+    let else_result = else_result.try_downcast::<T>()?;
+    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<AnyType>,
+    else_result: &Value<AnyType>,
+    data_type: &DataType,
+    len: usize,
+) -> Result<Option<Column>> {
+    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::<NullableType<$type>>(
+                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)) {

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::<Vec<_>>();
+        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                    │               │               │               │         │

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",

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<T0=Int64><Boolean NULL, T0, T0>(CAST<Boolean>(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<T0=Decimal(15, 2)><Boolean NULL, T0, T0>(CAST<Boolean>(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<T0=Int64 NULL><Boolean NULL, T0, Boolean NULL, T0, T0>(CAST<Boolean>(cond_a AS Boolean NULL), CAST<Int64>(expr_a AS Int64 NULL), cond_b, CAST<Int64>(expr_b AS Int64 NULL), expr_else)