Skip to content

Commit 1d19bd9

Browse files
dantengskydantengsky
committed
fix(query): optimize numeric if branch selection
Extracted from commit 87eb181 after separating the reproducible benchmark baseline.
1 parent 58be093 commit 1d19bd9

3 files changed

Lines changed: 227 additions & 0 deletions

File tree

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

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,13 +40,28 @@ use crate::type_check::format_function_argument_mismatch_hint;
4040
use crate::type_check::get_simple_cast_function;
4141
use crate::types::BooleanType;
4242
use crate::types::DataType;
43+
use crate::types::Decimal64Type;
44+
use crate::types::Decimal128Type;
45+
use crate::types::Decimal256Type;
4346
use crate::types::DecimalColumn;
47+
use crate::types::DecimalDataKind;
4448
use crate::types::DecimalDataType;
4549
use crate::types::F32;
50+
use crate::types::Float32Type;
51+
use crate::types::Float64Type;
52+
use crate::types::Int8Type;
53+
use crate::types::Int16Type;
54+
use crate::types::Int32Type;
55+
use crate::types::Int64Type;
4656
use crate::types::NullableType;
57+
use crate::types::NumberDataType;
4758
use crate::types::NumberScalar;
4859
use crate::types::ReturnType;
4960
use crate::types::StringType;
61+
use crate::types::UInt8Type;
62+
use crate::types::UInt16Type;
63+
use crate::types::UInt32Type;
64+
use crate::types::UInt64Type;
5065
use crate::types::ValueType;
5166
use crate::types::VariantType;
5267
use crate::types::VectorDataType;
@@ -101,6 +116,131 @@ impl<'a> EvaluateOptions<'a> {
101116
}
102117
}
103118

119+
fn select_binary_if_typed<T: ValueType>(
120+
flag: &Bitmap,
121+
then_result: &Value<AnyType>,
122+
else_result: &Value<AnyType>,
123+
data_type: &DataType,
124+
len: usize,
125+
) -> Result<Column> {
126+
let then_result = then_result.try_downcast::<T>()?;
127+
let else_result = else_result.try_downcast::<T>()?;
128+
let mut output_builder = ColumnBuilder::with_capacity(data_type, len);
129+
130+
{
131+
let mut output = T::downcast_builder(&mut output_builder);
132+
match (&then_result, &else_result) {
133+
(Value::Scalar(then_value), Value::Scalar(else_value)) => {
134+
let then_value = T::to_scalar_ref(then_value);
135+
let else_value = T::to_scalar_ref(else_value);
136+
for take_then in flag.iter() {
137+
let value = if take_then {
138+
then_value.clone()
139+
} else {
140+
else_value.clone()
141+
};
142+
T::push_item_mut(&mut output, value);
143+
}
144+
}
145+
(Value::Scalar(then_value), Value::Column(else_column)) => {
146+
let then_value = T::to_scalar_ref(then_value);
147+
for (row, take_then) in flag.iter().enumerate() {
148+
let value = if take_then {
149+
then_value.clone()
150+
} else {
151+
unsafe { T::index_column_unchecked(else_column, row) }
152+
};
153+
T::push_item_mut(&mut output, value);
154+
}
155+
}
156+
(Value::Column(then_column), Value::Scalar(else_value)) => {
157+
let else_value = T::to_scalar_ref(else_value);
158+
for (row, take_then) in flag.iter().enumerate() {
159+
let value = if take_then {
160+
unsafe { T::index_column_unchecked(then_column, row) }
161+
} else {
162+
else_value.clone()
163+
};
164+
T::push_item_mut(&mut output, value);
165+
}
166+
}
167+
(Value::Column(then_column), Value::Column(else_column)) => {
168+
for (row, take_then) in flag.iter().enumerate() {
169+
let value = if take_then {
170+
unsafe { T::index_column_unchecked(then_column, row) }
171+
} else {
172+
unsafe { T::index_column_unchecked(else_column, row) }
173+
};
174+
T::push_item_mut(&mut output, value);
175+
}
176+
}
177+
}
178+
}
179+
180+
Ok(output_builder.build())
181+
}
182+
183+
fn select_binary_numeric_if(
184+
flag: &Bitmap,
185+
then_result: &Value<AnyType>,
186+
else_result: &Value<AnyType>,
187+
data_type: &DataType,
188+
len: usize,
189+
) -> Result<Option<Column>> {
190+
macro_rules! select_plain {
191+
($type:ty) => {
192+
select_binary_if_typed::<$type>(flag, then_result, else_result, data_type, len)
193+
};
194+
}
195+
macro_rules! select_nullable {
196+
($type:ty) => {
197+
select_binary_if_typed::<NullableType<$type>>(
198+
flag,
199+
then_result,
200+
else_result,
201+
data_type,
202+
len,
203+
)
204+
};
205+
}
206+
macro_rules! select_number {
207+
($number_type:expr, $select:ident) => {
208+
match $number_type {
209+
NumberDataType::UInt8 => $select!(UInt8Type),
210+
NumberDataType::UInt16 => $select!(UInt16Type),
211+
NumberDataType::UInt32 => $select!(UInt32Type),
212+
NumberDataType::UInt64 => $select!(UInt64Type),
213+
NumberDataType::Int8 => $select!(Int8Type),
214+
NumberDataType::Int16 => $select!(Int16Type),
215+
NumberDataType::Int32 => $select!(Int32Type),
216+
NumberDataType::Int64 => $select!(Int64Type),
217+
NumberDataType::Float32 => $select!(Float32Type),
218+
NumberDataType::Float64 => $select!(Float64Type),
219+
}
220+
};
221+
}
222+
223+
let column = match data_type {
224+
DataType::Number(number_type) => select_number!(number_type, select_plain)?,
225+
DataType::Decimal(size) => match size.data_kind() {
226+
DecimalDataKind::Decimal64 => select_plain!(Decimal64Type)?,
227+
DecimalDataKind::Decimal128 => select_plain!(Decimal128Type)?,
228+
DecimalDataKind::Decimal256 => select_plain!(Decimal256Type)?,
229+
},
230+
DataType::Nullable(inner) => match inner.as_ref() {
231+
DataType::Number(number_type) => select_number!(number_type, select_nullable)?,
232+
DataType::Decimal(size) => match size.data_kind() {
233+
DecimalDataKind::Decimal64 => select_nullable!(Decimal64Type)?,
234+
DecimalDataKind::Decimal128 => select_nullable!(Decimal128Type)?,
235+
DecimalDataKind::Decimal256 => select_nullable!(Decimal256Type)?,
236+
},
237+
_ => return Ok(None),
238+
},
239+
_ => return Ok(None),
240+
};
241+
Ok(Some(column))
242+
}
243+
104244
pub struct Evaluator<'a> {
105245
data_block: &'a DataBlock,
106246
func_ctx: &'a FunctionContext,
@@ -1641,6 +1781,13 @@ impl<'a> Evaluator<'a> {
16411781
.all_equal()
16421782
);
16431783

1784+
if let (Some(len), [flag], [then_result]) = (len, flags.as_slice(), results.as_slice())
1785+
&& let Some(column) =
1786+
select_binary_numeric_if(flag, then_result, &else_result, &generics[0], len)?
1787+
{
1788+
return Ok(Value::Column(column));
1789+
}
1790+
16441791
// Pick the results from the result branches depending on the condition.
16451792
let mut output_builder = ColumnBuilder::with_capacity(&generics[0], len.unwrap_or(1));
16461793
for row_idx in 0..(len.unwrap_or(1)) {

‎src/query/functions/tests/it/scalars/control.rs‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,36 @@ fn test_if(file: &mut impl Write) {
6363
Int64Type::from_data_with_validity(vec![5i64, 6, 7, 8], vec![true, true, false, false]),
6464
),
6565
]);
66+
// Cover the non-nullable numeric fast path for a binary `if`.
67+
run_ast(file, "if(cond_a, expr_true, expr_else)", &[
68+
(
69+
"cond_a",
70+
BooleanType::from_data(vec![true, false, true, false]),
71+
),
72+
("expr_true", Int64Type::from_data(vec![1i64, 2, 3, 4])),
73+
("expr_else", Int64Type::from_data(vec![5i64, 6, 7, 8])),
74+
]);
75+
// Cover the decimal fast path while preserving precision and scale.
76+
run_ast(file, "if(cond_a, expr_true, expr_else)", &[
77+
(
78+
"cond_a",
79+
BooleanType::from_data(vec![true, false, true, false]),
80+
),
81+
(
82+
"expr_true",
83+
Decimal64Type::from_data_with_size(
84+
[100i64, 200, 300, 400],
85+
Some(DecimalSize::new_unchecked(15, 2)),
86+
),
87+
),
88+
(
89+
"expr_else",
90+
Decimal64Type::from_data_with_size(
91+
[500i64, 600, 700, 800],
92+
Some(DecimalSize::new_unchecked(15, 2)),
93+
),
94+
),
95+
]);
6696
run_ast(file, "if(cond_a, expr_a, cond_b, expr_b, expr_else)", &[
6797
(
6898
"cond_a",

‎src/query/functions/tests/it/scalars/testdata/control.txt‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,56 @@ evaluation (internal):
137137
+-----------+--------------------------------------------------------------------------------+
138138

139139

140+
ast : if(cond_a, expr_true, expr_else)
141+
raw expr : if(cond_a::Boolean, expr_true::Int64, expr_else::Int64)
142+
checked expr : if<T0=Int64><Boolean NULL, T0, T0>(CAST<Boolean>(cond_a AS Boolean NULL), expr_true, expr_else)
143+
evaluation:
144+
+--------+---------------+-----------+-----------+---------+
145+
| | cond_a | expr_true | expr_else | Output |
146+
+--------+---------------+-----------+-----------+---------+
147+
| Type | Boolean | Int64 | Int64 | Int64 |
148+
| Domain | {FALSE, TRUE} | {1..=4} | {5..=8} | {1..=8} |
149+
| Row 0 | true | 1 | 5 | 1 |
150+
| Row 1 | false | 2 | 6 | 6 |
151+
| Row 2 | true | 3 | 7 | 3 |
152+
| Row 3 | false | 4 | 8 | 8 |
153+
+--------+---------------+-----------+-----------+---------+
154+
evaluation (internal):
155+
+-----------+-------------------------------+
156+
| Column | Data |
157+
+-----------+-------------------------------+
158+
| cond_a | Column(Boolean([0b____0101])) |
159+
| expr_true | Column(Int64([1, 2, 3, 4])) |
160+
| expr_else | Column(Int64([5, 6, 7, 8])) |
161+
| Output | Int64([1, 6, 3, 8]) |
162+
+-----------+-------------------------------+
163+
164+
165+
ast : if(cond_a, expr_true, expr_else)
166+
raw expr : if(cond_a::Boolean, expr_true::Decimal(15, 2), expr_else::Decimal(15, 2))
167+
checked expr : if<T0=Decimal(15, 2)><Boolean NULL, T0, T0>(CAST<Boolean>(cond_a AS Boolean NULL), expr_true, expr_else)
168+
evaluation:
169+
+--------+---------------+----------------+----------------+----------------+
170+
| | cond_a | expr_true | expr_else | Output |
171+
+--------+---------------+----------------+----------------+----------------+
172+
| Type | Boolean | Decimal(15, 2) | Decimal(15, 2) | Decimal(15, 2) |
173+
| Domain | {FALSE, TRUE} | {1.00..=4.00} | {5.00..=8.00} | {1.00..=8.00} |
174+
| Row 0 | true | 1.00 | 5.00 | 1.00 |
175+
| Row 1 | false | 2.00 | 6.00 | 6.00 |
176+
| Row 2 | true | 3.00 | 7.00 | 3.00 |
177+
| Row 3 | false | 4.00 | 8.00 | 8.00 |
178+
+--------+---------------+----------------+----------------+----------------+
179+
evaluation (internal):
180+
+-----------+---------------------------------------------+
181+
| Column | Data |
182+
+-----------+---------------------------------------------+
183+
| cond_a | Column(Boolean([0b____0101])) |
184+
| expr_true | Column(Decimal64([1.00, 2.00, 3.00, 4.00])) |
185+
| expr_else | Column(Decimal64([5.00, 6.00, 7.00, 8.00])) |
186+
| Output | Decimal64([1.00, 6.00, 3.00, 8.00]) |
187+
+-----------+---------------------------------------------+
188+
189+
140190
ast : if(cond_a, expr_a, cond_b, expr_b, expr_else)
141191
raw expr : if(cond_a::Boolean, expr_a::Int64, cond_b::Boolean NULL, expr_b::Int64, expr_else::Int64 NULL)
142192
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)

0 commit comments

Comments
 (0)