reifydb_routine/function/duration/
weeks.rs1use reifydb_core::value::column::{Column, columns::Columns, data::ColumnData};
5use reifydb_type::value::{container::temporal::TemporalContainer, duration::Duration, r#type::Type};
6
7use crate::function::{Function, FunctionCapability, FunctionContext, FunctionInfo, error::FunctionError};
8
9pub struct DurationWeeks {
10 info: FunctionInfo,
11}
12
13impl Default for DurationWeeks {
14 fn default() -> Self {
15 Self::new()
16 }
17}
18
19impl DurationWeeks {
20 pub fn new() -> Self {
21 Self {
22 info: FunctionInfo::new("duration::weeks"),
23 }
24 }
25}
26
27fn extract_i64(data: &ColumnData, i: usize) -> Option<i64> {
28 match data {
29 ColumnData::Int1(c) => c.get(i).map(|&v| v as i64),
30 ColumnData::Int2(c) => c.get(i).map(|&v| v as i64),
31 ColumnData::Int4(c) => c.get(i).map(|&v| v as i64),
32 ColumnData::Int8(c) => c.get(i).copied(),
33 ColumnData::Int16(c) => c.get(i).map(|&v| v as i64),
34 ColumnData::Uint1(c) => c.get(i).map(|&v| v as i64),
35 ColumnData::Uint2(c) => c.get(i).map(|&v| v as i64),
36 ColumnData::Uint4(c) => c.get(i).map(|&v| v as i64),
37 ColumnData::Uint8(c) => c.get(i).map(|&v| v as i64),
38 ColumnData::Uint16(c) => c.get(i).map(|&v| v as i64),
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 DurationWeeks {
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::Duration
68 }
69
70 fn execute(&self, ctx: &FunctionContext, args: &Columns) -> Result<Columns, FunctionError> {
71 if args.len() != 1 {
72 return Err(FunctionError::ArityMismatch {
73 function: ctx.fragment.clone(),
74 expected: 1,
75 actual: args.len(),
76 });
77 }
78
79 let column = &args[0];
80 let (data, bitvec) = column.data().unwrap_option();
81 let row_count = data.len();
82
83 if !is_integer_type(data) {
84 return Err(FunctionError::InvalidArgumentType {
85 function: ctx.fragment.clone(),
86 argument_index: 0,
87 expected: vec![
88 Type::Int1,
89 Type::Int2,
90 Type::Int4,
91 Type::Int8,
92 Type::Int16,
93 Type::Uint1,
94 Type::Uint2,
95 Type::Uint4,
96 Type::Uint8,
97 Type::Uint16,
98 ],
99 actual: data.get_type(),
100 });
101 }
102
103 let mut container = TemporalContainer::with_capacity(row_count);
104
105 for i in 0..row_count {
106 if let Some(val) = extract_i64(data, i) {
107 container.push(Duration::from_weeks(val)?);
108 } else {
109 container.push_default();
110 }
111 }
112
113 let mut result_data = ColumnData::Duration(container);
114 if let Some(bv) = bitvec {
115 result_data = ColumnData::Option {
116 inner: Box::new(result_data),
117 bitvec: bv.clone(),
118 };
119 }
120 Ok(Columns::new(vec![Column::new(ctx.fragment.clone(), result_data)]))
121 }
122}