Skip to content

Commit 2d61b2f

Browse files
sundy-lisundy-li
authored
feat(query): support decimal to int (#12967)
* feat(query): support decimal to int * feat(query): test
1 parent 80d9e78 commit 2d61b2f

File tree

5 files changed

+415
-192
lines changed

5 files changed

+415
-192
lines changed

‎src/query/expression/src/types/decimal.rs

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ use common_exception::Result;
2222
use enum_as_inner::EnumAsInner;
2323
use ethnum::i256;
2424
use itertools::Itertools;
25+
use num_traits::NumCast;
2526
use num_traits::ToPrimitive;
2627
use serde::Deserialize;
2728
use serde::Serialize;
@@ -285,6 +286,8 @@ pub trait Decimal:
285286
fn to_float32(self, scale: u8) -> f32;
286287
fn to_float64(self, scale: u8) -> f64;
287288

289+
fn to_int<U: NumCast>(self, scale: u8) -> Option<U>;
290+
288291
fn try_downcast_column(column: &Column) -> Option<(Buffer<Self>, DecimalSize)>;
289292
fn try_downcast_builder<'a>(builder: &'a mut ColumnBuilder) -> Option<&'a mut Vec<Self>>;
290293

@@ -414,6 +417,11 @@ impl Decimal for i128 {
414417
self as f64 / div
415418
}
416419

420+
fn to_int<U: NumCast>(self, scale: u8) -> Option<U> {
421+
let div = 10i128.checked_pow(scale as u32)?;
422+
num_traits::cast(self / div)
423+
}
424+
417425
fn to_scalar(self, size: DecimalSize) -> DecimalScalar {
418426
DecimalScalar::Decimal128(self, size)
419427
}
@@ -563,6 +571,12 @@ impl Decimal for i256 {
563571
self.as_f64() / div
564572
}
565573

574+
fn to_int<U: NumCast>(self, scale: u8) -> Option<U> {
575+
let div = i256::from(10).checked_pow(scale as u32)?;
576+
let (h, l) = (self / div).into_words();
577+
if h > 0 { None } else { l.to_int(scale) }
578+
}
579+
566580
fn to_scalar(self, size: DecimalSize) -> DecimalScalar {
567581
DecimalScalar::Decimal256(self, size)
568582
}

‎src/query/functions/src/scalars/arithmetic.rs

Lines changed: 71 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@ use num_traits::AsPrimitive;
7474
use super::arithmetic_modulo::vectorize_modulo;
7575
use super::decimal::register_decimal_to_float32;
7676
use super::decimal::register_decimal_to_float64;
77+
use super::decimal::register_decimal_to_int;
7778
use crate::scalars::decimal::register_decimal_arithmetic;
7879

7980
pub fn register(registry: &mut FunctionRegistry) {
@@ -582,92 +583,96 @@ pub fn register_number_to_number(registry: &mut FunctionRegistry) {
582583
let name = format!("to_{dest_type}").to_lowercase();
583584
if src_type.can_lossless_cast_to(*dest_type) {
584585
registry.register_1_arg::<NumberType<SRC_TYPE>, NumberType<DEST_TYPE>, _, _>(
585-
&name,
586-
|_, domain| {
587-
let (domain, overflowing) = domain.overflow_cast();
588-
debug_assert!(!overflowing);
589-
FunctionDomain::Domain(domain)
590-
},
591-
|val, _| {
592-
val.as_()
593-
},
594-
);
586+
&name,
587+
|_, domain| {
588+
let (domain, overflowing) = domain.overflow_cast();
589+
debug_assert!(!overflowing);
590+
FunctionDomain::Domain(domain)
591+
},
592+
|val, _| {
593+
val.as_()
594+
},
595+
);
595596
} else {
596597
registry.register_passthrough_nullable_1_arg::<NumberType<SRC_TYPE>, NumberType<DEST_TYPE>, _, _>(
597-
&name,
598-
|_, domain| {
599-
let (domain, overflowing) = domain.overflow_cast();
600-
if overflowing {
601-
FunctionDomain::MayThrow
602-
} else {
603-
FunctionDomain::Domain(domain)
604-
}
605-
},
606-
vectorize_with_builder_1_arg::<NumberType<SRC_TYPE>, NumberType<DEST_TYPE>>(
607-
move |val, output, ctx| {
608-
match num_traits::cast::cast(val) {
609-
Some(val) => output.push(val),
610-
None => {
611-
ctx.set_error(output.len(),"number overflowed");
612-
output.push(DEST_TYPE::default());
598+
&name,
599+
|_, domain| {
600+
let (domain, overflowing) = domain.overflow_cast();
601+
if overflowing {
602+
FunctionDomain::MayThrow
603+
} else {
604+
FunctionDomain::Domain(domain)
605+
}
613606
},
614-
}
615-
}
616-
),
617-
);
607+
vectorize_with_builder_1_arg::<NumberType<SRC_TYPE>, NumberType<DEST_TYPE>>(
608+
move |val, output, ctx| {
609+
match num_traits::cast::cast(val) {
610+
Some(val) => output.push(val),
611+
None => {
612+
ctx.set_error(output.len(),"number overflowed");
613+
output.push(DEST_TYPE::default());
614+
},
615+
}
616+
}
617+
),
618+
);
618619
}
619620

620621
let name = format!("try_to_{dest_type}").to_lowercase();
621622
if src_type.can_lossless_cast_to(*dest_type) {
622623
registry.register_combine_nullable_1_arg::<NumberType<SRC_TYPE>, NumberType<DEST_TYPE>, _, _>(
623-
&name,
624-
|_, domain| {
625-
let (domain, overflowing) = domain.overflow_cast();
626-
debug_assert!(!overflowing);
627-
FunctionDomain::Domain(NullableDomain {
628-
has_null: false,
629-
value: Some(Box::new(
630-
domain,
631-
)),
632-
})
633-
},
634-
vectorize_1_arg::<NumberType<SRC_TYPE>, NullableType<NumberType<DEST_TYPE>>>(|val, _| {
635-
Some(val.as_())
636-
})
637-
);
624+
&name,
625+
|_, domain| {
626+
let (domain, overflowing) = domain.overflow_cast();
627+
debug_assert!(!overflowing);
628+
FunctionDomain::Domain(NullableDomain {
629+
has_null: false,
630+
value: Some(Box::new(
631+
domain,
632+
)),
633+
})
634+
},
635+
vectorize_1_arg::<NumberType<SRC_TYPE>, NullableType<NumberType<DEST_TYPE>>>(|val, _| {
636+
Some(val.as_())
637+
})
638+
);
638639
} else {
639640
registry.register_combine_nullable_1_arg::<NumberType<SRC_TYPE>, NumberType<DEST_TYPE>, _, _>(
640-
&name,
641-
|_, domain| {
642-
let (domain, overflowing) = domain.overflow_cast();
643-
FunctionDomain::Domain(NullableDomain {
644-
has_null: overflowing,
645-
value: Some(Box::new(
646-
domain,
647-
)),
648-
})
649-
},
650-
vectorize_with_builder_1_arg::<NumberType<SRC_TYPE>, NullableType<NumberType<DEST_TYPE>>>(
651-
|val, output, _| {
652-
if let Some(new_val) = num_traits::cast::cast(val) {
653-
output.push(new_val);
654-
} else {
655-
output.push_null();
656-
}
657-
}
658-
),
659-
);
641+
&name,
642+
|_, domain| {
643+
let (domain, overflowing) = domain.overflow_cast();
644+
FunctionDomain::Domain(NullableDomain {
645+
has_null: overflowing,
646+
value: Some(Box::new(
647+
domain,
648+
)),
649+
})
650+
},
651+
vectorize_with_builder_1_arg::<NumberType<SRC_TYPE>, NullableType<NumberType<DEST_TYPE>>>(
652+
|val, output, _| {
653+
if let Some(new_val) = num_traits::cast::cast(val) {
654+
output.push(new_val);
655+
} else {
656+
output.push_null();
657+
}
658+
}
659+
),
660+
);
660661
}
661662
}
662663
}),
663664
NumberClass::Decimal128 => {
664-
// todo(youngsofun): add decimal try_cast and decimal to int
665+
// todo(youngsofun): add decimal try_cast and decimal to int and float
665666
if matches!(dest_type, NumberDataType::Float32) {
666667
register_decimal_to_float32(registry);
667668
}
668669
if matches!(dest_type, NumberDataType::Float64) {
669670
register_decimal_to_float64(registry);
670671
}
672+
673+
with_number_mapped_type!(|DEST_TYPE| match dest_type {
674+
NumberDataType::DEST_TYPE => register_decimal_to_int::<DEST_TYPE>(registry),
675+
})
671676
}
672677
NumberClass::Decimal256 => {
673678
// already registered in Decimal128 branch

0 commit comments

Comments
 (0)