Skip to main content

reifydb_routine/function/time/
add.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use 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}