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