Skip to main content

reifydb_function/date/
new.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2// Copyright (c) 2025 ReifyDB
3
4use reifydb_core::value::column::data::ColumnData;
5use reifydb_type::value::{container::temporal::TemporalContainer, date::Date, r#type::Type};
6
7use crate::{
8	ScalarFunction, ScalarFunctionContext,
9	error::{ScalarFunctionError, ScalarFunctionResult},
10	propagate_options,
11};
12
13pub struct DateNew;
14
15impl DateNew {
16	pub fn new() -> Self {
17		Self
18	}
19}
20
21fn extract_i32(data: &ColumnData, i: usize) -> Option<i32> {
22	match data {
23		ColumnData::Int1(c) => c.get(i).map(|&v| v as i32),
24		ColumnData::Int2(c) => c.get(i).map(|&v| v as i32),
25		ColumnData::Int4(c) => c.get(i).copied(),
26		ColumnData::Int8(c) => c.get(i).map(|&v| v as i32),
27		ColumnData::Int16(c) => c.get(i).map(|&v| v as i32),
28		ColumnData::Uint1(c) => c.get(i).map(|&v| v as i32),
29		ColumnData::Uint2(c) => c.get(i).map(|&v| v as i32),
30		ColumnData::Uint4(c) => c.get(i).map(|&v| v as i32),
31		ColumnData::Uint8(c) => c.get(i).map(|&v| v as i32),
32		ColumnData::Uint16(c) => c.get(i).map(|&v| v as i32),
33		_ => None,
34	}
35}
36
37fn is_integer_type(data: &ColumnData) -> bool {
38	matches!(
39		data,
40		ColumnData::Int1(_)
41			| ColumnData::Int2(_) | ColumnData::Int4(_)
42			| ColumnData::Int8(_) | ColumnData::Int16(_)
43			| ColumnData::Uint1(_)
44			| ColumnData::Uint2(_)
45			| ColumnData::Uint4(_)
46			| ColumnData::Uint8(_)
47			| ColumnData::Uint16(_)
48	)
49}
50
51impl ScalarFunction for DateNew {
52	fn scalar(&self, ctx: ScalarFunctionContext) -> ScalarFunctionResult<ColumnData> {
53		if let Some(result) = propagate_options(self, &ctx) {
54			return result;
55		}
56
57		let columns = ctx.columns;
58		let row_count = ctx.row_count;
59
60		if columns.len() != 3 {
61			return Err(ScalarFunctionError::ArityMismatch {
62				function: ctx.fragment.clone(),
63				expected: 3,
64				actual: columns.len(),
65			});
66		}
67
68		let year_col = columns.get(0).unwrap();
69		let month_col = columns.get(1).unwrap();
70		let day_col = columns.get(2).unwrap();
71
72		if !is_integer_type(year_col.data()) {
73			return Err(ScalarFunctionError::InvalidArgumentType {
74				function: ctx.fragment.clone(),
75				argument_index: 0,
76				expected: vec![
77					Type::Int1,
78					Type::Int2,
79					Type::Int4,
80					Type::Int8,
81					Type::Int16,
82					Type::Uint1,
83					Type::Uint2,
84					Type::Uint4,
85					Type::Uint8,
86					Type::Uint16,
87				],
88				actual: year_col.data().get_type(),
89			});
90		}
91		if !is_integer_type(month_col.data()) {
92			return Err(ScalarFunctionError::InvalidArgumentType {
93				function: ctx.fragment.clone(),
94				argument_index: 1,
95				expected: vec![
96					Type::Int1,
97					Type::Int2,
98					Type::Int4,
99					Type::Int8,
100					Type::Int16,
101					Type::Uint1,
102					Type::Uint2,
103					Type::Uint4,
104					Type::Uint8,
105					Type::Uint16,
106				],
107				actual: month_col.data().get_type(),
108			});
109		}
110		if !is_integer_type(day_col.data()) {
111			return Err(ScalarFunctionError::InvalidArgumentType {
112				function: ctx.fragment.clone(),
113				argument_index: 2,
114				expected: vec![
115					Type::Int1,
116					Type::Int2,
117					Type::Int4,
118					Type::Int8,
119					Type::Int16,
120					Type::Uint1,
121					Type::Uint2,
122					Type::Uint4,
123					Type::Uint8,
124					Type::Uint16,
125				],
126				actual: day_col.data().get_type(),
127			});
128		}
129
130		let mut container = TemporalContainer::with_capacity(row_count);
131
132		for i in 0..row_count {
133			let year = extract_i32(year_col.data(), i);
134			let month = extract_i32(month_col.data(), i);
135			let day = extract_i32(day_col.data(), i);
136
137			match (year, month, day) {
138				(Some(y), Some(m), Some(d)) => {
139					if m >= 1 && d >= 1 {
140						match Date::new(y, m as u32, d as u32) {
141							Some(date) => container.push(date),
142							None => container.push_default(),
143						}
144					} else {
145						container.push_default();
146					}
147				}
148				_ => container.push_default(),
149			}
150		}
151
152		Ok(ColumnData::Date(container))
153	}
154
155	fn return_type(&self, _input_types: &[Type]) -> Type {
156		Type::Date
157	}
158}