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