reifydb_function/duration/
months.rs1use reifydb_core::value::column::data::ColumnData;
5use reifydb_type::value::{container::temporal::TemporalContainer, duration::Duration, r#type::Type};
6
7use crate::{
8 ScalarFunction, ScalarFunctionContext,
9 error::{ScalarFunctionError, ScalarFunctionResult},
10 propagate_options,
11};
12
13pub struct DurationMonths;
14
15impl DurationMonths {
16 pub fn new() -> Self {
17 Self
18 }
19}
20
21fn extract_i64(data: &ColumnData, i: usize) -> Option<i64> {
22 match data {
23 ColumnData::Int1(c) => c.get(i).map(|&v| v as i64),
24 ColumnData::Int2(c) => c.get(i).map(|&v| v as i64),
25 ColumnData::Int4(c) => c.get(i).map(|&v| v as i64),
26 ColumnData::Int8(c) => c.get(i).copied(),
27 ColumnData::Int16(c) => c.get(i).map(|&v| v as i64),
28 ColumnData::Uint1(c) => c.get(i).map(|&v| v as i64),
29 ColumnData::Uint2(c) => c.get(i).map(|&v| v as i64),
30 ColumnData::Uint4(c) => c.get(i).map(|&v| v as i64),
31 ColumnData::Uint8(c) => c.get(i).map(|&v| v as i64),
32 ColumnData::Uint16(c) => c.get(i).map(|&v| v as i64),
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 DurationMonths {
52 fn scalar(&self, ctx: ScalarFunctionContext) -> ScalarFunctionResult<ColumnData> {
53 if let Some(result) = propagate_options(self, &ctx) {
54 return result;
55 }
56 let columns = ctx.columns;
57 let row_count = ctx.row_count;
58
59 if columns.len() != 1 {
60 return Err(ScalarFunctionError::ArityMismatch {
61 function: ctx.fragment.clone(),
62 expected: 1,
63 actual: columns.len(),
64 });
65 }
66
67 let col = columns.get(0).unwrap();
68
69 if !is_integer_type(col.data()) {
70 return Err(ScalarFunctionError::InvalidArgumentType {
71 function: ctx.fragment.clone(),
72 argument_index: 0,
73 expected: vec![
74 Type::Int1,
75 Type::Int2,
76 Type::Int4,
77 Type::Int8,
78 Type::Int16,
79 Type::Uint1,
80 Type::Uint2,
81 Type::Uint4,
82 Type::Uint8,
83 Type::Uint16,
84 ],
85 actual: col.data().get_type(),
86 });
87 }
88
89 let mut container = TemporalContainer::with_capacity(row_count);
90
91 for i in 0..row_count {
92 if let Some(val) = extract_i64(col.data(), i) {
93 container.push(Duration::from_months(val));
94 } else {
95 container.push_default();
96 }
97 }
98
99 Ok(ColumnData::Duration(container))
100 }
101
102 fn return_type(&self, _input_types: &[Type]) -> Type {
103 Type::Duration
104 }
105}