Skip to main content

reifydb_function/math/scalar/
lcm.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2// Copyright (c) 2025 ReifyDB
3
4use reifydb_core::value::column::data::ColumnData;
5use reifydb_type::value::r#type::Type;
6
7use crate::{
8	ScalarFunction, ScalarFunctionContext,
9	error::{ScalarFunctionError, ScalarFunctionResult},
10	propagate_options,
11};
12
13pub struct Lcm;
14
15impl Lcm {
16	pub fn new() -> Self {
17		Self
18	}
19}
20
21fn numeric_to_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		_ => None,
33	}
34}
35
36fn compute_gcd(mut a: i64, mut b: i64) -> i64 {
37	a = a.abs();
38	b = b.abs();
39	while b != 0 {
40		let t = b;
41		b = a % b;
42		a = t;
43	}
44	a
45}
46
47fn compute_lcm(a: i64, b: i64) -> i64 {
48	if a == 0 || b == 0 {
49		return 0;
50	}
51	(a.abs() / compute_gcd(a, b)) * b.abs()
52}
53
54impl ScalarFunction for Lcm {
55	fn scalar(&self, ctx: ScalarFunctionContext) -> ScalarFunctionResult<ColumnData> {
56		if let Some(result) = propagate_options(self, &ctx) {
57			return result;
58		}
59		let columns = ctx.columns;
60		let row_count = ctx.row_count;
61
62		if columns.len() != 2 {
63			return Err(ScalarFunctionError::ArityMismatch {
64				function: ctx.fragment.clone(),
65				expected: 2,
66				actual: columns.len(),
67			});
68		}
69
70		let a_col = columns.get(0).unwrap();
71		let b_col = columns.get(1).unwrap();
72
73		if !a_col.data().get_type().is_number() {
74			return Err(ScalarFunctionError::InvalidArgumentType {
75				function: ctx.fragment.clone(),
76				argument_index: 0,
77				expected: vec![
78					Type::Int1,
79					Type::Int2,
80					Type::Int4,
81					Type::Int8,
82					Type::Uint1,
83					Type::Uint2,
84					Type::Uint4,
85					Type::Uint8,
86				],
87				actual: a_col.data().get_type(),
88			});
89		}
90
91		if !b_col.data().get_type().is_number() {
92			return Err(ScalarFunctionError::InvalidArgumentType {
93				function: ctx.fragment.clone(),
94				argument_index: 1,
95				expected: vec![
96					Type::Int1,
97					Type::Int2,
98					Type::Int4,
99					Type::Int8,
100					Type::Uint1,
101					Type::Uint2,
102					Type::Uint4,
103					Type::Uint8,
104				],
105				actual: b_col.data().get_type(),
106			});
107		}
108
109		let mut result = Vec::with_capacity(row_count);
110		let mut bitvec = Vec::with_capacity(row_count);
111
112		for i in 0..row_count {
113			match (numeric_to_i64(a_col.data(), i), numeric_to_i64(b_col.data(), i)) {
114				(Some(a), Some(b)) => {
115					result.push(compute_lcm(a, b));
116					bitvec.push(true);
117				}
118				_ => {
119					result.push(0);
120					bitvec.push(false);
121				}
122			}
123		}
124
125		Ok(ColumnData::int8_with_bitvec(result, bitvec))
126	}
127
128	fn return_type(&self, _input_types: &[Type]) -> Type {
129		Type::Int8
130	}
131}