fix(query): fix months_between overflow error (#17403)

databendlabs · Feb 4, 2025 · 9997961 · 9997961
1 parent 855ecc8
commit 9997961
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 14 deletions.
diff --git a/src/query/expression/src/utils/date_helper.rs b/src/query/expression/src/utils/date_helper.rs
@@ -230,25 +230,17 @@ impl EvalMonthsImpl {
             .checked_add(SignedDuration::from_hours(date_b as i64 * 24))
             .unwrap();
 
-        let year_diff = date_a.year() - date_b.year();
-        let month_diff = date_a.month() as i16 - date_b.month() as i16;
+        let year_diff = (date_a.year() - date_b.year()) as i64;
+        let month_diff = date_a.month() as i64 - date_b.month() as i64;
 
         // Calculate total months difference
         let total_months_diff = year_diff * 12 + month_diff;
 
         // Determine if special case for fractional part applies
         let is_same_day_of_month = date_a.day() == date_b.day();
-        let are_both_end_of_month = date_a
-            .checked_add(SignedDuration::from_hours(24))
-            .unwrap()
-            .month()
-            != date_a.month()
-            && date_b
-                .checked_add(SignedDuration::from_hours(24))
-                .unwrap()
-                .month()
-                != date_b.month();
 
+        let are_both_end_of_month =
+            date_a.last_of_month() == date_a && date_b.last_of_month() == date_b;
         let day_fraction = if is_same_day_of_month || are_both_end_of_month {
             0.0
         } else {

diff --git a/src/query/functions/src/scalars/timestamp/src/datetime.rs b/src/query/functions/src/scalars/timestamp/src/datetime.rs
@@ -1227,9 +1227,11 @@ fn register_diff_functions(registry: &mut FunctionRegistry) {
             let rm = rhs.max;
             let rn = rhs.min;
 
+            let min = EvalMonthsImpl::months_between(ln, rm);
+            let max = EvalMonthsImpl::months_between(lm, rn);
             FunctionDomain::Domain(SimpleDomain::<F64> {
-                min: EvalMonthsImpl::months_between(ln, rm).into(),
-                max: EvalMonthsImpl::months_between(lm, rn).into(),
+                min: min.into(),
+                max: max.into(),
             })
         },
         vectorize_2_arg::<DateType, DateType, Float64Type>(|a, b, _ctx| {

diff --git a/tests/sqllogictests/suites/query/functions/02_0012_function_datetimes.test b/tests/sqllogictests/suites/query/functions/02_0012_function_datetimes.test
@@ -257,6 +257,21 @@ select to_date('2023-01-01') - 100 = to_date('2022-09-23')
 ----
 1
 
+statement ok
+create or replace table t(mine_date date, birth_date date);
+
+statement ok
+insert into t values('2022-01-01', '2022-01-01'),('2022-01-01', '2022-02-01'),('2022-02-01', '2022-01-01');
+
+query T
+SELECT max(MONTHS_BETWEEN(mine_date, birth_date))/12 as demo_age FROM t GROUP BY  mine_date order by demo_age;
+----
+0.0
+0.08333333333333333
+
+statement ok
+drop table if exists t;
+
 query FF
 SELECT
     MONTHS_BETWEEN('2019-03-15'::DATE,