reifydb_function/date/
trunc.rs1use reifydb_core::value::column::data::ColumnData;
5use reifydb_type::value::{container::temporal::TemporalContainer, date::Date, r#type::Type};
6
7use crate::{ScalarFunction, ScalarFunctionContext, error::ScalarFunctionError, propagate_options};
8
9pub struct DateTrunc;
10
11impl DateTrunc {
12 pub fn new() -> Self {
13 Self
14 }
15}
16
17impl ScalarFunction for DateTrunc {
18 fn scalar(&self, ctx: ScalarFunctionContext) -> crate::error::ScalarFunctionResult<ColumnData> {
19 if let Some(result) = propagate_options(self, &ctx) {
20 return result;
21 }
22
23 let columns = ctx.columns;
24 let row_count = ctx.row_count;
25
26 if columns.len() != 2 {
27 return Err(ScalarFunctionError::ArityMismatch {
28 function: ctx.fragment.clone(),
29 expected: 2,
30 actual: columns.len(),
31 });
32 }
33
34 let date_col = columns.get(0).unwrap();
35 let prec_col = columns.get(1).unwrap();
36
37 match (date_col.data(), prec_col.data()) {
38 (
39 ColumnData::Date(date_container),
40 ColumnData::Utf8 {
41 container: prec_container,
42 ..
43 },
44 ) => {
45 let mut container = TemporalContainer::with_capacity(row_count);
46
47 for i in 0..row_count {
48 match (date_container.get(i), prec_container.is_defined(i)) {
49 (Some(d), true) => {
50 let precision = &prec_container[i];
51 let truncated =
52 match precision.as_str() {
53 "year" => Date::new(d.year(), 1, 1),
54 "month" => Date::new(d.year(), d.month(), 1),
55 other => {
56 return Err(ScalarFunctionError::ExecutionFailed {
57 function: ctx.fragment.clone(),
58 reason: format!("invalid precision: '{}'", other),
59 });
60 }
61 };
62 match truncated {
63 Some(val) => container.push(val),
64 None => container.push_default(),
65 }
66 }
67 _ => container.push_default(),
68 }
69 }
70
71 Ok(ColumnData::Date(container))
72 }
73 (ColumnData::Date(_), other) => Err(ScalarFunctionError::InvalidArgumentType {
74 function: ctx.fragment.clone(),
75 argument_index: 1,
76 expected: vec![Type::Utf8],
77 actual: other.get_type(),
78 }),
79 (other, _) => Err(ScalarFunctionError::InvalidArgumentType {
80 function: ctx.fragment.clone(),
81 argument_index: 0,
82 expected: vec![Type::Date],
83 actual: other.get_type(),
84 }),
85 }
86 }
87
88 fn return_type(&self, _input_types: &[Type]) -> Type {
89 Type::Date
90 }
91}