Skip to content

Commit 9e6cafc

Browse files
committed
Fix constant decimal array overflow for decimals bigger than i128
Signed-off-by: Robert Kruszewski <github@robertk.io>
1 parent d42e7c9 commit 9e6cafc

1 file changed

Lines changed: 59 additions & 8 deletions

File tree

‎vortex-array/src/aggregate_fn/fns/sum/constant.rs‎

Lines changed: 59 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use vortex_error::vortex_err;
88
use crate::dtype::DType;
99
use crate::dtype::Nullability;
1010
use crate::dtype::PType;
11+
use crate::dtype::i256;
1112
use crate::scalar::DecimalValue;
1213
use crate::scalar::Scalar;
1314

@@ -70,14 +71,19 @@ pub(crate) fn multiply_constant(
7071
.as_decimal()
7172
.decimal_value()
7273
.ok_or_else(|| vortex_err!("Expected non-null decimal scalar for sum"))?;
73-
let len_decimal = DecimalValue::from(len as i128);
74-
match val.checked_mul(&len_decimal) {
75-
Some(product) => {
76-
let ret_decimal = *return_dtype
77-
.as_decimal_opt()
78-
.ok_or_else(|| vortex_err!("Expected decimal return dtype"))?;
79-
Scalar::decimal(product, ret_decimal, Nullability::Nullable)
80-
}
74+
let ret_decimal = *return_dtype
75+
.as_decimal_opt()
76+
.ok_or_else(|| vortex_err!("Expected decimal return dtype"))?;
77+
78+
// Multiply in i256, the widest storage type, because the operands can be narrower
79+
// than the return type. Then check the product against the return precision.
80+
let val = DecimalValue::I256(val.as_i256());
81+
let len_decimal = DecimalValue::I256(i256::from_i128(len as i128));
82+
match val
83+
.checked_mul(&len_decimal)
84+
.and_then(|product| product.normalize(ret_decimal))
85+
{
86+
Some(product) => Scalar::decimal(product, ret_decimal, Nullability::Nullable),
8187
None => Scalar::null(return_dtype.as_nullable()),
8288
}
8389
}
@@ -216,4 +222,49 @@ mod tests {
216222
);
217223
Ok(())
218224
}
225+
226+
/// Regression test for <https://github.com/vortex-data/vortex/issues/10447>: an `I128`
227+
/// value times the length overflows `i128`, but fits the wider return type.
228+
#[test]
229+
fn sum_constant_decimal_i128_value_overflows_i128() -> VortexResult<()> {
230+
let value: i128 = 112_285_607_173_948_541_047_296_740_235_646_205_996;
231+
let array = ConstantArray::new(
232+
Scalar::decimal(
233+
DecimalValue::I128(value),
234+
DecimalDType::new(39, 39),
235+
Nullability::NonNullable,
236+
),
237+
2,
238+
)
239+
.into_array();
240+
241+
let result = sum(&array, &mut array_session().create_execution_ctx())?;
242+
243+
assert_eq!(
244+
result.as_decimal().decimal_value(),
245+
Some(DecimalValue::I256(
246+
i256::from_i128(value) + i256::from_i128(value)
247+
))
248+
);
249+
Ok(())
250+
}
251+
252+
#[test]
253+
fn sum_constant_decimal_exceeds_return_precision() -> VortexResult<()> {
254+
let max_value = i256::from_i128(10).checked_pow(76).unwrap() - i256::from_i128(1);
255+
let array = ConstantArray::new(
256+
Scalar::decimal(
257+
DecimalValue::I256(max_value),
258+
DecimalDType::new(76, 0),
259+
Nullability::NonNullable,
260+
),
261+
2,
262+
)
263+
.into_array();
264+
265+
let result = sum(&array, &mut array_session().create_execution_ctx())?;
266+
267+
assert!(result.is_null());
268+
Ok(())
269+
}
219270
}

0 commit comments

Comments
 (0)