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