reifydb_routine/function/datetime/
add.rs1use reifydb_core::value::column::{Column, columns::Columns, data::ColumnData};
5use reifydb_type::value::{container::temporal::TemporalContainer, r#type::Type};
6
7use crate::function::{Function, FunctionCapability, FunctionContext, FunctionInfo, error::FunctionError};
8
9pub struct DateTimeAdd {
10 info: FunctionInfo,
11}
12
13impl Default for DateTimeAdd {
14 fn default() -> Self {
15 Self::new()
16 }
17}
18
19impl DateTimeAdd {
20 pub fn new() -> Self {
21 Self {
22 info: FunctionInfo::new("datetime::add"),
23 }
24 }
25}
26
27impl Function for DateTimeAdd {
28 fn info(&self) -> &FunctionInfo {
29 &self.info
30 }
31
32 fn capabilities(&self) -> &[FunctionCapability] {
33 &[FunctionCapability::Scalar]
34 }
35
36 fn return_type(&self, _input_types: &[Type]) -> Type {
37 Type::DateTime
38 }
39
40 fn execute(&self, ctx: &FunctionContext, args: &Columns) -> Result<Columns, FunctionError> {
41 if args.len() != 2 {
42 return Err(FunctionError::ArityMismatch {
43 function: ctx.fragment.clone(),
44 expected: 2,
45 actual: args.len(),
46 });
47 }
48
49 let dt_col = &args[0];
50 let dur_col = &args[1];
51 let (dt_data, dt_bitvec) = dt_col.data().unwrap_option();
52 let (dur_data, dur_bitvec) = dur_col.data().unwrap_option();
53 let row_count = dt_data.len();
54
55 let result_data = match (dt_data, dur_data) {
56 (ColumnData::DateTime(dt_container), ColumnData::Duration(dur_container)) => {
57 let mut container = TemporalContainer::with_capacity(row_count);
58
59 for i in 0..row_count {
60 match (dt_container.get(i), dur_container.get(i)) {
61 (Some(dt), Some(dur)) => match dt.add_duration(dur) {
62 Ok(result) => container.push(result),
63 Err(err) => {
64 return Err(FunctionError::ExecutionFailed {
65 function: ctx.fragment.clone(),
66 reason: format!("{}", err),
67 });
68 }
69 },
70 _ => container.push_default(),
71 }
72 }
73
74 ColumnData::DateTime(container)
75 }
76 (ColumnData::DateTime(_), other) => {
77 return Err(FunctionError::InvalidArgumentType {
78 function: ctx.fragment.clone(),
79 argument_index: 1,
80 expected: vec![Type::Duration],
81 actual: other.get_type(),
82 });
83 }
84 (other, _) => {
85 return Err(FunctionError::InvalidArgumentType {
86 function: ctx.fragment.clone(),
87 argument_index: 0,
88 expected: vec![Type::DateTime],
89 actual: other.get_type(),
90 });
91 }
92 };
93
94 let final_data = match (dt_bitvec, dur_bitvec) {
95 (Some(bv), _) | (_, Some(bv)) => ColumnData::Option {
96 inner: Box::new(result_data),
97 bitvec: bv.clone(),
98 },
99 _ => result_data,
100 };
101
102 Ok(Columns::new(vec![Column::new(ctx.fragment.clone(), final_data)]))
103 }
104}