reifydb_function/date/
add.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2// Copyright (c) 2025 ReifyDB
3
4use reifydb_core::value::column::data::ColumnData;
5use reifydb_type::value::{container::temporal::TemporalContainer, date::Date, r#type::Type};
6
7use crate::{
8	ScalarFunction, ScalarFunctionContext,
9	error::{ScalarFunctionError, ScalarFunctionResult},
10	propagate_options,
11};
12
13pub struct DateAdd;
14
15impl DateAdd {
16	pub fn new() -> Self {
17		Self
18	}
19}
20
21impl ScalarFunction for DateAdd {
22	fn scalar(&self, ctx: ScalarFunctionContext) -> ScalarFunctionResult<ColumnData> {
23		if let Some(result) = propagate_options(self, &ctx) {
24			return result;
25		}
26
27		let columns = ctx.columns;
28		let row_count = ctx.row_count;
29
30		if columns.len() != 2 {
31			return Err(ScalarFunctionError::ArityMismatch {
32				function: ctx.fragment.clone(),
33				expected: 2,
34				actual: columns.len(),
35			});
36		}
37
38		let date_col = columns.get(0).unwrap();
39		let dur_col = columns.get(1).unwrap();
40
41		match (date_col.data(), dur_col.data()) {
42			(ColumnData::Date(date_container), ColumnData::Duration(dur_container)) => {
43				let mut container = TemporalContainer::with_capacity(row_count);
44
45				for i in 0..row_count {
46					match (date_container.get(i), dur_container.get(i)) {
47						(Some(date), Some(dur)) => {
48							let mut year = date.year();
49							let mut month = date.month() as i32;
50							let mut day = date.day();
51
52							// Add months component
53							let total_months = month + dur.get_months();
54							year += (total_months - 1).div_euclid(12);
55							month = (total_months - 1).rem_euclid(12) + 1;
56
57							// Clamp day to valid range for the new month
58							let max_day = days_in_month(year, month as u32);
59							if day > max_day {
60								day = max_day;
61							}
62
63							// Convert to days_since_epoch and add days component
64							if let Some(base) = Date::new(year, month as u32, day) {
65								let total_days = base.to_days_since_epoch()
66									+ dur.get_days() + (dur.get_nanos()
67									/ 86_400_000_000_000)
68									as i32;
69								match Date::from_days_since_epoch(total_days) {
70									Some(result) => container.push(result),
71									None => container.push_default(),
72								}
73							} else {
74								container.push_default();
75							}
76						}
77						_ => container.push_default(),
78					}
79				}
80
81				Ok(ColumnData::Date(container))
82			}
83			(ColumnData::Date(_), other) => Err(ScalarFunctionError::InvalidArgumentType {
84				function: ctx.fragment.clone(),
85				argument_index: 1,
86				expected: vec![Type::Duration],
87				actual: other.get_type(),
88			}),
89			(other, _) => Err(ScalarFunctionError::InvalidArgumentType {
90				function: ctx.fragment.clone(),
91				argument_index: 0,
92				expected: vec![Type::Date],
93				actual: other.get_type(),
94			}),
95		}
96	}
97
98	fn return_type(&self, _input_types: &[Type]) -> Type {
99		Type::Date
100	}
101}
102
103fn days_in_month(year: i32, month: u32) -> u32 {
104	match month {
105		1 | 3 | 5 | 7 | 8 | 10 | 12 => 31,
106		4 | 6 | 9 | 11 => 30,
107		2 => {
108			if (year % 4 == 0 && year % 100 != 0) || (year % 400 == 0) {
109				29
110			} else {
111				28
112			}
113		}
114		_ => 0,
115	}
116}
reifydb_function/date/add.rs

reifydb_function/date/
add.rs
