Skip to main content

reifydb_function/datetime/
format.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::{constraint::bytes::MaxBytes, container::utf8::Utf8Container, date::Date, r#type::Type};
6
7use crate::{
8	ScalarFunction, ScalarFunctionContext,
9	error::{ScalarFunctionError, ScalarFunctionResult},
10	propagate_options,
11};
12
13pub struct DateTimeFormat;
14
15impl DateTimeFormat {
16	pub fn new() -> Self {
17		Self
18	}
19}
20
21/// Compute day of year from year/month/day
22fn compute_day_of_year(year: i32, month: u32, day: u32) -> u32 {
23	let mut doy = 0u32;
24	for m in 1..month {
25		doy += Date::days_in_month(year, m);
26	}
27	doy + day
28}
29
30fn format_datetime(
31	year: i32,
32	month: u32,
33	day: u32,
34	hour: u32,
35	minute: u32,
36	second: u32,
37	nanosecond: u32,
38	fmt: &str,
39) -> Result<String, String> {
40	let mut result = String::new();
41	let mut chars = fmt.chars().peekable();
42
43	while let Some(ch) = chars.next() {
44		if ch == '%' {
45			match chars.peek() {
46				Some('Y') => {
47					chars.next();
48					result.push_str(&format!("{:04}", year));
49				}
50				Some('m') => {
51					chars.next();
52					result.push_str(&format!("{:02}", month));
53				}
54				Some('d') => {
55					chars.next();
56					result.push_str(&format!("{:02}", day));
57				}
58				Some('j') => {
59					chars.next();
60					let doy = compute_day_of_year(year, month, day);
61					result.push_str(&format!("{:03}", doy));
62				}
63				Some('H') => {
64					chars.next();
65					result.push_str(&format!("{:02}", hour));
66				}
67				Some('M') => {
68					chars.next();
69					result.push_str(&format!("{:02}", minute));
70				}
71				Some('S') => {
72					chars.next();
73					result.push_str(&format!("{:02}", second));
74				}
75				Some('f') => {
76					chars.next();
77					result.push_str(&format!("{:09}", nanosecond));
78				}
79				Some('3') => {
80					chars.next();
81					if chars.peek() == Some(&'f') {
82						chars.next();
83						result.push_str(&format!("{:03}", nanosecond / 1_000_000));
84					} else {
85						return Err(
86							"invalid format specifier: '%3' (expected '%3f')".to_string()
87						);
88					}
89				}
90				Some('6') => {
91					chars.next();
92					if chars.peek() == Some(&'f') {
93						chars.next();
94						result.push_str(&format!("{:06}", nanosecond / 1_000));
95					} else {
96						return Err(
97							"invalid format specifier: '%6' (expected '%6f')".to_string()
98						);
99					}
100				}
101				Some('%') => {
102					chars.next();
103					result.push('%');
104				}
105				Some(c) => {
106					let c = *c;
107					return Err(format!("invalid format specifier: '%{}'", c));
108				}
109				None => return Err("unexpected end of format string after '%'".to_string()),
110			}
111		} else {
112			result.push(ch);
113		}
114	}
115
116	Ok(result)
117}
118
119impl ScalarFunction for DateTimeFormat {
120	fn scalar(&self, ctx: ScalarFunctionContext) -> ScalarFunctionResult<ColumnData> {
121		if let Some(result) = propagate_options(self, &ctx) {
122			return result;
123		}
124		let columns = ctx.columns;
125		let row_count = ctx.row_count;
126
127		if columns.len() != 2 {
128			return Err(ScalarFunctionError::ArityMismatch {
129				function: ctx.fragment.clone(),
130				expected: 2,
131				actual: columns.len(),
132			});
133		}
134
135		let dt_col = columns.get(0).unwrap();
136		let fmt_col = columns.get(1).unwrap();
137
138		match (dt_col.data(), fmt_col.data()) {
139			(
140				ColumnData::DateTime(dt_container),
141				ColumnData::Utf8 {
142					container: fmt_container,
143					..
144				},
145			) => {
146				let mut result_data = Vec::with_capacity(row_count);
147
148				for i in 0..row_count {
149					match (dt_container.get(i), fmt_container.is_defined(i)) {
150						(Some(dt), true) => {
151							let fmt_str = &fmt_container[i];
152							match format_datetime(
153								dt.year(),
154								dt.month(),
155								dt.day(),
156								dt.hour(),
157								dt.minute(),
158								dt.second(),
159								dt.nanosecond(),
160								fmt_str,
161							) {
162								Ok(formatted) => {
163									result_data.push(formatted);
164								}
165								Err(reason) => {
166									return Err(
167										ScalarFunctionError::ExecutionFailed {
168											function: ctx.fragment.clone(),
169											reason,
170										},
171									);
172								}
173							}
174						}
175						_ => {
176							result_data.push(String::new());
177						}
178					}
179				}
180
181				Ok(ColumnData::Utf8 {
182					container: Utf8Container::new(result_data),
183					max_bytes: MaxBytes::MAX,
184				})
185			}
186			(ColumnData::DateTime(_), other) => Err(ScalarFunctionError::InvalidArgumentType {
187				function: ctx.fragment.clone(),
188				argument_index: 1,
189				expected: vec![Type::Utf8],
190				actual: other.get_type(),
191			}),
192			(other, _) => Err(ScalarFunctionError::InvalidArgumentType {
193				function: ctx.fragment.clone(),
194				argument_index: 0,
195				expected: vec![Type::DateTime],
196				actual: other.get_type(),
197			}),
198		}
199	}
200
201	fn return_type(&self, _input_types: &[Type]) -> Type {
202		Type::Utf8
203	}
204}