reifydb_routine/function/date/
new.rs1use reifydb_core::value::column::{ColumnWithName, buffer::ColumnBuffer, columns::Columns};
5use reifydb_type::value::{container::temporal::TemporalContainer, date::Date, r#type::Type};
6
7use crate::routine::{Function, FunctionKind, Routine, RoutineInfo, context::FunctionContext, error::RoutineError};
8
9pub struct DateNew {
10 info: RoutineInfo,
11}
12
13impl Default for DateNew {
14 fn default() -> Self {
15 Self::new()
16 }
17}
18
19impl DateNew {
20 pub fn new() -> Self {
21 Self {
22 info: RoutineInfo::new("date::new"),
23 }
24 }
25}
26
27fn extract_i32(data: &ColumnBuffer, i: usize) -> Option<i32> {
28 match data {
29 ColumnBuffer::Int1(c) => c.get(i).map(|&v| v as i32),
30 ColumnBuffer::Int2(c) => c.get(i).map(|&v| v as i32),
31 ColumnBuffer::Int4(c) => c.get(i).copied(),
32 ColumnBuffer::Int8(c) => c.get(i).map(|&v| v as i32),
33 ColumnBuffer::Int16(c) => c.get(i).map(|&v| v as i32),
34 ColumnBuffer::Uint1(c) => c.get(i).map(|&v| v as i32),
35 ColumnBuffer::Uint2(c) => c.get(i).map(|&v| v as i32),
36 ColumnBuffer::Uint4(c) => c.get(i).map(|&v| v as i32),
37 ColumnBuffer::Uint8(c) => c.get(i).map(|&v| v as i32),
38 ColumnBuffer::Uint16(c) => c.get(i).map(|&v| v as i32),
39 _ => None,
40 }
41}
42
43fn is_integer_type(data: &ColumnBuffer) -> bool {
44 matches!(
45 data,
46 ColumnBuffer::Int1(_)
47 | ColumnBuffer::Int2(_)
48 | ColumnBuffer::Int4(_)
49 | ColumnBuffer::Int8(_)
50 | ColumnBuffer::Int16(_)
51 | ColumnBuffer::Uint1(_)
52 | ColumnBuffer::Uint2(_)
53 | ColumnBuffer::Uint4(_)
54 | ColumnBuffer::Uint8(_)
55 | ColumnBuffer::Uint16(_)
56 )
57}
58
59impl<'a> Routine<FunctionContext<'a>> for DateNew {
60 fn info(&self) -> &RoutineInfo {
61 &self.info
62 }
63
64 fn return_type(&self, _input_types: &[Type]) -> Type {
65 Type::Date
66 }
67
68 fn execute(&self, ctx: &mut FunctionContext<'a>, args: &Columns) -> Result<Columns, RoutineError> {
69 if args.len() != 3 {
70 return Err(RoutineError::FunctionArityMismatch {
71 function: ctx.fragment.clone(),
72 expected: 3,
73 actual: args.len(),
74 });
75 }
76
77 let year_col = &args[0];
78 let month_col = &args[1];
79 let day_col = &args[2];
80 let (year_data, _) = year_col.unwrap_option();
81 let (month_data, _) = month_col.unwrap_option();
82 let (day_data, _) = day_col.unwrap_option();
83 let row_count = year_data.len();
84
85 if !is_integer_type(year_data) {
86 return Err(RoutineError::FunctionInvalidArgumentType {
87 function: ctx.fragment.clone(),
88 argument_index: 0,
89 expected: vec![
90 Type::Int1,
91 Type::Int2,
92 Type::Int4,
93 Type::Int8,
94 Type::Int16,
95 Type::Uint1,
96 Type::Uint2,
97 Type::Uint4,
98 Type::Uint8,
99 Type::Uint16,
100 ],
101 actual: year_data.get_type(),
102 });
103 }
104 if !is_integer_type(month_data) {
105 return Err(RoutineError::FunctionInvalidArgumentType {
106 function: ctx.fragment.clone(),
107 argument_index: 1,
108 expected: vec![
109 Type::Int1,
110 Type::Int2,
111 Type::Int4,
112 Type::Int8,
113 Type::Int16,
114 Type::Uint1,
115 Type::Uint2,
116 Type::Uint4,
117 Type::Uint8,
118 Type::Uint16,
119 ],
120 actual: month_data.get_type(),
121 });
122 }
123 if !is_integer_type(day_data) {
124 return Err(RoutineError::FunctionInvalidArgumentType {
125 function: ctx.fragment.clone(),
126 argument_index: 2,
127 expected: vec![
128 Type::Int1,
129 Type::Int2,
130 Type::Int4,
131 Type::Int8,
132 Type::Int16,
133 Type::Uint1,
134 Type::Uint2,
135 Type::Uint4,
136 Type::Uint8,
137 Type::Uint16,
138 ],
139 actual: day_data.get_type(),
140 });
141 }
142
143 let mut container = TemporalContainer::with_capacity(row_count);
144
145 for i in 0..row_count {
146 let year = extract_i32(year_data, i);
147 let month = extract_i32(month_data, i);
148 let day = extract_i32(day_data, i);
149
150 match (year, month, day) {
151 (Some(y), Some(m), Some(d)) => {
152 if m >= 1 && d >= 1 {
153 match Date::new(y, m as u32, d as u32) {
154 Some(date) => container.push(date),
155 None => container.push_default(),
156 }
157 } else {
158 container.push_default();
159 }
160 }
161 _ => container.push_default(),
162 }
163 }
164
165 Ok(Columns::new(vec![ColumnWithName::new(ctx.fragment.clone(), ColumnBuffer::Date(container))]))
166 }
167}
168
169impl Function for DateNew {
170 fn kinds(&self) -> &[FunctionKind] {
171 &[FunctionKind::Scalar]
172 }
173}