reifydb_routine/function/date/
subtract.rs1use reifydb_core::value::column::{ColumnWithName, buffer::ColumnBuffer, columns::Columns};
5use reifydb_type::value::{container::temporal::TemporalContainer, date::Date, r#type::Type};
6
7use crate::routine::{Function, FunctionKind, Routine, RoutineInfo, context::FunctionContext, error::RoutineError};
8
9pub struct DateSubtract {
10 info: RoutineInfo,
11}
12
13impl Default for DateSubtract {
14 fn default() -> Self {
15 Self::new()
16 }
17}
18
19impl DateSubtract {
20 pub fn new() -> Self {
21 Self {
22 info: RoutineInfo::new("date::subtract"),
23 }
24 }
25}
26
27impl<'a> Routine<FunctionContext<'a>> for DateSubtract {
28 fn info(&self) -> &RoutineInfo {
29 &self.info
30 }
31
32 fn return_type(&self, _input_types: &[Type]) -> Type {
33 Type::Date
34 }
35
36 fn execute(&self, ctx: &mut FunctionContext<'a>, args: &Columns) -> Result<Columns, RoutineError> {
37 if args.len() != 2 {
38 return Err(RoutineError::FunctionArityMismatch {
39 function: ctx.fragment.clone(),
40 expected: 2,
41 actual: args.len(),
42 });
43 }
44
45 let date_col = &args[0];
46 let dur_col = &args[1];
47 let (date_data, date_bitvec) = date_col.unwrap_option();
48 let (dur_data, dur_bitvec) = dur_col.unwrap_option();
49 let row_count = date_data.len();
50
51 let result_data = match (date_data, dur_data) {
52 (ColumnBuffer::Date(date_container), ColumnBuffer::Duration(dur_container)) => {
53 let mut container = TemporalContainer::with_capacity(row_count);
54
55 for i in 0..row_count {
56 match (date_container.get(i), dur_container.get(i)) {
57 (Some(date), Some(dur)) => {
58 let mut year = date.year();
59 let mut month = date.month() as i32;
60 let mut day = date.day();
61
62 let total_months = month - dur.get_months();
64 year += (total_months - 1).div_euclid(12);
65 month = (total_months - 1).rem_euclid(12) + 1;
66
67 let max_day = days_in_month(year, month as u32);
69 if day > max_day {
70 day = max_day;
71 }
72
73 if let Some(base) = Date::new(year, month as u32, day) {
75 let total_days = base.to_days_since_epoch()
76 - dur.get_days() - (dur.get_nanos()
77 / 86_400_000_000_000)
78 as i32;
79 match Date::from_days_since_epoch(total_days) {
80 Some(result) => container.push(result),
81 None => container.push_default(),
82 }
83 } else {
84 container.push_default();
85 }
86 }
87 _ => container.push_default(),
88 }
89 }
90
91 ColumnBuffer::Date(container)
92 }
93 (ColumnBuffer::Date(_), other) => {
94 return Err(RoutineError::FunctionInvalidArgumentType {
95 function: ctx.fragment.clone(),
96 argument_index: 1,
97 expected: vec![Type::Duration],
98 actual: other.get_type(),
99 });
100 }
101 (other, _) => {
102 return Err(RoutineError::FunctionInvalidArgumentType {
103 function: ctx.fragment.clone(),
104 argument_index: 0,
105 expected: vec![Type::Date],
106 actual: other.get_type(),
107 });
108 }
109 };
110
111 let final_data = match (date_bitvec, dur_bitvec) {
112 (Some(bv), _) | (_, Some(bv)) => ColumnBuffer::Option {
113 inner: Box::new(result_data),
114 bitvec: bv.clone(),
115 },
116 _ => result_data,
117 };
118
119 Ok(Columns::new(vec![ColumnWithName::new(ctx.fragment.clone(), final_data)]))
120 }
121}
122
123impl Function for DateSubtract {
124 fn kinds(&self) -> &[FunctionKind] {
125 &[FunctionKind::Scalar]
126 }
127}
128
129fn days_in_month(year: i32, month: u32) -> u32 {
130 match month {
131 1 | 3 | 5 | 7 | 8 | 10 | 12 => 31,
132 4 | 6 | 9 | 11 => 30,
133 2 => {
134 if (year % 4 == 0 && year % 100 != 0) || (year % 400 == 0) {
135 29
136 } else {
137 28
138 }
139 }
140 _ => 0,
141 }
142}