reifydb_function/duration/
scale.rs1use reifydb_core::value::column::data::ColumnData;
5use reifydb_type::value::{container::temporal::TemporalContainer, r#type::Type};
6
7use crate::{
8 ScalarFunction, ScalarFunctionContext,
9 error::{ScalarFunctionError, ScalarFunctionResult},
10 propagate_options,
11};
12
13pub struct DurationScale;
14
15impl DurationScale {
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 DurationScale {
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() != 2 {
60 return Err(ScalarFunctionError::ArityMismatch {
61 function: ctx.fragment.clone(),
62 expected: 2,
63 actual: columns.len(),
64 });
65 }
66
67 let dur_col = columns.get(0).unwrap();
68 let scalar_col = columns.get(1).unwrap();
69
70 match dur_col.data() {
71 ColumnData::Duration(dur_container) => {
72 if !is_integer_type(scalar_col.data()) {
73 return Err(ScalarFunctionError::InvalidArgumentType {
74 function: ctx.fragment.clone(),
75 argument_index: 1,
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: scalar_col.data().get_type(),
89 });
90 }
91
92 let mut container = TemporalContainer::with_capacity(row_count);
93
94 for i in 0..row_count {
95 match (dur_container.get(i), extract_i64(scalar_col.data(), i)) {
96 (Some(dur), Some(scalar)) => {
97 container.push(*dur * scalar);
98 }
99 _ => container.push_default(),
100 }
101 }
102
103 Ok(ColumnData::Duration(container))
104 }
105 other => Err(ScalarFunctionError::InvalidArgumentType {
106 function: ctx.fragment.clone(),
107 argument_index: 0,
108 expected: vec![Type::Duration],
109 actual: other.get_type(),
110 }),
111 }
112 }
113
114 fn return_type(&self, _input_types: &[Type]) -> Type {
115 Type::Duration
116 }
117}