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