Skip to content

Commit c759865

Browse files
authored
Support arithmetic scalar operation with DictionaryArray (apache#5151)
* Support arithmetic dyn scalar * For review: removing unnecessary macro parameter, adding one more type coercion pattern * For review: modify comment
1 parent 224c682 commit c759865

3 files changed

Lines changed: 459 additions & 60 deletions

File tree

datafusion/expr/src/type_coercion/binary.rs

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -341,6 +341,15 @@ fn mathematics_numerical_coercion(
341341
(Null, dec_type @ Decimal128(_, _)) | (dec_type @ Decimal128(_, _), Null) => {
342342
Some(dec_type.clone())
343343
}
344+
(Dictionary(key_type, value_type), _) => {
345+
let value_type =
346+
mathematics_numerical_coercion(mathematics_op, value_type, rhs_type);
347+
value_type
348+
.map(|value_type| Dictionary(key_type.clone(), Box::new(value_type)))
349+
}
350+
(_, Dictionary(_, value_type)) => {
351+
mathematics_numerical_coercion(mathematics_op, lhs_type, value_type)
352+
}
344353
(Decimal128(_, _), Float32 | Float64) => Some(Float64),
345354
(Float32 | Float64, Decimal128(_, _)) => Some(Float64),
346355
(Decimal128(_, _), _) => {
@@ -439,6 +448,16 @@ fn both_numeric_or_null_and_numeric(lhs_type: &DataType, rhs_type: &DataType) ->
439448
match (lhs_type, rhs_type) {
440449
(_, DataType::Null) => is_numeric(lhs_type),
441450
(DataType::Null, _) => is_numeric(rhs_type),
451+
(
452+
DataType::Dictionary(_, lhs_value_type),
453+
DataType::Dictionary(_, rhs_value_type),
454+
) => is_numeric(lhs_value_type) && is_numeric(rhs_value_type),
455+
(DataType::Dictionary(_, value_type), _) => {
456+
is_numeric(value_type) && is_numeric(rhs_type)
457+
}
458+
(_, DataType::Dictionary(_, value_type)) => {
459+
is_numeric(lhs_type) && is_numeric(value_type)
460+
}
442461
_ => is_numeric(lhs_type) && is_numeric(rhs_type),
443462
}
444463
}

0 commit comments

Comments
 (0)