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