Skip to main content

reifydb_routine/function/math/
abs.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2025 ReifyDB
3
4use num_traits::sign::Signed;
5use reifydb_core::value::column::{ColumnWithName, buffer::ColumnBuffer, columns::Columns};
6use reifydb_type::value::{container::number::NumberContainer, decimal::Decimal, int::Int, r#type::Type, uint::Uint};
7
8use crate::routine::{Function, FunctionKind, Routine, RoutineInfo, context::FunctionContext, error::RoutineError};
9
10pub struct Abs {
11	info: RoutineInfo,
12}
13
14impl Default for Abs {
15	fn default() -> Self {
16		Self::new()
17	}
18}
19
20impl Abs {
21	pub fn new() -> Self {
22		Self {
23			info: RoutineInfo::new("math::abs"),
24		}
25	}
26}
27
28impl<'a> Routine<FunctionContext<'a>> for Abs {
29	fn info(&self) -> &RoutineInfo {
30		&self.info
31	}
32
33	fn return_type(&self, input_types: &[Type]) -> Type {
34		input_types.first().cloned().unwrap_or(Type::Float8)
35	}
36
37	fn execute(&self, ctx: &mut FunctionContext<'a>, args: &Columns) -> Result<Columns, RoutineError> {
38		if args.len() != 1 {
39			return Err(RoutineError::FunctionArityMismatch {
40				function: ctx.fragment.clone(),
41				expected: 1,
42				actual: args.len(),
43			});
44		}
45
46		let column = &args[0];
47		let (data, bitvec) = column.unwrap_option();
48		let row_count = data.len();
49
50		let result_data = match data {
51			ColumnBuffer::Int1(container) => {
52				let mut data = Vec::with_capacity(row_count);
53				let mut res_bitvec = Vec::with_capacity(row_count);
54				for i in 0..row_count {
55					if let Some(&value) = container.get(i) {
56						data.push(value.abs());
57						res_bitvec.push(true);
58					} else {
59						data.push(0);
60						res_bitvec.push(false);
61					}
62				}
63				ColumnBuffer::int1_with_bitvec(data, res_bitvec)
64			}
65			ColumnBuffer::Int2(container) => {
66				let mut data = Vec::with_capacity(row_count);
67				let mut res_bitvec = Vec::with_capacity(row_count);
68				for i in 0..row_count {
69					if let Some(&value) = container.get(i) {
70						data.push(value.abs());
71						res_bitvec.push(true);
72					} else {
73						data.push(0);
74						res_bitvec.push(false);
75					}
76				}
77				ColumnBuffer::int2_with_bitvec(data, res_bitvec)
78			}
79			ColumnBuffer::Int4(container) => {
80				let mut data = Vec::with_capacity(row_count);
81				let mut res_bitvec = Vec::with_capacity(row_count);
82				for i in 0..row_count {
83					if let Some(&value) = container.get(i) {
84						data.push(value.abs());
85						res_bitvec.push(true);
86					} else {
87						data.push(0);
88						res_bitvec.push(false);
89					}
90				}
91				ColumnBuffer::int4_with_bitvec(data, res_bitvec)
92			}
93			ColumnBuffer::Int8(container) => {
94				let mut data = Vec::with_capacity(row_count);
95				let mut res_bitvec = Vec::with_capacity(row_count);
96				for i in 0..row_count {
97					if let Some(&value) = container.get(i) {
98						data.push(value.abs());
99						res_bitvec.push(true);
100					} else {
101						data.push(0);
102						res_bitvec.push(false);
103					}
104				}
105				ColumnBuffer::int8_with_bitvec(data, res_bitvec)
106			}
107			ColumnBuffer::Int16(container) => {
108				let mut data = Vec::with_capacity(row_count);
109				let mut res_bitvec = Vec::with_capacity(row_count);
110				for i in 0..row_count {
111					if let Some(&value) = container.get(i) {
112						data.push(value.abs());
113						res_bitvec.push(true);
114					} else {
115						data.push(0);
116						res_bitvec.push(false);
117					}
118				}
119				ColumnBuffer::int16_with_bitvec(data, res_bitvec)
120			}
121			ColumnBuffer::Uint1(container) => {
122				let mut data = Vec::with_capacity(row_count);
123				let mut res_bitvec = Vec::with_capacity(row_count);
124				for i in 0..row_count {
125					if let Some(&value) = container.get(i) {
126						data.push(value);
127						res_bitvec.push(true);
128					} else {
129						data.push(0);
130						res_bitvec.push(false);
131					}
132				}
133				ColumnBuffer::uint1_with_bitvec(data, res_bitvec)
134			}
135			ColumnBuffer::Uint2(container) => {
136				let mut data = Vec::with_capacity(row_count);
137				let mut res_bitvec = Vec::with_capacity(row_count);
138				for i in 0..row_count {
139					if let Some(&value) = container.get(i) {
140						data.push(value);
141						res_bitvec.push(true);
142					} else {
143						data.push(0);
144						res_bitvec.push(false);
145					}
146				}
147				ColumnBuffer::uint2_with_bitvec(data, res_bitvec)
148			}
149			ColumnBuffer::Uint4(container) => {
150				let mut data = Vec::with_capacity(row_count);
151				let mut res_bitvec = Vec::with_capacity(row_count);
152				for i in 0..row_count {
153					if let Some(&value) = container.get(i) {
154						data.push(value);
155						res_bitvec.push(true);
156					} else {
157						data.push(0);
158						res_bitvec.push(false);
159					}
160				}
161				ColumnBuffer::uint4_with_bitvec(data, res_bitvec)
162			}
163			ColumnBuffer::Uint8(container) => {
164				let mut data = Vec::with_capacity(row_count);
165				let mut res_bitvec = Vec::with_capacity(row_count);
166				for i in 0..row_count {
167					if let Some(&value) = container.get(i) {
168						data.push(value);
169						res_bitvec.push(true);
170					} else {
171						data.push(0);
172						res_bitvec.push(false);
173					}
174				}
175				ColumnBuffer::uint8_with_bitvec(data, res_bitvec)
176			}
177			ColumnBuffer::Uint16(container) => {
178				let mut data = Vec::with_capacity(row_count);
179				let mut res_bitvec = Vec::with_capacity(row_count);
180				for i in 0..row_count {
181					if let Some(&value) = container.get(i) {
182						data.push(value);
183						res_bitvec.push(true);
184					} else {
185						data.push(0);
186						res_bitvec.push(false);
187					}
188				}
189				ColumnBuffer::uint16_with_bitvec(data, res_bitvec)
190			}
191			ColumnBuffer::Float4(container) => {
192				let mut data = Vec::with_capacity(row_count);
193				let mut res_bitvec = Vec::with_capacity(row_count);
194				for i in 0..row_count {
195					if let Some(&value) = container.get(i) {
196						data.push(value.abs());
197						res_bitvec.push(true);
198					} else {
199						data.push(0.0);
200						res_bitvec.push(false);
201					}
202				}
203				ColumnBuffer::float4_with_bitvec(data, res_bitvec)
204			}
205			ColumnBuffer::Float8(container) => {
206				let mut data = Vec::with_capacity(row_count);
207				let mut res_bitvec = Vec::with_capacity(row_count);
208				for i in 0..row_count {
209					if let Some(&value) = container.get(i) {
210						data.push(value.abs());
211						res_bitvec.push(true);
212					} else {
213						data.push(0.0);
214						res_bitvec.push(false);
215					}
216				}
217				ColumnBuffer::float8_with_bitvec(data, res_bitvec)
218			}
219			ColumnBuffer::Int {
220				container,
221				max_bytes,
222			} => {
223				let mut data = Vec::with_capacity(row_count);
224				for i in 0..row_count {
225					if let Some(value) = container.get(i) {
226						data.push(Int::from(value.0.clone().abs()));
227					} else {
228						data.push(Int::default());
229					}
230				}
231				ColumnBuffer::Int {
232					container: NumberContainer::new(data),
233					max_bytes: *max_bytes,
234				}
235			}
236			ColumnBuffer::Uint {
237				container,
238				max_bytes,
239			} => {
240				let mut data = Vec::with_capacity(row_count);
241				for i in 0..row_count {
242					if let Some(value) = container.get(i) {
243						data.push(value.clone());
244					} else {
245						data.push(Uint::default());
246					}
247				}
248				ColumnBuffer::Uint {
249					container: NumberContainer::new(data),
250					max_bytes: *max_bytes,
251				}
252			}
253			ColumnBuffer::Decimal {
254				container,
255				precision,
256				scale,
257			} => {
258				let mut data = Vec::with_capacity(row_count);
259				for i in 0..row_count {
260					if let Some(value) = container.get(i) {
261						data.push(Decimal::from(value.0.clone().abs()));
262					} else {
263						data.push(Decimal::default());
264					}
265				}
266				ColumnBuffer::Decimal {
267					container: NumberContainer::new(data),
268					precision: *precision,
269					scale: *scale,
270				}
271			}
272			other => {
273				return Err(RoutineError::FunctionInvalidArgumentType {
274					function: ctx.fragment.clone(),
275					argument_index: 0,
276					expected: vec![
277						Type::Int1,
278						Type::Int2,
279						Type::Int4,
280						Type::Int8,
281						Type::Int16,
282						Type::Uint1,
283						Type::Uint2,
284						Type::Uint4,
285						Type::Uint8,
286						Type::Uint16,
287						Type::Float4,
288						Type::Float8,
289						Type::Int,
290						Type::Uint,
291						Type::Decimal,
292					],
293					actual: other.get_type(),
294				});
295			}
296		};
297
298		let final_data = if let Some(bv) = bitvec {
299			ColumnBuffer::Option {
300				inner: Box::new(result_data),
301				bitvec: bv.clone(),
302			}
303		} else {
304			result_data
305		};
306
307		Ok(Columns::new(vec![ColumnWithName::new(ctx.fragment.clone(), final_data)]))
308	}
309}
310
311impl Function for Abs {
312	fn kinds(&self) -> &[FunctionKind] {
313		&[FunctionKind::Scalar]
314	}
315}