reifydb_routine/function/duration/
trunc.rs1use reifydb_core::value::column::{ColumnWithName, buffer::ColumnBuffer, columns::Columns};
5use reifydb_type::value::{container::temporal::TemporalContainer, duration::Duration, r#type::Type};
6
7use crate::routine::{Function, FunctionKind, Routine, RoutineInfo, context::FunctionContext, error::RoutineError};
8
9pub struct DurationTrunc {
10 info: RoutineInfo,
11}
12
13impl Default for DurationTrunc {
14 fn default() -> Self {
15 Self::new()
16 }
17}
18
19impl DurationTrunc {
20 pub fn new() -> Self {
21 Self {
22 info: RoutineInfo::new("duration::trunc"),
23 }
24 }
25}
26
27impl<'a> Routine<FunctionContext<'a>> for DurationTrunc {
28 fn info(&self) -> &RoutineInfo {
29 &self.info
30 }
31
32 fn return_type(&self, _input_types: &[Type]) -> Type {
33 Type::Duration
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 dur_col = &args[0];
46 let prec_col = &args[1];
47
48 let (dur_data, dur_bv) = dur_col.unwrap_option();
49 let (prec_data, _) = prec_col.unwrap_option();
50
51 match (dur_data, prec_data) {
52 (
53 ColumnBuffer::Duration(dur_container),
54 ColumnBuffer::Utf8 {
55 container: prec_container,
56 ..
57 },
58 ) => {
59 let row_count = dur_data.len();
60 let mut container = TemporalContainer::with_capacity(row_count);
61
62 for i in 0..row_count {
63 match (dur_container.get(i), prec_container.is_defined(i)) {
64 (Some(dur), true) => {
65 let precision = prec_container.get(i).unwrap();
66 let months = dur.get_months();
67 let days = dur.get_days();
68 let nanos = dur.get_nanos();
69
70 let truncated = match precision {
71 "year" => Duration::new((months / 12) * 12, 0, 0)?,
72 "month" => Duration::new(months, 0, 0)?,
73 "day" => Duration::new(months, days, 0)?,
74 "hour" => Duration::new(
75 months,
76 days,
77 (nanos / 3_600_000_000_000) * 3_600_000_000_000,
78 )?,
79 "minute" => Duration::new(
80 months,
81 days,
82 (nanos / 60_000_000_000) * 60_000_000_000,
83 )?,
84 "second" => Duration::new(
85 months,
86 days,
87 (nanos / 1_000_000_000) * 1_000_000_000,
88 )?,
89 "millis" => Duration::new(
90 months,
91 days,
92 (nanos / 1_000_000) * 1_000_000,
93 )?,
94 other => {
95 return Err(
96 RoutineError::FunctionExecutionFailed {
97 function: ctx.fragment.clone(),
98 reason: format!(
99 "invalid precision: '{}'",
100 other
101 ),
102 },
103 );
104 }
105 };
106 container.push(truncated);
107 }
108 _ => container.push_default(),
109 }
110 }
111
112 let mut result_data = ColumnBuffer::Duration(container);
113 if let Some(bv) = dur_bv {
114 result_data = ColumnBuffer::Option {
115 inner: Box::new(result_data),
116 bitvec: bv.clone(),
117 };
118 }
119 Ok(Columns::new(vec![ColumnWithName::new(ctx.fragment.clone(), result_data)]))
120 }
121 (ColumnBuffer::Duration(_), other) => Err(RoutineError::FunctionInvalidArgumentType {
122 function: ctx.fragment.clone(),
123 argument_index: 1,
124 expected: vec![Type::Utf8],
125 actual: other.get_type(),
126 }),
127 (other, _) => Err(RoutineError::FunctionInvalidArgumentType {
128 function: ctx.fragment.clone(),
129 argument_index: 0,
130 expected: vec![Type::Duration],
131 actual: other.get_type(),
132 }),
133 }
134 }
135}
136
137impl Function for DurationTrunc {
138 fn kinds(&self) -> &[FunctionKind] {
139 &[FunctionKind::Scalar]
140 }
141}