reifydb_function/duration/
trunc.rs1use reifydb_core::value::column::data::ColumnData;
5use reifydb_type::value::{container::temporal::TemporalContainer, duration::Duration, r#type::Type};
6
7use crate::{
8 ScalarFunction, ScalarFunctionContext,
9 error::{ScalarFunctionError, ScalarFunctionResult},
10 propagate_options,
11};
12
13pub struct DurationTrunc;
14
15impl DurationTrunc {
16 pub fn new() -> Self {
17 Self
18 }
19}
20
21impl ScalarFunction for DurationTrunc {
22 fn scalar(&self, ctx: ScalarFunctionContext) -> ScalarFunctionResult<ColumnData> {
23 if let Some(result) = propagate_options(self, &ctx) {
24 return result;
25 }
26 let columns = ctx.columns;
27 let row_count = ctx.row_count;
28
29 if columns.len() != 2 {
30 return Err(ScalarFunctionError::ArityMismatch {
31 function: ctx.fragment.clone(),
32 expected: 2,
33 actual: columns.len(),
34 });
35 }
36
37 let dur_col = columns.get(0).unwrap();
38 let prec_col = columns.get(1).unwrap();
39
40 match (dur_col.data(), prec_col.data()) {
41 (
42 ColumnData::Duration(dur_container),
43 ColumnData::Utf8 {
44 container: prec_container,
45 ..
46 },
47 ) => {
48 let mut container = TemporalContainer::with_capacity(row_count);
49
50 for i in 0..row_count {
51 match (dur_container.get(i), prec_container.is_defined(i)) {
52 (Some(dur), true) => {
53 let precision = &prec_container[i];
54 let months = dur.get_months();
55 let days = dur.get_days();
56 let nanos = dur.get_nanos();
57
58 let truncated = match precision.as_str() {
59 "year" => Duration::new((months / 12) * 12, 0, 0),
60 "month" => Duration::new(months, 0, 0),
61 "day" => Duration::new(months, days, 0),
62 "hour" => Duration::new(
63 months,
64 days,
65 (nanos / 3_600_000_000_000) * 3_600_000_000_000,
66 ),
67 "minute" => Duration::new(
68 months,
69 days,
70 (nanos / 60_000_000_000) * 60_000_000_000,
71 ),
72 "second" => Duration::new(
73 months,
74 days,
75 (nanos / 1_000_000_000) * 1_000_000_000,
76 ),
77 "millis" => Duration::new(
78 months,
79 days,
80 (nanos / 1_000_000) * 1_000_000,
81 ),
82 other => {
83 return Err(
84 ScalarFunctionError::ExecutionFailed {
85 function: ctx.fragment.clone(),
86 reason: format!(
87 "invalid precision: '{}'",
88 other
89 ),
90 },
91 );
92 }
93 };
94 container.push(truncated);
95 }
96 _ => container.push_default(),
97 }
98 }
99
100 Ok(ColumnData::Duration(container))
101 }
102 (ColumnData::Duration(_), other) => Err(ScalarFunctionError::InvalidArgumentType {
103 function: ctx.fragment.clone(),
104 argument_index: 1,
105 expected: vec![Type::Utf8],
106 actual: other.get_type(),
107 }),
108 (other, _) => Err(ScalarFunctionError::InvalidArgumentType {
109 function: ctx.fragment.clone(),
110 argument_index: 0,
111 expected: vec![Type::Duration],
112 actual: other.get_type(),
113 }),
114 }
115 }
116
117 fn return_type(&self, _input_types: &[Type]) -> Type {
118 Type::Duration
119 }
120}