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::{
8 ScalarFunction, ScalarFunctionContext,
9 error::{ScalarFunctionError, ScalarFunctionResult},
10 propagate_options,
11};
12
13pub struct DateTimeTrunc;
14
15impl DateTimeTrunc {
16 pub fn new() -> Self {
17 Self
18 }
19}
20
21impl ScalarFunction for DateTimeTrunc {
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 dt_col = columns.get(0).unwrap();
38 let prec_col = columns.get(1).unwrap();
39
40 match (dt_col.data(), prec_col.data()) {
41 (
42 ColumnData::DateTime(dt_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 (dt_container.get(i), prec_container.is_defined(i)) {
52 (Some(dt), true) => {
53 let precision = &prec_container[i];
54 let truncated = match precision.as_str() {
55 "year" => DateTime::new(dt.year(), 1, 1, 0, 0, 0, 0),
56 "month" => DateTime::new(
57 dt.year(),
58 dt.month(),
59 1,
60 0,
61 0,
62 0,
63 0,
64 ),
65 "day" => DateTime::new(
66 dt.year(),
67 dt.month(),
68 dt.day(),
69 0,
70 0,
71 0,
72 0,
73 ),
74 "hour" => DateTime::new(
75 dt.year(),
76 dt.month(),
77 dt.day(),
78 dt.hour(),
79 0,
80 0,
81 0,
82 ),
83 "minute" => DateTime::new(
84 dt.year(),
85 dt.month(),
86 dt.day(),
87 dt.hour(),
88 dt.minute(),
89 0,
90 0,
91 ),
92 "second" => DateTime::new(
93 dt.year(),
94 dt.month(),
95 dt.day(),
96 dt.hour(),
97 dt.minute(),
98 dt.second(),
99 0,
100 ),
101 other => {
102 return Err(
103 ScalarFunctionError::ExecutionFailed {
104 function: ctx.fragment.clone(),
105 reason: format!(
106 "invalid precision: '{}'",
107 other
108 ),
109 },
110 );
111 }
112 };
113 match truncated {
114 Some(val) => container.push(val),
115 None => container.push_default(),
116 }
117 }
118 _ => container.push_default(),
119 }
120 }
121
122 Ok(ColumnData::DateTime(container))
123 }
124 (ColumnData::DateTime(_), other) => Err(ScalarFunctionError::InvalidArgumentType {
125 function: ctx.fragment.clone(),
126 argument_index: 1,
127 expected: vec![Type::Utf8],
128 actual: other.get_type(),
129 }),
130 (other, _) => Err(ScalarFunctionError::InvalidArgumentType {
131 function: ctx.fragment.clone(),
132 argument_index: 0,
133 expected: vec![Type::DateTime],
134 actual: other.get_type(),
135 }),
136 }
137 }
138
139 fn return_type(&self, _input_types: &[Type]) -> Type {
140 Type::DateTime
141 }
142}