Skip to main content

reifydb_function/math/scalar/
abs.rs

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