reifydb_routine/function/duration/
scale.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use reifydb_core::value::column::{ColumnWithName, buffer::ColumnBuffer, columns::Columns};
5use reifydb_type::value::{container::temporal::TemporalContainer, r#type::Type};
6
7use crate::routine::{Function, FunctionKind, Routine, RoutineInfo, context::FunctionContext, error::RoutineError};
8
9pub struct DurationScale {
10	info: RoutineInfo,
11}
12
13impl Default for DurationScale {
14	fn default() -> Self {
15		Self::new()
16	}
17}
18
19impl DurationScale {
20	pub fn new() -> Self {
21		Self {
22			info: RoutineInfo::new("duration::scale"),
23		}
24	}
25}
26
27fn extract_i64(data: &ColumnBuffer, i: usize) -> Option<i64> {
28	match data {
29		ColumnBuffer::Int1(c) => c.get(i).map(|&v| v as i64),
30		ColumnBuffer::Int2(c) => c.get(i).map(|&v| v as i64),
31		ColumnBuffer::Int4(c) => c.get(i).map(|&v| v as i64),
32		ColumnBuffer::Int8(c) => c.get(i).copied(),
33		ColumnBuffer::Int16(c) => c.get(i).map(|&v| v as i64),
34		ColumnBuffer::Uint1(c) => c.get(i).map(|&v| v as i64),
35		ColumnBuffer::Uint2(c) => c.get(i).map(|&v| v as i64),
36		ColumnBuffer::Uint4(c) => c.get(i).map(|&v| v as i64),
37		ColumnBuffer::Uint8(c) => c.get(i).map(|&v| v as i64),
38		ColumnBuffer::Uint16(c) => c.get(i).map(|&v| v as i64),
39		_ => None,
40	}
41}
42
43fn is_integer_type(data: &ColumnBuffer) -> bool {
44	matches!(
45		data,
46		ColumnBuffer::Int1(_)
47			| ColumnBuffer::Int2(_)
48			| ColumnBuffer::Int4(_)
49			| ColumnBuffer::Int8(_)
50			| ColumnBuffer::Int16(_)
51			| ColumnBuffer::Uint1(_)
52			| ColumnBuffer::Uint2(_)
53			| ColumnBuffer::Uint4(_)
54			| ColumnBuffer::Uint8(_)
55			| ColumnBuffer::Uint16(_)
56	)
57}
58
59impl<'a> Routine<FunctionContext<'a>> for DurationScale {
60	fn info(&self) -> &RoutineInfo {
61		&self.info
62	}
63
64	fn return_type(&self, _input_types: &[Type]) -> Type {
65		Type::Duration
66	}
67
68	fn execute(&self, ctx: &mut FunctionContext<'a>, args: &Columns) -> Result<Columns, RoutineError> {
69		if args.len() != 2 {
70			return Err(RoutineError::FunctionArityMismatch {
71				function: ctx.fragment.clone(),
72				expected: 2,
73				actual: args.len(),
74			});
75		}
76
77		let dur_col = &args[0];
78		let scalar_col = &args[1];
79
80		let (dur_data, dur_bv) = dur_col.unwrap_option();
81		let (scalar_data, scalar_bv) = scalar_col.unwrap_option();
82
83		match dur_data {
84			ColumnBuffer::Duration(dur_container) => {
85				if !is_integer_type(scalar_data) {
86					return Err(RoutineError::FunctionInvalidArgumentType {
87						function: ctx.fragment.clone(),
88						argument_index: 1,
89						expected: vec![
90							Type::Int1,
91							Type::Int2,
92							Type::Int4,
93							Type::Int8,
94							Type::Int16,
95							Type::Uint1,
96							Type::Uint2,
97							Type::Uint4,
98							Type::Uint8,
99							Type::Uint16,
100						],
101						actual: scalar_data.get_type(),
102					});
103				}
104
105				let row_count = dur_data.len();
106				let mut container = TemporalContainer::with_capacity(row_count);
107
108				for i in 0..row_count {
109					match (dur_container.get(i), extract_i64(scalar_data, i)) {
110						(Some(dur), Some(scalar)) => {
111							container.push(*dur * scalar);
112						}
113						_ => container.push_default(),
114					}
115				}
116
117				let mut result_data = ColumnBuffer::Duration(container);
118				if let Some(bv) = dur_bv {
119					result_data = ColumnBuffer::Option {
120						inner: Box::new(result_data),
121						bitvec: bv.clone(),
122					};
123				} else if let Some(bv) = scalar_bv {
124					result_data = ColumnBuffer::Option {
125						inner: Box::new(result_data),
126						bitvec: bv.clone(),
127					};
128				}
129				Ok(Columns::new(vec![ColumnWithName::new(ctx.fragment.clone(), result_data)]))
130			}
131			other => Err(RoutineError::FunctionInvalidArgumentType {
132				function: ctx.fragment.clone(),
133				argument_index: 0,
134				expected: vec![Type::Duration],
135				actual: other.get_type(),
136			}),
137		}
138	}
139}
140
141impl Function for DurationScale {
142	fn kinds(&self) -> &[FunctionKind] {
143		&[FunctionKind::Scalar]
144	}
145}
reifydb_routine/function/duration/scale.rs

reifydb_routine/function/duration/
scale.rs
